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 _ => false,
374 }
375}
376
377fn moe_prewarm_enabled() -> bool {
380 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
381 *E.get_or_init(|| {
382 std::env::var("MEMRA_MOE_PREWARM")
383 .map(|v| v != "0")
384 .unwrap_or(true)
385 })
386}
387
388fn cpu_expert_profile_admit_enabled() -> bool {
392 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
393 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
394}
395
396pub const PRIME_MIN_T: usize = 16;
400const PRIME_PIPE_MICROBATCHES: usize = 8;
401const PRIME_PIPE_MIN_CHUNK: usize = 128;
402const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
403const PRIME_PIPE_LINEAR_WORK: usize = 8;
404
405fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
406 crate::pp::prime_pp_on()
407 && !crate::pp::pp2_streams_off()
408 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
409}
410
411pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
415 if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
416 let parsed = value
417 .parse::<usize>()
418 .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
419 return if crate::cache::swa_ring_on() {
420 if parsed == 0 {
421 crate::cache::PRIME_CHUNK_MAX_TOKENS
422 } else {
423 parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
424 }
425 } else {
426 parsed
427 };
428 }
429 let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
430 if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
431 chunk.min(
432 t.div_ceil(PRIME_PIPE_MICROBATCHES)
433 .max(PRIME_PIPE_MIN_CHUNK),
434 )
435 } else {
436 chunk
437 }
438}
439
440fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
441 fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
442}
443
444fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
445 if chunk == 0 || t <= chunk {
446 return vec![(0, t)];
447 }
448 let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
449 let mut start = 0usize;
450 while start < t {
451 let mut end = (start + chunk).min(t);
452 if t - end > 0 && t - end < PRIME_MIN_T {
453 if ring_on {
454 let shifted = t - PRIME_MIN_T;
455 end = if shifted > start { shifted } else { t };
456 } else {
457 end = t;
458 }
459 }
460 ranges.push((start, end));
461 start = end;
462 }
463 ranges
464}
465
466fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
467 let prefix = prefix as u128;
468 prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
469}
470
471fn dynamic_prime_chunk_ranges(
472 t: usize,
473 fixed_chunk: usize,
474 fixed: &[(usize, usize)],
475) -> Vec<(usize, usize)> {
476 let n = fixed.len();
477 if n < 3 {
478 return fixed.to_vec();
479 }
480
481 let max_first = t - (n - 1) * PRIME_MIN_T;
482 let first = fixed_chunk
483 .div_ceil(2)
484 .max(PRIME_PIPE_EDGE_MIN_CHUNK)
485 .min(max_first);
486 let mut ranges = Vec::with_capacity(n);
487 ranges.push((0, first));
488
489 let first_work = prime_chunk_work(first, t);
490 let work_span = prime_chunk_work(t, t) - first_work;
491 let denominator = (n - 1) as u128;
492 let mut previous = first;
493 for boundary in 1..n - 1 {
494 let target = first_work * denominator + work_span * (boundary as u128);
495 let remaining = n - 1 - boundary;
496 let mut low = previous + PRIME_MIN_T;
497 let mut high = t - remaining * PRIME_MIN_T;
498 while low < high {
499 let mid = low + (high - low) / 2;
500 if prime_chunk_work(mid, t) * denominator >= target {
501 high = mid;
502 } else {
503 low = mid + 1;
504 }
505 }
506 ranges.push((previous, low));
507 previous = low;
508 }
509 ranges.push((previous, t));
510 ranges
511}
512
513pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
517 let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
518 let chunk = prime_chunk_tokens(t, n_layers);
519 let fixed = fixed_prime_chunk_ranges(t, chunk);
520 let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
521 Ok(value) => value == "dynamic",
522 Err(_) => true,
523 };
524 if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
525 fixed
526 } else {
527 dynamic_prime_chunk_ranges(t, chunk, &fixed)
528 }
529}
530
531impl HybridModel {
532 fn prime_trace_path() -> Option<&'static str> {
537 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
538 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
539 .as_deref()
540 }
541
542 pub fn forward(
544 &self,
545 e: &Engine,
546 tokens: &[u32],
547 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
548 if self.is_gemma4_e4b() {
549 return self.gemma4_e4b_forward(e, tokens, false);
550 }
551 if self.cfg.gemma4.is_some() {
552 return self.gemma4_forward(e, tokens, false);
553 }
554 let cfg = &self.cfg;
555 let n_embd = cfg.n_embd as usize;
556 let t = tokens.len();
557 let eps = cfg.rms_eps;
558 let pos: Vec<i32> = (0..t as i32).collect();
559 let pos_d = e.htod_i32(&pos)?;
560
561 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
564 let mut h = e.uninit(t * n_embd)?;
566 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
567
568 let mixed = match &layer.mixer {
569 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
570 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
571 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
572 };
573
574 let mut x1 = e.uninit(t * n_embd)?;
576 e.add(&x, &mixed, &mut x1, t * n_embd)?;
577
578 let mut z = e.uninit(t * n_embd)?;
580 e.rms_norm(
581 &x1,
582 layer.post_attn_norm.float_data(),
583 &mut z,
584 n_embd,
585 t,
586 eps,
587 )?;
588 let ffn_out = match &layer.ffn {
589 crate::hybrid::Ffn::Dense {
590 ffn_gate,
591 ffn_up,
592 ffn_down,
593 } => {
594 let n_ff = ffn_gate.out_features();
595 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
596 let up = g2.pop().unwrap();
597 let gate = g2.pop().unwrap();
598 let mut act = e.uninit(t * n_ff)?;
599 Self::ffn_act_lim(
604 e,
605 &self.cfg,
606 &gate,
607 &up,
608 1.0,
609 1.0,
610 self.cfg.clamp_shexp_at(il as u32),
611 &mut act,
612 t * n_ff,
613 )?;
614 e.matmul(ffn_down, &act, t)?
615 }
616 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
617 };
618 let mut x2 = e.uninit(t * n_embd)?;
619 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
620 x = x2;
621 }
622
623 let mut hn = e.uninit(t * n_embd)?;
624 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
625 let logits = e.matmul(&self.output, &hn, t)?;
626 Ok(e.dtoh(&logits)?)
627 }
628
629 pub fn forward_last(
635 &self,
636 e: &Engine,
637 tokens: &[u32],
638 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
639 if self.cfg.gemma4.is_some() {
640 return self.gemma4_forward(e, tokens, true);
641 }
642 let cfg = &self.cfg;
643 let n_embd = cfg.n_embd as usize;
644 let t = tokens.len();
645 let eps = cfg.rms_eps;
646 let pos: Vec<i32> = (0..t as i32).collect();
647 let pos_d = e.htod_i32(&pos)?;
648
649 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
653 for (il, layer) in self.layers.iter().enumerate() {
654 let mut h = e.uninit(t * n_embd)?;
655 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
656 if probe {
657 e.stream().synchronize()?;
658 eprintln!("[probe] L{il} norm ok");
659 }
660 let mixed = match &layer.mixer {
661 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
662 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
663 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
664 };
665 if probe {
666 e.stream().synchronize()?;
667 eprintln!("[probe] L{il} mixer ok");
668 }
669 let mut x1 = e.uninit(t * n_embd)?;
670 e.add(&x, &mixed, &mut x1, t * n_embd)?;
671 let mut z = e.uninit(t * n_embd)?;
672 e.rms_norm(
673 &x1,
674 layer.post_attn_norm.float_data(),
675 &mut z,
676 n_embd,
677 t,
678 eps,
679 )?;
680 let ffn_out = match &layer.ffn {
681 crate::hybrid::Ffn::Dense {
682 ffn_gate,
683 ffn_up,
684 ffn_down,
685 } => {
686 let n_ff = ffn_gate.out_features();
687 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
688 let up = g2.pop().unwrap();
689 let gate = g2.pop().unwrap();
690 let mut act = e.uninit(t * n_ff)?;
691 Self::ffn_act_lim(
693 e,
694 &self.cfg,
695 &gate,
696 &up,
697 1.0,
698 1.0,
699 self.cfg.clamp_shexp_at(il as u32),
700 &mut act,
701 t * n_ff,
702 )?;
703 e.matmul(ffn_down, &act, t)?
704 }
705 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
706 };
707 if probe {
708 e.stream().synchronize()?;
709 eprintln!("[probe] L{il} ffn ok");
710 }
711 let mut x2 = e.uninit(t * n_embd)?;
712 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
713 x = x2;
714 }
715 let mut hn = e.uninit(t * n_embd)?;
717 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
718 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)?;
721 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
722 let logits = e.matmul(&self.output, &hlast, 1)?; Ok(e.dtoh(&logits)?)
724 }
725
726 pub fn prime_cache(
758 &self,
759 e: &Engine,
760 tokens: &[u32],
761 cache: &mut Cache,
762 queued_after: usize,
763 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
764 self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
765 }
766
767 pub fn prime_cache_overlaid(
773 &self,
774 e: &Engine,
775 tokens: &[u32],
776 cache: &mut Cache,
777 queued_after: usize,
778 overlay: Option<&crate::vision::EmbedOverlay>,
779 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
780 let n_embd = self.cfg.n_embd as usize;
781 let t = tokens.len();
782 assert!(
786 t >= PRIME_MIN_T,
787 "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
788 );
789 assert!(
790 cache.pos + t <= cache.max_ctx,
791 "prime_cache: prompt exceeds cache max_ctx"
792 );
793
794 if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
806 if self.is_gemma4_e4b() {
807 if overlay.is_some() {
808 return Err(
809 "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
810 );
811 }
812 return self.gemma4_e4b_prime(e, tokens, cache);
813 }
814 return self.gemma4_prime(e, tokens, cache, overlay);
819 }
820 let ranges = prime_chunk_ranges(t, self.layers.len());
821 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
857 let seq_end = if legacy_calllocal {
858 cache.pos + t
859 } else {
860 cache.pos + t + queued_after
861 };
862 if ranges.len() == 1 {
863 return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
864 }
865 if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
870 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
871 if overlay.is_some() {
872 return Err(
873 "vision embedding overlay + pipelined PP prime unsupported (v1); \
874 run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
875 .into(),
876 );
877 }
878 if crate::pp::pp_multi_stream_same_device() {
879 return Err(
880 "prime chunk pipeline refused with 2 stage streams on one device — \
881 that concurrent-stream placement remains quarantined by the deferred \
882 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
883 the serial split."
884 .into(),
885 );
886 }
887 return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
888 }
889 }
890 let mut hiddens = e.uninit(t * n_embd)?;
891 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
892 for &(start, end) in &ranges {
893 if let Some(taps) = cache.dflash_taps.as_mut() {
895 taps.base = start;
896 }
897 let (l, hs, x) =
898 self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
899 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
900 last = Some((l, hs));
901 }
902 let (logits, h_seed) = last.unwrap();
903 Ok((logits, h_seed, hiddens))
904 }
905
906 fn prime_cache_pp2_pipelined(
911 &self,
912 e: &Engine,
913 tokens: &[u32],
914 cache: &mut Cache,
915 seq_end: usize,
916 ranges: &[(usize, usize)],
917 fence: &[usize],
918 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
919 debug_assert_eq!(fence.len(), 3);
920 debug_assert!(ranges.len() >= 2);
921 let rt = crate::pp::PpNRt::get(e)?;
922 assert_eq!(
923 rt.n_stages(),
924 2,
925 "prime pipeline requires exactly two PP stages"
926 );
927 let n_embd = self.cfg.n_embd as usize;
928 let t = tokens.len();
929 let initial_base = cache.pos;
930 let caller_stream = e.stream();
931
932 rt.fence_stages_behind(&caller_stream)?;
937 let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
938 rt.prepare_overlap_slots(0, max_payload)?;
939
940 let mut hiddens = e.uninit(t * n_embd)?;
941 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
942 let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
943 let (cache0, cache1) = stage_caches.parts();
944 let (first_start, first_end) = ranges[0];
945 let mut slot = self.prime_pp2_stage0_enqueue(
946 e,
947 rt,
948 &tokens[first_start..first_end],
949 cache0,
950 seq_end,
951 fence,
952 initial_base + first_start,
953 true,
954 )?;
955 cache0.pos = initial_base + first_end;
956
957 for (i, &(start, end)) in ranges.iter().enumerate() {
958 let base = initial_base + start;
959 debug_assert_eq!(
960 cache1.pos, base,
961 "stage 1 must drain chunks in original position order"
962 );
963 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
964 let next_base = initial_base + next_start;
965 debug_assert_eq!(
966 cache0.pos, next_base,
967 "stage 0 must issue chunks in original position order"
968 );
969 let cache0_stage = &mut *cache0;
970 std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
975 let stage0 = scope.spawn(move || -> Result<usize, String> {
976 let next = self
977 .prime_pp2_stage0_enqueue(
978 e,
979 rt,
980 &tokens[next_start..next_end],
981 cache0_stage,
982 seq_end,
983 fence,
984 next_base,
985 true,
986 )
987 .map_err(|err| err.to_string())?;
988 cache0_stage.pos = initial_base + next_end;
989 Ok(next)
990 });
991 let x = self.prime_pp2_stage1_enqueue(
992 e,
993 rt,
994 slot,
995 end - start,
996 cache1,
997 seq_end,
998 fence,
999 base,
1000 true,
1001 )?;
1002 let out = {
1003 rt.bind_stage(1)?;
1004 let _st1 = rt.enter(1);
1005 let e1 = rt.engine(1, e);
1006 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1007 };
1008 let next = stage0
1009 .join()
1010 .map_err(|_| "pipeprime stage-0 host walker panicked")?
1011 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1012 Ok((out, Some(next)))
1013 })?
1014 } else {
1015 let x = self.prime_pp2_stage1_enqueue(
1016 e,
1017 rt,
1018 slot,
1019 end - start,
1020 cache1,
1021 seq_end,
1022 fence,
1023 base,
1024 true,
1025 )?;
1026 let out = {
1027 rt.bind_stage(1)?;
1028 let _st1 = rt.enter(1);
1029 let e1 = rt.engine(1, e);
1030 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1031 };
1032 (out, None)
1033 };
1034
1035 rt.publish_to(1, &caller_stream)?;
1036 e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1037 last = Some((out.0, out.1));
1038 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1039
1040 if let Some(next) = next_slot {
1041 rt.fence_stages_behind(&caller_stream)?;
1046 slot = next;
1047 }
1048 }
1049
1050 debug_assert_eq!(cache0.pos, initial_base + t);
1051 debug_assert_eq!(cache1.pos, initial_base + t);
1052 let (logits, h_seed) = last.unwrap();
1053 Ok((logits, h_seed, hiddens))
1054 }
1055
1056 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1063 if Engine::gdn_db_on()
1064 && Engine::gdn_chunked_enabled()
1065 && t >= 16
1066 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1067 && num_k * 2 == num_v
1068 {
1069 num_k
1070 } else {
1071 num_v
1072 }
1073 }
1074
1075 fn f16out_on(e: &Engine, t: usize) -> bool {
1080 crate::f16_ffi::pp_f16_enabled()
1081 && t >= 16
1082 && !e.verify_exact_on()
1083 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1084 }
1085
1086 pub fn prime_slabs_get(
1094 &self,
1095 e: &Engine,
1096 t: usize,
1097 n_embd: usize,
1098 n_ff_max: usize,
1099 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1100 let mut slabs = self.prime_slabs.lock().unwrap();
1101 let dev = e.ctx().ordinal();
1102 let need_new = match slabs.get(&dev) {
1103 None => true,
1104 Some(sl) => sl.lock().unwrap().t_cap < t,
1105 };
1106 if need_new {
1107 slabs.insert(
1108 dev,
1109 std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1110 t_cap: t,
1111 h: e.uninit(t * n_embd)?,
1112 x1: e.uninit(t * n_embd)?,
1113 z: e.uninit(t * n_embd)?,
1114 act: e.uninit(t * n_ff_max)?,
1115 xa: e.uninit(t * n_embd)?,
1116 xb: e.uninit(t * n_embd)?,
1117 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1118 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1119 gate: e.uninit(t * n_ff_max)?,
1120 up: e.uninit(t * n_ff_max)?,
1121 ffn_out: e.uninit(t * n_embd)?,
1122 seg_glue: Vec::new(),
1123 mixed: e.uninit(t * n_embd)?,
1124 seg_mid: Vec::new(),
1125 seg_t: 0,
1126 })),
1127 );
1128 }
1129 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1130 }
1131
1132 fn prime_chunk(
1136 &self,
1137 e: &Engine,
1138 tokens: &[u32],
1139 cache: &mut Cache,
1140 seq_end: usize,
1141 chunk_off: usize,
1142 overlay: Option<&crate::vision::EmbedOverlay>,
1143 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1144 if crate::pp::pp_host_bounce_active()
1145 && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1146 {
1147 return Err(
1148 "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1149 has no active prime stage split and would peer-read remote weights; keep \
1150 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1151 .into(),
1152 );
1153 }
1154 if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1163 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1164 if overlay.is_some() {
1165 return Err("vision embedding overlay + PP prime unsupported (v1); \
1166 run single-device or MEMRA_PRIME_PP=0"
1167 .into());
1168 }
1169 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1170 }
1171 }
1172 if crate::pp::pp_host_bounce_active() {
1173 return Err(
1174 "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1175 refusing an unsplit remote-weight walk"
1176 .into(),
1177 );
1178 }
1179 let t = tokens.len();
1180 let base = cache.pos;
1181 debug_assert!(
1182 seq_end >= base + t,
1183 "prime_chunk: seq_end must cover this chunk"
1184 );
1185 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1186 let pos_d = e.htod_i32(&pos)?;
1187
1188 let mut x_embed = self.embed(e, tokens)?; if let Some(ov) = overlay {
1190 let n_embd = self.cfg.n_embd as usize;
1194 for &(pos, row_off, n_rows) in &ov.spans {
1195 let lo = pos.max(chunk_off);
1196 let hi = (pos + n_rows).min(chunk_off + t);
1197 if lo < hi {
1198 let src_row = row_off + (lo - pos);
1199 let view = ov
1200 .rows
1201 .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
1202 e.copy_view_into(
1203 &mut x_embed,
1204 (lo - chunk_off) * n_embd,
1205 &view,
1206 (hi - lo) * n_embd,
1207 )?;
1208 }
1209 }
1210 }
1211 let x = self.prime_layers(
1212 e,
1213 x_embed,
1214 0,
1215 self.layers.len(),
1216 &pos_d,
1217 t,
1218 base,
1219 cache,
1220 seq_end,
1221 )?;
1222 self.prime_chunk_epilogue(e, x, t, cache)
1223 }
1224
1225 #[allow(clippy::too_many_arguments)]
1241 fn prime_layers(
1242 &self,
1243 e: &Engine,
1244 x_in: CudaSlice<f32>,
1245 lo: usize,
1246 hi: usize,
1247 pos_d: &CudaSlice<i32>,
1248 t: usize,
1249 base: usize,
1250 cache: &mut Cache,
1251 seq_end: usize,
1252 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1253 let cfg = &self.cfg;
1254 let n_embd = cfg.n_embd as usize;
1255 let eps = cfg.rms_eps;
1256 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1260 let n_ff_max = self
1266 .layers
1267 .iter()
1268 .map(|l| match &l.ffn {
1269 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1270 _ => n_embd,
1271 })
1272 .max()
1273 .unwrap_or(n_embd)
1274 .max(n_embd);
1275 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1276 let slab = if use_slabs {
1277 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1278 } else {
1279 None
1280 };
1281 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1282 let mut x_own; type SlabRefs<'a> = (
1284 &'a mut CudaSlice<f32>,
1285 &'a mut CudaSlice<f32>,
1286 &'a mut CudaSlice<f32>,
1287 &'a mut CudaSlice<f32>,
1288 &'a mut CudaSlice<u8>,
1289 &'a mut CudaSlice<u8>,
1290 &'a mut CudaSlice<f32>,
1291 &'a mut CudaSlice<f32>,
1292 &'a mut CudaSlice<f32>,
1293 );
1294 let (mut x_cur, mut x_nxt, sl): (
1295 &mut CudaSlice<f32>,
1296 &mut CudaSlice<f32>,
1297 Option<SlabRefs>,
1298 );
1299 let mut seg: Option<(
1300 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1301 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1302 &mut CudaSlice<f32>,
1303 &mut usize,
1304 )> = None;
1305 let mut x_own2;
1306 match slab_guard.as_mut() {
1307 Some(g) => {
1308 let slabs = &mut **g;
1309 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1310 let PrimeSlabs {
1311 xa,
1312 xb,
1313 h,
1314 x1,
1315 z,
1316 act,
1317 h16,
1318 z16,
1319 gate,
1320 up,
1321 ffn_out,
1322 seg_glue,
1323 mixed,
1324 seg_mid,
1325 seg_t,
1326 ..
1327 } = slabs;
1328 x_cur = xa;
1329 x_nxt = xb;
1330 seg = Some((seg_glue, seg_mid, mixed, seg_t));
1331 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1332 }
1333 None => {
1334 x_own = x_in;
1335 x_own2 = e.uninit(t * n_embd)?;
1336 x_cur = &mut x_own;
1337 x_nxt = &mut x_own2;
1338 sl = None;
1339 }
1340 }
1341 let mut alloc_h;
1342 let mut alloc_x1;
1343 let mut alloc_z;
1344 let mut alloc_act;
1345 let mut alloc_h16;
1346 let mut alloc_z16;
1347 let mut alloc_gate;
1348 let mut alloc_up;
1349 let mut alloc_fo;
1350 let (h, x1, z, act): (
1351 &mut CudaSlice<f32>,
1352 &mut CudaSlice<f32>,
1353 &mut CudaSlice<f32>,
1354 &mut CudaSlice<f32>,
1355 );
1356 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1357 let (sl_gate, sl_up, sl_fo): (
1358 &mut CudaSlice<f32>,
1359 &mut CudaSlice<f32>,
1360 &mut CudaSlice<f32>,
1361 );
1362 match sl {
1363 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1364 h = a;
1365 x1 = b;
1366 z = c;
1367 act = d;
1368 h16 = e16;
1369 z16 = f16b;
1370 sl_gate = g;
1371 sl_up = u;
1372 sl_fo = fo;
1373 }
1374 None => {
1375 alloc_h = e.uninit(t * n_embd)?;
1376 alloc_x1 = e.uninit(t * n_embd)?;
1377 alloc_z = e.uninit(t * n_embd)?;
1378 alloc_act = e.uninit(t * n_ff_max)?;
1379 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1380 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1381 alloc_gate = e.uninit(t * n_ff_max)?;
1382 alloc_up = e.uninit(t * n_ff_max)?;
1383 alloc_fo = e.uninit(t * n_embd)?;
1384 h = &mut alloc_h;
1385 x1 = &mut alloc_x1;
1386 z = &mut alloc_z;
1387 act = &mut alloc_act;
1388 h16 = &mut alloc_h16;
1389 z16 = &mut alloc_z16;
1390 sl_gate = &mut alloc_gate;
1391 sl_up = &mut alloc_up;
1392 sl_fo = &mut alloc_fo;
1393 }
1394 }
1395 let n_layers = self.layers.len();
1400 let use_seg = f16fuse
1410 && seg.is_some()
1411 && self.cfg.step35.is_none()
1412 && lo == 0
1413 && hi == n_layers
1414 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1415 if let Some((sg, sm, _, st)) = seg.as_mut() {
1416 if **st != t {
1417 sg.clear();
1418 sg.extend((0..n_layers).map(|_| None));
1419 sm.clear();
1420 sm.extend((0..n_layers).map(|_| None));
1421 **st = t;
1422 }
1423 }
1424 {
1425 let layer_lo = &self.layers[lo];
1426 if f16fuse {
1427 e.rms_norm_f16out(
1428 x_cur,
1429 layer_lo.attn_norm.float_data(),
1430 h,
1431 h16,
1432 n_embd,
1433 t,
1434 eps,
1435 )?;
1436 } else {
1437 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1438 }
1439 }
1440 for il in lo..hi {
1441 let layer = &self.layers[il];
1442 let hx16 = if f16fuse { Some(&*h16) } else { None };
1443 if use_seg {
1444 let (pre, pre16, w_out) = match &layer.mixer {
1447 Mixer::Full(fa) => {
1448 let g3 = match hx16 {
1449 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1450 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1451 };
1452 let (pre, pre16) =
1453 self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1454 (pre, pre16, &fa.wo)
1455 }
1456 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1457 Mixer::Linear(la) => {
1458 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1459 let g4 = match hx16 {
1460 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1461 None => e.matmul_group(&ws, h, t)?,
1462 };
1463 let (pre, pre16) =
1464 self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1465 (pre, pre16, &la.ssm_out)
1466 }
1467 };
1468 {
1469 let (_, sm, mslab, _) = seg.as_mut().unwrap();
1470 let pre_n = pre.len() / t;
1471 let xh_pre = match pre16 {
1472 Some(x) => x,
1473 None => e.f16_act(&pre, t * pre_n, pre_n)?,
1474 };
1475 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1476 let y = e.matmul(w_out, &pre, t)?;
1477 e.copy_into(mslab, 0, &y, t * n_embd)?;
1478 }
1479 if sm[il].is_none() {
1480 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1481 let w_post = layer.post_attn_norm.float_data();
1482 e.stream().synchronize()?;
1483 e.stream()
1484 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1485 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1486 e.add(x_cur, mslab, x1, t * n_embd)?;
1487 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1488 Ok(())
1489 })();
1490 let g = e.stream().end_capture(
1491 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1492 r?;
1493 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1494 }
1495 sm[il].as_ref().unwrap().launch()?;
1496 }
1497 } else {
1498 let mixed = match &layer.mixer {
1499 Mixer::Full(fa) => {
1500 self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?
1501 }
1502 Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1503 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1504 };
1505 if f16fuse {
1506 e.add_rms_norm_f16out(
1509 x_cur,
1510 &mixed,
1511 layer.post_attn_norm.float_data(),
1512 x1,
1513 z,
1514 z16,
1515 n_embd,
1516 t,
1517 eps,
1518 )?;
1519 } else {
1520 e.add(x_cur, &mixed, x1, t * n_embd)?;
1521 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1522 }
1523 }
1524 let zx16 = if f16fuse { Some(&*z16) } else { None };
1525 match &layer.ffn {
1526 crate::hybrid::Ffn::Dense {
1527 ffn_gate,
1528 ffn_up,
1529 ffn_down,
1530 } => {
1531 let n_ff = ffn_gate.out_features();
1532 let mut into_ok = false;
1535 if let Some(xh) = zx16 {
1536 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1537 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1538 }
1539 if !into_ok {
1540 let mut g2 = match zx16 {
1541 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1542 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1543 };
1544 let up_y = g2.pop().unwrap();
1545 let gate_y = g2.pop().unwrap();
1546 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1547 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1548 }
1549 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1554 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1555 {
1556 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1557 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1558 Some(a16)
1559 } else {
1560 Self::ffn_act_lim(
1561 e,
1562 &self.cfg,
1563 sl_gate,
1564 sl_up,
1565 1.0,
1566 1.0,
1567 d_lim,
1568 act,
1569 t * n_ff,
1570 )?;
1571 None
1572 };
1573 let xh_act = match act16 {
1575 Some(x) => x,
1576 None => e.f16_act(act, t * n_ff, n_ff)?,
1577 };
1578 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1579 let y = e.matmul(ffn_down, &*act, t)?;
1580 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1581 }
1582 }
1583 crate::hybrid::Ffn::Moe(m) => {
1584 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1585 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1586 }
1587 }
1588 if use_seg && il + 1 < hi {
1589 let w_next = self.layers[il + 1].attn_norm.float_data();
1591 let (sg, _, _, _) = seg.as_mut().unwrap();
1592 if sg[il].is_none() {
1593 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1594 e.stream().synchronize()?;
1595 e.stream()
1596 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1597 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1598 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1599 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1600 Ok(())
1601 })();
1602 let g = e.stream().end_capture(
1603 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1604 );
1605 r?;
1606 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1607 }
1608 sg[il].as_ref().unwrap().launch()?;
1609 } else {
1610 if il + 1 < hi {
1611 let w_next = self.layers[il + 1].attn_norm.float_data();
1612 if f16fuse {
1613 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1614 } else {
1615 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1616 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1617 }
1618 } else {
1619 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1620 }
1621 }
1622 if let Some(path) = Self::prime_trace_path() {
1628 let row = (base + t - 1) as usize;
1629 let host = e.dtoh(x_nxt)?;
1630 let last = &host[(t - 1) * n_embd..t * n_embd];
1631 use std::io::Write as _;
1632 let mut f = std::fs::OpenOptions::new()
1633 .create(true)
1634 .append(true)
1635 .open(path)?;
1636 let mut h64: u64 = 0xcbf29ce484222325;
1637 for v in last {
1638 h64 ^= v.to_bits() as u64;
1639 h64 = h64.wrapping_mul(0x100000001b3);
1640 }
1641 writeln!(
1642 f,
1643 "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1644 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1645 last[0], last[1], last[2]
1646 )?;
1647 }
1648 self.dflash_tap(e, cache, il, x_nxt, t)?;
1651 std::mem::swap(&mut x_cur, &mut x_nxt);
1652 }
1653 let mut x = e.uninit(t * n_embd)?;
1655 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1656 drop(slab_guard);
1657 Ok(x)
1658 }
1659
1660 fn prime_chunk_epilogue(
1665 &self,
1666 e: &Engine,
1667 x: CudaSlice<f32>,
1668 t: usize,
1669 cache: &mut Cache,
1670 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1671 let n_embd = self.cfg.n_embd as usize;
1672 let eps = self.cfg.rms_eps;
1673 let mut h_seed = e.uninit(n_embd)?;
1677 if !crate::spec::spec_hpost() {
1678 e.copy_view_into(
1679 &mut h_seed,
1680 0,
1681 &x.slice((t - 1) * n_embd..t * n_embd),
1682 n_embd,
1683 )?;
1684 }
1685 let mut hn = e.uninit(t * n_embd)?;
1687 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1688 if crate::spec::spec_hpost() {
1689 e.copy_view_into(
1690 &mut h_seed,
1691 0,
1692 &hn.slice((t - 1) * n_embd..t * n_embd),
1693 n_embd,
1694 )?;
1695 }
1696 let last = e.view(&hn, t * n_embd);
1697 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1698 let mut hlast = e.uninit(n_embd)?;
1699 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1700 let logits = e.matmul(&self.output, &hlast, 1)?;
1701 cache.pos += t;
1702 Ok((
1705 e.dtoh(&logits)?,
1706 h_seed,
1707 if crate::spec::spec_hpost() { hn } else { x },
1708 ))
1709 }
1710
1711 fn prime_chunk_ppn(
1735 &self,
1736 e: &Engine,
1737 tokens: &[u32],
1738 cache: &mut Cache,
1739 seq_end: usize,
1740 fence: &[usize],
1741 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1742 let rt = crate::pp::PpNRt::get(e)?;
1743 let n_st = fence.len() - 1;
1744 assert_eq!(
1745 rt.n_stages(),
1746 n_st,
1747 "PpNRt stage count {} != fence stages {n_st}",
1748 rt.n_stages()
1749 );
1750 let n_embd = self.cfg.n_embd as usize;
1751 let t = tokens.len();
1752 let base = cache.pos;
1753 debug_assert!(
1754 seq_end >= base + t,
1755 "prime_chunk_ppn: seq_end must cover this chunk"
1756 );
1757 let payload = t * n_embd;
1758 let caller_stream = e.stream();
1762 rt.fence_stages_behind(&caller_stream)?;
1763
1764 if n_st == 2 {
1765 let slot =
1766 self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1767 let x =
1768 self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1769 let out = {
1770 rt.bind_stage(1)?;
1771 let _st1 = rt.enter(1);
1772 let e1 = rt.engine(1, e);
1773 self.prime_chunk_epilogue(e1, x, t, cache)?
1774 };
1775 rt.publish_to(1, &caller_stream)?;
1776 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1777 return Ok(out);
1778 }
1779
1780 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1781
1782 let mut slot = {
1784 let _st0 = rt.enter(0);
1785 let e0 = rt.engine(0, e);
1786 let pos_d = e0.htod_i32(&pos)?;
1787 let x = self.embed(e0, tokens)?;
1788 let x =
1789 self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1790 rt.tx(0, &x, payload)?
1791 };
1793
1794 for s in 1..n_st - 1 {
1796 let _st = rt.enter(s);
1797 let es = rt.engine(s, e);
1798 let pos_d = es.htod_i32(&pos)?;
1799 let x = rt.rx(s - 1, slot, payload)?;
1800 let x = self.prime_layers(
1801 es,
1802 x,
1803 fence[s],
1804 fence[s + 1],
1805 &pos_d,
1806 t,
1807 base,
1808 cache,
1809 seq_end,
1810 )?;
1811 slot = rt.tx(s, &x, payload)?;
1812 }
1813
1814 let _stl = rt.enter(n_st - 1);
1816 let el = rt.engine(n_st - 1, e);
1817 let pos_d = el.htod_i32(&pos)?;
1818 let x = rt.rx(n_st - 2, slot, payload)?;
1819 let x = self.prime_layers(
1820 el,
1821 x,
1822 fence[n_st - 1],
1823 fence[n_st],
1824 &pos_d,
1825 t,
1826 base,
1827 cache,
1828 seq_end,
1829 )?;
1830 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1831 rt.publish_to(n_st - 1, &caller_stream)?;
1837 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1838 Ok(out)
1839 }
1840
1841 fn prime_pp2_stage0_enqueue(
1842 &self,
1843 e: &Engine,
1844 rt: &crate::pp::PpNRt,
1845 tokens: &[u32],
1846 cache: &mut Cache,
1847 seq_end: usize,
1848 fence: &[usize],
1849 base: usize,
1850 pipelined: bool,
1851 ) -> Result<usize, Box<dyn std::error::Error>> {
1852 let t = tokens.len();
1853 let n_embd = self.cfg.n_embd as usize;
1854 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1855 rt.bind_stage(0)?;
1856 let _st0 = rt.enter(0);
1857 let e0 = rt.engine(0, e);
1858 let pos_d = e0.htod_i32(&pos)?;
1859 let x = self.embed(e0, tokens)?;
1860 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1861 let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1862 if pipelined {
1863 rt.tx_pipelined(0, &x, t * n_embd)
1864 } else {
1865 rt.tx(0, &x, t * n_embd)
1866 }
1867 }
1868
1869 fn prime_pp2_stage1_enqueue(
1870 &self,
1871 e: &Engine,
1872 rt: &crate::pp::PpNRt,
1873 slot: usize,
1874 t: usize,
1875 cache: &mut Cache,
1876 seq_end: usize,
1877 fence: &[usize],
1878 base: usize,
1879 pipelined: bool,
1880 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1881 let n_embd = self.cfg.n_embd as usize;
1882 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1883 rt.bind_stage(1)?;
1884 let _st1 = rt.enter(1);
1885 let e1 = rt.engine(1, e);
1886 let pos_d = e1.htod_i32(&pos)?;
1887 let x = rt.rx(0, slot, t * n_embd)?;
1888 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1889 self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
1890 }
1891
1892 pub fn prime_chunk_captured(
1908 &self,
1909 e: &Engine,
1910 x_in: &CudaSlice<f32>,
1911 pos_d: &CudaSlice<i32>,
1912 t: usize,
1913 cache: &mut Cache,
1914 len_d: &CudaSlice<i32>,
1915 logits_out: &mut CudaSlice<f32>,
1916 h_seed_out: &mut CudaSlice<f32>,
1917 ) -> Result<(), Box<dyn std::error::Error>> {
1918 let cfg = &self.cfg;
1919 let n_embd = cfg.n_embd as usize;
1920 let eps = cfg.rms_eps;
1921 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1922 let mut x = e.uninit(t * n_embd)?;
1923 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1924 for (il, layer) in self.layers.iter().enumerate() {
1925 let mut h = e.uninit(t * n_embd)?;
1926 let mut hx16: Option<CudaSlice<u8>> = None;
1927 if f16fuse {
1928 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1929 e.rms_norm_f16out(
1930 &x,
1931 layer.attn_norm.float_data(),
1932 &mut h,
1933 &mut b16,
1934 n_embd,
1935 t,
1936 eps,
1937 )?;
1938 hx16 = Some(b16);
1939 } else {
1940 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1941 }
1942 let mixed = match &layer.mixer {
1943 Mixer::Full(fa) => {
1947 self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
1948 }
1949 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1950 Mixer::Linear(la) => {
1951 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1952 let g4 = match hx16.as_ref() {
1953 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1954 None => e.matmul_group(&ws, &h, t)?,
1955 };
1956 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1957 }
1958 };
1959 let mut x1 = e.uninit(t * n_embd)?;
1960 e.add(&x, &mixed, &mut x1, t * n_embd)?;
1961 let mut z = e.uninit(t * n_embd)?;
1962 let mut zx16: Option<CudaSlice<u8>> = None;
1963 if f16fuse {
1964 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1965 e.rms_norm_f16out(
1966 &x1,
1967 layer.post_attn_norm.float_data(),
1968 &mut z,
1969 &mut b16,
1970 n_embd,
1971 t,
1972 eps,
1973 )?;
1974 zx16 = Some(b16);
1975 } else {
1976 e.rms_norm(
1977 &x1,
1978 layer.post_attn_norm.float_data(),
1979 &mut z,
1980 n_embd,
1981 t,
1982 eps,
1983 )?;
1984 }
1985 let ffn_out = match &layer.ffn {
1986 crate::hybrid::Ffn::Dense {
1987 ffn_gate,
1988 ffn_up,
1989 ffn_down,
1990 } => {
1991 let n_ff = ffn_gate.out_features();
1992 let mut g2 = match &zx16 {
1993 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1994 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1995 };
1996 let up = g2.pop().unwrap();
1997 let gate = g2.pop().unwrap();
1998 let mut act = e.uninit(t * n_ff)?;
1999 Self::ffn_act_lim(
2001 e,
2002 &self.cfg,
2003 &gate,
2004 &up,
2005 1.0,
2006 1.0,
2007 self.cfg.clamp_shexp_at(il as u32),
2008 &mut act,
2009 t * n_ff,
2010 )?;
2011 e.matmul(ffn_down, &act, t)?
2012 }
2013 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
2014 };
2015 let mut x2 = e.uninit(t * n_embd)?;
2016 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
2017 x = x2;
2018 }
2019 if !crate::spec::spec_hpost() {
2021 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
2022 }
2023 let mut hn = e.uninit(t * n_embd)?;
2024 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
2025 if crate::spec::spec_hpost() {
2026 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
2027 }
2028 let mut hlast = e.uninit(n_embd)?;
2029 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
2030 let logits = e.matmul(&self.output, &hlast, 1)?;
2031 let nv = logits.len();
2032 e.copy_into(logits_out, 0, &logits, nv)?;
2033 Ok(())
2034 }
2035
2036 fn step35_prime_batch_on() -> bool {
2037 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
2038 }
2039
2040 #[allow(clippy::too_many_arguments)]
2043 fn step35_prime_batch_layers(
2044 &self,
2045 e: &Engine,
2046 mut x: CudaSlice<f32>,
2047 lo: usize,
2048 hi: usize,
2049 ts: &[usize],
2050 offs: &[usize],
2051 pos_ds: &[CudaSlice<i32>],
2052 caches: &mut [&mut Cache],
2053 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2054 let cfg = &self.cfg;
2055 let n_embd = cfg.n_embd as usize;
2056 let eps = cfg.rms_eps;
2057 let b = ts.len();
2058 let total: usize = ts.iter().sum();
2059 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
2060
2061 let split = |e: &Engine,
2062 y: &CudaSlice<f32>,
2063 dim: usize|
2064 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2065 let mut out = Vec::with_capacity(b);
2066 for s in 0..b {
2067 let mut ys = e.uninit(ts[s] * dim)?;
2068 e.copy_view_into(
2069 &mut ys,
2070 0,
2071 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2072 ts[s] * dim,
2073 )?;
2074 out.push(ys);
2075 }
2076 Ok(out)
2077 };
2078
2079 for il in lo..hi {
2080 let layer = &self.layers[il];
2081 let Mixer::Full(fa) = &layer.mixer else {
2082 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2083 };
2084
2085 let mut h = e.uninit(total * n_embd)?;
2086 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2087 if f16fuse {
2088 e.rms_norm_f16out(
2089 &x,
2090 layer.attn_norm.float_data(),
2091 &mut h,
2092 &mut hx16,
2093 n_embd,
2094 total,
2095 eps,
2096 )?;
2097 } else {
2098 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2099 }
2100
2101 let gate_w = fa
2105 .attn_gate
2106 .as_ref()
2107 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2108 let mut g4 = if f16fuse {
2109 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2110 } else {
2111 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2112 };
2113 let gate = g4.pop().unwrap();
2114 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2115 (0..b).map(|_| Vec::with_capacity(3)).collect();
2116 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2117 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2118 parts[s].push(ys);
2119 }
2120 }
2121 let gates = split(e, &gate, gate_w.out_features())?;
2122 let geometry = self.step35_geom(il);
2123 let hd = geometry.head_dim_k as usize;
2124 let nh = geometry.n_head as usize;
2125 let mut ag_cat = e.uninit(total * nh * hd)?;
2126 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2127 let ag = self.step35_attn_pre_wo(
2128 e,
2129 fa,
2130 g3s,
2131 None,
2132 Some(&gate),
2133 &pos_ds[s],
2134 ts[s],
2135 Some(&mut *caches[s]),
2136 il,
2137 ts[s],
2138 )?;
2139 e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2140 }
2141 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2142
2143 let mut x1 = e.uninit(total * n_embd)?;
2144 let mut z = e.uninit(total * n_embd)?;
2145 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2146 if f16fuse {
2147 e.add_rms_norm_f16out(
2148 &x,
2149 &mixed,
2150 layer.post_attn_norm.float_data(),
2151 &mut x1,
2152 &mut z,
2153 &mut zx16,
2154 n_embd,
2155 total,
2156 eps,
2157 )?;
2158 } else {
2159 e.add(&x, &mixed, &mut x1, total * n_embd)?;
2160 e.rms_norm(
2161 &x1,
2162 layer.post_attn_norm.float_data(),
2163 &mut z,
2164 n_embd,
2165 total,
2166 eps,
2167 )?;
2168 }
2169
2170 let ffn_out = match &layer.ffn {
2171 crate::hybrid::Ffn::Dense {
2172 ffn_gate,
2173 ffn_up,
2174 ffn_down,
2175 } => {
2176 let n_ff = ffn_gate.out_features();
2177 let mut g2 = if f16fuse {
2178 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2179 } else {
2180 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2181 };
2182 let up = g2.pop().unwrap();
2183 let gate = g2.pop().unwrap();
2184 let mut act = e.uninit(total * n_ff)?;
2185 let d_lim = cfg.clamp_shexp_at(il as u32);
2186 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2187 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2188 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2189 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2190 Some(y) => y,
2191 None => e.matmul(ffn_down, &act, total)?,
2192 }
2193 } else {
2194 Self::ffn_act_lim(
2195 e,
2196 cfg,
2197 &gate,
2198 &up,
2199 1.0,
2200 1.0,
2201 d_lim,
2202 &mut act,
2203 total * n_ff,
2204 )?;
2205 e.matmul(ffn_down, &act, total)?
2206 }
2207 }
2208 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2209 };
2210 let mut x2 = e.uninit(total * n_embd)?;
2211 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2212 x = x2;
2213 }
2214 Ok(x)
2215 }
2216
2217 fn step35_prime_batch_epilogue(
2218 &self,
2219 e: &Engine,
2220 x: CudaSlice<f32>,
2221 ts: &[usize],
2222 offs: &[usize],
2223 caches: &mut [&mut Cache],
2224 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2225 let n_embd = self.cfg.n_embd as usize;
2226 let total: usize = ts.iter().sum();
2227 let mut hn = e.uninit(total * n_embd)?;
2228 e.rms_norm(
2229 &x,
2230 self.output_norm.float_data(),
2231 &mut hn,
2232 n_embd,
2233 total,
2234 self.cfg.rms_eps,
2235 )?;
2236
2237 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2238 let mut out = Vec::with_capacity(ts.len());
2239 for s in 0..ts.len() {
2240 let mut hidden = e.uninit(ts[s] * n_embd)?;
2241 e.copy_view_into(
2242 &mut hidden,
2243 0,
2244 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2245 ts[s] * n_embd,
2246 )?;
2247 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2248 let mut h_seed = e.uninit(n_embd)?;
2249 e.copy_view_into(
2250 &mut h_seed,
2251 0,
2252 &hidden_src.slice(last0..last0 + n_embd),
2253 n_embd,
2254 )?;
2255 let mut hlast = e.uninit(n_embd)?;
2257 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2258 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2259 caches[s].pos += ts[s];
2260 out.push((logits, h_seed, hidden));
2261 }
2262 Ok(out)
2263 }
2264
2265 fn step35_prime_cache_batch(
2266 &self,
2267 e: &Engine,
2268 prompts: &[&[u32]],
2269 caches: &mut [&mut Cache],
2270 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2271 if crate::pp::pp_host_bounce_active()
2272 && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2273 {
2274 return Err(
2275 "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2276 stage split; refusing an unsplit remote-weight walk"
2277 .into(),
2278 );
2279 }
2280 if !Self::step35_prime_batch_on() {
2281 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2282 }
2283 if caches.iter().any(|c| c.pos != 0) {
2284 return Err(
2285 "step35 batched prime currently supports complete fresh prompts only; \
2286 continuation/tick chunks require per-request queued_after"
2287 .into(),
2288 );
2289 }
2290
2291 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2292 for &t in &ts {
2293 assert!(
2294 t >= PRIME_MIN_T,
2295 "step35 batched prime needs T >= {PRIME_MIN_T}"
2296 );
2297 }
2298 for (s, c) in caches.iter().enumerate() {
2299 assert!(
2300 ts[s] <= c.max_ctx,
2301 "step35 batched prime exceeds cache max_ctx"
2302 );
2303 }
2304 let offs: Vec<usize> = ts
2305 .iter()
2306 .scan(0usize, |a, &t| {
2307 let o = *a;
2308 *a += t;
2309 Some(o)
2310 })
2311 .collect();
2312 let total: usize = ts.iter().sum();
2313 let payload = total * self.cfg.n_embd as usize;
2314 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2315 let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2316 let upload_positions =
2317 |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2318 positions
2319 .iter()
2320 .map(|p| e.htod_i32(p))
2321 .collect::<Result<_, _>>()
2322 };
2323
2324 static ONCE: std::sync::Once = std::sync::Once::new();
2325 ONCE.call_once(|| {
2326 eprintln!(
2327 "[step35-prime-batch] first concat prime: B={} tokens={total}",
2328 prompts.len()
2329 );
2330 });
2331
2332 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2333 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2334 let rt = crate::pp::PpNRt::get(e)?;
2335 let n_st = fence.len() - 1;
2336 assert_eq!(
2337 rt.n_stages(),
2338 n_st,
2339 "step35 prime batch stage count mismatch"
2340 );
2341 let caller_stream = e.stream();
2342 rt.fence_stages_behind(&caller_stream)?;
2343
2344 let mut slot = {
2345 let _st0 = rt.enter(0);
2346 let e0 = rt.engine(0, e);
2347 let pos_ds = upload_positions(e0)?;
2348 let x = self.embed(e0, &cat_tokens)?;
2349 let x = self.step35_prime_batch_layers(
2350 e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2351 )?;
2352 rt.tx(0, &x, payload)?
2353 };
2354 for s in 1..n_st - 1 {
2355 let _st = rt.enter(s);
2356 let es = rt.engine(s, e);
2357 let pos_ds = upload_positions(es)?;
2358 let x = rt.rx(s - 1, slot, payload)?;
2359 let x = self.step35_prime_batch_layers(
2360 es,
2361 x,
2362 fence[s],
2363 fence[s + 1],
2364 &ts,
2365 &offs,
2366 &pos_ds,
2367 caches,
2368 )?;
2369 slot = rt.tx(s, &x, payload)?;
2370 }
2371
2372 let _stl = rt.enter(n_st - 1);
2373 let el = rt.engine(n_st - 1, e);
2374 let pos_ds = upload_positions(el)?;
2375 let x = rt.rx(n_st - 2, slot, payload)?;
2376 let x = self.step35_prime_batch_layers(
2377 el,
2378 x,
2379 fence[n_st - 1],
2380 fence[n_st],
2381 &ts,
2382 &offs,
2383 &pos_ds,
2384 caches,
2385 )?;
2386 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2387 rt.publish_to(n_st - 1, &caller_stream)?;
2388 crate::pp::STEP35_PRIME_BATCH_SPLITS
2389 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2390 out
2391 } else {
2392 let pos_ds = upload_positions(e)?;
2393 let x = self.embed(e, &cat_tokens)?;
2394 let x = self.step35_prime_batch_layers(
2395 e,
2396 x,
2397 0,
2398 self.layers.len(),
2399 &ts,
2400 &offs,
2401 &pos_ds,
2402 caches,
2403 )?;
2404 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2405 }
2406 } else {
2407 let pos_ds = upload_positions(e)?;
2408 let x = self.embed(e, &cat_tokens)?;
2409 let x = self.step35_prime_batch_layers(
2410 e,
2411 x,
2412 0,
2413 self.layers.len(),
2414 &ts,
2415 &offs,
2416 &pos_ds,
2417 caches,
2418 )?;
2419 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2420 };
2421 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2422 Ok(out)
2423 }
2424
2425 pub fn prime_cache_batch(
2442 &self,
2443 e: &Engine,
2444 prompts: &[&[u32]],
2445 caches: &mut [&mut Cache],
2446 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2447 let cfg = &self.cfg;
2448 let n_embd = cfg.n_embd as usize;
2449 let eps = cfg.rms_eps;
2450 let b = prompts.len();
2451 assert!(b >= 1 && b == caches.len());
2452 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2453 let carried = pos0s.iter().any(|&p| p > 0);
2454 if cfg.gemma4.is_some() {
2460 return Err(
2461 "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2462 swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2463 .into(),
2464 );
2465 }
2466 if cfg.step35.is_some() {
2469 return self.step35_prime_cache_batch(e, prompts, caches);
2470 }
2471 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2472 for &t in &ts {
2473 assert!(
2474 t >= PRIME_MIN_T,
2475 "prime_cache_batch needs T >= {PRIME_MIN_T}"
2476 );
2477 }
2478 for (s, c) in caches.iter().enumerate() {
2479 assert!(
2480 c.pos + ts[s] <= c.max_ctx,
2481 "prime_cache_batch: prompt exceeds cache max_ctx"
2482 );
2483 }
2484 let total: usize = ts.iter().sum();
2485 let offs: Vec<usize> = ts
2486 .iter()
2487 .scan(0usize, |a, &t| {
2488 let o = *a;
2489 *a += t;
2490 Some(o)
2491 })
2492 .collect();
2493 let pos_ds: Vec<CudaSlice<i32>> = ts
2495 .iter()
2496 .zip(&pos0s)
2497 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2498 .collect::<Result<_, _>>()?;
2499 let split = |e: &Engine,
2501 y: &CudaSlice<f32>,
2502 dim: usize|
2503 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2504 let mut out = Vec::with_capacity(b);
2505 for s in 0..b {
2506 let mut ys = e.uninit(ts[s] * dim)?;
2507 e.copy_view_into(
2508 &mut ys,
2509 0,
2510 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2511 ts[s] * dim,
2512 )?;
2513 out.push(ys);
2514 }
2515 Ok(out)
2516 };
2517
2518 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2519 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
2521 let mut h = e.uninit(total * n_embd)?;
2522 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2523 e.rms_norm_f16out(
2524 &x,
2525 layer.attn_norm.float_data(),
2526 &mut h,
2527 &mut hx16,
2528 n_embd,
2529 total,
2530 eps,
2531 )?;
2532 let mut mixed = e.uninit(total * n_embd)?;
2534 match &layer.mixer {
2535 Mixer::Full(fa) => {
2536 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2537 let geometry = self.cfg.full_attention_geometry_at(il as u32);
2543 let (n_head, n_head_kv, head_dim) = (
2544 geometry.n_head as usize,
2545 geometry.n_head_kv as usize,
2546 geometry.head_dim_k as usize,
2547 );
2548 let fa_scale = geometry.attention_scale();
2549 let use_favl = !carried
2550 && (2..=8).contains(&b)
2551 && (head_dim == 256 || head_dim == 128)
2552 && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2553 && std::env::var("MEMRA_NOFA").is_err()
2554 && std::env::var("MEMRA_FA_FLOOR").is_err()
2555 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2556 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2557 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2558 if use_favl {
2559 let (qf_w, kf_w, vf_w) = (
2560 fa.wq.out_features(),
2561 fa.wk.out_features(),
2562 fa.wv.out_features(),
2563 );
2564 memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
2569 struct APre {
2570 q: CudaSlice<f32>,
2571 gate: Option<CudaSlice<f32>>,
2572 qn: CudaSlice<f32>,
2573 kn: CudaSlice<f32>,
2574 }
2575 let mut aps = Vec::with_capacity(b);
2576 for &t in ts.iter().take(b) {
2577 aps.push(APre {
2578 q: e.uninit(t * n_head * head_dim)?,
2579 gate: Some(e.uninit(t * n_head * head_dim)?),
2580 qn: e.uninit(t * n_head * head_dim)?,
2581 kn: e.uninit(t * n_head_kv * head_dim)?,
2582 });
2583 }
2584 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2585 let kvl = caches[0].kv[il].as_ref().unwrap();
2586 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2587 };
2588 let pargs: Vec<crate::AttnPreVl> = (0..b)
2589 .map(|s| {
2590 let (o, t) = (offs[s], ts[s]);
2591 let kvl = caches[s].kv[il].as_ref().unwrap();
2592 assert!(
2593 kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2594 "prime_cache_batch attn vl: fresh + capacity"
2595 );
2596 crate::AttnPreVl {
2597 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2598 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2599 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2600 q: e.addr_f32(&aps[s].q),
2601 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2602 qn: e.addr_f32(&aps[s].qn),
2603 kn: e.addr_f32(&aps[s].kn),
2604 kc: e.addr_u8(&kvl.k),
2605 vc: e.addr_u8(&kvl.v),
2606 t: t as i32,
2607 pad: 0,
2608 }
2609 })
2610 .collect();
2611 e.attn_pre_vl8(
2612 &pargs,
2613 fa.q_norm.float_data(),
2614 fa.k_norm.float_data(),
2615 head_dim,
2616 geometry.n_rot as usize,
2617 n_head,
2618 n_head_kv,
2619 self.cfg.rms_eps,
2620 geometry.rope_base,
2621 1.0,
2622 kv_dim_k,
2623 kv_dim_v,
2624 ktb,
2625 vtb,
2626 )?;
2627 for s in 0..b {
2628 let kvl = caches[s].kv[il].as_mut().unwrap();
2629 kvl.len += ts[s];
2630 let new_len = kvl.len as i32;
2631 e.set_i32_one(&mut kvl.len_d, new_len)?;
2632 }
2633 let mut attns = Vec::with_capacity(b);
2634 let mut mirrors = Vec::with_capacity(b);
2635 for &t in ts.iter().take(b) {
2636 attns.push(e.uninit(t * n_head * head_dim)?);
2637 let n = t * n_head_kv * head_dim;
2638 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2639 }
2640 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2643 Ok("0") => false,
2644 Ok("1") => true,
2645 _ => cfg!(memra_hopper_mma),
2646 };
2647 if fa3_on {
2648 let mut q16s = Vec::with_capacity(b);
2649 let mut v16s = Vec::with_capacity(b);
2650 for s in 0..b {
2651 let t = ts[s];
2652 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2653 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2654 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2655 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2656 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2657 e.f32_to_bf16_v(
2658 &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2659 &mut v16,
2660 t * n_head_kv * head_dim,
2661 )?;
2662 q16s.push(q16);
2663 v16s.push((k16, v16));
2664 }
2665 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2666 let mut kp = qp;
2667 let mut vp = qp;
2668 let mut op = [core::ptr::null_mut::<f32>(); 8];
2669 let mut tsv = [0i32; 8];
2670 for s in 0..b {
2671 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2672 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2673 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2674 op[s] = e.addr_f32(&attns[s]) as *mut f32;
2675 tsv[s] = ts[s] as i32;
2676 }
2677 let rc = unsafe {
2678 crate::fa3_vl_raw(
2679 qp.as_ptr(),
2680 kp.as_ptr(),
2681 vp.as_ptr(),
2682 op.as_ptr(),
2683 tsv.as_ptr(),
2684 b as i32,
2685 n_head as i32,
2686 n_head_kv as i32,
2687 head_dim as i32,
2688 fa_scale,
2689 e.stream().cu_stream() as *mut core::ffi::c_void,
2690 )
2691 };
2692 if rc != 0 {
2693 return Err(format!("memra_fa3_vl rc={rc}").into());
2694 }
2695 } else {
2696 let fargs: Vec<crate::FaSeqVl> = (0..b)
2697 .map(|s| crate::FaSeqVl {
2698 q: e.addr_f32(&aps[s].qn),
2699 k16: e.addr_u8(&mirrors[s].0),
2700 v16: e.addr_u8(&mirrors[s].1),
2701 o: e.addr_f32(&attns[s]),
2702 kf: e.addr_f32(&aps[s].kn),
2703 vf: e.addr_f32v(
2704 &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2705 ),
2706 t: ts[s] as i32,
2707 pad: 0,
2708 })
2709 .collect();
2710 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2711 }
2712 for (s, attn) in attns.into_iter().enumerate() {
2713 let (attn_g, ag16) = self.full_attn_prime_post_fa(
2714 e,
2715 attn,
2716 &aps[s].gate,
2717 ts[s],
2718 n_head,
2719 head_dim,
2720 )?;
2721 let mut done = false;
2722 if let Some(xh) = &ag16 {
2723 done = e.try_f16_gemm_pre_into_off(
2724 &fa.wo,
2725 xh,
2726 ts[s],
2727 &mut mixed,
2728 offs[s] * n_embd,
2729 )?;
2730 }
2731 if !done {
2732 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2733 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2734 }
2735 }
2736 } else {
2737 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2738 (0..b).map(|_| Vec::new()).collect();
2739 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2740 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2741 parts[s].push(ys);
2742 }
2743 }
2744 for (s, g3s) in parts.into_iter().enumerate() {
2745 let (attn_g, ag16) = self.full_attn_prime_core_inner(
2747 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2748 )?;
2749 let mut done = false;
2750 if let Some(xh) = &ag16 {
2751 done = e.try_f16_gemm_pre_into_off(
2752 &fa.wo,
2753 xh,
2754 ts[s],
2755 &mut mixed,
2756 offs[s] * n_embd,
2757 )?;
2758 }
2759 if !done {
2760 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2761 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2762 }
2763 }
2764 }
2765 }
2766 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2767 Mixer::Linear(la) => {
2768 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2773 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2774 let outs =
2775 self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2776 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2777 let (o, t) = (offs[s], ts[s]);
2778 let mut done = false;
2779 if let Some(xh) = &gn16 {
2780 done = e.try_f16_gemm_pre_into_off(
2781 &la.ssm_out,
2782 xh,
2783 t,
2784 &mut mixed,
2785 o * n_embd,
2786 )?;
2787 }
2788 if !done {
2789 let m = e.matmul(&la.ssm_out, &gn, t)?;
2790 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2791 }
2792 }
2793 }
2794 }
2795 let mut x1 = e.uninit(total * n_embd)?;
2796 let mut z = e.uninit(total * n_embd)?;
2797 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2798 e.add_rms_norm_f16out(
2799 &x,
2800 &mixed,
2801 layer.post_attn_norm.float_data(),
2802 &mut x1,
2803 &mut z,
2804 &mut zx16,
2805 n_embd,
2806 total,
2807 eps,
2808 )?;
2809 let ffn_out = match &layer.ffn {
2810 crate::hybrid::Ffn::Dense {
2811 ffn_gate,
2812 ffn_up,
2813 ffn_down,
2814 } => {
2815 let n_ff = ffn_gate.out_features();
2816 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2817 let up = g2.pop().unwrap();
2818 let gate = g2.pop().unwrap();
2819 let mut act = e.uninit(total * n_ff)?;
2820 let d_lim = self.cfg.clamp_shexp_at(il as u32);
2824 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2825 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2826 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2827 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2828 Some(y) => y,
2829 None => e.matmul(ffn_down, &act, total)?,
2830 }
2831 } else {
2832 Self::ffn_act_lim(
2833 e,
2834 &self.cfg,
2835 &gate,
2836 &up,
2837 1.0,
2838 1.0,
2839 d_lim,
2840 &mut act,
2841 total * n_ff,
2842 )?;
2843 e.matmul(ffn_down, &act, total)?
2844 }
2845 }
2846 crate::hybrid::Ffn::Moe(m) => {
2847 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2848 }
2849 };
2850 let mut x2 = e.uninit(total * n_embd)?;
2851 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2852 x = x2;
2853 }
2854 let mut hn = e.uninit(total * n_embd)?;
2856 e.rms_norm(
2857 &x,
2858 self.output_norm.float_data(),
2859 &mut hn,
2860 n_embd,
2861 total,
2862 eps,
2863 )?;
2864 let mut hcat = e.uninit(b * n_embd)?;
2870 for s in 0..b {
2871 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2872 e.copy_view_into(
2873 &mut hcat,
2874 s * n_embd,
2875 &hn.slice(last0..last0 + n_embd),
2876 n_embd,
2877 )?;
2878 }
2879 let logits_cat = if b >= 2 {
2880 e.try_f16_gemm(&self.output, &hcat, b)?
2881 } else {
2882 None
2883 };
2884 let logits_host: Option<Vec<f32>> = match &logits_cat {
2885 Some(lc) => Some(e.dtoh(lc)?),
2886 None => None,
2887 };
2888 let n_vocab = self.output.out_features();
2889 let mut hidden_all = if crate::spec::spec_hpost() {
2890 split(e, &hn, n_embd)?
2891 } else {
2892 split(e, &x, n_embd)?
2893 };
2894 let mut out = Vec::with_capacity(b);
2895 for s in 0..b {
2896 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2897 let mut h_seed = e.uninit(n_embd)?;
2898 if !crate::spec::spec_hpost() {
2899 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2900 } else {
2901 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2902 }
2903 let logits = match &logits_host {
2904 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2905 None => {
2906 let mut hlast = e.uninit(n_embd)?;
2907 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2908 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2909 }
2910 };
2911 caches[s].pos += ts[s];
2912 out.push((logits, h_seed, hidden_all.remove(0)));
2913 }
2914 Ok(out)
2915 }
2916
2917 #[allow(clippy::too_many_arguments)]
2928 fn full_attn_prime(
2929 &self,
2930 e: &Engine,
2931 fa: &FullAttnLayer,
2932 h: &CudaSlice<f32>,
2933 hx: Option<&CudaSlice<u8>>,
2934 pos_d: &CudaSlice<i32>,
2935 t: usize,
2936 cache: &mut Cache,
2937 il: usize,
2938 seq_end: usize,
2939 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2940 if self.cfg.step35.is_some() {
2941 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2942 }
2943 let g3 = match hx {
2948 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2949 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2950 };
2951 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2952 }
2953
2954 fn full_attn_prime_core(
2958 &self,
2959 e: &Engine,
2960 fa: &FullAttnLayer,
2961 g3: Vec<CudaSlice<f32>>,
2962 pos_d: &CudaSlice<i32>,
2963 t: usize,
2964 cache: &mut Cache,
2965 il: usize,
2966 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2967 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2968 if let Some(xh) = &ag16 {
2969 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2970 return Ok(y);
2971 }
2972 }
2973 Ok(e.matmul(&fa.wo, &attn_g, t)?)
2974 }
2975
2976 fn full_attn_prime_core_inner(
2977 &self,
2978 e: &Engine,
2979 fa: &FullAttnLayer,
2980 g3: Vec<CudaSlice<f32>>,
2981 pos_d: &CudaSlice<i32>,
2982 t: usize,
2983 cache: &mut Cache,
2984 il: usize,
2985 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2986 let cfg = &self.cfg;
2987 let geometry = cfg.full_attention_geometry_at(il as u32);
2988 let n_head = geometry.n_head as usize;
2989 let n_head_kv = geometry.n_head_kv as usize;
2990 let head_dim = geometry.head_dim_k as usize;
2991 let scale = geometry.attention_scale();
2992 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2993 let AttnPre { q, k, v, gate } = pre;
2994 let mut attn = e.uninit(t * n_head * head_dim)?;
2995 self.full_attn_prime_fa_dispatch(
2996 e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
2997 )?;
2998 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2999 }
3000
3001 #[allow(clippy::type_complexity)]
3005 fn full_attn_prime_pre_fa(
3006 &self,
3007 e: &Engine,
3008 fa: &FullAttnLayer,
3009 mut g3: Vec<CudaSlice<f32>>,
3010 pos_d: &CudaSlice<i32>,
3011 t: usize,
3012 cache: &mut Cache,
3013 il: usize,
3014 ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
3015 let cfg = &self.cfg;
3016 let geometry = cfg.full_attention_geometry_at(il as u32);
3017 let n_head = geometry.n_head as usize;
3018 let n_head_kv = geometry.n_head_kv as usize;
3019 let head_dim = geometry.head_dim_k as usize;
3020 let eps = cfg.rms_eps;
3021
3022 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3026 let v = g3.pop().unwrap();
3027 let mut k = g3.pop().unwrap();
3028 let qf = g3.pop().unwrap();
3029 let (mut q, gate) = if gated {
3030 let mut q = e.uninit(t * n_head * head_dim)?;
3031 let mut gate = e.uninit(t * n_head * head_dim)?;
3032 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3033 (q, Some(gate))
3034 } else {
3035 (qf, None)
3036 };
3037
3038 let mut qn = e.uninit(t * n_head * head_dim)?;
3039 e.rms_norm(
3040 &q,
3041 fa.q_norm.float_data(),
3042 &mut qn,
3043 head_dim,
3044 n_head * t,
3045 eps,
3046 )?;
3047 q = qn;
3048 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3049 e.rms_norm(
3050 &k,
3051 fa.k_norm.float_data(),
3052 &mut kn,
3053 head_dim,
3054 n_head_kv * t,
3055 eps,
3056 )?;
3057 k = kn;
3058 let rope_dims = geometry.n_rot as usize;
3059 e.rope_neox(
3060 &mut q,
3061 pos_d,
3062 head_dim,
3063 rope_dims,
3064 n_head,
3065 t,
3066 geometry.rope_base,
3067 1.0,
3068 )?;
3069 e.rope_neox(
3070 &mut k,
3071 pos_d,
3072 head_dim,
3073 rope_dims,
3074 n_head_kv,
3075 t,
3076 geometry.rope_base,
3077 1.0,
3078 )?;
3079
3080 {
3083 let kvl = cache.kv[il].as_mut().unwrap();
3084 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3085 e.append_kv_quantized_rows(
3086 &k,
3087 &v,
3088 &mut kvl.k,
3089 &mut kvl.v,
3090 kvl.len,
3091 t,
3092 kvl.kv_dim_k,
3093 kvl.kv_dim_v,
3094 kvl.k_tok_bytes,
3095 kvl.v_tok_bytes,
3096 crate::Engine::kv_fp8_on(),
3097 )?;
3098 kvl.len += t;
3099 let new_len = kvl.len as i32;
3100 e.set_i32_one(&mut kvl.len_d, new_len)?;
3101 }
3102
3103 let base_len = {
3104 let kvl = cache.kv[il].as_ref().unwrap();
3105 kvl.len - t };
3107 Ok((AttnPre { q, k, v, gate }, base_len))
3108 }
3109
3110 #[allow(clippy::too_many_arguments)]
3117 fn full_attn_prime_fa_dispatch(
3118 &self,
3119 e: &Engine,
3120 q: &CudaSlice<f32>,
3121 k: &CudaSlice<f32>,
3122 v: &CudaSlice<f32>,
3123 attn: &mut CudaSlice<f32>,
3124 base_len: usize,
3125 t: usize,
3126 cache: &mut Cache,
3127 il: usize,
3128 head_dim: usize,
3129 n_head: usize,
3130 n_head_kv: usize,
3131 scale: f32,
3132 ) -> Result<(), Box<dyn std::error::Error>> {
3133 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3146 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3147 e.sdpa_naive(
3148 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3149 )?;
3150 } else {
3151 e.fa_prefill(
3152 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3153 )?;
3154 }
3155 return Ok(());
3156 }
3157 let kvl = cache.kv[il].as_ref().unwrap();
3158 let t_kv = base_len + t;
3159 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3160 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3161 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3165 e.sdpa_naive_quantized_view(
3166 q,
3167 &k_view,
3168 &v_view,
3169 attn,
3170 head_dim,
3171 n_head,
3172 n_head_kv,
3173 t,
3174 t_kv,
3175 scale,
3176 true,
3177 kvl.k_tok_bytes,
3178 kvl.v_tok_bytes,
3179 )?;
3180 return Ok(());
3181 }
3182 let deqw = std::env::var("MEMRA_PRIME_DEQW")
3190 .map(|v| v != "0")
3191 .unwrap_or(true);
3192 if deqw {
3193 e.fa_prefill_view_ws(
3194 q,
3195 &k_view,
3196 &v_view,
3197 attn,
3198 head_dim,
3199 n_head,
3200 n_head_kv,
3201 t,
3202 t_kv,
3203 scale,
3204 true,
3205 kvl.k_tok_bytes,
3206 kvl.v_tok_bytes,
3207 crate::Engine::kv_fp8_on(),
3208 )?;
3209 } else {
3210 e.fa_prefill_view(
3211 q,
3212 &k_view,
3213 &v_view,
3214 attn,
3215 head_dim,
3216 n_head,
3217 n_head_kv,
3218 t,
3219 t_kv,
3220 scale,
3221 true,
3222 kvl.k_tok_bytes,
3223 kvl.v_tok_bytes,
3224 crate::Engine::kv_fp8_on(),
3225 )?;
3226 }
3227 Ok(())
3228 }
3229
3230 fn full_attn_prime_post_fa(
3233 &self,
3234 e: &Engine,
3235 attn: CudaSlice<f32>,
3236 gate: &Option<CudaSlice<f32>>,
3237 t: usize,
3238 n_head: usize,
3239 head_dim: usize,
3240 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3241 let (attn_g, ag16) = match gate {
3242 Some(gate) => {
3243 let n = t * n_head * head_dim;
3244 let mut ag = e.uninit(n)?;
3245 if Self::f16out_on(e, t) {
3246 let mut a16 = e.alloc_u8_uninit(n * 2)?;
3247 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3248 (ag, Some(a16))
3249 } else {
3250 let mut gsig = e.uninit(n)?;
3251 e.sigmoid(gate, &mut gsig, n)?;
3252 e.mul(&attn, &gsig, &mut ag, n)?;
3253 (ag, None)
3254 }
3255 }
3256 None => (attn, None),
3257 };
3258 Ok((attn_g, ag16))
3259 }
3260
3261 fn linear_attn_prime(
3268 &self,
3269 e: &Engine,
3270 la: &LinearAttnLayer,
3271 h: &CudaSlice<f32>,
3272 hx: Option<&CudaSlice<u8>>,
3273 t: usize,
3274 cache: &mut Cache,
3275 il: usize,
3276 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3277 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3279 let g4 = match hx {
3280 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3281 None => e.matmul_group(&ws, h, t)?,
3282 };
3283 self.linear_attn_prime_core(e, la, g4, t, cache, il)
3284 }
3285
3286 fn linear_attn_prime_core(
3288 &self,
3289 e: &Engine,
3290 la: &LinearAttnLayer,
3291 mut g4: Vec<CudaSlice<f32>>,
3292 t: usize,
3293 cache: &mut Cache,
3294 il: usize,
3295 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3296 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3297 }
3298
3299 #[allow(clippy::too_many_arguments)]
3303 fn linear_attn_prime_core_pad_inner(
3304 &self,
3305 e: &Engine,
3306 la: &LinearAttnLayer,
3307 mut g4: Vec<CudaSlice<f32>>,
3308 t: usize,
3309 cache: &mut Cache,
3310 il: usize,
3311 pad_len: Option<&CudaSlice<i32>>,
3312 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3313 let ssm = self.cfg.ssm.as_ref().unwrap();
3315 let d_state = ssm.state_size as usize;
3316 let num_k = ssm.group_count as usize;
3317 let num_v = ssm.time_step_rank as usize;
3318 let key_dim = d_state * num_k;
3319 let value_dim = d_state * num_v;
3320 let conv_dim = key_dim * 2 + value_dim;
3321 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(
3326 e,
3327 la,
3328 &qkv_mixed.slice(0..t * conv_dim),
3329 &z.slice(0..t * value_dim),
3330 &beta_raw.slice(0..t * num_v),
3331 &alpha.slice(0..t * num_v),
3332 t,
3333 cache,
3334 il,
3335 pad_len,
3336 )
3337 }
3338
3339 #[allow(clippy::too_many_arguments)]
3342 fn linear_attn_gdn_prep(
3343 &self,
3344 e: &Engine,
3345 la: &LinearAttnLayer,
3346 qkv_mixed: &cudarc::driver::CudaView<f32>,
3347 beta_raw: &cudarc::driver::CudaView<f32>,
3348 alpha: &cudarc::driver::CudaView<f32>,
3349 t: usize,
3350 cache: &mut Cache,
3351 il: usize,
3352 pad_len: Option<&CudaSlice<i32>>,
3353 ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3354 let cfg = &self.cfg;
3355 let ssm = cfg.ssm.as_ref().unwrap();
3356 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;
3364 debug_assert!(
3365 t >= d_conv - 1,
3366 "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3367 );
3368
3369 let rl = cache.recur[il].as_mut().unwrap();
3374 let hk = Self::gdn_hk(e, t, num_v, num_k);
3375 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3376 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
3378 let mut k_g = e.uninit(d_state * hk * t)?;
3379 let mut v_g = e.uninit(d_state * num_v * t)?;
3380 if conv_fuse {
3381 e.ssm_conv1d_gdn_state_pad(
3382 qkv_mixed,
3383 &mut rl.conv_state,
3384 la.ssm_conv1d.float_data(),
3385 &mut q_g,
3386 &mut k_g,
3387 &mut v_g,
3388 conv_dim,
3389 t,
3390 d_conv,
3391 d_state,
3392 num_v,
3393 num_k,
3394 key_dim,
3395 hk,
3396 pad_len,
3397 )?;
3398 } else {
3399 let mut conv_out = e.uninit(conv_dim * t)?; e.ssm_conv1d_tm_state_pad_v(
3401 qkv_mixed,
3402 &mut rl.conv_state,
3403 la.ssm_conv1d.float_data(),
3404 &mut conv_out,
3405 conv_dim,
3406 t,
3407 d_conv,
3408 pad_len,
3409 )?;
3410 e.qkv_to_gdn_repack(
3411 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3412 )?;
3413 }
3414 let mut q_l2 = e.uninit(d_state * hk * t)?;
3415 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3419 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3420 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3421 Some(qb)
3422 } else {
3423 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3424 None
3425 };
3426 let mut k_l2 = e.uninit(d_state * hk * t)?;
3427 let kb16 = if Engine::l2_v2_on(d_state) {
3429 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3430 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3431 Some(kb)
3432 } else {
3433 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3434 None
3435 };
3436 let mut beta = e.uninit(t * num_v)?;
3437 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3438 let mut g_log = e.uninit(t * num_v)?;
3439 e.gdn_glog_v(
3440 alpha,
3441 la.ssm_dt.float_data(),
3442 la.ssm_a.float_data(),
3443 &mut g_log,
3444 num_v,
3445 t,
3446 )?;
3447 if let Some(len_d) = pad_len {
3448 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3449 }
3450 Ok(GdnPrep {
3451 hk,
3452 q_l2,
3453 k_l2,
3454 v_g,
3455 beta,
3456 g_log,
3457 kb16,
3458 qb16,
3459 })
3460 }
3461
3462 #[allow(clippy::too_many_arguments)]
3467 fn linear_attn_prime_core_batch(
3468 &self,
3469 e: &Engine,
3470 la: &LinearAttnLayer,
3471 g4: &[CudaSlice<f32>],
3472 offs: &[usize],
3473 ts: &[usize],
3474 caches: &mut [&mut Cache],
3475 il: usize,
3476 ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3477 let ssm = self.cfg.ssm.as_ref().unwrap();
3478 let d_state = ssm.state_size as usize;
3479 let num_k = ssm.group_count as usize;
3480 let num_v = ssm.time_step_rank as usize;
3481 let key_dim = d_state * num_k;
3482 let value_dim = d_state * num_v;
3483 let conv_dim = key_dim * 2 + value_dim;
3484 let eps = self.cfg.rms_eps;
3485 let scale = 1.0 / (d_state as f32).sqrt();
3486 let b = ts.len();
3487 let c = Engine::gdn_chunk_size();
3488 let carried = caches.iter().any(|c| c.pos > 0);
3491 let use_vl = !carried
3492 && (2..=8).contains(&b)
3493 && Engine::gdn_chunked_enabled()
3494 && ts.iter().all(|&t| t >= 16)
3495 && e.gdn_mma_enabled(c)
3496 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3497 if !use_vl {
3498 return (0..b)
3499 .map(|s| {
3500 let (o, t) = (offs[s], ts[s]);
3501 self.linear_attn_prime_core_pad_view(
3502 e,
3503 la,
3504 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3505 &g4[1].slice(o * value_dim..(o + t) * value_dim),
3506 &g4[2].slice(o * num_v..(o + t) * num_v),
3507 &g4[3].slice(o * num_v..(o + t) * num_v),
3508 t,
3509 caches[s],
3510 il,
3511 None,
3512 )
3513 })
3514 .collect();
3515 }
3516 struct SeqBufs {
3520 conv_out: CudaSlice<f32>,
3521 q_g: CudaSlice<f32>,
3522 k_g: CudaSlice<f32>,
3523 v_g: CudaSlice<f32>,
3524 q_l2: CudaSlice<f32>,
3525 k_l2: CudaSlice<f32>,
3526 beta: CudaSlice<f32>,
3527 g_log: CudaSlice<f32>,
3528 gn: CudaSlice<f32>,
3529 gn16: CudaSlice<u8>,
3530 }
3531 let d_conv = ssm.conv_kernel as usize;
3532 let f16o = Self::f16out_on(e, 16);
3533 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
3535 let mut pres = Vec::with_capacity(b);
3536 for &t in ts.iter().take(b) {
3537 sb.push(SeqBufs {
3538 conv_out: e.uninit(conv_dim * t)?,
3539 q_g: e.uninit(d_state * hk * t)?,
3540 k_g: e.uninit(d_state * hk * t)?,
3541 v_g: e.uninit(d_state * num_v * t)?,
3542 q_l2: e.uninit(d_state * hk * t)?,
3543 k_l2: e.uninit(d_state * hk * t)?,
3544 beta: e.uninit(t * num_v)?,
3545 g_log: e.uninit(t * num_v)?,
3546 gn: e.uninit(d_state * num_v * t)?,
3547 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3548 });
3549 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3550 }
3551 let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3552 .map(|s| {
3553 let (o, t) = (offs[s], ts[s]);
3554 let rl = caches[s].recur[il].as_ref().unwrap();
3555 crate::GdnPrepVl {
3556 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3557 conv_state: e.addr_f32(&rl.conv_state),
3558 conv_out: e.addr_f32(&sb[s].conv_out),
3559 q_g: e.addr_f32(&sb[s].q_g),
3560 k_g: e.addr_f32(&sb[s].k_g),
3561 v_g: e.addr_f32(&sb[s].v_g),
3562 q_l2: e.addr_f32(&sb[s].q_l2),
3563 k_l2: e.addr_f32(&sb[s].k_l2),
3564 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3565 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3566 beta: e.addr_f32(&sb[s].beta),
3567 g_log: e.addr_f32(&sb[s].g_log),
3568 o: e.addr_f32(&pres[s].o),
3569 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3570 gn: e.addr_f32(&sb[s].gn),
3571 gn16: e.addr_u8(&sb[s].gn16),
3572 kb16: if Engine::l2_v2_on(d_state) {
3573 e.addr_u8(&pres[s].kb16)
3574 } else {
3575 0
3576 },
3577 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3578 e.addr_u8(&pres[s].qb16)
3579 } else {
3580 0
3581 },
3582 t: t as i32,
3583 pad: 0,
3584 }
3585 })
3586 .collect();
3587 let args: Vec<crate::GdnSeqVl> = (0..b)
3588 .map(|s| {
3589 let rl = caches[s].recur[il].as_ref().unwrap();
3590 crate::GdnSeqVl {
3591 kb16: e.addr_u8(&pres[s].kb16),
3592 gcum: e.addr_f32(&pres[s].gcum),
3593 beta: e.addr_f32(&sb[s].beta),
3594 u: e.addr_f32(&pres[s].u),
3595 wb16: e.addr_u8(&pres[s].wb16),
3596 y: e.addr_u8(&pres[s].y16),
3597 ssnap: e.addr_u8(&pres[s].ssnap16),
3598 state_in: e.addr_f32(&rl.ssm_state),
3599 state_out: e.addr_f32(&rl.ssm_state_alt),
3600 q: e.addr_f32(&sb[s].q_l2),
3601 p: e.addr_f32(&pres[s].p),
3602 o: e.addr_f32(&pres[s].o),
3603 k: e.addr_f32(&sb[s].k_l2),
3604 v: e.addr_f32(&sb[s].v_g),
3605 g: e.addr_f32(&sb[s].g_log),
3606 a: e.addr_f32(&pres[s].a),
3607 w: e.addr_f32(&pres[s].w),
3608 t: ts[s] as i32,
3609 nc: pres[s].nc as i32,
3610 }
3611 })
3612 .collect();
3613 e.gdn_prep_vl8(
3614 &prep_args,
3615 la.ssm_conv1d.float_data(),
3616 la.ssm_dt.float_data(),
3617 la.ssm_a.float_data(),
3618 conv_dim,
3619 d_conv,
3620 d_state,
3621 num_v,
3622 num_k,
3623 key_dim,
3624 hk,
3625 eps,
3626 )?;
3627 if !Engine::l2_v2_on(d_state) {
3630 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3631 }
3632 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3634 if !Engine::l2_v2_on(d_state) {
3636 for s in 0..b {
3637 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3638 }
3639 }
3640 let mut wa = [crate::GdnWVl::default(); 8];
3641 for s in 0..b {
3642 wa[s] = crate::GdnWVl {
3643 qb16: e.addr_u8(&pres[s].qb16),
3644 pb16: e.addr_u8(&pres[s].pb16),
3645 };
3646 }
3647 Some(crate::GdnWVl8(wa))
3648 } else {
3649 None
3650 };
3651 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3652 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3653 if f16o {
3654 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3655 }
3656 let mut out = Vec::with_capacity(b);
3658 for (s, bufs) in sb.into_iter().enumerate() {
3659 let rl = caches[s].recur[il].as_mut().unwrap();
3660 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3661 let (o, t) = (offs[s], ts[s]);
3662 let SeqBufs { mut gn, gn16, .. } = bufs;
3663 if f16o {
3664 out.push((gn, Some(gn16)));
3665 } else {
3666 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3667 e.gated_rmsnorm_zv(
3668 &pres[s].o,
3669 la.ssm_norm.float_data(),
3670 &z_v,
3671 &mut gn,
3672 d_state,
3673 num_v * t,
3674 eps,
3675 )?;
3676 out.push((gn, None));
3677 }
3678 }
3679 Ok(out)
3680 }
3681
3682 #[allow(clippy::too_many_arguments)]
3686 fn linear_attn_prime_core_pad_view(
3687 &self,
3688 e: &Engine,
3689 la: &LinearAttnLayer,
3690 qkv_mixed: &cudarc::driver::CudaView<f32>,
3691 z: &cudarc::driver::CudaView<f32>,
3692 beta_raw: &cudarc::driver::CudaView<f32>,
3693 alpha: &cudarc::driver::CudaView<f32>,
3694 t: usize,
3695 cache: &mut Cache,
3696 il: usize,
3697 pad_len: Option<&CudaSlice<i32>>,
3698 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3699 let cfg = &self.cfg;
3700 let ssm = cfg.ssm.as_ref().unwrap();
3701 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
3704 let scale = 1.0 / (d_state as f32).sqrt();
3705
3706 let prep =
3707 self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3708
3709 let mut o = e.uninit(d_state * num_v * t)?;
3715 let rl = cache.recur[il].as_mut().unwrap();
3716 {
3717 let crate::cache::RecurLayer {
3718 ssm_state,
3719 ssm_state_alt,
3720 ..
3721 } = rl;
3722 e.gdn_scan_prefill(
3723 &prep.q_l2,
3724 &prep.k_l2,
3725 &prep.v_g,
3726 &prep.g_log,
3727 &prep.beta,
3728 prep.kb16.as_ref(),
3729 prep.qb16.as_ref(),
3730 ssm_state,
3731 ssm_state_alt,
3732 &mut o,
3733 num_v,
3734 t,
3735 scale,
3736 prep.hk,
3737 )?;
3738 }
3739 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3740
3741 let mut gn = e.uninit(d_state * num_v * t)?;
3744 let gn16 = if Self::f16out_on(e, t) {
3745 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3746 e.gated_rmsnorm_f16out_zv(
3747 &o,
3748 la.ssm_norm.float_data(),
3749 z,
3750 &mut gn,
3751 &mut g16,
3752 d_state,
3753 num_v * t,
3754 eps,
3755 )?;
3756 Some(g16)
3757 } else {
3758 e.gated_rmsnorm_zv(
3759 &o,
3760 la.ssm_norm.float_data(),
3761 z,
3762 &mut gn,
3763 d_state,
3764 num_v * t,
3765 eps,
3766 )?;
3767 None
3768 };
3769 Ok((gn, gn16))
3770 }
3771
3772 #[allow(clippy::too_many_arguments)]
3774 fn linear_attn_prime_core_pad(
3775 &self,
3776 e: &Engine,
3777 la: &LinearAttnLayer,
3778 g4: Vec<CudaSlice<f32>>,
3779 t: usize,
3780 cache: &mut Cache,
3781 il: usize,
3782 pad_len: Option<&CudaSlice<i32>>,
3783 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3784 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3785 if let Some(xh) = &gn16 {
3786 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3787 return Ok(y);
3788 }
3789 }
3790 Ok(e.matmul(&la.ssm_out, &gn, t)?)
3791 }
3792
3793 pub fn full_attn(
3798 &self,
3799 e: &Engine,
3800 fa: &FullAttnLayer,
3801 h: &CudaSlice<f32>,
3802 pos_d: &CudaSlice<i32>,
3803 t: usize,
3804 il: usize,
3805 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3806 if self.cfg.step35.is_some() {
3807 return self.step35_attn(e, fa, h, pos_d, t, il);
3808 }
3809 let cfg = &self.cfg;
3810 let _n_embd = cfg.n_embd as usize;
3811 let geometry = cfg.full_attention_geometry_at(il as u32);
3812 let n_head = geometry.n_head as usize;
3813 let n_head_kv = geometry.n_head_kv as usize;
3814 let head_dim = geometry.head_dim_k as usize;
3815 let eps = cfg.rms_eps;
3816 let scale = geometry.attention_scale();
3817
3818 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3821 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3823 let v = g3.pop().unwrap();
3824 let mut k = g3.pop().unwrap();
3825 let qf = g3.pop().unwrap();
3826 let (mut q, gate) = if gated {
3827 let mut q = e.uninit(t * n_head * head_dim)?;
3828 let mut gate = e.uninit(t * n_head * head_dim)?;
3829 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3830 (q, Some(gate))
3831 } else {
3832 (qf, None)
3833 };
3834
3835 let mut qn = e.uninit(t * n_head * head_dim)?;
3837 e.rms_norm(
3838 &q,
3839 fa.q_norm.float_data(),
3840 &mut qn,
3841 head_dim,
3842 n_head * t,
3843 eps,
3844 )?;
3845 q = qn;
3846 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3847 e.rms_norm(
3848 &k,
3849 fa.k_norm.float_data(),
3850 &mut kn,
3851 head_dim,
3852 n_head_kv * t,
3853 eps,
3854 )?;
3855 k = kn;
3856 let rope_dims = geometry.n_rot as usize;
3857 e.rope_neox(
3858 &mut q,
3859 pos_d,
3860 head_dim,
3861 rope_dims,
3862 n_head,
3863 t,
3864 geometry.rope_base,
3865 1.0,
3866 )?;
3867 e.rope_neox(
3868 &mut k,
3869 pos_d,
3870 head_dim,
3871 rope_dims,
3872 n_head_kv,
3873 t,
3874 geometry.rope_base,
3875 1.0,
3876 )?;
3877
3878 let mut attn = e.uninit(t * n_head * head_dim)?;
3880 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3883 e.sdpa_naive(
3885 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3886 )?;
3887 } else {
3888 e.fa_prefill(
3889 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3890 )?;
3891 }
3892
3893 let attn_g = match &gate {
3895 Some(gate) => {
3896 let mut gsig = e.uninit(t * n_head * head_dim)?;
3897 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3898 let mut ag = e.uninit(t * n_head * head_dim)?;
3899 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3900 ag
3901 }
3902 None => attn,
3903 };
3904
3905 let o = e.matmul(&fa.wo, &attn_g, t)?;
3907 Ok(o)
3908 }
3909
3910 pub fn linear_attn(
3912 &self,
3913 e: &Engine,
3914 la: &LinearAttnLayer,
3915 h: &CudaSlice<f32>,
3916 t: usize,
3917 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3918 let cfg = &self.cfg;
3919 let _n_embd = cfg.n_embd as usize;
3920 let ssm = cfg.ssm.as_ref().unwrap();
3921 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;
3926 let head_v = d_state;
3927 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;
3931 let scale = 1.0 / (d_state as f32).sqrt();
3932
3933 let mut g4 = e.matmul_group(
3936 &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
3937 h,
3938 t,
3939 )?;
3940 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);
3952 let mut q_g = e.uninit(d_state * num_v * t)?;
3953 let mut k_g = e.uninit(d_state * num_v * t)?;
3954 let mut v_g = e.uninit(d_state * num_v * t)?;
3955 e.ssm_conv1d_gdn(
3956 &qkv_mixed,
3957 la.ssm_conv1d.float_data(),
3958 &mut q_g,
3959 &mut k_g,
3960 &mut v_g,
3961 conv_dim,
3962 t,
3963 d_conv,
3964 d_state,
3965 num_v,
3966 num_k,
3967 key_dim,
3968 )?;
3969 let mut q_l2 = e.uninit(d_state * num_v * t)?;
3971 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3972 let mut k_l2 = e.uninit(d_state * num_v * t)?;
3973 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3974 let v_gd = v_g;
3975
3976 let mut beta = e.uninit(t * num_v)?;
3979 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3980 let mut g_log = e.uninit(t * num_v)?;
3982 e.gdn_glog(
3983 &alpha,
3984 la.ssm_dt.float_data(),
3985 la.ssm_a.float_data(),
3986 &mut g_log,
3987 num_v,
3988 t,
3989 )?;
3990
3991 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
3994 let mut o = e.uninit(d_state * num_v * t)?;
3995 e.gdn_scan_prefill(
3996 &q_l2,
3997 &k_l2,
3998 &v_gd,
3999 &g_log,
4000 &beta,
4001 None,
4002 None,
4003 &state_in,
4004 &mut state_out,
4005 &mut o,
4006 num_v,
4007 t,
4008 scale,
4009 num_v,
4010 )?;
4011
4012 let mut gn = e.uninit(d_state * num_v * t)?;
4017 e.gated_rmsnorm(
4018 &o,
4019 la.ssm_norm.float_data(),
4020 &z,
4021 &mut gn,
4022 d_state,
4023 num_v * t,
4024 eps,
4025 )?;
4026
4027 let out = e.matmul(&la.ssm_out, &gn, t)?;
4031 Ok(out)
4032 }
4033}
4034
4035impl HybridModel {
4036 pub fn moe_ffn_il(
4047 &self,
4048 e: &Engine,
4049 m: &MoeWeights,
4050 z: &CudaSlice<f32>,
4051 t: usize,
4052 il: u16,
4053 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4054 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
4055 }
4056
4057 pub fn moe_ffn_il_prefill(
4060 &self,
4061 e: &Engine,
4062 m: &MoeWeights,
4063 z: &CudaSlice<f32>,
4064 t: usize,
4065 il: u16,
4066 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4067 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
4068 }
4069
4070 pub fn moe_ffn_il_zq8(
4074 &self,
4075 e: &Engine,
4076 m: &MoeWeights,
4077 z: &CudaSlice<f32>,
4078 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4079 t: usize,
4080 il: u16,
4081 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4082 Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4083 }
4084
4085 pub(crate) fn moe_ffn(
4093 e: &Engine,
4094 m: &MoeWeights,
4095 z: &CudaSlice<f32>,
4096 t: usize,
4097 cfg: &ModelConfig,
4098 il: u16,
4099 max_block: usize,
4100 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4101 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4102 }
4103
4104 #[allow(clippy::too_many_arguments)]
4105 pub(crate) fn moe_ffn_inner(
4106 e: &Engine,
4107 m: &MoeWeights,
4108 z: &CudaSlice<f32>,
4109 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4110 t: usize,
4111 cfg: &ModelConfig,
4112 il: u16,
4113 max_block: usize,
4114 prefill: bool,
4115 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4116 let worker_io = crate::spill_pread::worker_enabled();
4117 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4118 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4119 e.with_moe_cache(max_block, |cache, _| {
4120 cache.begin_forward_epoch(il, t);
4121 if worker_io {
4122 cache.begin_worker_scope();
4123 }
4124 Ok(())
4125 })?;
4126 }
4127 if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4128 let moe = cfg.moe.as_ref().unwrap();
4129 let n_expert = moe.expert_count as usize;
4130 let n_used = moe.expert_used_count as usize;
4131 let sigmoid = cfg.sigmoid_router().unwrap();
4132 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4133 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4134 return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4135 }
4136 if t > 1 && moe_grouped_enabled(cfg, prefill) {
4139 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4140 if std::env::var("MEMRA_MOE_GATE").is_ok() {
4145 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4146 let g_host = e.dtoh(&grouped_out)?;
4147 let s_host = e.dtoh(&seq_out)?;
4148 let g_bytes: &[u8] = unsafe {
4149 std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4150 };
4151 let s_bytes: &[u8] = unsafe {
4152 std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4153 };
4154 if g_bytes == s_bytes {
4155 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4156 } else {
4157 let diffs = g_host
4158 .iter()
4159 .zip(s_host.iter())
4160 .enumerate()
4161 .filter(|(_, (a, b))| a != b)
4162 .count();
4163 let maxdiff = g_host
4164 .iter()
4165 .zip(s_host.iter())
4166 .map(|(a, b)| (a - b).abs())
4167 .fold(0.0f32, f32::max);
4168 panic!(
4169 "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4170 g_host.len()
4171 );
4172 }
4173 }
4174 return Ok(grouped_out);
4175 }
4176 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4177 }
4178
4179 fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4180 let Some(moe) = cfg.moe.as_ref() else {
4181 return false;
4182 };
4183 static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4186 let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4187 std::env::var("MEMRA_MOE_STATS").is_ok()
4188 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4189 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4190 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4191 || std::env::var("MEMRA_MOE_GATE").is_ok()
4192 });
4193 cfg.step35.is_some()
4194 && sigmoid_router_enabled()
4195 && moe_dev_enabled()
4196 && moe_slab_enabled()
4197 && !observation_mode
4198 && moe.expert_used_count <= 8
4199 && m.has_uniform_expert_layout()
4200 && m.gate_exps.macros.is_none()
4201 && m.up_exps.macros.is_none()
4202 && m.down_exps.macros.is_none()
4203 && !m.has_macros
4204 && moe_q8_enabled()
4205 && q8_expert_supported(m.gate_exps.qtype)
4206 && q8_expert_supported(m.up_exps.qtype)
4207 && q8_expert_supported(m.down_exps.qtype)
4208 && m.dev_exps
4209 .as_ref()
4210 .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4211 }
4212
4213 pub(crate) fn moe_ffn_sequential(
4215 e: &Engine,
4216 m: &MoeWeights,
4217 z: &CudaSlice<f32>,
4218 t: usize,
4219 cfg: &ModelConfig,
4220 il: u16,
4221 max_block: usize,
4222 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4223 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4224 }
4225
4226 fn moe_router_logits(
4230 e: &Engine,
4231 m: &MoeWeights,
4232 z: &CudaSlice<f32>,
4233 t: usize,
4234 cfg: &ModelConfig,
4235 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4236 if t < PRIME_MIN_T {
4237 if crate::router_kernel_on() {
4239 e.router_gemv(
4240 m.gate_inp.float_data(),
4241 z,
4242 cfg.n_embd as usize,
4243 m.gate_exps.n_expert,
4244 t,
4245 )
4246 } else {
4247 e.matmul_decode_exact(&m.gate_inp, z, t)
4248 }
4249 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4250 e.router_gemv(
4251 m.gate_inp.float_data(),
4252 z,
4253 cfg.n_embd as usize,
4254 m.gate_exps.n_expert,
4255 t,
4256 )
4257 } else {
4258 e.matmul(&m.gate_inp, z, t)
4259 }
4260 }
4261
4262 fn trace_moe_routes(
4266 il: u16,
4267 t: usize,
4268 sel_all: &[u32],
4269 weights: &[f32],
4270 ) -> Result<(), Box<dyn std::error::Error>> {
4271 use std::io::Write as _;
4272 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4273 let mut f = std::fs::OpenOptions::new()
4274 .create(true)
4275 .append(true)
4276 .open(path)?;
4277 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4278 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4279 }
4280 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4281 let mut f = std::fs::OpenOptions::new()
4282 .create(true)
4283 .append(true)
4284 .open(path)?;
4285 let pairs: Vec<String> = sel_all
4286 .iter()
4287 .zip(weights)
4288 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4289 .collect();
4290 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4291 }
4292 Ok(())
4293 }
4294
4295 #[allow(clippy::too_many_arguments)]
4296 fn trace_sigmoid_router_logits(
4297 e: &Engine,
4298 il: u16,
4299 t: usize,
4300 n_expert: usize,
4301 n_used: usize,
4302 logits: &CudaSlice<f32>,
4303 m: &MoeWeights,
4304 (scaling_factor, route_norm): (f32, bool),
4305 ) -> Result<(), Box<dyn std::error::Error>> {
4306 if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4307 return Ok(());
4308 }
4309 let logits = e.dtoh(logits)?;
4310 let active: Vec<u8> = m
4311 .active_experts
4312 .as_ref()
4313 .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4314 .unwrap_or_else(|| vec![1; n_expert]);
4315 let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4316 crate::sigrouter_contract::capture_served_logits(
4317 il as u32,
4318 t,
4319 n_expert,
4320 n_used,
4321 scaling_factor,
4322 route_norm,
4323 &active,
4324 &bias,
4325 &logits,
4326 )?;
4327 Ok(())
4328 }
4329
4330 fn trace_moe_input(
4335 e: &Engine,
4336 il: u16,
4337 t: usize,
4338 n_embd: usize,
4339 z: &CudaSlice<f32>,
4340 ) -> Result<(), Box<dyn std::error::Error>> {
4341 use std::io::Write as _;
4342 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4343 return Ok(());
4344 };
4345 let host = e.dtoh(z)?;
4346 if host.len() != t * n_embd {
4347 return Err(format!(
4348 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4349 host.len(),
4350 t,
4351 n_embd
4352 )
4353 .into());
4354 }
4355 let bytes = unsafe {
4356 std::slice::from_raw_parts(
4357 host.as_ptr().cast::<u8>(),
4358 host.len() * std::mem::size_of::<f32>(),
4359 )
4360 };
4361 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4362 let mut state = state
4363 .lock()
4364 .map_err(|_| "MoE input trace writer lock is poisoned")?;
4365 if state.is_none() {
4366 let dir = std::path::PathBuf::from(&dir);
4367 std::fs::create_dir_all(&dir)?;
4368 let index = std::fs::OpenOptions::new()
4369 .create(true)
4370 .append(true)
4371 .open(dir.join("index.jsonl"))?;
4372 *state = Some(MoeInputTraceWriter {
4373 dir,
4374 index,
4375 payloads: std::collections::HashMap::new(),
4376 });
4377 }
4378 let writer = state.as_mut().unwrap();
4379 if writer.dir != std::path::Path::new(&dir) {
4380 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4381 }
4382 let file_name = format!("layer-{il:03}.f32");
4383 if !writer.payloads.contains_key(&il) {
4384 let payload = std::fs::OpenOptions::new()
4385 .create(true)
4386 .append(true)
4387 .open(writer.dir.join(&file_name))?;
4388 let offset = payload.metadata()?.len();
4389 writer.payloads.insert(il, (payload, offset));
4390 }
4391 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4392 let row_offset = *offset;
4393 payload.write_all(bytes)?;
4394 *offset += bytes.len() as u64;
4395 writeln!(
4396 writer.index,
4397 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4398 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4399 \"payload_bytes\":{}}}",
4400 bytes.len()
4401 )?;
4402 Ok(())
4403 }
4404
4405 #[allow(clippy::too_many_arguments)]
4406 pub(crate) fn moe_ffn_sequential_zq8(
4407 e: &Engine,
4408 m: &MoeWeights,
4409 z: &CudaSlice<f32>,
4410 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4411 t: usize,
4412 cfg: &ModelConfig,
4413 il: u16,
4414 max_block: usize,
4415 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4416 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4417 let moe = cfg.moe.as_ref().unwrap();
4418 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);
4425 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4426 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);
4429
4430 let lim_exp = cfg.clamp_exp_at(il as u32);
4433 let lim_shexp = cfg.clamp_shexp_at(il as u32);
4434 let use_cache = Engine::moe_cache_enabled();
4435 let uniform_experts = m.has_uniform_expert_layout();
4436 let moe_q8 = uniform_experts
4437 && moe_q8_enabled()
4438 && q8_expert_supported(m.gate_exps.qtype)
4439 && q8_expert_supported(m.up_exps.qtype)
4440 && q8_expert_supported(m.down_exps.qtype);
4441 let cpu_expert_requested = crate::cpu_experts::configured();
4448 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4449 return Err(std::io::Error::other(
4450 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4451 )
4452 .into());
4453 }
4454 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4455 let freeze_cpu_residency = cpu_expert_requested
4461 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4462 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4463 .ok()
4464 .and_then(|value| value.parse::<usize>().ok())
4465 .is_some_and(|tokens| tokens > 0);
4466 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4467 e.freeze_moe_cache();
4468 }
4469 let cache_frozen = use_cache && e.moe_cache_frozen();
4470 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4471
4472 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4475 if let Some(sig) = cfg.sigmoid_router() {
4476 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4477 }
4478
4479 let no_exp_macros = m.gate_exps.macros.is_none()
4518 && m.up_exps.macros.is_none()
4519 && m.down_exps.macros.is_none();
4520 if cfg.sigmoid_router().is_none()
4524 && cfg.m3.is_none()
4525 && cfg.hy3.is_none()
4526 && !cfg.swiglu_clamped_at(il as u32)
4527 && no_exp_macros
4528 && t >= PRIME_MIN_T
4529 && m.dev_exps.is_some()
4530 && moe_q8_enabled()
4531 && q8_expert_supported(m.gate_exps.qtype)
4532 && q8_expert_supported(m.up_exps.qtype)
4533 && q8_expert_supported(m.down_exps.qtype)
4534 && std::env::var("MEMRA_MOE_PAIRS")
4535 .map(|v| v != "0")
4536 .unwrap_or(true)
4537 && std::env::var("MEMRA_MOE_STATS").is_err()
4538 {
4539 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4540 }
4541
4542 let dev_ok = uniform_experts
4560 && cfg.sigmoid_router().is_none()
4561 && cfg.m3.is_none()
4562 && cfg.hy3.is_none()
4563 && !cfg.swiglu_clamped_at(il as u32);
4564 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4568 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4569 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4570 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4571 if dev_ok
4572 && t < PRIME_MIN_T
4573 && m.dev_exps.is_some()
4574 && n_used <= 8
4575 && moe_dev_enabled()
4576 && !observe_routes
4577 {
4578 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4579 }
4580 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4581 let row_ok = e.with_moe_cache(max_block, |c, eng| {
4582 if moe_prewarm_enabled() {
4583 c.prewarm_layer(il, m, eng)?;
4584 }
4585 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4586 })?;
4587 if row_ok {
4588 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4589 }
4590 }
4591
4592 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4594 if cpu_hybrid {
4595 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4596 e,
4597 &logits,
4598 z,
4599 t,
4600 n_expert,
4601 n_used,
4602 m.exp_probs_b.as_deref(),
4603 sig,
4604 m.active_experts.as_deref(),
4605 )?;
4606 (sel, w, Some(input))
4607 } else {
4608 let (sel, w) =
4609 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4610 (sel, w, None)
4611 }
4612 } else {
4613 let (sel, w) =
4614 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4615 (sel, w, None)
4616 };
4617 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4618
4619 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4623 Self::trace_moe_input(e, il, t, n_embd, z)?;
4624
4625 let worker_disk_prefetch =
4637 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4638 let promote_worker_h2d =
4639 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4640 if promote_worker_h2d {
4641 let mut selected_blocks = Vec::with_capacity(n_used * 3);
4642 for &ex in sel_all.iter().take(n_used) {
4643 let ex = ex as u16;
4644 selected_blocks.extend([
4645 BlockId::new(il, PROJ_GATE, ex),
4646 BlockId::new(il, PROJ_UP, ex),
4647 BlockId::new(il, PROJ_DOWN, ex),
4648 ]);
4649 }
4650 for &ex in sel_all.iter().take(n_used) {
4651 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4652 }
4653 e.with_moe_cache(max_block, |cache, eng| {
4654 cache.promote_worker_reads_at_safe_boundary(
4655 &selected_blocks,
4656 &selected_blocks,
4657 eng,
4658 )?;
4659 Ok(())
4660 })?;
4661 }
4662
4663 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4666 let mut cnt = vec![0u32; n_expert];
4667 for &s in sel_all.iter() {
4668 cnt[s as usize] += 1;
4669 }
4670 let total = sel_all.len() as f64;
4671 let mut h = 0.0f64;
4672 let mut active = 0usize;
4673 for &c in &cnt {
4674 if c > 0 {
4675 active += 1;
4676 let p = c as f64 / total;
4677 h -= p * p.log2();
4678 }
4679 }
4680 let maxc = cnt.iter().copied().max().unwrap_or(0);
4681 println!(
4682 "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4683 il,
4684 t,
4685 sel_all.len(),
4686 active,
4687 n_expert,
4688 h,
4689 (n_expert as f64).log2(),
4690 total / active.max(1) as f64,
4691 maxc
4692 );
4693 }
4694
4695 let gdec_may_fire = uniform_experts
4708 && use_cache
4709 && n_used <= 8
4710 && gdec_enabled()
4711 && !cfg.swiglu_clamped_at(il as u32);
4712 let slab_local = m
4728 .dev_exps
4729 .as_ref()
4730 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4731 let slab_bases = slab_local.map(|d| {
4732 use cudarc::driver::DevicePtr;
4733 let s = e.stream();
4734 let (pg, _g0) = d.gate.device_ptr(&s);
4735 let (pu, _g1) = d.up.device_ptr(&s);
4736 let (pd, _g2) = d.down.device_ptr(&s);
4737 (pg as u64, pu as u64, pd as u64)
4738 });
4739 let slab_fused_may_fire = slab_bases.is_some()
4749 && n_used <= 8
4750 && gdec_enabled()
4751 && !cfg.swiglu_clamped_at(il as u32)
4752 && cfg.m3.is_none()
4753 && no_exp_macros
4754 && moe_q8;
4755 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4758 e.uninit(t * n_embd)?
4759 } else {
4760 e.zeros(t * n_embd)?
4761 };
4762 let cpu_input = if cpu_hybrid {
4765 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4766 } else {
4767 None
4768 };
4769
4770 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;
4778 let mut scratch_u: Option<CudaSlice<u8>> = None;
4779 let mut scratch_d: Option<CudaSlice<u8>> = None;
4780 let page_window = moe_page_prefetch_window();
4788
4789 for tok in 0..t {
4792 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4793 let w = &w_all[tok * n_used..(tok + 1) * n_used];
4794 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4796
4797 let no_macros = m.gate_exps.macros.is_none()
4811 && m.up_exps.macros.is_none()
4812 && m.down_exps.macros.is_none();
4813 if slab_fused_may_fire {
4823 let (pg, pu, pd) = slab_bases.unwrap();
4824 let mut gp = [0u64; 8];
4825 let mut up = [0u64; 8];
4826 let mut dp = [0u64; 8];
4827 for (j, &ex) in sel.iter().enumerate() {
4828 let ex = ex as usize;
4829 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4830 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4831 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4832 }
4833 let mut wv = [0f32; 8];
4834 wv[..n_used].copy_from_slice(w);
4835 if tok_q8.is_none() {
4836 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4837 }
4838 let (zq, zd) = tok_q8.as_ref().unwrap();
4839 let act = e.moe_gate_up_silu8_q8(
4840 crate::WPtr8(gp),
4841 crate::WPtr8(up),
4842 zq,
4843 zd,
4844 n_embd,
4845 n_ff_exp,
4846 n_used,
4847 m.gate_exps.qtype,
4848 m.up_exps.qtype,
4849 m.gate_exps.row_bytes,
4850 m.up_exps.row_bytes,
4851 )?;
4852 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4853 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4854 e.moe_down8_fma_q8(
4855 crate::WPtr8(dp),
4856 crate::F32x8(wv),
4857 &aq2,
4858 &ad2,
4859 &mut dst,
4860 n_ff_exp,
4861 n_embd,
4862 n_used,
4863 m.down_exps.qtype,
4864 m.down_exps.row_bytes,
4865 )?;
4866 continue;
4867 }
4868 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4869 if tok_q8.is_none() {
4870 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4871 }
4872 let (zq, zd) = tok_q8.as_ref().unwrap();
4873 if Self::moe_gdec_token_q8(
4874 e,
4875 m,
4876 il,
4877 max_block,
4878 zq,
4879 zd,
4880 sel,
4881 w,
4882 &mut moe_out,
4883 tok,
4884 n_embd,
4885 n_ff_exp,
4886 n_used,
4887 )? {
4888 continue;
4889 }
4890 } else if gdec_may_fire
4891 && cfg.m3.is_none()
4892 && no_macros
4893 && Self::moe_gdec_token(
4894 e,
4895 m,
4896 il,
4897 max_block,
4898 &zt,
4899 sel,
4900 w,
4901 &mut moe_out,
4902 tok,
4903 n_embd,
4904 n_ff_exp,
4905 n_used,
4906 )?
4907 {
4908 continue;
4909 }
4910
4911 if gdec_may_fire || slab_fused_may_fire {
4917 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4918 e.memset_zeros_view(&mut row)?;
4919 }
4920
4921 let mut cpu_mask = vec![false; sel.len()];
4927 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
4928 let gpu_resident = if use_cache {
4929 e.with_moe_cache(max_block, |cache, _| {
4930 Ok(sel
4931 .iter()
4932 .map(|&expert| {
4933 let expert = expert as u16;
4934 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
4935 .into_iter()
4936 .filter(|&projection| {
4937 cache
4938 .resident(BlockId::new(il, projection, expert))
4939 .is_some()
4940 })
4941 .count()
4942 })
4943 .collect::<Vec<_>>())
4944 })?
4945 } else {
4946 vec![0; sel.len()]
4947 };
4948 let mut cpu_selected = Vec::new();
4949 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
4950 if gpu_resident[index] != 3 {
4951 cpu_mask[index] = true;
4952 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
4953 let expert = expert as usize;
4954 cpu_selected.push((expert, route_weight));
4955 }
4956 }
4957 if crate::cpu_experts::predictor_enabled() {
4958 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4962 crate::cpu_experts::predictor_submit(il, row);
4963 }
4964 if cpu_selected.is_empty() {
4965 None
4966 } else {
4967 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4968 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
4969 .map_err(std::io::Error::other)?;
4970 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
4971 }
4972 } else {
4973 None
4974 };
4975
4976 let worker_window = worker_disk_prefetch
4977 .then(worker_prefetch_window)
4978 .unwrap_or(0);
4979 for (j, &ex) in sel.iter().enumerate() {
4980 if cpu_mask[j] {
4981 continue;
4982 }
4983 let ex = ex as usize;
4984 if let Some(d) = slab_local {
4991 let gl = m.gate_exps.expert_layout(ex);
4992 let ul = m.up_exps.expert_layout(ex);
4993 let dl = m.down_exps.expert_layout(ex);
4994 let (g0, u0, d0) = (
4995 ex * m.gate_exps.expert_stride,
4996 ex * m.up_exps.expert_stride,
4997 ex * m.down_exps.expert_stride,
4998 );
4999 let (gate, up) = if moe_q8 {
5000 if tok_q8.is_none() {
5001 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5002 }
5003 let (zq, zd) = tok_q8.as_ref().unwrap();
5004 (
5005 e.qmatvec_expert_q8(
5006 &d.gate,
5007 g0..g0 + gl.len,
5008 zq,
5009 zd,
5010 1,
5011 m.gate_exps.in_f,
5012 m.gate_exps.out_f,
5013 gl.qtype,
5014 gl.row_bytes,
5015 )?,
5016 e.qmatvec_expert_q8(
5017 &d.up,
5018 u0..u0 + ul.len,
5019 zq,
5020 zd,
5021 1,
5022 m.up_exps.in_f,
5023 m.up_exps.out_f,
5024 ul.qtype,
5025 ul.row_bytes,
5026 )?,
5027 )
5028 } else {
5029 (
5030 e.qmatvec_view(
5031 &d.gate,
5032 g0..g0 + gl.len,
5033 &zt,
5034 1,
5035 m.gate_exps.in_f,
5036 m.gate_exps.out_f,
5037 gl.qtype,
5038 gl.row_bytes,
5039 )?,
5040 e.qmatvec_view(
5041 &d.up,
5042 u0..u0 + ul.len,
5043 &zt,
5044 1,
5045 m.up_exps.in_f,
5046 m.up_exps.out_f,
5047 ul.qtype,
5048 ul.row_bytes,
5049 )?,
5050 )
5051 };
5052 let mut act = e.uninit(n_ff_exp)?;
5053 Self::ffn_act_lim(
5054 e,
5055 cfg,
5056 &gate,
5057 &up,
5058 m.gate_exps.macro_scale(ex),
5059 m.up_exps.macro_scale(ex),
5060 lim_exp,
5061 &mut act,
5062 n_ff_exp,
5063 )?;
5064 let y = if moe_q8 {
5065 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5066 e.qmatvec_expert_q8(
5067 &d.down,
5068 d0..d0 + dl.len,
5069 &aq2,
5070 &ad2,
5071 1,
5072 m.down_exps.in_f,
5073 m.down_exps.out_f,
5074 dl.qtype,
5075 dl.row_bytes,
5076 )?
5077 } else {
5078 let actv = act.slice(0..n_ff_exp);
5079 e.qmatvec_view(
5080 &d.down,
5081 d0..d0 + dl.len,
5082 &actv,
5083 1,
5084 m.down_exps.in_f,
5085 m.down_exps.out_f,
5086 dl.qtype,
5087 dl.row_bytes,
5088 )?
5089 };
5090 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5091 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5092 continue;
5093 }
5094 for next in page_prefetch_positions(j, sel.len(), page_window) {
5095 Self::moe_prefetch_host_expert(sel[next] as usize, m);
5096 }
5097 let keep = [
5098 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5099 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5100 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5101 ];
5102 if worker_disk_prefetch && worker_window > 0 {
5103 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5104 Self::moe_prefetch_disk_expert(
5105 e,
5106 il,
5107 sel[next] as usize,
5108 m,
5109 max_block,
5110 &keep,
5111 )?;
5112 }
5113 } else if cache_dispatch
5114 && !cpu_hybrid
5115 && moe_prefetch_enabled()
5116 && j + 1 < sel.len()
5117 {
5118 let next = sel[j + 1] as usize;
5119 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5120 }
5121 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5122 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5123 if (gate_q8 || up_q8) && tok_q8.is_none() {
5126 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5127 }
5128 let gate = if gate_q8 {
5129 let (zq, zd) = tok_q8.as_ref().unwrap();
5130 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5131 } else {
5132 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5133 };
5134 let up = if up_q8 {
5135 let (zq, zd) = tok_q8.as_ref().unwrap();
5136 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5137 } else {
5138 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5139 };
5140 let mut act = e.uninit(n_ff_exp)?;
5141 Self::ffn_act_lim(
5142 e,
5143 cfg,
5144 &gate,
5145 &up,
5146 m.gate_exps.macro_scale(ex),
5147 m.up_exps.macro_scale(ex),
5148 lim_exp,
5149 &mut act,
5150 n_ff_exp,
5151 )?;
5152 let y = if down_q8 {
5153 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5154 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5155 } else {
5156 let actv = act.slice(0..n_ff_exp);
5157 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5158 };
5159 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5160 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5162 } else if cache_dispatch {
5163 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5168 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5169 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5171 e,
5172 cfg,
5173 &gate,
5174 &up,
5175 m.gate_exps.macro_scale(ex),
5176 m.up_exps.macro_scale(ex),
5177 lim_exp,
5178 &mut act,
5179 n_ff_exp,
5180 )?;
5181 let actv = act.slice(0..n_ff_exp);
5182 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5183 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5184 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5186 } else if cache_frozen {
5187 let gate = Self::moe_frozen_gemm(
5192 e,
5193 il,
5194 PROJ_GATE,
5195 ex,
5196 m,
5197 max_block,
5198 &zt,
5199 &mut scratch_g,
5200 g_len,
5201 )?;
5202 let up = Self::moe_frozen_gemm(
5203 e,
5204 il,
5205 PROJ_UP,
5206 ex,
5207 m,
5208 max_block,
5209 &zt,
5210 &mut scratch_u,
5211 u_len,
5212 )?;
5213 let mut act = e.uninit(n_ff_exp)?;
5214 Self::ffn_act_lim(
5215 e,
5216 cfg,
5217 &gate,
5218 &up,
5219 m.gate_exps.macro_scale(ex),
5220 m.up_exps.macro_scale(ex),
5221 lim_exp,
5222 &mut act,
5223 n_ff_exp,
5224 )?;
5225 let actv = act.slice(0..n_ff_exp);
5226 let y = Self::moe_frozen_gemm(
5227 e,
5228 il,
5229 PROJ_DOWN,
5230 ex,
5231 m,
5232 max_block,
5233 &actv,
5234 &mut scratch_d,
5235 d_len,
5236 )?;
5237 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5238 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5239 } else {
5240 if scratch_g.is_none() {
5244 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5245 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5246 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5247 }
5248 let (sg, su, sd) = (
5249 scratch_g.as_mut().unwrap(),
5250 scratch_u.as_mut().unwrap(),
5251 scratch_d.as_mut().unwrap(),
5252 );
5253 let gl = m.gate_exps.expert_layout(ex);
5254 let ul = m.up_exps.expert_layout(ex);
5255 let dl = m.down_exps.expert_layout(ex);
5256 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5257 let gate = e.qmatvec_view(
5258 sg,
5259 0..gl.len,
5260 &zt,
5261 1,
5262 m.gate_exps.in_f,
5263 m.gate_exps.out_f,
5264 gl.qtype,
5265 gl.row_bytes,
5266 )?;
5267
5268 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5269 let up = e.qmatvec_view(
5270 su,
5271 0..ul.len,
5272 &zt,
5273 1,
5274 m.up_exps.in_f,
5275 m.up_exps.out_f,
5276 ul.qtype,
5277 ul.row_bytes,
5278 )?;
5279
5280 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5282 e,
5283 cfg,
5284 &gate,
5285 &up,
5286 m.gate_exps.macro_scale(ex),
5287 m.up_exps.macro_scale(ex),
5288 lim_exp,
5289 &mut act,
5290 n_ff_exp,
5291 )?;
5292
5293 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5294 let actv = act.slice(0..n_ff_exp);
5295 let y = e.qmatvec_view(
5296 sd,
5297 0..dl.len,
5298 &actv,
5299 1,
5300 m.down_exps.in_f,
5301 m.down_exps.out_f,
5302 dl.qtype,
5303 dl.row_bytes,
5304 )?;
5305
5306 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5307 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5308 }
5309 }
5310 if let Some(worker) = cpu_worker {
5311 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5312 let cpu_output = e.htod(&cpu_output)?;
5313 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5314 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5315 }
5316 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5317 for (j, &ex) in sel.iter().enumerate() {
5318 if cpu_mask[j] {
5319 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5320 }
5321 }
5322 }
5323 }
5324
5325 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5330 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5331 {
5332 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
5341 let (sg_gate, sg_up) = if t == 1 {
5342 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5343 Some(pair) => pair,
5344 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5345 }
5346 } else if verify_t {
5347 (
5348 e.matmul_decode_exact(gate_shexp, z, t)?,
5349 e.matmul_decode_exact(up_shexp, z, t)?,
5350 )
5351 } else {
5352 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
5354 let mut sa = e.uninit(t * n_ff_sh)?; Self::ffn_act_lim(
5356 e,
5357 cfg,
5358 &sg_gate,
5359 &sg_up,
5360 1.0,
5361 1.0,
5362 lim_shexp,
5363 &mut sa,
5364 t * n_ff_sh,
5365 )?;
5366 let sh = if verify_t {
5367 e.matmul_decode_exact(down_shexp, &sa, t)?
5368 } else {
5369 e.matmul(down_shexp, &sa, t)?
5370 }; let g = match &m.gate_inp_shexp {
5384 Some(gate_inp_shexp) => {
5385 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5386 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5387 } else {
5388 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5389 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
5391 g
5392 }
5393 }
5394 None => e.htod(&vec![1.0f32; t])?,
5395 };
5396 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5398 }
5399
5400 Ok(moe_out)
5401 }
5402
5403 pub fn stage1_h2d_per_token(&self) -> u64 {
5406 use crate::hybrid::Ffn;
5407 let n_used = self
5408 .cfg
5409 .moe
5410 .as_ref()
5411 .map(|m| m.expert_used_count as u64)
5412 .unwrap_or(0);
5413 let mut bytes = 0u64;
5414 for l in self.layers.iter() {
5415 if let Ffn::Moe(m) = &l.ffn {
5416 bytes += n_used
5417 * (m.gate_exps.max_expert_bytes()
5418 + m.up_exps.max_expert_bytes()
5419 + m.down_exps.max_expert_bytes()) as u64;
5420 }
5421 }
5422 bytes
5423 }
5424
5425 pub(crate) fn max_moe_block(&self) -> usize {
5429 use crate::hybrid::Ffn;
5430 let mut mx = 0usize;
5431 let mut scan = |ffn: &Ffn| {
5432 if let Ffn::Moe(m) = ffn {
5433 mx = mx
5434 .max(m.gate_exps.max_expert_bytes())
5435 .max(m.up_exps.max_expert_bytes())
5436 .max(m.down_exps.max_expert_bytes());
5437 }
5438 };
5439 for l in self.layers.iter() {
5440 scan(&l.ffn);
5441 }
5442 if let Some(mtp) = self.mtp.as_ref() {
5443 scan(&mtp.ffn);
5444 }
5445 mx
5446 }
5447
5448 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5451 use crate::hybrid::Ffn;
5452 let mut sizes = Vec::new();
5453 let mut scan = |ffn: &Ffn| {
5454 let Ffn::Moe(m) = ffn else { return };
5455 for ex in 0..m.gate_exps.n_expert {
5456 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5457 continue;
5458 }
5459 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5460 let len = exps.expert_layout(ex).len;
5461 if len > 0 {
5462 sizes.push(len);
5463 }
5464 }
5465 }
5466 };
5467 for layer in &self.layers {
5468 scan(&layer.ffn);
5469 }
5470 if let Some(mtp) = &self.mtp {
5471 scan(&mtp.ffn);
5472 }
5473 sizes
5474 }
5475
5476 pub fn save_cpu_expert_residency_profile(
5482 &self,
5483 e: &Engine,
5484 path: &std::path::Path,
5485 ) -> Result<(), Box<dyn std::error::Error>> {
5486 let Some(ids) = e.export_moe_residency() else {
5487 return Err("no MoE residency cache to persist".into());
5488 };
5489 let mut body = format!(
5490 "memra-freeze-profile v1 max_block={} blocks={}\n",
5491 self.max_moe_block(),
5492 ids.len()
5493 );
5494 for (layer, proj, ex) in &ids {
5495 body.push_str(&format!("{layer} {proj} {ex}\n"));
5496 }
5497 let tmp = path.with_extension("tmp");
5498 std::fs::write(&tmp, body)?;
5499 std::fs::rename(&tmp, path)?;
5500 println!(
5501 "[moe-cache] freeze profile saved: {} blocks -> {}",
5502 ids.len(),
5503 path.display()
5504 );
5505 Ok(())
5506 }
5507
5508 pub fn restore_cpu_expert_residency_profile(
5512 &self,
5513 e: &Engine,
5514 path: &std::path::Path,
5515 ) -> Result<bool, Box<dyn std::error::Error>> {
5516 use crate::hybrid::Ffn;
5517 use crate::moe_cache::BlockId;
5518 let Ok(content) = std::fs::read_to_string(path) else {
5519 return Ok(false);
5520 };
5521 let mut lines = content.lines();
5522 let Some(header) = lines.next() else {
5523 return Ok(false);
5524 };
5525 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5526 if !header.starts_with(&expected) {
5527 println!(
5528 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5529 path.display()
5530 );
5531 return Ok(false);
5532 }
5533 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5534 std::collections::HashMap::new();
5535 for line in lines {
5536 let mut fields = line.split_whitespace();
5537 let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5538 else {
5539 continue;
5540 };
5541 let (Ok(layer), Ok(proj), Ok(ex)) =
5542 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5543 else {
5544 continue;
5545 };
5546 by_layer
5547 .entry(layer)
5548 .or_default()
5549 .push(BlockId::new(layer, proj, ex));
5550 }
5551 let requested: usize = by_layer.values().map(Vec::len).sum();
5552 if requested == 0 {
5553 return Ok(false);
5554 }
5555 let max_block = self.max_moe_block();
5556 let mut restaged = 0usize;
5557 let mut stage_layer =
5558 |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5559 let Ffn::Moe(m) = ffn else { return Ok(()) };
5560 let Some(ids) = by_layer.get(&layer_index) else {
5561 return Ok(());
5562 };
5563 e.with_moe_cache(max_block, |cache, eng| {
5564 for id in ids {
5565 if cache.restage_block(*id, m, eng)? {
5566 restaged += 1;
5567 }
5568 }
5569 Ok(())
5570 })
5571 };
5572 for (index, layer) in self.layers.iter().enumerate() {
5573 stage_layer(index as u16, &layer.ffn)?;
5574 }
5575 if let Some(mtp) = self.mtp.as_ref() {
5576 stage_layer(u16::MAX, &mtp.ffn)?;
5577 }
5578 e.freeze_moe_cache();
5579 println!(
5580 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5581 path.display()
5582 );
5583 Ok(true)
5584 }
5585
5586 pub fn freeze_cpu_expert_residency(
5588 &self,
5589 e: &Engine,
5590 ) -> Result<(), Box<dyn std::error::Error>> {
5591 e.freeze_moe_cache();
5592 Ok(())
5593 }
5594
5595 pub fn ffn_act(
5603 e: &Engine,
5604 cfg: &ModelConfig,
5605 gate: &CudaSlice<f32>,
5606 up: &CudaSlice<f32>,
5607 act: &mut CudaSlice<f32>,
5608 n: usize,
5609 ) -> Result<(), Box<dyn std::error::Error>> {
5610 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5611 }
5612
5613 #[allow(clippy::too_many_arguments)]
5617 pub(crate) fn ffn_act_scaled(
5618 e: &Engine,
5619 cfg: &ModelConfig,
5620 gate: &CudaSlice<f32>,
5621 up: &CudaSlice<f32>,
5622 gs: f32,
5623 us: f32,
5624 act: &mut CudaSlice<f32>,
5625 n: usize,
5626 ) -> Result<(), Box<dyn std::error::Error>> {
5627 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5628 }
5629
5630 #[allow(clippy::too_many_arguments)]
5639 pub(crate) fn ffn_act_lim(
5640 e: &Engine,
5641 cfg: &ModelConfig,
5642 gate: &CudaSlice<f32>,
5643 up: &CudaSlice<f32>,
5644 gs: f32,
5645 us: f32,
5646 limit: Option<f32>,
5647 act: &mut CudaSlice<f32>,
5648 n: usize,
5649 ) -> Result<(), Box<dyn std::error::Error>> {
5650 if let Some(m3) = cfg.m3.as_ref() {
5651 debug_assert!(
5652 limit.is_none(),
5653 "m3 swigluoai and step35 clamp are different archs"
5654 );
5655 return e.swigluoai_mul_scaled(
5656 gate,
5657 up,
5658 gs,
5659 us,
5660 m3.swiglu_alpha,
5661 m3.swiglu_limit,
5662 act,
5663 n,
5664 );
5665 }
5666 if let Some(l) = limit {
5667 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5668 }
5669 if gs == 1.0 && us == 1.0 {
5670 return e.silu_mul(gate, up, act, n);
5671 }
5672 e.silu_mul_scaled(gate, up, gs, us, act, n)
5673 }
5674
5675 fn moe_route(
5681 e: &Engine,
5682 logits: &CudaSlice<f32>,
5683 t: usize,
5684 n_expert: usize,
5685 n_used: usize,
5686 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5687 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5688 }
5689
5690 #[allow(clippy::too_many_arguments)]
5698 fn moe_route_sigmoid_cfg(
5699 e: &Engine,
5700 logits: &CudaSlice<f32>,
5701 t: usize,
5702 n_expert: usize,
5703 n_used: usize,
5704 m: &MoeWeights,
5705 (sf, route_norm): (f32, bool),
5706 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5707 if sigmoid_router_enabled() {
5708 return e.moe_router_sigmoid_topk_host(
5709 logits,
5710 t,
5711 n_expert,
5712 n_used,
5713 m.active_count(),
5714 &m.exp_probs_b_dev,
5715 &m.active_experts_dev,
5716 sf,
5717 route_norm,
5718 );
5719 }
5720 let lg = e.dtoh(logits)?;
5721 Self::moe_route_sigmoid_host(
5722 &lg,
5723 t,
5724 n_expert,
5725 n_used,
5726 m.exp_probs_b.as_deref(),
5727 sf,
5728 route_norm,
5729 m.active_experts.as_deref(),
5730 )
5731 }
5732
5733 fn moe_route_cfg(
5736 e: &Engine,
5737 logits: &CudaSlice<f32>,
5738 t: usize,
5739 n_expert: usize,
5740 n_used: usize,
5741 active: Option<&[bool]>,
5742 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5743 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5746 return e.moe_router_topk_host(logits, t, n_expert, n_used);
5747 }
5748 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
5751 let mut w_out = vec![0f32; t * n_used];
5752 for tok in 0..t {
5753 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5754 let maxl = row
5756 .iter()
5757 .enumerate()
5758 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5759 .map(|(_, &x)| x)
5760 .fold(f32::NEG_INFINITY, f32::max);
5761 let mut probs = vec![0f32; n_expert];
5762 let mut den = 0f32;
5763 for i in 0..n_expert {
5764 if active.is_some_and(|mask| !mask[i]) {
5765 continue;
5766 }
5767 let x = (row[i] - maxl).exp();
5768 probs[i] = x;
5769 den += x;
5770 }
5771 for p in probs.iter_mut() {
5772 *p /= den;
5773 }
5774 let mut idx: Vec<usize> = (0..n_expert)
5776 .filter(|&i| active.is_none_or(|mask| mask[i]))
5777 .collect();
5778 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5779 let sl = &idx[..n_used];
5780 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5781 let mut ws: f32 = wv.iter().sum();
5782 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() {
5784 *x /= ws;
5785 }
5786 for j in 0..n_used {
5787 sel[tok * n_used + j] = sl[j] as u32;
5788 w_out[tok * n_used + j] = wv[j];
5789 }
5790 }
5791 Ok((sel, w_out))
5792 }
5793
5794 #[allow(clippy::too_many_arguments)]
5795 fn moe_route_sigmoid_with_input(
5796 e: &Engine,
5797 logits: &CudaSlice<f32>,
5798 input: &CudaSlice<f32>,
5799 t: usize,
5800 n_expert: usize,
5801 n_used: usize,
5802 bias: Option<&[f32]>,
5803 (sf, route_norm): (f32, bool),
5804 active: Option<&[bool]>,
5805 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5806 let (lg, input) = e.dtoh_pair(logits, input)?;
5807 let (sel, w) =
5808 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5809 Ok((sel, w, input))
5810 }
5811
5812 pub fn start_moe_prefetch_predictor(
5817 &self,
5818 e: &Engine,
5819 cfg: &ModelConfig,
5820 ) -> Result<(), Box<dyn std::error::Error>> {
5821 use crate::hybrid::Ffn;
5822 let Some(sig) = cfg.sigmoid_router() else {
5823 return Err("prefetch predictor requires a sigmoid-router arch".into());
5824 };
5825 let resident: std::collections::HashSet<(u16, u8, u16)> = e
5826 .export_moe_residency()
5827 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5828 .into_iter()
5829 .collect();
5830 let mut layers = Vec::new();
5831 for (index, layer) in self.layers.iter().enumerate() {
5832 let Ffn::Moe(m) = &layer.ffn else { continue };
5833 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5834 continue;
5835 };
5836 let router = e.dtoh(data)?;
5837 let n_expert = m.gate_exps.n_expert;
5838 let n_embd = m.gate_exps.in_f;
5839 if router.len() != n_embd * n_expert {
5840 continue;
5841 }
5842 let build = |exps: &crate::model::HostExps| {
5843 (0..n_expert)
5844 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5845 .collect::<Vec<_>>()
5846 };
5847 layers.push((
5848 index as u16,
5849 crate::cpu_experts::PredictLayerInit {
5850 router,
5851 bias: m.exp_probs_b.clone(),
5852 active: m.active_experts.clone(),
5853 n_embd,
5854 n_used: cfg
5855 .moe
5856 .as_ref()
5857 .map(|moe| moe.expert_used_count as usize)
5858 .ok_or("prefetch predictor requires MoE config")?,
5859 sig,
5860 weights_n_expert: n_expert,
5861 gate: build(&m.gate_exps),
5862 up: build(&m.up_exps),
5863 down: build(&m.down_exps),
5864 },
5865 ));
5866 }
5867 crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5868 }
5869
5870 #[allow(clippy::too_many_arguments)]
5873 pub fn moe_route_sigmoid_host_public(
5874 logits: &[f32],
5875 t: usize,
5876 n_expert: usize,
5877 n_used: usize,
5878 bias: Option<&[f32]>,
5879 sf: f32,
5880 route_norm: bool,
5881 active: Option<&[bool]>,
5882 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5883 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
5884 }
5885
5886 #[allow(clippy::too_many_arguments)]
5887 fn moe_route_sigmoid_host(
5888 lg: &[f32],
5889 t: usize,
5890 n_expert: usize,
5891 n_used: usize,
5892 bias: Option<&[f32]>,
5893 sf: f32,
5894 route_norm: bool,
5895 active: Option<&[bool]>,
5896 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5897 let active_count = active
5898 .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
5899 .unwrap_or(n_expert);
5900 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5901 if lg.len() != t * n_expert {
5902 return Err(format!(
5903 "sigmoid router logits length mismatch: got {}, expected {}",
5904 lg.len(),
5905 t * n_expert,
5906 )
5907 .into());
5908 }
5909 let mut sel = vec![0u32; t * n_used];
5910 let mut w_out = vec![0f32; t * n_used];
5911 for tok in 0..t {
5912 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5913 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
5914 let selsc: Vec<f32> = match bias {
5916 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
5917 None => scores.clone(),
5918 };
5919 let mut idx: Vec<usize> = (0..n_expert)
5920 .filter(|&i| active.is_none_or(|mask| mask[i]))
5921 .collect();
5922 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
5923 let sl = &idx[..n_used];
5924 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
5925 if route_norm {
5926 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
5927 for x in wv.iter_mut() {
5928 *x = *x / ws * sf;
5929 }
5930 } else {
5931 for x in wv.iter_mut() {
5932 *x *= sf;
5933 }
5934 }
5935 for j in 0..n_used {
5936 sel[tok * n_used + j] = sl[j] as u32;
5937 w_out[tok * n_used + j] = wv[j];
5938 }
5939 }
5940 Ok((sel, w_out))
5941 }
5942
5943 #[allow(clippy::too_many_arguments)]
5947 fn moe_ffn_sigmoid_dev(
5948 e: &Engine,
5949 m: &MoeWeights,
5950 z: &CudaSlice<f32>,
5951 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
5952 logits: &CudaSlice<f32>,
5953 t: usize,
5954 cfg: &ModelConfig,
5955 il: u16,
5956 (scaling_factor, route_norm): (f32, bool),
5957 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5958 let moe = cfg.moe.as_ref().unwrap();
5959 let n_embd = cfg.n_embd as usize;
5960 let n_expert = moe.expert_count as usize;
5961 let n_used = moe.expert_used_count as usize;
5962 let n_ff_exp = moe.expert_ff_length as usize;
5963 let dev = m.dev_exps.as_ref().unwrap();
5964 debug_assert!(cfg.step35.is_some());
5965 debug_assert_eq!(dev.dev, e.ctx().ordinal());
5966 debug_assert!(m.has_uniform_expert_layout());
5967 debug_assert!(!m.has_macros);
5968
5969 let (sel_d, w_d) = e.moe_router_sigmoid_topk(
5970 logits,
5971 t,
5972 n_expert,
5973 n_used,
5974 m.active_count(),
5975 &m.exp_probs_b_dev,
5976 &m.active_experts_dev,
5977 scaling_factor,
5978 route_norm,
5979 )?;
5980 crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
5981 let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
5982 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5983 (combined, combined)
5984 } else {
5985 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5986 };
5987 let (zq, zd) = match (t, zq8) {
5988 (1, Some((q, d))) => (q.clone(), d.clone()),
5989 _ => e.quantize_q8_1(z, t, n_embd)?,
5990 };
5991 let n_pairs = t * n_used;
5992 let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
5993 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5997 let pair_tok_d = e.htod_i32(&pair_tok)?;
5998 let gate = e.moe_pairs_matvec_q8(
5999 &dev.ptr_row,
6000 0,
6001 &pair_tok_d,
6002 &sel_d,
6003 &zq,
6004 &zd,
6005 n_embd,
6006 n_ff_exp,
6007 n_expert,
6008 n_pairs,
6009 m.gate_exps.qtype,
6010 gate_row_bytes,
6011 )?;
6012 let up = e.moe_pairs_matvec_q8(
6013 &dev.ptr_row,
6014 1,
6015 &pair_tok_d,
6016 &sel_d,
6017 &zq,
6018 &zd,
6019 n_embd,
6020 n_ff_exp,
6021 n_expert,
6022 n_pairs,
6023 m.up_exps.qtype,
6024 up_row_bytes,
6025 )?;
6026 let mut act = e.uninit(n_pairs * n_ff_exp)?;
6027 Self::ffn_act_lim(
6028 e,
6029 cfg,
6030 &gate,
6031 &up,
6032 1.0,
6033 1.0,
6034 cfg.clamp_exp_at(il as u32),
6035 &mut act,
6036 n_pairs * n_ff_exp,
6037 )?;
6038 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6039 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6040 let pair_self_d = e.htod_i32(&pair_self)?;
6041 let down = e.moe_pairs_matvec_q8(
6042 &dev.ptr_row,
6043 2,
6044 &pair_self_d,
6045 &sel_d,
6046 &aq2,
6047 &ad2,
6048 n_ff_exp,
6049 n_embd,
6050 n_expert,
6051 n_pairs,
6052 m.down_exps.qtype,
6053 m.down_exps.row_bytes,
6054 )?;
6055 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6056 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6057 let tok_off_d = e.htod_i32(&tok_off)?;
6058 let tok_ids_d = e.htod_i32(&tok_ids)?;
6059 let mut output = e.uninit(t * n_embd)?;
6060 e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
6061 output
6062 } else {
6063 let act = e.moe_gate_up_silu8_dev_q8_rows(
6064 &dev.ptr_row,
6065 &sel_d,
6066 &zq,
6067 &zd,
6068 t,
6069 n_embd,
6070 n_ff_exp,
6071 n_used,
6072 n_expert,
6073 m.gate_exps.qtype,
6074 m.up_exps.qtype,
6075 gate_row_bytes,
6076 up_row_bytes,
6077 &m.dev_macros,
6078 )?;
6079 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6080 let mut output = e.uninit(t * n_embd)?;
6081 e.moe_down8_fma_dev_q8_rows_g(
6082 &dev.ptr_row,
6083 &sel_d,
6084 &w_d,
6085 &aq2,
6086 &ad2,
6087 &mut output,
6088 t,
6089 n_ff_exp,
6090 n_embd,
6091 n_used,
6092 n_expert,
6093 m.down_exps.qtype,
6094 m.down_exps.row_bytes,
6095 )?;
6096 output
6097 };
6098
6099 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6100 eprintln!(
6101 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6102 cfg.clamp_exp_at(il as u32).is_some(),
6103 dev.gu_il,
6104 );
6105 }
6106 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6107 Ok(moe_out)
6108 }
6109
6110 fn moe_ffn_pairs(
6119 e: &Engine,
6120 m: &MoeWeights,
6121 z: &CudaSlice<f32>,
6122 logits: &CudaSlice<f32>,
6123 t: usize,
6124 cfg: &ModelConfig,
6125 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6126 let moe = cfg.moe.as_ref().unwrap();
6127 let n_embd = cfg.n_embd as usize;
6128 let n_expert = moe.expert_count as usize;
6129 let n_used = moe.expert_used_count as usize;
6130 let n_ff_exp = moe.expert_ff_length as usize;
6131 debug_assert!(
6136 !cfg.swiglu_clamped_anywhere(),
6137 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6138 );
6139 let dev = m.dev_exps.as_ref().unwrap();
6140 let (rbg_d, rbu_d) = if dev.gu_il {
6142 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6143 (sxx, sxx)
6144 } else {
6145 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6146 };
6147
6148 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6149 let n_pairs = t * n_used;
6150 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6153 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6154 let pair_w: Vec<f32> = w_all.clone();
6155 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6156 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6157 let pt = e.htod_i32(&pair_tok)?;
6158 let px = e.htod_i32(&pair_ex)?;
6159 let pw = e.htod(&pair_w)?;
6160 let toff = e.htod_i32(&tok_off)?;
6161 let tids = e.htod_i32(&tok_ids)?;
6162
6163 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6167 for p in 0..n_pairs {
6168 by_ex[pair_ex[p] as usize].push(p as i32);
6169 }
6170 let mut ex_ids: Vec<i32> = Vec::new();
6171 let mut ex_off: Vec<i32> = vec![0];
6172 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6173 for (ex, list) in by_ex.iter().enumerate() {
6174 if list.is_empty() {
6175 continue;
6176 }
6177 ex_ids.push(ex as i32);
6178 ex_pairs.extend_from_slice(list);
6179 ex_off.push(ex_pairs.len() as i32);
6180 }
6181 let n_active = ex_ids.len();
6182 let exi = e.htod_i32(&ex_ids)?;
6183 let exo = e.htod_i32(&ex_off)?;
6184 let exp_d = e.htod_i32(&ex_pairs)?;
6185 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6206 let mma_t = *MMA_T.get_or_init(|| {
6207 std::env::var("MEMRA_MOE_MMA_T")
6208 .ok()
6209 .and_then(|v| v.parse().ok())
6210 .unwrap_or(16)
6211 });
6212 let use_mma = std::env::var("MEMRA_MOE_MMA")
6213 .map(|v| v != "0")
6214 .unwrap_or(true)
6215 && t >= mma_t
6216 && q8_expert_dec_supported(m.gate_exps.qtype)
6217 && q8_expert_dec_supported(m.up_exps.qtype)
6218 && q8_expert_dec_supported(m.down_exps.qtype)
6219 && n_embd % 256 == 0
6220 && n_ff_exp % 256 == 0;
6221 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6237 && q8_expert_dec_supported(m.up_exps.qtype)
6238 && q8_expert_dec_supported(m.down_exps.qtype)
6239 && n_embd % 256 == 0
6240 && n_ff_exp % 256 == 0;
6241 let f16g_mode = crate::moe_f16g_mode();
6242 let f16g = f16g_mode != 0
6243 && t >= mma_t
6244 && (f16g_mode != 3 || !mma_capable)
6245 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6246 && f16g_proj_ok(m.up_exps.qtype, n_embd)
6247 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6248 if use_mma || f16g {
6249 let y_down = if f16g {
6257 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6261 let csr_tok_d = e.htod_i32(&csr_tok)?;
6262 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6263 let g_csr = e.moe_f16_grouped(
6264 &dev.ptr_row,
6265 0,
6266 n_expert,
6267 &exi,
6268 &ex_off,
6269 &exo,
6270 &z_f16,
6271 &z_s,
6272 n_embd,
6273 n_ff_exp,
6274 n_active,
6275 n_pairs,
6276 m.gate_exps.qtype,
6277 rbg_d,
6278 )?;
6279 let u_csr = e.moe_f16_grouped(
6280 &dev.ptr_row,
6281 1,
6282 n_expert,
6283 &exi,
6284 &ex_off,
6285 &exo,
6286 &z_f16,
6287 &z_s,
6288 n_embd,
6289 n_ff_exp,
6290 n_active,
6291 n_pairs,
6292 m.up_exps.qtype,
6293 rbu_d,
6294 )?;
6295 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6296 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6297 let d_csr = e.moe_f16_grouped(
6298 &dev.ptr_row,
6299 2,
6300 n_expert,
6301 &exi,
6302 &ex_off,
6303 &exo,
6304 &a_f16,
6305 &a_s,
6306 n_ff_exp,
6307 n_embd,
6308 n_active,
6309 n_pairs,
6310 m.down_exps.qtype,
6311 m.down_exps.row_bytes,
6312 )?;
6313 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6314 } else {
6315 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6317 let gate = e.mmq_iq_experts(
6318 &dev.ptr_row,
6319 0,
6320 n_expert,
6321 &exi,
6322 &exo,
6323 &exp_d,
6324 &pt,
6325 &z_scr,
6326 n_embd,
6327 n_ff_exp,
6328 n_active,
6329 n_pairs,
6330 t,
6331 m.gate_exps.qtype,
6332 rbg_d,
6333 )?;
6334 let up = e.mmq_iq_experts(
6335 &dev.ptr_row,
6336 1,
6337 n_expert,
6338 &exi,
6339 &exo,
6340 &exp_d,
6341 &pt,
6342 &z_scr,
6343 n_embd,
6344 n_ff_exp,
6345 n_active,
6346 n_pairs,
6347 t,
6348 m.up_exps.qtype,
6349 rbu_d,
6350 )?;
6351 let a_scr = if crate::moe_fuse_actq_on() {
6357 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6358 } else {
6359 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6360 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6361 };
6362 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6363 let pself = e.htod_i32(&pair_self)?;
6364 e.mmq_iq_experts(
6365 &dev.ptr_row,
6366 2,
6367 n_expert,
6368 &exi,
6369 &exo,
6370 &exp_d,
6371 &pself,
6372 &a_scr,
6373 n_ff_exp,
6374 n_embd,
6375 n_active,
6376 n_pairs,
6377 n_pairs,
6378 m.down_exps.qtype,
6379 m.down_exps.row_bytes,
6380 )?
6381 };
6382 let mut moe_out = e.uninit(t * n_embd)?;
6383 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6384 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6385 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6386 {
6387 let n_ff_sh = gate_shexp.out_features();
6388 let sg_gate = e.matmul(gate_shexp, z, t)?;
6389 let sg_up = e.matmul(up_shexp, z, t)?;
6390 let mut sa = e.uninit(t * n_ff_sh)?;
6391 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6392 let sh = e.matmul(down_shexp, &sa, t)?;
6393 let g = match &m.gate_inp_shexp {
6399 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6400 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6401 }
6402 Some(gate_inp_shexp) => {
6403 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6404 let mut g = e.uninit(t)?;
6405 e.sigmoid(&gs, &mut g, t)?;
6406 g
6407 }
6408 None => e.htod(&vec![1.0f32; t])?,
6409 };
6410 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6411 }
6412 return Ok(moe_out);
6413 }
6414
6415 let dec = std::env::var("MEMRA_MOE_DEC")
6418 .map(|v| v != "0")
6419 .unwrap_or(true);
6420 let matvec = |proj,
6421 exi: &_,
6422 exo: &_,
6423 exp_d: &_,
6424 pt: &_,
6425 aq: &_,
6426 ad: &_,
6427 inf,
6428 outf,
6429 qtype,
6430 rb|
6431 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6432 let dec = dec && q8_expert_dec_supported(qtype);
6434 if dec {
6435 e.moe_pairs_matvec_q8_dec(
6436 &dev.ptr_row,
6437 proj,
6438 exi,
6439 exo,
6440 exp_d,
6441 pt,
6442 aq,
6443 ad,
6444 inf,
6445 outf,
6446 n_expert,
6447 n_active,
6448 n_pairs,
6449 qtype,
6450 rb,
6451 )
6452 } else {
6453 e.moe_pairs_matvec_q8_em(
6454 &dev.ptr_row,
6455 proj,
6456 exi,
6457 exo,
6458 exp_d,
6459 pt,
6460 aq,
6461 ad,
6462 inf,
6463 outf,
6464 n_expert,
6465 n_active,
6466 n_pairs,
6467 qtype,
6468 rb,
6469 )
6470 }
6471 };
6472 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6473 let gate = matvec(
6474 0,
6475 &exi,
6476 &exo,
6477 &exp_d,
6478 &pt,
6479 &zq,
6480 &zd,
6481 n_embd,
6482 n_ff_exp,
6483 m.gate_exps.qtype,
6484 rbg_d,
6485 )?;
6486 let up = matvec(
6487 1,
6488 &exi,
6489 &exo,
6490 &exp_d,
6491 &pt,
6492 &zq,
6493 &zd,
6494 n_embd,
6495 n_ff_exp,
6496 m.up_exps.qtype,
6497 rbu_d,
6498 )?;
6499 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6500 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6501 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6503 let pself = e.htod_i32(&pair_self)?;
6504 let y_down = matvec(
6505 2,
6506 &exi,
6507 &exo,
6508 &exp_d,
6509 &pself,
6510 &aq2,
6511 &ad2,
6512 n_ff_exp,
6513 n_embd,
6514 m.down_exps.qtype,
6515 m.down_exps.row_bytes,
6516 )?;
6517 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6519
6520 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6524 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6525 {
6526 let n_ff_sh = gate_shexp.out_features();
6527 let step_exact = cfg.step35.is_some();
6531 let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6532 let (sg_gate, sg_up) = if step_exact && t == 1 {
6533 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6534 Some(pair) => pair,
6535 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6536 }
6537 } else if verify_t {
6538 let mut fused = None;
6539 if crate::spec::spec_fused_t()
6540 && (2..=4).contains(&t)
6541 && e.uses_q8_1_fast(gate_shexp)
6542 && e.uses_q8_1_fast(up_shexp)
6543 {
6544 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6545 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6546 }
6547 match fused {
6548 Some(pair) => pair,
6549 None => (
6550 e.matmul_decode_exact(gate_shexp, z, t)?,
6551 e.matmul_decode_exact(up_shexp, z, t)?,
6552 ),
6553 }
6554 } else {
6555 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6556 };
6557 let mut sa = e.uninit(t * n_ff_sh)?;
6558 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6559 let sh = if verify_t {
6560 e.matmul_decode_exact(down_shexp, &sa, t)?
6561 } else {
6562 e.matmul(down_shexp, &sa, t)?
6563 };
6564 let g = match &m.gate_inp_shexp {
6569 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6570 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6571 }
6572 Some(gate_inp_shexp) => {
6573 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6574 let mut g = e.uninit(t)?;
6575 e.sigmoid(&gs, &mut g, t)?;
6576 g
6577 }
6578 None => e.htod(&vec![1.0f32; t])?,
6579 };
6580 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6581 }
6582 Ok(moe_out)
6583 }
6584
6585 #[allow(clippy::too_many_arguments)]
6587 #[allow(clippy::too_many_arguments)]
6588 fn moe_ffn_dev(
6589 e: &Engine,
6590 m: &MoeWeights,
6591 z: &CudaSlice<f32>,
6592 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6593 logits: &CudaSlice<f32>,
6594 t: usize,
6595 cfg: &ModelConfig,
6596 il: u16,
6597 max_block: usize,
6598 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6599 let moe = cfg.moe.as_ref().unwrap();
6600 let n_embd = cfg.n_embd as usize;
6601 let n_expert = moe.expert_count as usize;
6602 let n_used = moe.expert_used_count as usize;
6603 let n_ff_exp = moe.expert_ff_length as usize;
6604 debug_assert!(
6608 cfg.sigmoid_router().is_none(),
6609 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6610 );
6611 debug_assert!(
6612 !cfg.swiglu_clamped_at(il as u32),
6613 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6614 );
6615
6616 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6618 if m.has_macros {
6621 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6622 }
6623
6624 let mut moe_out = e.uninit(t * n_embd)?;
6626
6627 if let Some(dev) = m.dev_exps.as_ref() {
6630 let (rbg_d, rbu_d) = if dev.gu_il {
6633 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6634 (sxx, sxx)
6635 } else {
6636 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6637 };
6638 let q8 = moe_q8_enabled()
6639 && q8_expert_supported(m.gate_exps.qtype)
6640 && q8_expert_supported(m.up_exps.qtype)
6641 && q8_expert_supported(m.down_exps.qtype);
6642 let rows_arm = q8
6651 && t > 1
6652 && crate::spec::spec_m2()
6653 && n_ff_exp == 512
6654 && n_used <= 8
6655 && std::env::var("MEMRA_MOE_DEVQ8_GU")
6656 .map(|v| v.is_empty() || v == "v")
6657 .unwrap_or(true)
6658 && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6659 .map(|v| v.is_empty() || v == "w8h2v")
6660 .unwrap_or(true);
6661 let csr_mode = std::env::var("MEMRA_MOE_CSR")
6670 .ok()
6671 .and_then(|v| v.parse::<i32>().ok())
6672 .unwrap_or(1);
6673 let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6674 let csr_arm = rows_arm
6675 && csr_mode > 0
6676 && t <= 10
6677 && csr_qt(m.gate_exps.qtype)
6678 && csr_qt(m.up_exps.qtype)
6679 && csr_qt(m.down_exps.qtype);
6680 if csr_arm {
6681 if csr_mode == 2 {
6682 static ENGAGED: std::sync::Once = std::sync::Once::new();
6683 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6684 }
6685 let n_pairs = t * n_used;
6686 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6687 let act = e.moe_gate_up_silu8_dev_q8_csr(
6688 &dev.ptr_row,
6689 &sel_d,
6690 &zq,
6691 &zd,
6692 n_pairs,
6693 n_embd,
6694 n_ff_exp,
6695 n_used,
6696 n_expert,
6697 m.gate_exps.qtype,
6698 m.up_exps.qtype,
6699 rbg_d,
6700 rbu_d,
6701 )?;
6702 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6703 e.moe_down8_fma_dev_q8_rows(
6707 &dev.ptr_row,
6708 &sel_d,
6709 &w_d,
6710 &aq2,
6711 &ad2,
6712 &mut moe_out,
6713 t,
6714 n_ff_exp,
6715 n_embd,
6716 n_used,
6717 n_expert,
6718 m.down_exps.qtype,
6719 m.down_exps.row_bytes,
6720 )?;
6721 if csr_mode == 2 {
6722 let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6724 &dev.ptr_row,
6725 &sel_d,
6726 &zq,
6727 &zd,
6728 t,
6729 n_embd,
6730 n_ff_exp,
6731 n_used,
6732 n_expert,
6733 m.gate_exps.qtype,
6734 m.up_exps.qtype,
6735 rbg_d,
6736 rbu_d,
6737 &m.dev_macros,
6738 )?;
6739 let mut out_r = e.uninit(t * n_embd)?;
6740 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6741 e.moe_down8_fma_dev_q8_rows(
6742 &dev.ptr_row,
6743 &sel_d,
6744 &w_d,
6745 &aq2r,
6746 &ad2r,
6747 &mut out_r,
6748 t,
6749 n_ff_exp,
6750 n_embd,
6751 n_used,
6752 n_expert,
6753 m.down_exps.qtype,
6754 m.down_exps.row_bytes,
6755 )?;
6756 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6757 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6758 let ba = a1
6759 .iter()
6760 .zip(&a2)
6761 .filter(|(x, y)| x.to_bits() != y.to_bits())
6762 .count();
6763 let bo = o1
6764 .iter()
6765 .zip(&o2)
6766 .filter(|(x, y)| x.to_bits() != y.to_bits())
6767 .count();
6768 if ba + bo > 0 {
6769 eprintln!(
6770 "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6771 a1.len(),
6772 o1.len()
6773 );
6774 let sel_h = e.dtoh_i32(&sel_d)?;
6776 let mut shown = 0;
6777 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6778 if x.to_bits() != y.to_bits() && shown < 4 {
6779 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6780 let ex = sel_h[p];
6781 let npx = sel_h.iter().filter(|&&v| v == ex).count();
6782 eprintln!(
6783 " ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6784 );
6785 shown += 1;
6786 }
6787 }
6788 std::process::exit(3);
6789 }
6790 }
6791 } else if rows_arm {
6792 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6795 use std::sync::atomic::{AtomicU64, Ordering};
6796 static PAIRS: AtomicU64 = AtomicU64::new(0);
6797 static UNIQ: AtomicU64 = AtomicU64::new(0);
6798 static CALLS: AtomicU64 = AtomicU64::new(0);
6799 let sel_h = e.dtoh_i32(&sel_d)?;
6800 let mut u: Vec<i32> = sel_h.clone();
6801 u.sort_unstable();
6802 u.dedup();
6803 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6804 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6805 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6806 if c % 480 == 0 {
6807 let p = PAIRS.load(Ordering::Relaxed);
6808 let q = UNIQ.load(Ordering::Relaxed);
6809 eprintln!(
6810 "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6811 q as f64 / p as f64
6812 );
6813 }
6814 }
6815 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6816 let act = e.moe_gate_up_silu8_dev_q8_rows(
6817 &dev.ptr_row,
6818 &sel_d,
6819 &zq,
6820 &zd,
6821 t,
6822 n_embd,
6823 n_ff_exp,
6824 n_used,
6825 n_expert,
6826 m.gate_exps.qtype,
6827 m.up_exps.qtype,
6828 rbg_d,
6829 rbu_d,
6830 &m.dev_macros,
6831 )?;
6832 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6833 e.moe_down8_fma_dev_q8_rows(
6834 &dev.ptr_row,
6835 &sel_d,
6836 &w_d,
6837 &aq2,
6838 &ad2,
6839 &mut moe_out,
6840 t,
6841 n_ff_exp,
6842 n_embd,
6843 n_used,
6844 n_expert,
6845 m.down_exps.qtype,
6846 m.down_exps.row_bytes,
6847 )?;
6848 } else {
6849 for tok in 0..t {
6850 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6851 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6852 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6853 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6854 if q8 {
6855 let (zq, zd) = match (t, zq8) {
6856 (1, Some((q, d))) => (q.clone(), d.clone()),
6857 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6858 };
6859 let act = e.moe_gate_up_silu8_dev_q8(
6860 &dev.ptr_row,
6861 &selt,
6862 &zq,
6863 &zd,
6864 n_embd,
6865 n_ff_exp,
6866 n_used,
6867 n_expert,
6868 m.gate_exps.qtype,
6869 m.up_exps.qtype,
6870 rbg_d,
6871 rbu_d,
6872 &m.dev_macros,
6873 )?;
6874 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6875 e.moe_down8_fma_dev_q8(
6876 &dev.ptr_row,
6877 &selt,
6878 &wt,
6879 &aq2,
6880 &ad2,
6881 &mut dst,
6882 n_ff_exp,
6883 n_embd,
6884 n_used,
6885 n_expert,
6886 m.down_exps.qtype,
6887 m.down_exps.row_bytes,
6888 )?;
6889 } else {
6890 let act = e.moe_gate_up_silu8_dev(
6891 &dev.ptr_row,
6892 &selt,
6893 &zt,
6894 n_embd,
6895 n_ff_exp,
6896 n_used,
6897 n_expert,
6898 m.gate_exps.qtype,
6899 m.up_exps.qtype,
6900 rbg_d,
6901 rbu_d,
6902 &m.dev_macros,
6903 )?;
6904 e.moe_down8_fma_dev(
6905 &dev.ptr_row,
6906 &selt,
6907 &wt,
6908 &act,
6909 &mut dst,
6910 n_ff_exp,
6911 n_embd,
6912 n_used,
6913 n_expert,
6914 m.down_exps.qtype,
6915 m.down_exps.row_bytes,
6916 )?;
6917 }
6918 }
6919 }
6920 } else {
6921 let q8 = moe_q8_enabled()
6928 && q8_expert_supported(m.gate_exps.qtype)
6929 && q8_expert_supported(m.up_exps.qtype)
6930 && q8_expert_supported(m.down_exps.qtype);
6931 e.with_moe_cache(max_block, |c, eng| {
6932 let row = c
6933 .layer_dev_row(il, n_expert, eng)?
6934 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
6935 for tok in 0..t {
6936 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6937 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6938 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6939 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6940 if q8 {
6941 let (zq, zd) = match (t, zq8) {
6942 (1, Some((q, d))) => (q.clone(), d.clone()),
6943 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
6944 };
6945 let act = eng.moe_gate_up_silu8_dev_q8(
6946 row,
6947 &selt,
6948 &zq,
6949 &zd,
6950 n_embd,
6951 n_ff_exp,
6952 n_used,
6953 n_expert,
6954 m.gate_exps.qtype,
6955 m.up_exps.qtype,
6956 m.gate_exps.row_bytes,
6957 m.up_exps.row_bytes,
6958 &m.dev_macros,
6959 )?;
6960 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
6961 eng.moe_down8_fma_dev_q8(
6962 row,
6963 &selt,
6964 &wt,
6965 &aq2,
6966 &ad2,
6967 &mut dst,
6968 n_ff_exp,
6969 n_embd,
6970 n_used,
6971 n_expert,
6972 m.down_exps.qtype,
6973 m.down_exps.row_bytes,
6974 )?;
6975 } else {
6976 let act = eng.moe_gate_up_silu8_dev(
6977 row,
6978 &selt,
6979 &zt,
6980 n_embd,
6981 n_ff_exp,
6982 n_used,
6983 n_expert,
6984 m.gate_exps.qtype,
6985 m.up_exps.qtype,
6986 m.gate_exps.row_bytes,
6987 m.up_exps.row_bytes,
6988 &m.dev_macros,
6989 )?;
6990 eng.moe_down8_fma_dev(
6991 row,
6992 &selt,
6993 &wt,
6994 &act,
6995 &mut dst,
6996 n_ff_exp,
6997 n_embd,
6998 n_used,
6999 n_expert,
7000 m.down_exps.qtype,
7001 m.down_exps.row_bytes,
7002 )?;
7003 }
7004 }
7005 c.hits += (t * 3 * n_used) as u64;
7007 Ok(())
7008 })?;
7009 }
7010
7011 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7016 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7017 {
7018 let n_ff_sh = gate_shexp.out_features();
7019 let verify_t = t > 1 && t < PRIME_MIN_T;
7022 let (sg_gate, sg_up) = if t == 1 {
7023 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
7024 Some(pair) => pair,
7025 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
7026 }
7027 } else if verify_t {
7028 let mut fused = None;
7032 if crate::spec::spec_fused_t()
7033 && (2..=4).contains(&t)
7034 && e.uses_q8_1_fast(gate_shexp)
7035 && e.uses_q8_1_fast(up_shexp)
7036 {
7037 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7038 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
7039 }
7040 match fused {
7041 Some(pair) => pair,
7042 None => (
7043 e.matmul_decode_exact(gate_shexp, z, t)?,
7044 e.matmul_decode_exact(up_shexp, z, t)?,
7045 ),
7046 }
7047 } else {
7048 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
7049 };
7050 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
7052 let sh = if verify_t {
7053 e.matmul_decode_exact(down_shexp, &sa, t)?
7054 } else {
7055 e.matmul(down_shexp, &sa, t)?
7056 };
7057 let g = match &m.gate_inp_shexp {
7061 Some(gate_inp_shexp) => {
7062 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7065 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7066 } else {
7067 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7068 let mut g = e.uninit(t)?;
7069 e.sigmoid(&gs, &mut g, t)?;
7070 g
7071 }
7072 }
7073 None => e.htod(&vec![1.0f32; t])?,
7074 };
7075 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7076 }
7077
7078 Ok(moe_out)
7079 }
7080
7081 #[allow(clippy::too_many_arguments)]
7091 #[allow(clippy::too_many_arguments)]
7094 fn moe_gdec_token_q8(
7095 e: &Engine,
7096 m: &MoeWeights,
7097 il: u16,
7098 max_block: usize,
7099 zq: &CudaSlice<i8>,
7100 zd: &CudaSlice<f32>,
7101 sel: &[u32],
7102 w: &[f32],
7103 moe_out: &mut CudaSlice<f32>,
7104 tok: usize,
7105 n_embd: usize,
7106 n_ff_exp: usize,
7107 n_used: usize,
7108 ) -> Result<bool, Box<dyn std::error::Error>> {
7109 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7110 use cudarc::driver::DevicePtr;
7111 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7112 let mut g = [0u64; 8];
7113 let mut u = [0u64; 8];
7114 let mut d = [0u64; 8];
7115 for (j, &ex) in sel.iter().enumerate() {
7116 let ex = ex as u16;
7117 let (Some(sg), Some(su), Some(sd)) = (
7118 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7119 c.resident(BlockId::new(il, PROJ_UP, ex)),
7120 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7121 ) else {
7122 return Ok(None);
7123 };
7124 let __s = eng.stream();
7125 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7126 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7127 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7128 g[j] = pg as u64;
7129 u[j] = pu as u64;
7130 d[j] = pd as u64;
7131 }
7132 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7133 for &ex in sel {
7134 let ex = ex as u16;
7135 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7136 c.note_profile_hit(BlockId::new(il, proj, ex));
7137 }
7138 }
7139 }
7140 c.hits += (3 * n_used) as u64;
7141 Ok(Some((g, u, d)))
7142 })?;
7143 let Some((g, u, d)) = ptrs else {
7144 return Ok(false);
7145 };
7146 let mut wv = [0f32; 8];
7147 wv[..n_used].copy_from_slice(w);
7148 let act = e.moe_gate_up_silu8_q8(
7149 crate::WPtr8(g),
7150 crate::WPtr8(u),
7151 zq,
7152 zd,
7153 n_embd,
7154 n_ff_exp,
7155 n_used,
7156 m.gate_exps.qtype,
7157 m.up_exps.qtype,
7158 m.gate_exps.row_bytes,
7159 m.up_exps.row_bytes,
7160 )?;
7161 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7163 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7164 e.moe_down8_fma_q8(
7165 crate::WPtr8(d),
7166 crate::F32x8(wv),
7167 &aq2,
7168 &ad2,
7169 &mut dst,
7170 n_ff_exp,
7171 n_embd,
7172 n_used,
7173 m.down_exps.qtype,
7174 m.down_exps.row_bytes,
7175 )?;
7176 Ok(true)
7177 }
7178
7179 fn moe_gdec_token(
7180 e: &Engine,
7181 m: &MoeWeights,
7182 il: u16,
7183 max_block: usize,
7184 zt: &cudarc::driver::CudaView<f32>,
7185 sel: &[u32],
7186 w: &[f32],
7187 moe_out: &mut CudaSlice<f32>,
7188 tok: usize,
7189 n_embd: usize,
7190 n_ff_exp: usize,
7191 n_used: usize,
7192 ) -> Result<bool, Box<dyn std::error::Error>> {
7193 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7194 use cudarc::driver::DevicePtr;
7195 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7197 let mut g = [0u64; 8];
7198 let mut u = [0u64; 8];
7199 let mut d = [0u64; 8];
7200 for (j, &ex) in sel.iter().enumerate() {
7201 let ex = ex as u16;
7202 let (Some(sg), Some(su), Some(sd)) = (
7203 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7204 c.resident(BlockId::new(il, PROJ_UP, ex)),
7205 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7206 ) else {
7207 return Ok(None);
7208 };
7209 let __s = eng.stream();
7210 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7211 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7212 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7213 g[j] = pg as u64;
7214 u[j] = pu as u64;
7215 d[j] = pd as u64;
7216 }
7217 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7218 for &ex in sel {
7219 let ex = ex as u16;
7220 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7221 c.note_profile_hit(BlockId::new(il, proj, ex));
7222 }
7223 }
7224 }
7225 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
7227 })?;
7228 let Some((g, u, d)) = ptrs else {
7229 return Ok(false);
7230 };
7231 let mut wv = [0f32; 8];
7232 wv[..n_used].copy_from_slice(w);
7233 let act = e.moe_gate_up_silu8(
7235 crate::WPtr8(g),
7236 crate::WPtr8(u),
7237 zt,
7238 n_embd,
7239 n_ff_exp,
7240 n_used,
7241 m.gate_exps.qtype,
7242 m.up_exps.qtype,
7243 m.gate_exps.row_bytes,
7244 m.up_exps.row_bytes,
7245 )?;
7246 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7247 e.moe_down8_fma_into(
7248 crate::WPtr8(d),
7249 crate::F32x8(wv),
7250 &act,
7251 &mut dst,
7252 n_ff_exp,
7253 n_embd,
7254 n_used,
7255 m.down_exps.qtype,
7256 m.down_exps.row_bytes,
7257 )?;
7258 Ok(true)
7259 }
7260
7261 fn moe_cached_gemm_q8(
7266 e: &Engine,
7267 il: u16,
7268 proj: u8,
7269 ex: usize,
7270 m: &MoeWeights,
7271 max_block: usize,
7272 aq: &CudaSlice<i8>,
7273 ad: &CudaSlice<f32>,
7274 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7275 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7276 let exps = match proj {
7277 PROJ_GATE => &m.gate_exps,
7278 PROJ_UP => &m.up_exps,
7279 _ => &m.down_exps,
7280 };
7281 let layout = exps.expert_layout(ex);
7282 let id = BlockId::new(il, proj, ex as u16);
7283 let source = exps.expert_source(ex);
7284 e.with_moe_cache(max_block, |c, eng| {
7285 let slot = c.dispatch_source(id, source, eng)?;
7286 let DispatchSlot::Resident(sl) = slot;
7287 let buf = c.slot(sl);
7288 eng.qmatvec_expert_q8(
7289 buf,
7290 0..layout.len,
7291 aq,
7292 ad,
7293 1,
7294 exps.in_f,
7295 exps.out_f,
7296 layout.qtype,
7297 layout.row_bytes,
7298 )
7299 })
7300 }
7301
7302 fn moe_cached_gemm(
7303 e: &Engine,
7304 il: u16,
7305 proj: u8,
7306 ex: usize,
7307 m: &MoeWeights,
7308 max_block: usize,
7309 x: &cudarc::driver::CudaView<f32>,
7310 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7311 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7312 let exps = match proj {
7313 PROJ_GATE => &m.gate_exps,
7314 PROJ_UP => &m.up_exps,
7315 _ => &m.down_exps,
7316 };
7317 let layout = exps.expert_layout(ex);
7318 let id = BlockId::new(il, proj, ex as u16);
7319 let source = exps.expert_source(ex);
7320 e.with_moe_cache(max_block, |c, eng| {
7322 let slot = c.dispatch_source(id, source, eng)?;
7323 let DispatchSlot::Resident(sl) = slot;
7326 let buf = c.slot(sl);
7327 eng.qmatvec_view(
7328 buf,
7329 0..layout.len,
7330 x,
7331 1,
7332 exps.in_f,
7333 exps.out_f,
7334 layout.qtype,
7335 layout.row_bytes,
7336 )
7337 })
7338 }
7339
7340 fn moe_profile_admit_expert(
7344 e: &Engine,
7345 il: u16,
7346 ex: usize,
7347 m: &MoeWeights,
7348 max_block: usize,
7349 ) -> Result<(), Box<dyn std::error::Error>> {
7350 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7351 e.with_moe_cache(max_block, |cache, eng| {
7352 for (proj, exps) in [
7353 (PROJ_GATE, &m.gate_exps),
7354 (PROJ_UP, &m.up_exps),
7355 (PROJ_DOWN, &m.down_exps),
7356 ] {
7357 let id = BlockId::new(il, proj, ex as u16);
7358 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7359 }
7360 Ok(())
7361 })
7362 }
7363
7364 #[allow(clippy::too_many_arguments)]
7367 fn moe_frozen_gemm(
7368 e: &Engine,
7369 il: u16,
7370 proj: u8,
7371 ex: usize,
7372 m: &MoeWeights,
7373 max_block: usize,
7374 x: &cudarc::driver::CudaView<f32>,
7375 scratch: &mut Option<CudaSlice<u8>>,
7376 scratch_len: usize,
7377 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7378 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7379 let exps = match proj {
7380 PROJ_GATE => &m.gate_exps,
7381 PROJ_UP => &m.up_exps,
7382 _ => &m.down_exps,
7383 };
7384 let layout = exps.expert_layout(ex);
7385 let id = BlockId::new(il, proj, ex as u16);
7386 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7387 let Some(slot) = cache.resident(id) else {
7388 return Ok(None);
7389 };
7390 let buf = cache.slot(slot);
7391 Ok(Some(eng.qmatvec_view(
7392 buf,
7393 0..layout.len,
7394 x,
7395 1,
7396 exps.in_f,
7397 exps.out_f,
7398 layout.qtype,
7399 layout.row_bytes,
7400 )?))
7401 })? {
7402 return Ok(output);
7403 }
7404 if scratch.is_none() {
7405 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7406 }
7407 let scratch = scratch.as_mut().unwrap();
7408 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7409 e.qmatvec_view(
7410 scratch,
7411 0..layout.len,
7412 x,
7413 1,
7414 exps.in_f,
7415 exps.out_f,
7416 layout.qtype,
7417 layout.row_bytes,
7418 )
7419 }
7420
7421 fn moe_prefetch_expert(
7422 e: &Engine,
7423 il: u16,
7424 ex: usize,
7425 m: &MoeWeights,
7426 max_block: usize,
7427 keep: &[crate::moe_cache::BlockId],
7428 ) -> Result<(), Box<dyn std::error::Error>> {
7429 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7430 e.with_moe_cache(max_block, |c, eng| {
7431 for (proj, exps) in [
7432 (PROJ_GATE, &m.gate_exps),
7433 (PROJ_UP, &m.up_exps),
7434 (PROJ_DOWN, &m.down_exps),
7435 ] {
7436 let id = BlockId::new(il, proj, ex as u16);
7437 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7438 }
7439 Ok(())
7440 })
7441 }
7442
7443 fn moe_prefetch_disk_expert(
7446 e: &Engine,
7447 il: u16,
7448 ex: usize,
7449 m: &MoeWeights,
7450 max_block: usize,
7451 keep: &[crate::moe_cache::BlockId],
7452 ) -> Result<(), Box<dyn std::error::Error>> {
7453 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7454 e.with_moe_cache(max_block, |c, eng| {
7455 for (proj, exps) in [
7456 (PROJ_GATE, &m.gate_exps),
7457 (PROJ_UP, &m.up_exps),
7458 (PROJ_DOWN, &m.down_exps),
7459 ] {
7460 let source = exps.expert_source(ex);
7461 if let crate::model::ExpertSource::Disk { .. } = &source {
7462 let id = BlockId::new(il, proj, ex as u16);
7463 let _ = c.prefetch_source(id, source, keep, eng)?;
7464 }
7465 }
7466 Ok(())
7467 })
7468 }
7469
7470 #[inline]
7471 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7472 let _ = m.gate_exps.prefetch_expert_pages(ex);
7473 let _ = m.up_exps.prefetch_expert_pages(ex);
7474 let _ = m.down_exps.prefetch_expert_pages(ex);
7475 }
7476}
7477
7478impl HybridModel {
7495 #[allow(clippy::too_many_arguments)]
7499 fn moe_ffn_grouped_resident_q8(
7500 e: &Engine,
7501 m: &MoeWeights,
7502 z: &CudaSlice<f32>,
7503 t: usize,
7504 cfg: &ModelConfig,
7505 il: u16,
7506 sel_all: &[u32],
7507 w_all: &[f32],
7508 table: &CudaSlice<u64>,
7509 gu_il: bool,
7510 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7511 let moe = cfg.moe.as_ref().unwrap();
7512 let n_embd = cfg.n_embd as usize;
7513 let n_expert = moe.expert_count as usize;
7514 let n_used = moe.expert_used_count as usize;
7515 let n_ff_exp = moe.expert_ff_length as usize;
7516 let n_pairs = t * n_used;
7517 debug_assert_eq!(sel_all.len(), n_pairs);
7518 debug_assert_eq!(w_all.len(), n_pairs);
7519 debug_assert!(
7520 m.gate_exps.macros.is_none()
7521 && m.up_exps.macros.is_none()
7522 && m.down_exps.macros.is_none(),
7523 "resident grouped q8 does not fold per-expert macro scales",
7524 );
7525
7526 if !cfg.swiglu_clamped_at(il as u32) {
7532 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7533 let sel_d = e.htod_i32(&sel)?;
7534 let w_d = e.htod(w_all)?;
7535 let (gate_row_bytes, up_row_bytes) = if gu_il {
7536 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7537 (combined, combined)
7538 } else {
7539 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7540 };
7541 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7542 let act = e.moe_gate_up_silu8_dev_q8_rows(
7543 table,
7544 &sel_d,
7545 &zq,
7546 &zd,
7547 t,
7548 n_embd,
7549 n_ff_exp,
7550 n_used,
7551 n_expert,
7552 m.gate_exps.qtype,
7553 m.up_exps.qtype,
7554 gate_row_bytes,
7555 up_row_bytes,
7556 &m.dev_macros,
7557 )?;
7558 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7559 let mut moe_out = e.uninit(t * n_embd)?;
7560 e.moe_down8_fma_dev_q8_rows_g(
7561 table,
7562 &sel_d,
7563 &w_d,
7564 &aq2,
7565 &ad2,
7566 &mut moe_out,
7567 t,
7568 n_ff_exp,
7569 n_embd,
7570 n_used,
7571 n_expert,
7572 m.down_exps.qtype,
7573 m.down_exps.row_bytes,
7574 )?;
7575
7576 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7577 let mut counts = vec![0usize; n_expert];
7578 for &expert in sel_all {
7579 counts[expert as usize] += 1;
7580 }
7581 let mut sizes: Vec<usize> =
7582 counts.into_iter().filter(|&count| count != 0).collect();
7583 sizes.sort_unstable();
7584 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7585 println!(
7586 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7587 m_e: min={} median={} mean={mean:.1} max={}",
7588 sizes.len(),
7589 n_expert,
7590 sizes.first().copied().unwrap_or(0),
7591 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7592 sizes.last().copied().unwrap_or(0),
7593 );
7594 }
7595 return Ok(moe_out);
7596 }
7597
7598 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7602 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7603 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7604 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7605
7606 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7607 for (pair, &expert) in pair_ex.iter().enumerate() {
7608 by_expert[expert as usize].push(pair as i32);
7609 }
7610
7611 let pair_tok_d = e.htod_i32(&pair_tok)?;
7612 let pair_ex_d = e.htod_i32(&pair_ex)?;
7613 let pair_w_d = e.htod(w_all)?;
7614 let tok_off_d = e.htod_i32(&tok_off)?;
7615 let tok_ids_d = e.htod_i32(&tok_ids)?;
7616
7617 let matvec = |proj: i32,
7618 pair_rows: &CudaSlice<i32>,
7619 aq: &CudaSlice<i8>,
7620 ad: &CudaSlice<f32>,
7621 in_f: usize,
7622 out_f: usize,
7623 qtype: i32,
7624 row_bytes: usize|
7625 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7626 e.moe_pairs_matvec_q8(
7627 table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7628 row_bytes,
7629 )
7630 };
7631
7632 let (gate_row_bytes, up_row_bytes) = if gu_il {
7633 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7634 (combined, combined)
7635 } else {
7636 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7637 };
7638 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7639 let gate = matvec(
7640 0,
7641 &pair_tok_d,
7642 &zq,
7643 &zd,
7644 n_embd,
7645 n_ff_exp,
7646 m.gate_exps.qtype,
7647 gate_row_bytes,
7648 )?;
7649 let up = matvec(
7650 1,
7651 &pair_tok_d,
7652 &zq,
7653 &zd,
7654 n_embd,
7655 n_ff_exp,
7656 m.up_exps.qtype,
7657 up_row_bytes,
7658 )?;
7659 let mut act = e.uninit(n_pairs * n_ff_exp)?;
7660 Self::ffn_act_lim(
7661 e,
7662 cfg,
7663 &gate,
7664 &up,
7665 1.0,
7666 1.0,
7667 cfg.clamp_exp_at(il as u32),
7668 &mut act,
7669 n_pairs * n_ff_exp,
7670 )?;
7671 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7672 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7673 let pair_self_d = e.htod_i32(&pair_self)?;
7674 let down = matvec(
7675 2,
7676 &pair_self_d,
7677 &aq2,
7678 &ad2,
7679 n_ff_exp,
7680 n_embd,
7681 m.down_exps.qtype,
7682 m.down_exps.row_bytes,
7683 )?;
7684 let mut moe_out = e.uninit(t * n_embd)?;
7685 e.moe_pairs_scatter(
7686 &down,
7687 &pair_w_d,
7688 &tok_off_d,
7689 &tok_ids_d,
7690 &mut moe_out,
7691 t,
7692 n_embd,
7693 )?;
7694
7695 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7696 let mut sizes: Vec<usize> = by_expert
7697 .iter()
7698 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7699 .collect();
7700 sizes.sort_unstable();
7701 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7702 println!(
7703 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7704 m_e: min={} median={} mean={mean:.1} max={}",
7705 sizes.len(),
7706 n_expert,
7707 sizes.first().copied().unwrap_or(0),
7708 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7709 sizes.last().copied().unwrap_or(0),
7710 );
7711 }
7712 Ok(moe_out)
7713 }
7714
7715 fn moe_ffn_grouped_add_shared(
7716 e: &Engine,
7717 m: &MoeWeights,
7718 z: &CudaSlice<f32>,
7719 t: usize,
7720 cfg: &ModelConfig,
7721 il: u16,
7722 moe_out: &mut CudaSlice<f32>,
7723 ) -> Result<(), Box<dyn std::error::Error>> {
7724 let n_embd = cfg.n_embd as usize;
7725 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7726 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7727 {
7728 let n_ff_sh = gate_shexp.out_features();
7729 let sg_gate = e.matmul(gate_shexp, z, t)?;
7730 let sg_up = e.matmul(up_shexp, z, t)?;
7731 let mut sa = e.uninit(t * n_ff_sh)?;
7732 Self::ffn_act_lim(
7733 e,
7734 cfg,
7735 &sg_gate,
7736 &sg_up,
7737 1.0,
7738 1.0,
7739 cfg.clamp_shexp_at(il as u32),
7740 &mut sa,
7741 t * n_ff_sh,
7742 )?;
7743 let sh = e.matmul(down_shexp, &sa, t)?;
7744 let gate = match &m.gate_inp_shexp {
7745 Some(gate_inp_shexp) => {
7746 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7747 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7748 } else {
7749 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7750 let mut gate = e.uninit(t)?;
7751 e.sigmoid(&raw, &mut gate, t)?;
7752 gate
7753 }
7754 }
7755 None => e.htod(&vec![1.0f32; t])?,
7756 };
7757 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7758 }
7759 Ok(())
7760 }
7761
7762 pub(crate) fn moe_ffn_grouped(
7765 e: &Engine,
7766 m: &MoeWeights,
7767 z: &CudaSlice<f32>,
7768 t: usize,
7769 cfg: &ModelConfig,
7770 il: u16,
7771 max_block: usize,
7772 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7773 let moe = cfg.moe.as_ref().unwrap();
7774 let n_embd = cfg.n_embd as usize;
7775 let n_expert = moe.expert_count as usize;
7776 let n_used = moe.expert_used_count as usize;
7777 let n_ff_exp = moe.expert_ff_length as usize;
7778 let lim_exp = cfg.clamp_exp_at(il as u32);
7780
7781 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7785 if let Some(sig) = cfg.sigmoid_router() {
7786 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7787 }
7788 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7789 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7790 } else {
7791 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7792 };
7793 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7794 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7795 Self::trace_moe_input(e, il, t, n_embd, z)?;
7796
7797 let no_exp_macros = m.gate_exps.macros.is_none()
7802 && m.up_exps.macros.is_none()
7803 && m.down_exps.macros.is_none();
7804 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7805 m.has_uniform_expert_layout()
7806 && no_exp_macros
7807 && moe_q8_enabled()
7808 && q8_expert_supported(m.gate_exps.qtype)
7809 && q8_expert_supported(m.up_exps.qtype)
7810 && q8_expert_supported(m.down_exps.qtype)
7811 && moe_slab_enabled()
7812 && dev.dev == e.ctx().ordinal()
7813 });
7814 if let Some(dev) = resident_q8 {
7815 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7816 e,
7817 m,
7818 z,
7819 t,
7820 cfg,
7821 il,
7822 &sel_all,
7823 &w_all,
7824 &dev.ptr_row,
7825 dev.gu_il,
7826 )?;
7827 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7828 return Ok(moe_out);
7829 }
7830
7831 struct ExpertGroup {
7835 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
7839 let mut groups: Vec<ExpertGroup> = (0..n_expert)
7840 .map(|_| ExpertGroup {
7841 tok_indices: Vec::new(),
7842 slot_indices: Vec::new(),
7843 weights: Vec::new(),
7844 })
7845 .collect();
7846
7847 for tok in 0..t {
7848 for j in 0..n_used {
7849 let ex = sel_all[tok * n_used + j] as usize;
7850 let w = w_all[tok * n_used + j];
7851 groups[ex].tok_indices.push(tok as i32);
7852 groups[ex].slot_indices.push(j as i32);
7853 groups[ex].weights.push(w);
7854 }
7855 }
7856
7857 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7860 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
7864 let u_len = m.up_exps.max_expert_bytes();
7865 let d_len = m.down_exps.max_expert_bytes();
7866 let moe_q8 = m.has_uniform_expert_layout()
7867 && moe_q8_enabled()
7868 && q8_expert_supported(m.gate_exps.qtype)
7869 && q8_expert_supported(m.up_exps.qtype)
7870 && q8_expert_supported(m.down_exps.qtype);
7871 let slab_local = m
7874 .dev_exps
7875 .as_ref()
7876 .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
7877 let use_cache =
7878 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
7879 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
7882
7883 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
7885 (
7886 Some(e.alloc_u8(g_len)?),
7887 Some(e.alloc_u8(u_len)?),
7888 Some(e.alloc_u8(d_len)?),
7889 )
7890 } else {
7891 (None, None, None)
7892 };
7893
7894 let mut order: Vec<usize> = (0..n_expert)
7905 .filter(|&ex| !groups[ex].tok_indices.is_empty())
7906 .collect();
7907 order.sort_by(|&a, &b| {
7908 groups[b]
7909 .tok_indices
7910 .len()
7911 .cmp(&groups[a].tok_indices.len())
7912 .then(a.cmp(&b))
7913 });
7914 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
7916 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
7917 if worker_disk_prefetch {
7918 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
7919 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
7920 }
7921 }
7922 for (order_pos, &ex) in order.iter().enumerate() {
7923 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
7924 Self::moe_prefetch_host_expert(order[next], m);
7925 }
7926 if worker_disk_prefetch {
7927 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
7928 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7929 let keep = [
7930 BlockId::new(il, PROJ_GATE, ex as u16),
7931 BlockId::new(il, PROJ_UP, ex as u16),
7932 BlockId::new(il, PROJ_DOWN, ex as u16),
7933 ];
7934 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
7935 }
7936 }
7937 let grp = &groups[ex];
7938 let m_e = grp.tok_indices.len();
7939 m_dist.push(m_e);
7940 let gl = m.gate_exps.expert_layout(ex);
7941 let ul = m.up_exps.expert_layout(ex);
7942 let dl = m.down_exps.expert_layout(ex);
7943
7944 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
7948 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
7949 let dmac = m.down_exps.macro_scale(ex);
7950 let weight_d = if dmac == 1.0 {
7951 e.htod(&grp.weights)?
7952 } else {
7953 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
7954 e.htod(&scaled)?
7955 };
7956
7957 let mut gathered = e.zeros(m_e * n_embd)?;
7959 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
7960 let gv = gathered.slice(0..m_e * n_embd);
7961
7962 let y = if let Some(dev) = slab_local {
7965 let gate_start = ex * m.gate_exps.expert_stride;
7966 let up_start = ex * m.up_exps.expert_stride;
7967 let down_start = ex * m.down_exps.expert_stride;
7968 if grouped_q8 {
7969 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7970 let gate = e.qmatvec_expert_q8(
7971 &dev.gate,
7972 gate_start..gate_start + gl.len,
7973 &zq,
7974 &zd,
7975 m_e,
7976 m.gate_exps.in_f,
7977 m.gate_exps.out_f,
7978 gl.qtype,
7979 gl.row_bytes,
7980 )?;
7981 let up = e.qmatvec_expert_q8(
7982 &dev.up,
7983 up_start..up_start + ul.len,
7984 &zq,
7985 &zd,
7986 m_e,
7987 m.up_exps.in_f,
7988 m.up_exps.out_f,
7989 ul.qtype,
7990 ul.row_bytes,
7991 )?;
7992 let mut act = e.uninit(m_e * n_ff_exp)?;
7993 Self::ffn_act_lim(
7994 e,
7995 cfg,
7996 &gate,
7997 &up,
7998 m.gate_exps.macro_scale(ex),
7999 m.up_exps.macro_scale(ex),
8000 lim_exp,
8001 &mut act,
8002 m_e * n_ff_exp,
8003 )?;
8004 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8005 e.qmatvec_expert_q8(
8006 &dev.down,
8007 down_start..down_start + dl.len,
8008 &aq2,
8009 &ad2,
8010 m_e,
8011 m.down_exps.in_f,
8012 m.down_exps.out_f,
8013 dl.qtype,
8014 dl.row_bytes,
8015 )?
8016 } else {
8017 let gate = e.qmatvec_view(
8018 &dev.gate,
8019 gate_start..gate_start + gl.len,
8020 &gv,
8021 m_e,
8022 m.gate_exps.in_f,
8023 m.gate_exps.out_f,
8024 gl.qtype,
8025 gl.row_bytes,
8026 )?;
8027 let up = e.qmatvec_view(
8028 &dev.up,
8029 up_start..up_start + ul.len,
8030 &gv,
8031 m_e,
8032 m.up_exps.in_f,
8033 m.up_exps.out_f,
8034 ul.qtype,
8035 ul.row_bytes,
8036 )?;
8037 let mut act = e.uninit(m_e * n_ff_exp)?;
8038 Self::ffn_act_lim(
8039 e,
8040 cfg,
8041 &gate,
8042 &up,
8043 m.gate_exps.macro_scale(ex),
8044 m.up_exps.macro_scale(ex),
8045 lim_exp,
8046 &mut act,
8047 m_e * n_ff_exp,
8048 )?;
8049 let actv = act.slice(0..m_e * n_ff_exp);
8050 e.qmatvec_view(
8051 &dev.down,
8052 down_start..down_start + dl.len,
8053 &actv,
8054 m_e,
8055 m.down_exps.in_f,
8056 m.down_exps.out_f,
8057 dl.qtype,
8058 dl.row_bytes,
8059 )?
8060 }
8061 } else if use_cache {
8062 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8063 if grouped_q8 {
8064 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8065 let gate = e.with_moe_cache(max_block, |cache, eng| {
8066 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8067 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8068 eng.qmatvec_expert_q8(
8069 cache.buf(slot),
8070 0..gl.len,
8071 &zq,
8072 &zd,
8073 m_e,
8074 m.gate_exps.in_f,
8075 m.gate_exps.out_f,
8076 gl.qtype,
8077 gl.row_bytes,
8078 )
8079 })?;
8080 let up = e.with_moe_cache(max_block, |cache, eng| {
8081 let id = BlockId::new(il, PROJ_UP, ex as u16);
8082 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8083 eng.qmatvec_expert_q8(
8084 cache.buf(slot),
8085 0..ul.len,
8086 &zq,
8087 &zd,
8088 m_e,
8089 m.up_exps.in_f,
8090 m.up_exps.out_f,
8091 ul.qtype,
8092 ul.row_bytes,
8093 )
8094 })?;
8095 let mut act = e.uninit(m_e * n_ff_exp)?;
8096 Self::ffn_act_lim(
8097 e,
8098 cfg,
8099 &gate,
8100 &up,
8101 m.gate_exps.macro_scale(ex),
8102 m.up_exps.macro_scale(ex),
8103 lim_exp,
8104 &mut act,
8105 m_e * n_ff_exp,
8106 )?;
8107 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8108 e.with_moe_cache(max_block, |cache, eng| {
8109 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8110 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8111 eng.qmatvec_expert_q8(
8112 cache.buf(slot),
8113 0..dl.len,
8114 &aq2,
8115 &ad2,
8116 m_e,
8117 m.down_exps.in_f,
8118 m.down_exps.out_f,
8119 dl.qtype,
8120 dl.row_bytes,
8121 )
8122 })?
8123 } else {
8124 let gate = e.with_moe_cache(max_block, |cache, eng| {
8125 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8126 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8127 eng.qmatvec_view(
8128 cache.buf(slot),
8129 0..gl.len,
8130 &gv,
8131 m_e,
8132 m.gate_exps.in_f,
8133 m.gate_exps.out_f,
8134 gl.qtype,
8135 gl.row_bytes,
8136 )
8137 })?;
8138 let up = e.with_moe_cache(max_block, |cache, eng| {
8139 let id = BlockId::new(il, PROJ_UP, ex as u16);
8140 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8141 eng.qmatvec_view(
8142 cache.buf(slot),
8143 0..ul.len,
8144 &gv,
8145 m_e,
8146 m.up_exps.in_f,
8147 m.up_exps.out_f,
8148 ul.qtype,
8149 ul.row_bytes,
8150 )
8151 })?;
8152 let mut act = e.uninit(m_e * n_ff_exp)?;
8153 Self::ffn_act_lim(
8154 e,
8155 cfg,
8156 &gate,
8157 &up,
8158 m.gate_exps.macro_scale(ex),
8159 m.up_exps.macro_scale(ex),
8160 lim_exp,
8161 &mut act,
8162 m_e * n_ff_exp,
8163 )?;
8164 let actv = act.slice(0..m_e * n_ff_exp);
8165 e.with_moe_cache(max_block, |cache, eng| {
8166 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8167 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8168 eng.qmatvec_view(
8169 cache.buf(slot),
8170 0..dl.len,
8171 &actv,
8172 m_e,
8173 m.down_exps.in_f,
8174 m.down_exps.out_f,
8175 dl.qtype,
8176 dl.row_bytes,
8177 )
8178 })?
8179 }
8180 } else {
8181 let sg = scratch_g.as_mut().unwrap();
8182 let su = scratch_u.as_mut().unwrap();
8183 let sd = scratch_d.as_mut().unwrap();
8184 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8185 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8186 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8187 if grouped_q8 {
8188 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8189 let gate = e.qmatvec_expert_q8(
8190 sg,
8191 0..gl.len,
8192 &zq,
8193 &zd,
8194 m_e,
8195 m.gate_exps.in_f,
8196 m.gate_exps.out_f,
8197 gl.qtype,
8198 gl.row_bytes,
8199 )?;
8200 let up = e.qmatvec_expert_q8(
8201 su,
8202 0..ul.len,
8203 &zq,
8204 &zd,
8205 m_e,
8206 m.up_exps.in_f,
8207 m.up_exps.out_f,
8208 ul.qtype,
8209 ul.row_bytes,
8210 )?;
8211 let mut act = e.uninit(m_e * n_ff_exp)?;
8212 Self::ffn_act_lim(
8213 e,
8214 cfg,
8215 &gate,
8216 &up,
8217 m.gate_exps.macro_scale(ex),
8218 m.up_exps.macro_scale(ex),
8219 lim_exp,
8220 &mut act,
8221 m_e * n_ff_exp,
8222 )?;
8223 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8224 e.qmatvec_expert_q8(
8225 sd,
8226 0..dl.len,
8227 &aq2,
8228 &ad2,
8229 m_e,
8230 m.down_exps.in_f,
8231 m.down_exps.out_f,
8232 dl.qtype,
8233 dl.row_bytes,
8234 )?
8235 } else {
8236 let gate = e.qmatvec_view(
8237 sg,
8238 0..gl.len,
8239 &gv,
8240 m_e,
8241 m.gate_exps.in_f,
8242 m.gate_exps.out_f,
8243 gl.qtype,
8244 gl.row_bytes,
8245 )?;
8246 let up = e.qmatvec_view(
8247 su,
8248 0..ul.len,
8249 &gv,
8250 m_e,
8251 m.up_exps.in_f,
8252 m.up_exps.out_f,
8253 ul.qtype,
8254 ul.row_bytes,
8255 )?;
8256 let mut act = e.uninit(m_e * n_ff_exp)?;
8257 Self::ffn_act_lim(
8258 e,
8259 cfg,
8260 &gate,
8261 &up,
8262 m.gate_exps.macro_scale(ex),
8263 m.up_exps.macro_scale(ex),
8264 lim_exp,
8265 &mut act,
8266 m_e * n_ff_exp,
8267 )?;
8268 let actv = act.slice(0..m_e * n_ff_exp);
8269 e.qmatvec_view(
8270 sd,
8271 0..dl.len,
8272 &actv,
8273 m_e,
8274 m.down_exps.in_f,
8275 m.down_exps.out_f,
8276 dl.qtype,
8277 dl.row_bytes,
8278 )?
8279 }
8280 };
8281
8282 e.scatter_slot(
8284 &y,
8285 &tok_idx_d,
8286 &slot_idx_d,
8287 &weight_d,
8288 &mut slot_buf,
8289 &mut wbuf,
8290 n_embd,
8291 n_used,
8292 m_e,
8293 )?;
8294 }
8295
8296 let mut moe_out = e.zeros(t * n_embd)?;
8298 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8299
8300 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8302 m_dist.sort_unstable();
8303 let active = m_dist.len();
8304 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8305 let median = m_dist[active / 2];
8306 let max_m = *m_dist.last().unwrap();
8307 let min_m = m_dist[0];
8308 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8309 println!(
8310 "moe-grouped il={il} t={t} active={active}/{n_expert} \
8311 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8312 above_gemm_threshold(>=16)={above16}/{active}"
8313 );
8314 }
8315
8316 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8317 Ok(moe_out)
8318 }
8319
8320 pub(crate) fn moe_ffn_lockstep(
8327 &self,
8328 e: &Engine,
8329 m: &MoeWeights,
8330 zbatch: &CudaSlice<f32>,
8331 mrows: usize,
8332 il: u16,
8333 max_block: usize,
8334 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8335 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8336 let cfg = &self.cfg;
8337 let moe = cfg.moe.as_ref().unwrap();
8338 let n_embd = cfg.n_embd as usize;
8339 let n_expert = moe.expert_count as usize;
8340 let n_used = moe.expert_used_count as usize;
8341 let n_ff_exp = moe.expert_ff_length as usize;
8342 let lim_exp = cfg.clamp_exp_at(il as u32);
8344 let lim_shexp = cfg.clamp_shexp_at(il as u32);
8345
8346 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8347 if let Some(sig) = cfg.sigmoid_router() {
8348 Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8349 }
8350 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8351 Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8352 } else {
8353 Self::moe_route_cfg(
8354 e,
8355 &logits,
8356 mrows,
8357 n_expert,
8358 n_used,
8359 m.active_experts.as_deref(),
8360 )?
8361 };
8362 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8363
8364 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8366 Ok((0..n_expert)
8367 .map(|ex| {
8368 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8369 .into_iter()
8370 .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8371 })
8372 .collect())
8373 })?;
8374
8375 struct Group {
8376 rows: Vec<i32>,
8377 slots: Vec<i32>,
8378 weights: Vec<f32>,
8379 }
8380 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8381 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8382 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8383 Default::default();
8384 for row in 0..mrows {
8385 for j in 0..n_used {
8386 let ex = sel_all[row * n_used + j] as usize;
8387 let w = w_all[row * n_used + j];
8388 if resident_expert[ex] {
8389 let group = groups.entry(ex).or_insert_with(|| Group {
8390 rows: Vec::new(),
8391 slots: Vec::new(),
8392 weights: Vec::new(),
8393 });
8394 group.rows.push(row as i32);
8395 group.slots.push(j as i32);
8396 group.weights.push(w);
8397 } else {
8398 crate::cpu_experts::record_incomplete_gpu_residency(0);
8399 cpu_rows[row].push((ex, w));
8400 cpu_by_expert.entry(ex).or_default().push((row, w));
8401 }
8402 }
8403 }
8404
8405 let host_rows = e.dtoh(zbatch)?;
8411 let rows_ok = crate::cpu_experts::rows_supported();
8412 enum CpuPart {
8413 Single { row: usize },
8414 Rows { rows: Vec<usize> },
8415 }
8416 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8417 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8418 if rows_ok {
8419 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8420 .into_iter()
8421 .filter(|(_, rows)| rows.len() >= 2)
8422 .collect();
8423 shared.sort_by_key(|(ex, _)| *ex);
8424 for (ex, mut row_weights) in shared {
8425 row_weights.sort_by_key(|(row, _)| *row);
8426 let inputs: Vec<(&[f32], f32)> = row_weights
8427 .iter()
8428 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8429 .collect();
8430 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8431 .map_err(std::io::Error::other)?;
8432 for &(row, _) in &row_weights {
8433 rows_served.insert((row, ex));
8434 }
8435 tickets.push((
8436 CpuPart::Rows {
8437 rows: row_weights.iter().map(|&(row, _)| row).collect(),
8438 },
8439 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8440 ));
8441 }
8442 }
8443 for (row, selected) in cpu_rows.iter().enumerate() {
8444 let leftover: Vec<(usize, f32)> = selected
8445 .iter()
8446 .copied()
8447 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8448 .collect();
8449 if leftover.is_empty() {
8450 continue;
8451 }
8452 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8453 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8454 .map_err(std::io::Error::other)?;
8455 tickets.push((
8456 CpuPart::Single { row },
8457 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8458 ));
8459 }
8460
8461 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8462 let mut wbuf = e.zeros(mrows * n_used)?;
8463 let mut order: Vec<usize> = groups.keys().copied().collect();
8464 order.sort_by(|&a, &b| {
8465 groups[&b]
8466 .rows
8467 .len()
8468 .cmp(&groups[&a].rows.len())
8469 .then(a.cmp(&b))
8470 });
8471 for &ex in &order {
8472 let group = &groups[&ex];
8473 let m_e = group.rows.len();
8474 let gl = m.gate_exps.expert_layout(ex);
8475 let ul = m.up_exps.expert_layout(ex);
8476 let dl = m.down_exps.expert_layout(ex);
8477 let row_idx_d = e.htod_i32(&group.rows)?;
8478 let slot_idx_d = e.htod_i32(&group.slots)?;
8479 let dmac = m.down_exps.macro_scale(ex);
8480 let weight_d = if dmac == 1.0 {
8481 e.htod(&group.weights)?
8482 } else {
8483 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8484 e.htod(&scaled)?
8485 };
8486 let mut gathered = e.zeros(m_e * n_embd)?;
8487 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8488 let gv = gathered.slice(0..m_e * n_embd);
8489 let gate = e.with_moe_cache(max_block, |c, eng| {
8490 let slot = c
8491 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8492 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8493 eng.qmatvec_view(
8494 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8495 0..gl.len,
8496 &gv,
8497 m_e,
8498 m.gate_exps.in_f,
8499 m.gate_exps.out_f,
8500 gl.qtype,
8501 gl.row_bytes,
8502 )
8503 })?;
8504 let up = e.with_moe_cache(max_block, |c, eng| {
8505 let slot = c
8506 .resident(BlockId::new(il, PROJ_UP, ex as u16))
8507 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8508 eng.qmatvec_view(
8509 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8510 0..ul.len,
8511 &gv,
8512 m_e,
8513 m.up_exps.in_f,
8514 m.up_exps.out_f,
8515 ul.qtype,
8516 ul.row_bytes,
8517 )
8518 })?;
8519 let mut act = e.zeros(m_e * n_ff_exp)?;
8520 Self::ffn_act_lim(
8521 e,
8522 cfg,
8523 &gate,
8524 &up,
8525 m.gate_exps.macro_scale(ex),
8526 m.up_exps.macro_scale(ex),
8527 lim_exp,
8528 &mut act,
8529 m_e * n_ff_exp,
8530 )?;
8531 let actv = act.slice(0..m_e * n_ff_exp);
8532 let y = e.with_moe_cache(max_block, |c, eng| {
8533 let slot = c
8534 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8535 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8536 eng.qmatvec_view(
8537 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8538 0..dl.len,
8539 &actv,
8540 m_e,
8541 m.down_exps.in_f,
8542 m.down_exps.out_f,
8543 dl.qtype,
8544 dl.row_bytes,
8545 )
8546 })?;
8547 e.scatter_slot(
8548 &y,
8549 &row_idx_d,
8550 &slot_idx_d,
8551 &weight_d,
8552 &mut slot_buf,
8553 &mut wbuf,
8554 n_embd,
8555 n_used,
8556 m_e,
8557 )?;
8558 }
8559 let mut moe_out = e.zeros(mrows * n_embd)?;
8560 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8561
8562 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8564 for (part, ticket) in tickets {
8565 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8566 let mut add_row = |row: usize, chunk: &[f32]| {
8567 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8568 for (accumulator, value) in sum.iter_mut().zip(chunk) {
8569 *accumulator += value;
8570 }
8571 };
8572 match part {
8573 CpuPart::Single { row } => add_row(row, &cpu_output),
8574 CpuPart::Rows { rows } => {
8575 for (slot, row) in rows.into_iter().enumerate() {
8576 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8577 }
8578 }
8579 }
8580 }
8581 for (row, sum) in row_sums.into_iter().enumerate() {
8582 let Some(sum) = sum else { continue };
8583 let cpu_output = e.htod(&sum)?;
8584 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8585 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8586 }
8587
8588 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8589 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8590 {
8591 let n_ff_sh = gate_shexp.out_features();
8592 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8593 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8594 let mut sa = e.zeros(mrows * n_ff_sh)?;
8595 Self::ffn_act_lim(
8596 e,
8597 cfg,
8598 &sg_gate,
8599 &sg_up,
8600 1.0,
8601 1.0,
8602 lim_shexp,
8603 &mut sa,
8604 mrows * n_ff_sh,
8605 )?;
8606 let sh = e.matmul(down_shexp, &sa, mrows)?;
8607 let g = match &m.gate_inp_shexp {
8610 Some(gate_inp_shexp) => {
8611 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8612 }
8613 None => e.htod(&vec![1.0f32; mrows])?,
8614 };
8615 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8616 }
8617
8618 Ok(moe_out)
8619 }
8620}
8621
8622impl HybridModel {
8628 pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
8642 let g = self
8643 .cfg
8644 .gemma4
8645 .as_ref()
8646 .expect("gemma4_rope_dims on a non-gemma4 config");
8647 if g.swa_pattern[il] {
8648 g.rope_dims_swa as usize
8649 } else {
8650 g.rope_dims_global as usize
8651 }
8652 }
8653
8654 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8655 let g = self.cfg.gemma4.as_ref().unwrap();
8656 let swa = g.swa_pattern[il];
8657 let hd = if swa {
8658 g.key_length_swa
8659 } else {
8660 g.key_length_global
8661 } as usize;
8662 (
8666 hd,
8667 g.head_count_kv[il] as usize,
8668 self.cfg.n_head as usize,
8669 if swa {
8670 g.rope_base_swa
8671 } else {
8672 g.rope_base_global
8673 },
8674 1.0,
8675 swa,
8676 )
8677 }
8678
8679 pub(crate) fn gemma4_suppress(
8683 &self,
8684 e: &Engine,
8685 ld: &mut CudaSlice<f32>,
8686 t: usize,
8687 ) -> Result<(), Box<dyn std::error::Error>> {
8688 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8689 #[cfg(debug_assertions)]
8694 crate::debug_assert_tensor_stream_device(
8695 ids,
8696 &e.stream(),
8697 "gemma4_suppress.suppress_d",
8698 );
8699 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8700 }
8701 Ok(())
8702 }
8703
8704 #[allow(clippy::too_many_arguments)]
8709 fn gemma_fa_one_program() -> bool {
8718 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8719 *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
8720 }
8721
8722 fn gemma4_attn_prime(
8723 &self,
8724 e: &Engine,
8725 fa: &crate::hybrid::FullAttnLayer,
8726 il: usize,
8727 h: &CudaSlice<f32>,
8728 pos_d: &CudaSlice<i32>,
8729 t: usize,
8730 cache: Option<&mut Cache>,
8731 island: Option<&CudaSlice<i32>>,
8732 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8733 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8734 let eps = self.cfg.rms_eps;
8735 let aux = self.gemma4_aux.as_ref().unwrap();
8736 let ones = aux.ones(e);
8737 #[cfg(debug_assertions)]
8738 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8739
8740 e.mmq_act_begin();
8743 let q0 = e.matmul(&fa.wq, h, t)?; if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8745 let v = e.dtoh(&q0)?;
8746 let nan = v.iter().filter(|x| x.is_nan()).count();
8747 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8748 eprintln!(
8749 "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
8750 v.len()
8751 );
8752 }
8753 let k0 = e.matmul(&fa.wk, h, t)?; let v0 = if swa {
8757 e.matmul(&fa.wv, h, t)?
8758 } else {
8759 e.clone_dtod(&k0)?
8760 };
8761 if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8762 for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
8763 let v = e.dtoh(buf)?;
8764 let nan = v.iter().filter(|x| x.is_nan()).count();
8765 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8766 eprintln!(
8767 "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
8768 v.len()
8769 );
8770 }
8771 }
8772
8773 let mut q = e.uninit(t * nh * hd)?;
8774 let mut k = e.uninit(t * nkv * hd)?;
8775 let mut v = e.uninit(t * nkv * hd)?;
8777 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8781 let emit = island.is_none()
8784 && t >= 16
8785 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8786 && *EMIT.get_or_init(|| {
8787 std::env::var("MEMRA_FA_EMIT")
8788 .map(|s| s != "0")
8789 .unwrap_or(true)
8790 });
8791 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8792 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8793 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8794 let v_f16 = emit
8797 && crate::fa_f16pv_on()
8798 && match hd {
8799 512 => true,
8800 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8801 _ => false,
8802 };
8803 if emit {
8804 e.rms_norm_qkv_w4b(
8805 &q0,
8806 &k0,
8807 &v0,
8808 fa.q_norm.float_data(),
8809 fa.k_norm.float_data(),
8810 ones,
8811 &mut q,
8812 &mut k,
8813 &mut v,
8814 &mut vb,
8815 hd,
8816 nh * t,
8817 nkv * t,
8818 eps,
8819 v_f16,
8820 )?;
8821 } else {
8822 e.rms_norm_qkv(
8823 &q0,
8824 &k0,
8825 &v0,
8826 fa.q_norm.float_data(),
8827 fa.k_norm.float_data(),
8828 ones,
8829 &mut q,
8830 &mut k,
8831 &mut v,
8832 hd,
8833 nh * t,
8834 nkv * t,
8835 eps,
8836 )?;
8837 }
8838
8839 let ff = if swa {
8840 None
8841 } else {
8842 Some(
8843 aux.rope_freqs(e)
8844 .expect("gemma4 global rope needs rope_freqs.weight"),
8845 )
8846 };
8847 #[cfg(debug_assertions)]
8848 if let Some(ff) = ff {
8849 crate::debug_assert_tensor_stream_device(
8850 ff,
8851 &e.stream(),
8852 "gemma4_attn_prime.rope_freqs",
8853 );
8854 }
8855 if emit {
8856 e.rope_neox2_bf16e(
8857 &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8858 )?;
8859 } else {
8860 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8861 }
8862
8863 if let Some(cache) = cache {
8864 let kvl = cache.kv[il].as_mut().unwrap();
8865 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8866 e.append_kv_quantized_rows(
8867 &k,
8868 &v,
8869 &mut kvl.k,
8870 &mut kvl.v,
8871 kvl.len,
8872 t,
8873 kvl.kv_dim_k,
8874 kvl.kv_dim_v,
8875 kvl.k_tok_bytes,
8876 kvl.v_tok_bytes,
8877 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
8878 )?;
8879 kvl.len += t;
8880 }
8881 let mut attn = e.zeros(t * nh * hd)?;
8882 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8886 if let Some(span) = island {
8887 let w = if swa && t > win { win } else { 0 };
8892 e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
8893 } else if swa && (t > win || Self::gemma_fa_one_program()) {
8894 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8895 if emit {
8896 e.fa_prefill_w_pre(
8897 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
8898 )?;
8899 } else {
8900 e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8901 }
8902 } else {
8903 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8904 }
8905 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8906 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8907 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
8908 if emit {
8909 e.fa_prefill_hd512_pre(
8910 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
8911 )?;
8912 } else {
8913 e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8914 }
8915 } else {
8916 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8917 }
8918 Ok(e.matmul(&fa.wo, &attn, t)?)
8919 }
8920
8921 fn gemma4_attn(
8923 &self,
8924 e: &Engine,
8925 fa: &crate::hybrid::FullAttnLayer,
8926 il: usize,
8927 h: &CudaSlice<f32>,
8928 pos_d: &CudaSlice<i32>,
8929 t: usize,
8930 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8931 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
8932 }
8933
8934 fn gemma4_moe_q8(
8939 &self,
8940 e: &Engine,
8941 m: &crate::hybrid::MoeWeights,
8942 bits: &crate::hybrid::Gemma4MoeBits,
8943 mq: &(CudaSlice<i8>, CudaSlice<f32>),
8944 router_in: &CudaSlice<f32>,
8945 t: usize,
8946 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8947 let cfg = &self.cfg;
8948 let moe = cfg.moe.as_ref().unwrap();
8949 let n_embd = cfg.n_embd as usize;
8950 let n_expert = moe.expert_count as usize;
8951 let n_used = moe.expert_used_count as usize;
8952 let n_ff_exp = moe.expert_ff_length as usize;
8953 let logits = if crate::router_kernel_on() {
8957 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8958 } else {
8959 e.matmul(&m.gate_inp, router_in, t)?
8960 };
8961 let dev = m.dev_exps.as_ref().unwrap();
8962 let (sel_d, w_d) =
8963 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8964 let (zq, zd) = mq;
8965 if t == 1 {
8966 let selv = sel_d.slice(0..n_used);
8967 let wv = w_d.slice(0..n_used);
8968 let act = e.moe_gate_up_gelu8_dev_q8(
8969 &dev.ptr_row,
8970 &selv,
8971 zq,
8972 zd,
8973 n_embd,
8974 n_ff_exp,
8975 n_used,
8976 n_expert,
8977 m.gate_exps.qtype,
8978 m.up_exps.qtype,
8979 m.gate_exps.row_bytes,
8980 m.up_exps.row_bytes,
8981 )?;
8982 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8983 let mut moe_out = e.uninit(n_embd)?;
8984 e.moe_down8_fma_dev_q8(
8985 &dev.ptr_row,
8986 &selv,
8987 &wv,
8988 &aq2,
8989 &ad2,
8990 &mut moe_out.slice_mut(0..n_embd),
8991 n_ff_exp,
8992 n_embd,
8993 n_used,
8994 n_expert,
8995 m.down_exps.qtype,
8996 m.down_exps.row_bytes,
8997 )?;
8998 return Ok(moe_out);
8999 }
9000 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9001 let act = if csr {
9002 e.moe_gate_up_gelu8_dev_q8_csr(
9003 &dev.ptr_row,
9004 &sel_d,
9005 zq,
9006 zd,
9007 t * n_used,
9008 n_embd,
9009 n_ff_exp,
9010 n_used,
9011 n_expert,
9012 m.gate_exps.qtype,
9013 m.up_exps.qtype,
9014 m.gate_exps.row_bytes,
9015 m.up_exps.row_bytes,
9016 )?
9017 } else {
9018 e.moe_gate_up_gelu8_dev_q8_rows(
9019 &dev.ptr_row,
9020 &sel_d,
9021 zq,
9022 zd,
9023 t,
9024 n_embd,
9025 n_ff_exp,
9026 n_used,
9027 n_expert,
9028 m.gate_exps.qtype,
9029 m.up_exps.qtype,
9030 m.gate_exps.row_bytes,
9031 m.up_exps.row_bytes,
9032 )?
9033 };
9034 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9035 let mut moe_out = e.uninit(t * n_embd)?;
9036 e.moe_down8_fma_dev_q8_rows_g(
9039 &dev.ptr_row,
9040 &sel_d,
9041 &w_d,
9042 &aq2,
9043 &ad2,
9044 &mut moe_out,
9045 t,
9046 n_ff_exp,
9047 n_embd,
9048 n_used,
9049 n_expert,
9050 m.down_exps.qtype,
9051 m.down_exps.row_bytes,
9052 )?;
9053 Ok(moe_out)
9054 }
9055
9056 fn gemma4_moe(
9060 &self,
9061 e: &Engine,
9062 m: &crate::hybrid::MoeWeights,
9063 bits: &crate::hybrid::Gemma4MoeBits,
9064 moe_in: &CudaSlice<f32>,
9065 router_in: &CudaSlice<f32>,
9066 t: usize,
9067 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9068 let cfg = &self.cfg;
9069 let moe = cfg.moe.as_ref().unwrap();
9070 let n_embd = cfg.n_embd as usize;
9071 let n_expert = moe.expert_count as usize;
9072 let n_used = moe.expert_used_count as usize;
9073 let n_ff_exp = moe.expert_ff_length as usize;
9074
9075 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
9079 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9080 } else {
9081 e.matmul(&m.gate_inp, router_in, t)?
9082 };
9083
9084 if t < PRIME_MIN_T
9089 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9090 && expert_dp4a_supported(m.gate_exps.qtype)
9091 && expert_dp4a_supported(m.up_exps.qtype)
9092 && expert_dp4a_supported(m.down_exps.qtype)
9093 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9094 {
9095 let dev = m.dev_exps.as_ref().unwrap();
9096 let (sel_d, w_d) =
9097 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9098 if t == 1 {
9099 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
9100 let selv = sel_d.slice(0..n_used);
9101 let wv = w_d.slice(0..n_used);
9102 let act = e.moe_gate_up_gelu8_dev_q8(
9103 &dev.ptr_row,
9104 &selv,
9105 &zq,
9106 &zd,
9107 n_embd,
9108 n_ff_exp,
9109 n_used,
9110 n_expert,
9111 m.gate_exps.qtype,
9112 m.up_exps.qtype,
9113 m.gate_exps.row_bytes,
9114 m.up_exps.row_bytes,
9115 )?;
9116 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9117 let mut moe_out = e.uninit(n_embd)?;
9118 e.moe_down8_fma_dev_q8(
9119 &dev.ptr_row,
9120 &selv,
9121 &wv,
9122 &aq2,
9123 &ad2,
9124 &mut moe_out.slice_mut(0..n_embd),
9125 n_ff_exp,
9126 n_embd,
9127 n_used,
9128 n_expert,
9129 m.down_exps.qtype,
9130 m.down_exps.row_bytes,
9131 )?;
9132 return Ok(moe_out);
9133 }
9134 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9139 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9140 let act = if csr {
9141 e.moe_gate_up_gelu8_dev_q8_csr(
9142 &dev.ptr_row,
9143 &sel_d,
9144 &zq,
9145 &zd,
9146 t * n_used,
9147 n_embd,
9148 n_ff_exp,
9149 n_used,
9150 n_expert,
9151 m.gate_exps.qtype,
9152 m.up_exps.qtype,
9153 m.gate_exps.row_bytes,
9154 m.up_exps.row_bytes,
9155 )?
9156 } else {
9157 e.moe_gate_up_gelu8_dev_q8_rows(
9158 &dev.ptr_row,
9159 &sel_d,
9160 &zq,
9161 &zd,
9162 t,
9163 n_embd,
9164 n_ff_exp,
9165 n_used,
9166 n_expert,
9167 m.gate_exps.qtype,
9168 m.up_exps.qtype,
9169 m.gate_exps.row_bytes,
9170 m.up_exps.row_bytes,
9171 )?
9172 };
9173 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9174 let mut moe_out = e.uninit(t * n_embd)?;
9175 e.moe_down8_fma_dev_q8_rows_g(
9176 &dev.ptr_row,
9177 &sel_d,
9178 &w_d,
9179 &aq2,
9180 &ad2,
9181 &mut moe_out,
9182 t,
9183 n_ff_exp,
9184 n_embd,
9185 n_used,
9186 n_expert,
9187 m.down_exps.qtype,
9188 m.down_exps.row_bytes,
9189 )?;
9190 return Ok(moe_out);
9191 }
9192
9193 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9194 for (i, &sx) in sel_all.iter().enumerate() {
9195 w_all[i] *= bits.per_expert_scale[sx as usize];
9196 }
9197
9198 if t >= PRIME_MIN_T
9202 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9203 && expert_dp4a_supported(m.gate_exps.qtype)
9204 && expert_dp4a_supported(m.up_exps.qtype)
9205 && expert_dp4a_supported(m.down_exps.qtype)
9206 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9207 {
9208 let dev = m.dev_exps.as_ref().unwrap();
9209 let n_pairs = t * n_used;
9210 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9211 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9212 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9213 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9214 let pt = e.htod_i32(&pair_tok)?;
9215 let pw = e.htod(&w_all)?;
9216 let toff = e.htod_i32(&tok_off)?;
9217 let tids = e.htod_i32(&tok_ids)?;
9218 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9219 for p in 0..n_pairs {
9220 by_ex[pair_ex[p] as usize].push(p as i32);
9221 }
9222 let mut ex_ids: Vec<i32> = Vec::new();
9223 let mut ex_off: Vec<i32> = vec![0];
9224 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9225 for (ex, list) in by_ex.iter().enumerate() {
9226 if list.is_empty() {
9227 continue;
9228 }
9229 ex_ids.push(ex as i32);
9230 ex_pairs.extend_from_slice(list);
9231 ex_off.push(ex_pairs.len() as i32);
9232 }
9233 let n_active = ex_ids.len();
9234 let exi = e.htod_i32(&ex_ids)?;
9235 let exo = e.htod_i32(&ex_off)?;
9236 let exp_d = e.htod_i32(&ex_pairs)?;
9237 if crate::moe_f16g_gemma_on()
9245 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9246 && f16g_proj_ok(m.up_exps.qtype, n_embd)
9247 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9248 {
9249 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9250 let csr_tok_d = e.htod_i32(&csr_tok)?;
9251 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9252 let g_csr = e.moe_f16_grouped(
9253 &dev.ptr_row,
9254 0,
9255 n_expert,
9256 &exi,
9257 &ex_off,
9258 &exo,
9259 &z_f16,
9260 &z_s,
9261 n_embd,
9262 n_ff_exp,
9263 n_active,
9264 n_pairs,
9265 m.gate_exps.qtype,
9266 m.gate_exps.row_bytes,
9267 )?;
9268 let u_csr = e.moe_f16_grouped(
9269 &dev.ptr_row,
9270 1,
9271 n_expert,
9272 &exi,
9273 &ex_off,
9274 &exo,
9275 &z_f16,
9276 &z_s,
9277 n_embd,
9278 n_ff_exp,
9279 n_active,
9280 n_pairs,
9281 m.up_exps.qtype,
9282 m.up_exps.row_bytes,
9283 )?;
9284 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9285 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9286 let d_csr = e.moe_f16_grouped(
9287 &dev.ptr_row,
9288 2,
9289 n_expert,
9290 &exi,
9291 &ex_off,
9292 &exo,
9293 &a_f16,
9294 &a_s,
9295 n_ff_exp,
9296 n_embd,
9297 n_active,
9298 n_pairs,
9299 m.down_exps.qtype,
9300 m.down_exps.row_bytes,
9301 )?;
9302 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9303 let mut moe_out = e.uninit(t * n_embd)?;
9304 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9305 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9306 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9307 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9308 eprintln!(
9309 "[f16g-debug] post-permute bad={} post-scatter bad={}",
9310 scan(&yd),
9311 scan(&mo)
9312 );
9313 }
9314 return Ok(moe_out);
9315 }
9316 let mma =
9319 n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9320 let (gate, up) = if mma {
9321 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9322 (
9323 e.mmq_iq_experts(
9324 &dev.ptr_row,
9325 0,
9326 n_expert,
9327 &exi,
9328 &exo,
9329 &exp_d,
9330 &pt,
9331 &z_scr,
9332 n_embd,
9333 n_ff_exp,
9334 n_active,
9335 n_pairs,
9336 t,
9337 m.gate_exps.qtype,
9338 m.gate_exps.row_bytes,
9339 )?,
9340 e.mmq_iq_experts(
9341 &dev.ptr_row,
9342 1,
9343 n_expert,
9344 &exi,
9345 &exo,
9346 &exp_d,
9347 &pt,
9348 &z_scr,
9349 n_embd,
9350 n_ff_exp,
9351 n_active,
9352 n_pairs,
9353 t,
9354 m.up_exps.qtype,
9355 m.up_exps.row_bytes,
9356 )?,
9357 )
9358 } else {
9359 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9360 (
9361 e.moe_pairs_matvec_q8_dec(
9362 &dev.ptr_row,
9363 0,
9364 &exi,
9365 &exo,
9366 &exp_d,
9367 &pt,
9368 &zq,
9369 &zd,
9370 n_embd,
9371 n_ff_exp,
9372 n_expert,
9373 n_active,
9374 n_pairs,
9375 m.gate_exps.qtype,
9376 m.gate_exps.row_bytes,
9377 )?,
9378 e.moe_pairs_matvec_q8_dec(
9379 &dev.ptr_row,
9380 1,
9381 &exi,
9382 &exo,
9383 &exp_d,
9384 &pt,
9385 &zq,
9386 &zd,
9387 n_embd,
9388 n_ff_exp,
9389 n_expert,
9390 n_active,
9391 n_pairs,
9392 m.up_exps.qtype,
9393 m.up_exps.row_bytes,
9394 )?,
9395 )
9396 };
9397 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9398 let pself = e.htod_i32(&pair_self)?;
9399 let y_down = if mma {
9411 let in_pad = n_ff_exp.div_ceil(256) * 256;
9412 let a_scr = if crate::moe_fuse_actq_on() {
9413 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9414 } else {
9415 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9416 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9417 };
9418 e.mmq_iq_experts(
9419 &dev.ptr_row,
9420 2,
9421 n_expert,
9422 &exi,
9423 &exo,
9424 &exp_d,
9425 &pself,
9426 &a_scr,
9427 in_pad,
9428 n_embd,
9429 n_active,
9430 n_pairs,
9431 n_pairs,
9432 m.down_exps.qtype,
9433 m.down_exps.row_bytes,
9434 )?
9435 } else {
9436 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9437 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9438 e.moe_pairs_matvec_q8_dec(
9439 &dev.ptr_row,
9440 2,
9441 &exi,
9442 &exo,
9443 &exp_d,
9444 &pself,
9445 &aq2,
9446 &ad2,
9447 n_ff_exp,
9448 n_embd,
9449 n_expert,
9450 n_active,
9451 n_pairs,
9452 m.down_exps.qtype,
9453 m.down_exps.row_bytes,
9454 )?
9455 };
9456 let mut moe_out = e.uninit(t * n_embd)?;
9457 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9458 return Ok(moe_out);
9459 }
9460
9461 let g_len = m.gate_exps.expert_stride;
9462 let u_len = m.up_exps.expert_stride;
9463 let d_len = m.down_exps.expert_stride;
9464 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9468 let (mut sg, mut su, mut sd) = if dev.is_some() {
9469 (None, None, None)
9470 } else {
9471 (
9472 Some(e.alloc_u8_uninit(g_len)?),
9473 Some(e.alloc_u8_uninit(u_len)?),
9474 Some(e.alloc_u8_uninit(d_len)?),
9475 )
9476 };
9477 let mut moe_out = e.zeros(t * n_embd)?;
9478 for tok in 0..t {
9479 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9480 let w = &w_all[tok * n_used..(tok + 1) * n_used];
9481 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9482 for (j, &ex) in sel.iter().enumerate() {
9483 let ex = ex as usize;
9484 let gate = match dev {
9485 Some(d) => e.qmatvec_view(
9486 &d.gate,
9487 ex * g_len..(ex + 1) * g_len,
9488 &zt,
9489 1,
9490 m.gate_exps.in_f,
9491 m.gate_exps.out_f,
9492 m.gate_exps.qtype,
9493 m.gate_exps.row_bytes,
9494 )?,
9495 None => {
9496 let sg = sg.as_mut().unwrap();
9497 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9498 e.qmatvec_view(
9499 sg,
9500 0..g_len,
9501 &zt,
9502 1,
9503 m.gate_exps.in_f,
9504 m.gate_exps.out_f,
9505 m.gate_exps.qtype,
9506 m.gate_exps.row_bytes,
9507 )?
9508 }
9509 };
9510 let up = match dev {
9511 Some(d) => e.qmatvec_view(
9512 &d.up,
9513 ex * u_len..(ex + 1) * u_len,
9514 &zt,
9515 1,
9516 m.up_exps.in_f,
9517 m.up_exps.out_f,
9518 m.up_exps.qtype,
9519 m.up_exps.row_bytes,
9520 )?,
9521 None => {
9522 let su = su.as_mut().unwrap();
9523 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9524 e.qmatvec_view(
9525 su,
9526 0..u_len,
9527 &zt,
9528 1,
9529 m.up_exps.in_f,
9530 m.up_exps.out_f,
9531 m.up_exps.qtype,
9532 m.up_exps.row_bytes,
9533 )?
9534 }
9535 };
9536 let mut act = e.uninit(n_ff_exp)?;
9537 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9538 let actv = act.slice(0..n_ff_exp);
9539 let y = match dev {
9540 Some(d) => e.qmatvec_view(
9541 &d.down,
9542 ex * d_len..(ex + 1) * d_len,
9543 &actv,
9544 1,
9545 m.down_exps.in_f,
9546 m.down_exps.out_f,
9547 m.down_exps.qtype,
9548 m.down_exps.row_bytes,
9549 )?,
9550 None => {
9551 let sd = sd.as_mut().unwrap();
9552 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9553 e.qmatvec_view(
9554 sd,
9555 0..d_len,
9556 &actv,
9557 1,
9558 m.down_exps.in_f,
9559 m.down_exps.out_f,
9560 m.down_exps.qtype,
9561 m.down_exps.row_bytes,
9562 )?
9563 }
9564 };
9565 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9566 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9567 }
9568 }
9569 Ok(moe_out)
9570 }
9571
9572 fn gemma4_layer(
9574 &self,
9575 e: &Engine,
9576 il: usize,
9577 layer: &crate::hybrid::HybridLayer,
9578 x: &CudaSlice<f32>,
9579 pos_d: &CudaSlice<i32>,
9580 t: usize,
9581 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9582 let n_embd = self.cfg.n_embd as usize;
9583 let eps = self.cfg.rms_eps;
9584
9585 let mut h = e.zeros(t * n_embd)?;
9586 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9587 let Mixer::Full(fa) = &layer.mixer else {
9588 panic!("gemma4 layer {il} not full-attn")
9589 };
9590 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9591 let mut cur = e.zeros(t * n_embd)?;
9593 e.rms_norm(
9594 &o,
9595 layer.post_attn_norm.float_data(),
9596 &mut cur,
9597 n_embd,
9598 t,
9599 eps,
9600 )?;
9601 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9602 }
9603
9604 fn gemma4_layer_tail_add(
9608 &self,
9609 e: &Engine,
9610 layer: &crate::hybrid::HybridLayer,
9611 cur: &CudaSlice<f32>,
9612 x: &CudaSlice<f32>,
9613 t: usize,
9614 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9615 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9616 }
9617
9618 fn gemma4_layer_tail_add_n(
9621 &self,
9622 e: &Engine,
9623 layer: &crate::hybrid::HybridLayer,
9624 cur: &CudaSlice<f32>,
9625 x: &CudaSlice<f32>,
9626 t: usize,
9627 next_norm: Option<&CudaSlice<f32>>,
9628 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9629 let n_embd = self.cfg.n_embd as usize;
9630 let bits = layer.gemma4.as_ref().unwrap();
9631 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9632 let mut xn = e.uninit(t * n_embd)?;
9633 match next_norm {
9634 Some(w) => {
9635 let mut hn = e.uninit(t * n_embd)?;
9636 e.add_scale_rms_norm(
9637 &sn,
9638 &attn_out,
9639 bits.layer_scale,
9640 w,
9641 &mut xn,
9642 &mut hn,
9643 n_embd,
9644 t,
9645 self.cfg.rms_eps,
9646 )?;
9647 Ok((xn, Some(hn)))
9648 }
9649 None => {
9650 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9651 Ok((xn, None))
9652 }
9653 }
9654 }
9655
9656 fn gemma4_layer_tail_core(
9659 &self,
9660 e: &Engine,
9661 layer: &crate::hybrid::HybridLayer,
9662 cur: &CudaSlice<f32>,
9663 x: &CudaSlice<f32>,
9664 t: usize,
9665 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9666 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9667 }
9668
9669 fn gemma4_layer_tail_core_pn(
9676 &self,
9677 e: &Engine,
9678 layer: &crate::hybrid::HybridLayer,
9679 cur: &CudaSlice<f32>,
9680 x: &CudaSlice<f32>,
9681 t: usize,
9682 pre_norm: Option<&CudaSlice<f32>>,
9683 defer_post_norm: bool,
9684 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9685 let n_embd = self.cfg.n_embd as usize;
9686 let eps = self.cfg.rms_eps;
9687 let bits = layer.gemma4.as_ref().unwrap();
9688
9689 let Some(mbits) = bits.moe_bits.as_ref() else {
9692 let crate::hybrid::Ffn::Dense {
9693 ffn_gate,
9694 ffn_up,
9695 ffn_down,
9696 } = &layer.ffn
9697 else {
9698 panic!("gemma4 dense layer without Dense ffn")
9699 };
9700 let mut attn_out = e.uninit(t * n_embd)?;
9701 let mut zsh = e.uninit(t * n_embd)?;
9702 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9705 match pre_norm {
9706 Some(wa) if t == 1 => {
9707 zpair = Some(e.rms_pre_add_rms_norm_q8z(
9708 cur,
9709 wa,
9710 x,
9711 bits.ffn_norm.float_data(),
9712 &mut attn_out,
9713 &mut zsh,
9714 n_embd,
9715 t,
9716 eps,
9717 )?);
9718 }
9719 Some(wa) => e.rms_pre_add_rms_norm(
9720 cur,
9721 wa,
9722 x,
9723 bits.ffn_norm.float_data(),
9724 &mut attn_out,
9725 &mut zsh,
9726 n_embd,
9727 t,
9728 eps,
9729 )?,
9730 None => e.add_rms_norm(
9731 cur,
9732 x,
9733 bits.ffn_norm.float_data(),
9734 &mut attn_out,
9735 &mut zsh,
9736 n_embd,
9737 t,
9738 eps,
9739 )?,
9740 }
9741 let n_ff = ffn_gate.out_features();
9742 let (gate, up) = if t == 1 {
9748 let (zq, zd) = match zpair {
9749 Some(p) => p,
9750 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9751 };
9752 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9753 Some(p) => p,
9754 None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
9756 Some(p) => p,
9757 None => (
9758 e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9759 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9760 ),
9761 },
9762 }
9763 } else {
9764 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9769 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9770 let fused = if f2b {
9771 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9772 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9773 } else {
9774 None
9775 };
9776 match fused {
9777 Some(p) => p,
9778 None => {
9779 e.mmq_act_begin();
9781 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9782 }
9783 }
9784 };
9785 let mut act = e.uninit(t * n_ff)?;
9786 let f0 = if e.uses_q8_1_fast(ffn_down) {
9789 let upv = e.view(&up, t * n_ff);
9790 let up_all = upv.slice(0..t * n_ff);
9791 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9792 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9793 } else {
9794 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9795 e.matmul(ffn_down, &act, t)?
9796 };
9797 if defer_post_norm {
9798 return Ok((f0, attn_out));
9799 }
9800 let mut sn = e.uninit(t * n_embd)?;
9801 e.rms_norm(
9802 &f0,
9803 bits.post_ffw_norm.float_data(),
9804 &mut sn,
9805 n_embd,
9806 t,
9807 eps,
9808 )?;
9809 return Ok((sn, attn_out));
9810 };
9811
9812 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9813 let mut attn_out = e.uninit(t * n_embd)?;
9818 let mut router_in = e.uninit(t * n_embd)?;
9819 let fast_moe = match &layer.ffn {
9820 crate::hybrid::Ffn::Moe(m) => {
9821 m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9822 && expert_dp4a_supported(m.gate_exps.qtype)
9823 && expert_dp4a_supported(m.up_exps.qtype)
9824 && expert_dp4a_supported(m.down_exps.qtype)
9825 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9826 }
9827 _ => false,
9828 };
9829 let q8z = t < PRIME_MIN_T && fast_moe;
9830 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9831 let (z0, m2) = e.add_rms_norm3_q8z(
9832 cur,
9833 x,
9834 bits.ffn_norm.float_data(),
9835 &mbits.router_scale_pre,
9836 mbits.pre_ffw_norm_2.float_data(),
9837 &mut attn_out,
9838 &mut router_in,
9839 n_embd,
9840 t,
9841 eps,
9842 )?;
9843 (None, Some(z0), Some(m2))
9844 } else {
9845 let mut zsh = e.uninit(t * n_embd)?;
9846 let mut moe_in = e.uninit(t * n_embd)?;
9847 e.add_rms_norm3(
9848 cur,
9849 x,
9850 bits.ffn_norm.float_data(),
9851 &mbits.router_scale_pre,
9852 mbits.pre_ffw_norm_2.float_data(),
9853 &mut attn_out,
9854 &mut zsh,
9855 &mut router_in,
9856 &mut moe_in,
9857 n_embd,
9858 t,
9859 eps,
9860 )?;
9861 (Some((zsh, moe_in)), None, None)
9862 };
9863 let attn_out2 = attn_out;
9864 #[allow(unused_variables)]
9865 let attn_out = &attn_out2;
9866 let n_ff = mbits.shared_gate.out_features();
9867 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9868 if t == 1 {
9869 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9870 Some(p) => p,
9871 None => match e.matmul_nvfp4_fused2(
9872 &mbits.shared_gate,
9873 &mbits.shared_up,
9874 zq,
9875 zd,
9876 1,
9877 )? {
9878 Some(p) => p,
9879 None => {
9880 let h0 = e.zeros(0)?;
9881 (
9882 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
9883 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
9884 )
9885 }
9886 },
9887 }
9888 } else {
9889 let h0 = e.zeros(0)?;
9891 (
9892 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
9893 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
9894 )
9895 }
9896 } else {
9897 let (zsh, _) = zsh_f32.as_ref().unwrap();
9898 (
9899 e.matmul(&mbits.shared_gate, zsh, t)?,
9900 e.matmul(&mbits.shared_up, zsh, t)?,
9901 )
9902 };
9903 let mut act = e.uninit(t * n_ff)?;
9904 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9905 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
9906 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
9907 panic!("gemma4 layer not MoE")
9908 };
9909 let moe0 = match (&moe_q8, &zsh_f32) {
9910 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
9911 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
9912 _ => unreachable!(),
9913 };
9914 let mut mlp = e.uninit(t * n_embd)?;
9916 let mut moe = e.uninit(t * n_embd)?;
9917 e.rms_norm2x(
9918 &mlp0,
9919 &moe0,
9920 mbits.post_ffw_norm_1.float_data(),
9921 mbits.post_ffw_norm_2.float_data(),
9922 &mut mlp,
9923 &mut moe,
9924 n_embd,
9925 t,
9926 eps,
9927 )?;
9928
9929 let mut sum = e.uninit(t * n_embd)?;
9932 let mut sn = e.uninit(t * n_embd)?;
9933 e.add_rms_norm(
9934 &mlp,
9935 &moe,
9936 bits.post_ffw_norm.float_data(),
9937 &mut sum,
9938 &mut sn,
9939 n_embd,
9940 t,
9941 eps,
9942 )?;
9943 Ok((sn, attn_out2))
9944 }
9945
9946 pub(crate) fn gemma4_layer_tail_add_nq_pn(
9956 &self,
9957 e: &Engine,
9958 layer: &crate::hybrid::HybridLayer,
9959 o: &CudaSlice<f32>,
9960 x: &CudaSlice<f32>,
9961 t: usize,
9962 next_norm: Option<&CudaSlice<f32>>,
9963 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9964 {
9965 let n_embd = self.cfg.n_embd as usize;
9966 let eps = self.cfg.rms_eps;
9967 let bits = layer.gemma4.as_ref().unwrap();
9968 if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
9969 let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
9970 e,
9971 layer,
9972 o,
9973 x,
9974 t,
9975 Some(layer.post_attn_norm.float_data()),
9976 true,
9977 )?;
9978 let mut xn = e.uninit(t * n_embd)?;
9979 return match next_norm {
9980 Some(w) => {
9981 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
9982 &f0,
9983 bits.post_ffw_norm.float_data(),
9984 &attn_out,
9985 bits.layer_scale,
9986 w,
9987 &mut xn,
9988 n_embd,
9989 t,
9990 eps,
9991 )?;
9992 Ok((xn, Some(pair)))
9993 }
9994 None => {
9995 let mut sn = e.uninit(t * n_embd)?;
9996 e.rms_norm(
9997 &f0,
9998 bits.post_ffw_norm.float_data(),
9999 &mut sn,
10000 n_embd,
10001 t,
10002 eps,
10003 )?;
10004 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10005 Ok((xn, None))
10006 }
10007 };
10008 }
10009 let mut cur = e.uninit(t * n_embd)?;
10010 e.rms_norm(
10011 o,
10012 layer.post_attn_norm.float_data(),
10013 &mut cur,
10014 n_embd,
10015 t,
10016 eps,
10017 )?;
10018 self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
10019 }
10020
10021 pub(crate) fn gemma4_layer_tail_add_nq(
10022 &self,
10023 e: &Engine,
10024 layer: &crate::hybrid::HybridLayer,
10025 cur: &CudaSlice<f32>,
10026 x: &CudaSlice<f32>,
10027 t: usize,
10028 next_norm: Option<&CudaSlice<f32>>,
10029 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
10030 {
10031 let n_embd = self.cfg.n_embd as usize;
10032 let bits = layer.gemma4.as_ref().unwrap();
10033 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
10034 let mut xn = e.uninit(t * n_embd)?;
10035 match next_norm {
10036 Some(w) => {
10037 let pair = e.add_scale_rms_norm_q8_1(
10038 &sn,
10039 &attn_out,
10040 bits.layer_scale,
10041 w,
10042 &mut xn,
10043 n_embd,
10044 t,
10045 self.cfg.rms_eps,
10046 )?;
10047 Ok((xn, Some(pair)))
10048 }
10049 None => {
10050 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10051 Ok((xn, None))
10052 }
10053 }
10054 }
10055
10056 fn gemma4_forward(
10059 &self,
10060 e: &Engine,
10061 tokens: &[u32],
10062 last_only: bool,
10063 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
10064 if self.is_gemma4_e4b() {
10067 return self.gemma4_e4b_forward(e, tokens, last_only);
10068 }
10069 let n_embd = self.cfg.n_embd as usize;
10070 let t = tokens.len();
10071 let pos: Vec<i32> = (0..t as i32).collect();
10072 let pos_d = e.htod_i32(&pos)?;
10073
10074 let mut x = self.embed(e, tokens)?;
10075 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10076 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
10079 let stat =
10080 |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
10081 let h = e.dtoh(x)?;
10082 let bad = h.iter().filter(|v| !v.is_finite()).count();
10083 let mx = h
10084 .iter()
10085 .filter(|v| v.is_finite())
10086 .fold(0.0f32, |m, v| m.max(v.abs()));
10087 eprintln!(
10088 "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
10089 &h[..3]
10090 );
10091 Ok(())
10092 };
10093 if probe {
10094 stat(e, &x, "embed")?;
10095 }
10096 for (il, layer) in self.layers.iter().enumerate() {
10097 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
10098 if probe {
10099 stat(e, &x, &format!("L{il}"))?;
10100 }
10101 }
10102 let mut hn = e.zeros(t * n_embd)?;
10103 e.rms_norm(
10104 &x,
10105 self.output_norm.float_data(),
10106 &mut hn,
10107 n_embd,
10108 t,
10109 self.cfg.rms_eps,
10110 )?;
10111 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10112 let n_vocab = self.output.out_features();
10113 let logits = if last_only {
10114 let hv = e.view(&hn, t * n_embd);
10115 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
10116 let mut hlast = e.zeros(n_embd)?;
10117 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
10118 let mut ld = e.matmul(&self.output, &hlast, 1)?;
10119 e.softcap(&mut ld, cap, n_vocab)?;
10120 self.gemma4_suppress(e, &mut ld, 1)?;
10121 e.dtoh(&ld)?
10122 } else {
10123 let mut ld = e.matmul(&self.output, &hn, t)?;
10124 e.softcap(&mut ld, cap, t * n_vocab)?;
10125 self.gemma4_suppress(e, &mut ld, t)?;
10126 e.dtoh(&ld)?
10127 };
10128 Ok(logits)
10129 }
10130
10131 pub(crate) fn gemma4_prime(
10136 &self,
10137 e: &Engine,
10138 tokens: &[u32],
10139 cache: &mut Cache,
10140 overlay: Option<&crate::vision::EmbedOverlay>,
10141 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10142 if cache.pos != 0 {
10147 return Err(
10148 "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
10149 — prime the full prompt in one call or decode tokenwise"
10150 .into(),
10151 );
10152 }
10153 let n_embd = self.cfg.n_embd as usize;
10154 let eps = self.cfg.rms_eps;
10155 let t = tokens.len();
10156 let pos: Vec<i32> = (0..t as i32).collect();
10157 let pos_d = e.htod_i32(&pos)?;
10158 let mut x = self.embed(e, tokens)?;
10159 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10160 let island: Option<CudaSlice<i32>> = match overlay {
10167 Some(ov) => {
10168 let mut span_id = vec![-1i32; t];
10169 for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
10170 if pos + n_rows > t {
10171 return Err(format!(
10172 "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
10173 pos + n_rows
10174 )
10175 .into());
10176 }
10177 let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
10178 e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
10179 for s in span_id.iter_mut().skip(pos).take(n_rows) {
10180 *s = i as i32;
10181 }
10182 }
10183 if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
10187 None
10188 } else {
10189 Some(e.htod_i32(&span_id)?)
10190 }
10191 }
10192 None => None,
10193 };
10194 for (il, layer) in self.layers.iter().enumerate() {
10195 let mut h = e.zeros(t * n_embd)?;
10196 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
10197 let Mixer::Full(fa) = &layer.mixer else {
10198 panic!("gemma4 layer not full-attn")
10199 };
10200 let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
10201 if trace {
10202 let v = e.dtoh(&h)?;
10203 let nan = v.iter().filter(|x| x.is_nan()).count();
10204 eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
10205 }
10206 let o =
10207 self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
10208 if trace {
10209 let v = e.dtoh(&o)?;
10210 let nan = v.iter().filter(|x| x.is_nan()).count();
10211 eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
10212 }
10213 let mut cur = e.zeros(t * n_embd)?;
10214 e.rms_norm(
10215 &o,
10216 layer.post_attn_norm.float_data(),
10217 &mut cur,
10218 n_embd,
10219 t,
10220 eps,
10221 )?;
10222 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
10223 self.dflash_tap(e, cache, il, &x, t)?;
10224 if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
10226 let h = e.dtoh(&x)?;
10227 let nan = h.iter().filter(|v| v.is_nan()).count();
10228 let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
10229 eprintln!(
10230 "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
10231 h.len()
10232 );
10233 if nan > 0 {
10234 return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
10235 }
10236 }
10237 }
10238 cache.pos += t;
10239 let hiddens = e.clone_dtod(&x)?;
10240 let xv = e.view(&x, t * n_embd);
10241 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
10242 let mut h_seed = e.zeros(n_embd)?;
10243 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
10244 let mut hn = e.uninit(n_embd)?;
10245 e.rms_norm(
10246 &h_seed,
10247 self.output_norm.float_data(),
10248 &mut hn,
10249 n_embd,
10250 1,
10251 eps,
10252 )?;
10253 let mut ld = e.matmul(&self.output, &hn, 1)?;
10254 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10255 e.softcap(&mut ld, cap, self.output.out_features())?;
10256 self.gemma4_suppress(e, &mut ld, 1)?;
10257 let logits = e.dtoh(&ld)?;
10258 Ok((logits, h_seed, hiddens))
10259 }
10260
10261 fn gemma4_decode_attn(
10266 &self,
10267 e: &Engine,
10268 fa: &crate::hybrid::FullAttnLayer,
10269 il: usize,
10270 hq: &CudaSlice<i8>,
10271 hdq: &CudaSlice<f32>,
10272 pos_d: &CudaSlice<i32>,
10273 cache: &mut Cache,
10274 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10275 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10276 let eps = self.cfg.rms_eps;
10277 let aux = self.gemma4_aux.as_ref().unwrap();
10278 let ones = aux.ones(e);
10279 #[cfg(debug_assertions)]
10280 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
10281 let (hq, hdq) = (hq, hdq);
10282 let h0 = e.zeros(0)?;
10283 let h = &h0;
10284 let (q0, k0, v0) = if swa {
10285 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
10286 Some(t3) => t3,
10287 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10290 Some((q0, k0)) => {
10291 let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?;
10292 (q0, k0, v0)
10293 }
10294 None => (
10295 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10296 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10297 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10298 ),
10299 },
10300 }
10301 } else {
10302 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10303 Some(p) => p,
10304 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10305 Some(p) => p,
10306 None => (
10307 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10308 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10309 ),
10310 },
10311 };
10312 let v0 = e.clone_dtod(&k0)?;
10313 (q0, k0, v0)
10314 };
10315 let mut q = e.uninit(nh * hd)?;
10316 let mut k = e.uninit(nkv * hd)?;
10317 let mut v = e.uninit(nkv * hd)?;
10318 let ff = if swa {
10321 None
10322 } else {
10323 Some(
10324 aux.rope_freqs(e)
10325 .expect("gemma4 global rope needs rope_freqs.weight"),
10326 )
10327 };
10328 #[cfg(debug_assertions)]
10329 if let Some(ff) = ff {
10330 crate::debug_assert_tensor_stream_device(
10331 ff,
10332 &e.stream(),
10333 "gemma4_decode_attn.rope_freqs",
10334 );
10335 }
10336 let kvl = cache.kv[il].as_mut().unwrap();
10337 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10338 if crate::Engine::qkv_append_on() {
10339 e.rms_norm_qkv_rope_append(
10343 &q0,
10344 &k0,
10345 &v0,
10346 fa.q_norm.float_data(),
10347 fa.k_norm.float_data(),
10348 ones,
10349 &mut q,
10350 &mut k,
10351 &mut v,
10352 hd,
10353 self.gemma4_rope_dims(il),
10354 nh,
10355 nkv,
10356 pos_d,
10357 nh,
10358 nkv,
10359 base,
10360 1.0,
10361 ff,
10362 eps,
10363 &mut kvl.k,
10364 &mut kvl.v,
10365 kvl.len,
10366 kvl.k_tok_bytes,
10367 kvl.v_tok_bytes,
10368 kv_fp8,
10369 )?;
10370 } else {
10371 e.rms_norm_qkv_rope(
10372 &q0,
10373 &k0,
10374 &v0,
10375 fa.q_norm.float_data(),
10376 fa.k_norm.float_data(),
10377 ones,
10378 &mut q,
10379 &mut k,
10380 &mut v,
10381 hd,
10382 self.gemma4_rope_dims(il),
10383 nh,
10384 nkv,
10385 pos_d,
10386 nh,
10387 nkv,
10388 base,
10389 1.0,
10390 ff,
10391 eps,
10392 )?;
10393 e.append_kv_quantized(
10394 &k,
10395 &v,
10396 &mut kvl.k,
10397 &mut kvl.v,
10398 kvl.len,
10399 kvl.kv_dim_k,
10400 kvl.kv_dim_v,
10401 kvl.k_tok_bytes,
10402 kvl.v_tok_bytes,
10403 kv_fp8,
10404 )?;
10405 }
10406 kvl.len += 1;
10407 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10411 let mut attn = e.uninit(nh * hd)?;
10412 if !swa
10414 && hd == 512
10415 && kvl.len >= crate::fa512_min_tkv()
10416 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10417 {
10418 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10419 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10420 let base = kvl.len as i32;
10422 e.i32_set_k(&mut kvl.len_d, base)?;
10423 e.fa_decode_rows(
10424 &q,
10425 &kp,
10426 &vp,
10427 &mut attn,
10428 hd,
10429 nh,
10430 nkv,
10431 kvl.len - 1,
10432 1,
10433 scale,
10434 kvl.k_tok_bytes,
10435 kvl.v_tok_bytes,
10436 Some((&kvl.len_d, -1)),
10437 false,
10438 false,
10439 None,
10440 )?;
10441 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10442 }
10443 if swa
10445 && kvl.len > win
10446 && hd == 256
10447 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10448 {
10449 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10450 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10451 let base = kvl.len as i32;
10452 e.i32_set_k(&mut kvl.len_d, base)?;
10453 e.fa_decode_rows_w(
10454 &q,
10455 &kp,
10456 &vp,
10457 &mut attn,
10458 hd,
10459 nh,
10460 nkv,
10461 &kvl.len_d,
10462 -1,
10463 1,
10464 scale,
10465 win,
10466 kvl.k_tok_bytes,
10467 kvl.v_tok_bytes,
10468 None,
10469 )?;
10470 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10471 }
10472 let (off_tok, t_kv) = if swa && kvl.len > win {
10473 (kvl.len - win, win)
10474 } else {
10475 (0, kvl.len)
10476 };
10477 let k_view = e.view_u8_range(
10478 &kvl.k,
10479 off_tok * kvl.k_tok_bytes,
10480 (off_tok + t_kv) * kvl.k_tok_bytes,
10481 );
10482 let v_view = e.view_u8_range(
10483 &kvl.v,
10484 off_tok * kvl.v_tok_bytes,
10485 (off_tok + t_kv) * kvl.v_tok_bytes,
10486 );
10487 e.fa_decode_kvmod(
10488 &q,
10489 &k_view,
10490 &v_view,
10491 &mut attn,
10492 hd,
10493 nh,
10494 nkv,
10495 t_kv,
10496 scale,
10497 kvl.k_tok_bytes,
10498 kvl.v_tok_bytes,
10499 swa && crate::Engine::wkv_on(),
10500 )?;
10501 Ok(e.matmul(&fa.wo, &attn, 1)?)
10502 }
10503
10504 #[allow(clippy::too_many_arguments)]
10511 pub fn gemma4_decode_step_dc(
10512 &self,
10513 e: &Engine,
10514 token_d: &CudaSlice<u32>,
10515 pos_d: &mut CudaSlice<i32>,
10516 embd_gpu: &CudaSlice<u8>,
10517 embd_qt: i32,
10518 embd_rb: usize,
10519 cache: &mut Cache,
10520 n_vocab: usize,
10521 cap_bucket_max: Option<(usize, usize)>,
10522 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10523 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10524 self.gemma4_decode_step_dc_into(
10525 e,
10526 token_d,
10527 pos_d,
10528 embd_gpu,
10529 embd_qt,
10530 embd_rb,
10531 cache,
10532 n_vocab,
10533 cap_bucket_max,
10534 &mut tok_out,
10535 )?;
10536 Ok(tok_out)
10537 }
10538
10539 #[allow(clippy::too_many_arguments)]
10542 pub fn gemma4_decode_step_dc_into(
10543 &self,
10544 e: &Engine,
10545 token_d: &CudaSlice<u32>,
10546 pos_d: &mut CudaSlice<i32>,
10547 embd_gpu: &CudaSlice<u8>,
10548 embd_qt: i32,
10549 embd_rb: usize,
10550 cache: &mut Cache,
10551 n_vocab: usize,
10552 cap_bucket_max: Option<(usize, usize)>,
10553 tok_out: &mut CudaSlice<u32>,
10554 ) -> Result<(), Box<dyn std::error::Error>> {
10555 let n_embd = self.cfg.n_embd as usize;
10556 let eps = self.cfg.rms_eps;
10557 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10558 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10559 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10560 let n_layers = self.layers.len();
10561 for (il, layer) in self.layers.iter().enumerate() {
10562 let (hq, hdq) = match h_carry.take() {
10563 Some(p) => p,
10564 None => {
10565 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10566 }
10567 };
10568 let Mixer::Full(fa) = &layer.mixer else {
10569 panic!("gemma4 layer {il} not full-attn")
10570 };
10571 let o =
10572 self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10573 let next_norm = if il + 1 < n_layers {
10574 Some(self.layers[il + 1].attn_norm.float_data())
10575 } else {
10576 None
10577 };
10578 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
10579 x = xn;
10580 h_carry = hn;
10581 }
10582 let mut hn = e.uninit(n_embd)?;
10583 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10584 let mut logits = e.matmul(&self.output, &hn, 1)?;
10585 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10587 e.inc_seqlen(pos_d)?;
10588 if cap_bucket_max.is_none() {
10589 cache.pos += 1;
10590 }
10591 Ok(())
10592 }
10593
10594 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10601 let n_embd = self.cfg.n_embd as usize;
10602 let n_vocab = self.output.out_features();
10603 let n_layers = self.layers.len();
10604 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10605 for il in 0..n_layers {
10606 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10607 qmax = qmax.max(nh * hd);
10608 kvmax = kvmax.max(nkv * hd);
10609 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10610 ffmax = ffmax.max(ffn_gate.out_features());
10611 }
10612 }
10613 Ok(G4DcSlots {
10614 x: e.uninit(n_embd)?,
10615 xn: e.uninit(n_embd)?,
10616 cur: e.uninit(n_embd)?,
10617 hq: e.alloc_i8_uninit(n_embd)?,
10618 hd_: e.uninit(n_embd / 32)?,
10619 q0: e.uninit(qmax)?,
10620 k0: e.uninit(kvmax)?,
10621 v0: e.uninit(kvmax)?,
10622 q: e.uninit(qmax)?,
10623 k: e.uninit(kvmax)?,
10624 v: e.uninit(kvmax)?,
10625 attn: e.uninit(qmax)?,
10626 o: e.uninit(n_embd)?,
10627 attn_out: e.uninit(n_embd)?,
10628 zsh: e.uninit(n_embd)?,
10629 zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10632 zd: e.uninit(n_embd.max(qmax) / 32)?,
10633 gate: e.uninit(ffmax)?,
10634 up: e.uninit(ffmax)?,
10635 act: e.uninit(ffmax)?,
10636 actq: e.alloc_i8_uninit(ffmax)?,
10637 actd: e.uninit(ffmax / 32)?,
10638 f0: e.uninit(n_embd)?,
10639 sn: e.uninit(n_embd)?,
10640 hn: e.uninit(n_embd)?,
10641 logits: e.uninit(n_vocab)?,
10642 })
10643 }
10644
10645 fn g4_matvec_m1_into(
10648 &self,
10649 e: &Engine,
10650 w: &crate::model::GpuTensor,
10651 aq: &CudaSlice<i8>,
10652 ad: &CudaSlice<f32>,
10653 y: &mut CudaSlice<f32>,
10654 ) -> Result<(), Box<dyn std::error::Error>> {
10655 use crate::model::GpuTensor;
10656 let (bytes, qtype, row_bytes, scale, rp) = match w {
10657 GpuTensor::Quant {
10658 bytes,
10659 qtype,
10660 row_bytes,
10661 scale,
10662 rp,
10663 ..
10664 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10665 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10666 };
10667 let (mbytes, mrp) = match w {
10668 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10669 _ => (bytes, rp),
10670 };
10671 e.qmatvec_mmvq_into(
10672 mbytes,
10673 aq,
10674 ad,
10675 1,
10676 w.in_features(),
10677 w.out_features(),
10678 qtype,
10679 row_bytes,
10680 scale,
10681 mrp,
10682 y,
10683 )
10684 }
10685
10686 #[allow(clippy::too_many_arguments)]
10690 pub fn gemma4_decode_step_dc_slotted(
10691 &self,
10692 e: &Engine,
10693 token_d: &CudaSlice<u32>,
10694 pos_d: &mut CudaSlice<i32>,
10695 embd_gpu: &CudaSlice<u8>,
10696 embd_qt: i32,
10697 embd_rb: usize,
10698 cache: &mut Cache,
10699 n_vocab: usize,
10700 cap_bucket_max: Option<(usize, usize)>,
10701 sl: &mut G4DcSlots,
10702 tok_out: &mut CudaSlice<u32>,
10703 ring: Option<(&mut CudaSlice<u32>, usize)>,
10704 ) -> Result<(), Box<dyn std::error::Error>> {
10705 let n_embd = self.cfg.n_embd as usize;
10706 let eps = self.cfg.rms_eps;
10707 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10708 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10709 let n_layers = self.layers.len();
10710 let mut has_carry = false;
10711 for il in 0..n_layers {
10712 if !has_carry {
10713 e.rms_norm_q8_1_into(
10714 &sl.x,
10715 self.layers[il].attn_norm.float_data(),
10716 n_embd,
10717 1,
10718 eps,
10719 &mut sl.hq,
10720 &mut sl.hd_,
10721 )?;
10722 }
10723 has_carry = true;
10724 let layer = &self.layers[il];
10725 let Mixer::Full(fa) = &layer.mixer else {
10726 panic!("gemma4 layer {il} not full-attn")
10727 };
10728 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10729 if !Engine::g4_pnfold_on() {
10732 e.rms_norm(
10733 &sl.o,
10734 layer.post_attn_norm.float_data(),
10735 &mut sl.cur,
10736 n_embd,
10737 1,
10738 eps,
10739 )?;
10740 }
10741 let next_norm = if il + 1 < n_layers {
10742 Some(self.layers[il + 1].attn_norm.float_data())
10743 } else {
10744 None
10745 };
10746 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10747 std::mem::swap(&mut sl.x, &mut sl.xn);
10748 }
10749 e.rms_norm(
10750 &sl.x,
10751 self.output_norm.float_data(),
10752 &mut sl.hn,
10753 n_embd,
10754 1,
10755 eps,
10756 )?;
10757 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10758 {
10760 let (zq, zd) = (&sl.zq, &sl.zd);
10761 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10762 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10763 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10764 }
10765 self.gemma4_suppress(e, &mut sl.logits, 1)?;
10766 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10767 if let Some((ring, base)) = ring {
10768 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10772 }
10773 e.inc_seqlen(pos_d)?;
10774 if cap_bucket_max.is_none() {
10775 cache.pos += 1;
10776 }
10777 Ok(())
10778 }
10779
10780 #[allow(clippy::too_many_arguments)]
10782 fn gemma4_decode_attn_dc_slotted(
10783 &self,
10784 e: &Engine,
10785 fa: &crate::hybrid::FullAttnLayer,
10786 il: usize,
10787 pos_d: &CudaSlice<i32>,
10788 cache: &mut Cache,
10789 cap_bucket_max: Option<(usize, usize)>,
10790 sl: &mut G4DcSlots,
10791 ) -> Result<(), Box<dyn std::error::Error>> {
10792 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10793 let eps = self.cfg.rms_eps;
10794 let aux = self.gemma4_aux.as_ref().unwrap();
10795 let ones = aux.ones(e);
10796 #[cfg(debug_assertions)]
10797 crate::debug_assert_tensor_stream_device(
10798 ones,
10799 &e.stream(),
10800 "gemma4_decode_attn_dc_slotted.ones",
10801 );
10802 {
10803 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10804 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10805 if swa {
10806 if !e.matmul_q4_fused3_into(
10807 &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10808 )? {
10809 if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10813 {
10814 self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
10815 } else {
10816 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10817 }
10818 }
10819 } else {
10820 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10821 && !e
10822 .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10823 {
10824 return Err("slotted step: fused2 unavailable".into());
10825 }
10826 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10827 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10828 }
10829 }
10830 let ff = if swa {
10833 None
10834 } else {
10835 Some(
10836 aux.rope_freqs(e)
10837 .expect("gemma4 global rope needs rope_freqs.weight"),
10838 )
10839 };
10840 #[cfg(debug_assertions)]
10841 if let Some(ff) = ff {
10842 crate::debug_assert_tensor_stream_device(
10843 ff,
10844 &e.stream(),
10845 "gemma4_decode_attn_dc_slotted.rope_freqs",
10846 );
10847 }
10848 let kvl = cache.kv[il].as_mut().unwrap();
10849 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10850 if crate::Engine::qkv_append_on() {
10851 e.rms_norm_qkv_rope_append_dc(
10853 &sl.q0,
10854 &sl.k0,
10855 &sl.v0,
10856 fa.q_norm.float_data(),
10857 fa.k_norm.float_data(),
10858 ones,
10859 &mut sl.q,
10860 &mut sl.k,
10861 &mut sl.v,
10862 hd,
10863 self.gemma4_rope_dims(il),
10864 nh,
10865 nkv,
10866 pos_d,
10867 nh,
10868 nkv,
10869 base,
10870 1.0,
10871 ff,
10872 eps,
10873 &mut kvl.k,
10874 &mut kvl.v,
10875 &kvl.len_d,
10876 kvl.k_tok_bytes,
10877 kvl.v_tok_bytes,
10878 kv_fp8,
10879 )?;
10880 } else {
10881 e.rms_norm_qkv_rope(
10882 &sl.q0,
10883 &sl.k0,
10884 &sl.v0,
10885 fa.q_norm.float_data(),
10886 fa.k_norm.float_data(),
10887 ones,
10888 &mut sl.q,
10889 &mut sl.k,
10890 &mut sl.v,
10891 hd,
10892 self.gemma4_rope_dims(il),
10893 nh,
10894 nkv,
10895 pos_d,
10896 nh,
10897 nkv,
10898 base,
10899 1.0,
10900 ff,
10901 eps,
10902 )?;
10903 e.append_kv_quantized_dc(
10904 &sl.k,
10905 &sl.v,
10906 &mut kvl.k,
10907 &mut kvl.v,
10908 &kvl.len_d,
10909 kvl.kv_dim_k,
10910 kvl.kv_dim_v,
10911 kvl.k_tok_bytes,
10912 kvl.v_tok_bytes,
10913 kv_fp8,
10914 )?;
10915 }
10916 e.inc_seqlen(&mut kvl.len_d)?;
10917 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
10918 let k_view = e.view_u8(&kvl.k, kvl.k.len());
10919 let v_view = e.view_u8(&kvl.v, kvl.v.len());
10920 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10921 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10922 let mut fa_q8 = false;
10926 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10927 e.fa_decode_rows(
10928 &sl.q,
10929 &k_view,
10930 &v_view,
10931 &mut sl.attn,
10932 hd,
10933 nh,
10934 nkv,
10935 b_glob - 1,
10936 1,
10937 scale,
10938 kvl.k_tok_bytes,
10939 kvl.v_tok_bytes,
10940 Some((&kvl.len_d, -1)),
10941 false,
10942 false,
10943 Some((&mut sl.zq, &mut sl.zd)),
10944 )?;
10945 fa_q8 = true;
10946 } else if swa && b_swa > win && hd == 256 && rows_on {
10947 e.fa_decode_rows_w(
10948 &sl.q,
10949 &k_view,
10950 &v_view,
10951 &mut sl.attn,
10952 hd,
10953 nh,
10954 nkv,
10955 &kvl.len_d,
10956 -1,
10957 1,
10958 scale,
10959 win,
10960 kvl.k_tok_bytes,
10961 kvl.v_tok_bytes,
10962 Some((&mut sl.zq, &mut sl.zd)),
10963 )?;
10964 fa_q8 = true;
10965 } else {
10966 let b = if swa { b_swa } else { b_glob };
10967 e.fa_decode_dc(
10968 &sl.q,
10969 &k_view,
10970 &v_view,
10971 &mut sl.attn,
10972 hd,
10973 nh,
10974 nkv,
10975 &kvl.len_d,
10976 b,
10977 scale,
10978 kvl.k_tok_bytes,
10979 kvl.v_tok_bytes,
10980 swa && crate::Engine::wkv_on(),
10981 )?;
10982 }
10983 if !fa_q8 {
10984 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
10985 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
10986 }
10987 {
10988 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10989 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10990 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
10991 }
10992 Ok(())
10993 }
10994
10995 fn gemma4_layer_tail_slotted(
10998 &self,
10999 e: &Engine,
11000 layer: &crate::hybrid::HybridLayer,
11001 next_norm: Option<&CudaSlice<f32>>,
11002 sl: &mut G4DcSlots,
11003 ) -> Result<(), Box<dyn std::error::Error>> {
11004 let n_embd = self.cfg.n_embd as usize;
11005 let eps = self.cfg.rms_eps;
11006 let bits = layer.gemma4.as_ref().unwrap();
11007 let crate::hybrid::Ffn::Dense {
11008 ffn_gate,
11009 ffn_up,
11010 ffn_down,
11011 } = &layer.ffn
11012 else {
11013 return Err("slotted tail: dense ffn only".into());
11014 };
11015 let pnfold = Engine::g4_pnfold_on();
11016 if pnfold {
11017 let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
11020 let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
11021 e.rms_pre_add_rms_norm_q8z_into(
11022 or,
11023 layer.post_attn_norm.float_data(),
11024 xr,
11025 bits.ffn_norm.float_data(),
11026 &mut sl.attn_out,
11027 &mut sl.zsh,
11028 n_embd,
11029 1,
11030 eps,
11031 &mut sl.zq,
11032 &mut sl.zd,
11033 )?;
11034 } else {
11035 e.add_rms_norm(
11036 &sl.cur,
11037 &sl.x,
11038 bits.ffn_norm.float_data(),
11039 &mut sl.attn_out,
11040 &mut sl.zsh,
11041 n_embd,
11042 1,
11043 eps,
11044 )?;
11045 }
11046 let n_ff = ffn_gate.out_features();
11047 if !pnfold {
11048 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
11049 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
11050 }
11051 {
11052 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
11053 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11054 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
11055 && !e.matmul_nvfp4_fused2_into(
11056 ffn_gate,
11057 ffn_up,
11058 zq,
11059 zd,
11060 &mut sl.gate,
11061 &mut sl.up,
11062 )?
11063 {
11064 return Err("slotted tail: ffn fused2 unavailable".into());
11065 }
11066 }
11067 debug_assert!(e.uses_q8_1_fast(ffn_down));
11068 {
11069 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
11070 let upv = e.view(upr, n_ff);
11071 let up_all = upv.slice(0..n_ff);
11072 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
11073 e.gelu_tanh_mul_q8_1_into(
11074 gr,
11075 &up_all,
11076 &mut sl.act,
11077 n_ff,
11078 1,
11079 &mut sl.actq,
11080 &mut sl.actd,
11081 )?;
11082 }
11083 {
11084 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
11085 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
11086 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
11087 }
11088 if pnfold {
11089 if let Some(w) = next_norm {
11092 let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
11093 let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
11094 e.rms_pre_add_scale_rms_norm_q8_1_into(
11095 f0r,
11096 bits.post_ffw_norm.float_data(),
11097 aor,
11098 bits.layer_scale,
11099 w,
11100 &mut sl.xn,
11101 n_embd,
11102 1,
11103 eps,
11104 &mut sl.hq,
11105 &mut sl.hd_,
11106 )?;
11107 return Ok(());
11108 }
11109 }
11110 e.rms_norm(
11111 &sl.f0,
11112 bits.post_ffw_norm.float_data(),
11113 &mut sl.sn,
11114 n_embd,
11115 1,
11116 eps,
11117 )?;
11118 match next_norm {
11119 Some(w) => {
11120 e.add_scale_rms_norm_q8_1_into(
11121 &sl.sn,
11122 &sl.attn_out,
11123 bits.layer_scale,
11124 w,
11125 &mut sl.xn,
11126 n_embd,
11127 1,
11128 eps,
11129 &mut sl.hq,
11130 &mut sl.hd_,
11131 )?;
11132 }
11133 None => {
11134 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
11135 }
11136 }
11137 Ok(())
11138 }
11139
11140 #[allow(clippy::too_many_arguments)]
11142 fn gemma4_decode_attn_dc(
11143 &self,
11144 e: &Engine,
11145 fa: &crate::hybrid::FullAttnLayer,
11146 il: usize,
11147 hq: &CudaSlice<i8>,
11148 hdq: &CudaSlice<f32>,
11149 pos_d: &CudaSlice<i32>,
11150 cache: &mut Cache,
11151 cap_bucket_max: Option<(usize, usize)>,
11152 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11153 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11154 let eps = self.cfg.rms_eps;
11155 let aux = self.gemma4_aux.as_ref().unwrap();
11156 let ones = aux.ones(e);
11157 #[cfg(debug_assertions)]
11158 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
11159 let (q0, k0, v0) = if swa {
11160 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
11161 Some(t3) => t3,
11162 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11164 Some((q0, k0)) => {
11165 let h0 = e.zeros(0)?;
11166 let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
11167 (q0, k0, v0)
11168 }
11169 None => {
11170 let h0 = e.zeros(0)?;
11171 (
11172 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11173 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11174 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
11175 )
11176 }
11177 },
11178 }
11179 } else {
11180 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
11181 Some(p) => p,
11182 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11183 Some(p) => p,
11184 None => {
11185 let h0 = e.zeros(0)?;
11186 (
11187 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11188 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11189 )
11190 }
11191 },
11192 };
11193 let v0 = e.clone_dtod(&k0)?;
11194 (q0, k0, v0)
11195 };
11196 let mut q = e.uninit(nh * hd)?;
11197 let mut k = e.uninit(nkv * hd)?;
11198 let mut v = e.uninit(nkv * hd)?;
11199 let ff = if swa {
11201 None
11202 } else {
11203 Some(
11204 aux.rope_freqs(e)
11205 .expect("gemma4 global rope needs rope_freqs.weight"),
11206 )
11207 };
11208 #[cfg(debug_assertions)]
11209 if let Some(ff) = ff {
11210 crate::debug_assert_tensor_stream_device(
11211 ff,
11212 &e.stream(),
11213 "gemma4_decode_attn_dc.rope_freqs",
11214 );
11215 }
11216 let kvl = cache.kv[il].as_mut().unwrap();
11217 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
11218 if crate::Engine::qkv_append_on() {
11219 e.rms_norm_qkv_rope_append_dc(
11221 &q0,
11222 &k0,
11223 &v0,
11224 fa.q_norm.float_data(),
11225 fa.k_norm.float_data(),
11226 ones,
11227 &mut q,
11228 &mut k,
11229 &mut v,
11230 hd,
11231 self.gemma4_rope_dims(il),
11232 nh,
11233 nkv,
11234 pos_d,
11235 nh,
11236 nkv,
11237 base,
11238 1.0,
11239 ff,
11240 eps,
11241 &mut kvl.k,
11242 &mut kvl.v,
11243 &kvl.len_d,
11244 kvl.k_tok_bytes,
11245 kvl.v_tok_bytes,
11246 kv_fp8,
11247 )?;
11248 } else {
11249 e.rms_norm_qkv_rope(
11250 &q0,
11251 &k0,
11252 &v0,
11253 fa.q_norm.float_data(),
11254 fa.k_norm.float_data(),
11255 ones,
11256 &mut q,
11257 &mut k,
11258 &mut v,
11259 hd,
11260 self.gemma4_rope_dims(il),
11261 nh,
11262 nkv,
11263 pos_d,
11264 nh,
11265 nkv,
11266 base,
11267 1.0,
11268 ff,
11269 eps,
11270 )?;
11271 e.append_kv_quantized_dc(
11272 &k,
11273 &v,
11274 &mut kvl.k,
11275 &mut kvl.v,
11276 &kvl.len_d,
11277 kvl.kv_dim_k,
11278 kvl.kv_dim_v,
11279 kvl.k_tok_bytes,
11280 kvl.v_tok_bytes,
11281 kv_fp8,
11282 )?;
11283 }
11284 e.inc_seqlen(&mut kvl.len_d)?;
11285 let mut attn = e.uninit(nh * hd)?;
11286 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11289 match cap_bucket_max {
11294 None => {
11295 kvl.len += 1;
11299 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11300 if !swa
11301 && hd == 512
11302 && kvl.len >= crate::fa512_min_tkv()
11303 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11304 {
11305 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11308 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11309 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11310 e.fa_decode_rows(
11311 &q,
11312 &kp,
11313 &vp,
11314 &mut attn,
11315 hd,
11316 nh,
11317 nkv,
11318 kvl.len - 1,
11319 1,
11320 scale,
11321 kvl.k_tok_bytes,
11322 kvl.v_tok_bytes,
11323 Some((&kvl.len_d, -1)),
11324 false,
11325 false,
11326 Some((&mut aq8, &mut ad8)),
11327 )?;
11328 fa_q8 = Some((aq8, ad8));
11329 } else if swa
11330 && kvl.len > win
11331 && hd == 256
11332 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11333 {
11334 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11336 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11337 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11338 e.fa_decode_rows_w(
11339 &q,
11340 &kp,
11341 &vp,
11342 &mut attn,
11343 hd,
11344 nh,
11345 nkv,
11346 &kvl.len_d,
11347 -1,
11348 1,
11349 scale,
11350 win,
11351 kvl.k_tok_bytes,
11352 kvl.v_tok_bytes,
11353 Some((&mut aq8, &mut ad8)),
11354 )?;
11355 fa_q8 = Some((aq8, ad8));
11356 } else {
11357 let (off_tok, t_kv) = if swa && kvl.len > win {
11358 (kvl.len - win, win)
11359 } else {
11360 (0, kvl.len)
11361 };
11362 let k_view = e.view_u8_range(
11363 &kvl.k,
11364 off_tok * kvl.k_tok_bytes,
11365 (off_tok + t_kv) * kvl.k_tok_bytes,
11366 );
11367 let v_view = e.view_u8_range(
11368 &kvl.v,
11369 off_tok * kvl.v_tok_bytes,
11370 (off_tok + t_kv) * kvl.v_tok_bytes,
11371 );
11372 e.fa_decode_kvmod(
11373 &q,
11374 &k_view,
11375 &v_view,
11376 &mut attn,
11377 hd,
11378 nh,
11379 nkv,
11380 t_kv,
11381 scale,
11382 kvl.k_tok_bytes,
11383 kvl.v_tok_bytes,
11384 swa && crate::Engine::wkv_on(),
11385 )?;
11386 }
11387 }
11388 Some((b_swa, b_glob)) => {
11389 let k_view = e.view_u8(&kvl.k, kvl.k.len());
11395 let v_view = e.view_u8(&kvl.v, kvl.v.len());
11396 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11397 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11398 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11399 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11400 e.fa_decode_rows(
11401 &q,
11402 &k_view,
11403 &v_view,
11404 &mut attn,
11405 hd,
11406 nh,
11407 nkv,
11408 b_glob - 1,
11409 1,
11410 scale,
11411 kvl.k_tok_bytes,
11412 kvl.v_tok_bytes,
11413 Some((&kvl.len_d, -1)),
11414 false,
11415 false,
11416 Some((&mut aq8, &mut ad8)),
11417 )?;
11418 fa_q8 = Some((aq8, ad8));
11419 } else if swa && b_swa > win && hd == 256 && rows_on {
11420 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11421 e.fa_decode_rows_w(
11422 &q,
11423 &k_view,
11424 &v_view,
11425 &mut attn,
11426 hd,
11427 nh,
11428 nkv,
11429 &kvl.len_d,
11430 -1,
11431 1,
11432 scale,
11433 win,
11434 kvl.k_tok_bytes,
11435 kvl.v_tok_bytes,
11436 Some((&mut aq8, &mut ad8)),
11437 )?;
11438 fa_q8 = Some((aq8, ad8));
11439 } else {
11440 let b = if swa { b_swa } else { b_glob };
11441 e.fa_decode_dc(
11442 &q,
11443 &k_view,
11444 &v_view,
11445 &mut attn,
11446 hd,
11447 nh,
11448 nkv,
11449 &kvl.len_d,
11450 b,
11451 scale,
11452 kvl.k_tok_bytes,
11453 kvl.v_tok_bytes,
11454 swa && crate::Engine::wkv_on(),
11455 )?;
11456 }
11457 }
11458 }
11459 if let Some((aq8, ad8)) = fa_q8 {
11462 let mut y = e.uninit(fa.wo.out_features())?;
11463 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11464 return Ok(y);
11465 }
11466 Ok(e.matmul(&fa.wo, &attn, 1)?)
11467 }
11468
11469 pub fn gemma4_generate_graph(
11474 &self,
11475 e: &Engine,
11476 prompt_pos: usize,
11477 first_token: u32,
11478 cache: &mut Cache,
11479 max_new: usize,
11480 eos: &[u32],
11481 mut on_token: impl FnMut(u32) -> bool,
11482 ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11483 if self.is_gemma4_e4b() {
11484 return Err(
11485 "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11486 .into(),
11487 );
11488 }
11489 use crate::decode::StopReason;
11490 let n_vocab = self.output.out_features();
11491 let n_embd = self.cfg.n_embd as usize;
11492 let embd_gpu = self
11493 .embd_gpu
11494 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11495 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11496 for kvl in cache.kv.iter_mut().flatten() {
11497 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11498 }
11499 let mut token_d = e.stream().clone_htod(&[first_token])?;
11500 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11501 let g4 = self.cfg.gemma4.as_ref().unwrap();
11502 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11503 let nkv_s = g4
11505 .head_count_kv
11506 .iter()
11507 .zip(g4.swa_pattern.iter())
11508 .find(|p| *p.1)
11509 .map(|p| *p.0 as usize)
11510 .unwrap_or(8);
11511 let nkv_g = g4
11512 .head_count_kv
11513 .iter()
11514 .zip(g4.swa_pattern.iter())
11515 .find(|p| !*p.1)
11516 .map(|p| *p.0 as usize)
11517 .unwrap_or(2);
11518 let mut graphs: std::collections::HashMap<
11519 ((bool, usize), (bool, usize), bool, bool),
11520 (
11521 cudarc::driver::CudaGraph,
11522 Vec<Box<dyn std::any::Any + Send>>,
11523 ),
11524 > = Default::default();
11525 let mut slots = self.g4_dc_slots(e)?;
11528 const RING: usize = 64;
11531 const DRAIN: usize = 1;
11537 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11538 let ring_base = prompt_pos;
11539 let mut out = Vec::with_capacity(max_new);
11540 let mut reason = StopReason::MaxNew;
11541 let mut next = first_token;
11542 let mut captures = 0usize;
11543 for _ in 0..max_new {
11544 out.push(next);
11545 if eos.contains(&next) {
11546 reason = StopReason::Eos;
11547 break;
11548 }
11549 if !on_token(next) {
11550 reason = StopReason::Callback;
11551 break;
11552 }
11553 let t_kv = cache.pos + 1;
11554 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11562 let f512 = crate::fa512_min_tkv();
11563 let key_s = if t_kv > win {
11564 (true, usize::MAX)
11565 } else {
11566 e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11567 };
11568 let (key_g, rung_end) = if t_kv >= f512 {
11569 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11572 ((true, end), end)
11573 } else {
11574 (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11575 };
11576 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11577 if !graphs.contains_key(&key) {
11578 let bucket_max = (t_kv, rung_end);
11579 let snap = cache.snapshot(e)?;
11581 let pos_save = e.dtoh_i32_one(&pos_d)?;
11582 let len_save: Vec<Option<i32>> = cache
11583 .kv
11584 .iter()
11585 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11586 .collect();
11587 let tok_save = e.dtoh_u32_one(&token_d)?;
11588 let graph = {
11593 let tok_ref = &mut token_d;
11594 let pos_ref = &mut pos_d;
11595 let cache_ref = &mut *cache;
11596 let slots_ref = &mut slots;
11597 let ring_ref = &mut ring;
11598 e.capture_graph_retained_flags(
11599 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11600 |e| {
11601 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11603 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11604 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11605 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11606 cache_ref, n_vocab, Some(bucket_max),
11607 sl, tok_ref, Some((rg, ring_base)))
11608 })?
11609 };
11610 cache.rollback(e, &snap, 0)?;
11611 e.set_i32_one(&mut pos_d, pos_save)?;
11612 for (il, ls) in len_save.iter().enumerate() {
11613 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11614 e.set_i32_one(&mut kvl.len_d, *v)?;
11615 }
11616 }
11617 e.set_u32_one(&mut token_d, tok_save)?;
11618 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11619 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11620 eprintln!("[graph-census] {c:?}");
11621 }
11622 }
11623 graphs.insert(key, graph);
11624 captures += 1;
11625 }
11626 let mut chunk = 1usize;
11631 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11632 .ok()
11633 .and_then(|v| v.parse().ok())
11634 .unwrap_or(DRAIN);
11635 while chunk < drain_cap && out.len() + chunk < max_new {
11636 let t_next = cache.pos + 1 + chunk;
11637 let key_s2 = if t_next > win {
11638 (true, usize::MAX)
11639 } else {
11640 e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11641 };
11642 let key_g2 = if t_next >= f512 {
11643 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11644 } else {
11645 e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11646 };
11647 if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11648 break;
11649 }
11650 chunk += 1;
11651 }
11652 let g = &graphs.get(&key).unwrap().0;
11653 for _ in 0..chunk {
11654 g.launch()?;
11655 }
11656 e.stream().synchronize()?;
11657 let ringh = e.dtoh_u32(&ring)?;
11658 for j in 0..chunk {
11659 let pos_j = cache.pos + j;
11660 let tok_j = ringh[(pos_j - ring_base) % RING];
11661 cache.pos += 0; if j + 1 == chunk {
11663 next = tok_j;
11664 } else {
11665 out.push(tok_j);
11666 if eos.contains(&tok_j) || !on_token(tok_j) {
11667 reason = if eos.contains(&tok_j) {
11668 StopReason::Eos
11669 } else {
11670 StopReason::Callback
11671 };
11672 let keep = cache.pos + j + 1;
11674 e.set_i32_one(&mut pos_d, keep as i32)?;
11675 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11676 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11677 kvl.len = keep;
11678 }
11679 cache.pos = keep;
11680 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11681 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11682 }
11683 return Ok((out, reason));
11684 }
11685 }
11686 }
11687 cache.pos += chunk;
11688 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11689 kvl.len += chunk;
11690 }
11691 }
11692 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11693 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11694 }
11695 Ok((out, reason))
11696 }
11697
11698 pub(crate) fn gemma4_decode_step_t(
11704 &self,
11705 e: &Engine,
11706 tokens: &[u32],
11707 pos0: usize,
11708 cache: &mut Cache,
11709 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11710 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11711 }
11712
11713 pub(crate) fn gemma4_decode_step_t_am(
11717 &self,
11718 e: &Engine,
11719 tokens: &[u32],
11720 pos0: usize,
11721 cache: &mut Cache,
11722 ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11723 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11724 let t = tokens.len();
11725 let n_vocab = self.output.out_features();
11726 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11727 for i in 0..t {
11728 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11729 }
11730 Ok((e.dtoh_u32(&toks)?, hn))
11731 }
11732
11733 pub(crate) fn gemma4_decode_step_t_am_dev(
11736 &self,
11737 e: &Engine,
11738 tok_d: &CudaSlice<u32>,
11739 t: usize,
11740 pos0: usize,
11741 cache: &mut Cache,
11742 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11743 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11744 let n_vocab = self.output.out_features();
11745 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11746 for i in 0..t {
11747 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11748 }
11749 Ok((vam, hn))
11750 }
11751
11752 pub(crate) fn gemma4_decode_step_t_h(
11755 &self,
11756 e: &Engine,
11757 tokens: &[u32],
11758 pos0: usize,
11759 cache: &mut Cache,
11760 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11761 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11762 let t = tokens.len();
11763 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11764 e.softcap(&mut ld, cap, t * self.output.out_features())?;
11765 Ok((e.dtoh(&ld)?, hn))
11766 }
11767
11768 pub(crate) fn verify_stream_scratch(
11771 &self,
11772 e: &Engine,
11773 cap: usize,
11774 ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11775 Ok(VerifyStreamScratch {
11776 pos_d: e.htod_i32(&vec![0i32; cap])?,
11777 row_ctrs: (0..cap)
11778 .map(|_| e.htod_i32(&[0]))
11779 .collect::<Result<_, _>>()?,
11780 })
11781 }
11782
11783 pub(crate) fn gemma4_verify_t_am_stream(
11791 &self,
11792 e: &Engine,
11793 tok_d: &CudaSlice<u32>,
11794 t: usize,
11795 ctr: &CudaSlice<i32>,
11796 hint: usize,
11797 cache: &mut Cache,
11798 scr: &mut VerifyStreamScratch,
11799 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11800 let n_embd = self.cfg.n_embd as usize;
11801 let eps = self.cfg.rms_eps;
11802 assert!(t <= scr.row_ctrs.len() && t <= 64);
11803 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11804 for i in 0..t {
11805 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11806 }
11807 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11808 let embd_gpu = self
11809 .embd_gpu
11810 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11811 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11812 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11813 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11814 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11815 let n_layers = self.layers.len();
11816 for (il, layer) in self.layers.iter().enumerate() {
11817 let (hq, hdq) = match h_carry.take() {
11818 Some(p) => p,
11819 None => {
11820 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11821 }
11822 };
11823 let Mixer::Full(fa) = &layer.mixer else {
11824 panic!("gemma4 layer {il} not full-attn")
11825 };
11826 let o = self
11827 .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11828 let next_norm = if il + 1 < n_layers {
11829 Some(self.layers[il + 1].attn_norm.float_data())
11830 } else {
11831 None
11832 };
11833 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11834 x = xn;
11835 h_carry = hn;
11836 self.dflash_tap(e, cache, il, &x, t)?;
11837 }
11838 let mut hn = e.uninit(t * n_embd)?;
11839 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11840 let ld = e.matmul(&self.output, &hn, t)?;
11841 let n_vocab = self.output.out_features();
11842 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11843 for i in 0..t {
11844 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11845 }
11846 Ok((vam, hn))
11847 }
11848
11849 pub(crate) fn dflash_tap(
11856 &self,
11857 e: &Engine,
11858 cache: &mut Cache,
11859 il: usize,
11860 x: &CudaSlice<f32>,
11861 t: usize,
11862 ) -> Result<(), Box<dyn std::error::Error>> {
11863 let Some(taps) = cache.dflash_taps.as_mut() else {
11864 return Ok(());
11865 };
11866 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11867 return Ok(());
11868 };
11869 let h = taps.hidden;
11870 let n_taps = taps.layer_ids.len();
11871 let base = taps.base;
11872 debug_assert!(
11873 base + t <= taps.t,
11874 "tap window {base}+{t} exceeds sink {}",
11875 taps.t
11876 );
11877 let xv = e.view(x, t * h);
11878 for r in 0..t {
11879 let row = xv.slice(r * h..(r + 1) * h);
11880 e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
11881 }
11882 Ok(())
11883 }
11884
11885 fn gemma4_verify_trunk(
11886 &self,
11887 e: &Engine,
11888 tokens: &[u32],
11889 pos0: usize,
11890 cache: &mut Cache,
11891 tok_dev: Option<&CudaSlice<u32>>,
11892 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11893 let n_embd = self.cfg.n_embd as usize;
11894 let eps = self.cfg.rms_eps;
11895 let t = tokens.len();
11896 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
11897 let pos_d = e.htod_i32(&pos)?;
11898 let mut x = match tok_dev {
11899 Some(td) => {
11900 let embd_gpu = self
11901 .embd_gpu
11902 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11903 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11904 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
11905 }
11906 None => e.htod(&self.embd.gather(n_embd, tokens))?,
11907 };
11908 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11909 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11910 let n_layers = self.layers.len();
11911 for (il, layer) in self.layers.iter().enumerate() {
11912 let (hq, hdq) = match h_carry.take() {
11913 Some(p) => p,
11914 None => {
11915 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11916 }
11917 };
11918 let Mixer::Full(fa) = &layer.mixer else {
11919 panic!("gemma4 layer {il} not full-attn")
11920 };
11921 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
11922 let next_norm = if il + 1 < n_layers {
11923 Some(self.layers[il + 1].attn_norm.float_data())
11924 } else {
11925 None
11926 };
11927 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11928 x = xn;
11929 h_carry = hn;
11930 self.dflash_tap(e, cache, il, &x, t)?;
11931 }
11932 let mut hn = e.uninit(t * n_embd)?;
11933 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11934 let mut ld = e.matmul(&self.output, &hn, t)?;
11935 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
11937 Ok((ld, hn))
11938 }
11939
11940 #[allow(clippy::too_many_arguments)]
11948 fn gemma4_verify_attn_stream(
11949 &self,
11950 e: &Engine,
11951 fa: &crate::hybrid::FullAttnLayer,
11952 il: usize,
11953 hq: &CudaSlice<i8>,
11954 hdq: &CudaSlice<f32>,
11955 pos_d: &CudaSlice<i32>,
11956 t: usize,
11957 cache: &mut Cache,
11958 hint: usize,
11959 row_ctrs: &[CudaSlice<i32>],
11960 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11961 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11962 let eps = self.cfg.rms_eps;
11963 let aux = self.gemma4_aux.as_ref().unwrap();
11964 let ones = aux.ones(e);
11965 #[cfg(debug_assertions)]
11966 crate::debug_assert_tensor_stream_device(
11967 ones,
11968 &e.stream(),
11969 "gemma4_verify_attn_stream.ones",
11970 );
11971 let h0 = e.zeros(0)?;
11972 let h = &h0;
11973 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11976 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11977 let fused_qkv = if f2b {
11978 if swa {
11979 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11980 .map(|(a, b, c)| (a, b, Some(c)))
11981 } else {
11982 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11983 .map(|(a, b)| (a, b, None))
11984 }
11985 } else {
11986 None
11987 };
11988 let (q0, k0, v0) = match fused_qkv {
11989 Some((a, b, cv)) => {
11990 let v = match cv {
11991 Some(c) => c,
11992 None => e.clone_dtod(&b)?,
11993 };
11994 (a, b, v)
11995 }
11996 None => {
11997 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11998 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11999 let v0 = if swa {
12000 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12001 } else {
12002 e.clone_dtod(&k0)?
12003 };
12004 (q0, k0, v0)
12005 }
12006 };
12007 let mut q = e.uninit(t * nh * hd)?;
12008 let mut k = e.uninit(t * nkv * hd)?;
12009 let mut v = e.uninit(t * nkv * hd)?;
12010 let ff = if swa {
12013 None
12014 } else {
12015 Some(
12016 aux.rope_freqs(e)
12017 .expect("gemma4 global rope needs rope_freqs.weight"),
12018 )
12019 };
12020 #[cfg(debug_assertions)]
12021 if let Some(ff) = ff {
12022 crate::debug_assert_tensor_stream_device(
12023 ff,
12024 &e.stream(),
12025 "gemma4_verify_attn_stream.rope_freqs",
12026 );
12027 }
12028 e.rms_norm_qkv_rope(
12029 &q0,
12030 &k0,
12031 &v0,
12032 fa.q_norm.float_data(),
12033 fa.k_norm.float_data(),
12034 ones,
12035 &mut q,
12036 &mut k,
12037 &mut v,
12038 hd,
12039 self.gemma4_rope_dims(il),
12040 nh * t,
12041 nkv * t,
12042 pos_d,
12043 nh,
12044 nkv,
12045 base,
12046 1.0,
12047 ff,
12048 eps,
12049 )?;
12050 let kvl = cache.kv[il].as_mut().unwrap();
12051 e.append_kv_quantized_rows_dc(
12053 &k,
12054 &v,
12055 &mut kvl.k,
12056 &mut kvl.v,
12057 &kvl.len_d,
12058 t,
12059 kvl.kv_dim_k,
12060 kvl.kv_dim_v,
12061 kvl.k_tok_bytes,
12062 kvl.v_tok_bytes,
12063 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12064 )?;
12065 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12068 let mut attn = e.uninit(t * nh * hd)?;
12069 let k_view = e.view_u8(&kvl.k, kvl.k.len());
12070 let v_view = e.view_u8(&kvl.v, kvl.v.len());
12071 if swa && hint + 1 >= win {
12074 e.fa_decode_rows_w(
12077 &q,
12078 &k_view,
12079 &v_view,
12080 &mut attn,
12081 hd,
12082 nh,
12083 nkv,
12084 &kvl.len_d,
12085 0,
12086 t,
12087 scale,
12088 win,
12089 kvl.k_tok_bytes,
12090 kvl.v_tok_bytes,
12091 None,
12092 )?;
12093 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
12094 let bucket = (hint + t + 2)
12107 .next_power_of_two()
12108 .min(crate::fa512_min_tkv().saturating_sub(1));
12109 let qv = e.view(&q, t * nh * hd);
12110 for i in 0..t {
12111 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
12112 let mut q_one = e.uninit(nh * hd)?;
12113 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12114 let mut a_one = e.uninit(nh * hd)?;
12115 e.fa_decode_dc(
12116 &q_one,
12117 &k_view,
12118 &v_view,
12119 &mut a_one,
12120 hd,
12121 nh,
12122 nkv,
12123 &row_ctrs[i],
12124 bucket,
12125 scale,
12126 kvl.k_tok_bytes,
12127 kvl.v_tok_bytes,
12128 false,
12129 )?;
12130 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12131 }
12132 } else if hd == 512 {
12133 e.fa_decode_rows(
12136 &q,
12137 &k_view,
12138 &v_view,
12139 &mut attn,
12140 hd,
12141 nh,
12142 nkv,
12143 hint,
12144 t,
12145 scale,
12146 kvl.k_tok_bytes,
12147 kvl.v_tok_bytes,
12148 Some((&kvl.len_d, 0)),
12149 false,
12150 false,
12151 None,
12152 )?;
12153 } else {
12154 e.fa_decode_rows_dc(
12156 &q,
12157 &k_view,
12158 &v_view,
12159 &mut attn,
12160 hd,
12161 nh,
12162 nkv,
12163 &kvl.len_d,
12164 hint + t,
12165 t,
12166 scale,
12167 kvl.k_tok_bytes,
12168 kvl.v_tok_bytes,
12169 0,
12170 swa && crate::Engine::wkv_on(),
12171 )?;
12172 }
12173 Ok(e.matmul(&fa.wo, &attn, t)?)
12174 }
12175
12176 fn gemma4_verify_attn(
12177 &self,
12178 e: &Engine,
12179 fa: &crate::hybrid::FullAttnLayer,
12180 il: usize,
12181 hq: &CudaSlice<i8>,
12182 hdq: &CudaSlice<f32>,
12183 pos_d: &CudaSlice<i32>,
12184 t: usize,
12185 cache: &mut Cache,
12186 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12187 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12188 let eps = self.cfg.rms_eps;
12189 let aux = self.gemma4_aux.as_ref().unwrap();
12190 let ones = aux.ones(e);
12191 #[cfg(debug_assertions)]
12192 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
12193 let n_embd = self.cfg.n_embd as usize;
12194 let _ = n_embd;
12195
12196 let h0 = e.zeros(0)?;
12197 let h = &h0;
12198 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12201 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12202 let fused_qkv = if f2b {
12203 if swa {
12204 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12205 .map(|(a, b, c)| (a, b, Some(c)))
12206 } else {
12207 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12208 .map(|(a, b)| (a, b, None))
12209 }
12210 } else {
12211 None
12212 };
12213 let (q0, k0, v0) = match fused_qkv {
12214 Some((a, b, cv)) => {
12215 let v = match cv {
12216 Some(c) => c,
12217 None => e.clone_dtod(&b)?,
12218 };
12219 (a, b, v)
12220 }
12221 None => {
12222 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12223 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12224 let v0 = if swa {
12225 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12226 } else {
12227 e.clone_dtod(&k0)?
12228 };
12229 (q0, k0, v0)
12230 }
12231 };
12232 let mut q = e.uninit(t * nh * hd)?;
12233 let mut k = e.uninit(t * nkv * hd)?;
12234 let mut v = e.uninit(t * nkv * hd)?;
12235 let ff = if swa {
12238 None
12239 } else {
12240 Some(
12241 aux.rope_freqs(e)
12242 .expect("gemma4 global rope needs rope_freqs.weight"),
12243 )
12244 };
12245 #[cfg(debug_assertions)]
12246 if let Some(ff) = ff {
12247 crate::debug_assert_tensor_stream_device(
12248 ff,
12249 &e.stream(),
12250 "gemma4_verify_attn.rope_freqs",
12251 );
12252 }
12253 e.rms_norm_qkv_rope(
12254 &q0,
12255 &k0,
12256 &v0,
12257 fa.q_norm.float_data(),
12258 fa.k_norm.float_data(),
12259 ones,
12260 &mut q,
12261 &mut k,
12262 &mut v,
12263 hd,
12264 self.gemma4_rope_dims(il),
12265 nh * t,
12266 nkv * t,
12267 pos_d,
12268 nh,
12269 nkv,
12270 base,
12271 1.0,
12272 ff,
12273 eps,
12274 )?;
12275 let kvl = cache.kv[il].as_mut().unwrap();
12276 let base_len = kvl.len;
12277 e.append_kv_quantized_rows(
12278 &k,
12279 &v,
12280 &mut kvl.k,
12281 &mut kvl.v,
12282 base_len,
12283 t,
12284 kvl.kv_dim_k,
12285 kvl.kv_dim_v,
12286 kvl.k_tok_bytes,
12287 kvl.v_tok_bytes,
12288 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12289 )?;
12290 kvl.len += t;
12291 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12292 let mut attn = e.uninit(t * nh * hd)?;
12293 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
12296 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
12299 if rows_ok && (!swa || base_len + t <= win) {
12300 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12301 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12302 if hd == 512 {
12303 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12305 e.fa_decode_rows(
12306 &q,
12307 &k_view,
12308 &v_view,
12309 &mut attn,
12310 hd,
12311 nh,
12312 nkv,
12313 base_len,
12314 t,
12315 scale,
12316 kvl.k_tok_bytes,
12317 kvl.v_tok_bytes,
12318 Some((&kvl.len_d, 0)),
12319 false,
12320 swa && crate::Engine::wkv_on(),
12321 None,
12322 )?;
12323 } else {
12324 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12328 e.fa_decode_rows_dc(
12329 &q,
12330 &k_view,
12331 &v_view,
12332 &mut attn,
12333 hd,
12334 nh,
12335 nkv,
12336 &kvl.len_d,
12337 base_len + t,
12338 t,
12339 scale,
12340 kvl.k_tok_bytes,
12341 kvl.v_tok_bytes,
12342 0,
12343 swa && crate::Engine::wkv_on(),
12344 )?;
12345 }
12346 return Ok(e.matmul(&fa.wo, &attn, t)?);
12347 }
12348 if hd == 256
12356 && swa
12357 && base_len + 1 >= win
12358 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12359 {
12360 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12361 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12362 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12363 e.fa_decode_rows_w(
12364 &q,
12365 &k_view,
12366 &v_view,
12367 &mut attn,
12368 hd,
12369 nh,
12370 nkv,
12371 &kvl.len_d,
12372 0,
12373 t,
12374 scale,
12375 win,
12376 kvl.k_tok_bytes,
12377 kvl.v_tok_bytes,
12378 None,
12379 )?;
12380 return Ok(e.matmul(&fa.wo, &attn, t)?);
12381 }
12382 for i in 0..t {
12383 let avail = base_len + i + 1;
12384 let (off_tok, t_kv) = if swa && avail > win {
12385 (avail - win, win)
12386 } else {
12387 (0, avail)
12388 };
12389 let k_view = e.view_u8_range(
12390 &kvl.k,
12391 off_tok * kvl.k_tok_bytes,
12392 (off_tok + t_kv) * kvl.k_tok_bytes,
12393 );
12394 let v_view = e.view_u8_range(
12395 &kvl.v,
12396 off_tok * kvl.v_tok_bytes,
12397 (off_tok + t_kv) * kvl.v_tok_bytes,
12398 );
12399 let qi = e.view(&q, t * nh * hd);
12400 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12401 let mut q_one = e.uninit(nh * hd)?;
12402 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12403 let mut a_one = e.uninit(nh * hd)?;
12404 if swa
12408 && avail > win
12409 && hd == 256
12410 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12411 {
12412 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12413 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12414 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12415 e.fa_decode_rows_w(
12416 &q_one,
12417 &kp,
12418 &vp,
12419 &mut a_one,
12420 hd,
12421 nh,
12422 nkv,
12423 &kvl.len_d,
12424 0,
12425 1,
12426 scale,
12427 win,
12428 kvl.k_tok_bytes,
12429 kvl.v_tok_bytes,
12430 None,
12431 )?;
12432 } else if !swa
12433 && hd == 512
12434 && avail >= crate::fa512_min_tkv()
12435 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12436 {
12437 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12438 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12439 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12440 e.fa_decode_rows(
12441 &q_one,
12442 &kp,
12443 &vp,
12444 &mut a_one,
12445 hd,
12446 nh,
12447 nkv,
12448 avail - 1,
12449 1,
12450 scale,
12451 kvl.k_tok_bytes,
12452 kvl.v_tok_bytes,
12453 Some((&kvl.len_d, 0)),
12454 false,
12455 false,
12456 None,
12457 )?;
12458 } else {
12459 e.fa_decode_kvmod(
12460 &q_one,
12461 &k_view,
12462 &v_view,
12463 &mut a_one,
12464 hd,
12465 nh,
12466 nkv,
12467 t_kv,
12468 scale,
12469 kvl.k_tok_bytes,
12470 kvl.v_tok_bytes,
12471 swa && crate::Engine::wkv_on(),
12472 )?;
12473 }
12474 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12475 }
12476 Ok(e.matmul(&fa.wo, &attn, t)?)
12477 }
12478
12479 pub(crate) fn gemma4_decode_step_h(
12482 &self,
12483 e: &Engine,
12484 token: u32,
12485 cache: &mut Cache,
12486 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12487 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12492 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12493 }
12494 if crate::pp::pp_cuts(self.layers.len()).is_some() {
12495 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12496 }
12497 let n_embd = self.cfg.n_embd as usize;
12498 let eps = self.cfg.rms_eps;
12499 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12500 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12501 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12502 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12505 let n_layers = self.layers.len();
12506 for (il, layer) in self.layers.iter().enumerate() {
12507 let (hq, hdq) = match h_carry.take() {
12508 Some(p) => p,
12509 None => {
12510 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12511 }
12512 };
12513 let Mixer::Full(fa) = &layer.mixer else {
12514 panic!("gemma4 layer {il} not full-attn")
12515 };
12516 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12517 let next_norm = if il + 1 < n_layers {
12518 Some(self.layers[il + 1].attn_norm.float_data())
12519 } else {
12520 None
12521 };
12522 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12523 x = xn;
12524 h_carry = hn;
12525 }
12526 let mut hn = e.uninit(n_embd)?;
12527 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12528 let h_seed = e.clone_dtod(&x)?;
12529 let mut ld = e.matmul(&self.output, &hn, 1)?;
12530 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12531 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
12533 let logits = e.dtoh(&ld)?;
12534 cache.pos += 1;
12535 Ok((logits, h_seed))
12536 }
12537
12538 fn gemma4_decode_layers(
12546 &self,
12547 e: &Engine,
12548 mut x: CudaSlice<f32>,
12549 lo: usize,
12550 hi: usize,
12551 pos_d: &CudaSlice<i32>,
12552 cache: &mut Cache,
12553 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12554 let n_embd = self.cfg.n_embd as usize;
12555 let eps = self.cfg.rms_eps;
12556 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12557 for il in lo..hi {
12558 let layer = &self.layers[il];
12559 let (hq, hdq) = match h_carry.take() {
12560 Some(p) => p,
12561 None => {
12563 e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12564 }
12565 };
12566 let Mixer::Full(fa) = &layer.mixer else {
12567 panic!("gemma4 layer {il} not full-attn")
12568 };
12569 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12570 let next_norm = if il + 1 < hi {
12571 Some(self.layers[il + 1].attn_norm.float_data())
12572 } else {
12573 None
12574 };
12575 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12576 x = xn;
12577 h_carry = hn;
12578 }
12579 Ok(x)
12580 }
12581
12582 fn gemma4_decode_step_h_pp2(
12590 &self,
12591 e: &Engine,
12592 token: u32,
12593 cache: &mut Cache,
12594 split: usize,
12595 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12596 if crate::pp::pp2_streams_off() {
12597 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12598 }
12599 let rt = crate::pp::Pp2Rt::get(e)?;
12600 let e0 = rt.engine(0, e);
12601 let e1 = rt.engine(1, e);
12602 let n_embd = self.cfg.n_embd as usize;
12603 let eps = self.cfg.rms_eps;
12604 let pos = cache.pos as i32;
12605
12606 let slot = {
12608 let _st0 = rt.enter(0);
12609 let pos_d = e0.htod_i32(&[pos])?;
12610 #[cfg(debug_assertions)]
12611 crate::debug_assert_tensor_stream_device(
12612 &pos_d,
12613 &e0.stream(),
12614 "gemma4_decode_step_h_pp2.stage0.pos_d",
12615 );
12616 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12617 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12618 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12619 rt.tx(0, &x, n_embd)?
12620 };
12621
12622 let _st1 = rt.enter(1);
12624 let pos_d = e1.htod_i32(&[pos])?;
12625 #[cfg(debug_assertions)]
12626 crate::debug_assert_tensor_stream_device(
12627 &pos_d,
12628 &e1.stream(),
12629 "gemma4_decode_step_h_pp2.stage1.pos_d",
12630 );
12631 let x = rt.rx(0, slot, n_embd)?;
12632 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12633
12634 let mut hn = e1.uninit(n_embd)?;
12635 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12636 let h_seed = e1.clone_dtod(&x)?;
12637 let mut ld = e1.matmul(&self.output, &hn, 1)?;
12638 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12639 e1.softcap(&mut ld, cap, self.output.out_features())?;
12640 self.gemma4_suppress(e1, &mut ld, 1)?;
12641 let logits = e1.dtoh(&ld)?;
12642 cache.pos += 1;
12643 Ok((logits, h_seed))
12644 }
12645
12646 fn gemma4_decode_step_h_pp2_samestream(
12649 &self,
12650 e: &Engine,
12651 token: u32,
12652 cache: &mut Cache,
12653 split: usize,
12654 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12655 let n_embd = self.cfg.n_embd as usize;
12656 let eps = self.cfg.rms_eps;
12657 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12658
12659 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12661 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12662 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12663
12664 let boundary_tx = e.clone_dtod(&x)?;
12666 let boundary_rx = e.clone_dtod(&boundary_tx)?;
12667
12668 let x =
12670 self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12671
12672 let mut hn = e.uninit(n_embd)?;
12673 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12674 let h_seed = e.clone_dtod(&x)?;
12675 let mut ld = e.matmul(&self.output, &hn, 1)?;
12676 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12677 e.softcap(&mut ld, cap, self.output.out_features())?;
12678 self.gemma4_suppress(e, &mut ld, 1)?;
12679 let logits = e.dtoh(&ld)?;
12680 cache.pos += 1;
12681 Ok((logits, h_seed))
12682 }
12683}
12684
12685impl HybridModel {
12704 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12707 let geometry = self
12708 .cfg
12709 .layer_geometry(il as u32)
12710 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12711 debug_assert_eq!(
12712 geometry.attention_gate,
12713 memra_gguf::config::AttentionGateKind::SeparateHead
12714 );
12715 geometry
12716 }
12717
12718 #[allow(clippy::too_many_arguments)]
12778 fn step35_attn_pre_wo(
12779 &self,
12780 e: &Engine,
12781 fa: &FullAttnLayer,
12782 mut g3: Vec<CudaSlice<f32>>,
12783 hg: Option<&CudaSlice<f32>>,
12784 gt_pre: Option<&CudaSlice<f32>>,
12785 pos_d: &CudaSlice<i32>,
12786 t: usize,
12787 cache: Option<&mut Cache>,
12788 il: usize,
12789 seq_end: usize,
12790 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12791 let geometry = self.step35_geom(il);
12792 let hd = geometry.head_dim_k as usize;
12793 let nkv = geometry.n_head_kv as usize;
12794 let nh = geometry.n_head as usize;
12795 let rbase = geometry.rope_base;
12796 let scale = geometry.attention_scale();
12797 let swa = geometry.window.is_some();
12798 let eps = self.cfg.rms_eps;
12799 let win = geometry.window.unwrap_or(0) as usize;
12800 let n_rot = geometry.n_rot as usize;
12801
12802 let v = g3.pop().unwrap();
12803 let k0 = g3.pop().unwrap();
12804 let q0 = g3.pop().unwrap();
12805
12806 let mut q = e.uninit(t * nh * hd)?;
12810 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12811 let mut k = e.uninit(t * nkv * hd)?;
12812 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12813 let ff = if geometry.rope_factors {
12814 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12815 } else {
12816 None
12817 };
12818 #[cfg(debug_assertions)]
12819 if let Some(ff) = ff {
12820 crate::debug_assert_tensor_stream_device(
12821 ff,
12822 &e.stream(),
12823 "step35_attn_pre_wo.rope_freqs",
12824 );
12825 }
12826 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12827
12828 let mut attn = e.uninit(t * nh * hd)?;
12829 match cache {
12830 Some(cache) => {
12831 let base_len = cache.kv[il].as_ref().unwrap().len;
12832 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12834 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12835 let off = if swa {
12836 let raw = base_len.saturating_sub(win - 1);
12837 if legacy_tkv || legacy_calllocal {
12838 raw
12839 } else {
12840 raw & !31usize
12841 }
12842 } else {
12843 0
12844 };
12845 {
12846 let kvl = cache.kv[il].as_mut().unwrap();
12847 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12848 let write_row = e.prepare_kv_append(kvl, off, t)?;
12849 e.append_kv_quantized_rows(
12850 &k,
12851 &v,
12852 &mut kvl.k,
12853 &mut kvl.v,
12854 write_row,
12855 t,
12856 kvl.kv_dim_k,
12857 kvl.kv_dim_v,
12858 kvl.k_tok_bytes,
12859 kvl.v_tok_bytes,
12860 crate::Engine::kv_fp8_on(),
12861 )?;
12862 kvl.len += t;
12863 let new_len = kvl.len as i32;
12864 e.set_i32_one(&mut kvl.len_d, new_len)?;
12865 }
12866 let kvl = cache.kv[il].as_ref().unwrap();
12867 let t_kv = base_len + t - off;
12890 let physical = kvl.physical_rows(off, off + t_kv)?;
12891 let k_view = e.view_u8_range(
12892 &kvl.k,
12893 physical.start * kvl.k_tok_bytes,
12894 physical.end * kvl.k_tok_bytes,
12895 );
12896 let v_view = e.view_u8_range(
12897 &kvl.v,
12898 physical.start * kvl.v_tok_bytes,
12899 physical.end * kvl.v_tok_bytes,
12900 );
12901 let swa_naive = if legacy_tkv {
12913 t_kv > win
12914 } else {
12915 seq_end > win
12916 };
12917 if swa && swa_naive {
12918 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
12931 e.sdpa_naive_w_quantized_view(
12932 &q,
12933 &k_view,
12934 &v_view,
12935 &mut attn,
12936 hd,
12937 nh,
12938 nkv,
12939 t,
12940 t_kv,
12941 scale,
12942 true,
12943 win,
12944 kvl.k_tok_bytes,
12945 kvl.v_tok_bytes,
12946 )?;
12947 } else {
12948 e.fa_prefill_view_ws_w_hd128(
12949 &q,
12950 &k_view,
12951 &v_view,
12952 &mut attn,
12953 hd,
12954 nh,
12955 nkv,
12956 t,
12957 t_kv,
12958 scale,
12959 true,
12960 win,
12961 kvl.k_tok_bytes,
12962 kvl.v_tok_bytes,
12963 )?;
12964 }
12965 } else if std::env::var("MEMRA_NOFA").is_ok() {
12966 e.sdpa_naive_quantized_view(
12967 &q,
12968 &k_view,
12969 &v_view,
12970 &mut attn,
12971 hd,
12972 nh,
12973 nkv,
12974 t,
12975 t_kv,
12976 scale,
12977 true,
12978 kvl.k_tok_bytes,
12979 kvl.v_tok_bytes,
12980 )?;
12981 } else {
12982 e.fa_prefill_view_ws(
12987 &q,
12988 &k_view,
12989 &v_view,
12990 &mut attn,
12991 hd,
12992 nh,
12993 nkv,
12994 t,
12995 t_kv,
12996 scale,
12997 true,
12998 kvl.k_tok_bytes,
12999 kvl.v_tok_bytes,
13000 crate::Engine::kv_fp8_on(),
13001 )?;
13002 }
13003 }
13004 None => {
13005 debug_assert_eq!(
13010 seq_end, t,
13011 "step35 cacheless prefill is monolithic (seq_end == t)"
13012 );
13013 if swa && seq_end > win {
13014 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13015 } else if std::env::var("MEMRA_NOFA").is_ok() {
13016 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13017 } else {
13018 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13019 }
13020 }
13021 }
13022
13023 let gw = fa
13026 .attn_gate
13027 .as_ref()
13028 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13029 let gt_owned = if gt_pre.is_none() {
13030 Some(e.matmul(
13031 gw,
13032 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
13033 t,
13034 )?)
13035 } else {
13036 None
13037 };
13038 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
13039 let mut ag = e.uninit(t * nh * hd)?;
13040 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
13041 Ok(ag)
13042 }
13043
13044 pub(crate) fn step35_attn(
13047 &self,
13048 e: &Engine,
13049 fa: &FullAttnLayer,
13050 h: &CudaSlice<f32>,
13051 pos_d: &CudaSlice<i32>,
13052 t: usize,
13053 il: usize,
13054 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13055 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
13056 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
13058 Ok(e.matmul(&fa.wo, &ag, t)?)
13059 }
13060
13061 #[allow(clippy::too_many_arguments)]
13068 pub(crate) fn step35_attn_prime(
13069 &self,
13070 e: &Engine,
13071 fa: &FullAttnLayer,
13072 h: &CudaSlice<f32>,
13073 hx: Option<&CudaSlice<u8>>,
13074 pos_d: &CudaSlice<i32>,
13075 t: usize,
13076 cache: &mut Cache,
13077 il: usize,
13078 seq_end: usize,
13079 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13080 let g3 = match hx {
13081 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
13082 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
13083 };
13084 let ag =
13085 self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
13086 Ok(e.matmul(&fa.wo, &ag, t)?)
13087 }
13088
13089 #[allow(clippy::too_many_arguments)]
13099 pub(crate) fn step35_decode_attn(
13100 &self,
13101 e: &Engine,
13102 fa: &FullAttnLayer,
13103 il: usize,
13104 h: &CudaSlice<f32>,
13105 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
13106 pos_d: &CudaSlice<i32>,
13107 cache: &mut Cache,
13108 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13109 let geometry = self.step35_geom(il);
13110 let hd = geometry.head_dim_k as usize;
13111 let nkv = geometry.n_head_kv as usize;
13112 let nh = geometry.n_head as usize;
13113 let rbase = geometry.rope_base;
13114 let scale = geometry.attention_scale();
13115 let swa = geometry.window.is_some();
13116 let eps = self.cfg.rms_eps;
13117 let win = geometry.window.unwrap_or(0) as usize;
13118 let n_rot = geometry.n_rot as usize;
13119 let n_embd = self.cfg.n_embd as usize;
13120 let gw = fa
13121 .attn_gate
13122 .as_ref()
13123 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13124
13125 let (q0, k0, v0, gt) = match pre_q {
13126 Some((hq, hdq)) => {
13127 debug_assert!(
13128 e.uses_q8_1_fast(gw),
13129 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
13130 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
13131 );
13132 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
13133 Some(t3) => t3,
13134 None => (
13135 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
13136 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
13137 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
13138 ),
13139 };
13140 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
13141 (a, b, c, gt)
13142 }
13143 None => {
13144 if e.uses_q8_1_fast(&fa.wq)
13145 && e.uses_q8_1_fast(&fa.wk)
13146 && e.uses_q8_1_fast(&fa.wv)
13147 && e.uses_q8_1_fast(gw)
13148 {
13149 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
13150 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
13151 Some(t3) => t3,
13152 None => (
13153 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
13154 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
13155 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
13156 ),
13157 };
13158 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
13159 (a, b, c, gt)
13160 } else {
13161 (
13162 e.matmul(&fa.wq, h, 1)?,
13163 e.matmul(&fa.wk, h, 1)?,
13164 e.matmul(&fa.wv, h, 1)?,
13165 e.matmul(gw, h, 1)?,
13166 )
13167 }
13168 }
13169 };
13170
13171 let mut q = e.uninit(nh * hd)?;
13172 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
13173 let mut k = e.uninit(nkv * hd)?;
13174 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
13175 let ff = if swa {
13176 None
13177 } else {
13178 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
13179 };
13180 #[cfg(debug_assertions)]
13181 if let Some(ff) = ff {
13182 crate::debug_assert_tensor_stream_device(
13183 ff,
13184 &e.stream(),
13185 "step35_decode_attn.rope_freqs",
13186 );
13187 }
13188 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
13189
13190 if std::env::var("MEMRA_NOFA").is_ok() {
13191 return Err(
13192 "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
13193 cache; unset MEMRA_NOFA to use fa_decode"
13194 .into(),
13195 );
13196 }
13197 let kvl = cache.kv[il].as_mut().unwrap();
13198 let next_len = kvl.len + 1;
13199 let (off, t_kv) = if swa && next_len > win {
13200 (next_len - win, win)
13201 } else {
13202 (0, next_len)
13203 };
13204 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
13205 e.append_kv_quantized(
13206 &k,
13207 &v0,
13208 &mut kvl.k,
13209 &mut kvl.v,
13210 write_row,
13211 kvl.kv_dim_k,
13212 kvl.kv_dim_v,
13213 kvl.k_tok_bytes,
13214 kvl.v_tok_bytes,
13215 crate::Engine::kv_fp8_on(),
13216 )?;
13217 kvl.len = next_len;
13218 let physical = kvl.physical_rows(off, off + t_kv)?;
13219 let k_view = e.view_u8_range(
13220 &kvl.k,
13221 physical.start * kvl.k_tok_bytes,
13222 physical.end * kvl.k_tok_bytes,
13223 );
13224 let v_view = e.view_u8_range(
13225 &kvl.v,
13226 physical.start * kvl.v_tok_bytes,
13227 physical.end * kvl.v_tok_bytes,
13228 );
13229 let mut attn = e.uninit(nh * hd)?;
13230 e.fa_decode_kvmod(
13231 &q,
13232 &k_view,
13233 &v_view,
13234 &mut attn,
13235 hd,
13236 nh,
13237 nkv,
13238 t_kv,
13239 scale,
13240 kvl.k_tok_bytes,
13241 kvl.v_tok_bytes,
13242 crate::Engine::kv_fp8_on(),
13243 )?;
13244
13245 let mut ag = e.uninit(nh * hd)?;
13246 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
13247 Ok(e.matmul(&fa.wo, &ag, 1)?)
13248 }
13249}
13250
13251impl HybridModel {
13260 pub fn is_gemma4_e4b(&self) -> bool {
13261 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
13262 }
13263
13264 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
13268 let g = self.cfg.gemma4.as_ref().unwrap();
13269 let swa = g.swa_pattern[il];
13270 let hd = if swa {
13271 g.key_length_swa
13272 } else {
13273 g.key_length_global
13274 } as usize;
13275 let Mixer::Full(fa) = &self.layers[il].mixer else {
13276 panic!("e4b layer {il} not full-attn")
13277 };
13278 let nh = fa.wq.out_features() / hd;
13279 let nkv = fa.wk.out_features() / hd;
13280 (
13281 hd,
13282 nkv,
13283 nh,
13284 if swa {
13285 g.rope_base_swa
13286 } else {
13287 g.rope_base_global
13288 },
13289 1.0,
13290 swa,
13291 )
13292 }
13293
13294 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
13296 self.layers[il]
13297 .gemma4
13298 .as_ref()
13299 .and_then(|b| b.e4b.as_ref())
13300 .and_then(|e4| e4.kv_share.map(|t| t as usize))
13301 }
13302
13303 fn gemma4_e4b_inp_pl(
13308 &self,
13309 e: &Engine,
13310 tokens: &[u32],
13311 x_scaled: &CudaSlice<f32>,
13312 t: usize,
13313 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13314 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
13315 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
13316 }
13317
13318 fn gemma4_e4b_inp_pl_dev(
13320 &self,
13321 e: &Engine,
13322 tok_d: &CudaSlice<u32>,
13323 x_scaled: &CudaSlice<f32>,
13324 t: usize,
13325 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13326 let aux = self.gemma4_aux.as_ref().unwrap();
13327 let m = aux.e4b.as_ref().unwrap();
13328 let n_embd = self.cfg.n_embd as usize;
13329 let n_layer = self.layers.len();
13330 let width = m.n_epl * n_layer;
13331 let tbl = m.tok_tbl_gpu.get_or_init(|| {
13332 e.upload_u8(&m.tok_embd_bytes)
13333 .expect("e4b per-layer token table upload")
13334 });
13335 let mut a =
13336 e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
13337 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
13338 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
13339 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
13340 let mut pn = e.uninit(t * width)?;
13341 e.rms_norm(
13342 &p,
13343 m.proj_norm.float_data(),
13344 &mut pn,
13345 m.n_epl,
13346 t * n_layer,
13347 self.cfg.rms_eps,
13348 )?;
13349 let mut out = e.uninit(t * width)?;
13350 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
13351 Ok(out)
13352 }
13353
13354 #[allow(clippy::too_many_arguments)]
13359 fn gemma4_e4b_attn(
13360 &self,
13361 e: &Engine,
13362 il: usize,
13363 hq: &CudaSlice<i8>,
13364 hdq: &CudaSlice<f32>,
13365 pos_d: &CudaSlice<i32>,
13366 t: usize,
13367 cache: &mut Cache,
13368 dc_bucket: Option<usize>,
13369 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13370 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
13371 let eps = self.cfg.rms_eps;
13372 let aux = self.gemma4_aux.as_ref().unwrap();
13373 let ones = aux.ones(e);
13374 #[cfg(debug_assertions)]
13375 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
13376 let Mixer::Full(fa) = &self.layers[il].mixer else {
13377 unreachable!()
13378 };
13379 let h0 = e.zeros(0)?;
13383 let h = &h0;
13384
13385 let ff = if swa {
13386 None
13387 } else {
13388 Some(
13389 aux.rope_freqs(e)
13390 .expect("e4b global rope needs rope_freqs.weight"),
13391 )
13392 };
13393 #[cfg(debug_assertions)]
13394 if let Some(ff) = ff {
13395 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13396 }
13397 let share = self.gemma4_e4b_kv_target(il);
13398 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13400 let mut q;
13401 if let Some(_tgt) = share {
13402 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13403 q = e.uninit(t * nh * hd)?;
13404 let mut kdummy = e.uninit(1)?;
13407 let mut vdummy = e.uninit(1)?;
13408 e.rms_norm_qkv_rope(
13409 &q0,
13410 &q0,
13411 &q0,
13412 fa.q_norm.float_data(),
13413 fa.q_norm.float_data(),
13414 ones,
13415 &mut q,
13416 &mut kdummy,
13417 &mut vdummy,
13418 hd,
13419 self.gemma4_rope_dims(il),
13420 nh * t,
13421 0,
13422 pos_d,
13423 nh,
13424 1,
13425 base,
13426 1.0,
13427 ff,
13428 eps,
13429 )?;
13430 } else {
13431 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13435 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13436 q = e.uninit(t * nh * hd)?;
13437 let mut k = e.uninit(t * nkv * hd)?;
13438 let mut v = e.uninit(t * nkv * hd)?;
13439 if t == 1 && cat.is_some() {
13440 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13441 e.rms_norm_qkv_rope_cat(
13442 &qkv0,
13443 fa.q_norm.float_data(),
13444 fa.k_norm.float_data(),
13445 ones,
13446 &mut q,
13447 &mut k,
13448 &mut v,
13449 hd,
13450 self.gemma4_rope_dims(il),
13451 nh,
13452 nkv,
13453 pos_d,
13454 nh,
13455 nkv,
13456 base,
13457 1.0,
13458 ff,
13459 eps,
13460 )?;
13461 } else {
13462 let (q0, k0, v0) = match if t == 1 {
13463 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13464 } else {
13465 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13468 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13469 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13470 } else {
13471 None
13472 }
13473 } {
13474 Some(triple) => triple,
13475 None => (
13476 e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13477 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13478 e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13479 ), };
13481 e.rms_norm_qkv_rope(
13484 &q0,
13485 &k0,
13486 &v0,
13487 fa.q_norm.float_data(),
13488 fa.k_norm.float_data(),
13489 ones,
13490 &mut q,
13491 &mut k,
13492 &mut v,
13493 hd,
13494 self.gemma4_rope_dims(il),
13495 nh * t,
13496 nkv * t,
13497 pos_d,
13498 nh,
13499 nkv,
13500 base,
13501 1.0,
13502 ff,
13503 eps,
13504 )?;
13505 }
13506 let kvl = cache.kv[il].as_mut().unwrap();
13507 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13511 if dc_bucket.is_some() {
13512 debug_assert!(t == 1);
13517 e.append_kv_quantized_row_dc_inc(
13519 &k,
13520 &v,
13521 &mut kvl.k,
13522 &mut kvl.v,
13523 &mut kvl.len_d,
13524 kvl.kv_dim_k,
13525 kvl.kv_dim_v,
13526 kvl.k_tok_bytes,
13527 kvl.v_tok_bytes,
13528 cls,
13529 )?;
13530 } else {
13531 e.append_kv_quantized_rows(
13532 &k,
13533 &v,
13534 &mut kvl.k,
13535 &mut kvl.v,
13536 kvl.len,
13537 t,
13538 kvl.kv_dim_k,
13539 kvl.kv_dim_v,
13540 kvl.k_tok_bytes,
13541 kvl.v_tok_bytes,
13542 cls,
13543 )?;
13544 kvl.len += t;
13545 }
13546 kv_f32 = Some((k, v));
13547 }
13548 let kvl_idx = share.unwrap_or(il);
13551 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13552 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13554 let mut attn = e.uninit(t * nh * hd)?;
13555 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13567 if let Some((kf, vf)) = &kv_f32 {
13568 if hd == 256 && t <= win {
13569 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13570 return Ok(e.matmul(&fa.wo, &attn, t)?);
13571 }
13572 if hd == 256 && swa && t > win {
13573 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13574 return Ok(e.matmul(&fa.wo, &attn, t)?);
13575 }
13576 if hd == 512 && !swa {
13577 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13578 return Ok(e.matmul(&fa.wo, &attn, t)?);
13579 }
13580 } else if share.is_some() {
13581 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13582 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13583 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13584 if hd == 256 && (!swa || t <= win) {
13585 e.fa_prefill_view(
13587 &q,
13588 &k_view,
13589 &v_view,
13590 &mut attn,
13591 hd,
13592 nh,
13593 nkv,
13594 t,
13595 t,
13596 scale,
13597 true,
13598 kvl.k_tok_bytes,
13599 kvl.v_tok_bytes,
13600 g,
13601 )?;
13602 return Ok(e.matmul(&fa.wo, &attn, t)?);
13603 }
13604 let kv_dim = nkv * hd;
13607 let mut kf = e.uninit(t * kv_dim)?;
13608 let mut vf = e.uninit(t * kv_dim)?;
13609 e.fa_dequant_kv_view_f32(
13610 &k_view,
13611 &v_view,
13612 &mut kf,
13613 &mut vf,
13614 kv_dim,
13615 kv_dim,
13616 t,
13617 kvl.k_tok_bytes,
13618 kvl.v_tok_bytes,
13619 g,
13620 )?;
13621 if hd == 512 {
13622 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13623 } else {
13624 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13625 }
13626 return Ok(e.matmul(&fa.wo, &attn, t)?);
13627 }
13628 }
13629 if let Some(bucket) = dc_bucket {
13630 assert!(t == 1);
13635 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13641 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13642 } else {
13643 bucket
13644 };
13645 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13646 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13647 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13648 if crate::Engine::wpf_level() >= 1 {
13656 e.prefetch_weight_l2(&fa.wo)?;
13657 }
13658 if e.uses_q8_1_fast(&fa.wo) {
13661 let mut oq = e.alloc_i8_uninit(nh * hd)?;
13662 let mut od = e.zeros(nh * hd / 32)?;
13663 e.fa_decode_dc_q8(
13664 &q,
13665 &k_view,
13666 &v_view,
13667 &mut attn,
13668 hd,
13669 nh,
13670 nkv,
13671 &kvl.len_d,
13672 bucket,
13673 scale,
13674 kvl.k_tok_bytes,
13675 kvl.v_tok_bytes,
13676 g,
13677 Some((&mut oq, &mut od)),
13678 )?;
13679 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13680 }
13681 e.fa_decode_dc(
13682 &q,
13683 &k_view,
13684 &v_view,
13685 &mut attn,
13686 hd,
13687 nh,
13688 nkv,
13689 &kvl.len_d,
13690 bucket,
13691 scale,
13692 kvl.k_tok_bytes,
13693 kvl.v_tok_bytes,
13694 g,
13695 )?;
13696 return Ok(e.matmul(&fa.wo, &attn, t)?);
13697 }
13698 for i in 0..t {
13699 let avail = base_len + i + 1;
13700 let (off_tok, t_kv) = if swa && avail > win {
13701 (avail - win, win)
13702 } else {
13703 (0, avail)
13704 };
13705 let k_view = e.view_u8_range(
13706 &kvl.k,
13707 off_tok * kvl.k_tok_bytes,
13708 (off_tok + t_kv) * kvl.k_tok_bytes,
13709 );
13710 let v_view = e.view_u8_range(
13711 &kvl.v,
13712 off_tok * kvl.v_tok_bytes,
13713 (off_tok + t_kv) * kvl.v_tok_bytes,
13714 );
13715 let qv = e.view(&q, t * nh * hd);
13716 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13717 let mut q_one = e.uninit(nh * hd)?;
13718 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13719 let mut a_one = e.uninit(nh * hd)?;
13720 e.fa_decode_kvmod(
13724 &q_one,
13725 &k_view,
13726 &v_view,
13727 &mut a_one,
13728 hd,
13729 nh,
13730 nkv,
13731 t_kv,
13732 scale,
13733 kvl.k_tok_bytes,
13734 kvl.v_tok_bytes,
13735 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13736 )?;
13737 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13738 }
13739 Ok(e.matmul(&fa.wo, &attn, t)?)
13740 }
13741
13742 fn gemma4_e4b_trunk(
13747 &self,
13748 e: &Engine,
13749 tokens: &[u32],
13750 pos0: usize,
13751 cache: &mut Cache,
13752 head_last: bool,
13753 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13754 let n_embd = self.cfg.n_embd as usize;
13755 let t = tokens.len();
13756 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13757 let pos_d = e.htod_i32(&pos)?;
13758 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13759 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13760 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13761 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13762 }
13763
13764 fn gemma4_e4b_trunk_core(
13768 &self,
13769 e: &Engine,
13770 x_in: CudaSlice<f32>,
13771 inp_pl: CudaSlice<f32>,
13772 pos_d: &CudaSlice<i32>,
13773 t: usize,
13774 cache: &mut Cache,
13775 dc_bucket: Option<usize>,
13776 cap_logits: bool,
13777 head_last: bool,
13778 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13779 let n_embd = self.cfg.n_embd as usize;
13780 let eps = self.cfg.rms_eps;
13781 let n_layer = self.layers.len();
13782 let mut x = x_in;
13783 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13784 let n_epl = aux_e4b.n_epl;
13785
13786 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13792 for il in 0..n_layer {
13793 let layer = &self.layers[il];
13794 let (hq, hdq) = match h_carry.take() {
13795 Some(p) => p,
13796 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13797 };
13798 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13799 let bits = layer.gemma4.as_ref().unwrap();
13802 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13803 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13814 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13815 e,
13816 layer,
13817 &o,
13818 &x,
13819 t,
13820 Some(layer.post_attn_norm.float_data()),
13821 fuse_exit,
13822 )?;
13823 let mut resid = e.uninit(t * n_embd)?;
13824 let g = if fuse_exit {
13830 let (rq, rd) = e.rms_pre_add_q8_1(
13832 &sn,
13833 bits.post_ffw_norm.float_data(),
13834 &attn_out,
13835 &mut resid,
13836 n_embd,
13837 t,
13838 self.cfg.rms_eps,
13839 )?;
13840 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13841 } else {
13842 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13843 e.matmul(&e4b.inp_gate, &resid, t)?
13844 };
13845 let mut act = e.uninit(t * n_epl)?;
13846 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13847 let ipv = e.view(&inp_pl, n_epl * n_layer);
13848 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13849 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13850 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13851 } else {
13852 let mut inp_this = e.uninit(t * n_epl)?;
13853 e.copy_rows_strided(
13854 &inp_pl,
13855 &mut inp_this,
13856 n_epl,
13857 t,
13858 n_epl * n_layer,
13859 il * n_epl,
13860 )?;
13861 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13862 e.matmul(&e4b.proj, &act, t)?
13863 };
13864 let next_norm = if il + 1 < n_layer {
13867 self.layers[il + 1].attn_norm.float_data()
13868 } else {
13869 self.output_norm.float_data()
13870 };
13871 let mut xn = e.uninit(t * n_embd)?;
13872 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13873 &y,
13874 e4b.post_norm.float_data(),
13875 &resid,
13876 bits.layer_scale,
13877 next_norm,
13878 &mut xn,
13879 n_embd,
13880 t,
13881 eps,
13882 )?;
13883 h_carry = Some(pair);
13884 x = xn;
13885 }
13886 let (oq, odq) = h_carry.take().unwrap();
13890 let h0 = e.zeros(0)?;
13891 let hm = if head_last { 1 } else { t };
13892 let (hq, hd) = if head_last && t > 1 {
13893 let mut q1 = e.uninit_i8(n_embd)?;
13894 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
13895 let nb = n_embd / 32;
13896 let mut d1 = e.uninit(nb)?;
13897 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
13898 (q1, d1)
13899 } else {
13900 (oq, odq)
13901 };
13902 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
13903 if cap_logits {
13907 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
13908 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
13909 }
13910 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
13912 }
13913
13914 pub fn gemma4_e4b_decode_step_t_am_dev(
13921 &self,
13922 e: &Engine,
13923 tok_d: &CudaSlice<u32>,
13924 t: usize,
13925 pos0: usize,
13926 cache: &mut Cache,
13927 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13928 let n_embd = self.cfg.n_embd as usize;
13929 let eps = self.cfg.rms_eps;
13930 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13931 let pos_d = e.htod_i32(&pos)?;
13932 let embd_gpu = self
13933 .embd_gpu
13934 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
13935 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
13936 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
13937 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13938 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
13939 let (ld, xp) =
13940 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
13941 let n_vocab = self.output.out_features();
13944 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
13945 for i in 0..t {
13946 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
13947 }
13948 let mut hn = e.uninit(t * n_embd)?;
13949 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13950 cache.pos += t;
13951 Ok((vam, hn))
13952 }
13953
13954 pub(crate) fn gemma4_e4b_decode_step_t_h(
13957 &self,
13958 e: &Engine,
13959 tokens: &[u32],
13960 pos0: usize,
13961 cache: &mut Cache,
13962 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13963 let n_embd = self.cfg.n_embd as usize;
13964 let eps = self.cfg.rms_eps;
13965 let t = tokens.len();
13966 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
13967 let mut hn = e.uninit(t * n_embd)?;
13968 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13969 cache.pos += t;
13970 Ok((e.dtoh(&ld)?, hn))
13971 }
13972
13973 pub fn gemma4_e4b_decode_step_dcg(
13979 &self,
13980 e: &Engine,
13981 token_d: &mut CudaSlice<u32>,
13982 pos_d: &mut CudaSlice<i32>,
13983 embd_gpu: &CudaSlice<u8>,
13984 embd_qt: i32,
13985 embd_rb: usize,
13986 cache: &mut Cache,
13987 n_vocab: usize,
13988 bucket: usize,
13989 ) -> Result<(), Box<dyn std::error::Error>> {
13990 let n_embd = self.cfg.n_embd as usize;
13991 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13992 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13993 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13994 let (ld, _x) =
13995 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
13996 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
13997 e.inc_seqlen(pos_d)?;
13998 Ok(())
13999 }
14000
14001 #[allow(clippy::too_many_arguments)]
14009 pub fn gemma4_e4b_decode_step_dc(
14010 &self,
14011 e: &Engine,
14012 token_d: &CudaSlice<u32>,
14013 pos_d: &mut CudaSlice<i32>,
14014 embd_gpu: &CudaSlice<u8>,
14015 embd_qt: i32,
14016 embd_rb: usize,
14017 cache: &mut Cache,
14018 n_vocab: usize,
14019 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
14020 let n_embd = self.cfg.n_embd as usize;
14021 let eps = self.cfg.rms_eps;
14022 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
14023 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
14024 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
14025 let (ld, _x) =
14026 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
14027 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
14028 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
14029 e.inc_seqlen(pos_d)?;
14030 cache.pos += 1;
14031 let _ = eps;
14032 Ok(tok_out)
14033 }
14034
14035 pub(crate) fn gemma4_e4b_decode_step_h(
14038 &self,
14039 e: &Engine,
14040 token: u32,
14041 cache: &mut Cache,
14042 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14043 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
14044 let logits = e.dtoh(&ld)?;
14045 cache.pos += 1;
14046 Ok((logits, x))
14047 }
14048
14049 pub(crate) fn gemma4_e4b_prime(
14053 &self,
14054 e: &Engine,
14055 tokens: &[u32],
14056 cache: &mut Cache,
14057 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14058 if cache.pos != 0 {
14061 return Err(
14062 "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
14063 call or decode tokenwise"
14064 .into(),
14065 );
14066 }
14067 let n_embd = self.cfg.n_embd as usize;
14068 let t = tokens.len();
14069 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
14070 cache.pos += t;
14071 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
14073 let row = xv.slice((t - 1) * n_embd..t * n_embd);
14074 let mut h_seed = e.uninit(n_embd)?;
14075 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
14076 Ok((last, h_seed, x))
14077 }
14078
14079 pub(crate) fn gemma4_e4b_forward(
14081 &self,
14082 e: &Engine,
14083 tokens: &[u32],
14084 last_only: bool,
14085 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
14086 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
14087 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
14088 Ok(e.dtoh(&ld)?) }
14090}
14091
14092#[cfg(test)]
14093mod prime_chunk_schedule_tests {
14094 use super::{
14095 PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
14096 fixed_prime_chunk_ranges_for_ring,
14097 };
14098
14099 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
14100 ranges.iter().map(|(start, end)| end - start).collect()
14101 }
14102
14103 fn auto_chunk(t: usize) -> usize {
14104 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
14105 }
14106
14107 #[test]
14108 fn fixed_schedule_retains_measured_geometry() {
14109 assert_eq!(
14110 sizes(&fixed_prime_chunk_ranges(461, 128)),
14111 vec![128, 128, 128, 77]
14112 );
14113 assert_eq!(
14114 sizes(&fixed_prime_chunk_ranges(1833, 230)),
14115 vec![230, 230, 230, 230, 230, 230, 230, 223]
14116 );
14117 assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
14118 let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
14119 assert_eq!(capped, vec![4096, 4088, 16]);
14120 assert!(capped.iter().all(|&rows| rows <= 4096));
14121 assert_eq!(
14122 sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
14123 vec![4100],
14124 "flag-off schedule remains byte-for-byte the legacy monolithic tail",
14125 );
14126 }
14127
14128 #[test]
14129 fn dynamic_schedule_matches_registered_shapes() {
14130 let cases = [
14131 (461, vec![64, 141, 132, 124]),
14132 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
14133 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
14134 ];
14135 for (t, expected) in cases {
14136 let chunk = auto_chunk(t);
14137 let fixed = fixed_prime_chunk_ranges(t, chunk);
14138 assert_eq!(
14139 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
14140 expected
14141 );
14142 }
14143 }
14144
14145 #[test]
14146 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
14147 for t in 256..=8192 {
14148 let chunk = auto_chunk(t);
14149 let fixed = fixed_prime_chunk_ranges(t, chunk);
14150 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
14151 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
14152 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
14153 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
14154 for pair in dynamic.windows(2) {
14155 assert_eq!(pair[0].1, pair[1].0, "T={t}");
14156 }
14157 assert!(
14158 dynamic
14159 .iter()
14160 .all(|(start, end)| end - start >= PRIME_MIN_T),
14161 "T={t} sizes={:?}",
14162 sizes(&dynamic)
14163 );
14164 if dynamic.len() >= 3 {
14165 let chunk_sizes = sizes(&dynamic);
14166 assert!(
14167 chunk_sizes[0] < chunk_sizes[1],
14168 "T={t} sizes={chunk_sizes:?}"
14169 );
14170 assert!(
14171 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
14172 "T={t} sizes={chunk_sizes:?}"
14173 );
14174 }
14175 }
14176 }
14177}
14178
14179#[cfg(test)]
14180mod page_prefetch_tests {
14181 use super::{
14182 grouped_worker_prefetch_position, page_prefetch_positions,
14183 page_prefetch_window_from_values, worker_prefetch_positions,
14184 };
14185
14186 #[test]
14187 fn page_prefetch_window_keeps_existing_opt_in_default() {
14188 assert_eq!(page_prefetch_window_from_values(false, None), 0);
14189 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
14190 assert_eq!(page_prefetch_window_from_values(true, None), 1);
14191 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
14192 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
14193 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
14194 }
14195
14196 #[test]
14197 fn rolling_page_prefetch_advises_each_future_expert_once() {
14198 let advised: Vec<_> = (0..7)
14199 .flat_map(|position| page_prefetch_positions(position, 7, 3))
14200 .collect();
14201 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
14202
14203 let one_ahead: Vec<_> = (0..4)
14204 .flat_map(|position| page_prefetch_positions(position, 4, 1))
14205 .collect();
14206 assert_eq!(one_ahead, vec![1, 2, 3]);
14207 assert!(page_prefetch_positions(0, 4, 0).is_empty());
14208 }
14209
14210 #[test]
14211 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
14212 assert_eq!(grouped_worker_prefetch_position(0, None), None);
14213 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
14214 .chain(
14215 (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
14216 )
14217 .collect();
14218 assert_eq!(positions, vec![0, 1, 2, 3]);
14219 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
14220 }
14221
14222 #[test]
14223 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
14224 let queued: Vec<_> = (0..8)
14225 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
14226 .collect();
14227 assert_eq!(queued, (0..8).collect::<Vec<_>>());
14228
14229 let one_at_a_time: Vec<_> = (0..4)
14230 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
14231 .collect();
14232 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
14233 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
14234 }
14235}
14236
14237pub struct G4DcSlots {
14238 x: CudaSlice<f32>,
14239 xn: CudaSlice<f32>,
14240 cur: CudaSlice<f32>,
14241 hq: CudaSlice<i8>,
14242 hd_: CudaSlice<f32>,
14243 q0: CudaSlice<f32>,
14244 k0: CudaSlice<f32>,
14245 v0: CudaSlice<f32>,
14246 q: CudaSlice<f32>,
14247 k: CudaSlice<f32>,
14248 v: CudaSlice<f32>,
14249 attn: CudaSlice<f32>,
14250 o: CudaSlice<f32>,
14251 attn_out: CudaSlice<f32>,
14252 zsh: CudaSlice<f32>,
14253 zq: CudaSlice<i8>,
14254 zd: CudaSlice<f32>,
14255 gate: CudaSlice<f32>,
14256 up: CudaSlice<f32>,
14257 act: CudaSlice<f32>,
14258 actq: CudaSlice<i8>,
14259 actd: CudaSlice<f32>,
14260 f0: CudaSlice<f32>,
14261 sn: CudaSlice<f32>,
14262 hn: CudaSlice<f32>,
14263 logits: CudaSlice<f32>,
14264}