1use crate::Engine;
6use crate::cache::Cache;
7use cudarc::driver::CudaSlice;
8use memra_gguf::config::ModelConfig;
9
10pub struct PrimeSlabs {
13 pub t_cap: usize,
14 pub h: CudaSlice<f32>,
15 pub x1: CudaSlice<f32>,
16 pub z: CudaSlice<f32>,
17 pub act: CudaSlice<f32>,
18 pub xa: CudaSlice<f32>,
19 pub xb: CudaSlice<f32>,
20 pub h16: CudaSlice<u8>,
21 pub z16: CudaSlice<u8>,
22 pub gate: CudaSlice<f32>, pub up: CudaSlice<f32>, pub ffn_out: CudaSlice<f32>, pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
32 pub mixed: CudaSlice<f32>,
36 pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
37 pub seg_t: usize,
38}
39
40unsafe impl Send for PrimeSlabs {}
43
44fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
45 std::iter::repeat_with(|| None).take(n).collect()
46}
47
48struct PrimeCacheStages<'a> {
53 parent: &'a mut Cache,
54 cut: usize,
55 stage0: Cache,
56 stage1: Cache,
57}
58
59impl<'a> PrimeCacheStages<'a> {
60 fn new(parent: &'a mut Cache, cut: usize) -> Self {
61 let n = parent.kv.len();
62 assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
63 assert!(cut <= n, "PP-2 cache cut {cut} exceeds {n} layers");
64 let mut kv0 = empty_cache_layers(n);
65 let mut kv1 = empty_cache_layers(n);
66 let mut recur0 = empty_cache_layers(n);
67 let mut recur1 = empty_cache_layers(n);
68 for i in 0..cut {
69 kv0[i] = parent.kv[i].take();
70 recur0[i] = parent.recur[i].take();
71 }
72 for i in cut..n {
73 kv1[i] = parent.kv[i].take();
74 recur1[i] = parent.recur[i].take();
75 }
76 let pos = parent.pos;
77 let max_ctx = parent.max_ctx;
78 Self {
79 parent,
80 cut,
81 stage0: Cache {
82 kv: kv0,
83 recur: recur0,
84 pos,
85 max_ctx,
86 last_logits_dev: None,
87 dflash_taps: None,
88 },
89 stage1: Cache {
90 kv: kv1,
91 recur: recur1,
92 pos,
93 max_ctx,
94 last_logits_dev: None,
95 dflash_taps: None,
96 },
97 }
98 }
99
100 fn parts(&mut self) -> (&mut Cache, &mut Cache) {
101 (&mut self.stage0, &mut self.stage1)
102 }
103}
104
105impl Drop for PrimeCacheStages<'_> {
106 fn drop(&mut self) {
107 let n = self.parent.kv.len();
108 for i in 0..n {
109 let source = if i < self.cut {
110 &mut self.stage0
111 } else {
112 &mut self.stage1
113 };
114 debug_assert!(self.parent.kv[i].is_none());
115 debug_assert!(self.parent.recur[i].is_none());
116 self.parent.kv[i] = source.kv[i].take();
117 self.parent.recur[i] = source.recur[i].take();
118 }
119 self.parent.pos = self.stage0.pos.min(self.stage1.pos);
120 }
121}
122
123pub(crate) struct AttnPre {
125 pub q: cudarc::driver::CudaSlice<f32>,
126 pub k: cudarc::driver::CudaSlice<f32>,
127 pub v: cudarc::driver::CudaSlice<f32>,
128 pub gate: Option<cudarc::driver::CudaSlice<f32>>,
129}
130
131pub(crate) struct GdnPrep {
133 pub hk: usize,
134 pub q_l2: cudarc::driver::CudaSlice<f32>,
135 pub k_l2: cudarc::driver::CudaSlice<f32>,
136 pub v_g: cudarc::driver::CudaSlice<f32>,
137 pub beta: cudarc::driver::CudaSlice<f32>,
138 pub g_log: cudarc::driver::CudaSlice<f32>,
139 pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
140 pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
141}
142
143pub(crate) struct VerifyStreamScratch {
145 pub pos_d: CudaSlice<i32>,
146 pub row_ctrs: Vec<CudaSlice<i32>>,
147}
148use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
149
150struct MoeInputTraceWriter {
151 dir: std::path::PathBuf,
152 index: std::fs::File,
153 payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
154}
155
156static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
157 std::sync::OnceLock::new();
158
159fn gdec_enabled() -> bool {
162 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *E.get_or_init(|| {
164 std::env::var("MEMRA_MOE_GDEC")
165 .map(|v| v != "0")
166 .unwrap_or(true)
167 })
168}
169
170fn moe_slab_enabled() -> bool {
181 std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
182}
183
184fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
188 std::env::var("MEMRA_MOE_GROUPED")
189 .map(|value| value != "0")
190 .unwrap_or(false)
191}
192
193fn moe_prefetch_enabled() -> bool {
196 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
197 *E.get_or_init(|| {
198 std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
199 || crate::spill_pread::worker_enabled()
200 })
201}
202
203fn moe_page_prefetch_window() -> usize {
208 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
209 *W.get_or_init(|| {
210 page_prefetch_window_from_values(
211 std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
212 std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
213 .ok()
214 .as_deref(),
215 )
216 })
217}
218
219fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
220 if !enabled {
221 return 0;
222 }
223 raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
224}
225
226fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
230 if window == 0 || position >= len {
231 return len..len;
232 }
233 let (start, count) = if position == 0 {
234 (1, window)
235 } else {
236 (position.saturating_add(window), 1)
237 };
238 let start = start.min(len);
239 start..start.saturating_add(count).min(len)
240}
241
242fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
245 let position = current.map_or(0, |position| position.saturating_add(1));
246 (position < order_len).then_some(position)
247}
248
249fn worker_prefetch_window() -> usize {
254 static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
255 *WINDOW.get_or_init(|| {
256 let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
257 std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
258 .ok()
259 .and_then(|value| value.parse::<usize>().ok())
260 .unwrap_or(automatic.max(1))
261 })
262}
263
264fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
268 if window == 0 || position >= len {
269 return len..len;
270 }
271 let (start, count) = if position == 0 {
272 (0, window)
273 } else {
274 (position.saturating_add(window).saturating_sub(1), 1)
275 };
276 let start = start.min(len);
277 start..start.saturating_add(count).min(len)
278}
279
280fn moe_dev_enabled() -> bool {
285 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
286 *E.get_or_init(|| {
287 std::env::var("MEMRA_MOE_DEV")
288 .map(|v| v != "0")
289 .unwrap_or(true)
290 && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
291 })
292}
293
294fn sigmoid_router_enabled() -> bool {
297 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
298 *E.get_or_init(|| {
299 std::env::var("MEMRA_SIG_ROUTER")
300 .map(|v| v != "0")
301 .unwrap_or(true)
302 })
303}
304
305fn moe_q8_enabled() -> bool {
310 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
311 *E.get_or_init(|| {
312 std::env::var("MEMRA_MOE_Q8")
313 .map(|v| v != "0")
314 .unwrap_or(true)
315 })
316}
317
318fn expert_dp4a_supported(qt: i32) -> bool {
321 qt == crate::QT_Q4_0
322 || qt == crate::QT_IQ3_S
323 || qt == crate::QT_IQ4_XS
324 || qt == crate::QT_Q3_K
325 || qt == crate::QT_Q4_K
326 || qt == crate::QT_Q6_K
327}
328
329fn q8_expert_supported(qt: i32) -> bool {
330 static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
336 let kq = *KQ.get_or_init(|| {
337 std::env::var("MEMRA_MOE_Q8_KQ")
338 .map(|v| v != "0")
339 .unwrap_or(true)
340 });
341 let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
348 .map(|v| v != "0")
349 .unwrap_or(true);
350 qt == crate::QT_IQ3_S
351 || qt == crate::QT_IQ4_XS
352 || (nvfp4_q8 && qt == crate::QT_NVFP4)
353 || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
354}
355
356fn q8_expert_dec_supported(qt: i32) -> bool {
359 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
360}
361
362fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
368 match qt {
369 crate::QT_Q4_0 => in_f % 32 == 0,
370 crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
371 in_f % 256 == 0
372 }
373 crate::QT_NVFP4 => in_f % 64 == 0,
378 _ => false,
379 }
380}
381
382fn moe_prewarm_enabled() -> bool {
385 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386 *E.get_or_init(|| {
387 std::env::var("MEMRA_MOE_PREWARM")
388 .map(|v| v != "0")
389 .unwrap_or(true)
390 })
391}
392
393fn cpu_expert_profile_admit_enabled() -> bool {
397 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
398 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
399}
400
401pub const PRIME_MIN_T: usize = 16;
405const PRIME_PIPE_MICROBATCHES: usize = 8;
406const PRIME_PIPE_MIN_CHUNK: usize = 128;
407const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
408const PRIME_PIPE_LINEAR_WORK: usize = 8;
409
410fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
411 crate::pp::prime_pp_on()
412 && !crate::pp::pp2_streams_off()
413 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
414}
415
416pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
420 if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
421 let parsed = value
422 .parse::<usize>()
423 .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
424 return if crate::cache::swa_ring_on() {
425 if parsed == 0 {
426 crate::cache::PRIME_CHUNK_MAX_TOKENS
427 } else {
428 parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
429 }
430 } else {
431 parsed
432 };
433 }
434 let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
435 if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
436 chunk.min(
437 t.div_ceil(PRIME_PIPE_MICROBATCHES)
438 .max(PRIME_PIPE_MIN_CHUNK),
439 )
440 } else {
441 chunk
442 }
443}
444
445fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
446 fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
447}
448
449fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
450 if chunk == 0 || t <= chunk {
451 return vec![(0, t)];
452 }
453 let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
454 let mut start = 0usize;
455 while start < t {
456 let mut end = (start + chunk).min(t);
457 if t - end > 0 && t - end < PRIME_MIN_T {
458 if ring_on {
459 let shifted = t - PRIME_MIN_T;
460 end = if shifted > start { shifted } else { t };
461 } else {
462 end = t;
463 }
464 }
465 ranges.push((start, end));
466 start = end;
467 }
468 ranges
469}
470
471fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
472 let prefix = prefix as u128;
473 prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
474}
475
476fn dynamic_prime_chunk_ranges(
477 t: usize,
478 fixed_chunk: usize,
479 fixed: &[(usize, usize)],
480) -> Vec<(usize, usize)> {
481 let n = fixed.len();
482 if n < 3 {
483 return fixed.to_vec();
484 }
485
486 let max_first = t - (n - 1) * PRIME_MIN_T;
487 let first = fixed_chunk
488 .div_ceil(2)
489 .max(PRIME_PIPE_EDGE_MIN_CHUNK)
490 .min(max_first);
491 let mut ranges = Vec::with_capacity(n);
492 ranges.push((0, first));
493
494 let first_work = prime_chunk_work(first, t);
495 let work_span = prime_chunk_work(t, t) - first_work;
496 let denominator = (n - 1) as u128;
497 let mut previous = first;
498 for boundary in 1..n - 1 {
499 let target = first_work * denominator + work_span * (boundary as u128);
500 let remaining = n - 1 - boundary;
501 let mut low = previous + PRIME_MIN_T;
502 let mut high = t - remaining * PRIME_MIN_T;
503 while low < high {
504 let mid = low + (high - low) / 2;
505 if prime_chunk_work(mid, t) * denominator >= target {
506 high = mid;
507 } else {
508 low = mid + 1;
509 }
510 }
511 ranges.push((previous, low));
512 previous = low;
513 }
514 ranges.push((previous, t));
515 ranges
516}
517
518pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
522 let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
523 let chunk = prime_chunk_tokens(t, n_layers);
524 let fixed = fixed_prime_chunk_ranges(t, chunk);
525 let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
526 Ok(value) => value == "dynamic",
527 Err(_) => true,
528 };
529 if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
530 fixed
531 } else {
532 dynamic_prime_chunk_ranges(t, chunk, &fixed)
533 }
534}
535
536impl HybridModel {
537 fn prime_trace_path() -> Option<&'static str> {
542 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
543 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
544 .as_deref()
545 }
546
547 fn prime_anatomy_on() -> bool {
553 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
554 *E.get_or_init(|| std::env::var("MEMRA_PRIME_ANATOMY").as_deref() == Ok("1"))
555 }
556
557 fn prime_anatomy_slots() -> &'static [std::sync::atomic::AtomicU64; 5] {
558 static S: [std::sync::atomic::AtomicU64; 5] = [
559 std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), ];
565 &S
566 }
567
568 pub fn forward(
570 &self,
571 e: &Engine,
572 tokens: &[u32],
573 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
574 if self.is_gemma4_e4b() {
575 return self.gemma4_e4b_forward(e, tokens, false);
576 }
577 if self.cfg.gemma4.is_some() {
578 return self.gemma4_forward(e, tokens, false);
579 }
580 let cfg = &self.cfg;
581 let n_embd = cfg.n_embd as usize;
582 let t = tokens.len();
583 let eps = cfg.rms_eps;
584 let pos: Vec<i32> = (0..t as i32).collect();
585 let pos_d = e.htod_i32(&pos)?;
586
587 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
590 let mut h = e.uninit(t * n_embd)?;
592 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
593
594 let mixed = match &layer.mixer {
595 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
596 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
597 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
598 };
599
600 let mut x1 = e.uninit(t * n_embd)?;
602 e.add(&x, &mixed, &mut x1, t * n_embd)?;
603
604 let mut z = e.uninit(t * n_embd)?;
606 e.rms_norm(
607 &x1,
608 layer.post_attn_norm.float_data(),
609 &mut z,
610 n_embd,
611 t,
612 eps,
613 )?;
614 let ffn_out = match &layer.ffn {
615 crate::hybrid::Ffn::Dense {
616 ffn_gate,
617 ffn_up,
618 ffn_down,
619 } => {
620 let n_ff = ffn_gate.out_features();
621 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
622 let up = g2.pop().unwrap();
623 let gate = g2.pop().unwrap();
624 let mut act = e.uninit(t * n_ff)?;
625 Self::ffn_act_lim(
630 e,
631 &self.cfg,
632 &gate,
633 &up,
634 1.0,
635 1.0,
636 self.cfg.clamp_shexp_at(il as u32),
637 &mut act,
638 t * n_ff,
639 )?;
640 e.matmul(ffn_down, &act, t)?
641 }
642 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
643 };
644 let mut x2 = e.uninit(t * n_embd)?;
645 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
646 x = x2;
647 }
648
649 let mut hn = e.uninit(t * n_embd)?;
650 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
651 let logits = e.matmul(&self.output, &hn, t)?;
652 Ok(e.dtoh(&logits)?)
653 }
654
655 pub fn forward_last(
661 &self,
662 e: &Engine,
663 tokens: &[u32],
664 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
665 if self.cfg.gemma4.is_some() {
666 return self.gemma4_forward(e, tokens, true);
667 }
668 let cfg = &self.cfg;
669 let n_embd = cfg.n_embd as usize;
670 let t = tokens.len();
671 let eps = cfg.rms_eps;
672 let pos: Vec<i32> = (0..t as i32).collect();
673 let pos_d = e.htod_i32(&pos)?;
674
675 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
679 let anat = Self::prime_anatomy_on();
680 let mut anat_last = if anat {
681 e.stream().synchronize()?;
682 Some(std::time::Instant::now())
683 } else {
684 None
685 };
686 macro_rules! anat_mark {
687 ($slot:expr) => {
688 if let Some(ts) = anat_last.as_mut() {
689 e.stream().synchronize()?;
690 Self::prime_anatomy_slots()[$slot].fetch_add(
691 ts.elapsed().as_nanos() as u64,
692 std::sync::atomic::Ordering::Relaxed,
693 );
694 *ts = std::time::Instant::now();
695 }
696 };
697 }
698 for (il, layer) in self.layers.iter().enumerate() {
699 let mut h = e.uninit(t * n_embd)?;
700 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
701 if probe {
702 e.stream().synchronize()?;
703 eprintln!("[probe] L{il} norm ok");
704 }
705 anat_mark!(4);
706 let mixed = match &layer.mixer {
707 Mixer::Full(fa) => {
708 let y = self.full_attn(e, fa, &h, &pos_d, t, il)?;
709 anat_mark!(0);
710 y
711 }
712 Mixer::Linear(la) => {
713 let y = self.linear_attn(e, la, &h, t)?;
714 anat_mark!(1);
715 y
716 }
717 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
718 };
719 if probe {
720 e.stream().synchronize()?;
721 eprintln!("[probe] L{il} mixer ok");
722 }
723 let mut x1 = e.uninit(t * n_embd)?;
724 e.add(&x, &mixed, &mut x1, t * n_embd)?;
725 let mut z = e.uninit(t * n_embd)?;
726 e.rms_norm(
727 &x1,
728 layer.post_attn_norm.float_data(),
729 &mut z,
730 n_embd,
731 t,
732 eps,
733 )?;
734 anat_mark!(4);
735 let ffn_out = match &layer.ffn {
736 crate::hybrid::Ffn::Dense {
737 ffn_gate,
738 ffn_up,
739 ffn_down,
740 } => {
741 let n_ff = ffn_gate.out_features();
742 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
743 let up = g2.pop().unwrap();
744 let gate = g2.pop().unwrap();
745 let mut act = e.uninit(t * n_ff)?;
746 Self::ffn_act_lim(
748 e,
749 &self.cfg,
750 &gate,
751 &up,
752 1.0,
753 1.0,
754 self.cfg.clamp_shexp_at(il as u32),
755 &mut act,
756 t * n_ff,
757 )?;
758 let y = e.matmul(ffn_down, &act, t)?;
759 anat_mark!(3);
760 y
761 }
762 crate::hybrid::Ffn::Moe(m) => {
763 let y = self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?;
764 anat_mark!(2);
765 y
766 }
767 };
768 if probe {
769 e.stream().synchronize()?;
770 eprintln!("[probe] L{il} ffn ok");
771 }
772 let mut x2 = e.uninit(t * n_embd)?;
773 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
774 x = x2;
775 }
776 if anat {
777 let s = Self::prime_anatomy_slots();
778 let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
779 eprintln!(
780 "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
781 dense={:.1} norms_adds={:.1} (t={t}, forward_last)",
782 ms(0),
783 ms(1),
784 ms(2),
785 ms(3),
786 ms(4)
787 );
788 }
789 let mut hn = e.uninit(t * n_embd)?;
791 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
792 let last = e.view(&hn, t * n_embd); let last_row = last.slice((t - 1) * n_embd..t * n_embd); let mut hlast = e.uninit(n_embd)?;
795 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
796 let logits = e.matmul(&self.output, &hlast, 1)?; Ok(e.dtoh(&logits)?)
798 }
799
800 pub fn prime_cache(
832 &self,
833 e: &Engine,
834 tokens: &[u32],
835 cache: &mut Cache,
836 queued_after: usize,
837 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
838 self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
839 }
840
841 pub fn prime_cache_overlaid(
847 &self,
848 e: &Engine,
849 tokens: &[u32],
850 cache: &mut Cache,
851 queued_after: usize,
852 overlay: Option<&crate::vision::EmbedOverlay>,
853 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
854 let n_embd = self.cfg.n_embd as usize;
855 let t = tokens.len();
856 assert!(
860 t >= PRIME_MIN_T,
861 "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
862 );
863 assert!(
864 cache.pos + t <= cache.max_ctx,
865 "prime_cache: prompt exceeds cache max_ctx"
866 );
867
868 if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
880 if self.is_gemma4_e4b() {
881 if overlay.is_some() {
882 return Err(
883 "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
884 );
885 }
886 return self.gemma4_e4b_prime(e, tokens, cache);
887 }
888 return self.gemma4_prime(e, tokens, cache, overlay);
893 }
894 let ranges = prime_chunk_ranges(t, self.layers.len());
895 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
931 let seq_end = if legacy_calllocal {
932 cache.pos + t
933 } else {
934 cache.pos + t + queued_after
935 };
936 if ranges.len() == 1 {
937 return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
938 }
939 if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
944 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
945 if overlay.is_some() {
946 return Err(
947 "vision embedding overlay + pipelined PP prime unsupported (v1); \
948 run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
949 .into(),
950 );
951 }
952 if crate::pp::pp_multi_stream_same_device() {
953 return Err(
954 "prime chunk pipeline refused with 2 stage streams on one device — \
955 that concurrent-stream placement remains quarantined by the deferred \
956 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
957 the serial split."
958 .into(),
959 );
960 }
961 return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
962 }
963 }
964 let mut hiddens = e.uninit(t * n_embd)?;
965 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
966 for &(start, end) in &ranges {
967 if let Some(taps) = cache.dflash_taps.as_mut() {
969 taps.base = start;
970 }
971 let (l, hs, x) =
972 self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
973 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
974 last = Some((l, hs));
975 }
976 let (logits, h_seed) = last.unwrap();
977 Ok((logits, h_seed, hiddens))
978 }
979
980 fn prime_cache_pp2_pipelined(
985 &self,
986 e: &Engine,
987 tokens: &[u32],
988 cache: &mut Cache,
989 seq_end: usize,
990 ranges: &[(usize, usize)],
991 fence: &[usize],
992 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
993 debug_assert_eq!(fence.len(), 3);
994 debug_assert!(ranges.len() >= 2);
995 let rt = crate::pp::PpNRt::get(e)?;
996 assert_eq!(
997 rt.n_stages(),
998 2,
999 "prime pipeline requires exactly two PP stages"
1000 );
1001 let n_embd = self.cfg.n_embd as usize;
1002 let t = tokens.len();
1003 let initial_base = cache.pos;
1004 let caller_stream = e.stream();
1005
1006 rt.fence_stages_behind(&caller_stream)?;
1011 let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
1012 rt.prepare_overlap_slots(0, max_payload)?;
1013
1014 let mut hiddens = e.uninit(t * n_embd)?;
1015 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1016 let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
1017 let (cache0, cache1) = stage_caches.parts();
1018 let (first_start, first_end) = ranges[0];
1019 let mut slot = self.prime_pp2_stage0_enqueue(
1020 e,
1021 rt,
1022 &tokens[first_start..first_end],
1023 cache0,
1024 seq_end,
1025 fence,
1026 initial_base + first_start,
1027 true,
1028 )?;
1029 cache0.pos = initial_base + first_end;
1030
1031 for (i, &(start, end)) in ranges.iter().enumerate() {
1032 let base = initial_base + start;
1033 debug_assert_eq!(
1034 cache1.pos, base,
1035 "stage 1 must drain chunks in original position order"
1036 );
1037 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
1038 let next_base = initial_base + next_start;
1039 debug_assert_eq!(
1040 cache0.pos, next_base,
1041 "stage 0 must issue chunks in original position order"
1042 );
1043 let cache0_stage = &mut *cache0;
1044 std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
1049 let stage0 = scope.spawn(move || -> Result<usize, String> {
1050 let next = self
1051 .prime_pp2_stage0_enqueue(
1052 e,
1053 rt,
1054 &tokens[next_start..next_end],
1055 cache0_stage,
1056 seq_end,
1057 fence,
1058 next_base,
1059 true,
1060 )
1061 .map_err(|err| err.to_string())?;
1062 cache0_stage.pos = initial_base + next_end;
1063 Ok(next)
1064 });
1065 let x = self.prime_pp2_stage1_enqueue(
1066 e,
1067 rt,
1068 slot,
1069 end - start,
1070 cache1,
1071 seq_end,
1072 fence,
1073 base,
1074 true,
1075 )?;
1076 let out = {
1077 rt.bind_stage(1)?;
1078 let _st1 = rt.enter(1);
1079 let e1 = rt.engine(1, e);
1080 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1081 };
1082 let next = stage0
1083 .join()
1084 .map_err(|_| "pipeprime stage-0 host walker panicked")?
1085 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1086 Ok((out, Some(next)))
1087 })?
1088 } else {
1089 let x = self.prime_pp2_stage1_enqueue(
1090 e,
1091 rt,
1092 slot,
1093 end - start,
1094 cache1,
1095 seq_end,
1096 fence,
1097 base,
1098 true,
1099 )?;
1100 let out = {
1101 rt.bind_stage(1)?;
1102 let _st1 = rt.enter(1);
1103 let e1 = rt.engine(1, e);
1104 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1105 };
1106 (out, None)
1107 };
1108
1109 rt.publish_to(1, &caller_stream)?;
1110 e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1111 last = Some((out.0, out.1));
1112 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1113
1114 if let Some(next) = next_slot {
1115 rt.fence_stages_behind(&caller_stream)?;
1120 slot = next;
1121 }
1122 }
1123
1124 debug_assert_eq!(cache0.pos, initial_base + t);
1125 debug_assert_eq!(cache1.pos, initial_base + t);
1126 let (logits, h_seed) = last.unwrap();
1127 Ok((logits, h_seed, hiddens))
1128 }
1129
1130 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1137 if Engine::gdn_db_on()
1138 && Engine::gdn_chunked_enabled()
1139 && t >= 16
1140 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1141 && num_k * 2 == num_v
1142 {
1143 num_k
1144 } else {
1145 num_v
1146 }
1147 }
1148
1149 fn f16out_on(e: &Engine, t: usize) -> bool {
1154 crate::f16_ffi::pp_f16_enabled()
1155 && t >= 16
1156 && !e.verify_exact_on()
1157 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1158 }
1159
1160 pub fn prime_slabs_get(
1168 &self,
1169 e: &Engine,
1170 t: usize,
1171 n_embd: usize,
1172 n_ff_max: usize,
1173 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1174 let mut slabs = self.prime_slabs.lock().unwrap();
1175 let dev = e.ctx().ordinal();
1176 let need_new = match slabs.get(&dev) {
1177 None => true,
1178 Some(sl) => sl.lock().unwrap().t_cap < t,
1179 };
1180 if need_new {
1181 slabs.insert(
1182 dev,
1183 std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1184 t_cap: t,
1185 h: e.uninit(t * n_embd)?,
1186 x1: e.uninit(t * n_embd)?,
1187 z: e.uninit(t * n_embd)?,
1188 act: e.uninit(t * n_ff_max)?,
1189 xa: e.uninit(t * n_embd)?,
1190 xb: e.uninit(t * n_embd)?,
1191 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1192 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1193 gate: e.uninit(t * n_ff_max)?,
1194 up: e.uninit(t * n_ff_max)?,
1195 ffn_out: e.uninit(t * n_embd)?,
1196 seg_glue: Vec::new(),
1197 mixed: e.uninit(t * n_embd)?,
1198 seg_mid: Vec::new(),
1199 seg_t: 0,
1200 })),
1201 );
1202 }
1203 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1204 }
1205
1206 fn prime_chunk(
1210 &self,
1211 e: &Engine,
1212 tokens: &[u32],
1213 cache: &mut Cache,
1214 seq_end: usize,
1215 chunk_off: usize,
1216 overlay: Option<&crate::vision::EmbedOverlay>,
1217 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1218 if crate::pp::pp_host_bounce_active()
1219 && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1220 {
1221 return Err(
1222 "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1223 has no active prime stage split and would peer-read remote weights; keep \
1224 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1225 .into(),
1226 );
1227 }
1228 if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1237 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1238 if overlay.is_some() {
1239 return Err("vision embedding overlay + PP prime unsupported (v1); \
1240 run single-device or MEMRA_PRIME_PP=0"
1241 .into());
1242 }
1243 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1244 }
1245 }
1246 if crate::pp::pp_host_bounce_active() {
1247 return Err(
1248 "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1249 refusing an unsplit remote-weight walk"
1250 .into(),
1251 );
1252 }
1253 let t = tokens.len();
1254 let base = cache.pos;
1255 debug_assert!(
1256 seq_end >= base + t,
1257 "prime_chunk: seq_end must cover this chunk"
1258 );
1259 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1260 let pos_d = e.htod_i32(&pos)?;
1261
1262 let mut x_embed = self.embed(e, tokens)?; if let Some(ov) = overlay {
1264 let n_embd = self.cfg.n_embd as usize;
1268 for &(pos, row_off, n_rows) in &ov.spans {
1269 let lo = pos.max(chunk_off);
1270 let hi = (pos + n_rows).min(chunk_off + t);
1271 if lo < hi {
1272 let src_row = row_off + (lo - pos);
1273 let view = ov
1274 .rows
1275 .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
1276 e.copy_view_into(
1277 &mut x_embed,
1278 (lo - chunk_off) * n_embd,
1279 &view,
1280 (hi - lo) * n_embd,
1281 )?;
1282 }
1283 }
1284 }
1285 let x = self.prime_layers(
1286 e,
1287 x_embed,
1288 0,
1289 self.layers.len(),
1290 &pos_d,
1291 t,
1292 base,
1293 cache,
1294 seq_end,
1295 )?;
1296 self.prime_chunk_epilogue(e, x, t, cache)
1297 }
1298
1299 #[allow(clippy::too_many_arguments)]
1315 fn prime_layers(
1316 &self,
1317 e: &Engine,
1318 x_in: CudaSlice<f32>,
1319 lo: usize,
1320 hi: usize,
1321 pos_d: &CudaSlice<i32>,
1322 t: usize,
1323 base: usize,
1324 cache: &mut Cache,
1325 seq_end: usize,
1326 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1327 let cfg = &self.cfg;
1328 let n_embd = cfg.n_embd as usize;
1329 let eps = cfg.rms_eps;
1330 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1334 let n_ff_max = self
1340 .layers
1341 .iter()
1342 .map(|l| match &l.ffn {
1343 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1344 _ => n_embd,
1345 })
1346 .max()
1347 .unwrap_or(n_embd)
1348 .max(n_embd);
1349 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1350 let slab = if use_slabs {
1351 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1352 } else {
1353 None
1354 };
1355 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1356 let mut x_own; type SlabRefs<'a> = (
1358 &'a mut CudaSlice<f32>,
1359 &'a mut CudaSlice<f32>,
1360 &'a mut CudaSlice<f32>,
1361 &'a mut CudaSlice<f32>,
1362 &'a mut CudaSlice<u8>,
1363 &'a mut CudaSlice<u8>,
1364 &'a mut CudaSlice<f32>,
1365 &'a mut CudaSlice<f32>,
1366 &'a mut CudaSlice<f32>,
1367 );
1368 let (mut x_cur, mut x_nxt, sl): (
1369 &mut CudaSlice<f32>,
1370 &mut CudaSlice<f32>,
1371 Option<SlabRefs>,
1372 );
1373 let mut seg: Option<(
1374 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1375 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1376 &mut CudaSlice<f32>,
1377 &mut usize,
1378 )> = None;
1379 let mut x_own2;
1380 match slab_guard.as_mut() {
1381 Some(g) => {
1382 let slabs = &mut **g;
1383 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1384 let PrimeSlabs {
1385 xa,
1386 xb,
1387 h,
1388 x1,
1389 z,
1390 act,
1391 h16,
1392 z16,
1393 gate,
1394 up,
1395 ffn_out,
1396 seg_glue,
1397 mixed,
1398 seg_mid,
1399 seg_t,
1400 ..
1401 } = slabs;
1402 x_cur = xa;
1403 x_nxt = xb;
1404 seg = Some((seg_glue, seg_mid, mixed, seg_t));
1405 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1406 }
1407 None => {
1408 x_own = x_in;
1409 x_own2 = e.uninit(t * n_embd)?;
1410 x_cur = &mut x_own;
1411 x_nxt = &mut x_own2;
1412 sl = None;
1413 }
1414 }
1415 let mut alloc_h;
1416 let mut alloc_x1;
1417 let mut alloc_z;
1418 let mut alloc_act;
1419 let mut alloc_h16;
1420 let mut alloc_z16;
1421 let mut alloc_gate;
1422 let mut alloc_up;
1423 let mut alloc_fo;
1424 let (h, x1, z, act): (
1425 &mut CudaSlice<f32>,
1426 &mut CudaSlice<f32>,
1427 &mut CudaSlice<f32>,
1428 &mut CudaSlice<f32>,
1429 );
1430 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1431 let (sl_gate, sl_up, sl_fo): (
1432 &mut CudaSlice<f32>,
1433 &mut CudaSlice<f32>,
1434 &mut CudaSlice<f32>,
1435 );
1436 match sl {
1437 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1438 h = a;
1439 x1 = b;
1440 z = c;
1441 act = d;
1442 h16 = e16;
1443 z16 = f16b;
1444 sl_gate = g;
1445 sl_up = u;
1446 sl_fo = fo;
1447 }
1448 None => {
1449 alloc_h = e.uninit(t * n_embd)?;
1450 alloc_x1 = e.uninit(t * n_embd)?;
1451 alloc_z = e.uninit(t * n_embd)?;
1452 alloc_act = e.uninit(t * n_ff_max)?;
1453 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1454 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1455 alloc_gate = e.uninit(t * n_ff_max)?;
1456 alloc_up = e.uninit(t * n_ff_max)?;
1457 alloc_fo = e.uninit(t * n_embd)?;
1458 h = &mut alloc_h;
1459 x1 = &mut alloc_x1;
1460 z = &mut alloc_z;
1461 act = &mut alloc_act;
1462 h16 = &mut alloc_h16;
1463 z16 = &mut alloc_z16;
1464 sl_gate = &mut alloc_gate;
1465 sl_up = &mut alloc_up;
1466 sl_fo = &mut alloc_fo;
1467 }
1468 }
1469 let n_layers = self.layers.len();
1474 let use_seg = f16fuse
1484 && seg.is_some()
1485 && self.cfg.step35.is_none()
1486 && lo == 0
1487 && hi == n_layers
1488 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1489 if let Some((sg, sm, _, st)) = seg.as_mut() {
1490 if **st != t {
1491 sg.clear();
1492 sg.extend((0..n_layers).map(|_| None));
1493 sm.clear();
1494 sm.extend((0..n_layers).map(|_| None));
1495 **st = t;
1496 }
1497 }
1498 {
1499 let layer_lo = &self.layers[lo];
1500 if f16fuse {
1501 e.rms_norm_f16out(
1502 x_cur,
1503 layer_lo.attn_norm.float_data(),
1504 h,
1505 h16,
1506 n_embd,
1507 t,
1508 eps,
1509 )?;
1510 } else {
1511 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1512 }
1513 }
1514 let anat = Self::prime_anatomy_on();
1515 let mut anat_last = if anat {
1516 e.stream().synchronize()?;
1517 Some(std::time::Instant::now())
1518 } else {
1519 None
1520 };
1521 macro_rules! anat_mark {
1523 ($slot:expr) => {
1524 if let Some(ts) = anat_last.as_mut() {
1525 e.stream().synchronize()?;
1526 Self::prime_anatomy_slots()[$slot].fetch_add(
1527 ts.elapsed().as_nanos() as u64,
1528 std::sync::atomic::Ordering::Relaxed,
1529 );
1530 *ts = std::time::Instant::now();
1531 }
1532 };
1533 }
1534 for il in lo..hi {
1535 let layer = &self.layers[il];
1536 let hx16 = if f16fuse { Some(&*h16) } else { None };
1537 if use_seg {
1538 let (pre, pre16, w_out) = match &layer.mixer {
1541 Mixer::Full(fa) => {
1542 let g3 = match hx16 {
1543 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1544 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1545 };
1546 let (pre, pre16) =
1547 self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1548 (pre, pre16, &fa.wo)
1549 }
1550 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1551 Mixer::Linear(la) => {
1552 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1553 let g4 = match hx16 {
1554 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1555 None => e.matmul_group(&ws, h, t)?,
1556 };
1557 let (pre, pre16) =
1558 self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1559 (pre, pre16, &la.ssm_out)
1560 }
1561 };
1562 {
1563 let (_, sm, mslab, _) = seg.as_mut().unwrap();
1564 let pre_n = pre.len() / t;
1565 let xh_pre = match pre16 {
1566 Some(x) => x,
1567 None => e.f16_act(&pre, t * pre_n, pre_n)?,
1568 };
1569 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1570 let y = e.matmul(w_out, &pre, t)?;
1571 e.copy_into(mslab, 0, &y, t * n_embd)?;
1572 }
1573 if sm[il].is_none() {
1574 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1575 let w_post = layer.post_attn_norm.float_data();
1576 e.stream().synchronize()?;
1577 e.stream()
1578 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1579 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1580 e.add(x_cur, mslab, x1, t * n_embd)?;
1581 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1582 Ok(())
1583 })();
1584 let g = e.stream().end_capture(
1585 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1586 r?;
1587 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1588 }
1589 sm[il].as_ref().unwrap().launch()?;
1590 }
1591 } else {
1592 let mixed = match &layer.mixer {
1593 Mixer::Full(fa) => {
1594 let y =
1595 self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?;
1596 anat_mark!(0);
1597 y
1598 }
1599 Mixer::Linear(la) => {
1600 let y = self.linear_attn_prime(e, la, h, hx16, t, cache, il)?;
1601 anat_mark!(1);
1602 y
1603 }
1604 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1605 };
1606 if f16fuse {
1607 e.add_rms_norm_f16out(
1610 x_cur,
1611 &mixed,
1612 layer.post_attn_norm.float_data(),
1613 x1,
1614 z,
1615 z16,
1616 n_embd,
1617 t,
1618 eps,
1619 )?;
1620 } else {
1621 e.add(x_cur, &mixed, x1, t * n_embd)?;
1622 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1623 }
1624 anat_mark!(4);
1625 }
1626 let zx16 = if f16fuse { Some(&*z16) } else { None };
1627 match &layer.ffn {
1628 crate::hybrid::Ffn::Dense {
1629 ffn_gate,
1630 ffn_up,
1631 ffn_down,
1632 } => {
1633 let n_ff = ffn_gate.out_features();
1634 let mut into_ok = false;
1637 if let Some(xh) = zx16 {
1638 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1639 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1640 }
1641 if !into_ok {
1642 let mut g2 = match zx16 {
1643 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1644 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1645 };
1646 let up_y = g2.pop().unwrap();
1647 let gate_y = g2.pop().unwrap();
1648 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1649 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1650 }
1651 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1656 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1657 {
1658 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1659 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1660 Some(a16)
1661 } else {
1662 Self::ffn_act_lim(
1663 e,
1664 &self.cfg,
1665 sl_gate,
1666 sl_up,
1667 1.0,
1668 1.0,
1669 d_lim,
1670 act,
1671 t * n_ff,
1672 )?;
1673 None
1674 };
1675 let xh_act = match act16 {
1677 Some(x) => x,
1678 None => e.f16_act(act, t * n_ff, n_ff)?,
1679 };
1680 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1681 let y = e.matmul(ffn_down, &*act, t)?;
1682 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1683 }
1684 }
1685 crate::hybrid::Ffn::Moe(m) => {
1686 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1687 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1688 anat_mark!(2);
1689 }
1690 }
1691 if let (crate::hybrid::Ffn::Dense { .. }, true) = (&layer.ffn, anat) {
1692 anat_mark!(3);
1693 }
1694 if use_seg && il + 1 < hi {
1695 let w_next = self.layers[il + 1].attn_norm.float_data();
1697 let (sg, _, _, _) = seg.as_mut().unwrap();
1698 if sg[il].is_none() {
1699 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1700 e.stream().synchronize()?;
1701 e.stream()
1702 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1703 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1704 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1705 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1706 Ok(())
1707 })();
1708 let g = e.stream().end_capture(
1709 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1710 );
1711 r?;
1712 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1713 }
1714 sg[il].as_ref().unwrap().launch()?;
1715 } else {
1716 if il + 1 < hi {
1717 let w_next = self.layers[il + 1].attn_norm.float_data();
1718 if f16fuse {
1719 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1720 } else {
1721 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1722 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1723 }
1724 } else {
1725 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1726 }
1727 }
1728 anat_mark!(4);
1729 if let Some(path) = Self::prime_trace_path() {
1735 let row = (base + t - 1) as usize;
1736 let host = e.dtoh(x_nxt)?;
1737 let last = &host[(t - 1) * n_embd..t * n_embd];
1738 use std::io::Write as _;
1739 let mut f = std::fs::OpenOptions::new()
1740 .create(true)
1741 .append(true)
1742 .open(path)?;
1743 let mut h64: u64 = 0xcbf29ce484222325;
1744 for v in last {
1745 h64 ^= v.to_bits() as u64;
1746 h64 = h64.wrapping_mul(0x100000001b3);
1747 }
1748 writeln!(
1749 f,
1750 "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1751 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1752 last[0], last[1], last[2]
1753 )?;
1754 }
1755 self.dflash_tap(e, cache, il, x_nxt, t)?;
1758 std::mem::swap(&mut x_cur, &mut x_nxt);
1759 }
1760 if anat {
1761 let s = Self::prime_anatomy_slots();
1762 let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
1763 eprintln!(
1764 "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
1765 dense={:.1} norms_adds={:.1} (t={t}, layers {lo}..{hi})",
1766 ms(0),
1767 ms(1),
1768 ms(2),
1769 ms(3),
1770 ms(4)
1771 );
1772 }
1773 let mut x = e.uninit(t * n_embd)?;
1775 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1776 drop(slab_guard);
1777 Ok(x)
1778 }
1779
1780 fn prime_chunk_epilogue(
1785 &self,
1786 e: &Engine,
1787 x: CudaSlice<f32>,
1788 t: usize,
1789 cache: &mut Cache,
1790 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1791 let n_embd = self.cfg.n_embd as usize;
1792 let eps = self.cfg.rms_eps;
1793 let mut h_seed = e.uninit(n_embd)?;
1797 if !crate::spec::spec_hpost() {
1798 e.copy_view_into(
1799 &mut h_seed,
1800 0,
1801 &x.slice((t - 1) * n_embd..t * n_embd),
1802 n_embd,
1803 )?;
1804 }
1805 let mut hn = e.uninit(t * n_embd)?;
1807 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1808 if crate::spec::spec_hpost() {
1809 e.copy_view_into(
1810 &mut h_seed,
1811 0,
1812 &hn.slice((t - 1) * n_embd..t * n_embd),
1813 n_embd,
1814 )?;
1815 }
1816 let last = e.view(&hn, t * n_embd);
1817 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1818 let mut hlast = e.uninit(n_embd)?;
1819 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1820 let logits = e.matmul(&self.output, &hlast, 1)?;
1821 cache.pos += t;
1822 Ok((
1825 e.dtoh(&logits)?,
1826 h_seed,
1827 if crate::spec::spec_hpost() { hn } else { x },
1828 ))
1829 }
1830
1831 fn prime_chunk_ppn(
1855 &self,
1856 e: &Engine,
1857 tokens: &[u32],
1858 cache: &mut Cache,
1859 seq_end: usize,
1860 fence: &[usize],
1861 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1862 let rt = crate::pp::PpNRt::get(e)?;
1863 let n_st = fence.len() - 1;
1864 assert_eq!(
1865 rt.n_stages(),
1866 n_st,
1867 "PpNRt stage count {} != fence stages {n_st}",
1868 rt.n_stages()
1869 );
1870 let n_embd = self.cfg.n_embd as usize;
1871 let t = tokens.len();
1872 let base = cache.pos;
1873 debug_assert!(
1874 seq_end >= base + t,
1875 "prime_chunk_ppn: seq_end must cover this chunk"
1876 );
1877 let payload = t * n_embd;
1878 let caller_stream = e.stream();
1882 rt.fence_stages_behind(&caller_stream)?;
1883
1884 if n_st == 2 {
1885 let slot =
1886 self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1887 let x =
1888 self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1889 let out = {
1890 rt.bind_stage(1)?;
1891 let _st1 = rt.enter(1);
1892 let e1 = rt.engine(1, e);
1893 self.prime_chunk_epilogue(e1, x, t, cache)?
1894 };
1895 rt.publish_to(1, &caller_stream)?;
1896 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1897 return Ok(out);
1898 }
1899
1900 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1901
1902 let mut slot = {
1904 let _st0 = rt.enter(0);
1905 let e0 = rt.engine(0, e);
1906 let pos_d = e0.htod_i32(&pos)?;
1907 let x = self.embed(e0, tokens)?;
1908 let x =
1909 self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1910 rt.tx(0, &x, payload)?
1911 };
1913
1914 for s in 1..n_st - 1 {
1916 let _st = rt.enter(s);
1917 let es = rt.engine(s, e);
1918 let pos_d = es.htod_i32(&pos)?;
1919 let x = rt.rx(s - 1, slot, payload)?;
1920 let x = self.prime_layers(
1921 es,
1922 x,
1923 fence[s],
1924 fence[s + 1],
1925 &pos_d,
1926 t,
1927 base,
1928 cache,
1929 seq_end,
1930 )?;
1931 slot = rt.tx(s, &x, payload)?;
1932 }
1933
1934 let _stl = rt.enter(n_st - 1);
1936 let el = rt.engine(n_st - 1, e);
1937 let pos_d = el.htod_i32(&pos)?;
1938 let x = rt.rx(n_st - 2, slot, payload)?;
1939 let x = self.prime_layers(
1940 el,
1941 x,
1942 fence[n_st - 1],
1943 fence[n_st],
1944 &pos_d,
1945 t,
1946 base,
1947 cache,
1948 seq_end,
1949 )?;
1950 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1951 rt.publish_to(n_st - 1, &caller_stream)?;
1957 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1958 Ok(out)
1959 }
1960
1961 fn prime_pp2_stage0_enqueue(
1962 &self,
1963 e: &Engine,
1964 rt: &crate::pp::PpNRt,
1965 tokens: &[u32],
1966 cache: &mut Cache,
1967 seq_end: usize,
1968 fence: &[usize],
1969 base: usize,
1970 pipelined: bool,
1971 ) -> Result<usize, Box<dyn std::error::Error>> {
1972 let t = tokens.len();
1973 let n_embd = self.cfg.n_embd as usize;
1974 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1975 rt.bind_stage(0)?;
1976 let _st0 = rt.enter(0);
1977 let e0 = rt.engine(0, e);
1978 let pos_d = e0.htod_i32(&pos)?;
1979 let x = self.embed(e0, tokens)?;
1980 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1981 let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1982 if pipelined {
1983 rt.tx_pipelined(0, &x, t * n_embd)
1984 } else {
1985 rt.tx(0, &x, t * n_embd)
1986 }
1987 }
1988
1989 fn prime_pp2_stage1_enqueue(
1990 &self,
1991 e: &Engine,
1992 rt: &crate::pp::PpNRt,
1993 slot: usize,
1994 t: usize,
1995 cache: &mut Cache,
1996 seq_end: usize,
1997 fence: &[usize],
1998 base: usize,
1999 pipelined: bool,
2000 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2001 let n_embd = self.cfg.n_embd as usize;
2002 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
2003 rt.bind_stage(1)?;
2004 let _st1 = rt.enter(1);
2005 let e1 = rt.engine(1, e);
2006 let pos_d = e1.htod_i32(&pos)?;
2007 let x = rt.rx(0, slot, t * n_embd)?;
2008 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
2009 self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
2010 }
2011
2012 pub fn prime_chunk_captured(
2028 &self,
2029 e: &Engine,
2030 x_in: &CudaSlice<f32>,
2031 pos_d: &CudaSlice<i32>,
2032 t: usize,
2033 cache: &mut Cache,
2034 len_d: &CudaSlice<i32>,
2035 logits_out: &mut CudaSlice<f32>,
2036 h_seed_out: &mut CudaSlice<f32>,
2037 ) -> Result<(), Box<dyn std::error::Error>> {
2038 let cfg = &self.cfg;
2039 let n_embd = cfg.n_embd as usize;
2040 let eps = cfg.rms_eps;
2041 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
2042 let mut x = e.uninit(t * n_embd)?;
2043 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
2044 for (il, layer) in self.layers.iter().enumerate() {
2045 let mut h = e.uninit(t * n_embd)?;
2046 let mut hx16: Option<CudaSlice<u8>> = None;
2047 if f16fuse {
2048 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
2049 e.rms_norm_f16out(
2050 &x,
2051 layer.attn_norm.float_data(),
2052 &mut h,
2053 &mut b16,
2054 n_embd,
2055 t,
2056 eps,
2057 )?;
2058 hx16 = Some(b16);
2059 } else {
2060 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2061 }
2062 let mixed = match &layer.mixer {
2063 Mixer::Full(fa) => {
2067 self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
2068 }
2069 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2070 Mixer::Linear(la) => {
2071 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2072 let g4 = match hx16.as_ref() {
2073 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
2074 None => e.matmul_group(&ws, &h, t)?,
2075 };
2076 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
2077 }
2078 };
2079 let mut x1 = e.uninit(t * n_embd)?;
2080 e.add(&x, &mixed, &mut x1, t * n_embd)?;
2081 let mut z = e.uninit(t * n_embd)?;
2082 let mut zx16: Option<CudaSlice<u8>> = None;
2083 if f16fuse {
2084 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
2085 e.rms_norm_f16out(
2086 &x1,
2087 layer.post_attn_norm.float_data(),
2088 &mut z,
2089 &mut b16,
2090 n_embd,
2091 t,
2092 eps,
2093 )?;
2094 zx16 = Some(b16);
2095 } else {
2096 e.rms_norm(
2097 &x1,
2098 layer.post_attn_norm.float_data(),
2099 &mut z,
2100 n_embd,
2101 t,
2102 eps,
2103 )?;
2104 }
2105 let ffn_out = match &layer.ffn {
2106 crate::hybrid::Ffn::Dense {
2107 ffn_gate,
2108 ffn_up,
2109 ffn_down,
2110 } => {
2111 let n_ff = ffn_gate.out_features();
2112 let mut g2 = match &zx16 {
2113 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
2114 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
2115 };
2116 let up = g2.pop().unwrap();
2117 let gate = g2.pop().unwrap();
2118 let mut act = e.uninit(t * n_ff)?;
2119 Self::ffn_act_lim(
2121 e,
2122 &self.cfg,
2123 &gate,
2124 &up,
2125 1.0,
2126 1.0,
2127 self.cfg.clamp_shexp_at(il as u32),
2128 &mut act,
2129 t * n_ff,
2130 )?;
2131 e.matmul(ffn_down, &act, t)?
2132 }
2133 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
2134 };
2135 let mut x2 = e.uninit(t * n_embd)?;
2136 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
2137 x = x2;
2138 }
2139 if !crate::spec::spec_hpost() {
2141 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
2142 }
2143 let mut hn = e.uninit(t * n_embd)?;
2144 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
2145 if crate::spec::spec_hpost() {
2146 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
2147 }
2148 let mut hlast = e.uninit(n_embd)?;
2149 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
2150 let logits = e.matmul(&self.output, &hlast, 1)?;
2151 let nv = logits.len();
2152 e.copy_into(logits_out, 0, &logits, nv)?;
2153 Ok(())
2154 }
2155
2156 fn step35_prime_batch_on() -> bool {
2157 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
2158 }
2159
2160 #[allow(clippy::too_many_arguments)]
2163 fn step35_prime_batch_layers(
2164 &self,
2165 e: &Engine,
2166 mut x: CudaSlice<f32>,
2167 lo: usize,
2168 hi: usize,
2169 ts: &[usize],
2170 offs: &[usize],
2171 pos_ds: &[CudaSlice<i32>],
2172 caches: &mut [&mut Cache],
2173 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2174 let cfg = &self.cfg;
2175 let n_embd = cfg.n_embd as usize;
2176 let eps = cfg.rms_eps;
2177 let b = ts.len();
2178 let total: usize = ts.iter().sum();
2179 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
2180
2181 let split = |e: &Engine,
2182 y: &CudaSlice<f32>,
2183 dim: usize|
2184 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2185 let mut out = Vec::with_capacity(b);
2186 for s in 0..b {
2187 let mut ys = e.uninit(ts[s] * dim)?;
2188 e.copy_view_into(
2189 &mut ys,
2190 0,
2191 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2192 ts[s] * dim,
2193 )?;
2194 out.push(ys);
2195 }
2196 Ok(out)
2197 };
2198
2199 for il in lo..hi {
2200 let layer = &self.layers[il];
2201 let Mixer::Full(fa) = &layer.mixer else {
2202 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2203 };
2204
2205 let mut h = e.uninit(total * n_embd)?;
2206 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2207 if f16fuse {
2208 e.rms_norm_f16out(
2209 &x,
2210 layer.attn_norm.float_data(),
2211 &mut h,
2212 &mut hx16,
2213 n_embd,
2214 total,
2215 eps,
2216 )?;
2217 } else {
2218 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2219 }
2220
2221 let gate_w = fa
2225 .attn_gate
2226 .as_ref()
2227 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2228 let mut g4 = if f16fuse {
2229 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2230 } else {
2231 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2232 };
2233 let gate = g4.pop().unwrap();
2234 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2235 (0..b).map(|_| Vec::with_capacity(3)).collect();
2236 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2237 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2238 parts[s].push(ys);
2239 }
2240 }
2241 let gates = split(e, &gate, gate_w.out_features())?;
2242 let geometry = self.step35_geom(il);
2243 let hd = geometry.head_dim_k as usize;
2244 let nh = geometry.n_head as usize;
2245 let mut ag_cat = e.uninit(total * nh * hd)?;
2246 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2247 let ag = self.step35_attn_pre_wo(
2248 e,
2249 fa,
2250 g3s,
2251 None,
2252 Some(&gate),
2253 &pos_ds[s],
2254 ts[s],
2255 Some(&mut *caches[s]),
2256 il,
2257 ts[s],
2258 )?;
2259 e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2260 }
2261 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2262
2263 let mut x1 = e.uninit(total * n_embd)?;
2264 let mut z = e.uninit(total * n_embd)?;
2265 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2266 if f16fuse {
2267 e.add_rms_norm_f16out(
2268 &x,
2269 &mixed,
2270 layer.post_attn_norm.float_data(),
2271 &mut x1,
2272 &mut z,
2273 &mut zx16,
2274 n_embd,
2275 total,
2276 eps,
2277 )?;
2278 } else {
2279 e.add(&x, &mixed, &mut x1, total * n_embd)?;
2280 e.rms_norm(
2281 &x1,
2282 layer.post_attn_norm.float_data(),
2283 &mut z,
2284 n_embd,
2285 total,
2286 eps,
2287 )?;
2288 }
2289
2290 let ffn_out = match &layer.ffn {
2291 crate::hybrid::Ffn::Dense {
2292 ffn_gate,
2293 ffn_up,
2294 ffn_down,
2295 } => {
2296 let n_ff = ffn_gate.out_features();
2297 let mut g2 = if f16fuse {
2298 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2299 } else {
2300 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2301 };
2302 let up = g2.pop().unwrap();
2303 let gate = g2.pop().unwrap();
2304 let mut act = e.uninit(total * n_ff)?;
2305 let d_lim = cfg.clamp_shexp_at(il as u32);
2306 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2307 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2308 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2309 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2310 Some(y) => y,
2311 None => e.matmul(ffn_down, &act, total)?,
2312 }
2313 } else {
2314 Self::ffn_act_lim(
2315 e,
2316 cfg,
2317 &gate,
2318 &up,
2319 1.0,
2320 1.0,
2321 d_lim,
2322 &mut act,
2323 total * n_ff,
2324 )?;
2325 e.matmul(ffn_down, &act, total)?
2326 }
2327 }
2328 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2329 };
2330 let mut x2 = e.uninit(total * n_embd)?;
2331 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2332 x = x2;
2333 }
2334 Ok(x)
2335 }
2336
2337 fn step35_prime_batch_epilogue(
2338 &self,
2339 e: &Engine,
2340 x: CudaSlice<f32>,
2341 ts: &[usize],
2342 offs: &[usize],
2343 caches: &mut [&mut Cache],
2344 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2345 let n_embd = self.cfg.n_embd as usize;
2346 let total: usize = ts.iter().sum();
2347 let mut hn = e.uninit(total * n_embd)?;
2348 e.rms_norm(
2349 &x,
2350 self.output_norm.float_data(),
2351 &mut hn,
2352 n_embd,
2353 total,
2354 self.cfg.rms_eps,
2355 )?;
2356
2357 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2358 let mut out = Vec::with_capacity(ts.len());
2359 for s in 0..ts.len() {
2360 let mut hidden = e.uninit(ts[s] * n_embd)?;
2361 e.copy_view_into(
2362 &mut hidden,
2363 0,
2364 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2365 ts[s] * n_embd,
2366 )?;
2367 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2368 let mut h_seed = e.uninit(n_embd)?;
2369 e.copy_view_into(
2370 &mut h_seed,
2371 0,
2372 &hidden_src.slice(last0..last0 + n_embd),
2373 n_embd,
2374 )?;
2375 let mut hlast = e.uninit(n_embd)?;
2377 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2378 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2379 caches[s].pos += ts[s];
2380 out.push((logits, h_seed, hidden));
2381 }
2382 Ok(out)
2383 }
2384
2385 fn step35_prime_cache_batch(
2386 &self,
2387 e: &Engine,
2388 prompts: &[&[u32]],
2389 caches: &mut [&mut Cache],
2390 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2391 if crate::pp::pp_host_bounce_active()
2392 && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2393 {
2394 return Err(
2395 "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2396 stage split; refusing an unsplit remote-weight walk"
2397 .into(),
2398 );
2399 }
2400 if !Self::step35_prime_batch_on() {
2401 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2402 }
2403 if caches.iter().any(|c| c.pos != 0) {
2404 return Err(
2405 "step35 batched prime currently supports complete fresh prompts only; \
2406 continuation/tick chunks require per-request queued_after"
2407 .into(),
2408 );
2409 }
2410
2411 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2412 for &t in &ts {
2413 assert!(
2414 t >= PRIME_MIN_T,
2415 "step35 batched prime needs T >= {PRIME_MIN_T}"
2416 );
2417 }
2418 for (s, c) in caches.iter().enumerate() {
2419 assert!(
2420 ts[s] <= c.max_ctx,
2421 "step35 batched prime exceeds cache max_ctx"
2422 );
2423 }
2424 let offs: Vec<usize> = ts
2425 .iter()
2426 .scan(0usize, |a, &t| {
2427 let o = *a;
2428 *a += t;
2429 Some(o)
2430 })
2431 .collect();
2432 let total: usize = ts.iter().sum();
2433 let payload = total * self.cfg.n_embd as usize;
2434 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2435 let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2436 let upload_positions =
2437 |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2438 positions
2439 .iter()
2440 .map(|p| e.htod_i32(p))
2441 .collect::<Result<_, _>>()
2442 };
2443
2444 static ONCE: std::sync::Once = std::sync::Once::new();
2445 ONCE.call_once(|| {
2446 eprintln!(
2447 "[step35-prime-batch] first concat prime: B={} tokens={total}",
2448 prompts.len()
2449 );
2450 });
2451
2452 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2453 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2454 let rt = crate::pp::PpNRt::get(e)?;
2455 let n_st = fence.len() - 1;
2456 assert_eq!(
2457 rt.n_stages(),
2458 n_st,
2459 "step35 prime batch stage count mismatch"
2460 );
2461 let caller_stream = e.stream();
2462 rt.fence_stages_behind(&caller_stream)?;
2463
2464 let mut slot = {
2465 let _st0 = rt.enter(0);
2466 let e0 = rt.engine(0, e);
2467 let pos_ds = upload_positions(e0)?;
2468 let x = self.embed(e0, &cat_tokens)?;
2469 let x = self.step35_prime_batch_layers(
2470 e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2471 )?;
2472 rt.tx(0, &x, payload)?
2473 };
2474 for s in 1..n_st - 1 {
2475 let _st = rt.enter(s);
2476 let es = rt.engine(s, e);
2477 let pos_ds = upload_positions(es)?;
2478 let x = rt.rx(s - 1, slot, payload)?;
2479 let x = self.step35_prime_batch_layers(
2480 es,
2481 x,
2482 fence[s],
2483 fence[s + 1],
2484 &ts,
2485 &offs,
2486 &pos_ds,
2487 caches,
2488 )?;
2489 slot = rt.tx(s, &x, payload)?;
2490 }
2491
2492 let _stl = rt.enter(n_st - 1);
2493 let el = rt.engine(n_st - 1, e);
2494 let pos_ds = upload_positions(el)?;
2495 let x = rt.rx(n_st - 2, slot, payload)?;
2496 let x = self.step35_prime_batch_layers(
2497 el,
2498 x,
2499 fence[n_st - 1],
2500 fence[n_st],
2501 &ts,
2502 &offs,
2503 &pos_ds,
2504 caches,
2505 )?;
2506 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2507 rt.publish_to(n_st - 1, &caller_stream)?;
2508 crate::pp::STEP35_PRIME_BATCH_SPLITS
2509 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2510 out
2511 } else {
2512 let pos_ds = upload_positions(e)?;
2513 let x = self.embed(e, &cat_tokens)?;
2514 let x = self.step35_prime_batch_layers(
2515 e,
2516 x,
2517 0,
2518 self.layers.len(),
2519 &ts,
2520 &offs,
2521 &pos_ds,
2522 caches,
2523 )?;
2524 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2525 }
2526 } else {
2527 let pos_ds = upload_positions(e)?;
2528 let x = self.embed(e, &cat_tokens)?;
2529 let x = self.step35_prime_batch_layers(
2530 e,
2531 x,
2532 0,
2533 self.layers.len(),
2534 &ts,
2535 &offs,
2536 &pos_ds,
2537 caches,
2538 )?;
2539 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2540 };
2541 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2542 Ok(out)
2543 }
2544
2545 pub fn prime_cache_batch(
2562 &self,
2563 e: &Engine,
2564 prompts: &[&[u32]],
2565 caches: &mut [&mut Cache],
2566 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2567 let cfg = &self.cfg;
2568 let n_embd = cfg.n_embd as usize;
2569 let eps = cfg.rms_eps;
2570 let b = prompts.len();
2571 assert!(b >= 1 && b == caches.len());
2572 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2573 let carried = pos0s.iter().any(|&p| p > 0);
2574 if cfg.gemma4.is_some() {
2580 return Err(
2581 "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2582 swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2583 .into(),
2584 );
2585 }
2586 if cfg.step35.is_some() {
2589 return self.step35_prime_cache_batch(e, prompts, caches);
2590 }
2591 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2592 for &t in &ts {
2593 assert!(
2594 t >= PRIME_MIN_T,
2595 "prime_cache_batch needs T >= {PRIME_MIN_T}"
2596 );
2597 }
2598 for (s, c) in caches.iter().enumerate() {
2599 assert!(
2600 c.pos + ts[s] <= c.max_ctx,
2601 "prime_cache_batch: prompt exceeds cache max_ctx"
2602 );
2603 }
2604 let total: usize = ts.iter().sum();
2605 let offs: Vec<usize> = ts
2606 .iter()
2607 .scan(0usize, |a, &t| {
2608 let o = *a;
2609 *a += t;
2610 Some(o)
2611 })
2612 .collect();
2613 let pos_ds: Vec<CudaSlice<i32>> = ts
2615 .iter()
2616 .zip(&pos0s)
2617 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2618 .collect::<Result<_, _>>()?;
2619 let split = |e: &Engine,
2621 y: &CudaSlice<f32>,
2622 dim: usize|
2623 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2624 let mut out = Vec::with_capacity(b);
2625 for s in 0..b {
2626 let mut ys = e.uninit(ts[s] * dim)?;
2627 e.copy_view_into(
2628 &mut ys,
2629 0,
2630 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2631 ts[s] * dim,
2632 )?;
2633 out.push(ys);
2634 }
2635 Ok(out)
2636 };
2637
2638 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2639 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
2641 let mut h = e.uninit(total * n_embd)?;
2642 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2643 e.rms_norm_f16out(
2644 &x,
2645 layer.attn_norm.float_data(),
2646 &mut h,
2647 &mut hx16,
2648 n_embd,
2649 total,
2650 eps,
2651 )?;
2652 let mut mixed = e.uninit(total * n_embd)?;
2654 match &layer.mixer {
2655 Mixer::Full(fa) => {
2656 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2657 let geometry = self.cfg.full_attention_geometry_at(il as u32);
2663 let (n_head, n_head_kv, head_dim) = (
2664 geometry.n_head as usize,
2665 geometry.n_head_kv as usize,
2666 geometry.head_dim_k as usize,
2667 );
2668 let fa_scale = geometry.attention_scale();
2669 let use_favl = !carried
2670 && (2..=8).contains(&b)
2671 && (head_dim == 256 || head_dim == 128)
2672 && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2673 && std::env::var("MEMRA_NOFA").is_err()
2674 && std::env::var("MEMRA_FA_FLOOR").is_err()
2675 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2676 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2677 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2678 if use_favl {
2679 let (qf_w, kf_w, vf_w) = (
2680 fa.wq.out_features(),
2681 fa.wk.out_features(),
2682 fa.wv.out_features(),
2683 );
2684 memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
2689 struct APre {
2690 q: CudaSlice<f32>,
2691 gate: Option<CudaSlice<f32>>,
2692 qn: CudaSlice<f32>,
2693 kn: CudaSlice<f32>,
2694 }
2695 let mut aps = Vec::with_capacity(b);
2696 for &t in ts.iter().take(b) {
2697 aps.push(APre {
2698 q: e.uninit(t * n_head * head_dim)?,
2699 gate: Some(e.uninit(t * n_head * head_dim)?),
2700 qn: e.uninit(t * n_head * head_dim)?,
2701 kn: e.uninit(t * n_head_kv * head_dim)?,
2702 });
2703 }
2704 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2705 let kvl = caches[0].kv[il].as_ref().unwrap();
2706 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2707 };
2708 let pargs: Vec<crate::AttnPreVl> = (0..b)
2709 .map(|s| {
2710 let (o, t) = (offs[s], ts[s]);
2711 let kvl = caches[s].kv[il].as_ref().unwrap();
2712 assert!(
2713 kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2714 "prime_cache_batch attn vl: fresh + capacity"
2715 );
2716 crate::AttnPreVl {
2717 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2718 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2719 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2720 q: e.addr_f32(&aps[s].q),
2721 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2722 qn: e.addr_f32(&aps[s].qn),
2723 kn: e.addr_f32(&aps[s].kn),
2724 kc: e.addr_u8(&kvl.k),
2725 vc: e.addr_u8(&kvl.v),
2726 t: t as i32,
2727 pad: 0,
2728 }
2729 })
2730 .collect();
2731 e.attn_pre_vl8(
2732 &pargs,
2733 fa.q_norm.float_data(),
2734 fa.k_norm.float_data(),
2735 head_dim,
2736 geometry.n_rot as usize,
2737 n_head,
2738 n_head_kv,
2739 self.cfg.rms_eps,
2740 geometry.rope_base,
2741 1.0,
2742 kv_dim_k,
2743 kv_dim_v,
2744 ktb,
2745 vtb,
2746 )?;
2747 for s in 0..b {
2748 let kvl = caches[s].kv[il].as_mut().unwrap();
2749 kvl.len += ts[s];
2750 let new_len = kvl.len as i32;
2751 e.set_i32_one(&mut kvl.len_d, new_len)?;
2752 }
2753 let mut attns = Vec::with_capacity(b);
2754 let mut mirrors = Vec::with_capacity(b);
2755 for &t in ts.iter().take(b) {
2756 attns.push(e.uninit(t * n_head * head_dim)?);
2757 let n = t * n_head_kv * head_dim;
2758 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2759 }
2760 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2763 Ok("0") => false,
2764 Ok("1") => true,
2765 _ => cfg!(memra_hopper_mma),
2766 };
2767 if fa3_on {
2768 let mut q16s = Vec::with_capacity(b);
2769 let mut v16s = Vec::with_capacity(b);
2770 for s in 0..b {
2771 let t = ts[s];
2772 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2773 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2774 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2775 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2776 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2777 e.f32_to_bf16_v(
2778 &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2779 &mut v16,
2780 t * n_head_kv * head_dim,
2781 )?;
2782 q16s.push(q16);
2783 v16s.push((k16, v16));
2784 }
2785 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2786 let mut kp = qp;
2787 let mut vp = qp;
2788 let mut op = [core::ptr::null_mut::<f32>(); 8];
2789 let mut tsv = [0i32; 8];
2790 for s in 0..b {
2791 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2792 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2793 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2794 op[s] = e.addr_f32(&attns[s]) as *mut f32;
2795 tsv[s] = ts[s] as i32;
2796 }
2797 let rc = unsafe {
2798 crate::fa3_vl_raw(
2799 qp.as_ptr(),
2800 kp.as_ptr(),
2801 vp.as_ptr(),
2802 op.as_ptr(),
2803 tsv.as_ptr(),
2804 b as i32,
2805 n_head as i32,
2806 n_head_kv as i32,
2807 head_dim as i32,
2808 fa_scale,
2809 e.stream().cu_stream() as *mut core::ffi::c_void,
2810 )
2811 };
2812 if rc != 0 {
2813 return Err(format!("memra_fa3_vl rc={rc}").into());
2814 }
2815 } else {
2816 let fargs: Vec<crate::FaSeqVl> = (0..b)
2817 .map(|s| crate::FaSeqVl {
2818 q: e.addr_f32(&aps[s].qn),
2819 k16: e.addr_u8(&mirrors[s].0),
2820 v16: e.addr_u8(&mirrors[s].1),
2821 o: e.addr_f32(&attns[s]),
2822 kf: e.addr_f32(&aps[s].kn),
2823 vf: e.addr_f32v(
2824 &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2825 ),
2826 t: ts[s] as i32,
2827 pad: 0,
2828 })
2829 .collect();
2830 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2831 }
2832 for (s, attn) in attns.into_iter().enumerate() {
2833 let (attn_g, ag16) = self.full_attn_prime_post_fa(
2834 e,
2835 attn,
2836 &aps[s].gate,
2837 ts[s],
2838 n_head,
2839 head_dim,
2840 )?;
2841 let mut done = false;
2842 if let Some(xh) = &ag16 {
2843 done = e.try_f16_gemm_pre_into_off(
2844 &fa.wo,
2845 xh,
2846 ts[s],
2847 &mut mixed,
2848 offs[s] * n_embd,
2849 )?;
2850 }
2851 if !done {
2852 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2853 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2854 }
2855 }
2856 } else {
2857 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2858 (0..b).map(|_| Vec::new()).collect();
2859 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2860 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2861 parts[s].push(ys);
2862 }
2863 }
2864 for (s, g3s) in parts.into_iter().enumerate() {
2865 let (attn_g, ag16) = self.full_attn_prime_core_inner(
2867 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2868 )?;
2869 let mut done = false;
2870 if let Some(xh) = &ag16 {
2871 done = e.try_f16_gemm_pre_into_off(
2872 &fa.wo,
2873 xh,
2874 ts[s],
2875 &mut mixed,
2876 offs[s] * n_embd,
2877 )?;
2878 }
2879 if !done {
2880 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2881 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2882 }
2883 }
2884 }
2885 }
2886 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2887 Mixer::Linear(la) => {
2888 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2893 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2894 let outs =
2895 self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2896 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2897 let (o, t) = (offs[s], ts[s]);
2898 let mut done = false;
2899 if let Some(xh) = &gn16 {
2900 done = e.try_f16_gemm_pre_into_off(
2901 &la.ssm_out,
2902 xh,
2903 t,
2904 &mut mixed,
2905 o * n_embd,
2906 )?;
2907 }
2908 if !done {
2909 let m = e.matmul(&la.ssm_out, &gn, t)?;
2910 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2911 }
2912 }
2913 }
2914 }
2915 let mut x1 = e.uninit(total * n_embd)?;
2916 let mut z = e.uninit(total * n_embd)?;
2917 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2918 e.add_rms_norm_f16out(
2919 &x,
2920 &mixed,
2921 layer.post_attn_norm.float_data(),
2922 &mut x1,
2923 &mut z,
2924 &mut zx16,
2925 n_embd,
2926 total,
2927 eps,
2928 )?;
2929 let ffn_out = match &layer.ffn {
2930 crate::hybrid::Ffn::Dense {
2931 ffn_gate,
2932 ffn_up,
2933 ffn_down,
2934 } => {
2935 let n_ff = ffn_gate.out_features();
2936 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2937 let up = g2.pop().unwrap();
2938 let gate = g2.pop().unwrap();
2939 let mut act = e.uninit(total * n_ff)?;
2940 let d_lim = self.cfg.clamp_shexp_at(il as u32);
2944 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2945 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2946 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2947 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2948 Some(y) => y,
2949 None => e.matmul(ffn_down, &act, total)?,
2950 }
2951 } else {
2952 Self::ffn_act_lim(
2953 e,
2954 &self.cfg,
2955 &gate,
2956 &up,
2957 1.0,
2958 1.0,
2959 d_lim,
2960 &mut act,
2961 total * n_ff,
2962 )?;
2963 e.matmul(ffn_down, &act, total)?
2964 }
2965 }
2966 crate::hybrid::Ffn::Moe(m) => {
2967 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2968 }
2969 };
2970 let mut x2 = e.uninit(total * n_embd)?;
2971 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2972 x = x2;
2973 }
2974 let mut hn = e.uninit(total * n_embd)?;
2976 e.rms_norm(
2977 &x,
2978 self.output_norm.float_data(),
2979 &mut hn,
2980 n_embd,
2981 total,
2982 eps,
2983 )?;
2984 let mut hcat = e.uninit(b * n_embd)?;
2990 for s in 0..b {
2991 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2992 e.copy_view_into(
2993 &mut hcat,
2994 s * n_embd,
2995 &hn.slice(last0..last0 + n_embd),
2996 n_embd,
2997 )?;
2998 }
2999 let logits_cat = if b >= 2 {
3000 e.try_f16_gemm(&self.output, &hcat, b)?
3001 } else {
3002 None
3003 };
3004 let logits_host: Option<Vec<f32>> = match &logits_cat {
3005 Some(lc) => Some(e.dtoh(lc)?),
3006 None => None,
3007 };
3008 let n_vocab = self.output.out_features();
3009 let mut hidden_all = if crate::spec::spec_hpost() {
3010 split(e, &hn, n_embd)?
3011 } else {
3012 split(e, &x, n_embd)?
3013 };
3014 let mut out = Vec::with_capacity(b);
3015 for s in 0..b {
3016 let last0 = (offs[s] + ts[s] - 1) * n_embd;
3017 let mut h_seed = e.uninit(n_embd)?;
3018 if !crate::spec::spec_hpost() {
3019 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
3020 } else {
3021 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
3022 }
3023 let logits = match &logits_host {
3024 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
3025 None => {
3026 let mut hlast = e.uninit(n_embd)?;
3027 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
3028 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
3029 }
3030 };
3031 caches[s].pos += ts[s];
3032 out.push((logits, h_seed, hidden_all.remove(0)));
3033 }
3034 Ok(out)
3035 }
3036
3037 #[allow(clippy::too_many_arguments)]
3048 fn full_attn_prime(
3049 &self,
3050 e: &Engine,
3051 fa: &FullAttnLayer,
3052 h: &CudaSlice<f32>,
3053 hx: Option<&CudaSlice<u8>>,
3054 pos_d: &CudaSlice<i32>,
3055 t: usize,
3056 cache: &mut Cache,
3057 il: usize,
3058 seq_end: usize,
3059 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3060 if self.cfg.step35.is_some() {
3061 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
3062 }
3063 let g3 = match hx {
3068 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
3069 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
3070 };
3071 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
3072 }
3073
3074 fn full_attn_prime_core(
3078 &self,
3079 e: &Engine,
3080 fa: &FullAttnLayer,
3081 g3: Vec<CudaSlice<f32>>,
3082 pos_d: &CudaSlice<i32>,
3083 t: usize,
3084 cache: &mut Cache,
3085 il: usize,
3086 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3087 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
3088 if let Some(xh) = &ag16 {
3089 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
3090 return Ok(y);
3091 }
3092 }
3093 Ok(e.matmul(&fa.wo, &attn_g, t)?)
3094 }
3095
3096 fn full_attn_prime_core_inner(
3097 &self,
3098 e: &Engine,
3099 fa: &FullAttnLayer,
3100 g3: Vec<CudaSlice<f32>>,
3101 pos_d: &CudaSlice<i32>,
3102 t: usize,
3103 cache: &mut Cache,
3104 il: usize,
3105 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3106 let cfg = &self.cfg;
3107 let geometry = cfg.full_attention_geometry_at(il as u32);
3108 let n_head = geometry.n_head as usize;
3109 let n_head_kv = geometry.n_head_kv as usize;
3110 let head_dim = geometry.head_dim_k as usize;
3111 let scale = geometry.attention_scale();
3112 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
3113 let AttnPre { q, k, v, gate } = pre;
3114 let mut attn = e.uninit(t * n_head * head_dim)?;
3115 self.full_attn_prime_fa_dispatch(
3116 e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
3117 )?;
3118 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
3119 }
3120
3121 #[allow(clippy::type_complexity)]
3125 fn full_attn_prime_pre_fa(
3126 &self,
3127 e: &Engine,
3128 fa: &FullAttnLayer,
3129 mut g3: Vec<CudaSlice<f32>>,
3130 pos_d: &CudaSlice<i32>,
3131 t: usize,
3132 cache: &mut Cache,
3133 il: usize,
3134 ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
3135 let cfg = &self.cfg;
3136 let geometry = cfg.full_attention_geometry_at(il as u32);
3137 let n_head = geometry.n_head as usize;
3138 let n_head_kv = geometry.n_head_kv as usize;
3139 let head_dim = geometry.head_dim_k as usize;
3140 let eps = cfg.rms_eps;
3141
3142 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3146 let v = g3.pop().unwrap();
3147 let mut k = g3.pop().unwrap();
3148 let qf = g3.pop().unwrap();
3149 let (mut q, gate) = if gated {
3150 let mut q = e.uninit(t * n_head * head_dim)?;
3151 let mut gate = e.uninit(t * n_head * head_dim)?;
3152 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3153 (q, Some(gate))
3154 } else {
3155 (qf, None)
3156 };
3157
3158 let mut qn = e.uninit(t * n_head * head_dim)?;
3159 e.rms_norm(
3160 &q,
3161 fa.q_norm.float_data(),
3162 &mut qn,
3163 head_dim,
3164 n_head * t,
3165 eps,
3166 )?;
3167 q = qn;
3168 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3169 e.rms_norm(
3170 &k,
3171 fa.k_norm.float_data(),
3172 &mut kn,
3173 head_dim,
3174 n_head_kv * t,
3175 eps,
3176 )?;
3177 k = kn;
3178 let rope_dims = geometry.n_rot as usize;
3179 e.rope_neox(
3180 &mut q,
3181 pos_d,
3182 head_dim,
3183 rope_dims,
3184 n_head,
3185 t,
3186 geometry.rope_base,
3187 1.0,
3188 )?;
3189 e.rope_neox(
3190 &mut k,
3191 pos_d,
3192 head_dim,
3193 rope_dims,
3194 n_head_kv,
3195 t,
3196 geometry.rope_base,
3197 1.0,
3198 )?;
3199
3200 {
3203 let kvl = cache.kv[il].as_mut().unwrap();
3204 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3205 e.append_kv_quantized_rows(
3206 &k,
3207 &v,
3208 &mut kvl.k,
3209 &mut kvl.v,
3210 kvl.len,
3211 t,
3212 kvl.kv_dim_k,
3213 kvl.kv_dim_v,
3214 kvl.k_tok_bytes,
3215 kvl.v_tok_bytes,
3216 crate::Engine::kv_fp8_on(),
3217 )?;
3218 kvl.len += t;
3219 let new_len = kvl.len as i32;
3220 e.set_i32_one(&mut kvl.len_d, new_len)?;
3221 }
3222
3223 let base_len = {
3224 let kvl = cache.kv[il].as_ref().unwrap();
3225 kvl.len - t };
3227 Ok((AttnPre { q, k, v, gate }, base_len))
3228 }
3229
3230 #[allow(clippy::too_many_arguments)]
3237 fn full_attn_prime_fa_dispatch(
3238 &self,
3239 e: &Engine,
3240 q: &CudaSlice<f32>,
3241 k: &CudaSlice<f32>,
3242 v: &CudaSlice<f32>,
3243 attn: &mut CudaSlice<f32>,
3244 base_len: usize,
3245 t: usize,
3246 cache: &mut Cache,
3247 il: usize,
3248 head_dim: usize,
3249 n_head: usize,
3250 n_head_kv: usize,
3251 scale: f32,
3252 ) -> Result<(), Box<dyn std::error::Error>> {
3253 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3266 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3267 e.sdpa_naive(
3268 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3269 )?;
3270 } else {
3271 e.fa_prefill(
3272 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3273 )?;
3274 }
3275 return Ok(());
3276 }
3277 let kvl = cache.kv[il].as_ref().unwrap();
3278 let t_kv = base_len + t;
3279 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3280 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3281 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3285 e.sdpa_naive_quantized_view(
3286 q,
3287 &k_view,
3288 &v_view,
3289 attn,
3290 head_dim,
3291 n_head,
3292 n_head_kv,
3293 t,
3294 t_kv,
3295 scale,
3296 true,
3297 kvl.k_tok_bytes,
3298 kvl.v_tok_bytes,
3299 )?;
3300 return Ok(());
3301 }
3302 let deqw = std::env::var("MEMRA_PRIME_DEQW")
3310 .map(|v| v != "0")
3311 .unwrap_or(true);
3312 if deqw {
3313 e.fa_prefill_view_ws(
3314 q,
3315 &k_view,
3316 &v_view,
3317 attn,
3318 head_dim,
3319 n_head,
3320 n_head_kv,
3321 t,
3322 t_kv,
3323 scale,
3324 true,
3325 kvl.k_tok_bytes,
3326 kvl.v_tok_bytes,
3327 crate::Engine::kv_fp8_on(),
3328 )?;
3329 } else {
3330 e.fa_prefill_view(
3331 q,
3332 &k_view,
3333 &v_view,
3334 attn,
3335 head_dim,
3336 n_head,
3337 n_head_kv,
3338 t,
3339 t_kv,
3340 scale,
3341 true,
3342 kvl.k_tok_bytes,
3343 kvl.v_tok_bytes,
3344 crate::Engine::kv_fp8_on(),
3345 )?;
3346 }
3347 Ok(())
3348 }
3349
3350 fn full_attn_prime_post_fa(
3353 &self,
3354 e: &Engine,
3355 attn: CudaSlice<f32>,
3356 gate: &Option<CudaSlice<f32>>,
3357 t: usize,
3358 n_head: usize,
3359 head_dim: usize,
3360 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3361 let (attn_g, ag16) = match gate {
3362 Some(gate) => {
3363 let n = t * n_head * head_dim;
3364 let mut ag = e.uninit(n)?;
3365 if Self::f16out_on(e, t) {
3366 let mut a16 = e.alloc_u8_uninit(n * 2)?;
3367 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3368 (ag, Some(a16))
3369 } else {
3370 let mut gsig = e.uninit(n)?;
3371 e.sigmoid(gate, &mut gsig, n)?;
3372 e.mul(&attn, &gsig, &mut ag, n)?;
3373 (ag, None)
3374 }
3375 }
3376 None => (attn, None),
3377 };
3378 Ok((attn_g, ag16))
3379 }
3380
3381 fn linear_attn_prime(
3388 &self,
3389 e: &Engine,
3390 la: &LinearAttnLayer,
3391 h: &CudaSlice<f32>,
3392 hx: Option<&CudaSlice<u8>>,
3393 t: usize,
3394 cache: &mut Cache,
3395 il: usize,
3396 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3397 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3399 let g4 = match hx {
3400 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3401 None => e.matmul_group(&ws, h, t)?,
3402 };
3403 self.linear_attn_prime_core(e, la, g4, t, cache, il)
3404 }
3405
3406 fn linear_attn_prime_core(
3408 &self,
3409 e: &Engine,
3410 la: &LinearAttnLayer,
3411 mut g4: Vec<CudaSlice<f32>>,
3412 t: usize,
3413 cache: &mut Cache,
3414 il: usize,
3415 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3416 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3417 }
3418
3419 #[allow(clippy::too_many_arguments)]
3423 fn linear_attn_prime_core_pad_inner(
3424 &self,
3425 e: &Engine,
3426 la: &LinearAttnLayer,
3427 mut g4: Vec<CudaSlice<f32>>,
3428 t: usize,
3429 cache: &mut Cache,
3430 il: usize,
3431 pad_len: Option<&CudaSlice<i32>>,
3432 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3433 let ssm = self.cfg.ssm.as_ref().unwrap();
3435 let d_state = ssm.state_size as usize;
3436 let num_k = ssm.group_count as usize;
3437 let num_v = ssm.time_step_rank as usize;
3438 let key_dim = d_state * num_k;
3439 let value_dim = d_state * num_v;
3440 let conv_dim = key_dim * 2 + value_dim;
3441 let alpha = g4.pop().unwrap(); let beta_raw = g4.pop().unwrap(); let z = g4.pop().unwrap(); let qkv_mixed = g4.pop().unwrap(); self.linear_attn_prime_core_pad_view(
3446 e,
3447 la,
3448 &qkv_mixed.slice(0..t * conv_dim),
3449 &z.slice(0..t * value_dim),
3450 &beta_raw.slice(0..t * num_v),
3451 &alpha.slice(0..t * num_v),
3452 t,
3453 cache,
3454 il,
3455 pad_len,
3456 )
3457 }
3458
3459 #[allow(clippy::too_many_arguments)]
3462 fn linear_attn_gdn_prep(
3463 &self,
3464 e: &Engine,
3465 la: &LinearAttnLayer,
3466 qkv_mixed: &cudarc::driver::CudaView<f32>,
3467 beta_raw: &cudarc::driver::CudaView<f32>,
3468 alpha: &cudarc::driver::CudaView<f32>,
3469 t: usize,
3470 cache: &mut Cache,
3471 il: usize,
3472 pad_len: Option<&CudaSlice<i32>>,
3473 ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3474 let cfg = &self.cfg;
3475 let ssm = cfg.ssm.as_ref().unwrap();
3476 let d_state = ssm.state_size as usize; let num_k = ssm.group_count as usize; let num_v = ssm.time_step_rank as usize; let d_conv = ssm.conv_kernel as usize; let key_dim = d_state * num_k; let value_dim = d_state * num_v; let conv_dim = key_dim * 2 + value_dim; let eps = cfg.rms_eps;
3484 debug_assert!(
3485 t >= d_conv - 1,
3486 "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3487 );
3488
3489 let rl = cache.recur[il].as_mut().unwrap();
3494 let hk = Self::gdn_hk(e, t, num_v, num_k);
3495 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3496 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
3498 let mut k_g = e.uninit(d_state * hk * t)?;
3499 let mut v_g = e.uninit(d_state * num_v * t)?;
3500 if conv_fuse {
3501 e.ssm_conv1d_gdn_state_pad(
3502 qkv_mixed,
3503 &mut rl.conv_state,
3504 la.ssm_conv1d.float_data(),
3505 &mut q_g,
3506 &mut k_g,
3507 &mut v_g,
3508 conv_dim,
3509 t,
3510 d_conv,
3511 d_state,
3512 num_v,
3513 num_k,
3514 key_dim,
3515 hk,
3516 pad_len,
3517 )?;
3518 } else {
3519 let mut conv_out = e.uninit(conv_dim * t)?; e.ssm_conv1d_tm_state_pad_v(
3521 qkv_mixed,
3522 &mut rl.conv_state,
3523 la.ssm_conv1d.float_data(),
3524 &mut conv_out,
3525 conv_dim,
3526 t,
3527 d_conv,
3528 pad_len,
3529 )?;
3530 e.qkv_to_gdn_repack(
3531 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3532 )?;
3533 }
3534 let mut q_l2 = e.uninit(d_state * hk * t)?;
3535 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3539 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3540 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3541 Some(qb)
3542 } else {
3543 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3544 None
3545 };
3546 let mut k_l2 = e.uninit(d_state * hk * t)?;
3547 let kb16 = if Engine::l2_v2_on(d_state) {
3549 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3550 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3551 Some(kb)
3552 } else {
3553 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3554 None
3555 };
3556 let mut beta = e.uninit(t * num_v)?;
3557 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3558 let mut g_log = e.uninit(t * num_v)?;
3559 e.gdn_glog_v(
3560 alpha,
3561 la.ssm_dt.float_data(),
3562 la.ssm_a.float_data(),
3563 &mut g_log,
3564 num_v,
3565 t,
3566 )?;
3567 if let Some(len_d) = pad_len {
3568 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3569 }
3570 Ok(GdnPrep {
3571 hk,
3572 q_l2,
3573 k_l2,
3574 v_g,
3575 beta,
3576 g_log,
3577 kb16,
3578 qb16,
3579 })
3580 }
3581
3582 #[allow(clippy::too_many_arguments)]
3587 fn linear_attn_prime_core_batch(
3588 &self,
3589 e: &Engine,
3590 la: &LinearAttnLayer,
3591 g4: &[CudaSlice<f32>],
3592 offs: &[usize],
3593 ts: &[usize],
3594 caches: &mut [&mut Cache],
3595 il: usize,
3596 ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3597 let ssm = self.cfg.ssm.as_ref().unwrap();
3598 let d_state = ssm.state_size as usize;
3599 let num_k = ssm.group_count as usize;
3600 let num_v = ssm.time_step_rank as usize;
3601 let key_dim = d_state * num_k;
3602 let value_dim = d_state * num_v;
3603 let conv_dim = key_dim * 2 + value_dim;
3604 let eps = self.cfg.rms_eps;
3605 let scale = 1.0 / (d_state as f32).sqrt();
3606 let b = ts.len();
3607 let c = Engine::gdn_chunk_size();
3608 let carried = caches.iter().any(|c| c.pos > 0);
3611 let use_vl = !carried
3612 && (2..=8).contains(&b)
3613 && Engine::gdn_chunked_enabled()
3614 && ts.iter().all(|&t| t >= 16)
3615 && e.gdn_mma_enabled(c)
3616 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3617 if !use_vl {
3618 return (0..b)
3619 .map(|s| {
3620 let (o, t) = (offs[s], ts[s]);
3621 self.linear_attn_prime_core_pad_view(
3622 e,
3623 la,
3624 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3625 &g4[1].slice(o * value_dim..(o + t) * value_dim),
3626 &g4[2].slice(o * num_v..(o + t) * num_v),
3627 &g4[3].slice(o * num_v..(o + t) * num_v),
3628 t,
3629 caches[s],
3630 il,
3631 None,
3632 )
3633 })
3634 .collect();
3635 }
3636 struct SeqBufs {
3640 conv_out: CudaSlice<f32>,
3641 q_g: CudaSlice<f32>,
3642 k_g: CudaSlice<f32>,
3643 v_g: CudaSlice<f32>,
3644 q_l2: CudaSlice<f32>,
3645 k_l2: CudaSlice<f32>,
3646 beta: CudaSlice<f32>,
3647 g_log: CudaSlice<f32>,
3648 gn: CudaSlice<f32>,
3649 gn16: CudaSlice<u8>,
3650 }
3651 let d_conv = ssm.conv_kernel as usize;
3652 let f16o = Self::f16out_on(e, 16);
3653 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
3655 let mut pres = Vec::with_capacity(b);
3656 for &t in ts.iter().take(b) {
3657 sb.push(SeqBufs {
3658 conv_out: e.uninit(conv_dim * t)?,
3659 q_g: e.uninit(d_state * hk * t)?,
3660 k_g: e.uninit(d_state * hk * t)?,
3661 v_g: e.uninit(d_state * num_v * t)?,
3662 q_l2: e.uninit(d_state * hk * t)?,
3663 k_l2: e.uninit(d_state * hk * t)?,
3664 beta: e.uninit(t * num_v)?,
3665 g_log: e.uninit(t * num_v)?,
3666 gn: e.uninit(d_state * num_v * t)?,
3667 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3668 });
3669 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3670 }
3671 let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3672 .map(|s| {
3673 let (o, t) = (offs[s], ts[s]);
3674 let rl = caches[s].recur[il].as_ref().unwrap();
3675 crate::GdnPrepVl {
3676 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3677 conv_state: e.addr_f32(&rl.conv_state),
3678 conv_out: e.addr_f32(&sb[s].conv_out),
3679 q_g: e.addr_f32(&sb[s].q_g),
3680 k_g: e.addr_f32(&sb[s].k_g),
3681 v_g: e.addr_f32(&sb[s].v_g),
3682 q_l2: e.addr_f32(&sb[s].q_l2),
3683 k_l2: e.addr_f32(&sb[s].k_l2),
3684 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3685 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3686 beta: e.addr_f32(&sb[s].beta),
3687 g_log: e.addr_f32(&sb[s].g_log),
3688 o: e.addr_f32(&pres[s].o),
3689 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3690 gn: e.addr_f32(&sb[s].gn),
3691 gn16: e.addr_u8(&sb[s].gn16),
3692 kb16: if Engine::l2_v2_on(d_state) {
3693 e.addr_u8(&pres[s].kb16)
3694 } else {
3695 0
3696 },
3697 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3698 e.addr_u8(&pres[s].qb16)
3699 } else {
3700 0
3701 },
3702 t: t as i32,
3703 pad: 0,
3704 }
3705 })
3706 .collect();
3707 let args: Vec<crate::GdnSeqVl> = (0..b)
3708 .map(|s| {
3709 let rl = caches[s].recur[il].as_ref().unwrap();
3710 crate::GdnSeqVl {
3711 kb16: e.addr_u8(&pres[s].kb16),
3712 gcum: e.addr_f32(&pres[s].gcum),
3713 beta: e.addr_f32(&sb[s].beta),
3714 u: e.addr_f32(&pres[s].u),
3715 wb16: e.addr_u8(&pres[s].wb16),
3716 y: e.addr_u8(&pres[s].y16),
3717 ssnap: e.addr_u8(&pres[s].ssnap16),
3718 state_in: e.addr_f32(&rl.ssm_state),
3719 state_out: e.addr_f32(&rl.ssm_state_alt),
3720 q: e.addr_f32(&sb[s].q_l2),
3721 p: e.addr_f32(&pres[s].p),
3722 o: e.addr_f32(&pres[s].o),
3723 k: e.addr_f32(&sb[s].k_l2),
3724 v: e.addr_f32(&sb[s].v_g),
3725 g: e.addr_f32(&sb[s].g_log),
3726 a: e.addr_f32(&pres[s].a),
3727 w: e.addr_f32(&pres[s].w),
3728 t: ts[s] as i32,
3729 nc: pres[s].nc as i32,
3730 }
3731 })
3732 .collect();
3733 e.gdn_prep_vl8(
3734 &prep_args,
3735 la.ssm_conv1d.float_data(),
3736 la.ssm_dt.float_data(),
3737 la.ssm_a.float_data(),
3738 conv_dim,
3739 d_conv,
3740 d_state,
3741 num_v,
3742 num_k,
3743 key_dim,
3744 hk,
3745 eps,
3746 )?;
3747 if !Engine::l2_v2_on(d_state) {
3750 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3751 }
3752 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3754 if !Engine::l2_v2_on(d_state) {
3756 for s in 0..b {
3757 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3758 }
3759 }
3760 let mut wa = [crate::GdnWVl::default(); 8];
3761 for s in 0..b {
3762 wa[s] = crate::GdnWVl {
3763 qb16: e.addr_u8(&pres[s].qb16),
3764 pb16: e.addr_u8(&pres[s].pb16),
3765 };
3766 }
3767 Some(crate::GdnWVl8(wa))
3768 } else {
3769 None
3770 };
3771 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3772 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3773 if f16o {
3774 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3775 }
3776 let mut out = Vec::with_capacity(b);
3778 for (s, bufs) in sb.into_iter().enumerate() {
3779 let rl = caches[s].recur[il].as_mut().unwrap();
3780 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3781 let (o, t) = (offs[s], ts[s]);
3782 let SeqBufs { mut gn, gn16, .. } = bufs;
3783 if f16o {
3784 out.push((gn, Some(gn16)));
3785 } else {
3786 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3787 e.gated_rmsnorm_zv(
3788 &pres[s].o,
3789 la.ssm_norm.float_data(),
3790 &z_v,
3791 &mut gn,
3792 d_state,
3793 num_v * t,
3794 eps,
3795 )?;
3796 out.push((gn, None));
3797 }
3798 }
3799 Ok(out)
3800 }
3801
3802 #[allow(clippy::too_many_arguments)]
3806 fn linear_attn_prime_core_pad_view(
3807 &self,
3808 e: &Engine,
3809 la: &LinearAttnLayer,
3810 qkv_mixed: &cudarc::driver::CudaView<f32>,
3811 z: &cudarc::driver::CudaView<f32>,
3812 beta_raw: &cudarc::driver::CudaView<f32>,
3813 alpha: &cudarc::driver::CudaView<f32>,
3814 t: usize,
3815 cache: &mut Cache,
3816 il: usize,
3817 pad_len: Option<&CudaSlice<i32>>,
3818 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3819 let cfg = &self.cfg;
3820 let ssm = cfg.ssm.as_ref().unwrap();
3821 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
3824 let scale = 1.0 / (d_state as f32).sqrt();
3825
3826 let prep =
3827 self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3828
3829 let mut o = e.uninit(d_state * num_v * t)?;
3835 let rl = cache.recur[il].as_mut().unwrap();
3836 {
3837 let crate::cache::RecurLayer {
3838 ssm_state,
3839 ssm_state_alt,
3840 ..
3841 } = rl;
3842 e.gdn_scan_prefill(
3843 &prep.q_l2,
3844 &prep.k_l2,
3845 &prep.v_g,
3846 &prep.g_log,
3847 &prep.beta,
3848 prep.kb16.as_ref(),
3849 prep.qb16.as_ref(),
3850 ssm_state,
3851 ssm_state_alt,
3852 &mut o,
3853 num_v,
3854 t,
3855 scale,
3856 prep.hk,
3857 )?;
3858 }
3859 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3860
3861 let mut gn = e.uninit(d_state * num_v * t)?;
3864 let gn16 = if Self::f16out_on(e, t) {
3865 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3866 e.gated_rmsnorm_f16out_zv(
3867 &o,
3868 la.ssm_norm.float_data(),
3869 z,
3870 &mut gn,
3871 &mut g16,
3872 d_state,
3873 num_v * t,
3874 eps,
3875 )?;
3876 Some(g16)
3877 } else {
3878 e.gated_rmsnorm_zv(
3879 &o,
3880 la.ssm_norm.float_data(),
3881 z,
3882 &mut gn,
3883 d_state,
3884 num_v * t,
3885 eps,
3886 )?;
3887 None
3888 };
3889 Ok((gn, gn16))
3890 }
3891
3892 #[allow(clippy::too_many_arguments)]
3894 fn linear_attn_prime_core_pad(
3895 &self,
3896 e: &Engine,
3897 la: &LinearAttnLayer,
3898 g4: Vec<CudaSlice<f32>>,
3899 t: usize,
3900 cache: &mut Cache,
3901 il: usize,
3902 pad_len: Option<&CudaSlice<i32>>,
3903 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3904 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3905 if let Some(xh) = &gn16 {
3906 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3907 return Ok(y);
3908 }
3909 }
3910 Ok(e.matmul(&la.ssm_out, &gn, t)?)
3911 }
3912
3913 pub fn full_attn(
3918 &self,
3919 e: &Engine,
3920 fa: &FullAttnLayer,
3921 h: &CudaSlice<f32>,
3922 pos_d: &CudaSlice<i32>,
3923 t: usize,
3924 il: usize,
3925 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3926 if self.cfg.step35.is_some() {
3927 return self.step35_attn(e, fa, h, pos_d, t, il);
3928 }
3929 let cfg = &self.cfg;
3930 let _n_embd = cfg.n_embd as usize;
3931 let geometry = cfg.full_attention_geometry_at(il as u32);
3932 let n_head = geometry.n_head as usize;
3933 let n_head_kv = geometry.n_head_kv as usize;
3934 let head_dim = geometry.head_dim_k as usize;
3935 let eps = cfg.rms_eps;
3936 let scale = geometry.attention_scale();
3937
3938 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3941 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3943 let v = g3.pop().unwrap();
3944 let mut k = g3.pop().unwrap();
3945 let qf = g3.pop().unwrap();
3946 let (mut q, gate) = if gated {
3947 let mut q = e.uninit(t * n_head * head_dim)?;
3948 let mut gate = e.uninit(t * n_head * head_dim)?;
3949 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3950 (q, Some(gate))
3951 } else {
3952 (qf, None)
3953 };
3954
3955 let mut qn = e.uninit(t * n_head * head_dim)?;
3957 e.rms_norm(
3958 &q,
3959 fa.q_norm.float_data(),
3960 &mut qn,
3961 head_dim,
3962 n_head * t,
3963 eps,
3964 )?;
3965 q = qn;
3966 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3967 e.rms_norm(
3968 &k,
3969 fa.k_norm.float_data(),
3970 &mut kn,
3971 head_dim,
3972 n_head_kv * t,
3973 eps,
3974 )?;
3975 k = kn;
3976 let rope_dims = geometry.n_rot as usize;
3977 e.rope_neox(
3978 &mut q,
3979 pos_d,
3980 head_dim,
3981 rope_dims,
3982 n_head,
3983 t,
3984 geometry.rope_base,
3985 1.0,
3986 )?;
3987 e.rope_neox(
3988 &mut k,
3989 pos_d,
3990 head_dim,
3991 rope_dims,
3992 n_head_kv,
3993 t,
3994 geometry.rope_base,
3995 1.0,
3996 )?;
3997
3998 let mut attn = e.uninit(t * n_head * head_dim)?;
4000 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
4003 e.sdpa_naive(
4005 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
4006 )?;
4007 } else {
4008 e.fa_prefill(
4009 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
4010 )?;
4011 }
4012
4013 let attn_g = match &gate {
4015 Some(gate) => {
4016 let mut gsig = e.uninit(t * n_head * head_dim)?;
4017 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
4018 let mut ag = e.uninit(t * n_head * head_dim)?;
4019 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
4020 ag
4021 }
4022 None => attn,
4023 };
4024
4025 let o = e.matmul(&fa.wo, &attn_g, t)?;
4027 Ok(o)
4028 }
4029
4030 pub fn linear_attn(
4032 &self,
4033 e: &Engine,
4034 la: &LinearAttnLayer,
4035 h: &CudaSlice<f32>,
4036 t: usize,
4037 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4038 let cfg = &self.cfg;
4039 let _n_embd = cfg.n_embd as usize;
4040 let ssm = cfg.ssm.as_ref().unwrap();
4041 let d_state = ssm.state_size as usize; let num_k = ssm.group_count as usize; let num_v = ssm.time_step_rank as usize; let d_conv = ssm.conv_kernel as usize; let head_k = d_state;
4046 let head_v = d_state;
4047 let key_dim = head_k * num_k; let value_dim = head_v * num_v; let conv_dim = key_dim * 2 + value_dim; let eps = cfg.rms_eps;
4051 let scale = 1.0 / (d_state as f32).sqrt();
4052
4053 let mut g4 = e.matmul_group(
4056 &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
4057 h,
4058 t,
4059 )?;
4060 let alpha = g4.pop().unwrap(); let beta_raw = g4.pop().unwrap(); let z = g4.pop().unwrap(); let qkv_mixed = g4.pop().unwrap(); let _ = (head_k, head_v);
4072 let mut q_g = e.uninit(d_state * num_v * t)?;
4073 let mut k_g = e.uninit(d_state * num_v * t)?;
4074 let mut v_g = e.uninit(d_state * num_v * t)?;
4075 e.ssm_conv1d_gdn(
4076 &qkv_mixed,
4077 la.ssm_conv1d.float_data(),
4078 &mut q_g,
4079 &mut k_g,
4080 &mut v_g,
4081 conv_dim,
4082 t,
4083 d_conv,
4084 d_state,
4085 num_v,
4086 num_k,
4087 key_dim,
4088 )?;
4089 let mut q_l2 = e.uninit(d_state * num_v * t)?;
4091 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
4092 let mut k_l2 = e.uninit(d_state * num_v * t)?;
4093 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
4094 let v_gd = v_g;
4095
4096 let mut beta = e.uninit(t * num_v)?;
4099 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
4100 let mut g_log = e.uninit(t * num_v)?;
4102 e.gdn_glog(
4103 &alpha,
4104 la.ssm_dt.float_data(),
4105 la.ssm_a.float_data(),
4106 &mut g_log,
4107 num_v,
4108 t,
4109 )?;
4110
4111 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
4114 let mut o = e.uninit(d_state * num_v * t)?;
4115 e.gdn_scan_prefill(
4116 &q_l2,
4117 &k_l2,
4118 &v_gd,
4119 &g_log,
4120 &beta,
4121 None,
4122 None,
4123 &state_in,
4124 &mut state_out,
4125 &mut o,
4126 num_v,
4127 t,
4128 scale,
4129 num_v,
4130 )?;
4131
4132 let mut gn = e.uninit(d_state * num_v * t)?;
4137 e.gated_rmsnorm(
4138 &o,
4139 la.ssm_norm.float_data(),
4140 &z,
4141 &mut gn,
4142 d_state,
4143 num_v * t,
4144 eps,
4145 )?;
4146
4147 let out = e.matmul(&la.ssm_out, &gn, t)?;
4151 Ok(out)
4152 }
4153}
4154
4155impl HybridModel {
4156 pub fn moe_ffn_il(
4167 &self,
4168 e: &Engine,
4169 m: &MoeWeights,
4170 z: &CudaSlice<f32>,
4171 t: usize,
4172 il: u16,
4173 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4174 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
4175 }
4176
4177 pub fn moe_ffn_il_prefill(
4180 &self,
4181 e: &Engine,
4182 m: &MoeWeights,
4183 z: &CudaSlice<f32>,
4184 t: usize,
4185 il: u16,
4186 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4187 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
4188 }
4189
4190 pub fn moe_ffn_il_zq8(
4194 &self,
4195 e: &Engine,
4196 m: &MoeWeights,
4197 z: &CudaSlice<f32>,
4198 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4199 t: usize,
4200 il: u16,
4201 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4202 Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4203 }
4204
4205 pub(crate) fn moe_ffn(
4213 e: &Engine,
4214 m: &MoeWeights,
4215 z: &CudaSlice<f32>,
4216 t: usize,
4217 cfg: &ModelConfig,
4218 il: u16,
4219 max_block: usize,
4220 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4221 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4222 }
4223
4224 #[allow(clippy::too_many_arguments)]
4225 pub(crate) fn moe_ffn_inner(
4226 e: &Engine,
4227 m: &MoeWeights,
4228 z: &CudaSlice<f32>,
4229 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4230 t: usize,
4231 cfg: &ModelConfig,
4232 il: u16,
4233 max_block: usize,
4234 prefill: bool,
4235 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4236 let worker_io = crate::spill_pread::worker_enabled();
4237 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4238 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4239 e.with_moe_cache(max_block, |cache, _| {
4240 cache.begin_forward_epoch(il, t);
4241 if worker_io {
4242 cache.begin_worker_scope();
4243 }
4244 Ok(())
4245 })?;
4246 }
4247 if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4248 let moe = cfg.moe.as_ref().unwrap();
4249 let n_expert = moe.expert_count as usize;
4250 let n_used = moe.expert_used_count as usize;
4251 let sigmoid = cfg.sigmoid_router().unwrap();
4252 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4253 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4254 return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4255 }
4256 if t > 1 && moe_grouped_enabled(cfg, prefill) {
4259 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4260 if std::env::var("MEMRA_MOE_GATE").is_ok() {
4265 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4266 let g_host = e.dtoh(&grouped_out)?;
4267 let s_host = e.dtoh(&seq_out)?;
4268 let g_bytes: &[u8] = unsafe {
4269 std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4270 };
4271 let s_bytes: &[u8] = unsafe {
4272 std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4273 };
4274 if g_bytes == s_bytes {
4275 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4276 } else {
4277 let diffs = g_host
4278 .iter()
4279 .zip(s_host.iter())
4280 .enumerate()
4281 .filter(|(_, (a, b))| a != b)
4282 .count();
4283 let maxdiff = g_host
4284 .iter()
4285 .zip(s_host.iter())
4286 .map(|(a, b)| (a - b).abs())
4287 .fold(0.0f32, f32::max);
4288 panic!(
4289 "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4290 g_host.len()
4291 );
4292 }
4293 }
4294 return Ok(grouped_out);
4295 }
4296 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4297 }
4298
4299 fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4300 let Some(moe) = cfg.moe.as_ref() else {
4301 return false;
4302 };
4303 static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4306 let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4307 std::env::var("MEMRA_MOE_STATS").is_ok()
4308 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4309 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4310 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4311 || std::env::var("MEMRA_MOE_GATE").is_ok()
4312 });
4313 cfg.step35.is_some()
4314 && sigmoid_router_enabled()
4315 && moe_dev_enabled()
4316 && moe_slab_enabled()
4317 && !observation_mode
4318 && moe.expert_used_count <= 8
4319 && m.has_uniform_expert_layout()
4320 && m.gate_exps.macros.is_none()
4321 && m.up_exps.macros.is_none()
4322 && m.down_exps.macros.is_none()
4323 && !m.has_macros
4324 && moe_q8_enabled()
4325 && q8_expert_supported(m.gate_exps.qtype)
4326 && q8_expert_supported(m.up_exps.qtype)
4327 && q8_expert_supported(m.down_exps.qtype)
4328 && m.dev_exps
4329 .as_ref()
4330 .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4331 }
4332
4333 pub(crate) fn moe_ffn_sequential(
4335 e: &Engine,
4336 m: &MoeWeights,
4337 z: &CudaSlice<f32>,
4338 t: usize,
4339 cfg: &ModelConfig,
4340 il: u16,
4341 max_block: usize,
4342 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4343 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4344 }
4345
4346 fn moe_router_logits(
4350 e: &Engine,
4351 m: &MoeWeights,
4352 z: &CudaSlice<f32>,
4353 t: usize,
4354 cfg: &ModelConfig,
4355 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4356 if t < PRIME_MIN_T {
4357 if crate::router_kernel_on() {
4359 e.router_gemv(
4360 m.gate_inp.float_data(),
4361 z,
4362 cfg.n_embd as usize,
4363 m.gate_exps.n_expert,
4364 t,
4365 )
4366 } else {
4367 e.matmul_decode_exact(&m.gate_inp, z, t)
4368 }
4369 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4370 e.router_gemv(
4371 m.gate_inp.float_data(),
4372 z,
4373 cfg.n_embd as usize,
4374 m.gate_exps.n_expert,
4375 t,
4376 )
4377 } else {
4378 e.matmul(&m.gate_inp, z, t)
4379 }
4380 }
4381
4382 fn trace_moe_routes(
4386 il: u16,
4387 t: usize,
4388 sel_all: &[u32],
4389 weights: &[f32],
4390 ) -> Result<(), Box<dyn std::error::Error>> {
4391 use std::io::Write as _;
4392 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4393 let mut f = std::fs::OpenOptions::new()
4394 .create(true)
4395 .append(true)
4396 .open(path)?;
4397 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4398 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4399 }
4400 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4401 let mut f = std::fs::OpenOptions::new()
4402 .create(true)
4403 .append(true)
4404 .open(path)?;
4405 let pairs: Vec<String> = sel_all
4406 .iter()
4407 .zip(weights)
4408 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4409 .collect();
4410 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4411 }
4412 Ok(())
4413 }
4414
4415 #[allow(clippy::too_many_arguments)]
4416 fn trace_sigmoid_router_logits(
4417 e: &Engine,
4418 il: u16,
4419 t: usize,
4420 n_expert: usize,
4421 n_used: usize,
4422 logits: &CudaSlice<f32>,
4423 m: &MoeWeights,
4424 (scaling_factor, route_norm): (f32, bool),
4425 ) -> Result<(), Box<dyn std::error::Error>> {
4426 if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4427 return Ok(());
4428 }
4429 let logits = e.dtoh(logits)?;
4430 let active: Vec<u8> = m
4431 .active_experts
4432 .as_ref()
4433 .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4434 .unwrap_or_else(|| vec![1; n_expert]);
4435 let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4436 crate::sigrouter_contract::capture_served_logits(
4437 il as u32,
4438 t,
4439 n_expert,
4440 n_used,
4441 scaling_factor,
4442 route_norm,
4443 &active,
4444 &bias,
4445 &logits,
4446 )?;
4447 Ok(())
4448 }
4449
4450 fn trace_moe_input(
4455 e: &Engine,
4456 il: u16,
4457 t: usize,
4458 n_embd: usize,
4459 z: &CudaSlice<f32>,
4460 ) -> Result<(), Box<dyn std::error::Error>> {
4461 use std::io::Write as _;
4462 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4463 return Ok(());
4464 };
4465 let host = e.dtoh(z)?;
4466 if host.len() != t * n_embd {
4467 return Err(format!(
4468 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4469 host.len(),
4470 t,
4471 n_embd
4472 )
4473 .into());
4474 }
4475 let bytes = unsafe {
4476 std::slice::from_raw_parts(
4477 host.as_ptr().cast::<u8>(),
4478 host.len() * std::mem::size_of::<f32>(),
4479 )
4480 };
4481 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4482 let mut state = state
4483 .lock()
4484 .map_err(|_| "MoE input trace writer lock is poisoned")?;
4485 if state.is_none() {
4486 let dir = std::path::PathBuf::from(&dir);
4487 std::fs::create_dir_all(&dir)?;
4488 let index = std::fs::OpenOptions::new()
4489 .create(true)
4490 .append(true)
4491 .open(dir.join("index.jsonl"))?;
4492 *state = Some(MoeInputTraceWriter {
4493 dir,
4494 index,
4495 payloads: std::collections::HashMap::new(),
4496 });
4497 }
4498 let writer = state.as_mut().unwrap();
4499 if writer.dir != std::path::Path::new(&dir) {
4500 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4501 }
4502 let file_name = format!("layer-{il:03}.f32");
4503 if !writer.payloads.contains_key(&il) {
4504 let payload = std::fs::OpenOptions::new()
4505 .create(true)
4506 .append(true)
4507 .open(writer.dir.join(&file_name))?;
4508 let offset = payload.metadata()?.len();
4509 writer.payloads.insert(il, (payload, offset));
4510 }
4511 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4512 let row_offset = *offset;
4513 payload.write_all(bytes)?;
4514 *offset += bytes.len() as u64;
4515 writeln!(
4516 writer.index,
4517 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4518 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4519 \"payload_bytes\":{}}}",
4520 bytes.len()
4521 )?;
4522 Ok(())
4523 }
4524
4525 #[allow(clippy::too_many_arguments)]
4526 pub(crate) fn moe_ffn_sequential_zq8(
4527 e: &Engine,
4528 m: &MoeWeights,
4529 z: &CudaSlice<f32>,
4530 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4531 t: usize,
4532 cfg: &ModelConfig,
4533 il: u16,
4534 max_block: usize,
4535 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4536 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4537 let moe = cfg.moe.as_ref().unwrap();
4538 let n_embd = cfg.n_embd as usize; let n_expert = moe.expert_count as usize; let n_used = moe.expert_used_count as usize; let n_ff_exp = moe.expert_ff_length as usize; debug_assert_eq!(m.gate_exps.in_f, n_embd);
4545 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4546 debug_assert_eq!(m.down_exps.in_f, n_ff_exp); debug_assert_eq!(m.down_exps.out_f, n_embd); debug_assert_eq!(m.gate_exps.n_expert, n_expert);
4549
4550 let lim_exp = cfg.clamp_exp_at(il as u32);
4553 let lim_shexp = cfg.clamp_shexp_at(il as u32);
4554 let use_cache = Engine::moe_cache_enabled();
4555 let uniform_experts = m.has_uniform_expert_layout();
4556 let moe_q8 = uniform_experts
4557 && moe_q8_enabled()
4558 && q8_expert_supported(m.gate_exps.qtype)
4559 && q8_expert_supported(m.up_exps.qtype)
4560 && q8_expert_supported(m.down_exps.qtype);
4561 let cpu_expert_requested = crate::cpu_experts::configured();
4568 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4569 return Err(std::io::Error::other(
4570 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4571 )
4572 .into());
4573 }
4574 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4575 let freeze_cpu_residency = cpu_expert_requested
4581 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4582 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4583 .ok()
4584 .and_then(|value| value.parse::<usize>().ok())
4585 .is_some_and(|tokens| tokens > 0);
4586 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4587 e.freeze_moe_cache();
4588 }
4589 let cache_frozen = use_cache && e.moe_cache_frozen();
4590 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4591
4592 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4595 if let Some(sig) = cfg.sigmoid_router() {
4596 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4597 }
4598
4599 let no_exp_macros = m.gate_exps.macros.is_none()
4638 && m.up_exps.macros.is_none()
4639 && m.down_exps.macros.is_none();
4640 if cfg.sigmoid_router().is_none()
4644 && cfg.m3.is_none()
4645 && cfg.hy3.is_none()
4646 && !cfg.swiglu_clamped_at(il as u32)
4647 && no_exp_macros
4648 && t >= PRIME_MIN_T
4649 && m.dev_exps.is_some()
4650 && moe_q8_enabled()
4651 && q8_expert_supported(m.gate_exps.qtype)
4652 && q8_expert_supported(m.up_exps.qtype)
4653 && q8_expert_supported(m.down_exps.qtype)
4654 && std::env::var("MEMRA_MOE_PAIRS")
4655 .map(|v| v != "0")
4656 .unwrap_or(true)
4657 && std::env::var("MEMRA_MOE_STATS").is_err()
4658 {
4659 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4660 }
4661
4662 let dev_ok = uniform_experts
4680 && cfg.sigmoid_router().is_none()
4681 && cfg.m3.is_none()
4682 && cfg.hy3.is_none()
4683 && !cfg.swiglu_clamped_at(il as u32);
4684 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4688 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4689 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4690 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4691 if dev_ok
4692 && t < PRIME_MIN_T
4693 && m.dev_exps.is_some()
4694 && n_used <= 8
4695 && moe_dev_enabled()
4696 && !observe_routes
4697 {
4698 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4699 }
4700 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4701 let row_ok = e.with_moe_cache(max_block, |c, eng| {
4702 if moe_prewarm_enabled() {
4703 c.prewarm_layer(il, m, eng)?;
4704 }
4705 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4706 })?;
4707 if row_ok {
4708 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4709 }
4710 }
4711
4712 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4714 if cpu_hybrid {
4715 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4716 e,
4717 &logits,
4718 z,
4719 t,
4720 n_expert,
4721 n_used,
4722 m.exp_probs_b.as_deref(),
4723 sig,
4724 m.active_experts.as_deref(),
4725 )?;
4726 (sel, w, Some(input))
4727 } else {
4728 let (sel, w) =
4729 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4730 (sel, w, None)
4731 }
4732 } else {
4733 let (sel, w) =
4734 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4735 (sel, w, None)
4736 };
4737 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4738
4739 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4743 Self::trace_moe_input(e, il, t, n_embd, z)?;
4744
4745 let worker_disk_prefetch =
4757 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4758 let promote_worker_h2d =
4759 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4760 if promote_worker_h2d {
4761 let mut selected_blocks = Vec::with_capacity(n_used * 3);
4762 for &ex in sel_all.iter().take(n_used) {
4763 let ex = ex as u16;
4764 selected_blocks.extend([
4765 BlockId::new(il, PROJ_GATE, ex),
4766 BlockId::new(il, PROJ_UP, ex),
4767 BlockId::new(il, PROJ_DOWN, ex),
4768 ]);
4769 }
4770 for &ex in sel_all.iter().take(n_used) {
4771 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4772 }
4773 e.with_moe_cache(max_block, |cache, eng| {
4774 cache.promote_worker_reads_at_safe_boundary(
4775 &selected_blocks,
4776 &selected_blocks,
4777 eng,
4778 )?;
4779 Ok(())
4780 })?;
4781 }
4782
4783 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4786 let mut cnt = vec![0u32; n_expert];
4787 for &s in sel_all.iter() {
4788 cnt[s as usize] += 1;
4789 }
4790 let total = sel_all.len() as f64;
4791 let mut h = 0.0f64;
4792 let mut active = 0usize;
4793 for &c in &cnt {
4794 if c > 0 {
4795 active += 1;
4796 let p = c as f64 / total;
4797 h -= p * p.log2();
4798 }
4799 }
4800 let maxc = cnt.iter().copied().max().unwrap_or(0);
4801 println!(
4802 "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4803 il,
4804 t,
4805 sel_all.len(),
4806 active,
4807 n_expert,
4808 h,
4809 (n_expert as f64).log2(),
4810 total / active.max(1) as f64,
4811 maxc
4812 );
4813 }
4814
4815 let gdec_may_fire = uniform_experts
4828 && use_cache
4829 && n_used <= 8
4830 && gdec_enabled()
4831 && !cfg.swiglu_clamped_at(il as u32);
4832 let slab_local = m
4848 .dev_exps
4849 .as_ref()
4850 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4851 let slab_bases = slab_local.map(|d| {
4852 use cudarc::driver::DevicePtr;
4853 let s = e.stream();
4854 let (pg, _g0) = d.gate.device_ptr(&s);
4855 let (pu, _g1) = d.up.device_ptr(&s);
4856 let (pd, _g2) = d.down.device_ptr(&s);
4857 (pg as u64, pu as u64, pd as u64)
4858 });
4859 let slab_fused_may_fire = slab_bases.is_some()
4869 && n_used <= 8
4870 && gdec_enabled()
4871 && !cfg.swiglu_clamped_at(il as u32)
4872 && cfg.m3.is_none()
4873 && no_exp_macros
4874 && moe_q8;
4875 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4878 e.uninit(t * n_embd)?
4879 } else {
4880 e.zeros(t * n_embd)?
4881 };
4882 let cpu_input = if cpu_hybrid {
4885 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4886 } else {
4887 None
4888 };
4889
4890 let g_len = m.gate_exps.max_expert_bytes(); let u_len = m.up_exps.max_expert_bytes(); let d_len = m.down_exps.max_expert_bytes(); let mut scratch_g: Option<CudaSlice<u8>> = None;
4898 let mut scratch_u: Option<CudaSlice<u8>> = None;
4899 let mut scratch_d: Option<CudaSlice<u8>> = None;
4900 let page_window = moe_page_prefetch_window();
4908
4909 for tok in 0..t {
4912 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4913 let w = &w_all[tok * n_used..(tok + 1) * n_used];
4914 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4916
4917 let no_macros = m.gate_exps.macros.is_none()
4931 && m.up_exps.macros.is_none()
4932 && m.down_exps.macros.is_none();
4933 if slab_fused_may_fire {
4943 let (pg, pu, pd) = slab_bases.unwrap();
4944 let mut gp = [0u64; 8];
4945 let mut up = [0u64; 8];
4946 let mut dp = [0u64; 8];
4947 for (j, &ex) in sel.iter().enumerate() {
4948 let ex = ex as usize;
4949 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4950 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4951 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4952 }
4953 let mut wv = [0f32; 8];
4954 wv[..n_used].copy_from_slice(w);
4955 if tok_q8.is_none() {
4956 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4957 }
4958 let (zq, zd) = tok_q8.as_ref().unwrap();
4959 let act = e.moe_gate_up_silu8_q8(
4960 crate::WPtr8(gp),
4961 crate::WPtr8(up),
4962 zq,
4963 zd,
4964 n_embd,
4965 n_ff_exp,
4966 n_used,
4967 m.gate_exps.qtype,
4968 m.up_exps.qtype,
4969 m.gate_exps.row_bytes,
4970 m.up_exps.row_bytes,
4971 )?;
4972 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4973 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4974 e.moe_down8_fma_q8(
4975 crate::WPtr8(dp),
4976 crate::F32x8(wv),
4977 &aq2,
4978 &ad2,
4979 &mut dst,
4980 n_ff_exp,
4981 n_embd,
4982 n_used,
4983 m.down_exps.qtype,
4984 m.down_exps.row_bytes,
4985 )?;
4986 continue;
4987 }
4988 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4989 if tok_q8.is_none() {
4990 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4991 }
4992 let (zq, zd) = tok_q8.as_ref().unwrap();
4993 if Self::moe_gdec_token_q8(
4994 e,
4995 m,
4996 il,
4997 max_block,
4998 zq,
4999 zd,
5000 sel,
5001 w,
5002 &mut moe_out,
5003 tok,
5004 n_embd,
5005 n_ff_exp,
5006 n_used,
5007 )? {
5008 continue;
5009 }
5010 } else if gdec_may_fire
5011 && cfg.m3.is_none()
5012 && no_macros
5013 && Self::moe_gdec_token(
5014 e,
5015 m,
5016 il,
5017 max_block,
5018 &zt,
5019 sel,
5020 w,
5021 &mut moe_out,
5022 tok,
5023 n_embd,
5024 n_ff_exp,
5025 n_used,
5026 )?
5027 {
5028 continue;
5029 }
5030
5031 if gdec_may_fire || slab_fused_may_fire {
5037 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5038 e.memset_zeros_view(&mut row)?;
5039 }
5040
5041 let mut cpu_mask = vec![false; sel.len()];
5047 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
5048 let gpu_resident = if use_cache {
5049 e.with_moe_cache(max_block, |cache, _| {
5050 Ok(sel
5051 .iter()
5052 .map(|&expert| {
5053 let expert = expert as u16;
5054 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
5055 .into_iter()
5056 .filter(|&projection| {
5057 cache
5058 .resident(BlockId::new(il, projection, expert))
5059 .is_some()
5060 })
5061 .count()
5062 })
5063 .collect::<Vec<_>>())
5064 })?
5065 } else {
5066 vec![0; sel.len()]
5067 };
5068 let mut cpu_selected = Vec::new();
5069 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
5070 if gpu_resident[index] != 3 {
5071 cpu_mask[index] = true;
5072 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
5073 let expert = expert as usize;
5074 cpu_selected.push((expert, route_weight));
5075 }
5076 }
5077 if crate::cpu_experts::predictor_enabled() {
5078 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
5082 crate::cpu_experts::predictor_submit(il, row);
5083 }
5084 if cpu_selected.is_empty() {
5085 None
5086 } else {
5087 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
5088 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
5089 .map_err(std::io::Error::other)?;
5090 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
5091 }
5092 } else {
5093 None
5094 };
5095
5096 let worker_window = worker_disk_prefetch
5097 .then(worker_prefetch_window)
5098 .unwrap_or(0);
5099 for (j, &ex) in sel.iter().enumerate() {
5100 if cpu_mask[j] {
5101 continue;
5102 }
5103 let ex = ex as usize;
5104 if let Some(d) = slab_local {
5111 let gl = m.gate_exps.expert_layout(ex);
5112 let ul = m.up_exps.expert_layout(ex);
5113 let dl = m.down_exps.expert_layout(ex);
5114 let (g0, u0, d0) = (
5115 ex * m.gate_exps.expert_stride,
5116 ex * m.up_exps.expert_stride,
5117 ex * m.down_exps.expert_stride,
5118 );
5119 let (gate, up) = if moe_q8 {
5120 if tok_q8.is_none() {
5121 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5122 }
5123 let (zq, zd) = tok_q8.as_ref().unwrap();
5124 (
5125 e.qmatvec_expert_q8(
5126 &d.gate,
5127 g0..g0 + gl.len,
5128 zq,
5129 zd,
5130 1,
5131 m.gate_exps.in_f,
5132 m.gate_exps.out_f,
5133 gl.qtype,
5134 gl.row_bytes,
5135 )?,
5136 e.qmatvec_expert_q8(
5137 &d.up,
5138 u0..u0 + ul.len,
5139 zq,
5140 zd,
5141 1,
5142 m.up_exps.in_f,
5143 m.up_exps.out_f,
5144 ul.qtype,
5145 ul.row_bytes,
5146 )?,
5147 )
5148 } else {
5149 (
5150 e.qmatvec_view(
5151 &d.gate,
5152 g0..g0 + gl.len,
5153 &zt,
5154 1,
5155 m.gate_exps.in_f,
5156 m.gate_exps.out_f,
5157 gl.qtype,
5158 gl.row_bytes,
5159 )?,
5160 e.qmatvec_view(
5161 &d.up,
5162 u0..u0 + ul.len,
5163 &zt,
5164 1,
5165 m.up_exps.in_f,
5166 m.up_exps.out_f,
5167 ul.qtype,
5168 ul.row_bytes,
5169 )?,
5170 )
5171 };
5172 let mut act = e.uninit(n_ff_exp)?;
5173 Self::ffn_act_lim(
5174 e,
5175 cfg,
5176 &gate,
5177 &up,
5178 m.gate_exps.macro_scale(ex),
5179 m.up_exps.macro_scale(ex),
5180 lim_exp,
5181 &mut act,
5182 n_ff_exp,
5183 )?;
5184 let y = if moe_q8 {
5185 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5186 e.qmatvec_expert_q8(
5187 &d.down,
5188 d0..d0 + dl.len,
5189 &aq2,
5190 &ad2,
5191 1,
5192 m.down_exps.in_f,
5193 m.down_exps.out_f,
5194 dl.qtype,
5195 dl.row_bytes,
5196 )?
5197 } else {
5198 let actv = act.slice(0..n_ff_exp);
5199 e.qmatvec_view(
5200 &d.down,
5201 d0..d0 + dl.len,
5202 &actv,
5203 1,
5204 m.down_exps.in_f,
5205 m.down_exps.out_f,
5206 dl.qtype,
5207 dl.row_bytes,
5208 )?
5209 };
5210 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5211 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5212 continue;
5213 }
5214 for next in page_prefetch_positions(j, sel.len(), page_window) {
5215 Self::moe_prefetch_host_expert(sel[next] as usize, m);
5216 }
5217 let keep = [
5218 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5219 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5220 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5221 ];
5222 if worker_disk_prefetch && worker_window > 0 {
5223 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5224 Self::moe_prefetch_disk_expert(
5225 e,
5226 il,
5227 sel[next] as usize,
5228 m,
5229 max_block,
5230 &keep,
5231 )?;
5232 }
5233 } else if cache_dispatch
5234 && !cpu_hybrid
5235 && moe_prefetch_enabled()
5236 && j + 1 < sel.len()
5237 {
5238 let next = sel[j + 1] as usize;
5239 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5240 }
5241 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5242 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5243 if (gate_q8 || up_q8) && tok_q8.is_none() {
5246 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5247 }
5248 let gate = if gate_q8 {
5249 let (zq, zd) = tok_q8.as_ref().unwrap();
5250 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5251 } else {
5252 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5253 };
5254 let up = if up_q8 {
5255 let (zq, zd) = tok_q8.as_ref().unwrap();
5256 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5257 } else {
5258 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5259 };
5260 let mut act = e.uninit(n_ff_exp)?;
5261 Self::ffn_act_lim(
5262 e,
5263 cfg,
5264 &gate,
5265 &up,
5266 m.gate_exps.macro_scale(ex),
5267 m.up_exps.macro_scale(ex),
5268 lim_exp,
5269 &mut act,
5270 n_ff_exp,
5271 )?;
5272 let y = if down_q8 {
5273 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5274 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5275 } else {
5276 let actv = act.slice(0..n_ff_exp);
5277 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5278 };
5279 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5280 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5282 } else if cache_dispatch {
5283 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5288 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5289 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5291 e,
5292 cfg,
5293 &gate,
5294 &up,
5295 m.gate_exps.macro_scale(ex),
5296 m.up_exps.macro_scale(ex),
5297 lim_exp,
5298 &mut act,
5299 n_ff_exp,
5300 )?;
5301 let actv = act.slice(0..n_ff_exp);
5302 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5303 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5304 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5306 } else if cache_frozen {
5307 let gate = Self::moe_frozen_gemm(
5312 e,
5313 il,
5314 PROJ_GATE,
5315 ex,
5316 m,
5317 max_block,
5318 &zt,
5319 &mut scratch_g,
5320 g_len,
5321 )?;
5322 let up = Self::moe_frozen_gemm(
5323 e,
5324 il,
5325 PROJ_UP,
5326 ex,
5327 m,
5328 max_block,
5329 &zt,
5330 &mut scratch_u,
5331 u_len,
5332 )?;
5333 let mut act = e.uninit(n_ff_exp)?;
5334 Self::ffn_act_lim(
5335 e,
5336 cfg,
5337 &gate,
5338 &up,
5339 m.gate_exps.macro_scale(ex),
5340 m.up_exps.macro_scale(ex),
5341 lim_exp,
5342 &mut act,
5343 n_ff_exp,
5344 )?;
5345 let actv = act.slice(0..n_ff_exp);
5346 let y = Self::moe_frozen_gemm(
5347 e,
5348 il,
5349 PROJ_DOWN,
5350 ex,
5351 m,
5352 max_block,
5353 &actv,
5354 &mut scratch_d,
5355 d_len,
5356 )?;
5357 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5358 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5359 } else {
5360 if scratch_g.is_none() {
5364 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5365 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5366 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5367 }
5368 let (sg, su, sd) = (
5369 scratch_g.as_mut().unwrap(),
5370 scratch_u.as_mut().unwrap(),
5371 scratch_d.as_mut().unwrap(),
5372 );
5373 let gl = m.gate_exps.expert_layout(ex);
5374 let ul = m.up_exps.expert_layout(ex);
5375 let dl = m.down_exps.expert_layout(ex);
5376 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5377 let gate = e.qmatvec_view(
5378 sg,
5379 0..gl.len,
5380 &zt,
5381 1,
5382 m.gate_exps.in_f,
5383 m.gate_exps.out_f,
5384 gl.qtype,
5385 gl.row_bytes,
5386 )?;
5387
5388 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5389 let up = e.qmatvec_view(
5390 su,
5391 0..ul.len,
5392 &zt,
5393 1,
5394 m.up_exps.in_f,
5395 m.up_exps.out_f,
5396 ul.qtype,
5397 ul.row_bytes,
5398 )?;
5399
5400 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5402 e,
5403 cfg,
5404 &gate,
5405 &up,
5406 m.gate_exps.macro_scale(ex),
5407 m.up_exps.macro_scale(ex),
5408 lim_exp,
5409 &mut act,
5410 n_ff_exp,
5411 )?;
5412
5413 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5414 let actv = act.slice(0..n_ff_exp);
5415 let y = e.qmatvec_view(
5416 sd,
5417 0..dl.len,
5418 &actv,
5419 1,
5420 m.down_exps.in_f,
5421 m.down_exps.out_f,
5422 dl.qtype,
5423 dl.row_bytes,
5424 )?;
5425
5426 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5427 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5428 }
5429 }
5430 if let Some(worker) = cpu_worker {
5431 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5432 let cpu_output = e.htod(&cpu_output)?;
5433 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5434 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5435 }
5436 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5437 for (j, &ex) in sel.iter().enumerate() {
5438 if cpu_mask[j] {
5439 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5440 }
5441 }
5442 }
5443 }
5444
5445 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5450 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5451 {
5452 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
5461 let (sg_gate, sg_up) = if t == 1 {
5462 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5463 Some(pair) => pair,
5464 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5465 }
5466 } else if verify_t {
5467 (
5468 e.matmul_decode_exact(gate_shexp, z, t)?,
5469 e.matmul_decode_exact(up_shexp, z, t)?,
5470 )
5471 } else {
5472 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
5474 let mut sa = e.uninit(t * n_ff_sh)?; Self::ffn_act_lim(
5476 e,
5477 cfg,
5478 &sg_gate,
5479 &sg_up,
5480 1.0,
5481 1.0,
5482 lim_shexp,
5483 &mut sa,
5484 t * n_ff_sh,
5485 )?;
5486 let sh = if verify_t {
5487 e.matmul_decode_exact(down_shexp, &sa, t)?
5488 } else {
5489 e.matmul(down_shexp, &sa, t)?
5490 }; let g = match &m.gate_inp_shexp {
5504 Some(gate_inp_shexp) => {
5505 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5506 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5507 } else {
5508 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5509 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
5511 g
5512 }
5513 }
5514 None => e.htod(&vec![1.0f32; t])?,
5515 };
5516 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5518 }
5519
5520 Ok(moe_out)
5521 }
5522
5523 pub fn stage1_h2d_per_token(&self) -> u64 {
5526 use crate::hybrid::Ffn;
5527 let n_used = self
5528 .cfg
5529 .moe
5530 .as_ref()
5531 .map(|m| m.expert_used_count as u64)
5532 .unwrap_or(0);
5533 let mut bytes = 0u64;
5534 for l in self.layers.iter() {
5535 if let Ffn::Moe(m) = &l.ffn {
5536 bytes += n_used
5537 * (m.gate_exps.max_expert_bytes()
5538 + m.up_exps.max_expert_bytes()
5539 + m.down_exps.max_expert_bytes()) as u64;
5540 }
5541 }
5542 bytes
5543 }
5544
5545 pub(crate) fn max_moe_block(&self) -> usize {
5549 use crate::hybrid::Ffn;
5550 let mut mx = 0usize;
5551 let mut scan = |ffn: &Ffn| {
5552 if let Ffn::Moe(m) = ffn {
5553 mx = mx
5554 .max(m.gate_exps.max_expert_bytes())
5555 .max(m.up_exps.max_expert_bytes())
5556 .max(m.down_exps.max_expert_bytes());
5557 }
5558 };
5559 for l in self.layers.iter() {
5560 scan(&l.ffn);
5561 }
5562 if let Some(mtp) = self.mtp.as_ref() {
5563 scan(&mtp.ffn);
5564 }
5565 mx
5566 }
5567
5568 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5571 use crate::hybrid::Ffn;
5572 let mut sizes = Vec::new();
5573 let mut scan = |ffn: &Ffn| {
5574 let Ffn::Moe(m) = ffn else { return };
5575 for ex in 0..m.gate_exps.n_expert {
5576 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5577 continue;
5578 }
5579 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5580 let len = exps.expert_layout(ex).len;
5581 if len > 0 {
5582 sizes.push(len);
5583 }
5584 }
5585 }
5586 };
5587 for layer in &self.layers {
5588 scan(&layer.ffn);
5589 }
5590 if let Some(mtp) = &self.mtp {
5591 scan(&mtp.ffn);
5592 }
5593 sizes
5594 }
5595
5596 pub fn save_cpu_expert_residency_profile(
5602 &self,
5603 e: &Engine,
5604 path: &std::path::Path,
5605 ) -> Result<(), Box<dyn std::error::Error>> {
5606 let Some(ids) = e.export_moe_residency() else {
5607 return Err("no MoE residency cache to persist".into());
5608 };
5609 let mut body = format!(
5610 "memra-freeze-profile v1 max_block={} blocks={}\n",
5611 self.max_moe_block(),
5612 ids.len()
5613 );
5614 for (layer, proj, ex) in &ids {
5615 body.push_str(&format!("{layer} {proj} {ex}\n"));
5616 }
5617 let tmp = path.with_extension("tmp");
5618 std::fs::write(&tmp, body)?;
5619 std::fs::rename(&tmp, path)?;
5620 println!(
5621 "[moe-cache] freeze profile saved: {} blocks -> {}",
5622 ids.len(),
5623 path.display()
5624 );
5625 Ok(())
5626 }
5627
5628 pub fn restore_cpu_expert_residency_profile(
5632 &self,
5633 e: &Engine,
5634 path: &std::path::Path,
5635 ) -> Result<bool, Box<dyn std::error::Error>> {
5636 use crate::hybrid::Ffn;
5637 use crate::moe_cache::BlockId;
5638 let Ok(content) = std::fs::read_to_string(path) else {
5639 return Ok(false);
5640 };
5641 let mut lines = content.lines();
5642 let Some(header) = lines.next() else {
5643 return Ok(false);
5644 };
5645 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5646 if !header.starts_with(&expected) {
5647 println!(
5648 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5649 path.display()
5650 );
5651 return Ok(false);
5652 }
5653 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5654 std::collections::HashMap::new();
5655 for line in lines {
5656 let mut fields = line.split_whitespace();
5657 let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5658 else {
5659 continue;
5660 };
5661 let (Ok(layer), Ok(proj), Ok(ex)) =
5662 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5663 else {
5664 continue;
5665 };
5666 by_layer
5667 .entry(layer)
5668 .or_default()
5669 .push(BlockId::new(layer, proj, ex));
5670 }
5671 let requested: usize = by_layer.values().map(Vec::len).sum();
5672 if requested == 0 {
5673 return Ok(false);
5674 }
5675 let max_block = self.max_moe_block();
5676 let mut restaged = 0usize;
5677 let mut stage_layer =
5678 |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5679 let Ffn::Moe(m) = ffn else { return Ok(()) };
5680 let Some(ids) = by_layer.get(&layer_index) else {
5681 return Ok(());
5682 };
5683 e.with_moe_cache(max_block, |cache, eng| {
5684 for id in ids {
5685 if cache.restage_block(*id, m, eng)? {
5686 restaged += 1;
5687 }
5688 }
5689 Ok(())
5690 })
5691 };
5692 for (index, layer) in self.layers.iter().enumerate() {
5693 stage_layer(index as u16, &layer.ffn)?;
5694 }
5695 if let Some(mtp) = self.mtp.as_ref() {
5696 stage_layer(u16::MAX, &mtp.ffn)?;
5697 }
5698 e.freeze_moe_cache();
5699 println!(
5700 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5701 path.display()
5702 );
5703 Ok(true)
5704 }
5705
5706 pub fn freeze_cpu_expert_residency(
5708 &self,
5709 e: &Engine,
5710 ) -> Result<(), Box<dyn std::error::Error>> {
5711 e.freeze_moe_cache();
5712 Ok(())
5713 }
5714
5715 pub fn ffn_act(
5723 e: &Engine,
5724 cfg: &ModelConfig,
5725 gate: &CudaSlice<f32>,
5726 up: &CudaSlice<f32>,
5727 act: &mut CudaSlice<f32>,
5728 n: usize,
5729 ) -> Result<(), Box<dyn std::error::Error>> {
5730 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5731 }
5732
5733 #[allow(clippy::too_many_arguments)]
5737 pub(crate) fn ffn_act_scaled(
5738 e: &Engine,
5739 cfg: &ModelConfig,
5740 gate: &CudaSlice<f32>,
5741 up: &CudaSlice<f32>,
5742 gs: f32,
5743 us: f32,
5744 act: &mut CudaSlice<f32>,
5745 n: usize,
5746 ) -> Result<(), Box<dyn std::error::Error>> {
5747 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5748 }
5749
5750 #[allow(clippy::too_many_arguments)]
5759 pub(crate) fn ffn_act_lim(
5760 e: &Engine,
5761 cfg: &ModelConfig,
5762 gate: &CudaSlice<f32>,
5763 up: &CudaSlice<f32>,
5764 gs: f32,
5765 us: f32,
5766 limit: Option<f32>,
5767 act: &mut CudaSlice<f32>,
5768 n: usize,
5769 ) -> Result<(), Box<dyn std::error::Error>> {
5770 if let Some(m3) = cfg.m3.as_ref() {
5771 debug_assert!(
5772 limit.is_none(),
5773 "m3 swigluoai and step35 clamp are different archs"
5774 );
5775 return e.swigluoai_mul_scaled(
5776 gate,
5777 up,
5778 gs,
5779 us,
5780 m3.swiglu_alpha,
5781 m3.swiglu_limit,
5782 act,
5783 n,
5784 );
5785 }
5786 if let Some(l) = limit {
5787 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5788 }
5789 if gs == 1.0 && us == 1.0 {
5790 return e.silu_mul(gate, up, act, n);
5791 }
5792 e.silu_mul_scaled(gate, up, gs, us, act, n)
5793 }
5794
5795 fn moe_route(
5801 e: &Engine,
5802 logits: &CudaSlice<f32>,
5803 t: usize,
5804 n_expert: usize,
5805 n_used: usize,
5806 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5807 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5808 }
5809
5810 #[allow(clippy::too_many_arguments)]
5818 fn moe_route_sigmoid_cfg(
5819 e: &Engine,
5820 logits: &CudaSlice<f32>,
5821 t: usize,
5822 n_expert: usize,
5823 n_used: usize,
5824 m: &MoeWeights,
5825 (sf, route_norm): (f32, bool),
5826 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5827 if sigmoid_router_enabled() {
5828 return e.moe_router_sigmoid_topk_host(
5829 logits,
5830 t,
5831 n_expert,
5832 n_used,
5833 m.active_count(),
5834 &m.exp_probs_b_dev,
5835 &m.active_experts_dev,
5836 sf,
5837 route_norm,
5838 );
5839 }
5840 let lg = e.dtoh(logits)?;
5841 Self::moe_route_sigmoid_host(
5842 &lg,
5843 t,
5844 n_expert,
5845 n_used,
5846 m.exp_probs_b.as_deref(),
5847 sf,
5848 route_norm,
5849 m.active_experts.as_deref(),
5850 )
5851 }
5852
5853 fn moe_route_cfg(
5856 e: &Engine,
5857 logits: &CudaSlice<f32>,
5858 t: usize,
5859 n_expert: usize,
5860 n_used: usize,
5861 active: Option<&[bool]>,
5862 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5863 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5866 return e.moe_router_topk_host(logits, t, n_expert, n_used);
5867 }
5868 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
5871 let mut w_out = vec![0f32; t * n_used];
5872 for tok in 0..t {
5873 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5874 let maxl = row
5876 .iter()
5877 .enumerate()
5878 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5879 .map(|(_, &x)| x)
5880 .fold(f32::NEG_INFINITY, f32::max);
5881 let mut probs = vec![0f32; n_expert];
5882 let mut den = 0f32;
5883 for i in 0..n_expert {
5884 if active.is_some_and(|mask| !mask[i]) {
5885 continue;
5886 }
5887 let x = (row[i] - maxl).exp();
5888 probs[i] = x;
5889 den += x;
5890 }
5891 for p in probs.iter_mut() {
5892 *p /= den;
5893 }
5894 let mut idx: Vec<usize> = (0..n_expert)
5896 .filter(|&i| active.is_none_or(|mask| mask[i]))
5897 .collect();
5898 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5899 let sl = &idx[..n_used];
5900 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5901 let mut ws: f32 = wv.iter().sum();
5902 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() {
5904 *x /= ws;
5905 }
5906 for j in 0..n_used {
5907 sel[tok * n_used + j] = sl[j] as u32;
5908 w_out[tok * n_used + j] = wv[j];
5909 }
5910 }
5911 Ok((sel, w_out))
5912 }
5913
5914 #[allow(clippy::too_many_arguments)]
5915 fn moe_route_sigmoid_with_input(
5916 e: &Engine,
5917 logits: &CudaSlice<f32>,
5918 input: &CudaSlice<f32>,
5919 t: usize,
5920 n_expert: usize,
5921 n_used: usize,
5922 bias: Option<&[f32]>,
5923 (sf, route_norm): (f32, bool),
5924 active: Option<&[bool]>,
5925 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5926 let (lg, input) = e.dtoh_pair(logits, input)?;
5927 let (sel, w) =
5928 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5929 Ok((sel, w, input))
5930 }
5931
5932 pub fn start_moe_prefetch_predictor(
5937 &self,
5938 e: &Engine,
5939 cfg: &ModelConfig,
5940 ) -> Result<(), Box<dyn std::error::Error>> {
5941 use crate::hybrid::Ffn;
5942 let Some(sig) = cfg.sigmoid_router() else {
5943 return Err("prefetch predictor requires a sigmoid-router arch".into());
5944 };
5945 let resident: std::collections::HashSet<(u16, u8, u16)> = e
5946 .export_moe_residency()
5947 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5948 .into_iter()
5949 .collect();
5950 let mut layers = Vec::new();
5951 for (index, layer) in self.layers.iter().enumerate() {
5952 let Ffn::Moe(m) = &layer.ffn else { continue };
5953 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5954 continue;
5955 };
5956 let router = e.dtoh(data)?;
5957 let n_expert = m.gate_exps.n_expert;
5958 let n_embd = m.gate_exps.in_f;
5959 if router.len() != n_embd * n_expert {
5960 continue;
5961 }
5962 let build = |exps: &crate::model::HostExps| {
5963 (0..n_expert)
5964 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5965 .collect::<Vec<_>>()
5966 };
5967 layers.push((
5968 index as u16,
5969 crate::cpu_experts::PredictLayerInit {
5970 router,
5971 bias: m.exp_probs_b.clone(),
5972 active: m.active_experts.clone(),
5973 n_embd,
5974 n_used: cfg
5975 .moe
5976 .as_ref()
5977 .map(|moe| moe.expert_used_count as usize)
5978 .ok_or("prefetch predictor requires MoE config")?,
5979 sig,
5980 weights_n_expert: n_expert,
5981 gate: build(&m.gate_exps),
5982 up: build(&m.up_exps),
5983 down: build(&m.down_exps),
5984 },
5985 ));
5986 }
5987 crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5988 }
5989
5990 #[allow(clippy::too_many_arguments)]
5993 pub fn moe_route_sigmoid_host_public(
5994 logits: &[f32],
5995 t: usize,
5996 n_expert: usize,
5997 n_used: usize,
5998 bias: Option<&[f32]>,
5999 sf: f32,
6000 route_norm: bool,
6001 active: Option<&[bool]>,
6002 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6003 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
6004 }
6005
6006 #[allow(clippy::too_many_arguments)]
6007 fn moe_route_sigmoid_host(
6008 lg: &[f32],
6009 t: usize,
6010 n_expert: usize,
6011 n_used: usize,
6012 bias: Option<&[f32]>,
6013 sf: f32,
6014 route_norm: bool,
6015 active: Option<&[bool]>,
6016 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6017 let active_count = active
6018 .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
6019 .unwrap_or(n_expert);
6020 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6021 if lg.len() != t * n_expert {
6022 return Err(format!(
6023 "sigmoid router logits length mismatch: got {}, expected {}",
6024 lg.len(),
6025 t * n_expert,
6026 )
6027 .into());
6028 }
6029 let mut sel = vec![0u32; t * n_used];
6030 let mut w_out = vec![0f32; t * n_used];
6031 for tok in 0..t {
6032 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
6033 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
6034 let selsc: Vec<f32> = match bias {
6036 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
6037 None => scores.clone(),
6038 };
6039 let mut idx: Vec<usize> = (0..n_expert)
6040 .filter(|&i| active.is_none_or(|mask| mask[i]))
6041 .collect();
6042 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
6043 let sl = &idx[..n_used];
6044 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
6045 if route_norm {
6046 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
6047 for x in wv.iter_mut() {
6048 *x = *x / ws * sf;
6049 }
6050 } else {
6051 for x in wv.iter_mut() {
6052 *x *= sf;
6053 }
6054 }
6055 for j in 0..n_used {
6056 sel[tok * n_used + j] = sl[j] as u32;
6057 w_out[tok * n_used + j] = wv[j];
6058 }
6059 }
6060 Ok((sel, w_out))
6061 }
6062
6063 #[allow(clippy::too_many_arguments)]
6067 fn moe_ffn_sigmoid_dev(
6068 e: &Engine,
6069 m: &MoeWeights,
6070 z: &CudaSlice<f32>,
6071 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6072 logits: &CudaSlice<f32>,
6073 t: usize,
6074 cfg: &ModelConfig,
6075 il: u16,
6076 (scaling_factor, route_norm): (f32, bool),
6077 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6078 let moe = cfg.moe.as_ref().unwrap();
6079 let n_embd = cfg.n_embd as usize;
6080 let n_expert = moe.expert_count as usize;
6081 let n_used = moe.expert_used_count as usize;
6082 let n_ff_exp = moe.expert_ff_length as usize;
6083 let dev = m.dev_exps.as_ref().unwrap();
6084 debug_assert!(cfg.step35.is_some());
6085 debug_assert_eq!(dev.dev, e.ctx().ordinal());
6086 debug_assert!(m.has_uniform_expert_layout());
6087 debug_assert!(!m.has_macros);
6088
6089 let (sel_d, w_d) = e.moe_router_sigmoid_topk(
6090 logits,
6091 t,
6092 n_expert,
6093 n_used,
6094 m.active_count(),
6095 &m.exp_probs_b_dev,
6096 &m.active_experts_dev,
6097 scaling_factor,
6098 route_norm,
6099 )?;
6100 crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
6101 let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
6102 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6103 (combined, combined)
6104 } else {
6105 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6106 };
6107 let (zq, zd) = match (t, zq8) {
6108 (1, Some((q, d))) => (q.clone(), d.clone()),
6109 _ => e.quantize_q8_1(z, t, n_embd)?,
6110 };
6111 let n_pairs = t * n_used;
6112 let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
6113 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
6117 let pair_tok_d = e.htod_i32(&pair_tok)?;
6118 let gate = e.moe_pairs_matvec_q8(
6119 &dev.ptr_row,
6120 0,
6121 &pair_tok_d,
6122 &sel_d,
6123 &zq,
6124 &zd,
6125 n_embd,
6126 n_ff_exp,
6127 n_expert,
6128 n_pairs,
6129 m.gate_exps.qtype,
6130 gate_row_bytes,
6131 )?;
6132 let up = e.moe_pairs_matvec_q8(
6133 &dev.ptr_row,
6134 1,
6135 &pair_tok_d,
6136 &sel_d,
6137 &zq,
6138 &zd,
6139 n_embd,
6140 n_ff_exp,
6141 n_expert,
6142 n_pairs,
6143 m.up_exps.qtype,
6144 up_row_bytes,
6145 )?;
6146 let mut act = e.uninit(n_pairs * n_ff_exp)?;
6147 Self::ffn_act_lim(
6148 e,
6149 cfg,
6150 &gate,
6151 &up,
6152 1.0,
6153 1.0,
6154 cfg.clamp_exp_at(il as u32),
6155 &mut act,
6156 n_pairs * n_ff_exp,
6157 )?;
6158 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6159 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6160 let pair_self_d = e.htod_i32(&pair_self)?;
6161 let down = e.moe_pairs_matvec_q8(
6162 &dev.ptr_row,
6163 2,
6164 &pair_self_d,
6165 &sel_d,
6166 &aq2,
6167 &ad2,
6168 n_ff_exp,
6169 n_embd,
6170 n_expert,
6171 n_pairs,
6172 m.down_exps.qtype,
6173 m.down_exps.row_bytes,
6174 )?;
6175 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6176 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6177 let tok_off_d = e.htod_i32(&tok_off)?;
6178 let tok_ids_d = e.htod_i32(&tok_ids)?;
6179 let mut output = e.uninit(t * n_embd)?;
6180 e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
6181 output
6182 } else {
6183 let act = e.moe_gate_up_silu8_dev_q8_rows(
6184 &dev.ptr_row,
6185 &sel_d,
6186 &zq,
6187 &zd,
6188 t,
6189 n_embd,
6190 n_ff_exp,
6191 n_used,
6192 n_expert,
6193 m.gate_exps.qtype,
6194 m.up_exps.qtype,
6195 gate_row_bytes,
6196 up_row_bytes,
6197 &m.dev_macros,
6198 )?;
6199 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6200 let mut output = e.uninit(t * n_embd)?;
6201 e.moe_down8_fma_dev_q8_rows_g(
6202 &dev.ptr_row,
6203 &sel_d,
6204 &w_d,
6205 &aq2,
6206 &ad2,
6207 &mut output,
6208 t,
6209 n_ff_exp,
6210 n_embd,
6211 n_used,
6212 n_expert,
6213 m.down_exps.qtype,
6214 m.down_exps.row_bytes,
6215 )?;
6216 output
6217 };
6218
6219 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6220 eprintln!(
6221 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6222 cfg.clamp_exp_at(il as u32).is_some(),
6223 dev.gu_il,
6224 );
6225 }
6226 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6227 Ok(moe_out)
6228 }
6229
6230 fn moe_ffn_pairs(
6239 e: &Engine,
6240 m: &MoeWeights,
6241 z: &CudaSlice<f32>,
6242 logits: &CudaSlice<f32>,
6243 t: usize,
6244 cfg: &ModelConfig,
6245 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6246 let moe = cfg.moe.as_ref().unwrap();
6247 let n_embd = cfg.n_embd as usize;
6248 let n_expert = moe.expert_count as usize;
6249 let n_used = moe.expert_used_count as usize;
6250 let n_ff_exp = moe.expert_ff_length as usize;
6251 debug_assert!(
6256 !cfg.swiglu_clamped_anywhere(),
6257 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6258 );
6259 let dev = m.dev_exps.as_ref().unwrap();
6260 let (rbg_d, rbu_d) = if dev.gu_il {
6262 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6263 (sxx, sxx)
6264 } else {
6265 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6266 };
6267
6268 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6269 let n_pairs = t * n_used;
6270 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6273 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6274 let pair_w: Vec<f32> = w_all.clone();
6275 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6276 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6277 let pt = e.htod_i32(&pair_tok)?;
6278 let px = e.htod_i32(&pair_ex)?;
6279 let pw = e.htod(&pair_w)?;
6280 let toff = e.htod_i32(&tok_off)?;
6281 let tids = e.htod_i32(&tok_ids)?;
6282
6283 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6287 for p in 0..n_pairs {
6288 by_ex[pair_ex[p] as usize].push(p as i32);
6289 }
6290 let mut ex_ids: Vec<i32> = Vec::new();
6291 let mut ex_off: Vec<i32> = vec![0];
6292 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6293 for (ex, list) in by_ex.iter().enumerate() {
6294 if list.is_empty() {
6295 continue;
6296 }
6297 ex_ids.push(ex as i32);
6298 ex_pairs.extend_from_slice(list);
6299 ex_off.push(ex_pairs.len() as i32);
6300 }
6301 let n_active = ex_ids.len();
6302 let exi = e.htod_i32(&ex_ids)?;
6303 let exo = e.htod_i32(&ex_off)?;
6304 let exp_d = e.htod_i32(&ex_pairs)?;
6305 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6326 let mma_t = *MMA_T.get_or_init(|| {
6327 std::env::var("MEMRA_MOE_MMA_T")
6328 .ok()
6329 .and_then(|v| v.parse().ok())
6330 .unwrap_or(16)
6331 });
6332 let use_mma = std::env::var("MEMRA_MOE_MMA")
6333 .map(|v| v != "0")
6334 .unwrap_or(true)
6335 && t >= mma_t
6336 && q8_expert_dec_supported(m.gate_exps.qtype)
6337 && q8_expert_dec_supported(m.up_exps.qtype)
6338 && q8_expert_dec_supported(m.down_exps.qtype)
6339 && n_embd % 256 == 0
6340 && n_ff_exp % 256 == 0;
6341 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6357 && q8_expert_dec_supported(m.up_exps.qtype)
6358 && q8_expert_dec_supported(m.down_exps.qtype)
6359 && n_embd % 256 == 0
6360 && n_ff_exp % 256 == 0;
6361 let f16g_mode = crate::moe_f16g_mode();
6362 let f16g = f16g_mode != 0
6363 && t >= mma_t
6364 && (f16g_mode != 3 || !mma_capable)
6365 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6366 && f16g_proj_ok(m.up_exps.qtype, n_embd)
6367 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6368 if use_mma || f16g {
6369 let y_down = if f16g {
6377 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6381 let csr_tok_d = e.htod_i32(&csr_tok)?;
6382 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6383 let g_csr = e.moe_f16_grouped(
6384 &dev.ptr_row,
6385 0,
6386 n_expert,
6387 &exi,
6388 &ex_off,
6389 &exo,
6390 &z_f16,
6391 &z_s,
6392 n_embd,
6393 n_ff_exp,
6394 n_active,
6395 n_pairs,
6396 m.gate_exps.qtype,
6397 rbg_d,
6398 )?;
6399 let u_csr = e.moe_f16_grouped(
6400 &dev.ptr_row,
6401 1,
6402 n_expert,
6403 &exi,
6404 &ex_off,
6405 &exo,
6406 &z_f16,
6407 &z_s,
6408 n_embd,
6409 n_ff_exp,
6410 n_active,
6411 n_pairs,
6412 m.up_exps.qtype,
6413 rbu_d,
6414 )?;
6415 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6416 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6417 let d_csr = e.moe_f16_grouped(
6418 &dev.ptr_row,
6419 2,
6420 n_expert,
6421 &exi,
6422 &ex_off,
6423 &exo,
6424 &a_f16,
6425 &a_s,
6426 n_ff_exp,
6427 n_embd,
6428 n_active,
6429 n_pairs,
6430 m.down_exps.qtype,
6431 m.down_exps.row_bytes,
6432 )?;
6433 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6434 } else {
6435 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6437 let gate = e.mmq_iq_experts(
6438 &dev.ptr_row,
6439 0,
6440 n_expert,
6441 &exi,
6442 &exo,
6443 &exp_d,
6444 &pt,
6445 &z_scr,
6446 n_embd,
6447 n_ff_exp,
6448 n_active,
6449 n_pairs,
6450 t,
6451 m.gate_exps.qtype,
6452 rbg_d,
6453 )?;
6454 let up = e.mmq_iq_experts(
6455 &dev.ptr_row,
6456 1,
6457 n_expert,
6458 &exi,
6459 &exo,
6460 &exp_d,
6461 &pt,
6462 &z_scr,
6463 n_embd,
6464 n_ff_exp,
6465 n_active,
6466 n_pairs,
6467 t,
6468 m.up_exps.qtype,
6469 rbu_d,
6470 )?;
6471 let a_scr = if crate::moe_fuse_actq_on() {
6477 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6478 } else {
6479 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6480 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6481 };
6482 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6483 let pself = e.htod_i32(&pair_self)?;
6484 e.mmq_iq_experts(
6485 &dev.ptr_row,
6486 2,
6487 n_expert,
6488 &exi,
6489 &exo,
6490 &exp_d,
6491 &pself,
6492 &a_scr,
6493 n_ff_exp,
6494 n_embd,
6495 n_active,
6496 n_pairs,
6497 n_pairs,
6498 m.down_exps.qtype,
6499 m.down_exps.row_bytes,
6500 )?
6501 };
6502 let mut moe_out = e.uninit(t * n_embd)?;
6503 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6504 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6505 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6506 {
6507 let n_ff_sh = gate_shexp.out_features();
6508 let sg_gate = e.matmul(gate_shexp, z, t)?;
6509 let sg_up = e.matmul(up_shexp, z, t)?;
6510 let mut sa = e.uninit(t * n_ff_sh)?;
6511 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6512 let sh = e.matmul(down_shexp, &sa, t)?;
6513 let g = match &m.gate_inp_shexp {
6519 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6520 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6521 }
6522 Some(gate_inp_shexp) => {
6523 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6524 let mut g = e.uninit(t)?;
6525 e.sigmoid(&gs, &mut g, t)?;
6526 g
6527 }
6528 None => e.htod(&vec![1.0f32; t])?,
6529 };
6530 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6531 }
6532 return Ok(moe_out);
6533 }
6534
6535 let dec = std::env::var("MEMRA_MOE_DEC")
6538 .map(|v| v != "0")
6539 .unwrap_or(true);
6540 let matvec = |proj,
6541 exi: &_,
6542 exo: &_,
6543 exp_d: &_,
6544 pt: &_,
6545 aq: &_,
6546 ad: &_,
6547 inf,
6548 outf,
6549 qtype,
6550 rb|
6551 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6552 let dec = dec && q8_expert_dec_supported(qtype);
6554 if dec {
6555 e.moe_pairs_matvec_q8_dec(
6556 &dev.ptr_row,
6557 proj,
6558 exi,
6559 exo,
6560 exp_d,
6561 pt,
6562 aq,
6563 ad,
6564 inf,
6565 outf,
6566 n_expert,
6567 n_active,
6568 n_pairs,
6569 qtype,
6570 rb,
6571 )
6572 } else {
6573 e.moe_pairs_matvec_q8_em(
6574 &dev.ptr_row,
6575 proj,
6576 exi,
6577 exo,
6578 exp_d,
6579 pt,
6580 aq,
6581 ad,
6582 inf,
6583 outf,
6584 n_expert,
6585 n_active,
6586 n_pairs,
6587 qtype,
6588 rb,
6589 )
6590 }
6591 };
6592 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6593 let gate = matvec(
6594 0,
6595 &exi,
6596 &exo,
6597 &exp_d,
6598 &pt,
6599 &zq,
6600 &zd,
6601 n_embd,
6602 n_ff_exp,
6603 m.gate_exps.qtype,
6604 rbg_d,
6605 )?;
6606 let up = matvec(
6607 1,
6608 &exi,
6609 &exo,
6610 &exp_d,
6611 &pt,
6612 &zq,
6613 &zd,
6614 n_embd,
6615 n_ff_exp,
6616 m.up_exps.qtype,
6617 rbu_d,
6618 )?;
6619 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6620 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6621 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6623 let pself = e.htod_i32(&pair_self)?;
6624 let y_down = matvec(
6625 2,
6626 &exi,
6627 &exo,
6628 &exp_d,
6629 &pself,
6630 &aq2,
6631 &ad2,
6632 n_ff_exp,
6633 n_embd,
6634 m.down_exps.qtype,
6635 m.down_exps.row_bytes,
6636 )?;
6637 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6639
6640 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6644 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6645 {
6646 let n_ff_sh = gate_shexp.out_features();
6647 let step_exact = cfg.step35.is_some();
6651 let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6652 let (sg_gate, sg_up) = if step_exact && t == 1 {
6653 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6654 Some(pair) => pair,
6655 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6656 }
6657 } else if verify_t {
6658 let mut fused = None;
6659 if crate::spec::spec_fused_t()
6660 && (2..=4).contains(&t)
6661 && e.uses_q8_1_fast(gate_shexp)
6662 && e.uses_q8_1_fast(up_shexp)
6663 {
6664 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6665 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6666 }
6667 match fused {
6668 Some(pair) => pair,
6669 None => (
6670 e.matmul_decode_exact(gate_shexp, z, t)?,
6671 e.matmul_decode_exact(up_shexp, z, t)?,
6672 ),
6673 }
6674 } else {
6675 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6676 };
6677 let mut sa = e.uninit(t * n_ff_sh)?;
6678 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6679 let sh = if verify_t {
6680 e.matmul_decode_exact(down_shexp, &sa, t)?
6681 } else {
6682 e.matmul(down_shexp, &sa, t)?
6683 };
6684 let g = match &m.gate_inp_shexp {
6689 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6690 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6691 }
6692 Some(gate_inp_shexp) => {
6693 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6694 let mut g = e.uninit(t)?;
6695 e.sigmoid(&gs, &mut g, t)?;
6696 g
6697 }
6698 None => e.htod(&vec![1.0f32; t])?,
6699 };
6700 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6701 }
6702 Ok(moe_out)
6703 }
6704
6705 #[allow(clippy::too_many_arguments)]
6707 #[allow(clippy::too_many_arguments)]
6708 fn moe_ffn_dev(
6709 e: &Engine,
6710 m: &MoeWeights,
6711 z: &CudaSlice<f32>,
6712 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6713 logits: &CudaSlice<f32>,
6714 t: usize,
6715 cfg: &ModelConfig,
6716 il: u16,
6717 max_block: usize,
6718 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6719 let moe = cfg.moe.as_ref().unwrap();
6720 let n_embd = cfg.n_embd as usize;
6721 let n_expert = moe.expert_count as usize;
6722 let n_used = moe.expert_used_count as usize;
6723 let n_ff_exp = moe.expert_ff_length as usize;
6724 debug_assert!(
6728 cfg.sigmoid_router().is_none(),
6729 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6730 );
6731 debug_assert!(
6732 !cfg.swiglu_clamped_at(il as u32),
6733 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6734 );
6735
6736 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6738 if m.has_macros {
6741 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6742 }
6743
6744 let mut moe_out = e.uninit(t * n_embd)?;
6746
6747 if let Some(dev) = m.dev_exps.as_ref() {
6750 let (rbg_d, rbu_d) = if dev.gu_il {
6753 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6754 (sxx, sxx)
6755 } else {
6756 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6757 };
6758 let q8 = moe_q8_enabled()
6759 && q8_expert_supported(m.gate_exps.qtype)
6760 && q8_expert_supported(m.up_exps.qtype)
6761 && q8_expert_supported(m.down_exps.qtype);
6762 let rows_arm = q8
6771 && t > 1
6772 && crate::spec::spec_m2()
6773 && n_ff_exp == 512
6774 && n_used <= 8
6775 && std::env::var("MEMRA_MOE_DEVQ8_GU")
6776 .map(|v| v.is_empty() || v == "v")
6777 .unwrap_or(true)
6778 && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6779 .map(|v| v.is_empty() || v == "w8h2v")
6780 .unwrap_or(true);
6781 let csr_mode = std::env::var("MEMRA_MOE_CSR")
6790 .ok()
6791 .and_then(|v| v.parse::<i32>().ok())
6792 .unwrap_or(1);
6793 let csr_qt =
6797 |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S || qt == crate::QT_NVFP4;
6798 let csr_uniform = m.gate_exps.qtype == m.up_exps.qtype;
6799 let csr_arm = rows_arm
6800 && csr_mode > 0
6801 && t <= 10
6802 && csr_uniform
6803 && csr_qt(m.gate_exps.qtype)
6804 && csr_qt(m.up_exps.qtype)
6805 && csr_qt(m.down_exps.qtype);
6806 if csr_arm {
6807 if csr_mode == 2 {
6808 static ENGAGED: std::sync::Once = std::sync::Once::new();
6809 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6810 }
6811 let n_pairs = t * n_used;
6812 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6813 let act = e.moe_gate_up_silu8_dev_q8_csr(
6814 &dev.ptr_row,
6815 &sel_d,
6816 &zq,
6817 &zd,
6818 n_pairs,
6819 n_embd,
6820 n_ff_exp,
6821 n_used,
6822 n_expert,
6823 m.gate_exps.qtype,
6824 m.up_exps.qtype,
6825 rbg_d,
6826 rbu_d,
6827 )?;
6828 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6829 e.moe_down8_fma_dev_q8_rows(
6833 &dev.ptr_row,
6834 &sel_d,
6835 &w_d,
6836 &aq2,
6837 &ad2,
6838 &mut moe_out,
6839 t,
6840 n_ff_exp,
6841 n_embd,
6842 n_used,
6843 n_expert,
6844 m.down_exps.qtype,
6845 m.down_exps.row_bytes,
6846 )?;
6847 if csr_mode == 2 {
6848 let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6850 &dev.ptr_row,
6851 &sel_d,
6852 &zq,
6853 &zd,
6854 t,
6855 n_embd,
6856 n_ff_exp,
6857 n_used,
6858 n_expert,
6859 m.gate_exps.qtype,
6860 m.up_exps.qtype,
6861 rbg_d,
6862 rbu_d,
6863 &m.dev_macros,
6864 )?;
6865 let mut out_r = e.uninit(t * n_embd)?;
6866 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6867 e.moe_down8_fma_dev_q8_rows(
6868 &dev.ptr_row,
6869 &sel_d,
6870 &w_d,
6871 &aq2r,
6872 &ad2r,
6873 &mut out_r,
6874 t,
6875 n_ff_exp,
6876 n_embd,
6877 n_used,
6878 n_expert,
6879 m.down_exps.qtype,
6880 m.down_exps.row_bytes,
6881 )?;
6882 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6883 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6884 let ba = a1
6885 .iter()
6886 .zip(&a2)
6887 .filter(|(x, y)| x.to_bits() != y.to_bits())
6888 .count();
6889 let bo = o1
6890 .iter()
6891 .zip(&o2)
6892 .filter(|(x, y)| x.to_bits() != y.to_bits())
6893 .count();
6894 if ba + bo > 0 {
6895 eprintln!(
6896 "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6897 a1.len(),
6898 o1.len()
6899 );
6900 let sel_h = e.dtoh_i32(&sel_d)?;
6902 let mut shown = 0;
6903 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6904 if x.to_bits() != y.to_bits() && shown < 4 {
6905 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6906 let ex = sel_h[p];
6907 let npx = sel_h.iter().filter(|&&v| v == ex).count();
6908 eprintln!(
6909 " ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6910 );
6911 shown += 1;
6912 }
6913 }
6914 std::process::exit(3);
6915 }
6916 }
6917 } else if rows_arm {
6918 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6921 use std::sync::atomic::{AtomicU64, Ordering};
6922 static PAIRS: AtomicU64 = AtomicU64::new(0);
6923 static UNIQ: AtomicU64 = AtomicU64::new(0);
6924 static CALLS: AtomicU64 = AtomicU64::new(0);
6925 let sel_h = e.dtoh_i32(&sel_d)?;
6926 let mut u: Vec<i32> = sel_h.clone();
6927 u.sort_unstable();
6928 u.dedup();
6929 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6930 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6931 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6932 if c % 480 == 0 {
6933 let p = PAIRS.load(Ordering::Relaxed);
6934 let q = UNIQ.load(Ordering::Relaxed);
6935 eprintln!(
6936 "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6937 q as f64 / p as f64
6938 );
6939 }
6940 }
6941 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6942 let act = e.moe_gate_up_silu8_dev_q8_rows(
6943 &dev.ptr_row,
6944 &sel_d,
6945 &zq,
6946 &zd,
6947 t,
6948 n_embd,
6949 n_ff_exp,
6950 n_used,
6951 n_expert,
6952 m.gate_exps.qtype,
6953 m.up_exps.qtype,
6954 rbg_d,
6955 rbu_d,
6956 &m.dev_macros,
6957 )?;
6958 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6959 e.moe_down8_fma_dev_q8_rows(
6960 &dev.ptr_row,
6961 &sel_d,
6962 &w_d,
6963 &aq2,
6964 &ad2,
6965 &mut moe_out,
6966 t,
6967 n_ff_exp,
6968 n_embd,
6969 n_used,
6970 n_expert,
6971 m.down_exps.qtype,
6972 m.down_exps.row_bytes,
6973 )?;
6974 } else {
6975 for tok in 0..t {
6976 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6977 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6978 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6979 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6980 if q8 {
6981 let (zq, zd) = match (t, zq8) {
6982 (1, Some((q, d))) => (q.clone(), d.clone()),
6983 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6984 };
6985 let act = e.moe_gate_up_silu8_dev_q8(
6986 &dev.ptr_row,
6987 &selt,
6988 &zq,
6989 &zd,
6990 n_embd,
6991 n_ff_exp,
6992 n_used,
6993 n_expert,
6994 m.gate_exps.qtype,
6995 m.up_exps.qtype,
6996 rbg_d,
6997 rbu_d,
6998 &m.dev_macros,
6999 )?;
7000 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7001 e.moe_down8_fma_dev_q8(
7002 &dev.ptr_row,
7003 &selt,
7004 &wt,
7005 &aq2,
7006 &ad2,
7007 &mut dst,
7008 n_ff_exp,
7009 n_embd,
7010 n_used,
7011 n_expert,
7012 m.down_exps.qtype,
7013 m.down_exps.row_bytes,
7014 )?;
7015 } else {
7016 let act = e.moe_gate_up_silu8_dev(
7017 &dev.ptr_row,
7018 &selt,
7019 &zt,
7020 n_embd,
7021 n_ff_exp,
7022 n_used,
7023 n_expert,
7024 m.gate_exps.qtype,
7025 m.up_exps.qtype,
7026 rbg_d,
7027 rbu_d,
7028 &m.dev_macros,
7029 )?;
7030 e.moe_down8_fma_dev(
7031 &dev.ptr_row,
7032 &selt,
7033 &wt,
7034 &act,
7035 &mut dst,
7036 n_ff_exp,
7037 n_embd,
7038 n_used,
7039 n_expert,
7040 m.down_exps.qtype,
7041 m.down_exps.row_bytes,
7042 )?;
7043 }
7044 }
7045 }
7046 } else {
7047 let q8 = moe_q8_enabled()
7054 && q8_expert_supported(m.gate_exps.qtype)
7055 && q8_expert_supported(m.up_exps.qtype)
7056 && q8_expert_supported(m.down_exps.qtype);
7057 e.with_moe_cache(max_block, |c, eng| {
7058 let row = c
7059 .layer_dev_row(il, n_expert, eng)?
7060 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
7061 for tok in 0..t {
7062 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
7063 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
7064 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
7065 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7066 if q8 {
7067 let (zq, zd) = match (t, zq8) {
7068 (1, Some((q, d))) => (q.clone(), d.clone()),
7069 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
7070 };
7071 let act = eng.moe_gate_up_silu8_dev_q8(
7072 row,
7073 &selt,
7074 &zq,
7075 &zd,
7076 n_embd,
7077 n_ff_exp,
7078 n_used,
7079 n_expert,
7080 m.gate_exps.qtype,
7081 m.up_exps.qtype,
7082 m.gate_exps.row_bytes,
7083 m.up_exps.row_bytes,
7084 &m.dev_macros,
7085 )?;
7086 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
7087 eng.moe_down8_fma_dev_q8(
7088 row,
7089 &selt,
7090 &wt,
7091 &aq2,
7092 &ad2,
7093 &mut dst,
7094 n_ff_exp,
7095 n_embd,
7096 n_used,
7097 n_expert,
7098 m.down_exps.qtype,
7099 m.down_exps.row_bytes,
7100 )?;
7101 } else {
7102 let act = eng.moe_gate_up_silu8_dev(
7103 row,
7104 &selt,
7105 &zt,
7106 n_embd,
7107 n_ff_exp,
7108 n_used,
7109 n_expert,
7110 m.gate_exps.qtype,
7111 m.up_exps.qtype,
7112 m.gate_exps.row_bytes,
7113 m.up_exps.row_bytes,
7114 &m.dev_macros,
7115 )?;
7116 eng.moe_down8_fma_dev(
7117 row,
7118 &selt,
7119 &wt,
7120 &act,
7121 &mut dst,
7122 n_ff_exp,
7123 n_embd,
7124 n_used,
7125 n_expert,
7126 m.down_exps.qtype,
7127 m.down_exps.row_bytes,
7128 )?;
7129 }
7130 }
7131 c.hits += (t * 3 * n_used) as u64;
7133 Ok(())
7134 })?;
7135 }
7136
7137 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7142 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7143 {
7144 let n_ff_sh = gate_shexp.out_features();
7145 let verify_t = t > 1 && t < PRIME_MIN_T;
7148 let (sg_gate, sg_up) = if t == 1 {
7149 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
7150 Some(pair) => pair,
7151 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
7152 }
7153 } else if verify_t {
7154 let mut fused = None;
7158 if crate::spec::spec_fused_t()
7159 && (2..=4).contains(&t)
7160 && e.uses_q8_1_fast(gate_shexp)
7161 && e.uses_q8_1_fast(up_shexp)
7162 {
7163 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7164 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
7165 }
7166 match fused {
7167 Some(pair) => pair,
7168 None => (
7169 e.matmul_decode_exact(gate_shexp, z, t)?,
7170 e.matmul_decode_exact(up_shexp, z, t)?,
7171 ),
7172 }
7173 } else {
7174 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
7175 };
7176 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
7178 let sh = if verify_t {
7179 e.matmul_decode_exact(down_shexp, &sa, t)?
7180 } else {
7181 e.matmul(down_shexp, &sa, t)?
7182 };
7183 let g = match &m.gate_inp_shexp {
7187 Some(gate_inp_shexp) => {
7188 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7191 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7192 } else {
7193 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7194 let mut g = e.uninit(t)?;
7195 e.sigmoid(&gs, &mut g, t)?;
7196 g
7197 }
7198 }
7199 None => e.htod(&vec![1.0f32; t])?,
7200 };
7201 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7202 }
7203
7204 Ok(moe_out)
7205 }
7206
7207 #[allow(clippy::too_many_arguments)]
7217 #[allow(clippy::too_many_arguments)]
7220 fn moe_gdec_token_q8(
7221 e: &Engine,
7222 m: &MoeWeights,
7223 il: u16,
7224 max_block: usize,
7225 zq: &CudaSlice<i8>,
7226 zd: &CudaSlice<f32>,
7227 sel: &[u32],
7228 w: &[f32],
7229 moe_out: &mut CudaSlice<f32>,
7230 tok: usize,
7231 n_embd: usize,
7232 n_ff_exp: usize,
7233 n_used: usize,
7234 ) -> Result<bool, Box<dyn std::error::Error>> {
7235 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7236 use cudarc::driver::DevicePtr;
7237 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7238 let mut g = [0u64; 8];
7239 let mut u = [0u64; 8];
7240 let mut d = [0u64; 8];
7241 for (j, &ex) in sel.iter().enumerate() {
7242 let ex = ex as u16;
7243 let (Some(sg), Some(su), Some(sd)) = (
7244 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7245 c.resident(BlockId::new(il, PROJ_UP, ex)),
7246 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7247 ) else {
7248 return Ok(None);
7249 };
7250 let __s = eng.stream();
7251 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7252 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7253 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7254 g[j] = pg as u64;
7255 u[j] = pu as u64;
7256 d[j] = pd as u64;
7257 }
7258 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7259 for &ex in sel {
7260 let ex = ex as u16;
7261 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7262 c.note_profile_hit(BlockId::new(il, proj, ex));
7263 }
7264 }
7265 }
7266 c.hits += (3 * n_used) as u64;
7267 Ok(Some((g, u, d)))
7268 })?;
7269 let Some((g, u, d)) = ptrs else {
7270 return Ok(false);
7271 };
7272 let mut wv = [0f32; 8];
7273 wv[..n_used].copy_from_slice(w);
7274 let act = e.moe_gate_up_silu8_q8(
7275 crate::WPtr8(g),
7276 crate::WPtr8(u),
7277 zq,
7278 zd,
7279 n_embd,
7280 n_ff_exp,
7281 n_used,
7282 m.gate_exps.qtype,
7283 m.up_exps.qtype,
7284 m.gate_exps.row_bytes,
7285 m.up_exps.row_bytes,
7286 )?;
7287 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7289 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7290 e.moe_down8_fma_q8(
7291 crate::WPtr8(d),
7292 crate::F32x8(wv),
7293 &aq2,
7294 &ad2,
7295 &mut dst,
7296 n_ff_exp,
7297 n_embd,
7298 n_used,
7299 m.down_exps.qtype,
7300 m.down_exps.row_bytes,
7301 )?;
7302 Ok(true)
7303 }
7304
7305 fn moe_gdec_token(
7306 e: &Engine,
7307 m: &MoeWeights,
7308 il: u16,
7309 max_block: usize,
7310 zt: &cudarc::driver::CudaView<f32>,
7311 sel: &[u32],
7312 w: &[f32],
7313 moe_out: &mut CudaSlice<f32>,
7314 tok: usize,
7315 n_embd: usize,
7316 n_ff_exp: usize,
7317 n_used: usize,
7318 ) -> Result<bool, Box<dyn std::error::Error>> {
7319 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7320 use cudarc::driver::DevicePtr;
7321 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7323 let mut g = [0u64; 8];
7324 let mut u = [0u64; 8];
7325 let mut d = [0u64; 8];
7326 for (j, &ex) in sel.iter().enumerate() {
7327 let ex = ex as u16;
7328 let (Some(sg), Some(su), Some(sd)) = (
7329 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7330 c.resident(BlockId::new(il, PROJ_UP, ex)),
7331 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7332 ) else {
7333 return Ok(None);
7334 };
7335 let __s = eng.stream();
7336 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7337 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7338 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7339 g[j] = pg as u64;
7340 u[j] = pu as u64;
7341 d[j] = pd as u64;
7342 }
7343 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7344 for &ex in sel {
7345 let ex = ex as u16;
7346 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7347 c.note_profile_hit(BlockId::new(il, proj, ex));
7348 }
7349 }
7350 }
7351 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
7353 })?;
7354 let Some((g, u, d)) = ptrs else {
7355 return Ok(false);
7356 };
7357 let mut wv = [0f32; 8];
7358 wv[..n_used].copy_from_slice(w);
7359 let act = e.moe_gate_up_silu8(
7361 crate::WPtr8(g),
7362 crate::WPtr8(u),
7363 zt,
7364 n_embd,
7365 n_ff_exp,
7366 n_used,
7367 m.gate_exps.qtype,
7368 m.up_exps.qtype,
7369 m.gate_exps.row_bytes,
7370 m.up_exps.row_bytes,
7371 )?;
7372 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7373 e.moe_down8_fma_into(
7374 crate::WPtr8(d),
7375 crate::F32x8(wv),
7376 &act,
7377 &mut dst,
7378 n_ff_exp,
7379 n_embd,
7380 n_used,
7381 m.down_exps.qtype,
7382 m.down_exps.row_bytes,
7383 )?;
7384 Ok(true)
7385 }
7386
7387 fn moe_cached_gemm_q8(
7392 e: &Engine,
7393 il: u16,
7394 proj: u8,
7395 ex: usize,
7396 m: &MoeWeights,
7397 max_block: usize,
7398 aq: &CudaSlice<i8>,
7399 ad: &CudaSlice<f32>,
7400 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7401 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7402 let exps = match proj {
7403 PROJ_GATE => &m.gate_exps,
7404 PROJ_UP => &m.up_exps,
7405 _ => &m.down_exps,
7406 };
7407 let layout = exps.expert_layout(ex);
7408 let id = BlockId::new(il, proj, ex as u16);
7409 let source = exps.expert_source(ex);
7410 e.with_moe_cache(max_block, |c, eng| {
7411 let slot = c.dispatch_source(id, source, eng)?;
7412 let DispatchSlot::Resident(sl) = slot;
7413 let buf = c.slot(sl);
7414 eng.qmatvec_expert_q8(
7415 buf,
7416 0..layout.len,
7417 aq,
7418 ad,
7419 1,
7420 exps.in_f,
7421 exps.out_f,
7422 layout.qtype,
7423 layout.row_bytes,
7424 )
7425 })
7426 }
7427
7428 fn moe_cached_gemm(
7429 e: &Engine,
7430 il: u16,
7431 proj: u8,
7432 ex: usize,
7433 m: &MoeWeights,
7434 max_block: usize,
7435 x: &cudarc::driver::CudaView<f32>,
7436 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7437 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7438 let exps = match proj {
7439 PROJ_GATE => &m.gate_exps,
7440 PROJ_UP => &m.up_exps,
7441 _ => &m.down_exps,
7442 };
7443 let layout = exps.expert_layout(ex);
7444 let id = BlockId::new(il, proj, ex as u16);
7445 let source = exps.expert_source(ex);
7446 e.with_moe_cache(max_block, |c, eng| {
7448 let slot = c.dispatch_source(id, source, eng)?;
7449 let DispatchSlot::Resident(sl) = slot;
7452 let buf = c.slot(sl);
7453 eng.qmatvec_view(
7454 buf,
7455 0..layout.len,
7456 x,
7457 1,
7458 exps.in_f,
7459 exps.out_f,
7460 layout.qtype,
7461 layout.row_bytes,
7462 )
7463 })
7464 }
7465
7466 fn moe_profile_admit_expert(
7470 e: &Engine,
7471 il: u16,
7472 ex: usize,
7473 m: &MoeWeights,
7474 max_block: usize,
7475 ) -> Result<(), Box<dyn std::error::Error>> {
7476 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7477 e.with_moe_cache(max_block, |cache, eng| {
7478 for (proj, exps) in [
7479 (PROJ_GATE, &m.gate_exps),
7480 (PROJ_UP, &m.up_exps),
7481 (PROJ_DOWN, &m.down_exps),
7482 ] {
7483 let id = BlockId::new(il, proj, ex as u16);
7484 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7485 }
7486 Ok(())
7487 })
7488 }
7489
7490 #[allow(clippy::too_many_arguments)]
7493 fn moe_frozen_gemm(
7494 e: &Engine,
7495 il: u16,
7496 proj: u8,
7497 ex: usize,
7498 m: &MoeWeights,
7499 max_block: usize,
7500 x: &cudarc::driver::CudaView<f32>,
7501 scratch: &mut Option<CudaSlice<u8>>,
7502 scratch_len: usize,
7503 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7504 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7505 let exps = match proj {
7506 PROJ_GATE => &m.gate_exps,
7507 PROJ_UP => &m.up_exps,
7508 _ => &m.down_exps,
7509 };
7510 let layout = exps.expert_layout(ex);
7511 let id = BlockId::new(il, proj, ex as u16);
7512 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7513 let Some(slot) = cache.resident(id) else {
7514 return Ok(None);
7515 };
7516 let buf = cache.slot(slot);
7517 Ok(Some(eng.qmatvec_view(
7518 buf,
7519 0..layout.len,
7520 x,
7521 1,
7522 exps.in_f,
7523 exps.out_f,
7524 layout.qtype,
7525 layout.row_bytes,
7526 )?))
7527 })? {
7528 return Ok(output);
7529 }
7530 if scratch.is_none() {
7531 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7532 }
7533 let scratch = scratch.as_mut().unwrap();
7534 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7535 e.qmatvec_view(
7536 scratch,
7537 0..layout.len,
7538 x,
7539 1,
7540 exps.in_f,
7541 exps.out_f,
7542 layout.qtype,
7543 layout.row_bytes,
7544 )
7545 }
7546
7547 fn moe_prefetch_expert(
7548 e: &Engine,
7549 il: u16,
7550 ex: usize,
7551 m: &MoeWeights,
7552 max_block: usize,
7553 keep: &[crate::moe_cache::BlockId],
7554 ) -> Result<(), Box<dyn std::error::Error>> {
7555 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7556 e.with_moe_cache(max_block, |c, eng| {
7557 for (proj, exps) in [
7558 (PROJ_GATE, &m.gate_exps),
7559 (PROJ_UP, &m.up_exps),
7560 (PROJ_DOWN, &m.down_exps),
7561 ] {
7562 let id = BlockId::new(il, proj, ex as u16);
7563 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7564 }
7565 Ok(())
7566 })
7567 }
7568
7569 fn moe_prefetch_disk_expert(
7572 e: &Engine,
7573 il: u16,
7574 ex: usize,
7575 m: &MoeWeights,
7576 max_block: usize,
7577 keep: &[crate::moe_cache::BlockId],
7578 ) -> Result<(), Box<dyn std::error::Error>> {
7579 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7580 e.with_moe_cache(max_block, |c, eng| {
7581 for (proj, exps) in [
7582 (PROJ_GATE, &m.gate_exps),
7583 (PROJ_UP, &m.up_exps),
7584 (PROJ_DOWN, &m.down_exps),
7585 ] {
7586 let source = exps.expert_source(ex);
7587 if let crate::model::ExpertSource::Disk { .. } = &source {
7588 let id = BlockId::new(il, proj, ex as u16);
7589 let _ = c.prefetch_source(id, source, keep, eng)?;
7590 }
7591 }
7592 Ok(())
7593 })
7594 }
7595
7596 #[inline]
7597 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7598 let _ = m.gate_exps.prefetch_expert_pages(ex);
7599 let _ = m.up_exps.prefetch_expert_pages(ex);
7600 let _ = m.down_exps.prefetch_expert_pages(ex);
7601 }
7602}
7603
7604impl HybridModel {
7621 #[allow(clippy::too_many_arguments)]
7625 fn moe_ffn_grouped_resident_q8(
7626 e: &Engine,
7627 m: &MoeWeights,
7628 z: &CudaSlice<f32>,
7629 t: usize,
7630 cfg: &ModelConfig,
7631 il: u16,
7632 sel_all: &[u32],
7633 w_all: &[f32],
7634 table: &CudaSlice<u64>,
7635 gu_il: bool,
7636 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7637 let moe = cfg.moe.as_ref().unwrap();
7638 let n_embd = cfg.n_embd as usize;
7639 let n_expert = moe.expert_count as usize;
7640 let n_used = moe.expert_used_count as usize;
7641 let n_ff_exp = moe.expert_ff_length as usize;
7642 let n_pairs = t * n_used;
7643 debug_assert_eq!(sel_all.len(), n_pairs);
7644 debug_assert_eq!(w_all.len(), n_pairs);
7645 debug_assert!(
7646 m.gate_exps.macros.is_none()
7647 && m.up_exps.macros.is_none()
7648 && m.down_exps.macros.is_none(),
7649 "resident grouped q8 does not fold per-expert macro scales",
7650 );
7651
7652 if !cfg.swiglu_clamped_at(il as u32) {
7658 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7659 let sel_d = e.htod_i32(&sel)?;
7660 let w_d = e.htod(w_all)?;
7661 let (gate_row_bytes, up_row_bytes) = if gu_il {
7662 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7663 (combined, combined)
7664 } else {
7665 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7666 };
7667 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7668 let act = e.moe_gate_up_silu8_dev_q8_rows(
7669 table,
7670 &sel_d,
7671 &zq,
7672 &zd,
7673 t,
7674 n_embd,
7675 n_ff_exp,
7676 n_used,
7677 n_expert,
7678 m.gate_exps.qtype,
7679 m.up_exps.qtype,
7680 gate_row_bytes,
7681 up_row_bytes,
7682 &m.dev_macros,
7683 )?;
7684 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7685 let mut moe_out = e.uninit(t * n_embd)?;
7686 e.moe_down8_fma_dev_q8_rows_g(
7687 table,
7688 &sel_d,
7689 &w_d,
7690 &aq2,
7691 &ad2,
7692 &mut moe_out,
7693 t,
7694 n_ff_exp,
7695 n_embd,
7696 n_used,
7697 n_expert,
7698 m.down_exps.qtype,
7699 m.down_exps.row_bytes,
7700 )?;
7701
7702 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7703 let mut counts = vec![0usize; n_expert];
7704 for &expert in sel_all {
7705 counts[expert as usize] += 1;
7706 }
7707 let mut sizes: Vec<usize> =
7708 counts.into_iter().filter(|&count| count != 0).collect();
7709 sizes.sort_unstable();
7710 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7711 println!(
7712 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7713 m_e: min={} median={} mean={mean:.1} max={}",
7714 sizes.len(),
7715 n_expert,
7716 sizes.first().copied().unwrap_or(0),
7717 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7718 sizes.last().copied().unwrap_or(0),
7719 );
7720 }
7721 return Ok(moe_out);
7722 }
7723
7724 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7728 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7729 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7730 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7731
7732 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7733 for (pair, &expert) in pair_ex.iter().enumerate() {
7734 by_expert[expert as usize].push(pair as i32);
7735 }
7736
7737 let pair_tok_d = e.htod_i32(&pair_tok)?;
7738 let pair_ex_d = e.htod_i32(&pair_ex)?;
7739 let pair_w_d = e.htod(w_all)?;
7740 let tok_off_d = e.htod_i32(&tok_off)?;
7741 let tok_ids_d = e.htod_i32(&tok_ids)?;
7742
7743 let matvec = |proj: i32,
7744 pair_rows: &CudaSlice<i32>,
7745 aq: &CudaSlice<i8>,
7746 ad: &CudaSlice<f32>,
7747 in_f: usize,
7748 out_f: usize,
7749 qtype: i32,
7750 row_bytes: usize|
7751 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7752 e.moe_pairs_matvec_q8(
7753 table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7754 row_bytes,
7755 )
7756 };
7757
7758 let (gate_row_bytes, up_row_bytes) = if gu_il {
7759 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7760 (combined, combined)
7761 } else {
7762 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7763 };
7764 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7765 let gate = matvec(
7766 0,
7767 &pair_tok_d,
7768 &zq,
7769 &zd,
7770 n_embd,
7771 n_ff_exp,
7772 m.gate_exps.qtype,
7773 gate_row_bytes,
7774 )?;
7775 let up = matvec(
7776 1,
7777 &pair_tok_d,
7778 &zq,
7779 &zd,
7780 n_embd,
7781 n_ff_exp,
7782 m.up_exps.qtype,
7783 up_row_bytes,
7784 )?;
7785 let mut act = e.uninit(n_pairs * n_ff_exp)?;
7786 Self::ffn_act_lim(
7787 e,
7788 cfg,
7789 &gate,
7790 &up,
7791 1.0,
7792 1.0,
7793 cfg.clamp_exp_at(il as u32),
7794 &mut act,
7795 n_pairs * n_ff_exp,
7796 )?;
7797 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7798 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7799 let pair_self_d = e.htod_i32(&pair_self)?;
7800 let down = matvec(
7801 2,
7802 &pair_self_d,
7803 &aq2,
7804 &ad2,
7805 n_ff_exp,
7806 n_embd,
7807 m.down_exps.qtype,
7808 m.down_exps.row_bytes,
7809 )?;
7810 let mut moe_out = e.uninit(t * n_embd)?;
7811 e.moe_pairs_scatter(
7812 &down,
7813 &pair_w_d,
7814 &tok_off_d,
7815 &tok_ids_d,
7816 &mut moe_out,
7817 t,
7818 n_embd,
7819 )?;
7820
7821 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7822 let mut sizes: Vec<usize> = by_expert
7823 .iter()
7824 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7825 .collect();
7826 sizes.sort_unstable();
7827 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7828 println!(
7829 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7830 m_e: min={} median={} mean={mean:.1} max={}",
7831 sizes.len(),
7832 n_expert,
7833 sizes.first().copied().unwrap_or(0),
7834 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7835 sizes.last().copied().unwrap_or(0),
7836 );
7837 }
7838 Ok(moe_out)
7839 }
7840
7841 fn moe_ffn_grouped_add_shared(
7842 e: &Engine,
7843 m: &MoeWeights,
7844 z: &CudaSlice<f32>,
7845 t: usize,
7846 cfg: &ModelConfig,
7847 il: u16,
7848 moe_out: &mut CudaSlice<f32>,
7849 ) -> Result<(), Box<dyn std::error::Error>> {
7850 let n_embd = cfg.n_embd as usize;
7851 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7852 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7853 {
7854 let n_ff_sh = gate_shexp.out_features();
7855 let sg_gate = e.matmul(gate_shexp, z, t)?;
7856 let sg_up = e.matmul(up_shexp, z, t)?;
7857 let mut sa = e.uninit(t * n_ff_sh)?;
7858 Self::ffn_act_lim(
7859 e,
7860 cfg,
7861 &sg_gate,
7862 &sg_up,
7863 1.0,
7864 1.0,
7865 cfg.clamp_shexp_at(il as u32),
7866 &mut sa,
7867 t * n_ff_sh,
7868 )?;
7869 let sh = e.matmul(down_shexp, &sa, t)?;
7870 let gate = match &m.gate_inp_shexp {
7871 Some(gate_inp_shexp) => {
7872 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7873 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7874 } else {
7875 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7876 let mut gate = e.uninit(t)?;
7877 e.sigmoid(&raw, &mut gate, t)?;
7878 gate
7879 }
7880 }
7881 None => e.htod(&vec![1.0f32; t])?,
7882 };
7883 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7884 }
7885 Ok(())
7886 }
7887
7888 pub(crate) fn moe_ffn_grouped(
7891 e: &Engine,
7892 m: &MoeWeights,
7893 z: &CudaSlice<f32>,
7894 t: usize,
7895 cfg: &ModelConfig,
7896 il: u16,
7897 max_block: usize,
7898 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7899 let moe = cfg.moe.as_ref().unwrap();
7900 let n_embd = cfg.n_embd as usize;
7901 let n_expert = moe.expert_count as usize;
7902 let n_used = moe.expert_used_count as usize;
7903 let n_ff_exp = moe.expert_ff_length as usize;
7904 let lim_exp = cfg.clamp_exp_at(il as u32);
7906
7907 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7911 if let Some(sig) = cfg.sigmoid_router() {
7912 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7913 }
7914 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7915 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7916 } else {
7917 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7918 };
7919 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7920 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7921 Self::trace_moe_input(e, il, t, n_embd, z)?;
7922
7923 let no_exp_macros = m.gate_exps.macros.is_none()
7928 && m.up_exps.macros.is_none()
7929 && m.down_exps.macros.is_none();
7930 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7931 m.has_uniform_expert_layout()
7932 && no_exp_macros
7933 && moe_q8_enabled()
7934 && q8_expert_supported(m.gate_exps.qtype)
7935 && q8_expert_supported(m.up_exps.qtype)
7936 && q8_expert_supported(m.down_exps.qtype)
7937 && moe_slab_enabled()
7938 && dev.dev == e.ctx().ordinal()
7939 });
7940 if let Some(dev) = resident_q8 {
7941 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7942 e,
7943 m,
7944 z,
7945 t,
7946 cfg,
7947 il,
7948 &sel_all,
7949 &w_all,
7950 &dev.ptr_row,
7951 dev.gu_il,
7952 )?;
7953 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7954 return Ok(moe_out);
7955 }
7956
7957 struct ExpertGroup {
7961 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
7965 let mut groups: Vec<ExpertGroup> = (0..n_expert)
7966 .map(|_| ExpertGroup {
7967 tok_indices: Vec::new(),
7968 slot_indices: Vec::new(),
7969 weights: Vec::new(),
7970 })
7971 .collect();
7972
7973 for tok in 0..t {
7974 for j in 0..n_used {
7975 let ex = sel_all[tok * n_used + j] as usize;
7976 let w = w_all[tok * n_used + j];
7977 groups[ex].tok_indices.push(tok as i32);
7978 groups[ex].slot_indices.push(j as i32);
7979 groups[ex].weights.push(w);
7980 }
7981 }
7982
7983 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7986 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
7990 let u_len = m.up_exps.max_expert_bytes();
7991 let d_len = m.down_exps.max_expert_bytes();
7992 let moe_q8 = m.has_uniform_expert_layout()
7993 && moe_q8_enabled()
7994 && q8_expert_supported(m.gate_exps.qtype)
7995 && q8_expert_supported(m.up_exps.qtype)
7996 && q8_expert_supported(m.down_exps.qtype);
7997 let slab_local = m
8000 .dev_exps
8001 .as_ref()
8002 .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
8003 let use_cache =
8004 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
8005 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
8008
8009 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
8011 (
8012 Some(e.alloc_u8(g_len)?),
8013 Some(e.alloc_u8(u_len)?),
8014 Some(e.alloc_u8(d_len)?),
8015 )
8016 } else {
8017 (None, None, None)
8018 };
8019
8020 let mut order: Vec<usize> = (0..n_expert)
8031 .filter(|&ex| !groups[ex].tok_indices.is_empty())
8032 .collect();
8033 order.sort_by(|&a, &b| {
8034 groups[b]
8035 .tok_indices
8036 .len()
8037 .cmp(&groups[a].tok_indices.len())
8038 .then(a.cmp(&b))
8039 });
8040 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
8042 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
8043 if worker_disk_prefetch {
8044 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
8045 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
8046 }
8047 }
8048 for (order_pos, &ex) in order.iter().enumerate() {
8049 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
8050 Self::moe_prefetch_host_expert(order[next], m);
8051 }
8052 if worker_disk_prefetch {
8053 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
8054 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8055 let keep = [
8056 BlockId::new(il, PROJ_GATE, ex as u16),
8057 BlockId::new(il, PROJ_UP, ex as u16),
8058 BlockId::new(il, PROJ_DOWN, ex as u16),
8059 ];
8060 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
8061 }
8062 }
8063 let grp = &groups[ex];
8064 let m_e = grp.tok_indices.len();
8065 m_dist.push(m_e);
8066 let gl = m.gate_exps.expert_layout(ex);
8067 let ul = m.up_exps.expert_layout(ex);
8068 let dl = m.down_exps.expert_layout(ex);
8069
8070 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
8074 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
8075 let dmac = m.down_exps.macro_scale(ex);
8076 let weight_d = if dmac == 1.0 {
8077 e.htod(&grp.weights)?
8078 } else {
8079 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
8080 e.htod(&scaled)?
8081 };
8082
8083 let mut gathered = e.zeros(m_e * n_embd)?;
8085 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
8086 let gv = gathered.slice(0..m_e * n_embd);
8087
8088 let y = if let Some(dev) = slab_local {
8091 let gate_start = ex * m.gate_exps.expert_stride;
8092 let up_start = ex * m.up_exps.expert_stride;
8093 let down_start = ex * m.down_exps.expert_stride;
8094 if grouped_q8 {
8095 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8096 let gate = e.qmatvec_expert_q8(
8097 &dev.gate,
8098 gate_start..gate_start + gl.len,
8099 &zq,
8100 &zd,
8101 m_e,
8102 m.gate_exps.in_f,
8103 m.gate_exps.out_f,
8104 gl.qtype,
8105 gl.row_bytes,
8106 )?;
8107 let up = e.qmatvec_expert_q8(
8108 &dev.up,
8109 up_start..up_start + ul.len,
8110 &zq,
8111 &zd,
8112 m_e,
8113 m.up_exps.in_f,
8114 m.up_exps.out_f,
8115 ul.qtype,
8116 ul.row_bytes,
8117 )?;
8118 let mut act = e.uninit(m_e * n_ff_exp)?;
8119 Self::ffn_act_lim(
8120 e,
8121 cfg,
8122 &gate,
8123 &up,
8124 m.gate_exps.macro_scale(ex),
8125 m.up_exps.macro_scale(ex),
8126 lim_exp,
8127 &mut act,
8128 m_e * n_ff_exp,
8129 )?;
8130 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8131 e.qmatvec_expert_q8(
8132 &dev.down,
8133 down_start..down_start + dl.len,
8134 &aq2,
8135 &ad2,
8136 m_e,
8137 m.down_exps.in_f,
8138 m.down_exps.out_f,
8139 dl.qtype,
8140 dl.row_bytes,
8141 )?
8142 } else {
8143 let gate = e.qmatvec_view(
8144 &dev.gate,
8145 gate_start..gate_start + gl.len,
8146 &gv,
8147 m_e,
8148 m.gate_exps.in_f,
8149 m.gate_exps.out_f,
8150 gl.qtype,
8151 gl.row_bytes,
8152 )?;
8153 let up = e.qmatvec_view(
8154 &dev.up,
8155 up_start..up_start + ul.len,
8156 &gv,
8157 m_e,
8158 m.up_exps.in_f,
8159 m.up_exps.out_f,
8160 ul.qtype,
8161 ul.row_bytes,
8162 )?;
8163 let mut act = e.uninit(m_e * n_ff_exp)?;
8164 Self::ffn_act_lim(
8165 e,
8166 cfg,
8167 &gate,
8168 &up,
8169 m.gate_exps.macro_scale(ex),
8170 m.up_exps.macro_scale(ex),
8171 lim_exp,
8172 &mut act,
8173 m_e * n_ff_exp,
8174 )?;
8175 let actv = act.slice(0..m_e * n_ff_exp);
8176 e.qmatvec_view(
8177 &dev.down,
8178 down_start..down_start + dl.len,
8179 &actv,
8180 m_e,
8181 m.down_exps.in_f,
8182 m.down_exps.out_f,
8183 dl.qtype,
8184 dl.row_bytes,
8185 )?
8186 }
8187 } else if use_cache {
8188 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8189 if grouped_q8 {
8190 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8191 let gate = e.with_moe_cache(max_block, |cache, eng| {
8192 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8193 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8194 eng.qmatvec_expert_q8(
8195 cache.buf(slot),
8196 0..gl.len,
8197 &zq,
8198 &zd,
8199 m_e,
8200 m.gate_exps.in_f,
8201 m.gate_exps.out_f,
8202 gl.qtype,
8203 gl.row_bytes,
8204 )
8205 })?;
8206 let up = e.with_moe_cache(max_block, |cache, eng| {
8207 let id = BlockId::new(il, PROJ_UP, ex as u16);
8208 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8209 eng.qmatvec_expert_q8(
8210 cache.buf(slot),
8211 0..ul.len,
8212 &zq,
8213 &zd,
8214 m_e,
8215 m.up_exps.in_f,
8216 m.up_exps.out_f,
8217 ul.qtype,
8218 ul.row_bytes,
8219 )
8220 })?;
8221 let mut act = e.uninit(m_e * n_ff_exp)?;
8222 Self::ffn_act_lim(
8223 e,
8224 cfg,
8225 &gate,
8226 &up,
8227 m.gate_exps.macro_scale(ex),
8228 m.up_exps.macro_scale(ex),
8229 lim_exp,
8230 &mut act,
8231 m_e * n_ff_exp,
8232 )?;
8233 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8234 e.with_moe_cache(max_block, |cache, eng| {
8235 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8236 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8237 eng.qmatvec_expert_q8(
8238 cache.buf(slot),
8239 0..dl.len,
8240 &aq2,
8241 &ad2,
8242 m_e,
8243 m.down_exps.in_f,
8244 m.down_exps.out_f,
8245 dl.qtype,
8246 dl.row_bytes,
8247 )
8248 })?
8249 } else {
8250 let gate = e.with_moe_cache(max_block, |cache, eng| {
8251 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8252 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8253 eng.qmatvec_view(
8254 cache.buf(slot),
8255 0..gl.len,
8256 &gv,
8257 m_e,
8258 m.gate_exps.in_f,
8259 m.gate_exps.out_f,
8260 gl.qtype,
8261 gl.row_bytes,
8262 )
8263 })?;
8264 let up = e.with_moe_cache(max_block, |cache, eng| {
8265 let id = BlockId::new(il, PROJ_UP, ex as u16);
8266 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8267 eng.qmatvec_view(
8268 cache.buf(slot),
8269 0..ul.len,
8270 &gv,
8271 m_e,
8272 m.up_exps.in_f,
8273 m.up_exps.out_f,
8274 ul.qtype,
8275 ul.row_bytes,
8276 )
8277 })?;
8278 let mut act = e.uninit(m_e * n_ff_exp)?;
8279 Self::ffn_act_lim(
8280 e,
8281 cfg,
8282 &gate,
8283 &up,
8284 m.gate_exps.macro_scale(ex),
8285 m.up_exps.macro_scale(ex),
8286 lim_exp,
8287 &mut act,
8288 m_e * n_ff_exp,
8289 )?;
8290 let actv = act.slice(0..m_e * n_ff_exp);
8291 e.with_moe_cache(max_block, |cache, eng| {
8292 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8293 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8294 eng.qmatvec_view(
8295 cache.buf(slot),
8296 0..dl.len,
8297 &actv,
8298 m_e,
8299 m.down_exps.in_f,
8300 m.down_exps.out_f,
8301 dl.qtype,
8302 dl.row_bytes,
8303 )
8304 })?
8305 }
8306 } else {
8307 let sg = scratch_g.as_mut().unwrap();
8308 let su = scratch_u.as_mut().unwrap();
8309 let sd = scratch_d.as_mut().unwrap();
8310 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8311 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8312 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8313 if grouped_q8 {
8314 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8315 let gate = e.qmatvec_expert_q8(
8316 sg,
8317 0..gl.len,
8318 &zq,
8319 &zd,
8320 m_e,
8321 m.gate_exps.in_f,
8322 m.gate_exps.out_f,
8323 gl.qtype,
8324 gl.row_bytes,
8325 )?;
8326 let up = e.qmatvec_expert_q8(
8327 su,
8328 0..ul.len,
8329 &zq,
8330 &zd,
8331 m_e,
8332 m.up_exps.in_f,
8333 m.up_exps.out_f,
8334 ul.qtype,
8335 ul.row_bytes,
8336 )?;
8337 let mut act = e.uninit(m_e * n_ff_exp)?;
8338 Self::ffn_act_lim(
8339 e,
8340 cfg,
8341 &gate,
8342 &up,
8343 m.gate_exps.macro_scale(ex),
8344 m.up_exps.macro_scale(ex),
8345 lim_exp,
8346 &mut act,
8347 m_e * n_ff_exp,
8348 )?;
8349 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8350 e.qmatvec_expert_q8(
8351 sd,
8352 0..dl.len,
8353 &aq2,
8354 &ad2,
8355 m_e,
8356 m.down_exps.in_f,
8357 m.down_exps.out_f,
8358 dl.qtype,
8359 dl.row_bytes,
8360 )?
8361 } else {
8362 let gate = e.qmatvec_view(
8363 sg,
8364 0..gl.len,
8365 &gv,
8366 m_e,
8367 m.gate_exps.in_f,
8368 m.gate_exps.out_f,
8369 gl.qtype,
8370 gl.row_bytes,
8371 )?;
8372 let up = e.qmatvec_view(
8373 su,
8374 0..ul.len,
8375 &gv,
8376 m_e,
8377 m.up_exps.in_f,
8378 m.up_exps.out_f,
8379 ul.qtype,
8380 ul.row_bytes,
8381 )?;
8382 let mut act = e.uninit(m_e * n_ff_exp)?;
8383 Self::ffn_act_lim(
8384 e,
8385 cfg,
8386 &gate,
8387 &up,
8388 m.gate_exps.macro_scale(ex),
8389 m.up_exps.macro_scale(ex),
8390 lim_exp,
8391 &mut act,
8392 m_e * n_ff_exp,
8393 )?;
8394 let actv = act.slice(0..m_e * n_ff_exp);
8395 e.qmatvec_view(
8396 sd,
8397 0..dl.len,
8398 &actv,
8399 m_e,
8400 m.down_exps.in_f,
8401 m.down_exps.out_f,
8402 dl.qtype,
8403 dl.row_bytes,
8404 )?
8405 }
8406 };
8407
8408 e.scatter_slot(
8410 &y,
8411 &tok_idx_d,
8412 &slot_idx_d,
8413 &weight_d,
8414 &mut slot_buf,
8415 &mut wbuf,
8416 n_embd,
8417 n_used,
8418 m_e,
8419 )?;
8420 }
8421
8422 let mut moe_out = e.zeros(t * n_embd)?;
8424 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8425
8426 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8428 m_dist.sort_unstable();
8429 let active = m_dist.len();
8430 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8431 let median = m_dist[active / 2];
8432 let max_m = *m_dist.last().unwrap();
8433 let min_m = m_dist[0];
8434 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8435 println!(
8436 "moe-grouped il={il} t={t} active={active}/{n_expert} \
8437 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8438 above_gemm_threshold(>=16)={above16}/{active}"
8439 );
8440 }
8441
8442 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8443 Ok(moe_out)
8444 }
8445
8446 pub(crate) fn moe_ffn_lockstep(
8453 &self,
8454 e: &Engine,
8455 m: &MoeWeights,
8456 zbatch: &CudaSlice<f32>,
8457 mrows: usize,
8458 il: u16,
8459 max_block: usize,
8460 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8461 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8462 let cfg = &self.cfg;
8463 let moe = cfg.moe.as_ref().unwrap();
8464 let n_embd = cfg.n_embd as usize;
8465 let n_expert = moe.expert_count as usize;
8466 let n_used = moe.expert_used_count as usize;
8467 let n_ff_exp = moe.expert_ff_length as usize;
8468 let lim_exp = cfg.clamp_exp_at(il as u32);
8470 let lim_shexp = cfg.clamp_shexp_at(il as u32);
8471
8472 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8473 if let Some(sig) = cfg.sigmoid_router() {
8474 Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8475 }
8476 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8477 Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8478 } else {
8479 Self::moe_route_cfg(
8480 e,
8481 &logits,
8482 mrows,
8483 n_expert,
8484 n_used,
8485 m.active_experts.as_deref(),
8486 )?
8487 };
8488 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8489
8490 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8492 Ok((0..n_expert)
8493 .map(|ex| {
8494 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8495 .into_iter()
8496 .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8497 })
8498 .collect())
8499 })?;
8500
8501 struct Group {
8502 rows: Vec<i32>,
8503 slots: Vec<i32>,
8504 weights: Vec<f32>,
8505 }
8506 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8507 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8508 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8509 Default::default();
8510 for row in 0..mrows {
8511 for j in 0..n_used {
8512 let ex = sel_all[row * n_used + j] as usize;
8513 let w = w_all[row * n_used + j];
8514 if resident_expert[ex] {
8515 let group = groups.entry(ex).or_insert_with(|| Group {
8516 rows: Vec::new(),
8517 slots: Vec::new(),
8518 weights: Vec::new(),
8519 });
8520 group.rows.push(row as i32);
8521 group.slots.push(j as i32);
8522 group.weights.push(w);
8523 } else {
8524 crate::cpu_experts::record_incomplete_gpu_residency(0);
8525 cpu_rows[row].push((ex, w));
8526 cpu_by_expert.entry(ex).or_default().push((row, w));
8527 }
8528 }
8529 }
8530
8531 let host_rows = e.dtoh(zbatch)?;
8537 let rows_ok = crate::cpu_experts::rows_supported();
8538 enum CpuPart {
8539 Single { row: usize },
8540 Rows { rows: Vec<usize> },
8541 }
8542 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8543 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8544 if rows_ok {
8545 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8546 .into_iter()
8547 .filter(|(_, rows)| rows.len() >= 2)
8548 .collect();
8549 shared.sort_by_key(|(ex, _)| *ex);
8550 for (ex, mut row_weights) in shared {
8551 row_weights.sort_by_key(|(row, _)| *row);
8552 let inputs: Vec<(&[f32], f32)> = row_weights
8553 .iter()
8554 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8555 .collect();
8556 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8557 .map_err(std::io::Error::other)?;
8558 for &(row, _) in &row_weights {
8559 rows_served.insert((row, ex));
8560 }
8561 tickets.push((
8562 CpuPart::Rows {
8563 rows: row_weights.iter().map(|&(row, _)| row).collect(),
8564 },
8565 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8566 ));
8567 }
8568 }
8569 for (row, selected) in cpu_rows.iter().enumerate() {
8570 let leftover: Vec<(usize, f32)> = selected
8571 .iter()
8572 .copied()
8573 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8574 .collect();
8575 if leftover.is_empty() {
8576 continue;
8577 }
8578 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8579 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8580 .map_err(std::io::Error::other)?;
8581 tickets.push((
8582 CpuPart::Single { row },
8583 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8584 ));
8585 }
8586
8587 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8588 let mut wbuf = e.zeros(mrows * n_used)?;
8589 let mut order: Vec<usize> = groups.keys().copied().collect();
8590 order.sort_by(|&a, &b| {
8591 groups[&b]
8592 .rows
8593 .len()
8594 .cmp(&groups[&a].rows.len())
8595 .then(a.cmp(&b))
8596 });
8597 for &ex in &order {
8598 let group = &groups[&ex];
8599 let m_e = group.rows.len();
8600 let gl = m.gate_exps.expert_layout(ex);
8601 let ul = m.up_exps.expert_layout(ex);
8602 let dl = m.down_exps.expert_layout(ex);
8603 let row_idx_d = e.htod_i32(&group.rows)?;
8604 let slot_idx_d = e.htod_i32(&group.slots)?;
8605 let dmac = m.down_exps.macro_scale(ex);
8606 let weight_d = if dmac == 1.0 {
8607 e.htod(&group.weights)?
8608 } else {
8609 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8610 e.htod(&scaled)?
8611 };
8612 let mut gathered = e.zeros(m_e * n_embd)?;
8613 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8614 let gv = gathered.slice(0..m_e * n_embd);
8615 let gate = e.with_moe_cache(max_block, |c, eng| {
8616 let slot = c
8617 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8618 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8619 eng.qmatvec_view(
8620 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8621 0..gl.len,
8622 &gv,
8623 m_e,
8624 m.gate_exps.in_f,
8625 m.gate_exps.out_f,
8626 gl.qtype,
8627 gl.row_bytes,
8628 )
8629 })?;
8630 let up = e.with_moe_cache(max_block, |c, eng| {
8631 let slot = c
8632 .resident(BlockId::new(il, PROJ_UP, ex as u16))
8633 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8634 eng.qmatvec_view(
8635 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8636 0..ul.len,
8637 &gv,
8638 m_e,
8639 m.up_exps.in_f,
8640 m.up_exps.out_f,
8641 ul.qtype,
8642 ul.row_bytes,
8643 )
8644 })?;
8645 let mut act = e.zeros(m_e * n_ff_exp)?;
8646 Self::ffn_act_lim(
8647 e,
8648 cfg,
8649 &gate,
8650 &up,
8651 m.gate_exps.macro_scale(ex),
8652 m.up_exps.macro_scale(ex),
8653 lim_exp,
8654 &mut act,
8655 m_e * n_ff_exp,
8656 )?;
8657 let actv = act.slice(0..m_e * n_ff_exp);
8658 let y = e.with_moe_cache(max_block, |c, eng| {
8659 let slot = c
8660 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8661 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8662 eng.qmatvec_view(
8663 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8664 0..dl.len,
8665 &actv,
8666 m_e,
8667 m.down_exps.in_f,
8668 m.down_exps.out_f,
8669 dl.qtype,
8670 dl.row_bytes,
8671 )
8672 })?;
8673 e.scatter_slot(
8674 &y,
8675 &row_idx_d,
8676 &slot_idx_d,
8677 &weight_d,
8678 &mut slot_buf,
8679 &mut wbuf,
8680 n_embd,
8681 n_used,
8682 m_e,
8683 )?;
8684 }
8685 let mut moe_out = e.zeros(mrows * n_embd)?;
8686 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8687
8688 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8690 for (part, ticket) in tickets {
8691 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8692 let mut add_row = |row: usize, chunk: &[f32]| {
8693 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8694 for (accumulator, value) in sum.iter_mut().zip(chunk) {
8695 *accumulator += value;
8696 }
8697 };
8698 match part {
8699 CpuPart::Single { row } => add_row(row, &cpu_output),
8700 CpuPart::Rows { rows } => {
8701 for (slot, row) in rows.into_iter().enumerate() {
8702 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8703 }
8704 }
8705 }
8706 }
8707 for (row, sum) in row_sums.into_iter().enumerate() {
8708 let Some(sum) = sum else { continue };
8709 let cpu_output = e.htod(&sum)?;
8710 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8711 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8712 }
8713
8714 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8715 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8716 {
8717 let n_ff_sh = gate_shexp.out_features();
8718 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8719 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8720 let mut sa = e.zeros(mrows * n_ff_sh)?;
8721 Self::ffn_act_lim(
8722 e,
8723 cfg,
8724 &sg_gate,
8725 &sg_up,
8726 1.0,
8727 1.0,
8728 lim_shexp,
8729 &mut sa,
8730 mrows * n_ff_sh,
8731 )?;
8732 let sh = e.matmul(down_shexp, &sa, mrows)?;
8733 let g = match &m.gate_inp_shexp {
8736 Some(gate_inp_shexp) => {
8737 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8738 }
8739 None => e.htod(&vec![1.0f32; mrows])?,
8740 };
8741 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8742 }
8743
8744 Ok(moe_out)
8745 }
8746}
8747
8748impl HybridModel {
8754 pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
8768 let g = self
8769 .cfg
8770 .gemma4
8771 .as_ref()
8772 .expect("gemma4_rope_dims on a non-gemma4 config");
8773 if g.swa_pattern[il] {
8774 g.rope_dims_swa as usize
8775 } else {
8776 g.rope_dims_global as usize
8777 }
8778 }
8779
8780 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8781 let g = self.cfg.gemma4.as_ref().unwrap();
8782 let swa = g.swa_pattern[il];
8783 let hd = if swa {
8784 g.key_length_swa
8785 } else {
8786 g.key_length_global
8787 } as usize;
8788 (
8792 hd,
8793 g.head_count_kv[il] as usize,
8794 self.cfg.n_head as usize,
8795 if swa {
8796 g.rope_base_swa
8797 } else {
8798 g.rope_base_global
8799 },
8800 1.0,
8801 swa,
8802 )
8803 }
8804
8805 pub(crate) fn gemma4_suppress(
8809 &self,
8810 e: &Engine,
8811 ld: &mut CudaSlice<f32>,
8812 t: usize,
8813 ) -> Result<(), Box<dyn std::error::Error>> {
8814 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8815 #[cfg(debug_assertions)]
8820 crate::debug_assert_tensor_stream_device(
8821 ids,
8822 &e.stream(),
8823 "gemma4_suppress.suppress_d",
8824 );
8825 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8826 }
8827 Ok(())
8828 }
8829
8830 #[allow(clippy::too_many_arguments)]
8835 fn gemma_fa_one_program() -> bool {
8844 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8845 *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
8846 }
8847
8848 fn gemma4_attn_prime(
8849 &self,
8850 e: &Engine,
8851 fa: &crate::hybrid::FullAttnLayer,
8852 il: usize,
8853 h: &CudaSlice<f32>,
8854 pos_d: &CudaSlice<i32>,
8855 t: usize,
8856 cache: Option<&mut Cache>,
8857 island: Option<&CudaSlice<i32>>,
8858 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8859 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8860 let eps = self.cfg.rms_eps;
8861 let aux = self.gemma4_aux.as_ref().unwrap();
8862 let ones = aux.ones(e);
8863 #[cfg(debug_assertions)]
8864 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8865
8866 e.mmq_act_begin();
8869 let q0 = e.matmul(&fa.wq, h, t)?; if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8871 let v = e.dtoh(&q0)?;
8872 let nan = v.iter().filter(|x| x.is_nan()).count();
8873 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8874 eprintln!(
8875 "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
8876 v.len()
8877 );
8878 }
8879 let k0 = e.matmul(&fa.wk, h, t)?; let v0 = if swa {
8883 e.matmul(&fa.wv, h, t)?
8884 } else {
8885 e.clone_dtod(&k0)?
8886 };
8887 if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8888 for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
8889 let v = e.dtoh(buf)?;
8890 let nan = v.iter().filter(|x| x.is_nan()).count();
8891 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8892 eprintln!(
8893 "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
8894 v.len()
8895 );
8896 }
8897 }
8898
8899 let mut q = e.uninit(t * nh * hd)?;
8900 let mut k = e.uninit(t * nkv * hd)?;
8901 let mut v = e.uninit(t * nkv * hd)?;
8903 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8907 let emit = island.is_none()
8910 && t >= 16
8911 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8912 && *EMIT.get_or_init(|| {
8913 std::env::var("MEMRA_FA_EMIT")
8914 .map(|s| s != "0")
8915 .unwrap_or(true)
8916 });
8917 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8918 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8919 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8920 let v_f16 = emit
8923 && crate::fa_f16pv_on()
8924 && match hd {
8925 512 => true,
8926 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8927 _ => false,
8928 };
8929 if emit {
8930 e.rms_norm_qkv_w4b(
8931 &q0,
8932 &k0,
8933 &v0,
8934 fa.q_norm.float_data(),
8935 fa.k_norm.float_data(),
8936 ones,
8937 &mut q,
8938 &mut k,
8939 &mut v,
8940 &mut vb,
8941 hd,
8942 nh * t,
8943 nkv * t,
8944 eps,
8945 v_f16,
8946 )?;
8947 } else {
8948 e.rms_norm_qkv(
8949 &q0,
8950 &k0,
8951 &v0,
8952 fa.q_norm.float_data(),
8953 fa.k_norm.float_data(),
8954 ones,
8955 &mut q,
8956 &mut k,
8957 &mut v,
8958 hd,
8959 nh * t,
8960 nkv * t,
8961 eps,
8962 )?;
8963 }
8964
8965 let ff = if swa {
8966 None
8967 } else {
8968 Some(
8969 aux.rope_freqs(e)
8970 .expect("gemma4 global rope needs rope_freqs.weight"),
8971 )
8972 };
8973 #[cfg(debug_assertions)]
8974 if let Some(ff) = ff {
8975 crate::debug_assert_tensor_stream_device(
8976 ff,
8977 &e.stream(),
8978 "gemma4_attn_prime.rope_freqs",
8979 );
8980 }
8981 if emit {
8982 e.rope_neox2_bf16e(
8983 &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8984 )?;
8985 } else {
8986 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8987 }
8988
8989 if let Some(cache) = cache {
8990 let kvl = cache.kv[il].as_mut().unwrap();
8991 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8992 e.append_kv_quantized_rows(
8993 &k,
8994 &v,
8995 &mut kvl.k,
8996 &mut kvl.v,
8997 kvl.len,
8998 t,
8999 kvl.kv_dim_k,
9000 kvl.kv_dim_v,
9001 kvl.k_tok_bytes,
9002 kvl.v_tok_bytes,
9003 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
9004 )?;
9005 kvl.len += t;
9006 }
9007 let mut attn = e.zeros(t * nh * hd)?;
9008 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
9012 if let Some(span) = island {
9013 let w = if swa && t > win { win } else { 0 };
9018 e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
9019 } else if swa && (t > win || Self::gemma_fa_one_program()) {
9020 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
9021 if emit {
9022 e.fa_prefill_w_pre(
9023 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
9024 )?;
9025 } else {
9026 e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
9027 }
9028 } else {
9029 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
9030 }
9031 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
9032 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9033 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
9034 if emit {
9035 e.fa_prefill_hd512_pre(
9036 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
9037 )?;
9038 } else {
9039 e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9040 }
9041 } else {
9042 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9043 }
9044 Ok(e.matmul(&fa.wo, &attn, t)?)
9045 }
9046
9047 fn gemma4_attn(
9049 &self,
9050 e: &Engine,
9051 fa: &crate::hybrid::FullAttnLayer,
9052 il: usize,
9053 h: &CudaSlice<f32>,
9054 pos_d: &CudaSlice<i32>,
9055 t: usize,
9056 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9057 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
9058 }
9059
9060 fn gemma4_moe_q8(
9065 &self,
9066 e: &Engine,
9067 m: &crate::hybrid::MoeWeights,
9068 bits: &crate::hybrid::Gemma4MoeBits,
9069 mq: &(CudaSlice<i8>, CudaSlice<f32>),
9070 router_in: &CudaSlice<f32>,
9071 t: usize,
9072 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9073 let cfg = &self.cfg;
9074 let moe = cfg.moe.as_ref().unwrap();
9075 let n_embd = cfg.n_embd as usize;
9076 let n_expert = moe.expert_count as usize;
9077 let n_used = moe.expert_used_count as usize;
9078 let n_ff_exp = moe.expert_ff_length as usize;
9079 let logits = if crate::router_kernel_on() {
9083 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9084 } else {
9085 e.matmul(&m.gate_inp, router_in, t)?
9086 };
9087 let dev = m.dev_exps.as_ref().unwrap();
9088 let (sel_d, w_d) =
9089 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9090 let (zq, zd) = mq;
9091 if t == 1 {
9092 let selv = sel_d.slice(0..n_used);
9093 let wv = w_d.slice(0..n_used);
9094 let act = e.moe_gate_up_gelu8_dev_q8(
9095 &dev.ptr_row,
9096 &selv,
9097 zq,
9098 zd,
9099 n_embd,
9100 n_ff_exp,
9101 n_used,
9102 n_expert,
9103 m.gate_exps.qtype,
9104 m.up_exps.qtype,
9105 m.gate_exps.row_bytes,
9106 m.up_exps.row_bytes,
9107 )?;
9108 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9109 let mut moe_out = e.uninit(n_embd)?;
9110 e.moe_down8_fma_dev_q8(
9111 &dev.ptr_row,
9112 &selv,
9113 &wv,
9114 &aq2,
9115 &ad2,
9116 &mut moe_out.slice_mut(0..n_embd),
9117 n_ff_exp,
9118 n_embd,
9119 n_used,
9120 n_expert,
9121 m.down_exps.qtype,
9122 m.down_exps.row_bytes,
9123 )?;
9124 return Ok(moe_out);
9125 }
9126 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9127 let act = if csr {
9128 e.moe_gate_up_gelu8_dev_q8_csr(
9129 &dev.ptr_row,
9130 &sel_d,
9131 zq,
9132 zd,
9133 t * n_used,
9134 n_embd,
9135 n_ff_exp,
9136 n_used,
9137 n_expert,
9138 m.gate_exps.qtype,
9139 m.up_exps.qtype,
9140 m.gate_exps.row_bytes,
9141 m.up_exps.row_bytes,
9142 )?
9143 } else {
9144 e.moe_gate_up_gelu8_dev_q8_rows(
9145 &dev.ptr_row,
9146 &sel_d,
9147 zq,
9148 zd,
9149 t,
9150 n_embd,
9151 n_ff_exp,
9152 n_used,
9153 n_expert,
9154 m.gate_exps.qtype,
9155 m.up_exps.qtype,
9156 m.gate_exps.row_bytes,
9157 m.up_exps.row_bytes,
9158 )?
9159 };
9160 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9161 let mut moe_out = e.uninit(t * n_embd)?;
9162 e.moe_down8_fma_dev_q8_rows_g(
9165 &dev.ptr_row,
9166 &sel_d,
9167 &w_d,
9168 &aq2,
9169 &ad2,
9170 &mut moe_out,
9171 t,
9172 n_ff_exp,
9173 n_embd,
9174 n_used,
9175 n_expert,
9176 m.down_exps.qtype,
9177 m.down_exps.row_bytes,
9178 )?;
9179 Ok(moe_out)
9180 }
9181
9182 fn gemma4_moe(
9186 &self,
9187 e: &Engine,
9188 m: &crate::hybrid::MoeWeights,
9189 bits: &crate::hybrid::Gemma4MoeBits,
9190 moe_in: &CudaSlice<f32>,
9191 router_in: &CudaSlice<f32>,
9192 t: usize,
9193 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9194 let cfg = &self.cfg;
9195 let moe = cfg.moe.as_ref().unwrap();
9196 let n_embd = cfg.n_embd as usize;
9197 let n_expert = moe.expert_count as usize;
9198 let n_used = moe.expert_used_count as usize;
9199 let n_ff_exp = moe.expert_ff_length as usize;
9200
9201 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
9205 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9206 } else {
9207 e.matmul(&m.gate_inp, router_in, t)?
9208 };
9209
9210 if t < PRIME_MIN_T
9215 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9216 && expert_dp4a_supported(m.gate_exps.qtype)
9217 && expert_dp4a_supported(m.up_exps.qtype)
9218 && expert_dp4a_supported(m.down_exps.qtype)
9219 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9220 {
9221 let dev = m.dev_exps.as_ref().unwrap();
9222 let (sel_d, w_d) =
9223 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9224 if t == 1 {
9225 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
9226 let selv = sel_d.slice(0..n_used);
9227 let wv = w_d.slice(0..n_used);
9228 let act = e.moe_gate_up_gelu8_dev_q8(
9229 &dev.ptr_row,
9230 &selv,
9231 &zq,
9232 &zd,
9233 n_embd,
9234 n_ff_exp,
9235 n_used,
9236 n_expert,
9237 m.gate_exps.qtype,
9238 m.up_exps.qtype,
9239 m.gate_exps.row_bytes,
9240 m.up_exps.row_bytes,
9241 )?;
9242 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9243 let mut moe_out = e.uninit(n_embd)?;
9244 e.moe_down8_fma_dev_q8(
9245 &dev.ptr_row,
9246 &selv,
9247 &wv,
9248 &aq2,
9249 &ad2,
9250 &mut moe_out.slice_mut(0..n_embd),
9251 n_ff_exp,
9252 n_embd,
9253 n_used,
9254 n_expert,
9255 m.down_exps.qtype,
9256 m.down_exps.row_bytes,
9257 )?;
9258 return Ok(moe_out);
9259 }
9260 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9265 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9266 let act = if csr {
9267 e.moe_gate_up_gelu8_dev_q8_csr(
9268 &dev.ptr_row,
9269 &sel_d,
9270 &zq,
9271 &zd,
9272 t * n_used,
9273 n_embd,
9274 n_ff_exp,
9275 n_used,
9276 n_expert,
9277 m.gate_exps.qtype,
9278 m.up_exps.qtype,
9279 m.gate_exps.row_bytes,
9280 m.up_exps.row_bytes,
9281 )?
9282 } else {
9283 e.moe_gate_up_gelu8_dev_q8_rows(
9284 &dev.ptr_row,
9285 &sel_d,
9286 &zq,
9287 &zd,
9288 t,
9289 n_embd,
9290 n_ff_exp,
9291 n_used,
9292 n_expert,
9293 m.gate_exps.qtype,
9294 m.up_exps.qtype,
9295 m.gate_exps.row_bytes,
9296 m.up_exps.row_bytes,
9297 )?
9298 };
9299 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9300 let mut moe_out = e.uninit(t * n_embd)?;
9301 e.moe_down8_fma_dev_q8_rows_g(
9302 &dev.ptr_row,
9303 &sel_d,
9304 &w_d,
9305 &aq2,
9306 &ad2,
9307 &mut moe_out,
9308 t,
9309 n_ff_exp,
9310 n_embd,
9311 n_used,
9312 n_expert,
9313 m.down_exps.qtype,
9314 m.down_exps.row_bytes,
9315 )?;
9316 return Ok(moe_out);
9317 }
9318
9319 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9320 for (i, &sx) in sel_all.iter().enumerate() {
9321 w_all[i] *= bits.per_expert_scale[sx as usize];
9322 }
9323
9324 if t >= PRIME_MIN_T
9328 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9329 && expert_dp4a_supported(m.gate_exps.qtype)
9330 && expert_dp4a_supported(m.up_exps.qtype)
9331 && expert_dp4a_supported(m.down_exps.qtype)
9332 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9333 {
9334 let dev = m.dev_exps.as_ref().unwrap();
9335 let n_pairs = t * n_used;
9336 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9337 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9338 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9339 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9340 let pt = e.htod_i32(&pair_tok)?;
9341 let pw = e.htod(&w_all)?;
9342 let toff = e.htod_i32(&tok_off)?;
9343 let tids = e.htod_i32(&tok_ids)?;
9344 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9345 for p in 0..n_pairs {
9346 by_ex[pair_ex[p] as usize].push(p as i32);
9347 }
9348 let mut ex_ids: Vec<i32> = Vec::new();
9349 let mut ex_off: Vec<i32> = vec![0];
9350 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9351 for (ex, list) in by_ex.iter().enumerate() {
9352 if list.is_empty() {
9353 continue;
9354 }
9355 ex_ids.push(ex as i32);
9356 ex_pairs.extend_from_slice(list);
9357 ex_off.push(ex_pairs.len() as i32);
9358 }
9359 let n_active = ex_ids.len();
9360 let exi = e.htod_i32(&ex_ids)?;
9361 let exo = e.htod_i32(&ex_off)?;
9362 let exp_d = e.htod_i32(&ex_pairs)?;
9363 if crate::moe_f16g_gemma_on()
9371 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9372 && f16g_proj_ok(m.up_exps.qtype, n_embd)
9373 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9374 {
9375 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9376 let csr_tok_d = e.htod_i32(&csr_tok)?;
9377 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9378 let g_csr = e.moe_f16_grouped(
9379 &dev.ptr_row,
9380 0,
9381 n_expert,
9382 &exi,
9383 &ex_off,
9384 &exo,
9385 &z_f16,
9386 &z_s,
9387 n_embd,
9388 n_ff_exp,
9389 n_active,
9390 n_pairs,
9391 m.gate_exps.qtype,
9392 m.gate_exps.row_bytes,
9393 )?;
9394 let u_csr = e.moe_f16_grouped(
9395 &dev.ptr_row,
9396 1,
9397 n_expert,
9398 &exi,
9399 &ex_off,
9400 &exo,
9401 &z_f16,
9402 &z_s,
9403 n_embd,
9404 n_ff_exp,
9405 n_active,
9406 n_pairs,
9407 m.up_exps.qtype,
9408 m.up_exps.row_bytes,
9409 )?;
9410 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9411 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9412 let d_csr = e.moe_f16_grouped(
9413 &dev.ptr_row,
9414 2,
9415 n_expert,
9416 &exi,
9417 &ex_off,
9418 &exo,
9419 &a_f16,
9420 &a_s,
9421 n_ff_exp,
9422 n_embd,
9423 n_active,
9424 n_pairs,
9425 m.down_exps.qtype,
9426 m.down_exps.row_bytes,
9427 )?;
9428 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9429 let mut moe_out = e.uninit(t * n_embd)?;
9430 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9431 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9432 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9433 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9434 eprintln!(
9435 "[f16g-debug] post-permute bad={} post-scatter bad={}",
9436 scan(&yd),
9437 scan(&mo)
9438 );
9439 }
9440 return Ok(moe_out);
9441 }
9442 let mma =
9445 n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9446 let (gate, up) = if mma {
9447 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9448 (
9449 e.mmq_iq_experts(
9450 &dev.ptr_row,
9451 0,
9452 n_expert,
9453 &exi,
9454 &exo,
9455 &exp_d,
9456 &pt,
9457 &z_scr,
9458 n_embd,
9459 n_ff_exp,
9460 n_active,
9461 n_pairs,
9462 t,
9463 m.gate_exps.qtype,
9464 m.gate_exps.row_bytes,
9465 )?,
9466 e.mmq_iq_experts(
9467 &dev.ptr_row,
9468 1,
9469 n_expert,
9470 &exi,
9471 &exo,
9472 &exp_d,
9473 &pt,
9474 &z_scr,
9475 n_embd,
9476 n_ff_exp,
9477 n_active,
9478 n_pairs,
9479 t,
9480 m.up_exps.qtype,
9481 m.up_exps.row_bytes,
9482 )?,
9483 )
9484 } else {
9485 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9486 (
9487 e.moe_pairs_matvec_q8_dec(
9488 &dev.ptr_row,
9489 0,
9490 &exi,
9491 &exo,
9492 &exp_d,
9493 &pt,
9494 &zq,
9495 &zd,
9496 n_embd,
9497 n_ff_exp,
9498 n_expert,
9499 n_active,
9500 n_pairs,
9501 m.gate_exps.qtype,
9502 m.gate_exps.row_bytes,
9503 )?,
9504 e.moe_pairs_matvec_q8_dec(
9505 &dev.ptr_row,
9506 1,
9507 &exi,
9508 &exo,
9509 &exp_d,
9510 &pt,
9511 &zq,
9512 &zd,
9513 n_embd,
9514 n_ff_exp,
9515 n_expert,
9516 n_active,
9517 n_pairs,
9518 m.up_exps.qtype,
9519 m.up_exps.row_bytes,
9520 )?,
9521 )
9522 };
9523 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9524 let pself = e.htod_i32(&pair_self)?;
9525 let y_down = if mma {
9537 let in_pad = n_ff_exp.div_ceil(256) * 256;
9538 let a_scr = if crate::moe_fuse_actq_on() {
9539 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9540 } else {
9541 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9542 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9543 };
9544 e.mmq_iq_experts(
9545 &dev.ptr_row,
9546 2,
9547 n_expert,
9548 &exi,
9549 &exo,
9550 &exp_d,
9551 &pself,
9552 &a_scr,
9553 in_pad,
9554 n_embd,
9555 n_active,
9556 n_pairs,
9557 n_pairs,
9558 m.down_exps.qtype,
9559 m.down_exps.row_bytes,
9560 )?
9561 } else {
9562 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9563 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9564 e.moe_pairs_matvec_q8_dec(
9565 &dev.ptr_row,
9566 2,
9567 &exi,
9568 &exo,
9569 &exp_d,
9570 &pself,
9571 &aq2,
9572 &ad2,
9573 n_ff_exp,
9574 n_embd,
9575 n_expert,
9576 n_active,
9577 n_pairs,
9578 m.down_exps.qtype,
9579 m.down_exps.row_bytes,
9580 )?
9581 };
9582 let mut moe_out = e.uninit(t * n_embd)?;
9583 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9584 return Ok(moe_out);
9585 }
9586
9587 let g_len = m.gate_exps.expert_stride;
9588 let u_len = m.up_exps.expert_stride;
9589 let d_len = m.down_exps.expert_stride;
9590 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9594 let (mut sg, mut su, mut sd) = if dev.is_some() {
9595 (None, None, None)
9596 } else {
9597 (
9598 Some(e.alloc_u8_uninit(g_len)?),
9599 Some(e.alloc_u8_uninit(u_len)?),
9600 Some(e.alloc_u8_uninit(d_len)?),
9601 )
9602 };
9603 let mut moe_out = e.zeros(t * n_embd)?;
9604 for tok in 0..t {
9605 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9606 let w = &w_all[tok * n_used..(tok + 1) * n_used];
9607 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9608 for (j, &ex) in sel.iter().enumerate() {
9609 let ex = ex as usize;
9610 let gate = match dev {
9611 Some(d) => e.qmatvec_view(
9612 &d.gate,
9613 ex * g_len..(ex + 1) * g_len,
9614 &zt,
9615 1,
9616 m.gate_exps.in_f,
9617 m.gate_exps.out_f,
9618 m.gate_exps.qtype,
9619 m.gate_exps.row_bytes,
9620 )?,
9621 None => {
9622 let sg = sg.as_mut().unwrap();
9623 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9624 e.qmatvec_view(
9625 sg,
9626 0..g_len,
9627 &zt,
9628 1,
9629 m.gate_exps.in_f,
9630 m.gate_exps.out_f,
9631 m.gate_exps.qtype,
9632 m.gate_exps.row_bytes,
9633 )?
9634 }
9635 };
9636 let up = match dev {
9637 Some(d) => e.qmatvec_view(
9638 &d.up,
9639 ex * u_len..(ex + 1) * u_len,
9640 &zt,
9641 1,
9642 m.up_exps.in_f,
9643 m.up_exps.out_f,
9644 m.up_exps.qtype,
9645 m.up_exps.row_bytes,
9646 )?,
9647 None => {
9648 let su = su.as_mut().unwrap();
9649 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9650 e.qmatvec_view(
9651 su,
9652 0..u_len,
9653 &zt,
9654 1,
9655 m.up_exps.in_f,
9656 m.up_exps.out_f,
9657 m.up_exps.qtype,
9658 m.up_exps.row_bytes,
9659 )?
9660 }
9661 };
9662 let mut act = e.uninit(n_ff_exp)?;
9663 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9664 let actv = act.slice(0..n_ff_exp);
9665 let y = match dev {
9666 Some(d) => e.qmatvec_view(
9667 &d.down,
9668 ex * d_len..(ex + 1) * d_len,
9669 &actv,
9670 1,
9671 m.down_exps.in_f,
9672 m.down_exps.out_f,
9673 m.down_exps.qtype,
9674 m.down_exps.row_bytes,
9675 )?,
9676 None => {
9677 let sd = sd.as_mut().unwrap();
9678 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9679 e.qmatvec_view(
9680 sd,
9681 0..d_len,
9682 &actv,
9683 1,
9684 m.down_exps.in_f,
9685 m.down_exps.out_f,
9686 m.down_exps.qtype,
9687 m.down_exps.row_bytes,
9688 )?
9689 }
9690 };
9691 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9692 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9693 }
9694 }
9695 Ok(moe_out)
9696 }
9697
9698 fn gemma4_layer(
9700 &self,
9701 e: &Engine,
9702 il: usize,
9703 layer: &crate::hybrid::HybridLayer,
9704 x: &CudaSlice<f32>,
9705 pos_d: &CudaSlice<i32>,
9706 t: usize,
9707 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9708 let n_embd = self.cfg.n_embd as usize;
9709 let eps = self.cfg.rms_eps;
9710
9711 let mut h = e.zeros(t * n_embd)?;
9712 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9713 let Mixer::Full(fa) = &layer.mixer else {
9714 panic!("gemma4 layer {il} not full-attn")
9715 };
9716 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9717 let mut cur = e.zeros(t * n_embd)?;
9719 e.rms_norm(
9720 &o,
9721 layer.post_attn_norm.float_data(),
9722 &mut cur,
9723 n_embd,
9724 t,
9725 eps,
9726 )?;
9727 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9728 }
9729
9730 fn gemma4_layer_tail_add(
9734 &self,
9735 e: &Engine,
9736 layer: &crate::hybrid::HybridLayer,
9737 cur: &CudaSlice<f32>,
9738 x: &CudaSlice<f32>,
9739 t: usize,
9740 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9741 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9742 }
9743
9744 fn gemma4_layer_tail_add_n(
9747 &self,
9748 e: &Engine,
9749 layer: &crate::hybrid::HybridLayer,
9750 cur: &CudaSlice<f32>,
9751 x: &CudaSlice<f32>,
9752 t: usize,
9753 next_norm: Option<&CudaSlice<f32>>,
9754 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9755 let n_embd = self.cfg.n_embd as usize;
9756 let bits = layer.gemma4.as_ref().unwrap();
9757 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9758 let mut xn = e.uninit(t * n_embd)?;
9759 match next_norm {
9760 Some(w) => {
9761 let mut hn = e.uninit(t * n_embd)?;
9762 e.add_scale_rms_norm(
9763 &sn,
9764 &attn_out,
9765 bits.layer_scale,
9766 w,
9767 &mut xn,
9768 &mut hn,
9769 n_embd,
9770 t,
9771 self.cfg.rms_eps,
9772 )?;
9773 Ok((xn, Some(hn)))
9774 }
9775 None => {
9776 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9777 Ok((xn, None))
9778 }
9779 }
9780 }
9781
9782 fn gemma4_layer_tail_core(
9785 &self,
9786 e: &Engine,
9787 layer: &crate::hybrid::HybridLayer,
9788 cur: &CudaSlice<f32>,
9789 x: &CudaSlice<f32>,
9790 t: usize,
9791 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9792 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9793 }
9794
9795 fn gemma4_layer_tail_core_pn(
9802 &self,
9803 e: &Engine,
9804 layer: &crate::hybrid::HybridLayer,
9805 cur: &CudaSlice<f32>,
9806 x: &CudaSlice<f32>,
9807 t: usize,
9808 pre_norm: Option<&CudaSlice<f32>>,
9809 defer_post_norm: bool,
9810 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9811 let n_embd = self.cfg.n_embd as usize;
9812 let eps = self.cfg.rms_eps;
9813 let bits = layer.gemma4.as_ref().unwrap();
9814
9815 let Some(mbits) = bits.moe_bits.as_ref() else {
9818 let crate::hybrid::Ffn::Dense {
9819 ffn_gate,
9820 ffn_up,
9821 ffn_down,
9822 } = &layer.ffn
9823 else {
9824 panic!("gemma4 dense layer without Dense ffn")
9825 };
9826 let mut attn_out = e.uninit(t * n_embd)?;
9827 let mut zsh = e.uninit(t * n_embd)?;
9828 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9831 match pre_norm {
9832 Some(wa) if t == 1 => {
9833 zpair = Some(e.rms_pre_add_rms_norm_q8z(
9834 cur,
9835 wa,
9836 x,
9837 bits.ffn_norm.float_data(),
9838 &mut attn_out,
9839 &mut zsh,
9840 n_embd,
9841 t,
9842 eps,
9843 )?);
9844 }
9845 Some(wa) => e.rms_pre_add_rms_norm(
9846 cur,
9847 wa,
9848 x,
9849 bits.ffn_norm.float_data(),
9850 &mut attn_out,
9851 &mut zsh,
9852 n_embd,
9853 t,
9854 eps,
9855 )?,
9856 None => e.add_rms_norm(
9857 cur,
9858 x,
9859 bits.ffn_norm.float_data(),
9860 &mut attn_out,
9861 &mut zsh,
9862 n_embd,
9863 t,
9864 eps,
9865 )?,
9866 }
9867 let n_ff = ffn_gate.out_features();
9868 let (gate, up) = if t == 1 {
9874 let (zq, zd) = match zpair {
9875 Some(p) => p,
9876 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9877 };
9878 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9879 Some(p) => p,
9880 None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
9882 Some(p) => p,
9883 None => (
9884 e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9885 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9886 ),
9887 },
9888 }
9889 } else {
9890 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9895 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9896 let fused = if f2b {
9897 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9898 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9899 } else {
9900 None
9901 };
9902 match fused {
9903 Some(p) => p,
9904 None => {
9905 e.mmq_act_begin();
9907 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9908 }
9909 }
9910 };
9911 let mut act = e.uninit(t * n_ff)?;
9912 let f0 = if e.uses_q8_1_fast(ffn_down) {
9915 let upv = e.view(&up, t * n_ff);
9916 let up_all = upv.slice(0..t * n_ff);
9917 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9918 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9919 } else {
9920 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9921 e.matmul(ffn_down, &act, t)?
9922 };
9923 if defer_post_norm {
9924 return Ok((f0, attn_out));
9925 }
9926 let mut sn = e.uninit(t * n_embd)?;
9927 e.rms_norm(
9928 &f0,
9929 bits.post_ffw_norm.float_data(),
9930 &mut sn,
9931 n_embd,
9932 t,
9933 eps,
9934 )?;
9935 return Ok((sn, attn_out));
9936 };
9937
9938 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9939 let mut attn_out = e.uninit(t * n_embd)?;
9944 let mut router_in = e.uninit(t * n_embd)?;
9945 let fast_moe = match &layer.ffn {
9946 crate::hybrid::Ffn::Moe(m) => {
9947 m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9948 && expert_dp4a_supported(m.gate_exps.qtype)
9949 && expert_dp4a_supported(m.up_exps.qtype)
9950 && expert_dp4a_supported(m.down_exps.qtype)
9951 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9952 }
9953 _ => false,
9954 };
9955 let q8z = t < PRIME_MIN_T && fast_moe;
9956 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9957 let (z0, m2) = e.add_rms_norm3_q8z(
9958 cur,
9959 x,
9960 bits.ffn_norm.float_data(),
9961 &mbits.router_scale_pre,
9962 mbits.pre_ffw_norm_2.float_data(),
9963 &mut attn_out,
9964 &mut router_in,
9965 n_embd,
9966 t,
9967 eps,
9968 )?;
9969 (None, Some(z0), Some(m2))
9970 } else {
9971 let mut zsh = e.uninit(t * n_embd)?;
9972 let mut moe_in = e.uninit(t * n_embd)?;
9973 e.add_rms_norm3(
9974 cur,
9975 x,
9976 bits.ffn_norm.float_data(),
9977 &mbits.router_scale_pre,
9978 mbits.pre_ffw_norm_2.float_data(),
9979 &mut attn_out,
9980 &mut zsh,
9981 &mut router_in,
9982 &mut moe_in,
9983 n_embd,
9984 t,
9985 eps,
9986 )?;
9987 (Some((zsh, moe_in)), None, None)
9988 };
9989 let attn_out2 = attn_out;
9990 #[allow(unused_variables)]
9991 let attn_out = &attn_out2;
9992 let n_ff = mbits.shared_gate.out_features();
9993 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9994 if t == 1 {
9995 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9996 Some(p) => p,
9997 None => match e.matmul_nvfp4_fused2(
9998 &mbits.shared_gate,
9999 &mbits.shared_up,
10000 zq,
10001 zd,
10002 1,
10003 )? {
10004 Some(p) => p,
10005 None => {
10006 let h0 = e.zeros(0)?;
10007 (
10008 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
10009 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
10010 )
10011 }
10012 },
10013 }
10014 } else {
10015 let h0 = e.zeros(0)?;
10017 (
10018 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
10019 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
10020 )
10021 }
10022 } else {
10023 let (zsh, _) = zsh_f32.as_ref().unwrap();
10024 (
10025 e.matmul(&mbits.shared_gate, zsh, t)?,
10026 e.matmul(&mbits.shared_up, zsh, t)?,
10027 )
10028 };
10029 let mut act = e.uninit(t * n_ff)?;
10030 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
10031 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
10032 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
10033 panic!("gemma4 layer not MoE")
10034 };
10035 let moe0 = match (&moe_q8, &zsh_f32) {
10036 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
10037 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
10038 _ => unreachable!(),
10039 };
10040 let mut mlp = e.uninit(t * n_embd)?;
10042 let mut moe = e.uninit(t * n_embd)?;
10043 e.rms_norm2x(
10044 &mlp0,
10045 &moe0,
10046 mbits.post_ffw_norm_1.float_data(),
10047 mbits.post_ffw_norm_2.float_data(),
10048 &mut mlp,
10049 &mut moe,
10050 n_embd,
10051 t,
10052 eps,
10053 )?;
10054
10055 let mut sum = e.uninit(t * n_embd)?;
10058 let mut sn = e.uninit(t * n_embd)?;
10059 e.add_rms_norm(
10060 &mlp,
10061 &moe,
10062 bits.post_ffw_norm.float_data(),
10063 &mut sum,
10064 &mut sn,
10065 n_embd,
10066 t,
10067 eps,
10068 )?;
10069 Ok((sn, attn_out2))
10070 }
10071
10072 pub(crate) fn gemma4_layer_tail_add_nq_pn(
10082 &self,
10083 e: &Engine,
10084 layer: &crate::hybrid::HybridLayer,
10085 o: &CudaSlice<f32>,
10086 x: &CudaSlice<f32>,
10087 t: usize,
10088 next_norm: Option<&CudaSlice<f32>>,
10089 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
10090 {
10091 let n_embd = self.cfg.n_embd as usize;
10092 let eps = self.cfg.rms_eps;
10093 let bits = layer.gemma4.as_ref().unwrap();
10094 if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
10095 let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
10096 e,
10097 layer,
10098 o,
10099 x,
10100 t,
10101 Some(layer.post_attn_norm.float_data()),
10102 true,
10103 )?;
10104 let mut xn = e.uninit(t * n_embd)?;
10105 return match next_norm {
10106 Some(w) => {
10107 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
10108 &f0,
10109 bits.post_ffw_norm.float_data(),
10110 &attn_out,
10111 bits.layer_scale,
10112 w,
10113 &mut xn,
10114 n_embd,
10115 t,
10116 eps,
10117 )?;
10118 Ok((xn, Some(pair)))
10119 }
10120 None => {
10121 let mut sn = e.uninit(t * n_embd)?;
10122 e.rms_norm(
10123 &f0,
10124 bits.post_ffw_norm.float_data(),
10125 &mut sn,
10126 n_embd,
10127 t,
10128 eps,
10129 )?;
10130 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10131 Ok((xn, None))
10132 }
10133 };
10134 }
10135 let mut cur = e.uninit(t * n_embd)?;
10136 e.rms_norm(
10137 o,
10138 layer.post_attn_norm.float_data(),
10139 &mut cur,
10140 n_embd,
10141 t,
10142 eps,
10143 )?;
10144 self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
10145 }
10146
10147 pub(crate) fn gemma4_layer_tail_add_nq(
10148 &self,
10149 e: &Engine,
10150 layer: &crate::hybrid::HybridLayer,
10151 cur: &CudaSlice<f32>,
10152 x: &CudaSlice<f32>,
10153 t: usize,
10154 next_norm: Option<&CudaSlice<f32>>,
10155 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
10156 {
10157 let n_embd = self.cfg.n_embd as usize;
10158 let bits = layer.gemma4.as_ref().unwrap();
10159 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
10160 let mut xn = e.uninit(t * n_embd)?;
10161 match next_norm {
10162 Some(w) => {
10163 let pair = e.add_scale_rms_norm_q8_1(
10164 &sn,
10165 &attn_out,
10166 bits.layer_scale,
10167 w,
10168 &mut xn,
10169 n_embd,
10170 t,
10171 self.cfg.rms_eps,
10172 )?;
10173 Ok((xn, Some(pair)))
10174 }
10175 None => {
10176 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10177 Ok((xn, None))
10178 }
10179 }
10180 }
10181
10182 fn gemma4_forward(
10185 &self,
10186 e: &Engine,
10187 tokens: &[u32],
10188 last_only: bool,
10189 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
10190 if self.is_gemma4_e4b() {
10193 return self.gemma4_e4b_forward(e, tokens, last_only);
10194 }
10195 let n_embd = self.cfg.n_embd as usize;
10196 let t = tokens.len();
10197 let pos: Vec<i32> = (0..t as i32).collect();
10198 let pos_d = e.htod_i32(&pos)?;
10199
10200 let mut x = self.embed(e, tokens)?;
10201 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10202 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
10205 let stat =
10206 |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
10207 let h = e.dtoh(x)?;
10208 let bad = h.iter().filter(|v| !v.is_finite()).count();
10209 let mx = h
10210 .iter()
10211 .filter(|v| v.is_finite())
10212 .fold(0.0f32, |m, v| m.max(v.abs()));
10213 eprintln!(
10214 "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
10215 &h[..3]
10216 );
10217 Ok(())
10218 };
10219 if probe {
10220 stat(e, &x, "embed")?;
10221 }
10222 for (il, layer) in self.layers.iter().enumerate() {
10223 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
10224 if probe {
10225 stat(e, &x, &format!("L{il}"))?;
10226 }
10227 }
10228 let mut hn = e.zeros(t * n_embd)?;
10229 e.rms_norm(
10230 &x,
10231 self.output_norm.float_data(),
10232 &mut hn,
10233 n_embd,
10234 t,
10235 self.cfg.rms_eps,
10236 )?;
10237 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10238 let n_vocab = self.output.out_features();
10239 let logits = if last_only {
10240 let hv = e.view(&hn, t * n_embd);
10241 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
10242 let mut hlast = e.zeros(n_embd)?;
10243 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
10244 let mut ld = e.matmul(&self.output, &hlast, 1)?;
10245 e.softcap(&mut ld, cap, n_vocab)?;
10246 self.gemma4_suppress(e, &mut ld, 1)?;
10247 e.dtoh(&ld)?
10248 } else {
10249 let mut ld = e.matmul(&self.output, &hn, t)?;
10250 e.softcap(&mut ld, cap, t * n_vocab)?;
10251 self.gemma4_suppress(e, &mut ld, t)?;
10252 e.dtoh(&ld)?
10253 };
10254 Ok(logits)
10255 }
10256
10257 pub(crate) fn gemma4_prime(
10262 &self,
10263 e: &Engine,
10264 tokens: &[u32],
10265 cache: &mut Cache,
10266 overlay: Option<&crate::vision::EmbedOverlay>,
10267 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10268 if cache.pos != 0 {
10273 return Err(
10274 "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
10275 — prime the full prompt in one call or decode tokenwise"
10276 .into(),
10277 );
10278 }
10279 let n_embd = self.cfg.n_embd as usize;
10280 let eps = self.cfg.rms_eps;
10281 let t = tokens.len();
10282 let pos: Vec<i32> = (0..t as i32).collect();
10283 let pos_d = e.htod_i32(&pos)?;
10284 let mut x = self.embed(e, tokens)?;
10285 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10286 let island: Option<CudaSlice<i32>> = match overlay {
10293 Some(ov) => {
10294 let mut span_id = vec![-1i32; t];
10295 for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
10296 if pos + n_rows > t {
10297 return Err(format!(
10298 "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
10299 pos + n_rows
10300 )
10301 .into());
10302 }
10303 let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
10304 e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
10305 for s in span_id.iter_mut().skip(pos).take(n_rows) {
10306 *s = i as i32;
10307 }
10308 }
10309 if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
10313 None
10314 } else {
10315 Some(e.htod_i32(&span_id)?)
10316 }
10317 }
10318 None => None,
10319 };
10320 for (il, layer) in self.layers.iter().enumerate() {
10321 let mut h = e.zeros(t * n_embd)?;
10322 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
10323 let Mixer::Full(fa) = &layer.mixer else {
10324 panic!("gemma4 layer not full-attn")
10325 };
10326 let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
10327 if trace {
10328 let v = e.dtoh(&h)?;
10329 let nan = v.iter().filter(|x| x.is_nan()).count();
10330 eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
10331 }
10332 let o =
10333 self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
10334 if trace {
10335 let v = e.dtoh(&o)?;
10336 let nan = v.iter().filter(|x| x.is_nan()).count();
10337 eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
10338 }
10339 let mut cur = e.zeros(t * n_embd)?;
10340 e.rms_norm(
10341 &o,
10342 layer.post_attn_norm.float_data(),
10343 &mut cur,
10344 n_embd,
10345 t,
10346 eps,
10347 )?;
10348 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
10349 self.dflash_tap(e, cache, il, &x, t)?;
10350 if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
10352 let h = e.dtoh(&x)?;
10353 let nan = h.iter().filter(|v| v.is_nan()).count();
10354 let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
10355 eprintln!(
10356 "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
10357 h.len()
10358 );
10359 if nan > 0 {
10360 return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
10361 }
10362 }
10363 }
10364 cache.pos += t;
10365 let hiddens = e.clone_dtod(&x)?;
10366 let xv = e.view(&x, t * n_embd);
10367 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
10368 let mut h_seed = e.zeros(n_embd)?;
10369 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
10370 let mut hn = e.uninit(n_embd)?;
10371 e.rms_norm(
10372 &h_seed,
10373 self.output_norm.float_data(),
10374 &mut hn,
10375 n_embd,
10376 1,
10377 eps,
10378 )?;
10379 let mut ld = e.matmul(&self.output, &hn, 1)?;
10380 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10381 e.softcap(&mut ld, cap, self.output.out_features())?;
10382 self.gemma4_suppress(e, &mut ld, 1)?;
10383 let logits = e.dtoh(&ld)?;
10384 Ok((logits, h_seed, hiddens))
10385 }
10386
10387 fn gemma4_decode_attn(
10392 &self,
10393 e: &Engine,
10394 fa: &crate::hybrid::FullAttnLayer,
10395 il: usize,
10396 hq: &CudaSlice<i8>,
10397 hdq: &CudaSlice<f32>,
10398 pos_d: &CudaSlice<i32>,
10399 cache: &mut Cache,
10400 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10401 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10402 let eps = self.cfg.rms_eps;
10403 let aux = self.gemma4_aux.as_ref().unwrap();
10404 let ones = aux.ones(e);
10405 #[cfg(debug_assertions)]
10406 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
10407 let (hq, hdq) = (hq, hdq);
10408 let h0 = e.zeros(0)?;
10409 let h = &h0;
10410 let (q0, k0, v0) = if swa {
10411 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
10412 Some(t3) => t3,
10413 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10416 Some((q0, k0)) => {
10417 let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?;
10418 (q0, k0, v0)
10419 }
10420 None => (
10421 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10422 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10423 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10424 ),
10425 },
10426 }
10427 } else {
10428 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10429 Some(p) => p,
10430 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10431 Some(p) => p,
10432 None => (
10433 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10434 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10435 ),
10436 },
10437 };
10438 let v0 = e.clone_dtod(&k0)?;
10439 (q0, k0, v0)
10440 };
10441 let mut q = e.uninit(nh * hd)?;
10442 let mut k = e.uninit(nkv * hd)?;
10443 let mut v = e.uninit(nkv * hd)?;
10444 let ff = if swa {
10447 None
10448 } else {
10449 Some(
10450 aux.rope_freqs(e)
10451 .expect("gemma4 global rope needs rope_freqs.weight"),
10452 )
10453 };
10454 #[cfg(debug_assertions)]
10455 if let Some(ff) = ff {
10456 crate::debug_assert_tensor_stream_device(
10457 ff,
10458 &e.stream(),
10459 "gemma4_decode_attn.rope_freqs",
10460 );
10461 }
10462 let kvl = cache.kv[il].as_mut().unwrap();
10463 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10464 if crate::Engine::qkv_append_on() {
10465 e.rms_norm_qkv_rope_append(
10469 &q0,
10470 &k0,
10471 &v0,
10472 fa.q_norm.float_data(),
10473 fa.k_norm.float_data(),
10474 ones,
10475 &mut q,
10476 &mut k,
10477 &mut v,
10478 hd,
10479 self.gemma4_rope_dims(il),
10480 nh,
10481 nkv,
10482 pos_d,
10483 nh,
10484 nkv,
10485 base,
10486 1.0,
10487 ff,
10488 eps,
10489 &mut kvl.k,
10490 &mut kvl.v,
10491 kvl.len,
10492 kvl.k_tok_bytes,
10493 kvl.v_tok_bytes,
10494 kv_fp8,
10495 )?;
10496 } else {
10497 e.rms_norm_qkv_rope(
10498 &q0,
10499 &k0,
10500 &v0,
10501 fa.q_norm.float_data(),
10502 fa.k_norm.float_data(),
10503 ones,
10504 &mut q,
10505 &mut k,
10506 &mut v,
10507 hd,
10508 self.gemma4_rope_dims(il),
10509 nh,
10510 nkv,
10511 pos_d,
10512 nh,
10513 nkv,
10514 base,
10515 1.0,
10516 ff,
10517 eps,
10518 )?;
10519 e.append_kv_quantized(
10520 &k,
10521 &v,
10522 &mut kvl.k,
10523 &mut kvl.v,
10524 kvl.len,
10525 kvl.kv_dim_k,
10526 kvl.kv_dim_v,
10527 kvl.k_tok_bytes,
10528 kvl.v_tok_bytes,
10529 kv_fp8,
10530 )?;
10531 }
10532 kvl.len += 1;
10533 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10537 let mut attn = e.uninit(nh * hd)?;
10538 if !swa
10540 && hd == 512
10541 && kvl.len >= crate::fa512_min_tkv()
10542 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10543 {
10544 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10545 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10546 let base = kvl.len as i32;
10548 e.i32_set_k(&mut kvl.len_d, base)?;
10549 e.fa_decode_rows(
10550 &q,
10551 &kp,
10552 &vp,
10553 &mut attn,
10554 hd,
10555 nh,
10556 nkv,
10557 kvl.len - 1,
10558 1,
10559 scale,
10560 kvl.k_tok_bytes,
10561 kvl.v_tok_bytes,
10562 Some((&kvl.len_d, -1)),
10563 false,
10564 false,
10565 None,
10566 )?;
10567 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10568 }
10569 if swa
10571 && kvl.len > win
10572 && hd == 256
10573 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10574 {
10575 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10576 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10577 let base = kvl.len as i32;
10578 e.i32_set_k(&mut kvl.len_d, base)?;
10579 e.fa_decode_rows_w(
10580 &q,
10581 &kp,
10582 &vp,
10583 &mut attn,
10584 hd,
10585 nh,
10586 nkv,
10587 &kvl.len_d,
10588 -1,
10589 1,
10590 scale,
10591 win,
10592 kvl.k_tok_bytes,
10593 kvl.v_tok_bytes,
10594 None,
10595 )?;
10596 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10597 }
10598 let (off_tok, t_kv) = if swa && kvl.len > win {
10599 (kvl.len - win, win)
10600 } else {
10601 (0, kvl.len)
10602 };
10603 let k_view = e.view_u8_range(
10604 &kvl.k,
10605 off_tok * kvl.k_tok_bytes,
10606 (off_tok + t_kv) * kvl.k_tok_bytes,
10607 );
10608 let v_view = e.view_u8_range(
10609 &kvl.v,
10610 off_tok * kvl.v_tok_bytes,
10611 (off_tok + t_kv) * kvl.v_tok_bytes,
10612 );
10613 e.fa_decode_kvmod(
10614 &q,
10615 &k_view,
10616 &v_view,
10617 &mut attn,
10618 hd,
10619 nh,
10620 nkv,
10621 t_kv,
10622 scale,
10623 kvl.k_tok_bytes,
10624 kvl.v_tok_bytes,
10625 swa && crate::Engine::wkv_on(),
10626 )?;
10627 Ok(e.matmul(&fa.wo, &attn, 1)?)
10628 }
10629
10630 #[allow(clippy::too_many_arguments)]
10637 pub fn gemma4_decode_step_dc(
10638 &self,
10639 e: &Engine,
10640 token_d: &CudaSlice<u32>,
10641 pos_d: &mut CudaSlice<i32>,
10642 embd_gpu: &CudaSlice<u8>,
10643 embd_qt: i32,
10644 embd_rb: usize,
10645 cache: &mut Cache,
10646 n_vocab: usize,
10647 cap_bucket_max: Option<(usize, usize)>,
10648 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10649 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10650 self.gemma4_decode_step_dc_into(
10651 e,
10652 token_d,
10653 pos_d,
10654 embd_gpu,
10655 embd_qt,
10656 embd_rb,
10657 cache,
10658 n_vocab,
10659 cap_bucket_max,
10660 &mut tok_out,
10661 )?;
10662 Ok(tok_out)
10663 }
10664
10665 #[allow(clippy::too_many_arguments)]
10668 pub fn gemma4_decode_step_dc_into(
10669 &self,
10670 e: &Engine,
10671 token_d: &CudaSlice<u32>,
10672 pos_d: &mut CudaSlice<i32>,
10673 embd_gpu: &CudaSlice<u8>,
10674 embd_qt: i32,
10675 embd_rb: usize,
10676 cache: &mut Cache,
10677 n_vocab: usize,
10678 cap_bucket_max: Option<(usize, usize)>,
10679 tok_out: &mut CudaSlice<u32>,
10680 ) -> Result<(), Box<dyn std::error::Error>> {
10681 let n_embd = self.cfg.n_embd as usize;
10682 let eps = self.cfg.rms_eps;
10683 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10684 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10685 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10686 let n_layers = self.layers.len();
10687 for (il, layer) in self.layers.iter().enumerate() {
10688 let (hq, hdq) = match h_carry.take() {
10689 Some(p) => p,
10690 None => {
10691 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10692 }
10693 };
10694 let Mixer::Full(fa) = &layer.mixer else {
10695 panic!("gemma4 layer {il} not full-attn")
10696 };
10697 let o =
10698 self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10699 let next_norm = if il + 1 < n_layers {
10700 Some(self.layers[il + 1].attn_norm.float_data())
10701 } else {
10702 None
10703 };
10704 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
10705 x = xn;
10706 h_carry = hn;
10707 }
10708 let mut hn = e.uninit(n_embd)?;
10709 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10710 let mut logits = e.matmul(&self.output, &hn, 1)?;
10711 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10713 e.inc_seqlen(pos_d)?;
10714 if cap_bucket_max.is_none() {
10715 cache.pos += 1;
10716 }
10717 Ok(())
10718 }
10719
10720 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10727 let n_embd = self.cfg.n_embd as usize;
10728 let n_vocab = self.output.out_features();
10729 let n_layers = self.layers.len();
10730 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10731 for il in 0..n_layers {
10732 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10733 qmax = qmax.max(nh * hd);
10734 kvmax = kvmax.max(nkv * hd);
10735 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10736 ffmax = ffmax.max(ffn_gate.out_features());
10737 }
10738 }
10739 Ok(G4DcSlots {
10740 x: e.uninit(n_embd)?,
10741 xn: e.uninit(n_embd)?,
10742 cur: e.uninit(n_embd)?,
10743 hq: e.alloc_i8_uninit(n_embd)?,
10744 hd_: e.uninit(n_embd / 32)?,
10745 q0: e.uninit(qmax)?,
10746 k0: e.uninit(kvmax)?,
10747 v0: e.uninit(kvmax)?,
10748 q: e.uninit(qmax)?,
10749 k: e.uninit(kvmax)?,
10750 v: e.uninit(kvmax)?,
10751 attn: e.uninit(qmax)?,
10752 o: e.uninit(n_embd)?,
10753 attn_out: e.uninit(n_embd)?,
10754 zsh: e.uninit(n_embd)?,
10755 zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10758 zd: e.uninit(n_embd.max(qmax) / 32)?,
10759 gate: e.uninit(ffmax)?,
10760 up: e.uninit(ffmax)?,
10761 act: e.uninit(ffmax)?,
10762 actq: e.alloc_i8_uninit(ffmax)?,
10763 actd: e.uninit(ffmax / 32)?,
10764 f0: e.uninit(n_embd)?,
10765 sn: e.uninit(n_embd)?,
10766 hn: e.uninit(n_embd)?,
10767 logits: e.uninit(n_vocab)?,
10768 })
10769 }
10770
10771 fn g4_matvec_m1_into(
10774 &self,
10775 e: &Engine,
10776 w: &crate::model::GpuTensor,
10777 aq: &CudaSlice<i8>,
10778 ad: &CudaSlice<f32>,
10779 y: &mut CudaSlice<f32>,
10780 ) -> Result<(), Box<dyn std::error::Error>> {
10781 use crate::model::GpuTensor;
10782 let (bytes, qtype, row_bytes, scale, rp) = match w {
10783 GpuTensor::Quant {
10784 bytes,
10785 qtype,
10786 row_bytes,
10787 scale,
10788 rp,
10789 ..
10790 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10791 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10792 };
10793 let (mbytes, mrp) = match w {
10794 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10795 _ => (bytes, rp),
10796 };
10797 e.qmatvec_mmvq_into(
10798 mbytes,
10799 aq,
10800 ad,
10801 1,
10802 w.in_features(),
10803 w.out_features(),
10804 qtype,
10805 row_bytes,
10806 scale,
10807 mrp,
10808 y,
10809 )
10810 }
10811
10812 #[allow(clippy::too_many_arguments)]
10816 pub fn gemma4_decode_step_dc_slotted(
10817 &self,
10818 e: &Engine,
10819 token_d: &CudaSlice<u32>,
10820 pos_d: &mut CudaSlice<i32>,
10821 embd_gpu: &CudaSlice<u8>,
10822 embd_qt: i32,
10823 embd_rb: usize,
10824 cache: &mut Cache,
10825 n_vocab: usize,
10826 cap_bucket_max: Option<(usize, usize)>,
10827 sl: &mut G4DcSlots,
10828 tok_out: &mut CudaSlice<u32>,
10829 ring: Option<(&mut CudaSlice<u32>, usize)>,
10830 ) -> Result<(), Box<dyn std::error::Error>> {
10831 let n_embd = self.cfg.n_embd as usize;
10832 let eps = self.cfg.rms_eps;
10833 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10834 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10835 let n_layers = self.layers.len();
10836 let mut has_carry = false;
10837 for il in 0..n_layers {
10838 if !has_carry {
10839 e.rms_norm_q8_1_into(
10840 &sl.x,
10841 self.layers[il].attn_norm.float_data(),
10842 n_embd,
10843 1,
10844 eps,
10845 &mut sl.hq,
10846 &mut sl.hd_,
10847 )?;
10848 }
10849 has_carry = true;
10850 let layer = &self.layers[il];
10851 let Mixer::Full(fa) = &layer.mixer else {
10852 panic!("gemma4 layer {il} not full-attn")
10853 };
10854 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10855 if !Engine::g4_pnfold_on() {
10858 e.rms_norm(
10859 &sl.o,
10860 layer.post_attn_norm.float_data(),
10861 &mut sl.cur,
10862 n_embd,
10863 1,
10864 eps,
10865 )?;
10866 }
10867 let next_norm = if il + 1 < n_layers {
10868 Some(self.layers[il + 1].attn_norm.float_data())
10869 } else {
10870 None
10871 };
10872 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10873 std::mem::swap(&mut sl.x, &mut sl.xn);
10874 }
10875 e.rms_norm(
10876 &sl.x,
10877 self.output_norm.float_data(),
10878 &mut sl.hn,
10879 n_embd,
10880 1,
10881 eps,
10882 )?;
10883 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10884 {
10886 let (zq, zd) = (&sl.zq, &sl.zd);
10887 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10888 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10889 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10890 }
10891 self.gemma4_suppress(e, &mut sl.logits, 1)?;
10892 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10893 if let Some((ring, base)) = ring {
10894 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10898 }
10899 e.inc_seqlen(pos_d)?;
10900 if cap_bucket_max.is_none() {
10901 cache.pos += 1;
10902 }
10903 Ok(())
10904 }
10905
10906 #[allow(clippy::too_many_arguments)]
10908 fn gemma4_decode_attn_dc_slotted(
10909 &self,
10910 e: &Engine,
10911 fa: &crate::hybrid::FullAttnLayer,
10912 il: usize,
10913 pos_d: &CudaSlice<i32>,
10914 cache: &mut Cache,
10915 cap_bucket_max: Option<(usize, usize)>,
10916 sl: &mut G4DcSlots,
10917 ) -> Result<(), Box<dyn std::error::Error>> {
10918 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10919 let eps = self.cfg.rms_eps;
10920 let aux = self.gemma4_aux.as_ref().unwrap();
10921 let ones = aux.ones(e);
10922 #[cfg(debug_assertions)]
10923 crate::debug_assert_tensor_stream_device(
10924 ones,
10925 &e.stream(),
10926 "gemma4_decode_attn_dc_slotted.ones",
10927 );
10928 {
10929 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10930 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10931 if swa {
10932 if !e.matmul_q4_fused3_into(
10933 &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10934 )? {
10935 if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10939 {
10940 self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
10941 } else {
10942 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10943 }
10944 }
10945 } else {
10946 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10947 && !e
10948 .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10949 {
10950 return Err("slotted step: fused2 unavailable".into());
10951 }
10952 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10953 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10954 }
10955 }
10956 let ff = if swa {
10959 None
10960 } else {
10961 Some(
10962 aux.rope_freqs(e)
10963 .expect("gemma4 global rope needs rope_freqs.weight"),
10964 )
10965 };
10966 #[cfg(debug_assertions)]
10967 if let Some(ff) = ff {
10968 crate::debug_assert_tensor_stream_device(
10969 ff,
10970 &e.stream(),
10971 "gemma4_decode_attn_dc_slotted.rope_freqs",
10972 );
10973 }
10974 let kvl = cache.kv[il].as_mut().unwrap();
10975 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10976 if crate::Engine::qkv_append_on() {
10977 e.rms_norm_qkv_rope_append_dc(
10979 &sl.q0,
10980 &sl.k0,
10981 &sl.v0,
10982 fa.q_norm.float_data(),
10983 fa.k_norm.float_data(),
10984 ones,
10985 &mut sl.q,
10986 &mut sl.k,
10987 &mut sl.v,
10988 hd,
10989 self.gemma4_rope_dims(il),
10990 nh,
10991 nkv,
10992 pos_d,
10993 nh,
10994 nkv,
10995 base,
10996 1.0,
10997 ff,
10998 eps,
10999 &mut kvl.k,
11000 &mut kvl.v,
11001 &kvl.len_d,
11002 kvl.k_tok_bytes,
11003 kvl.v_tok_bytes,
11004 kv_fp8,
11005 )?;
11006 } else {
11007 e.rms_norm_qkv_rope(
11008 &sl.q0,
11009 &sl.k0,
11010 &sl.v0,
11011 fa.q_norm.float_data(),
11012 fa.k_norm.float_data(),
11013 ones,
11014 &mut sl.q,
11015 &mut sl.k,
11016 &mut sl.v,
11017 hd,
11018 self.gemma4_rope_dims(il),
11019 nh,
11020 nkv,
11021 pos_d,
11022 nh,
11023 nkv,
11024 base,
11025 1.0,
11026 ff,
11027 eps,
11028 )?;
11029 e.append_kv_quantized_dc(
11030 &sl.k,
11031 &sl.v,
11032 &mut kvl.k,
11033 &mut kvl.v,
11034 &kvl.len_d,
11035 kvl.kv_dim_k,
11036 kvl.kv_dim_v,
11037 kvl.k_tok_bytes,
11038 kvl.v_tok_bytes,
11039 kv_fp8,
11040 )?;
11041 }
11042 e.inc_seqlen(&mut kvl.len_d)?;
11043 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
11044 let k_view = e.view_u8(&kvl.k, kvl.k.len());
11045 let v_view = e.view_u8(&kvl.v, kvl.v.len());
11046 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11047 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11048 let mut fa_q8 = false;
11052 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11053 e.fa_decode_rows(
11054 &sl.q,
11055 &k_view,
11056 &v_view,
11057 &mut sl.attn,
11058 hd,
11059 nh,
11060 nkv,
11061 b_glob - 1,
11062 1,
11063 scale,
11064 kvl.k_tok_bytes,
11065 kvl.v_tok_bytes,
11066 Some((&kvl.len_d, -1)),
11067 false,
11068 false,
11069 Some((&mut sl.zq, &mut sl.zd)),
11070 )?;
11071 fa_q8 = true;
11072 } else if swa && b_swa > win && hd == 256 && rows_on {
11073 e.fa_decode_rows_w(
11074 &sl.q,
11075 &k_view,
11076 &v_view,
11077 &mut sl.attn,
11078 hd,
11079 nh,
11080 nkv,
11081 &kvl.len_d,
11082 -1,
11083 1,
11084 scale,
11085 win,
11086 kvl.k_tok_bytes,
11087 kvl.v_tok_bytes,
11088 Some((&mut sl.zq, &mut sl.zd)),
11089 )?;
11090 fa_q8 = true;
11091 } else {
11092 let b = if swa { b_swa } else { b_glob };
11093 e.fa_decode_dc(
11094 &sl.q,
11095 &k_view,
11096 &v_view,
11097 &mut sl.attn,
11098 hd,
11099 nh,
11100 nkv,
11101 &kvl.len_d,
11102 b,
11103 scale,
11104 kvl.k_tok_bytes,
11105 kvl.v_tok_bytes,
11106 swa && crate::Engine::wkv_on(),
11107 )?;
11108 }
11109 if !fa_q8 {
11110 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
11111 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
11112 }
11113 {
11114 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
11115 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11116 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
11117 }
11118 Ok(())
11119 }
11120
11121 fn gemma4_layer_tail_slotted(
11124 &self,
11125 e: &Engine,
11126 layer: &crate::hybrid::HybridLayer,
11127 next_norm: Option<&CudaSlice<f32>>,
11128 sl: &mut G4DcSlots,
11129 ) -> Result<(), Box<dyn std::error::Error>> {
11130 let n_embd = self.cfg.n_embd as usize;
11131 let eps = self.cfg.rms_eps;
11132 let bits = layer.gemma4.as_ref().unwrap();
11133 let crate::hybrid::Ffn::Dense {
11134 ffn_gate,
11135 ffn_up,
11136 ffn_down,
11137 } = &layer.ffn
11138 else {
11139 return Err("slotted tail: dense ffn only".into());
11140 };
11141 let pnfold = Engine::g4_pnfold_on();
11142 if pnfold {
11143 let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
11146 let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
11147 e.rms_pre_add_rms_norm_q8z_into(
11148 or,
11149 layer.post_attn_norm.float_data(),
11150 xr,
11151 bits.ffn_norm.float_data(),
11152 &mut sl.attn_out,
11153 &mut sl.zsh,
11154 n_embd,
11155 1,
11156 eps,
11157 &mut sl.zq,
11158 &mut sl.zd,
11159 )?;
11160 } else {
11161 e.add_rms_norm(
11162 &sl.cur,
11163 &sl.x,
11164 bits.ffn_norm.float_data(),
11165 &mut sl.attn_out,
11166 &mut sl.zsh,
11167 n_embd,
11168 1,
11169 eps,
11170 )?;
11171 }
11172 let n_ff = ffn_gate.out_features();
11173 if !pnfold {
11174 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
11175 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
11176 }
11177 {
11178 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
11179 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11180 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
11181 && !e.matmul_nvfp4_fused2_into(
11182 ffn_gate,
11183 ffn_up,
11184 zq,
11185 zd,
11186 &mut sl.gate,
11187 &mut sl.up,
11188 )?
11189 {
11190 return Err("slotted tail: ffn fused2 unavailable".into());
11191 }
11192 }
11193 debug_assert!(e.uses_q8_1_fast(ffn_down));
11194 {
11195 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
11196 let upv = e.view(upr, n_ff);
11197 let up_all = upv.slice(0..n_ff);
11198 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
11199 e.gelu_tanh_mul_q8_1_into(
11200 gr,
11201 &up_all,
11202 &mut sl.act,
11203 n_ff,
11204 1,
11205 &mut sl.actq,
11206 &mut sl.actd,
11207 )?;
11208 }
11209 {
11210 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
11211 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
11212 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
11213 }
11214 if pnfold {
11215 if let Some(w) = next_norm {
11218 let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
11219 let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
11220 e.rms_pre_add_scale_rms_norm_q8_1_into(
11221 f0r,
11222 bits.post_ffw_norm.float_data(),
11223 aor,
11224 bits.layer_scale,
11225 w,
11226 &mut sl.xn,
11227 n_embd,
11228 1,
11229 eps,
11230 &mut sl.hq,
11231 &mut sl.hd_,
11232 )?;
11233 return Ok(());
11234 }
11235 }
11236 e.rms_norm(
11237 &sl.f0,
11238 bits.post_ffw_norm.float_data(),
11239 &mut sl.sn,
11240 n_embd,
11241 1,
11242 eps,
11243 )?;
11244 match next_norm {
11245 Some(w) => {
11246 e.add_scale_rms_norm_q8_1_into(
11247 &sl.sn,
11248 &sl.attn_out,
11249 bits.layer_scale,
11250 w,
11251 &mut sl.xn,
11252 n_embd,
11253 1,
11254 eps,
11255 &mut sl.hq,
11256 &mut sl.hd_,
11257 )?;
11258 }
11259 None => {
11260 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
11261 }
11262 }
11263 Ok(())
11264 }
11265
11266 #[allow(clippy::too_many_arguments)]
11268 fn gemma4_decode_attn_dc(
11269 &self,
11270 e: &Engine,
11271 fa: &crate::hybrid::FullAttnLayer,
11272 il: usize,
11273 hq: &CudaSlice<i8>,
11274 hdq: &CudaSlice<f32>,
11275 pos_d: &CudaSlice<i32>,
11276 cache: &mut Cache,
11277 cap_bucket_max: Option<(usize, usize)>,
11278 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11279 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11280 let eps = self.cfg.rms_eps;
11281 let aux = self.gemma4_aux.as_ref().unwrap();
11282 let ones = aux.ones(e);
11283 #[cfg(debug_assertions)]
11284 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
11285 let (q0, k0, v0) = if swa {
11286 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
11287 Some(t3) => t3,
11288 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11290 Some((q0, k0)) => {
11291 let h0 = e.zeros(0)?;
11292 let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
11293 (q0, k0, v0)
11294 }
11295 None => {
11296 let h0 = e.zeros(0)?;
11297 (
11298 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11299 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11300 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
11301 )
11302 }
11303 },
11304 }
11305 } else {
11306 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
11307 Some(p) => p,
11308 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11309 Some(p) => p,
11310 None => {
11311 let h0 = e.zeros(0)?;
11312 (
11313 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11314 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11315 )
11316 }
11317 },
11318 };
11319 let v0 = e.clone_dtod(&k0)?;
11320 (q0, k0, v0)
11321 };
11322 let mut q = e.uninit(nh * hd)?;
11323 let mut k = e.uninit(nkv * hd)?;
11324 let mut v = e.uninit(nkv * hd)?;
11325 let ff = if swa {
11327 None
11328 } else {
11329 Some(
11330 aux.rope_freqs(e)
11331 .expect("gemma4 global rope needs rope_freqs.weight"),
11332 )
11333 };
11334 #[cfg(debug_assertions)]
11335 if let Some(ff) = ff {
11336 crate::debug_assert_tensor_stream_device(
11337 ff,
11338 &e.stream(),
11339 "gemma4_decode_attn_dc.rope_freqs",
11340 );
11341 }
11342 let kvl = cache.kv[il].as_mut().unwrap();
11343 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
11344 if crate::Engine::qkv_append_on() {
11345 e.rms_norm_qkv_rope_append_dc(
11347 &q0,
11348 &k0,
11349 &v0,
11350 fa.q_norm.float_data(),
11351 fa.k_norm.float_data(),
11352 ones,
11353 &mut q,
11354 &mut k,
11355 &mut v,
11356 hd,
11357 self.gemma4_rope_dims(il),
11358 nh,
11359 nkv,
11360 pos_d,
11361 nh,
11362 nkv,
11363 base,
11364 1.0,
11365 ff,
11366 eps,
11367 &mut kvl.k,
11368 &mut kvl.v,
11369 &kvl.len_d,
11370 kvl.k_tok_bytes,
11371 kvl.v_tok_bytes,
11372 kv_fp8,
11373 )?;
11374 } else {
11375 e.rms_norm_qkv_rope(
11376 &q0,
11377 &k0,
11378 &v0,
11379 fa.q_norm.float_data(),
11380 fa.k_norm.float_data(),
11381 ones,
11382 &mut q,
11383 &mut k,
11384 &mut v,
11385 hd,
11386 self.gemma4_rope_dims(il),
11387 nh,
11388 nkv,
11389 pos_d,
11390 nh,
11391 nkv,
11392 base,
11393 1.0,
11394 ff,
11395 eps,
11396 )?;
11397 e.append_kv_quantized_dc(
11398 &k,
11399 &v,
11400 &mut kvl.k,
11401 &mut kvl.v,
11402 &kvl.len_d,
11403 kvl.kv_dim_k,
11404 kvl.kv_dim_v,
11405 kvl.k_tok_bytes,
11406 kvl.v_tok_bytes,
11407 kv_fp8,
11408 )?;
11409 }
11410 e.inc_seqlen(&mut kvl.len_d)?;
11411 let mut attn = e.uninit(nh * hd)?;
11412 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11415 match cap_bucket_max {
11420 None => {
11421 kvl.len += 1;
11425 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11426 if !swa
11427 && hd == 512
11428 && kvl.len >= crate::fa512_min_tkv()
11429 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11430 {
11431 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11434 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11435 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11436 e.fa_decode_rows(
11437 &q,
11438 &kp,
11439 &vp,
11440 &mut attn,
11441 hd,
11442 nh,
11443 nkv,
11444 kvl.len - 1,
11445 1,
11446 scale,
11447 kvl.k_tok_bytes,
11448 kvl.v_tok_bytes,
11449 Some((&kvl.len_d, -1)),
11450 false,
11451 false,
11452 Some((&mut aq8, &mut ad8)),
11453 )?;
11454 fa_q8 = Some((aq8, ad8));
11455 } else if swa
11456 && kvl.len > win
11457 && hd == 256
11458 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11459 {
11460 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11462 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11463 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11464 e.fa_decode_rows_w(
11465 &q,
11466 &kp,
11467 &vp,
11468 &mut attn,
11469 hd,
11470 nh,
11471 nkv,
11472 &kvl.len_d,
11473 -1,
11474 1,
11475 scale,
11476 win,
11477 kvl.k_tok_bytes,
11478 kvl.v_tok_bytes,
11479 Some((&mut aq8, &mut ad8)),
11480 )?;
11481 fa_q8 = Some((aq8, ad8));
11482 } else {
11483 let (off_tok, t_kv) = if swa && kvl.len > win {
11484 (kvl.len - win, win)
11485 } else {
11486 (0, kvl.len)
11487 };
11488 let k_view = e.view_u8_range(
11489 &kvl.k,
11490 off_tok * kvl.k_tok_bytes,
11491 (off_tok + t_kv) * kvl.k_tok_bytes,
11492 );
11493 let v_view = e.view_u8_range(
11494 &kvl.v,
11495 off_tok * kvl.v_tok_bytes,
11496 (off_tok + t_kv) * kvl.v_tok_bytes,
11497 );
11498 e.fa_decode_kvmod(
11499 &q,
11500 &k_view,
11501 &v_view,
11502 &mut attn,
11503 hd,
11504 nh,
11505 nkv,
11506 t_kv,
11507 scale,
11508 kvl.k_tok_bytes,
11509 kvl.v_tok_bytes,
11510 swa && crate::Engine::wkv_on(),
11511 )?;
11512 }
11513 }
11514 Some((b_swa, b_glob)) => {
11515 let k_view = e.view_u8(&kvl.k, kvl.k.len());
11521 let v_view = e.view_u8(&kvl.v, kvl.v.len());
11522 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11523 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11524 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11525 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11526 e.fa_decode_rows(
11527 &q,
11528 &k_view,
11529 &v_view,
11530 &mut attn,
11531 hd,
11532 nh,
11533 nkv,
11534 b_glob - 1,
11535 1,
11536 scale,
11537 kvl.k_tok_bytes,
11538 kvl.v_tok_bytes,
11539 Some((&kvl.len_d, -1)),
11540 false,
11541 false,
11542 Some((&mut aq8, &mut ad8)),
11543 )?;
11544 fa_q8 = Some((aq8, ad8));
11545 } else if swa && b_swa > win && hd == 256 && rows_on {
11546 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11547 e.fa_decode_rows_w(
11548 &q,
11549 &k_view,
11550 &v_view,
11551 &mut attn,
11552 hd,
11553 nh,
11554 nkv,
11555 &kvl.len_d,
11556 -1,
11557 1,
11558 scale,
11559 win,
11560 kvl.k_tok_bytes,
11561 kvl.v_tok_bytes,
11562 Some((&mut aq8, &mut ad8)),
11563 )?;
11564 fa_q8 = Some((aq8, ad8));
11565 } else {
11566 let b = if swa { b_swa } else { b_glob };
11567 e.fa_decode_dc(
11568 &q,
11569 &k_view,
11570 &v_view,
11571 &mut attn,
11572 hd,
11573 nh,
11574 nkv,
11575 &kvl.len_d,
11576 b,
11577 scale,
11578 kvl.k_tok_bytes,
11579 kvl.v_tok_bytes,
11580 swa && crate::Engine::wkv_on(),
11581 )?;
11582 }
11583 }
11584 }
11585 if let Some((aq8, ad8)) = fa_q8 {
11588 let mut y = e.uninit(fa.wo.out_features())?;
11589 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11590 return Ok(y);
11591 }
11592 Ok(e.matmul(&fa.wo, &attn, 1)?)
11593 }
11594
11595 pub fn gemma4_generate_graph(
11600 &self,
11601 e: &Engine,
11602 prompt_pos: usize,
11603 first_token: u32,
11604 cache: &mut Cache,
11605 max_new: usize,
11606 eos: &[u32],
11607 mut on_token: impl FnMut(u32) -> bool,
11608 ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11609 if self.is_gemma4_e4b() {
11610 return Err(
11611 "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11612 .into(),
11613 );
11614 }
11615 use crate::decode::StopReason;
11616 let n_vocab = self.output.out_features();
11617 let n_embd = self.cfg.n_embd as usize;
11618 let embd_gpu = self
11619 .embd_gpu
11620 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11621 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11622 for kvl in cache.kv.iter_mut().flatten() {
11623 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11624 }
11625 let mut token_d = e.stream().clone_htod(&[first_token])?;
11626 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11627 let g4 = self.cfg.gemma4.as_ref().unwrap();
11628 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11629 let nkv_s = g4
11631 .head_count_kv
11632 .iter()
11633 .zip(g4.swa_pattern.iter())
11634 .find(|p| *p.1)
11635 .map(|p| *p.0 as usize)
11636 .unwrap_or(8);
11637 let nkv_g = g4
11638 .head_count_kv
11639 .iter()
11640 .zip(g4.swa_pattern.iter())
11641 .find(|p| !*p.1)
11642 .map(|p| *p.0 as usize)
11643 .unwrap_or(2);
11644 let mut graphs: std::collections::HashMap<
11645 ((bool, usize), (bool, usize), bool, bool),
11646 (
11647 cudarc::driver::CudaGraph,
11648 Vec<Box<dyn std::any::Any + Send>>,
11649 ),
11650 > = Default::default();
11651 let mut slots = self.g4_dc_slots(e)?;
11654 const RING: usize = 64;
11657 const DRAIN: usize = 1;
11663 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11664 let ring_base = prompt_pos;
11665 let mut out = Vec::with_capacity(max_new);
11666 let mut reason = StopReason::MaxNew;
11667 let mut next = first_token;
11668 let mut captures = 0usize;
11669 for _ in 0..max_new {
11670 out.push(next);
11671 if eos.contains(&next) {
11672 reason = StopReason::Eos;
11673 break;
11674 }
11675 if !on_token(next) {
11676 reason = StopReason::Callback;
11677 break;
11678 }
11679 let t_kv = cache.pos + 1;
11680 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11688 let f512 = crate::fa512_min_tkv();
11689 let key_s = if t_kv > win {
11690 (true, usize::MAX)
11691 } else {
11692 e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11693 };
11694 let (key_g, rung_end) = if t_kv >= f512 {
11695 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11698 ((true, end), end)
11699 } else {
11700 (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11701 };
11702 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11703 if !graphs.contains_key(&key) {
11704 let bucket_max = (t_kv, rung_end);
11705 let snap = cache.snapshot(e)?;
11707 let pos_save = e.dtoh_i32_one(&pos_d)?;
11708 let len_save: Vec<Option<i32>> = cache
11709 .kv
11710 .iter()
11711 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11712 .collect();
11713 let tok_save = e.dtoh_u32_one(&token_d)?;
11714 let graph = {
11719 let tok_ref = &mut token_d;
11720 let pos_ref = &mut pos_d;
11721 let cache_ref = &mut *cache;
11722 let slots_ref = &mut slots;
11723 let ring_ref = &mut ring;
11724 e.capture_graph_retained_flags(
11725 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11726 |e| {
11727 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11729 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11730 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11731 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11732 cache_ref, n_vocab, Some(bucket_max),
11733 sl, tok_ref, Some((rg, ring_base)))
11734 })?
11735 };
11736 cache.rollback(e, &snap, 0)?;
11737 e.set_i32_one(&mut pos_d, pos_save)?;
11738 for (il, ls) in len_save.iter().enumerate() {
11739 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11740 e.set_i32_one(&mut kvl.len_d, *v)?;
11741 }
11742 }
11743 e.set_u32_one(&mut token_d, tok_save)?;
11744 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11745 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11746 eprintln!("[graph-census] {c:?}");
11747 }
11748 }
11749 graphs.insert(key, graph);
11750 captures += 1;
11751 }
11752 let mut chunk = 1usize;
11757 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11758 .ok()
11759 .and_then(|v| v.parse().ok())
11760 .unwrap_or(DRAIN);
11761 while chunk < drain_cap && out.len() + chunk < max_new {
11762 let t_next = cache.pos + 1 + chunk;
11763 let key_s2 = if t_next > win {
11764 (true, usize::MAX)
11765 } else {
11766 e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11767 };
11768 let key_g2 = if t_next >= f512 {
11769 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11770 } else {
11771 e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11772 };
11773 if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11774 break;
11775 }
11776 chunk += 1;
11777 }
11778 let g = &graphs.get(&key).unwrap().0;
11779 for _ in 0..chunk {
11780 g.launch()?;
11781 }
11782 e.stream().synchronize()?;
11783 let ringh = e.dtoh_u32(&ring)?;
11784 for j in 0..chunk {
11785 let pos_j = cache.pos + j;
11786 let tok_j = ringh[(pos_j - ring_base) % RING];
11787 cache.pos += 0; if j + 1 == chunk {
11789 next = tok_j;
11790 } else {
11791 out.push(tok_j);
11792 if eos.contains(&tok_j) || !on_token(tok_j) {
11793 reason = if eos.contains(&tok_j) {
11794 StopReason::Eos
11795 } else {
11796 StopReason::Callback
11797 };
11798 let keep = cache.pos + j + 1;
11800 e.set_i32_one(&mut pos_d, keep as i32)?;
11801 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11802 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11803 kvl.len = keep;
11804 }
11805 cache.pos = keep;
11806 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11807 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11808 }
11809 return Ok((out, reason));
11810 }
11811 }
11812 }
11813 cache.pos += chunk;
11814 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11815 kvl.len += chunk;
11816 }
11817 }
11818 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11819 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11820 }
11821 Ok((out, reason))
11822 }
11823
11824 pub(crate) fn gemma4_decode_step_t(
11830 &self,
11831 e: &Engine,
11832 tokens: &[u32],
11833 pos0: usize,
11834 cache: &mut Cache,
11835 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11836 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11837 }
11838
11839 pub(crate) fn gemma4_decode_step_t_am(
11843 &self,
11844 e: &Engine,
11845 tokens: &[u32],
11846 pos0: usize,
11847 cache: &mut Cache,
11848 ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11849 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11850 let t = tokens.len();
11851 let n_vocab = self.output.out_features();
11852 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11853 for i in 0..t {
11854 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11855 }
11856 Ok((e.dtoh_u32(&toks)?, hn))
11857 }
11858
11859 pub(crate) fn gemma4_decode_step_t_am_dev(
11862 &self,
11863 e: &Engine,
11864 tok_d: &CudaSlice<u32>,
11865 t: usize,
11866 pos0: usize,
11867 cache: &mut Cache,
11868 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11869 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11870 let n_vocab = self.output.out_features();
11871 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11872 for i in 0..t {
11873 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11874 }
11875 Ok((vam, hn))
11876 }
11877
11878 pub(crate) fn gemma4_decode_step_t_h(
11881 &self,
11882 e: &Engine,
11883 tokens: &[u32],
11884 pos0: usize,
11885 cache: &mut Cache,
11886 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11887 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11888 let t = tokens.len();
11889 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11890 e.softcap(&mut ld, cap, t * self.output.out_features())?;
11891 Ok((e.dtoh(&ld)?, hn))
11892 }
11893
11894 pub(crate) fn verify_stream_scratch(
11897 &self,
11898 e: &Engine,
11899 cap: usize,
11900 ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11901 Ok(VerifyStreamScratch {
11902 pos_d: e.htod_i32(&vec![0i32; cap])?,
11903 row_ctrs: (0..cap)
11904 .map(|_| e.htod_i32(&[0]))
11905 .collect::<Result<_, _>>()?,
11906 })
11907 }
11908
11909 pub(crate) fn gemma4_verify_t_am_stream(
11917 &self,
11918 e: &Engine,
11919 tok_d: &CudaSlice<u32>,
11920 t: usize,
11921 ctr: &CudaSlice<i32>,
11922 hint: usize,
11923 cache: &mut Cache,
11924 scr: &mut VerifyStreamScratch,
11925 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11926 let n_embd = self.cfg.n_embd as usize;
11927 let eps = self.cfg.rms_eps;
11928 assert!(t <= scr.row_ctrs.len() && t <= 64);
11929 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11930 for i in 0..t {
11931 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11932 }
11933 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11934 let embd_gpu = self
11935 .embd_gpu
11936 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11937 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11938 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11939 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11940 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11941 let n_layers = self.layers.len();
11942 for (il, layer) in self.layers.iter().enumerate() {
11943 let (hq, hdq) = match h_carry.take() {
11944 Some(p) => p,
11945 None => {
11946 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11947 }
11948 };
11949 let Mixer::Full(fa) = &layer.mixer else {
11950 panic!("gemma4 layer {il} not full-attn")
11951 };
11952 let o = self
11953 .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11954 let next_norm = if il + 1 < n_layers {
11955 Some(self.layers[il + 1].attn_norm.float_data())
11956 } else {
11957 None
11958 };
11959 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11960 x = xn;
11961 h_carry = hn;
11962 self.dflash_tap(e, cache, il, &x, t)?;
11963 }
11964 let mut hn = e.uninit(t * n_embd)?;
11965 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11966 let ld = e.matmul(&self.output, &hn, t)?;
11967 let n_vocab = self.output.out_features();
11968 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11969 for i in 0..t {
11970 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11971 }
11972 Ok((vam, hn))
11973 }
11974
11975 pub(crate) fn dflash_tap(
11982 &self,
11983 e: &Engine,
11984 cache: &mut Cache,
11985 il: usize,
11986 x: &CudaSlice<f32>,
11987 t: usize,
11988 ) -> Result<(), Box<dyn std::error::Error>> {
11989 let Some(taps) = cache.dflash_taps.as_mut() else {
11990 return Ok(());
11991 };
11992 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11993 return Ok(());
11994 };
11995 let h = taps.hidden;
11996 let n_taps = taps.layer_ids.len();
11997 let base = taps.base;
11998 debug_assert!(
11999 base + t <= taps.t,
12000 "tap window {base}+{t} exceeds sink {}",
12001 taps.t
12002 );
12003 let xv = e.view(x, t * h);
12004 for r in 0..t {
12005 let row = xv.slice(r * h..(r + 1) * h);
12006 e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
12007 }
12008 Ok(())
12009 }
12010
12011 fn gemma4_verify_trunk(
12012 &self,
12013 e: &Engine,
12014 tokens: &[u32],
12015 pos0: usize,
12016 cache: &mut Cache,
12017 tok_dev: Option<&CudaSlice<u32>>,
12018 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12019 let n_embd = self.cfg.n_embd as usize;
12020 let eps = self.cfg.rms_eps;
12021 let t = tokens.len();
12022 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
12023 let pos_d = e.htod_i32(&pos)?;
12024 let mut x = match tok_dev {
12025 Some(td) => {
12026 let embd_gpu = self
12027 .embd_gpu
12028 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
12029 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
12030 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
12031 }
12032 None => e.htod(&self.embd.gather(n_embd, tokens))?,
12033 };
12034 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
12035 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12036 let n_layers = self.layers.len();
12037 for (il, layer) in self.layers.iter().enumerate() {
12038 let (hq, hdq) = match h_carry.take() {
12039 Some(p) => p,
12040 None => {
12041 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
12042 }
12043 };
12044 let Mixer::Full(fa) = &layer.mixer else {
12045 panic!("gemma4 layer {il} not full-attn")
12046 };
12047 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
12048 let next_norm = if il + 1 < n_layers {
12049 Some(self.layers[il + 1].attn_norm.float_data())
12050 } else {
12051 None
12052 };
12053 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
12054 x = xn;
12055 h_carry = hn;
12056 self.dflash_tap(e, cache, il, &x, t)?;
12057 }
12058 let mut hn = e.uninit(t * n_embd)?;
12059 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
12060 let mut ld = e.matmul(&self.output, &hn, t)?;
12061 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
12063 Ok((ld, hn))
12064 }
12065
12066 #[allow(clippy::too_many_arguments)]
12074 fn gemma4_verify_attn_stream(
12075 &self,
12076 e: &Engine,
12077 fa: &crate::hybrid::FullAttnLayer,
12078 il: usize,
12079 hq: &CudaSlice<i8>,
12080 hdq: &CudaSlice<f32>,
12081 pos_d: &CudaSlice<i32>,
12082 t: usize,
12083 cache: &mut Cache,
12084 hint: usize,
12085 row_ctrs: &[CudaSlice<i32>],
12086 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12087 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12088 let eps = self.cfg.rms_eps;
12089 let aux = self.gemma4_aux.as_ref().unwrap();
12090 let ones = aux.ones(e);
12091 #[cfg(debug_assertions)]
12092 crate::debug_assert_tensor_stream_device(
12093 ones,
12094 &e.stream(),
12095 "gemma4_verify_attn_stream.ones",
12096 );
12097 let h0 = e.zeros(0)?;
12098 let h = &h0;
12099 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12102 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12103 let fused_qkv = if f2b {
12104 if swa {
12105 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12106 .map(|(a, b, c)| (a, b, Some(c)))
12107 } else {
12108 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12109 .map(|(a, b)| (a, b, None))
12110 }
12111 } else {
12112 None
12113 };
12114 let (q0, k0, v0) = match fused_qkv {
12115 Some((a, b, cv)) => {
12116 let v = match cv {
12117 Some(c) => c,
12118 None => e.clone_dtod(&b)?,
12119 };
12120 (a, b, v)
12121 }
12122 None => {
12123 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12124 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12125 let v0 = if swa {
12126 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12127 } else {
12128 e.clone_dtod(&k0)?
12129 };
12130 (q0, k0, v0)
12131 }
12132 };
12133 let mut q = e.uninit(t * nh * hd)?;
12134 let mut k = e.uninit(t * nkv * hd)?;
12135 let mut v = e.uninit(t * nkv * hd)?;
12136 let ff = if swa {
12139 None
12140 } else {
12141 Some(
12142 aux.rope_freqs(e)
12143 .expect("gemma4 global rope needs rope_freqs.weight"),
12144 )
12145 };
12146 #[cfg(debug_assertions)]
12147 if let Some(ff) = ff {
12148 crate::debug_assert_tensor_stream_device(
12149 ff,
12150 &e.stream(),
12151 "gemma4_verify_attn_stream.rope_freqs",
12152 );
12153 }
12154 e.rms_norm_qkv_rope(
12155 &q0,
12156 &k0,
12157 &v0,
12158 fa.q_norm.float_data(),
12159 fa.k_norm.float_data(),
12160 ones,
12161 &mut q,
12162 &mut k,
12163 &mut v,
12164 hd,
12165 self.gemma4_rope_dims(il),
12166 nh * t,
12167 nkv * t,
12168 pos_d,
12169 nh,
12170 nkv,
12171 base,
12172 1.0,
12173 ff,
12174 eps,
12175 )?;
12176 let kvl = cache.kv[il].as_mut().unwrap();
12177 e.append_kv_quantized_rows_dc(
12179 &k,
12180 &v,
12181 &mut kvl.k,
12182 &mut kvl.v,
12183 &kvl.len_d,
12184 t,
12185 kvl.kv_dim_k,
12186 kvl.kv_dim_v,
12187 kvl.k_tok_bytes,
12188 kvl.v_tok_bytes,
12189 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12190 )?;
12191 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12194 let mut attn = e.uninit(t * nh * hd)?;
12195 let k_view = e.view_u8(&kvl.k, kvl.k.len());
12196 let v_view = e.view_u8(&kvl.v, kvl.v.len());
12197 if swa && hint + 1 >= win {
12200 e.fa_decode_rows_w(
12203 &q,
12204 &k_view,
12205 &v_view,
12206 &mut attn,
12207 hd,
12208 nh,
12209 nkv,
12210 &kvl.len_d,
12211 0,
12212 t,
12213 scale,
12214 win,
12215 kvl.k_tok_bytes,
12216 kvl.v_tok_bytes,
12217 None,
12218 )?;
12219 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
12220 let bucket = (hint + t + 2)
12233 .next_power_of_two()
12234 .min(crate::fa512_min_tkv().saturating_sub(1));
12235 let qv = e.view(&q, t * nh * hd);
12236 for i in 0..t {
12237 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
12238 let mut q_one = e.uninit(nh * hd)?;
12239 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12240 let mut a_one = e.uninit(nh * hd)?;
12241 e.fa_decode_dc(
12242 &q_one,
12243 &k_view,
12244 &v_view,
12245 &mut a_one,
12246 hd,
12247 nh,
12248 nkv,
12249 &row_ctrs[i],
12250 bucket,
12251 scale,
12252 kvl.k_tok_bytes,
12253 kvl.v_tok_bytes,
12254 false,
12255 )?;
12256 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12257 }
12258 } else if hd == 512 {
12259 e.fa_decode_rows(
12262 &q,
12263 &k_view,
12264 &v_view,
12265 &mut attn,
12266 hd,
12267 nh,
12268 nkv,
12269 hint,
12270 t,
12271 scale,
12272 kvl.k_tok_bytes,
12273 kvl.v_tok_bytes,
12274 Some((&kvl.len_d, 0)),
12275 false,
12276 false,
12277 None,
12278 )?;
12279 } else {
12280 e.fa_decode_rows_dc(
12282 &q,
12283 &k_view,
12284 &v_view,
12285 &mut attn,
12286 hd,
12287 nh,
12288 nkv,
12289 &kvl.len_d,
12290 hint + t,
12291 t,
12292 scale,
12293 kvl.k_tok_bytes,
12294 kvl.v_tok_bytes,
12295 0,
12296 swa && crate::Engine::wkv_on(),
12297 )?;
12298 }
12299 Ok(e.matmul(&fa.wo, &attn, t)?)
12300 }
12301
12302 fn gemma4_verify_attn(
12303 &self,
12304 e: &Engine,
12305 fa: &crate::hybrid::FullAttnLayer,
12306 il: usize,
12307 hq: &CudaSlice<i8>,
12308 hdq: &CudaSlice<f32>,
12309 pos_d: &CudaSlice<i32>,
12310 t: usize,
12311 cache: &mut Cache,
12312 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12313 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12314 let eps = self.cfg.rms_eps;
12315 let aux = self.gemma4_aux.as_ref().unwrap();
12316 let ones = aux.ones(e);
12317 #[cfg(debug_assertions)]
12318 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
12319 let n_embd = self.cfg.n_embd as usize;
12320 let _ = n_embd;
12321
12322 let h0 = e.zeros(0)?;
12323 let h = &h0;
12324 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12327 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12328 let fused_qkv = if f2b {
12329 if swa {
12330 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12331 .map(|(a, b, c)| (a, b, Some(c)))
12332 } else {
12333 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12334 .map(|(a, b)| (a, b, None))
12335 }
12336 } else {
12337 None
12338 };
12339 let (q0, k0, v0) = match fused_qkv {
12340 Some((a, b, cv)) => {
12341 let v = match cv {
12342 Some(c) => c,
12343 None => e.clone_dtod(&b)?,
12344 };
12345 (a, b, v)
12346 }
12347 None => {
12348 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12349 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12350 let v0 = if swa {
12351 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12352 } else {
12353 e.clone_dtod(&k0)?
12354 };
12355 (q0, k0, v0)
12356 }
12357 };
12358 let mut q = e.uninit(t * nh * hd)?;
12359 let mut k = e.uninit(t * nkv * hd)?;
12360 let mut v = e.uninit(t * nkv * hd)?;
12361 let ff = if swa {
12364 None
12365 } else {
12366 Some(
12367 aux.rope_freqs(e)
12368 .expect("gemma4 global rope needs rope_freqs.weight"),
12369 )
12370 };
12371 #[cfg(debug_assertions)]
12372 if let Some(ff) = ff {
12373 crate::debug_assert_tensor_stream_device(
12374 ff,
12375 &e.stream(),
12376 "gemma4_verify_attn.rope_freqs",
12377 );
12378 }
12379 e.rms_norm_qkv_rope(
12380 &q0,
12381 &k0,
12382 &v0,
12383 fa.q_norm.float_data(),
12384 fa.k_norm.float_data(),
12385 ones,
12386 &mut q,
12387 &mut k,
12388 &mut v,
12389 hd,
12390 self.gemma4_rope_dims(il),
12391 nh * t,
12392 nkv * t,
12393 pos_d,
12394 nh,
12395 nkv,
12396 base,
12397 1.0,
12398 ff,
12399 eps,
12400 )?;
12401 let kvl = cache.kv[il].as_mut().unwrap();
12402 let base_len = kvl.len;
12403 e.append_kv_quantized_rows(
12404 &k,
12405 &v,
12406 &mut kvl.k,
12407 &mut kvl.v,
12408 base_len,
12409 t,
12410 kvl.kv_dim_k,
12411 kvl.kv_dim_v,
12412 kvl.k_tok_bytes,
12413 kvl.v_tok_bytes,
12414 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12415 )?;
12416 kvl.len += t;
12417 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12418 let mut attn = e.uninit(t * nh * hd)?;
12419 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
12422 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
12425 if rows_ok && (!swa || base_len + t <= win) {
12426 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12427 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12428 if hd == 512 {
12429 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12431 e.fa_decode_rows(
12432 &q,
12433 &k_view,
12434 &v_view,
12435 &mut attn,
12436 hd,
12437 nh,
12438 nkv,
12439 base_len,
12440 t,
12441 scale,
12442 kvl.k_tok_bytes,
12443 kvl.v_tok_bytes,
12444 Some((&kvl.len_d, 0)),
12445 false,
12446 swa && crate::Engine::wkv_on(),
12447 None,
12448 )?;
12449 } else {
12450 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12454 e.fa_decode_rows_dc(
12455 &q,
12456 &k_view,
12457 &v_view,
12458 &mut attn,
12459 hd,
12460 nh,
12461 nkv,
12462 &kvl.len_d,
12463 base_len + t,
12464 t,
12465 scale,
12466 kvl.k_tok_bytes,
12467 kvl.v_tok_bytes,
12468 0,
12469 swa && crate::Engine::wkv_on(),
12470 )?;
12471 }
12472 return Ok(e.matmul(&fa.wo, &attn, t)?);
12473 }
12474 if hd == 256
12482 && swa
12483 && base_len + 1 >= win
12484 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12485 {
12486 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12487 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12488 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12489 e.fa_decode_rows_w(
12490 &q,
12491 &k_view,
12492 &v_view,
12493 &mut attn,
12494 hd,
12495 nh,
12496 nkv,
12497 &kvl.len_d,
12498 0,
12499 t,
12500 scale,
12501 win,
12502 kvl.k_tok_bytes,
12503 kvl.v_tok_bytes,
12504 None,
12505 )?;
12506 return Ok(e.matmul(&fa.wo, &attn, t)?);
12507 }
12508 for i in 0..t {
12509 let avail = base_len + i + 1;
12510 let (off_tok, t_kv) = if swa && avail > win {
12511 (avail - win, win)
12512 } else {
12513 (0, avail)
12514 };
12515 let k_view = e.view_u8_range(
12516 &kvl.k,
12517 off_tok * kvl.k_tok_bytes,
12518 (off_tok + t_kv) * kvl.k_tok_bytes,
12519 );
12520 let v_view = e.view_u8_range(
12521 &kvl.v,
12522 off_tok * kvl.v_tok_bytes,
12523 (off_tok + t_kv) * kvl.v_tok_bytes,
12524 );
12525 let qi = e.view(&q, t * nh * hd);
12526 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12527 let mut q_one = e.uninit(nh * hd)?;
12528 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12529 let mut a_one = e.uninit(nh * hd)?;
12530 if swa
12534 && avail > win
12535 && hd == 256
12536 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12537 {
12538 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12539 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12540 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12541 e.fa_decode_rows_w(
12542 &q_one,
12543 &kp,
12544 &vp,
12545 &mut a_one,
12546 hd,
12547 nh,
12548 nkv,
12549 &kvl.len_d,
12550 0,
12551 1,
12552 scale,
12553 win,
12554 kvl.k_tok_bytes,
12555 kvl.v_tok_bytes,
12556 None,
12557 )?;
12558 } else if !swa
12559 && hd == 512
12560 && avail >= crate::fa512_min_tkv()
12561 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12562 {
12563 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12564 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12565 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12566 e.fa_decode_rows(
12567 &q_one,
12568 &kp,
12569 &vp,
12570 &mut a_one,
12571 hd,
12572 nh,
12573 nkv,
12574 avail - 1,
12575 1,
12576 scale,
12577 kvl.k_tok_bytes,
12578 kvl.v_tok_bytes,
12579 Some((&kvl.len_d, 0)),
12580 false,
12581 false,
12582 None,
12583 )?;
12584 } else {
12585 e.fa_decode_kvmod(
12586 &q_one,
12587 &k_view,
12588 &v_view,
12589 &mut a_one,
12590 hd,
12591 nh,
12592 nkv,
12593 t_kv,
12594 scale,
12595 kvl.k_tok_bytes,
12596 kvl.v_tok_bytes,
12597 swa && crate::Engine::wkv_on(),
12598 )?;
12599 }
12600 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12601 }
12602 Ok(e.matmul(&fa.wo, &attn, t)?)
12603 }
12604
12605 pub(crate) fn gemma4_decode_step_h(
12608 &self,
12609 e: &Engine,
12610 token: u32,
12611 cache: &mut Cache,
12612 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12613 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12618 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12619 }
12620 if crate::pp::pp_cuts(self.layers.len()).is_some() {
12621 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12622 }
12623 let n_embd = self.cfg.n_embd as usize;
12624 let eps = self.cfg.rms_eps;
12625 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12626 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12627 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12628 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12631 let n_layers = self.layers.len();
12632 for (il, layer) in self.layers.iter().enumerate() {
12633 let (hq, hdq) = match h_carry.take() {
12634 Some(p) => p,
12635 None => {
12636 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12637 }
12638 };
12639 let Mixer::Full(fa) = &layer.mixer else {
12640 panic!("gemma4 layer {il} not full-attn")
12641 };
12642 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12643 let next_norm = if il + 1 < n_layers {
12644 Some(self.layers[il + 1].attn_norm.float_data())
12645 } else {
12646 None
12647 };
12648 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12649 x = xn;
12650 h_carry = hn;
12651 }
12652 let mut hn = e.uninit(n_embd)?;
12653 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12654 let h_seed = e.clone_dtod(&x)?;
12655 let mut ld = e.matmul(&self.output, &hn, 1)?;
12656 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12657 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
12659 let logits = e.dtoh(&ld)?;
12660 cache.pos += 1;
12661 Ok((logits, h_seed))
12662 }
12663
12664 fn gemma4_decode_layers(
12672 &self,
12673 e: &Engine,
12674 mut x: CudaSlice<f32>,
12675 lo: usize,
12676 hi: usize,
12677 pos_d: &CudaSlice<i32>,
12678 cache: &mut Cache,
12679 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12680 let n_embd = self.cfg.n_embd as usize;
12681 let eps = self.cfg.rms_eps;
12682 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12683 for il in lo..hi {
12684 let layer = &self.layers[il];
12685 let (hq, hdq) = match h_carry.take() {
12686 Some(p) => p,
12687 None => {
12689 e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12690 }
12691 };
12692 let Mixer::Full(fa) = &layer.mixer else {
12693 panic!("gemma4 layer {il} not full-attn")
12694 };
12695 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12696 let next_norm = if il + 1 < hi {
12697 Some(self.layers[il + 1].attn_norm.float_data())
12698 } else {
12699 None
12700 };
12701 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12702 x = xn;
12703 h_carry = hn;
12704 }
12705 Ok(x)
12706 }
12707
12708 fn gemma4_decode_step_h_pp2(
12716 &self,
12717 e: &Engine,
12718 token: u32,
12719 cache: &mut Cache,
12720 split: usize,
12721 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12722 if crate::pp::pp2_streams_off() {
12723 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12724 }
12725 let rt = crate::pp::Pp2Rt::get(e)?;
12726 let e0 = rt.engine(0, e);
12727 let e1 = rt.engine(1, e);
12728 let n_embd = self.cfg.n_embd as usize;
12729 let eps = self.cfg.rms_eps;
12730 let pos = cache.pos as i32;
12731
12732 let slot = {
12734 let _st0 = rt.enter(0);
12735 let pos_d = e0.htod_i32(&[pos])?;
12736 #[cfg(debug_assertions)]
12737 crate::debug_assert_tensor_stream_device(
12738 &pos_d,
12739 &e0.stream(),
12740 "gemma4_decode_step_h_pp2.stage0.pos_d",
12741 );
12742 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12743 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12744 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12745 rt.tx(0, &x, n_embd)?
12746 };
12747
12748 let _st1 = rt.enter(1);
12750 let pos_d = e1.htod_i32(&[pos])?;
12751 #[cfg(debug_assertions)]
12752 crate::debug_assert_tensor_stream_device(
12753 &pos_d,
12754 &e1.stream(),
12755 "gemma4_decode_step_h_pp2.stage1.pos_d",
12756 );
12757 let x = rt.rx(0, slot, n_embd)?;
12758 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12759
12760 let mut hn = e1.uninit(n_embd)?;
12761 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12762 let h_seed = e1.clone_dtod(&x)?;
12763 let mut ld = e1.matmul(&self.output, &hn, 1)?;
12764 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12765 e1.softcap(&mut ld, cap, self.output.out_features())?;
12766 self.gemma4_suppress(e1, &mut ld, 1)?;
12767 let logits = e1.dtoh(&ld)?;
12768 cache.pos += 1;
12769 Ok((logits, h_seed))
12770 }
12771
12772 fn gemma4_decode_step_h_pp2_samestream(
12775 &self,
12776 e: &Engine,
12777 token: u32,
12778 cache: &mut Cache,
12779 split: usize,
12780 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12781 let n_embd = self.cfg.n_embd as usize;
12782 let eps = self.cfg.rms_eps;
12783 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12784
12785 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12787 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12788 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12789
12790 let boundary_tx = e.clone_dtod(&x)?;
12792 let boundary_rx = e.clone_dtod(&boundary_tx)?;
12793
12794 let x =
12796 self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12797
12798 let mut hn = e.uninit(n_embd)?;
12799 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12800 let h_seed = e.clone_dtod(&x)?;
12801 let mut ld = e.matmul(&self.output, &hn, 1)?;
12802 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12803 e.softcap(&mut ld, cap, self.output.out_features())?;
12804 self.gemma4_suppress(e, &mut ld, 1)?;
12805 let logits = e.dtoh(&ld)?;
12806 cache.pos += 1;
12807 Ok((logits, h_seed))
12808 }
12809}
12810
12811impl HybridModel {
12830 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12833 let geometry = self
12834 .cfg
12835 .layer_geometry(il as u32)
12836 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12837 debug_assert_eq!(
12838 geometry.attention_gate,
12839 memra_gguf::config::AttentionGateKind::SeparateHead
12840 );
12841 geometry
12842 }
12843
12844 #[allow(clippy::too_many_arguments)]
12904 fn step35_attn_pre_wo(
12905 &self,
12906 e: &Engine,
12907 fa: &FullAttnLayer,
12908 mut g3: Vec<CudaSlice<f32>>,
12909 hg: Option<&CudaSlice<f32>>,
12910 gt_pre: Option<&CudaSlice<f32>>,
12911 pos_d: &CudaSlice<i32>,
12912 t: usize,
12913 cache: Option<&mut Cache>,
12914 il: usize,
12915 seq_end: usize,
12916 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12917 let geometry = self.step35_geom(il);
12918 let hd = geometry.head_dim_k as usize;
12919 let nkv = geometry.n_head_kv as usize;
12920 let nh = geometry.n_head as usize;
12921 let rbase = geometry.rope_base;
12922 let scale = geometry.attention_scale();
12923 let swa = geometry.window.is_some();
12924 let eps = self.cfg.rms_eps;
12925 let win = geometry.window.unwrap_or(0) as usize;
12926 let n_rot = geometry.n_rot as usize;
12927
12928 let v = g3.pop().unwrap();
12929 let k0 = g3.pop().unwrap();
12930 let q0 = g3.pop().unwrap();
12931
12932 let mut q = e.uninit(t * nh * hd)?;
12936 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12937 let mut k = e.uninit(t * nkv * hd)?;
12938 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12939 let ff = if geometry.rope_factors {
12940 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12941 } else {
12942 None
12943 };
12944 #[cfg(debug_assertions)]
12945 if let Some(ff) = ff {
12946 crate::debug_assert_tensor_stream_device(
12947 ff,
12948 &e.stream(),
12949 "step35_attn_pre_wo.rope_freqs",
12950 );
12951 }
12952 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12953
12954 let mut attn = e.uninit(t * nh * hd)?;
12955 match cache {
12956 Some(cache) => {
12957 let base_len = cache.kv[il].as_ref().unwrap().len;
12958 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12960 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12961 let off = if swa {
12962 let raw = base_len.saturating_sub(win - 1);
12963 if legacy_tkv || legacy_calllocal {
12964 raw
12965 } else {
12966 raw & !31usize
12967 }
12968 } else {
12969 0
12970 };
12971 {
12972 let kvl = cache.kv[il].as_mut().unwrap();
12973 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12974 let write_row = e.prepare_kv_append(kvl, off, t)?;
12975 e.append_kv_quantized_rows(
12976 &k,
12977 &v,
12978 &mut kvl.k,
12979 &mut kvl.v,
12980 write_row,
12981 t,
12982 kvl.kv_dim_k,
12983 kvl.kv_dim_v,
12984 kvl.k_tok_bytes,
12985 kvl.v_tok_bytes,
12986 crate::Engine::kv_fp8_on(),
12987 )?;
12988 kvl.len += t;
12989 let new_len = kvl.len as i32;
12990 e.set_i32_one(&mut kvl.len_d, new_len)?;
12991 }
12992 let kvl = cache.kv[il].as_ref().unwrap();
12993 let t_kv = base_len + t - off;
13016 let physical = kvl.physical_rows(off, off + t_kv)?;
13017 let k_view = e.view_u8_range(
13018 &kvl.k,
13019 physical.start * kvl.k_tok_bytes,
13020 physical.end * kvl.k_tok_bytes,
13021 );
13022 let v_view = e.view_u8_range(
13023 &kvl.v,
13024 physical.start * kvl.v_tok_bytes,
13025 physical.end * kvl.v_tok_bytes,
13026 );
13027 let swa_naive = if legacy_tkv {
13039 t_kv > win
13040 } else {
13041 seq_end > win
13042 };
13043 if swa && swa_naive {
13044 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
13057 e.sdpa_naive_w_quantized_view(
13058 &q,
13059 &k_view,
13060 &v_view,
13061 &mut attn,
13062 hd,
13063 nh,
13064 nkv,
13065 t,
13066 t_kv,
13067 scale,
13068 true,
13069 win,
13070 kvl.k_tok_bytes,
13071 kvl.v_tok_bytes,
13072 )?;
13073 } else {
13074 e.fa_prefill_view_ws_w_hd128(
13075 &q,
13076 &k_view,
13077 &v_view,
13078 &mut attn,
13079 hd,
13080 nh,
13081 nkv,
13082 t,
13083 t_kv,
13084 scale,
13085 true,
13086 win,
13087 kvl.k_tok_bytes,
13088 kvl.v_tok_bytes,
13089 )?;
13090 }
13091 } else if std::env::var("MEMRA_NOFA").is_ok() {
13092 e.sdpa_naive_quantized_view(
13093 &q,
13094 &k_view,
13095 &v_view,
13096 &mut attn,
13097 hd,
13098 nh,
13099 nkv,
13100 t,
13101 t_kv,
13102 scale,
13103 true,
13104 kvl.k_tok_bytes,
13105 kvl.v_tok_bytes,
13106 )?;
13107 } else {
13108 e.fa_prefill_view_ws(
13113 &q,
13114 &k_view,
13115 &v_view,
13116 &mut attn,
13117 hd,
13118 nh,
13119 nkv,
13120 t,
13121 t_kv,
13122 scale,
13123 true,
13124 kvl.k_tok_bytes,
13125 kvl.v_tok_bytes,
13126 crate::Engine::kv_fp8_on(),
13127 )?;
13128 }
13129 }
13130 None => {
13131 debug_assert_eq!(
13136 seq_end, t,
13137 "step35 cacheless prefill is monolithic (seq_end == t)"
13138 );
13139 if swa && seq_end > win {
13140 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13141 } else if std::env::var("MEMRA_NOFA").is_ok() {
13142 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13143 } else {
13144 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13145 }
13146 }
13147 }
13148
13149 let gw = fa
13152 .attn_gate
13153 .as_ref()
13154 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13155 let gt_owned = if gt_pre.is_none() {
13156 Some(e.matmul(
13157 gw,
13158 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
13159 t,
13160 )?)
13161 } else {
13162 None
13163 };
13164 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
13165 let mut ag = e.uninit(t * nh * hd)?;
13166 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
13167 Ok(ag)
13168 }
13169
13170 pub(crate) fn step35_attn(
13173 &self,
13174 e: &Engine,
13175 fa: &FullAttnLayer,
13176 h: &CudaSlice<f32>,
13177 pos_d: &CudaSlice<i32>,
13178 t: usize,
13179 il: usize,
13180 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13181 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
13182 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
13184 Ok(e.matmul(&fa.wo, &ag, t)?)
13185 }
13186
13187 #[allow(clippy::too_many_arguments)]
13194 pub(crate) fn step35_attn_prime(
13195 &self,
13196 e: &Engine,
13197 fa: &FullAttnLayer,
13198 h: &CudaSlice<f32>,
13199 hx: Option<&CudaSlice<u8>>,
13200 pos_d: &CudaSlice<i32>,
13201 t: usize,
13202 cache: &mut Cache,
13203 il: usize,
13204 seq_end: usize,
13205 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13206 let g3 = match hx {
13207 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
13208 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
13209 };
13210 let ag =
13211 self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
13212 Ok(e.matmul(&fa.wo, &ag, t)?)
13213 }
13214
13215 #[allow(clippy::too_many_arguments)]
13225 pub(crate) fn step35_decode_attn(
13226 &self,
13227 e: &Engine,
13228 fa: &FullAttnLayer,
13229 il: usize,
13230 h: &CudaSlice<f32>,
13231 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
13232 pos_d: &CudaSlice<i32>,
13233 cache: &mut Cache,
13234 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13235 let geometry = self.step35_geom(il);
13236 let hd = geometry.head_dim_k as usize;
13237 let nkv = geometry.n_head_kv as usize;
13238 let nh = geometry.n_head as usize;
13239 let rbase = geometry.rope_base;
13240 let scale = geometry.attention_scale();
13241 let swa = geometry.window.is_some();
13242 let eps = self.cfg.rms_eps;
13243 let win = geometry.window.unwrap_or(0) as usize;
13244 let n_rot = geometry.n_rot as usize;
13245 let n_embd = self.cfg.n_embd as usize;
13246 let gw = fa
13247 .attn_gate
13248 .as_ref()
13249 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13250
13251 let (q0, k0, v0, gt) = match pre_q {
13252 Some((hq, hdq)) => {
13253 debug_assert!(
13254 e.uses_q8_1_fast(gw),
13255 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
13256 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
13257 );
13258 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
13259 Some(t3) => t3,
13260 None => (
13261 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
13262 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
13263 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
13264 ),
13265 };
13266 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
13267 (a, b, c, gt)
13268 }
13269 None => {
13270 if e.uses_q8_1_fast(&fa.wq)
13271 && e.uses_q8_1_fast(&fa.wk)
13272 && e.uses_q8_1_fast(&fa.wv)
13273 && e.uses_q8_1_fast(gw)
13274 {
13275 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
13276 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
13277 Some(t3) => t3,
13278 None => (
13279 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
13280 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
13281 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
13282 ),
13283 };
13284 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
13285 (a, b, c, gt)
13286 } else {
13287 (
13288 e.matmul(&fa.wq, h, 1)?,
13289 e.matmul(&fa.wk, h, 1)?,
13290 e.matmul(&fa.wv, h, 1)?,
13291 e.matmul(gw, h, 1)?,
13292 )
13293 }
13294 }
13295 };
13296
13297 let mut q = e.uninit(nh * hd)?;
13298 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
13299 let mut k = e.uninit(nkv * hd)?;
13300 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
13301 let ff = if swa {
13302 None
13303 } else {
13304 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
13305 };
13306 #[cfg(debug_assertions)]
13307 if let Some(ff) = ff {
13308 crate::debug_assert_tensor_stream_device(
13309 ff,
13310 &e.stream(),
13311 "step35_decode_attn.rope_freqs",
13312 );
13313 }
13314 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
13315
13316 if std::env::var("MEMRA_NOFA").is_ok() {
13317 return Err(
13318 "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
13319 cache; unset MEMRA_NOFA to use fa_decode"
13320 .into(),
13321 );
13322 }
13323 let kvl = cache.kv[il].as_mut().unwrap();
13324 let next_len = kvl.len + 1;
13325 let (off, t_kv) = if swa && next_len > win {
13326 (next_len - win, win)
13327 } else {
13328 (0, next_len)
13329 };
13330 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
13331 e.append_kv_quantized(
13332 &k,
13333 &v0,
13334 &mut kvl.k,
13335 &mut kvl.v,
13336 write_row,
13337 kvl.kv_dim_k,
13338 kvl.kv_dim_v,
13339 kvl.k_tok_bytes,
13340 kvl.v_tok_bytes,
13341 crate::Engine::kv_fp8_on(),
13342 )?;
13343 kvl.len = next_len;
13344 let physical = kvl.physical_rows(off, off + t_kv)?;
13345 let k_view = e.view_u8_range(
13346 &kvl.k,
13347 physical.start * kvl.k_tok_bytes,
13348 physical.end * kvl.k_tok_bytes,
13349 );
13350 let v_view = e.view_u8_range(
13351 &kvl.v,
13352 physical.start * kvl.v_tok_bytes,
13353 physical.end * kvl.v_tok_bytes,
13354 );
13355 let mut attn = e.uninit(nh * hd)?;
13356 e.fa_decode_kvmod(
13357 &q,
13358 &k_view,
13359 &v_view,
13360 &mut attn,
13361 hd,
13362 nh,
13363 nkv,
13364 t_kv,
13365 scale,
13366 kvl.k_tok_bytes,
13367 kvl.v_tok_bytes,
13368 crate::Engine::kv_fp8_on(),
13369 )?;
13370
13371 let mut ag = e.uninit(nh * hd)?;
13372 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
13373 Ok(e.matmul(&fa.wo, &ag, 1)?)
13374 }
13375}
13376
13377impl HybridModel {
13386 pub fn is_gemma4_e4b(&self) -> bool {
13387 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
13388 }
13389
13390 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
13394 let g = self.cfg.gemma4.as_ref().unwrap();
13395 let swa = g.swa_pattern[il];
13396 let hd = if swa {
13397 g.key_length_swa
13398 } else {
13399 g.key_length_global
13400 } as usize;
13401 let Mixer::Full(fa) = &self.layers[il].mixer else {
13402 panic!("e4b layer {il} not full-attn")
13403 };
13404 let nh = fa.wq.out_features() / hd;
13405 let nkv = fa.wk.out_features() / hd;
13406 (
13407 hd,
13408 nkv,
13409 nh,
13410 if swa {
13411 g.rope_base_swa
13412 } else {
13413 g.rope_base_global
13414 },
13415 1.0,
13416 swa,
13417 )
13418 }
13419
13420 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
13422 self.layers[il]
13423 .gemma4
13424 .as_ref()
13425 .and_then(|b| b.e4b.as_ref())
13426 .and_then(|e4| e4.kv_share.map(|t| t as usize))
13427 }
13428
13429 fn gemma4_e4b_inp_pl(
13434 &self,
13435 e: &Engine,
13436 tokens: &[u32],
13437 x_scaled: &CudaSlice<f32>,
13438 t: usize,
13439 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13440 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
13441 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
13442 }
13443
13444 fn gemma4_e4b_inp_pl_dev(
13446 &self,
13447 e: &Engine,
13448 tok_d: &CudaSlice<u32>,
13449 x_scaled: &CudaSlice<f32>,
13450 t: usize,
13451 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13452 let aux = self.gemma4_aux.as_ref().unwrap();
13453 let m = aux.e4b.as_ref().unwrap();
13454 let n_embd = self.cfg.n_embd as usize;
13455 let n_layer = self.layers.len();
13456 let width = m.n_epl * n_layer;
13457 let tbl = m.tok_tbl_gpu.get_or_init(|| {
13458 e.upload_u8(&m.tok_embd_bytes)
13459 .expect("e4b per-layer token table upload")
13460 });
13461 let mut a =
13462 e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
13463 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
13464 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
13465 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
13466 let mut pn = e.uninit(t * width)?;
13467 e.rms_norm(
13468 &p,
13469 m.proj_norm.float_data(),
13470 &mut pn,
13471 m.n_epl,
13472 t * n_layer,
13473 self.cfg.rms_eps,
13474 )?;
13475 let mut out = e.uninit(t * width)?;
13476 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
13477 Ok(out)
13478 }
13479
13480 #[allow(clippy::too_many_arguments)]
13485 fn gemma4_e4b_attn(
13486 &self,
13487 e: &Engine,
13488 il: usize,
13489 hq: &CudaSlice<i8>,
13490 hdq: &CudaSlice<f32>,
13491 pos_d: &CudaSlice<i32>,
13492 t: usize,
13493 cache: &mut Cache,
13494 dc_bucket: Option<usize>,
13495 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13496 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
13497 let eps = self.cfg.rms_eps;
13498 let aux = self.gemma4_aux.as_ref().unwrap();
13499 let ones = aux.ones(e);
13500 #[cfg(debug_assertions)]
13501 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
13502 let Mixer::Full(fa) = &self.layers[il].mixer else {
13503 unreachable!()
13504 };
13505 let h0 = e.zeros(0)?;
13509 let h = &h0;
13510
13511 let ff = if swa {
13512 None
13513 } else {
13514 Some(
13515 aux.rope_freqs(e)
13516 .expect("e4b global rope needs rope_freqs.weight"),
13517 )
13518 };
13519 #[cfg(debug_assertions)]
13520 if let Some(ff) = ff {
13521 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13522 }
13523 let share = self.gemma4_e4b_kv_target(il);
13524 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13526 let mut q;
13527 if let Some(_tgt) = share {
13528 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13529 q = e.uninit(t * nh * hd)?;
13530 let mut kdummy = e.uninit(1)?;
13533 let mut vdummy = e.uninit(1)?;
13534 e.rms_norm_qkv_rope(
13535 &q0,
13536 &q0,
13537 &q0,
13538 fa.q_norm.float_data(),
13539 fa.q_norm.float_data(),
13540 ones,
13541 &mut q,
13542 &mut kdummy,
13543 &mut vdummy,
13544 hd,
13545 self.gemma4_rope_dims(il),
13546 nh * t,
13547 0,
13548 pos_d,
13549 nh,
13550 1,
13551 base,
13552 1.0,
13553 ff,
13554 eps,
13555 )?;
13556 } else {
13557 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13561 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13562 q = e.uninit(t * nh * hd)?;
13563 let mut k = e.uninit(t * nkv * hd)?;
13564 let mut v = e.uninit(t * nkv * hd)?;
13565 if t == 1 && cat.is_some() {
13566 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13567 e.rms_norm_qkv_rope_cat(
13568 &qkv0,
13569 fa.q_norm.float_data(),
13570 fa.k_norm.float_data(),
13571 ones,
13572 &mut q,
13573 &mut k,
13574 &mut v,
13575 hd,
13576 self.gemma4_rope_dims(il),
13577 nh,
13578 nkv,
13579 pos_d,
13580 nh,
13581 nkv,
13582 base,
13583 1.0,
13584 ff,
13585 eps,
13586 )?;
13587 } else {
13588 let (q0, k0, v0) = match if t == 1 {
13589 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13590 } else {
13591 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13594 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13595 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13596 } else {
13597 None
13598 }
13599 } {
13600 Some(triple) => triple,
13601 None => (
13602 e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13603 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13604 e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13605 ), };
13607 e.rms_norm_qkv_rope(
13610 &q0,
13611 &k0,
13612 &v0,
13613 fa.q_norm.float_data(),
13614 fa.k_norm.float_data(),
13615 ones,
13616 &mut q,
13617 &mut k,
13618 &mut v,
13619 hd,
13620 self.gemma4_rope_dims(il),
13621 nh * t,
13622 nkv * t,
13623 pos_d,
13624 nh,
13625 nkv,
13626 base,
13627 1.0,
13628 ff,
13629 eps,
13630 )?;
13631 }
13632 let kvl = cache.kv[il].as_mut().unwrap();
13633 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13637 if dc_bucket.is_some() {
13638 debug_assert!(t == 1);
13643 e.append_kv_quantized_row_dc_inc(
13645 &k,
13646 &v,
13647 &mut kvl.k,
13648 &mut kvl.v,
13649 &mut kvl.len_d,
13650 kvl.kv_dim_k,
13651 kvl.kv_dim_v,
13652 kvl.k_tok_bytes,
13653 kvl.v_tok_bytes,
13654 cls,
13655 )?;
13656 } else {
13657 e.append_kv_quantized_rows(
13658 &k,
13659 &v,
13660 &mut kvl.k,
13661 &mut kvl.v,
13662 kvl.len,
13663 t,
13664 kvl.kv_dim_k,
13665 kvl.kv_dim_v,
13666 kvl.k_tok_bytes,
13667 kvl.v_tok_bytes,
13668 cls,
13669 )?;
13670 kvl.len += t;
13671 }
13672 kv_f32 = Some((k, v));
13673 }
13674 let kvl_idx = share.unwrap_or(il);
13677 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13678 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13680 let mut attn = e.uninit(t * nh * hd)?;
13681 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13693 if let Some((kf, vf)) = &kv_f32 {
13694 if hd == 256 && t <= win {
13695 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13696 return Ok(e.matmul(&fa.wo, &attn, t)?);
13697 }
13698 if hd == 256 && swa && t > win {
13699 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13700 return Ok(e.matmul(&fa.wo, &attn, t)?);
13701 }
13702 if hd == 512 && !swa {
13703 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13704 return Ok(e.matmul(&fa.wo, &attn, t)?);
13705 }
13706 } else if share.is_some() {
13707 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13708 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13709 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13710 if hd == 256 && (!swa || t <= win) {
13711 e.fa_prefill_view(
13713 &q,
13714 &k_view,
13715 &v_view,
13716 &mut attn,
13717 hd,
13718 nh,
13719 nkv,
13720 t,
13721 t,
13722 scale,
13723 true,
13724 kvl.k_tok_bytes,
13725 kvl.v_tok_bytes,
13726 g,
13727 )?;
13728 return Ok(e.matmul(&fa.wo, &attn, t)?);
13729 }
13730 let kv_dim = nkv * hd;
13733 let mut kf = e.uninit(t * kv_dim)?;
13734 let mut vf = e.uninit(t * kv_dim)?;
13735 e.fa_dequant_kv_view_f32(
13736 &k_view,
13737 &v_view,
13738 &mut kf,
13739 &mut vf,
13740 kv_dim,
13741 kv_dim,
13742 t,
13743 kvl.k_tok_bytes,
13744 kvl.v_tok_bytes,
13745 g,
13746 )?;
13747 if hd == 512 {
13748 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13749 } else {
13750 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13751 }
13752 return Ok(e.matmul(&fa.wo, &attn, t)?);
13753 }
13754 }
13755 if let Some(bucket) = dc_bucket {
13756 assert!(t == 1);
13761 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13767 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13768 } else {
13769 bucket
13770 };
13771 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13772 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13773 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13774 if crate::Engine::wpf_level() >= 1 {
13782 e.prefetch_weight_l2(&fa.wo)?;
13783 }
13784 if e.uses_q8_1_fast(&fa.wo) {
13787 let mut oq = e.alloc_i8_uninit(nh * hd)?;
13788 let mut od = e.zeros(nh * hd / 32)?;
13789 e.fa_decode_dc_q8(
13790 &q,
13791 &k_view,
13792 &v_view,
13793 &mut attn,
13794 hd,
13795 nh,
13796 nkv,
13797 &kvl.len_d,
13798 bucket,
13799 scale,
13800 kvl.k_tok_bytes,
13801 kvl.v_tok_bytes,
13802 g,
13803 Some((&mut oq, &mut od)),
13804 )?;
13805 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13806 }
13807 e.fa_decode_dc(
13808 &q,
13809 &k_view,
13810 &v_view,
13811 &mut attn,
13812 hd,
13813 nh,
13814 nkv,
13815 &kvl.len_d,
13816 bucket,
13817 scale,
13818 kvl.k_tok_bytes,
13819 kvl.v_tok_bytes,
13820 g,
13821 )?;
13822 return Ok(e.matmul(&fa.wo, &attn, t)?);
13823 }
13824 for i in 0..t {
13825 let avail = base_len + i + 1;
13826 let (off_tok, t_kv) = if swa && avail > win {
13827 (avail - win, win)
13828 } else {
13829 (0, avail)
13830 };
13831 let k_view = e.view_u8_range(
13832 &kvl.k,
13833 off_tok * kvl.k_tok_bytes,
13834 (off_tok + t_kv) * kvl.k_tok_bytes,
13835 );
13836 let v_view = e.view_u8_range(
13837 &kvl.v,
13838 off_tok * kvl.v_tok_bytes,
13839 (off_tok + t_kv) * kvl.v_tok_bytes,
13840 );
13841 let qv = e.view(&q, t * nh * hd);
13842 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13843 let mut q_one = e.uninit(nh * hd)?;
13844 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13845 let mut a_one = e.uninit(nh * hd)?;
13846 e.fa_decode_kvmod(
13850 &q_one,
13851 &k_view,
13852 &v_view,
13853 &mut a_one,
13854 hd,
13855 nh,
13856 nkv,
13857 t_kv,
13858 scale,
13859 kvl.k_tok_bytes,
13860 kvl.v_tok_bytes,
13861 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13862 )?;
13863 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13864 }
13865 Ok(e.matmul(&fa.wo, &attn, t)?)
13866 }
13867
13868 fn gemma4_e4b_trunk(
13873 &self,
13874 e: &Engine,
13875 tokens: &[u32],
13876 pos0: usize,
13877 cache: &mut Cache,
13878 head_last: bool,
13879 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13880 let n_embd = self.cfg.n_embd as usize;
13881 let t = tokens.len();
13882 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13883 let pos_d = e.htod_i32(&pos)?;
13884 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13885 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13886 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13887 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13888 }
13889
13890 fn gemma4_e4b_trunk_core(
13894 &self,
13895 e: &Engine,
13896 x_in: CudaSlice<f32>,
13897 inp_pl: CudaSlice<f32>,
13898 pos_d: &CudaSlice<i32>,
13899 t: usize,
13900 cache: &mut Cache,
13901 dc_bucket: Option<usize>,
13902 cap_logits: bool,
13903 head_last: bool,
13904 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13905 let n_embd = self.cfg.n_embd as usize;
13906 let eps = self.cfg.rms_eps;
13907 let n_layer = self.layers.len();
13908 let mut x = x_in;
13909 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13910 let n_epl = aux_e4b.n_epl;
13911
13912 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13918 for il in 0..n_layer {
13919 let layer = &self.layers[il];
13920 let (hq, hdq) = match h_carry.take() {
13921 Some(p) => p,
13922 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13923 };
13924 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13925 let bits = layer.gemma4.as_ref().unwrap();
13928 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13929 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13940 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13941 e,
13942 layer,
13943 &o,
13944 &x,
13945 t,
13946 Some(layer.post_attn_norm.float_data()),
13947 fuse_exit,
13948 )?;
13949 let mut resid = e.uninit(t * n_embd)?;
13950 let g = if fuse_exit {
13956 let (rq, rd) = e.rms_pre_add_q8_1(
13958 &sn,
13959 bits.post_ffw_norm.float_data(),
13960 &attn_out,
13961 &mut resid,
13962 n_embd,
13963 t,
13964 self.cfg.rms_eps,
13965 )?;
13966 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13967 } else {
13968 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13969 e.matmul(&e4b.inp_gate, &resid, t)?
13970 };
13971 let mut act = e.uninit(t * n_epl)?;
13972 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13973 let ipv = e.view(&inp_pl, n_epl * n_layer);
13974 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13975 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13976 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13977 } else {
13978 let mut inp_this = e.uninit(t * n_epl)?;
13979 e.copy_rows_strided(
13980 &inp_pl,
13981 &mut inp_this,
13982 n_epl,
13983 t,
13984 n_epl * n_layer,
13985 il * n_epl,
13986 )?;
13987 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13988 e.matmul(&e4b.proj, &act, t)?
13989 };
13990 let next_norm = if il + 1 < n_layer {
13993 self.layers[il + 1].attn_norm.float_data()
13994 } else {
13995 self.output_norm.float_data()
13996 };
13997 let mut xn = e.uninit(t * n_embd)?;
13998 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13999 &y,
14000 e4b.post_norm.float_data(),
14001 &resid,
14002 bits.layer_scale,
14003 next_norm,
14004 &mut xn,
14005 n_embd,
14006 t,
14007 eps,
14008 )?;
14009 h_carry = Some(pair);
14010 x = xn;
14011 }
14012 let (oq, odq) = h_carry.take().unwrap();
14016 let h0 = e.zeros(0)?;
14017 let hm = if head_last { 1 } else { t };
14018 let (hq, hd) = if head_last && t > 1 {
14019 let mut q1 = e.uninit_i8(n_embd)?;
14020 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
14021 let nb = n_embd / 32;
14022 let mut d1 = e.uninit(nb)?;
14023 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
14024 (q1, d1)
14025 } else {
14026 (oq, odq)
14027 };
14028 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
14029 if cap_logits {
14033 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
14034 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
14035 }
14036 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
14038 }
14039
14040 pub fn gemma4_e4b_decode_step_t_am_dev(
14047 &self,
14048 e: &Engine,
14049 tok_d: &CudaSlice<u32>,
14050 t: usize,
14051 pos0: usize,
14052 cache: &mut Cache,
14053 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14054 let n_embd = self.cfg.n_embd as usize;
14055 let eps = self.cfg.rms_eps;
14056 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
14057 let pos_d = e.htod_i32(&pos)?;
14058 let embd_gpu = self
14059 .embd_gpu
14060 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
14061 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
14062 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
14063 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
14064 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
14065 let (ld, xp) =
14066 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
14067 let n_vocab = self.output.out_features();
14070 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
14071 for i in 0..t {
14072 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
14073 }
14074 let mut hn = e.uninit(t * n_embd)?;
14075 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
14076 cache.pos += t;
14077 Ok((vam, hn))
14078 }
14079
14080 pub(crate) fn gemma4_e4b_decode_step_t_h(
14083 &self,
14084 e: &Engine,
14085 tokens: &[u32],
14086 pos0: usize,
14087 cache: &mut Cache,
14088 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14089 let n_embd = self.cfg.n_embd as usize;
14090 let eps = self.cfg.rms_eps;
14091 let t = tokens.len();
14092 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
14093 let mut hn = e.uninit(t * n_embd)?;
14094 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
14095 cache.pos += t;
14096 Ok((e.dtoh(&ld)?, hn))
14097 }
14098
14099 pub fn gemma4_e4b_decode_step_dcg(
14105 &self,
14106 e: &Engine,
14107 token_d: &mut CudaSlice<u32>,
14108 pos_d: &mut CudaSlice<i32>,
14109 embd_gpu: &CudaSlice<u8>,
14110 embd_qt: i32,
14111 embd_rb: usize,
14112 cache: &mut Cache,
14113 n_vocab: usize,
14114 bucket: usize,
14115 ) -> Result<(), Box<dyn std::error::Error>> {
14116 let n_embd = self.cfg.n_embd as usize;
14117 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
14118 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
14119 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
14120 let (ld, _x) =
14121 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
14122 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
14123 e.inc_seqlen(pos_d)?;
14124 Ok(())
14125 }
14126
14127 #[allow(clippy::too_many_arguments)]
14135 pub fn gemma4_e4b_decode_step_dc(
14136 &self,
14137 e: &Engine,
14138 token_d: &CudaSlice<u32>,
14139 pos_d: &mut CudaSlice<i32>,
14140 embd_gpu: &CudaSlice<u8>,
14141 embd_qt: i32,
14142 embd_rb: usize,
14143 cache: &mut Cache,
14144 n_vocab: usize,
14145 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
14146 let n_embd = self.cfg.n_embd as usize;
14147 let eps = self.cfg.rms_eps;
14148 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
14149 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
14150 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
14151 let (ld, _x) =
14152 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
14153 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
14154 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
14155 e.inc_seqlen(pos_d)?;
14156 cache.pos += 1;
14157 let _ = eps;
14158 Ok(tok_out)
14159 }
14160
14161 pub(crate) fn gemma4_e4b_decode_step_h(
14164 &self,
14165 e: &Engine,
14166 token: u32,
14167 cache: &mut Cache,
14168 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14169 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
14170 let logits = e.dtoh(&ld)?;
14171 cache.pos += 1;
14172 Ok((logits, x))
14173 }
14174
14175 pub(crate) fn gemma4_e4b_prime(
14179 &self,
14180 e: &Engine,
14181 tokens: &[u32],
14182 cache: &mut Cache,
14183 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14184 if cache.pos != 0 {
14187 return Err(
14188 "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
14189 call or decode tokenwise"
14190 .into(),
14191 );
14192 }
14193 let n_embd = self.cfg.n_embd as usize;
14194 let t = tokens.len();
14195 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
14196 cache.pos += t;
14197 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
14199 let row = xv.slice((t - 1) * n_embd..t * n_embd);
14200 let mut h_seed = e.uninit(n_embd)?;
14201 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
14202 Ok((last, h_seed, x))
14203 }
14204
14205 pub(crate) fn gemma4_e4b_forward(
14207 &self,
14208 e: &Engine,
14209 tokens: &[u32],
14210 last_only: bool,
14211 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
14212 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
14213 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
14214 Ok(e.dtoh(&ld)?) }
14216}
14217
14218#[cfg(test)]
14219mod prime_chunk_schedule_tests {
14220 use super::{
14221 PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
14222 fixed_prime_chunk_ranges_for_ring,
14223 };
14224
14225 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
14226 ranges.iter().map(|(start, end)| end - start).collect()
14227 }
14228
14229 fn auto_chunk(t: usize) -> usize {
14230 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
14231 }
14232
14233 #[test]
14234 fn fixed_schedule_retains_measured_geometry() {
14235 assert_eq!(
14236 sizes(&fixed_prime_chunk_ranges(461, 128)),
14237 vec![128, 128, 128, 77]
14238 );
14239 assert_eq!(
14240 sizes(&fixed_prime_chunk_ranges(1833, 230)),
14241 vec![230, 230, 230, 230, 230, 230, 230, 223]
14242 );
14243 assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
14244 let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
14245 assert_eq!(capped, vec![4096, 4088, 16]);
14246 assert!(capped.iter().all(|&rows| rows <= 4096));
14247 assert_eq!(
14248 sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
14249 vec![4100],
14250 "flag-off schedule remains byte-for-byte the legacy monolithic tail",
14251 );
14252 }
14253
14254 #[test]
14255 fn dynamic_schedule_matches_registered_shapes() {
14256 let cases = [
14257 (461, vec![64, 141, 132, 124]),
14258 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
14259 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
14260 ];
14261 for (t, expected) in cases {
14262 let chunk = auto_chunk(t);
14263 let fixed = fixed_prime_chunk_ranges(t, chunk);
14264 assert_eq!(
14265 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
14266 expected
14267 );
14268 }
14269 }
14270
14271 #[test]
14272 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
14273 for t in 256..=8192 {
14274 let chunk = auto_chunk(t);
14275 let fixed = fixed_prime_chunk_ranges(t, chunk);
14276 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
14277 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
14278 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
14279 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
14280 for pair in dynamic.windows(2) {
14281 assert_eq!(pair[0].1, pair[1].0, "T={t}");
14282 }
14283 assert!(
14284 dynamic
14285 .iter()
14286 .all(|(start, end)| end - start >= PRIME_MIN_T),
14287 "T={t} sizes={:?}",
14288 sizes(&dynamic)
14289 );
14290 if dynamic.len() >= 3 {
14291 let chunk_sizes = sizes(&dynamic);
14292 assert!(
14293 chunk_sizes[0] < chunk_sizes[1],
14294 "T={t} sizes={chunk_sizes:?}"
14295 );
14296 assert!(
14297 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
14298 "T={t} sizes={chunk_sizes:?}"
14299 );
14300 }
14301 }
14302 }
14303}
14304
14305#[cfg(test)]
14306mod page_prefetch_tests {
14307 use super::{
14308 grouped_worker_prefetch_position, page_prefetch_positions,
14309 page_prefetch_window_from_values, worker_prefetch_positions,
14310 };
14311
14312 #[test]
14313 fn page_prefetch_window_keeps_existing_opt_in_default() {
14314 assert_eq!(page_prefetch_window_from_values(false, None), 0);
14315 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
14316 assert_eq!(page_prefetch_window_from_values(true, None), 1);
14317 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
14318 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
14319 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
14320 }
14321
14322 #[test]
14323 fn rolling_page_prefetch_advises_each_future_expert_once() {
14324 let advised: Vec<_> = (0..7)
14325 .flat_map(|position| page_prefetch_positions(position, 7, 3))
14326 .collect();
14327 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
14328
14329 let one_ahead: Vec<_> = (0..4)
14330 .flat_map(|position| page_prefetch_positions(position, 4, 1))
14331 .collect();
14332 assert_eq!(one_ahead, vec![1, 2, 3]);
14333 assert!(page_prefetch_positions(0, 4, 0).is_empty());
14334 }
14335
14336 #[test]
14337 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
14338 assert_eq!(grouped_worker_prefetch_position(0, None), None);
14339 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
14340 .chain(
14341 (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
14342 )
14343 .collect();
14344 assert_eq!(positions, vec![0, 1, 2, 3]);
14345 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
14346 }
14347
14348 #[test]
14349 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
14350 let queued: Vec<_> = (0..8)
14351 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
14352 .collect();
14353 assert_eq!(queued, (0..8).collect::<Vec<_>>());
14354
14355 let one_at_a_time: Vec<_> = (0..4)
14356 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
14357 .collect();
14358 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
14359 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
14360 }
14361}
14362
14363pub struct G4DcSlots {
14364 x: CudaSlice<f32>,
14365 xn: CudaSlice<f32>,
14366 cur: CudaSlice<f32>,
14367 hq: CudaSlice<i8>,
14368 hd_: CudaSlice<f32>,
14369 q0: CudaSlice<f32>,
14370 k0: CudaSlice<f32>,
14371 v0: CudaSlice<f32>,
14372 q: CudaSlice<f32>,
14373 k: CudaSlice<f32>,
14374 v: CudaSlice<f32>,
14375 attn: CudaSlice<f32>,
14376 o: CudaSlice<f32>,
14377 attn_out: CudaSlice<f32>,
14378 zsh: CudaSlice<f32>,
14379 zq: CudaSlice<i8>,
14380 zd: CudaSlice<f32>,
14381 gate: CudaSlice<f32>,
14382 up: CudaSlice<f32>,
14383 act: CudaSlice<f32>,
14384 actq: CudaSlice<i8>,
14385 actd: CudaSlice<f32>,
14386 f0: CudaSlice<f32>,
14387 sn: CudaSlice<f32>,
14388 hn: CudaSlice<f32>,
14389 logits: CudaSlice<f32>,
14390}