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 let n_embd = self.cfg.n_embd as usize;
765 let t = tokens.len();
766 assert!(
770 t >= PRIME_MIN_T,
771 "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
772 );
773 assert!(
774 cache.pos + t <= cache.max_ctx,
775 "prime_cache: prompt exceeds cache max_ctx"
776 );
777
778 if self.is_gemma4_e4b() {
790 return self.gemma4_e4b_prime(e, tokens, cache);
791 }
792 if self.cfg.gemma4.is_some() {
793 return self.gemma4_prime(e, tokens, cache);
795 }
796 let ranges = prime_chunk_ranges(t, self.layers.len());
797 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
833 let seq_end = if legacy_calllocal {
834 cache.pos + t
835 } else {
836 cache.pos + t + queued_after
837 };
838 if ranges.len() == 1 {
839 return self.prime_chunk(e, tokens, cache, seq_end);
840 }
841 if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
846 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
847 if crate::pp::pp_multi_stream_same_device() {
848 return Err(
849 "prime chunk pipeline refused with 2 stage streams on one device — \
850 that concurrent-stream placement remains quarantined by the deferred \
851 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
852 the serial split."
853 .into(),
854 );
855 }
856 return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
857 }
858 }
859 let mut hiddens = e.uninit(t * n_embd)?;
860 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
861 for &(start, end) in &ranges {
862 let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
863 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
864 last = Some((l, hs));
865 }
866 let (logits, h_seed) = last.unwrap();
867 Ok((logits, h_seed, hiddens))
868 }
869
870 fn prime_cache_pp2_pipelined(
875 &self,
876 e: &Engine,
877 tokens: &[u32],
878 cache: &mut Cache,
879 seq_end: usize,
880 ranges: &[(usize, usize)],
881 fence: &[usize],
882 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
883 debug_assert_eq!(fence.len(), 3);
884 debug_assert!(ranges.len() >= 2);
885 let rt = crate::pp::PpNRt::get(e)?;
886 assert_eq!(
887 rt.n_stages(),
888 2,
889 "prime pipeline requires exactly two PP stages"
890 );
891 let n_embd = self.cfg.n_embd as usize;
892 let t = tokens.len();
893 let initial_base = cache.pos;
894 let caller_stream = e.stream();
895
896 rt.fence_stages_behind(&caller_stream)?;
901 let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
902 rt.prepare_overlap_slots(0, max_payload)?;
903
904 let mut hiddens = e.uninit(t * n_embd)?;
905 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
906 let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
907 let (cache0, cache1) = stage_caches.parts();
908 let (first_start, first_end) = ranges[0];
909 let mut slot = self.prime_pp2_stage0_enqueue(
910 e,
911 rt,
912 &tokens[first_start..first_end],
913 cache0,
914 seq_end,
915 fence,
916 initial_base + first_start,
917 true,
918 )?;
919 cache0.pos = initial_base + first_end;
920
921 for (i, &(start, end)) in ranges.iter().enumerate() {
922 let base = initial_base + start;
923 debug_assert_eq!(
924 cache1.pos, base,
925 "stage 1 must drain chunks in original position order"
926 );
927 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
928 let next_base = initial_base + next_start;
929 debug_assert_eq!(
930 cache0.pos, next_base,
931 "stage 0 must issue chunks in original position order"
932 );
933 let cache0_stage = &mut *cache0;
934 std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
939 let stage0 = scope.spawn(move || -> Result<usize, String> {
940 let next = self
941 .prime_pp2_stage0_enqueue(
942 e,
943 rt,
944 &tokens[next_start..next_end],
945 cache0_stage,
946 seq_end,
947 fence,
948 next_base,
949 true,
950 )
951 .map_err(|err| err.to_string())?;
952 cache0_stage.pos = initial_base + next_end;
953 Ok(next)
954 });
955 let x = self.prime_pp2_stage1_enqueue(
956 e,
957 rt,
958 slot,
959 end - start,
960 cache1,
961 seq_end,
962 fence,
963 base,
964 true,
965 )?;
966 let out = {
967 rt.bind_stage(1)?;
968 let _st1 = rt.enter(1);
969 let e1 = rt.engine(1, e);
970 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
971 };
972 let next = stage0
973 .join()
974 .map_err(|_| "pipeprime stage-0 host walker panicked")?
975 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
976 Ok((out, Some(next)))
977 })?
978 } else {
979 let x = self.prime_pp2_stage1_enqueue(
980 e,
981 rt,
982 slot,
983 end - start,
984 cache1,
985 seq_end,
986 fence,
987 base,
988 true,
989 )?;
990 let out = {
991 rt.bind_stage(1)?;
992 let _st1 = rt.enter(1);
993 let e1 = rt.engine(1, e);
994 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
995 };
996 (out, None)
997 };
998
999 rt.publish_to(1, &caller_stream)?;
1000 e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1001 last = Some((out.0, out.1));
1002 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1003
1004 if let Some(next) = next_slot {
1005 rt.fence_stages_behind(&caller_stream)?;
1010 slot = next;
1011 }
1012 }
1013
1014 debug_assert_eq!(cache0.pos, initial_base + t);
1015 debug_assert_eq!(cache1.pos, initial_base + t);
1016 let (logits, h_seed) = last.unwrap();
1017 Ok((logits, h_seed, hiddens))
1018 }
1019
1020 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1027 if Engine::gdn_db_on()
1028 && Engine::gdn_chunked_enabled()
1029 && t >= 16
1030 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1031 && num_k * 2 == num_v
1032 {
1033 num_k
1034 } else {
1035 num_v
1036 }
1037 }
1038
1039 fn f16out_on(e: &Engine, t: usize) -> bool {
1044 crate::f16_ffi::pp_f16_enabled()
1045 && t >= 16
1046 && !e.verify_exact_on()
1047 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1048 }
1049
1050 pub fn prime_slabs_get(
1058 &self,
1059 e: &Engine,
1060 t: usize,
1061 n_embd: usize,
1062 n_ff_max: usize,
1063 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1064 let mut slabs = self.prime_slabs.lock().unwrap();
1065 let dev = e.ctx().ordinal();
1066 let need_new = match slabs.get(&dev) {
1067 None => true,
1068 Some(sl) => sl.lock().unwrap().t_cap < t,
1069 };
1070 if need_new {
1071 slabs.insert(
1072 dev,
1073 std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1074 t_cap: t,
1075 h: e.uninit(t * n_embd)?,
1076 x1: e.uninit(t * n_embd)?,
1077 z: e.uninit(t * n_embd)?,
1078 act: e.uninit(t * n_ff_max)?,
1079 xa: e.uninit(t * n_embd)?,
1080 xb: e.uninit(t * n_embd)?,
1081 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1082 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1083 gate: e.uninit(t * n_ff_max)?,
1084 up: e.uninit(t * n_ff_max)?,
1085 ffn_out: e.uninit(t * n_embd)?,
1086 seg_glue: Vec::new(),
1087 mixed: e.uninit(t * n_embd)?,
1088 seg_mid: Vec::new(),
1089 seg_t: 0,
1090 })),
1091 );
1092 }
1093 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1094 }
1095
1096 fn prime_chunk(
1100 &self,
1101 e: &Engine,
1102 tokens: &[u32],
1103 cache: &mut Cache,
1104 seq_end: usize,
1105 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1106 if crate::pp::pp_host_bounce_active()
1107 && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1108 {
1109 return Err(
1110 "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1111 has no active prime stage split and would peer-read remote weights; keep \
1112 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1113 .into(),
1114 );
1115 }
1116 if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1125 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1126 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1127 }
1128 }
1129 if crate::pp::pp_host_bounce_active() {
1130 return Err(
1131 "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1132 refusing an unsplit remote-weight walk"
1133 .into(),
1134 );
1135 }
1136 let t = tokens.len();
1137 let base = cache.pos;
1138 debug_assert!(
1139 seq_end >= base + t,
1140 "prime_chunk: seq_end must cover this chunk"
1141 );
1142 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1143 let pos_d = e.htod_i32(&pos)?;
1144
1145 let x_embed = self.embed(e, tokens)?; let x = self.prime_layers(
1147 e,
1148 x_embed,
1149 0,
1150 self.layers.len(),
1151 &pos_d,
1152 t,
1153 base,
1154 cache,
1155 seq_end,
1156 )?;
1157 self.prime_chunk_epilogue(e, x, t, cache)
1158 }
1159
1160 #[allow(clippy::too_many_arguments)]
1176 fn prime_layers(
1177 &self,
1178 e: &Engine,
1179 x_in: CudaSlice<f32>,
1180 lo: usize,
1181 hi: usize,
1182 pos_d: &CudaSlice<i32>,
1183 t: usize,
1184 base: usize,
1185 cache: &mut Cache,
1186 seq_end: usize,
1187 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1188 let cfg = &self.cfg;
1189 let n_embd = cfg.n_embd as usize;
1190 let eps = cfg.rms_eps;
1191 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1195 let n_ff_max = self
1201 .layers
1202 .iter()
1203 .map(|l| match &l.ffn {
1204 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1205 _ => n_embd,
1206 })
1207 .max()
1208 .unwrap_or(n_embd)
1209 .max(n_embd);
1210 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1211 let slab = if use_slabs {
1212 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1213 } else {
1214 None
1215 };
1216 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1217 let mut x_own; type SlabRefs<'a> = (
1219 &'a mut CudaSlice<f32>,
1220 &'a mut CudaSlice<f32>,
1221 &'a mut CudaSlice<f32>,
1222 &'a mut CudaSlice<f32>,
1223 &'a mut CudaSlice<u8>,
1224 &'a mut CudaSlice<u8>,
1225 &'a mut CudaSlice<f32>,
1226 &'a mut CudaSlice<f32>,
1227 &'a mut CudaSlice<f32>,
1228 );
1229 let (mut x_cur, mut x_nxt, sl): (
1230 &mut CudaSlice<f32>,
1231 &mut CudaSlice<f32>,
1232 Option<SlabRefs>,
1233 );
1234 let mut seg: Option<(
1235 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1236 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1237 &mut CudaSlice<f32>,
1238 &mut usize,
1239 )> = None;
1240 let mut x_own2;
1241 match slab_guard.as_mut() {
1242 Some(g) => {
1243 let slabs = &mut **g;
1244 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1245 let PrimeSlabs {
1246 xa,
1247 xb,
1248 h,
1249 x1,
1250 z,
1251 act,
1252 h16,
1253 z16,
1254 gate,
1255 up,
1256 ffn_out,
1257 seg_glue,
1258 mixed,
1259 seg_mid,
1260 seg_t,
1261 ..
1262 } = slabs;
1263 x_cur = xa;
1264 x_nxt = xb;
1265 seg = Some((seg_glue, seg_mid, mixed, seg_t));
1266 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1267 }
1268 None => {
1269 x_own = x_in;
1270 x_own2 = e.uninit(t * n_embd)?;
1271 x_cur = &mut x_own;
1272 x_nxt = &mut x_own2;
1273 sl = None;
1274 }
1275 }
1276 let mut alloc_h;
1277 let mut alloc_x1;
1278 let mut alloc_z;
1279 let mut alloc_act;
1280 let mut alloc_h16;
1281 let mut alloc_z16;
1282 let mut alloc_gate;
1283 let mut alloc_up;
1284 let mut alloc_fo;
1285 let (h, x1, z, act): (
1286 &mut CudaSlice<f32>,
1287 &mut CudaSlice<f32>,
1288 &mut CudaSlice<f32>,
1289 &mut CudaSlice<f32>,
1290 );
1291 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1292 let (sl_gate, sl_up, sl_fo): (
1293 &mut CudaSlice<f32>,
1294 &mut CudaSlice<f32>,
1295 &mut CudaSlice<f32>,
1296 );
1297 match sl {
1298 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1299 h = a;
1300 x1 = b;
1301 z = c;
1302 act = d;
1303 h16 = e16;
1304 z16 = f16b;
1305 sl_gate = g;
1306 sl_up = u;
1307 sl_fo = fo;
1308 }
1309 None => {
1310 alloc_h = e.uninit(t * n_embd)?;
1311 alloc_x1 = e.uninit(t * n_embd)?;
1312 alloc_z = e.uninit(t * n_embd)?;
1313 alloc_act = e.uninit(t * n_ff_max)?;
1314 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1315 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1316 alloc_gate = e.uninit(t * n_ff_max)?;
1317 alloc_up = e.uninit(t * n_ff_max)?;
1318 alloc_fo = e.uninit(t * n_embd)?;
1319 h = &mut alloc_h;
1320 x1 = &mut alloc_x1;
1321 z = &mut alloc_z;
1322 act = &mut alloc_act;
1323 h16 = &mut alloc_h16;
1324 z16 = &mut alloc_z16;
1325 sl_gate = &mut alloc_gate;
1326 sl_up = &mut alloc_up;
1327 sl_fo = &mut alloc_fo;
1328 }
1329 }
1330 let n_layers = self.layers.len();
1335 let use_seg = f16fuse
1345 && seg.is_some()
1346 && self.cfg.step35.is_none()
1347 && lo == 0
1348 && hi == n_layers
1349 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1350 if let Some((sg, sm, _, st)) = seg.as_mut() {
1351 if **st != t {
1352 sg.clear();
1353 sg.extend((0..n_layers).map(|_| None));
1354 sm.clear();
1355 sm.extend((0..n_layers).map(|_| None));
1356 **st = t;
1357 }
1358 }
1359 {
1360 let layer_lo = &self.layers[lo];
1361 if f16fuse {
1362 e.rms_norm_f16out(
1363 x_cur,
1364 layer_lo.attn_norm.float_data(),
1365 h,
1366 h16,
1367 n_embd,
1368 t,
1369 eps,
1370 )?;
1371 } else {
1372 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1373 }
1374 }
1375 for il in lo..hi {
1376 let layer = &self.layers[il];
1377 let hx16 = if f16fuse { Some(&*h16) } else { None };
1378 if use_seg {
1379 let (pre, pre16, w_out) = match &layer.mixer {
1382 Mixer::Full(fa) => {
1383 let g3 = match hx16 {
1384 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1385 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1386 };
1387 let (pre, pre16) =
1388 self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1389 (pre, pre16, &fa.wo)
1390 }
1391 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1392 Mixer::Linear(la) => {
1393 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1394 let g4 = match hx16 {
1395 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1396 None => e.matmul_group(&ws, h, t)?,
1397 };
1398 let (pre, pre16) =
1399 self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1400 (pre, pre16, &la.ssm_out)
1401 }
1402 };
1403 {
1404 let (_, sm, mslab, _) = seg.as_mut().unwrap();
1405 let pre_n = pre.len() / t;
1406 let xh_pre = match pre16 {
1407 Some(x) => x,
1408 None => e.f16_act(&pre, t * pre_n, pre_n)?,
1409 };
1410 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1411 let y = e.matmul(w_out, &pre, t)?;
1412 e.copy_into(mslab, 0, &y, t * n_embd)?;
1413 }
1414 if sm[il].is_none() {
1415 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1416 let w_post = layer.post_attn_norm.float_data();
1417 e.stream().synchronize()?;
1418 e.stream()
1419 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1420 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1421 e.add(x_cur, mslab, x1, t * n_embd)?;
1422 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1423 Ok(())
1424 })();
1425 let g = e.stream().end_capture(
1426 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1427 r?;
1428 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1429 }
1430 sm[il].as_ref().unwrap().launch()?;
1431 }
1432 } else {
1433 let mixed = match &layer.mixer {
1434 Mixer::Full(fa) => {
1435 self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?
1436 }
1437 Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1438 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1439 };
1440 if f16fuse {
1441 e.add_rms_norm_f16out(
1444 x_cur,
1445 &mixed,
1446 layer.post_attn_norm.float_data(),
1447 x1,
1448 z,
1449 z16,
1450 n_embd,
1451 t,
1452 eps,
1453 )?;
1454 } else {
1455 e.add(x_cur, &mixed, x1, t * n_embd)?;
1456 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1457 }
1458 }
1459 let zx16 = if f16fuse { Some(&*z16) } else { None };
1460 match &layer.ffn {
1461 crate::hybrid::Ffn::Dense {
1462 ffn_gate,
1463 ffn_up,
1464 ffn_down,
1465 } => {
1466 let n_ff = ffn_gate.out_features();
1467 let mut into_ok = false;
1470 if let Some(xh) = zx16 {
1471 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1472 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1473 }
1474 if !into_ok {
1475 let mut g2 = match zx16 {
1476 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1477 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1478 };
1479 let up_y = g2.pop().unwrap();
1480 let gate_y = g2.pop().unwrap();
1481 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1482 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1483 }
1484 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1489 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1490 {
1491 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1492 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1493 Some(a16)
1494 } else {
1495 Self::ffn_act_lim(
1496 e,
1497 &self.cfg,
1498 sl_gate,
1499 sl_up,
1500 1.0,
1501 1.0,
1502 d_lim,
1503 act,
1504 t * n_ff,
1505 )?;
1506 None
1507 };
1508 let xh_act = match act16 {
1510 Some(x) => x,
1511 None => e.f16_act(act, t * n_ff, n_ff)?,
1512 };
1513 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1514 let y = e.matmul(ffn_down, &*act, t)?;
1515 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1516 }
1517 }
1518 crate::hybrid::Ffn::Moe(m) => {
1519 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1520 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1521 }
1522 }
1523 if use_seg && il + 1 < hi {
1524 let w_next = self.layers[il + 1].attn_norm.float_data();
1526 let (sg, _, _, _) = seg.as_mut().unwrap();
1527 if sg[il].is_none() {
1528 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1529 e.stream().synchronize()?;
1530 e.stream()
1531 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1532 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1533 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1534 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1535 Ok(())
1536 })();
1537 let g = e.stream().end_capture(
1538 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1539 );
1540 r?;
1541 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1542 }
1543 sg[il].as_ref().unwrap().launch()?;
1544 } else {
1545 if il + 1 < hi {
1546 let w_next = self.layers[il + 1].attn_norm.float_data();
1547 if f16fuse {
1548 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1549 } else {
1550 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1551 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1552 }
1553 } else {
1554 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1555 }
1556 }
1557 if let Some(path) = Self::prime_trace_path() {
1563 let row = (base + t - 1) as usize;
1564 let host = e.dtoh(x_nxt)?;
1565 let last = &host[(t - 1) * n_embd..t * n_embd];
1566 use std::io::Write as _;
1567 let mut f = std::fs::OpenOptions::new()
1568 .create(true)
1569 .append(true)
1570 .open(path)?;
1571 let mut h64: u64 = 0xcbf29ce484222325;
1572 for v in last {
1573 h64 ^= v.to_bits() as u64;
1574 h64 = h64.wrapping_mul(0x100000001b3);
1575 }
1576 writeln!(
1577 f,
1578 "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1579 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1580 last[0], last[1], last[2]
1581 )?;
1582 }
1583 std::mem::swap(&mut x_cur, &mut x_nxt);
1584 }
1585 let mut x = e.uninit(t * n_embd)?;
1587 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1588 drop(slab_guard);
1589 Ok(x)
1590 }
1591
1592 fn prime_chunk_epilogue(
1597 &self,
1598 e: &Engine,
1599 x: CudaSlice<f32>,
1600 t: usize,
1601 cache: &mut Cache,
1602 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1603 let n_embd = self.cfg.n_embd as usize;
1604 let eps = self.cfg.rms_eps;
1605 let mut h_seed = e.uninit(n_embd)?;
1609 if !crate::spec::spec_hpost() {
1610 e.copy_view_into(
1611 &mut h_seed,
1612 0,
1613 &x.slice((t - 1) * n_embd..t * n_embd),
1614 n_embd,
1615 )?;
1616 }
1617 let mut hn = e.uninit(t * n_embd)?;
1619 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1620 if crate::spec::spec_hpost() {
1621 e.copy_view_into(
1622 &mut h_seed,
1623 0,
1624 &hn.slice((t - 1) * n_embd..t * n_embd),
1625 n_embd,
1626 )?;
1627 }
1628 let last = e.view(&hn, t * n_embd);
1629 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1630 let mut hlast = e.uninit(n_embd)?;
1631 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1632 let logits = e.matmul(&self.output, &hlast, 1)?;
1633 cache.pos += t;
1634 Ok((
1637 e.dtoh(&logits)?,
1638 h_seed,
1639 if crate::spec::spec_hpost() { hn } else { x },
1640 ))
1641 }
1642
1643 fn prime_chunk_ppn(
1667 &self,
1668 e: &Engine,
1669 tokens: &[u32],
1670 cache: &mut Cache,
1671 seq_end: usize,
1672 fence: &[usize],
1673 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1674 let rt = crate::pp::PpNRt::get(e)?;
1675 let n_st = fence.len() - 1;
1676 assert_eq!(
1677 rt.n_stages(),
1678 n_st,
1679 "PpNRt stage count {} != fence stages {n_st}",
1680 rt.n_stages()
1681 );
1682 let n_embd = self.cfg.n_embd as usize;
1683 let t = tokens.len();
1684 let base = cache.pos;
1685 debug_assert!(
1686 seq_end >= base + t,
1687 "prime_chunk_ppn: seq_end must cover this chunk"
1688 );
1689 let payload = t * n_embd;
1690 let caller_stream = e.stream();
1694 rt.fence_stages_behind(&caller_stream)?;
1695
1696 if n_st == 2 {
1697 let slot =
1698 self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1699 let x =
1700 self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1701 let out = {
1702 rt.bind_stage(1)?;
1703 let _st1 = rt.enter(1);
1704 let e1 = rt.engine(1, e);
1705 self.prime_chunk_epilogue(e1, x, t, cache)?
1706 };
1707 rt.publish_to(1, &caller_stream)?;
1708 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1709 return Ok(out);
1710 }
1711
1712 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1713
1714 let mut slot = {
1716 let _st0 = rt.enter(0);
1717 let e0 = rt.engine(0, e);
1718 let pos_d = e0.htod_i32(&pos)?;
1719 let x = self.embed(e0, tokens)?;
1720 let x =
1721 self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1722 rt.tx(0, &x, payload)?
1723 };
1725
1726 for s in 1..n_st - 1 {
1728 let _st = rt.enter(s);
1729 let es = rt.engine(s, e);
1730 let pos_d = es.htod_i32(&pos)?;
1731 let x = rt.rx(s - 1, slot, payload)?;
1732 let x = self.prime_layers(
1733 es,
1734 x,
1735 fence[s],
1736 fence[s + 1],
1737 &pos_d,
1738 t,
1739 base,
1740 cache,
1741 seq_end,
1742 )?;
1743 slot = rt.tx(s, &x, payload)?;
1744 }
1745
1746 let _stl = rt.enter(n_st - 1);
1748 let el = rt.engine(n_st - 1, e);
1749 let pos_d = el.htod_i32(&pos)?;
1750 let x = rt.rx(n_st - 2, slot, payload)?;
1751 let x = self.prime_layers(
1752 el,
1753 x,
1754 fence[n_st - 1],
1755 fence[n_st],
1756 &pos_d,
1757 t,
1758 base,
1759 cache,
1760 seq_end,
1761 )?;
1762 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1763 rt.publish_to(n_st - 1, &caller_stream)?;
1769 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1770 Ok(out)
1771 }
1772
1773 fn prime_pp2_stage0_enqueue(
1774 &self,
1775 e: &Engine,
1776 rt: &crate::pp::PpNRt,
1777 tokens: &[u32],
1778 cache: &mut Cache,
1779 seq_end: usize,
1780 fence: &[usize],
1781 base: usize,
1782 pipelined: bool,
1783 ) -> Result<usize, Box<dyn std::error::Error>> {
1784 let t = tokens.len();
1785 let n_embd = self.cfg.n_embd as usize;
1786 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1787 rt.bind_stage(0)?;
1788 let _st0 = rt.enter(0);
1789 let e0 = rt.engine(0, e);
1790 let pos_d = e0.htod_i32(&pos)?;
1791 let x = self.embed(e0, tokens)?;
1792 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1793 let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1794 if pipelined {
1795 rt.tx_pipelined(0, &x, t * n_embd)
1796 } else {
1797 rt.tx(0, &x, t * n_embd)
1798 }
1799 }
1800
1801 fn prime_pp2_stage1_enqueue(
1802 &self,
1803 e: &Engine,
1804 rt: &crate::pp::PpNRt,
1805 slot: usize,
1806 t: usize,
1807 cache: &mut Cache,
1808 seq_end: usize,
1809 fence: &[usize],
1810 base: usize,
1811 pipelined: bool,
1812 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1813 let n_embd = self.cfg.n_embd as usize;
1814 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1815 rt.bind_stage(1)?;
1816 let _st1 = rt.enter(1);
1817 let e1 = rt.engine(1, e);
1818 let pos_d = e1.htod_i32(&pos)?;
1819 let x = rt.rx(0, slot, t * n_embd)?;
1820 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1821 self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
1822 }
1823
1824 pub fn prime_chunk_captured(
1840 &self,
1841 e: &Engine,
1842 x_in: &CudaSlice<f32>,
1843 pos_d: &CudaSlice<i32>,
1844 t: usize,
1845 cache: &mut Cache,
1846 len_d: &CudaSlice<i32>,
1847 logits_out: &mut CudaSlice<f32>,
1848 h_seed_out: &mut CudaSlice<f32>,
1849 ) -> Result<(), Box<dyn std::error::Error>> {
1850 let cfg = &self.cfg;
1851 let n_embd = cfg.n_embd as usize;
1852 let eps = cfg.rms_eps;
1853 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1854 let mut x = e.uninit(t * n_embd)?;
1855 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1856 for (il, layer) in self.layers.iter().enumerate() {
1857 let mut h = e.uninit(t * n_embd)?;
1858 let mut hx16: Option<CudaSlice<u8>> = None;
1859 if f16fuse {
1860 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1861 e.rms_norm_f16out(
1862 &x,
1863 layer.attn_norm.float_data(),
1864 &mut h,
1865 &mut b16,
1866 n_embd,
1867 t,
1868 eps,
1869 )?;
1870 hx16 = Some(b16);
1871 } else {
1872 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1873 }
1874 let mixed = match &layer.mixer {
1875 Mixer::Full(fa) => {
1879 self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
1880 }
1881 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1882 Mixer::Linear(la) => {
1883 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1884 let g4 = match hx16.as_ref() {
1885 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1886 None => e.matmul_group(&ws, &h, t)?,
1887 };
1888 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1889 }
1890 };
1891 let mut x1 = e.uninit(t * n_embd)?;
1892 e.add(&x, &mixed, &mut x1, t * n_embd)?;
1893 let mut z = e.uninit(t * n_embd)?;
1894 let mut zx16: Option<CudaSlice<u8>> = None;
1895 if f16fuse {
1896 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1897 e.rms_norm_f16out(
1898 &x1,
1899 layer.post_attn_norm.float_data(),
1900 &mut z,
1901 &mut b16,
1902 n_embd,
1903 t,
1904 eps,
1905 )?;
1906 zx16 = Some(b16);
1907 } else {
1908 e.rms_norm(
1909 &x1,
1910 layer.post_attn_norm.float_data(),
1911 &mut z,
1912 n_embd,
1913 t,
1914 eps,
1915 )?;
1916 }
1917 let ffn_out = match &layer.ffn {
1918 crate::hybrid::Ffn::Dense {
1919 ffn_gate,
1920 ffn_up,
1921 ffn_down,
1922 } => {
1923 let n_ff = ffn_gate.out_features();
1924 let mut g2 = match &zx16 {
1925 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1926 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1927 };
1928 let up = g2.pop().unwrap();
1929 let gate = g2.pop().unwrap();
1930 let mut act = e.uninit(t * n_ff)?;
1931 Self::ffn_act_lim(
1933 e,
1934 &self.cfg,
1935 &gate,
1936 &up,
1937 1.0,
1938 1.0,
1939 self.cfg.clamp_shexp_at(il as u32),
1940 &mut act,
1941 t * n_ff,
1942 )?;
1943 e.matmul(ffn_down, &act, t)?
1944 }
1945 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
1946 };
1947 let mut x2 = e.uninit(t * n_embd)?;
1948 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1949 x = x2;
1950 }
1951 if !crate::spec::spec_hpost() {
1953 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1954 }
1955 let mut hn = e.uninit(t * n_embd)?;
1956 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1957 if crate::spec::spec_hpost() {
1958 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1959 }
1960 let mut hlast = e.uninit(n_embd)?;
1961 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1962 let logits = e.matmul(&self.output, &hlast, 1)?;
1963 let nv = logits.len();
1964 e.copy_into(logits_out, 0, &logits, nv)?;
1965 Ok(())
1966 }
1967
1968 fn step35_prime_batch_on() -> bool {
1969 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
1970 }
1971
1972 #[allow(clippy::too_many_arguments)]
1975 fn step35_prime_batch_layers(
1976 &self,
1977 e: &Engine,
1978 mut x: CudaSlice<f32>,
1979 lo: usize,
1980 hi: usize,
1981 ts: &[usize],
1982 offs: &[usize],
1983 pos_ds: &[CudaSlice<i32>],
1984 caches: &mut [&mut Cache],
1985 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1986 let cfg = &self.cfg;
1987 let n_embd = cfg.n_embd as usize;
1988 let eps = cfg.rms_eps;
1989 let b = ts.len();
1990 let total: usize = ts.iter().sum();
1991 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
1992
1993 let split = |e: &Engine,
1994 y: &CudaSlice<f32>,
1995 dim: usize|
1996 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1997 let mut out = Vec::with_capacity(b);
1998 for s in 0..b {
1999 let mut ys = e.uninit(ts[s] * dim)?;
2000 e.copy_view_into(
2001 &mut ys,
2002 0,
2003 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2004 ts[s] * dim,
2005 )?;
2006 out.push(ys);
2007 }
2008 Ok(out)
2009 };
2010
2011 for il in lo..hi {
2012 let layer = &self.layers[il];
2013 let Mixer::Full(fa) = &layer.mixer else {
2014 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2015 };
2016
2017 let mut h = e.uninit(total * n_embd)?;
2018 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2019 if f16fuse {
2020 e.rms_norm_f16out(
2021 &x,
2022 layer.attn_norm.float_data(),
2023 &mut h,
2024 &mut hx16,
2025 n_embd,
2026 total,
2027 eps,
2028 )?;
2029 } else {
2030 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2031 }
2032
2033 let gate_w = fa
2037 .attn_gate
2038 .as_ref()
2039 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2040 let mut g4 = if f16fuse {
2041 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2042 } else {
2043 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2044 };
2045 let gate = g4.pop().unwrap();
2046 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2047 (0..b).map(|_| Vec::with_capacity(3)).collect();
2048 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2049 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2050 parts[s].push(ys);
2051 }
2052 }
2053 let gates = split(e, &gate, gate_w.out_features())?;
2054 let geometry = self.step35_geom(il);
2055 let hd = geometry.head_dim_k as usize;
2056 let nh = geometry.n_head as usize;
2057 let mut ag_cat = e.uninit(total * nh * hd)?;
2058 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2059 let ag = self.step35_attn_pre_wo(
2060 e,
2061 fa,
2062 g3s,
2063 None,
2064 Some(&gate),
2065 &pos_ds[s],
2066 ts[s],
2067 Some(&mut *caches[s]),
2068 il,
2069 ts[s],
2070 )?;
2071 e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2072 }
2073 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2074
2075 let mut x1 = e.uninit(total * n_embd)?;
2076 let mut z = e.uninit(total * n_embd)?;
2077 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2078 if f16fuse {
2079 e.add_rms_norm_f16out(
2080 &x,
2081 &mixed,
2082 layer.post_attn_norm.float_data(),
2083 &mut x1,
2084 &mut z,
2085 &mut zx16,
2086 n_embd,
2087 total,
2088 eps,
2089 )?;
2090 } else {
2091 e.add(&x, &mixed, &mut x1, total * n_embd)?;
2092 e.rms_norm(
2093 &x1,
2094 layer.post_attn_norm.float_data(),
2095 &mut z,
2096 n_embd,
2097 total,
2098 eps,
2099 )?;
2100 }
2101
2102 let ffn_out = match &layer.ffn {
2103 crate::hybrid::Ffn::Dense {
2104 ffn_gate,
2105 ffn_up,
2106 ffn_down,
2107 } => {
2108 let n_ff = ffn_gate.out_features();
2109 let mut g2 = if f16fuse {
2110 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2111 } else {
2112 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2113 };
2114 let up = g2.pop().unwrap();
2115 let gate = g2.pop().unwrap();
2116 let mut act = e.uninit(total * n_ff)?;
2117 let d_lim = cfg.clamp_shexp_at(il as u32);
2118 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2119 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2120 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2121 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2122 Some(y) => y,
2123 None => e.matmul(ffn_down, &act, total)?,
2124 }
2125 } else {
2126 Self::ffn_act_lim(
2127 e,
2128 cfg,
2129 &gate,
2130 &up,
2131 1.0,
2132 1.0,
2133 d_lim,
2134 &mut act,
2135 total * n_ff,
2136 )?;
2137 e.matmul(ffn_down, &act, total)?
2138 }
2139 }
2140 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2141 };
2142 let mut x2 = e.uninit(total * n_embd)?;
2143 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2144 x = x2;
2145 }
2146 Ok(x)
2147 }
2148
2149 fn step35_prime_batch_epilogue(
2150 &self,
2151 e: &Engine,
2152 x: CudaSlice<f32>,
2153 ts: &[usize],
2154 offs: &[usize],
2155 caches: &mut [&mut Cache],
2156 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2157 let n_embd = self.cfg.n_embd as usize;
2158 let total: usize = ts.iter().sum();
2159 let mut hn = e.uninit(total * n_embd)?;
2160 e.rms_norm(
2161 &x,
2162 self.output_norm.float_data(),
2163 &mut hn,
2164 n_embd,
2165 total,
2166 self.cfg.rms_eps,
2167 )?;
2168
2169 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2170 let mut out = Vec::with_capacity(ts.len());
2171 for s in 0..ts.len() {
2172 let mut hidden = e.uninit(ts[s] * n_embd)?;
2173 e.copy_view_into(
2174 &mut hidden,
2175 0,
2176 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2177 ts[s] * n_embd,
2178 )?;
2179 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2180 let mut h_seed = e.uninit(n_embd)?;
2181 e.copy_view_into(
2182 &mut h_seed,
2183 0,
2184 &hidden_src.slice(last0..last0 + n_embd),
2185 n_embd,
2186 )?;
2187 let mut hlast = e.uninit(n_embd)?;
2189 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2190 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2191 caches[s].pos += ts[s];
2192 out.push((logits, h_seed, hidden));
2193 }
2194 Ok(out)
2195 }
2196
2197 fn step35_prime_cache_batch(
2198 &self,
2199 e: &Engine,
2200 prompts: &[&[u32]],
2201 caches: &mut [&mut Cache],
2202 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2203 if crate::pp::pp_host_bounce_active()
2204 && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2205 {
2206 return Err(
2207 "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2208 stage split; refusing an unsplit remote-weight walk"
2209 .into(),
2210 );
2211 }
2212 if !Self::step35_prime_batch_on() {
2213 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2214 }
2215 if caches.iter().any(|c| c.pos != 0) {
2216 return Err(
2217 "step35 batched prime currently supports complete fresh prompts only; \
2218 continuation/tick chunks require per-request queued_after"
2219 .into(),
2220 );
2221 }
2222
2223 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2224 for &t in &ts {
2225 assert!(
2226 t >= PRIME_MIN_T,
2227 "step35 batched prime needs T >= {PRIME_MIN_T}"
2228 );
2229 }
2230 for (s, c) in caches.iter().enumerate() {
2231 assert!(
2232 ts[s] <= c.max_ctx,
2233 "step35 batched prime exceeds cache max_ctx"
2234 );
2235 }
2236 let offs: Vec<usize> = ts
2237 .iter()
2238 .scan(0usize, |a, &t| {
2239 let o = *a;
2240 *a += t;
2241 Some(o)
2242 })
2243 .collect();
2244 let total: usize = ts.iter().sum();
2245 let payload = total * self.cfg.n_embd as usize;
2246 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2247 let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2248 let upload_positions =
2249 |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2250 positions
2251 .iter()
2252 .map(|p| e.htod_i32(p))
2253 .collect::<Result<_, _>>()
2254 };
2255
2256 static ONCE: std::sync::Once = std::sync::Once::new();
2257 ONCE.call_once(|| {
2258 eprintln!(
2259 "[step35-prime-batch] first concat prime: B={} tokens={total}",
2260 prompts.len()
2261 );
2262 });
2263
2264 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2265 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2266 let rt = crate::pp::PpNRt::get(e)?;
2267 let n_st = fence.len() - 1;
2268 assert_eq!(
2269 rt.n_stages(),
2270 n_st,
2271 "step35 prime batch stage count mismatch"
2272 );
2273 let caller_stream = e.stream();
2274 rt.fence_stages_behind(&caller_stream)?;
2275
2276 let mut slot = {
2277 let _st0 = rt.enter(0);
2278 let e0 = rt.engine(0, e);
2279 let pos_ds = upload_positions(e0)?;
2280 let x = self.embed(e0, &cat_tokens)?;
2281 let x = self.step35_prime_batch_layers(
2282 e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2283 )?;
2284 rt.tx(0, &x, payload)?
2285 };
2286 for s in 1..n_st - 1 {
2287 let _st = rt.enter(s);
2288 let es = rt.engine(s, e);
2289 let pos_ds = upload_positions(es)?;
2290 let x = rt.rx(s - 1, slot, payload)?;
2291 let x = self.step35_prime_batch_layers(
2292 es,
2293 x,
2294 fence[s],
2295 fence[s + 1],
2296 &ts,
2297 &offs,
2298 &pos_ds,
2299 caches,
2300 )?;
2301 slot = rt.tx(s, &x, payload)?;
2302 }
2303
2304 let _stl = rt.enter(n_st - 1);
2305 let el = rt.engine(n_st - 1, e);
2306 let pos_ds = upload_positions(el)?;
2307 let x = rt.rx(n_st - 2, slot, payload)?;
2308 let x = self.step35_prime_batch_layers(
2309 el,
2310 x,
2311 fence[n_st - 1],
2312 fence[n_st],
2313 &ts,
2314 &offs,
2315 &pos_ds,
2316 caches,
2317 )?;
2318 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2319 rt.publish_to(n_st - 1, &caller_stream)?;
2320 crate::pp::STEP35_PRIME_BATCH_SPLITS
2321 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2322 out
2323 } else {
2324 let pos_ds = upload_positions(e)?;
2325 let x = self.embed(e, &cat_tokens)?;
2326 let x = self.step35_prime_batch_layers(
2327 e,
2328 x,
2329 0,
2330 self.layers.len(),
2331 &ts,
2332 &offs,
2333 &pos_ds,
2334 caches,
2335 )?;
2336 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2337 }
2338 } else {
2339 let pos_ds = upload_positions(e)?;
2340 let x = self.embed(e, &cat_tokens)?;
2341 let x = self.step35_prime_batch_layers(
2342 e,
2343 x,
2344 0,
2345 self.layers.len(),
2346 &ts,
2347 &offs,
2348 &pos_ds,
2349 caches,
2350 )?;
2351 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2352 };
2353 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2354 Ok(out)
2355 }
2356
2357 pub fn prime_cache_batch(
2374 &self,
2375 e: &Engine,
2376 prompts: &[&[u32]],
2377 caches: &mut [&mut Cache],
2378 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2379 let cfg = &self.cfg;
2380 let n_embd = cfg.n_embd as usize;
2381 let eps = cfg.rms_eps;
2382 let b = prompts.len();
2383 assert!(b >= 1 && b == caches.len());
2384 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2385 let carried = pos0s.iter().any(|&p| p > 0);
2386 if cfg.gemma4.is_some() {
2392 return Err(
2393 "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2394 swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2395 .into(),
2396 );
2397 }
2398 if cfg.step35.is_some() {
2401 return self.step35_prime_cache_batch(e, prompts, caches);
2402 }
2403 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2404 for &t in &ts {
2405 assert!(
2406 t >= PRIME_MIN_T,
2407 "prime_cache_batch needs T >= {PRIME_MIN_T}"
2408 );
2409 }
2410 for (s, c) in caches.iter().enumerate() {
2411 assert!(
2412 c.pos + ts[s] <= c.max_ctx,
2413 "prime_cache_batch: prompt exceeds cache max_ctx"
2414 );
2415 }
2416 let total: usize = ts.iter().sum();
2417 let offs: Vec<usize> = ts
2418 .iter()
2419 .scan(0usize, |a, &t| {
2420 let o = *a;
2421 *a += t;
2422 Some(o)
2423 })
2424 .collect();
2425 let pos_ds: Vec<CudaSlice<i32>> = ts
2427 .iter()
2428 .zip(&pos0s)
2429 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2430 .collect::<Result<_, _>>()?;
2431 let split = |e: &Engine,
2433 y: &CudaSlice<f32>,
2434 dim: usize|
2435 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2436 let mut out = Vec::with_capacity(b);
2437 for s in 0..b {
2438 let mut ys = e.uninit(ts[s] * dim)?;
2439 e.copy_view_into(
2440 &mut ys,
2441 0,
2442 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2443 ts[s] * dim,
2444 )?;
2445 out.push(ys);
2446 }
2447 Ok(out)
2448 };
2449
2450 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2451 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
2453 let mut h = e.uninit(total * n_embd)?;
2454 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2455 e.rms_norm_f16out(
2456 &x,
2457 layer.attn_norm.float_data(),
2458 &mut h,
2459 &mut hx16,
2460 n_embd,
2461 total,
2462 eps,
2463 )?;
2464 let mut mixed = e.uninit(total * n_embd)?;
2466 match &layer.mixer {
2467 Mixer::Full(fa) => {
2468 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2469 let geometry = self.cfg.full_attention_geometry_at(il as u32);
2475 let (n_head, n_head_kv, head_dim) = (
2476 geometry.n_head as usize,
2477 geometry.n_head_kv as usize,
2478 geometry.head_dim_k as usize,
2479 );
2480 let fa_scale = geometry.attention_scale();
2481 let use_favl = !carried
2482 && (2..=8).contains(&b)
2483 && (head_dim == 256 || head_dim == 128)
2484 && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2485 && std::env::var("MEMRA_NOFA").is_err()
2486 && std::env::var("MEMRA_FA_FLOOR").is_err()
2487 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2488 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2489 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2490 if use_favl {
2491 let (qf_w, kf_w, vf_w) = (
2492 fa.wq.out_features(),
2493 fa.wk.out_features(),
2494 fa.wv.out_features(),
2495 );
2496 struct APre {
2497 q: CudaSlice<f32>,
2498 gate: Option<CudaSlice<f32>>,
2499 qn: CudaSlice<f32>,
2500 kn: CudaSlice<f32>,
2501 }
2502 let mut aps = Vec::with_capacity(b);
2503 for &t in ts.iter().take(b) {
2504 aps.push(APre {
2505 q: e.uninit(t * n_head * head_dim)?,
2506 gate: Some(e.uninit(t * n_head * head_dim)?),
2507 qn: e.uninit(t * n_head * head_dim)?,
2508 kn: e.uninit(t * n_head_kv * head_dim)?,
2509 });
2510 }
2511 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2512 let kvl = caches[0].kv[il].as_ref().unwrap();
2513 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2514 };
2515 let pargs: Vec<crate::AttnPreVl> = (0..b)
2516 .map(|s| {
2517 let (o, t) = (offs[s], ts[s]);
2518 let kvl = caches[s].kv[il].as_ref().unwrap();
2519 assert!(
2520 kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2521 "prime_cache_batch attn vl: fresh + capacity"
2522 );
2523 crate::AttnPreVl {
2524 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2525 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2526 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2527 q: e.addr_f32(&aps[s].q),
2528 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2529 qn: e.addr_f32(&aps[s].qn),
2530 kn: e.addr_f32(&aps[s].kn),
2531 kc: e.addr_u8(&kvl.k),
2532 vc: e.addr_u8(&kvl.v),
2533 t: t as i32,
2534 pad: 0,
2535 }
2536 })
2537 .collect();
2538 e.attn_pre_vl8(
2539 &pargs,
2540 fa.q_norm.float_data(),
2541 fa.k_norm.float_data(),
2542 head_dim,
2543 geometry.n_rot as usize,
2544 n_head,
2545 n_head_kv,
2546 self.cfg.rms_eps,
2547 geometry.rope_base,
2548 1.0,
2549 kv_dim_k,
2550 kv_dim_v,
2551 ktb,
2552 vtb,
2553 )?;
2554 for s in 0..b {
2555 let kvl = caches[s].kv[il].as_mut().unwrap();
2556 kvl.len += ts[s];
2557 let new_len = kvl.len as i32;
2558 e.set_i32_one(&mut kvl.len_d, new_len)?;
2559 }
2560 let mut attns = Vec::with_capacity(b);
2561 let mut mirrors = Vec::with_capacity(b);
2562 for &t in ts.iter().take(b) {
2563 attns.push(e.uninit(t * n_head * head_dim)?);
2564 let n = t * n_head_kv * head_dim;
2565 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2566 }
2567 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2570 Ok("0") => false,
2571 Ok("1") => true,
2572 _ => cfg!(memra_hopper_mma),
2573 };
2574 if fa3_on {
2575 let mut q16s = Vec::with_capacity(b);
2576 let mut v16s = Vec::with_capacity(b);
2577 for s in 0..b {
2578 let t = ts[s];
2579 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2580 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2581 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2582 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2583 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2584 e.f32_to_bf16_v(
2585 &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2586 &mut v16,
2587 t * n_head_kv * head_dim,
2588 )?;
2589 q16s.push(q16);
2590 v16s.push((k16, v16));
2591 }
2592 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2593 let mut kp = qp;
2594 let mut vp = qp;
2595 let mut op = [core::ptr::null_mut::<f32>(); 8];
2596 let mut tsv = [0i32; 8];
2597 for s in 0..b {
2598 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2599 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2600 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2601 op[s] = e.addr_f32(&attns[s]) as *mut f32;
2602 tsv[s] = ts[s] as i32;
2603 }
2604 let rc = unsafe {
2605 crate::fa3_vl_raw(
2606 qp.as_ptr(),
2607 kp.as_ptr(),
2608 vp.as_ptr(),
2609 op.as_ptr(),
2610 tsv.as_ptr(),
2611 b as i32,
2612 n_head as i32,
2613 n_head_kv as i32,
2614 head_dim as i32,
2615 fa_scale,
2616 e.stream().cu_stream() as *mut core::ffi::c_void,
2617 )
2618 };
2619 if rc != 0 {
2620 return Err(format!("memra_fa3_vl rc={rc}").into());
2621 }
2622 } else {
2623 let fargs: Vec<crate::FaSeqVl> = (0..b)
2624 .map(|s| crate::FaSeqVl {
2625 q: e.addr_f32(&aps[s].qn),
2626 k16: e.addr_u8(&mirrors[s].0),
2627 v16: e.addr_u8(&mirrors[s].1),
2628 o: e.addr_f32(&attns[s]),
2629 kf: e.addr_f32(&aps[s].kn),
2630 vf: e.addr_f32v(
2631 &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2632 ),
2633 t: ts[s] as i32,
2634 pad: 0,
2635 })
2636 .collect();
2637 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2638 }
2639 for (s, attn) in attns.into_iter().enumerate() {
2640 let (attn_g, ag16) = self.full_attn_prime_post_fa(
2641 e,
2642 attn,
2643 &aps[s].gate,
2644 ts[s],
2645 n_head,
2646 head_dim,
2647 )?;
2648 let mut done = false;
2649 if let Some(xh) = &ag16 {
2650 done = e.try_f16_gemm_pre_into_off(
2651 &fa.wo,
2652 xh,
2653 ts[s],
2654 &mut mixed,
2655 offs[s] * n_embd,
2656 )?;
2657 }
2658 if !done {
2659 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2660 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2661 }
2662 }
2663 } else {
2664 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2665 (0..b).map(|_| Vec::new()).collect();
2666 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2667 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2668 parts[s].push(ys);
2669 }
2670 }
2671 for (s, g3s) in parts.into_iter().enumerate() {
2672 let (attn_g, ag16) = self.full_attn_prime_core_inner(
2674 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2675 )?;
2676 let mut done = false;
2677 if let Some(xh) = &ag16 {
2678 done = e.try_f16_gemm_pre_into_off(
2679 &fa.wo,
2680 xh,
2681 ts[s],
2682 &mut mixed,
2683 offs[s] * n_embd,
2684 )?;
2685 }
2686 if !done {
2687 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2688 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2689 }
2690 }
2691 }
2692 }
2693 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2694 Mixer::Linear(la) => {
2695 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2700 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2701 let outs =
2702 self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2703 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2704 let (o, t) = (offs[s], ts[s]);
2705 let mut done = false;
2706 if let Some(xh) = &gn16 {
2707 done = e.try_f16_gemm_pre_into_off(
2708 &la.ssm_out,
2709 xh,
2710 t,
2711 &mut mixed,
2712 o * n_embd,
2713 )?;
2714 }
2715 if !done {
2716 let m = e.matmul(&la.ssm_out, &gn, t)?;
2717 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2718 }
2719 }
2720 }
2721 }
2722 let mut x1 = e.uninit(total * n_embd)?;
2723 let mut z = e.uninit(total * n_embd)?;
2724 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2725 e.add_rms_norm_f16out(
2726 &x,
2727 &mixed,
2728 layer.post_attn_norm.float_data(),
2729 &mut x1,
2730 &mut z,
2731 &mut zx16,
2732 n_embd,
2733 total,
2734 eps,
2735 )?;
2736 let ffn_out = match &layer.ffn {
2737 crate::hybrid::Ffn::Dense {
2738 ffn_gate,
2739 ffn_up,
2740 ffn_down,
2741 } => {
2742 let n_ff = ffn_gate.out_features();
2743 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2744 let up = g2.pop().unwrap();
2745 let gate = g2.pop().unwrap();
2746 let mut act = e.uninit(total * n_ff)?;
2747 let d_lim = self.cfg.clamp_shexp_at(il as u32);
2751 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2752 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2753 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2754 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2755 Some(y) => y,
2756 None => e.matmul(ffn_down, &act, total)?,
2757 }
2758 } else {
2759 Self::ffn_act_lim(
2760 e,
2761 &self.cfg,
2762 &gate,
2763 &up,
2764 1.0,
2765 1.0,
2766 d_lim,
2767 &mut act,
2768 total * n_ff,
2769 )?;
2770 e.matmul(ffn_down, &act, total)?
2771 }
2772 }
2773 crate::hybrid::Ffn::Moe(m) => {
2774 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2775 }
2776 };
2777 let mut x2 = e.uninit(total * n_embd)?;
2778 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2779 x = x2;
2780 }
2781 let mut hn = e.uninit(total * n_embd)?;
2783 e.rms_norm(
2784 &x,
2785 self.output_norm.float_data(),
2786 &mut hn,
2787 n_embd,
2788 total,
2789 eps,
2790 )?;
2791 let mut hcat = e.uninit(b * n_embd)?;
2797 for s in 0..b {
2798 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2799 e.copy_view_into(
2800 &mut hcat,
2801 s * n_embd,
2802 &hn.slice(last0..last0 + n_embd),
2803 n_embd,
2804 )?;
2805 }
2806 let logits_cat = if b >= 2 {
2807 e.try_f16_gemm(&self.output, &hcat, b)?
2808 } else {
2809 None
2810 };
2811 let logits_host: Option<Vec<f32>> = match &logits_cat {
2812 Some(lc) => Some(e.dtoh(lc)?),
2813 None => None,
2814 };
2815 let n_vocab = self.output.out_features();
2816 let mut hidden_all = if crate::spec::spec_hpost() {
2817 split(e, &hn, n_embd)?
2818 } else {
2819 split(e, &x, n_embd)?
2820 };
2821 let mut out = Vec::with_capacity(b);
2822 for s in 0..b {
2823 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2824 let mut h_seed = e.uninit(n_embd)?;
2825 if !crate::spec::spec_hpost() {
2826 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2827 } else {
2828 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2829 }
2830 let logits = match &logits_host {
2831 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2832 None => {
2833 let mut hlast = e.uninit(n_embd)?;
2834 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2835 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2836 }
2837 };
2838 caches[s].pos += ts[s];
2839 out.push((logits, h_seed, hidden_all.remove(0)));
2840 }
2841 Ok(out)
2842 }
2843
2844 #[allow(clippy::too_many_arguments)]
2855 fn full_attn_prime(
2856 &self,
2857 e: &Engine,
2858 fa: &FullAttnLayer,
2859 h: &CudaSlice<f32>,
2860 hx: Option<&CudaSlice<u8>>,
2861 pos_d: &CudaSlice<i32>,
2862 t: usize,
2863 cache: &mut Cache,
2864 il: usize,
2865 seq_end: usize,
2866 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2867 if self.cfg.step35.is_some() {
2868 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2869 }
2870 let g3 = match hx {
2875 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2876 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2877 };
2878 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2879 }
2880
2881 fn full_attn_prime_core(
2885 &self,
2886 e: &Engine,
2887 fa: &FullAttnLayer,
2888 g3: Vec<CudaSlice<f32>>,
2889 pos_d: &CudaSlice<i32>,
2890 t: usize,
2891 cache: &mut Cache,
2892 il: usize,
2893 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2894 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2895 if let Some(xh) = &ag16 {
2896 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2897 return Ok(y);
2898 }
2899 }
2900 Ok(e.matmul(&fa.wo, &attn_g, t)?)
2901 }
2902
2903 fn full_attn_prime_core_inner(
2904 &self,
2905 e: &Engine,
2906 fa: &FullAttnLayer,
2907 g3: Vec<CudaSlice<f32>>,
2908 pos_d: &CudaSlice<i32>,
2909 t: usize,
2910 cache: &mut Cache,
2911 il: usize,
2912 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2913 let cfg = &self.cfg;
2914 let geometry = cfg.full_attention_geometry_at(il as u32);
2915 let n_head = geometry.n_head as usize;
2916 let n_head_kv = geometry.n_head_kv as usize;
2917 let head_dim = geometry.head_dim_k as usize;
2918 let scale = geometry.attention_scale();
2919 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2920 let AttnPre { q, k, v, gate } = pre;
2921 let mut attn = e.uninit(t * n_head * head_dim)?;
2922 self.full_attn_prime_fa_dispatch(
2923 e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
2924 )?;
2925 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2926 }
2927
2928 #[allow(clippy::type_complexity)]
2932 fn full_attn_prime_pre_fa(
2933 &self,
2934 e: &Engine,
2935 fa: &FullAttnLayer,
2936 mut g3: Vec<CudaSlice<f32>>,
2937 pos_d: &CudaSlice<i32>,
2938 t: usize,
2939 cache: &mut Cache,
2940 il: usize,
2941 ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
2942 let cfg = &self.cfg;
2943 let geometry = cfg.full_attention_geometry_at(il as u32);
2944 let n_head = geometry.n_head as usize;
2945 let n_head_kv = geometry.n_head_kv as usize;
2946 let head_dim = geometry.head_dim_k as usize;
2947 let eps = cfg.rms_eps;
2948
2949 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
2953 let v = g3.pop().unwrap();
2954 let mut k = g3.pop().unwrap();
2955 let qf = g3.pop().unwrap();
2956 let (mut q, gate) = if gated {
2957 let mut q = e.uninit(t * n_head * head_dim)?;
2958 let mut gate = e.uninit(t * n_head * head_dim)?;
2959 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2960 (q, Some(gate))
2961 } else {
2962 (qf, None)
2963 };
2964
2965 let mut qn = e.uninit(t * n_head * head_dim)?;
2966 e.rms_norm(
2967 &q,
2968 fa.q_norm.float_data(),
2969 &mut qn,
2970 head_dim,
2971 n_head * t,
2972 eps,
2973 )?;
2974 q = qn;
2975 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2976 e.rms_norm(
2977 &k,
2978 fa.k_norm.float_data(),
2979 &mut kn,
2980 head_dim,
2981 n_head_kv * t,
2982 eps,
2983 )?;
2984 k = kn;
2985 let rope_dims = geometry.n_rot as usize;
2986 e.rope_neox(
2987 &mut q,
2988 pos_d,
2989 head_dim,
2990 rope_dims,
2991 n_head,
2992 t,
2993 geometry.rope_base,
2994 1.0,
2995 )?;
2996 e.rope_neox(
2997 &mut k,
2998 pos_d,
2999 head_dim,
3000 rope_dims,
3001 n_head_kv,
3002 t,
3003 geometry.rope_base,
3004 1.0,
3005 )?;
3006
3007 {
3010 let kvl = cache.kv[il].as_mut().unwrap();
3011 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3012 e.append_kv_quantized_rows(
3013 &k,
3014 &v,
3015 &mut kvl.k,
3016 &mut kvl.v,
3017 kvl.len,
3018 t,
3019 kvl.kv_dim_k,
3020 kvl.kv_dim_v,
3021 kvl.k_tok_bytes,
3022 kvl.v_tok_bytes,
3023 crate::Engine::kv_fp8_on(),
3024 )?;
3025 kvl.len += t;
3026 let new_len = kvl.len as i32;
3027 e.set_i32_one(&mut kvl.len_d, new_len)?;
3028 }
3029
3030 let base_len = {
3031 let kvl = cache.kv[il].as_ref().unwrap();
3032 kvl.len - t };
3034 Ok((AttnPre { q, k, v, gate }, base_len))
3035 }
3036
3037 #[allow(clippy::too_many_arguments)]
3044 fn full_attn_prime_fa_dispatch(
3045 &self,
3046 e: &Engine,
3047 q: &CudaSlice<f32>,
3048 k: &CudaSlice<f32>,
3049 v: &CudaSlice<f32>,
3050 attn: &mut CudaSlice<f32>,
3051 base_len: usize,
3052 t: usize,
3053 cache: &mut Cache,
3054 il: usize,
3055 head_dim: usize,
3056 n_head: usize,
3057 n_head_kv: usize,
3058 scale: f32,
3059 ) -> Result<(), Box<dyn std::error::Error>> {
3060 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3073 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3074 e.sdpa_naive(
3075 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3076 )?;
3077 } else {
3078 e.fa_prefill(
3079 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3080 )?;
3081 }
3082 return Ok(());
3083 }
3084 let kvl = cache.kv[il].as_ref().unwrap();
3085 let t_kv = base_len + t;
3086 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3087 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3088 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3092 e.sdpa_naive_quantized_view(
3093 q,
3094 &k_view,
3095 &v_view,
3096 attn,
3097 head_dim,
3098 n_head,
3099 n_head_kv,
3100 t,
3101 t_kv,
3102 scale,
3103 true,
3104 kvl.k_tok_bytes,
3105 kvl.v_tok_bytes,
3106 )?;
3107 return Ok(());
3108 }
3109 let deqw = std::env::var("MEMRA_PRIME_DEQW")
3117 .map(|v| v != "0")
3118 .unwrap_or(true);
3119 if deqw {
3120 e.fa_prefill_view_ws(
3121 q,
3122 &k_view,
3123 &v_view,
3124 attn,
3125 head_dim,
3126 n_head,
3127 n_head_kv,
3128 t,
3129 t_kv,
3130 scale,
3131 true,
3132 kvl.k_tok_bytes,
3133 kvl.v_tok_bytes,
3134 crate::Engine::kv_fp8_on(),
3135 )?;
3136 } else {
3137 e.fa_prefill_view(
3138 q,
3139 &k_view,
3140 &v_view,
3141 attn,
3142 head_dim,
3143 n_head,
3144 n_head_kv,
3145 t,
3146 t_kv,
3147 scale,
3148 true,
3149 kvl.k_tok_bytes,
3150 kvl.v_tok_bytes,
3151 crate::Engine::kv_fp8_on(),
3152 )?;
3153 }
3154 Ok(())
3155 }
3156
3157 fn full_attn_prime_post_fa(
3160 &self,
3161 e: &Engine,
3162 attn: CudaSlice<f32>,
3163 gate: &Option<CudaSlice<f32>>,
3164 t: usize,
3165 n_head: usize,
3166 head_dim: usize,
3167 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3168 let (attn_g, ag16) = match gate {
3169 Some(gate) => {
3170 let n = t * n_head * head_dim;
3171 let mut ag = e.uninit(n)?;
3172 if Self::f16out_on(e, t) {
3173 let mut a16 = e.alloc_u8_uninit(n * 2)?;
3174 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3175 (ag, Some(a16))
3176 } else {
3177 let mut gsig = e.uninit(n)?;
3178 e.sigmoid(gate, &mut gsig, n)?;
3179 e.mul(&attn, &gsig, &mut ag, n)?;
3180 (ag, None)
3181 }
3182 }
3183 None => (attn, None),
3184 };
3185 Ok((attn_g, ag16))
3186 }
3187
3188 fn linear_attn_prime(
3195 &self,
3196 e: &Engine,
3197 la: &LinearAttnLayer,
3198 h: &CudaSlice<f32>,
3199 hx: Option<&CudaSlice<u8>>,
3200 t: usize,
3201 cache: &mut Cache,
3202 il: usize,
3203 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3204 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3206 let g4 = match hx {
3207 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3208 None => e.matmul_group(&ws, h, t)?,
3209 };
3210 self.linear_attn_prime_core(e, la, g4, t, cache, il)
3211 }
3212
3213 fn linear_attn_prime_core(
3215 &self,
3216 e: &Engine,
3217 la: &LinearAttnLayer,
3218 mut g4: Vec<CudaSlice<f32>>,
3219 t: usize,
3220 cache: &mut Cache,
3221 il: usize,
3222 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3223 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3224 }
3225
3226 #[allow(clippy::too_many_arguments)]
3230 fn linear_attn_prime_core_pad_inner(
3231 &self,
3232 e: &Engine,
3233 la: &LinearAttnLayer,
3234 mut g4: Vec<CudaSlice<f32>>,
3235 t: usize,
3236 cache: &mut Cache,
3237 il: usize,
3238 pad_len: Option<&CudaSlice<i32>>,
3239 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3240 let ssm = self.cfg.ssm.as_ref().unwrap();
3242 let d_state = ssm.state_size as usize;
3243 let num_k = ssm.group_count as usize;
3244 let num_v = ssm.time_step_rank as usize;
3245 let key_dim = d_state * num_k;
3246 let value_dim = d_state * num_v;
3247 let conv_dim = key_dim * 2 + value_dim;
3248 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(
3253 e,
3254 la,
3255 &qkv_mixed.slice(0..t * conv_dim),
3256 &z.slice(0..t * value_dim),
3257 &beta_raw.slice(0..t * num_v),
3258 &alpha.slice(0..t * num_v),
3259 t,
3260 cache,
3261 il,
3262 pad_len,
3263 )
3264 }
3265
3266 #[allow(clippy::too_many_arguments)]
3269 fn linear_attn_gdn_prep(
3270 &self,
3271 e: &Engine,
3272 la: &LinearAttnLayer,
3273 qkv_mixed: &cudarc::driver::CudaView<f32>,
3274 beta_raw: &cudarc::driver::CudaView<f32>,
3275 alpha: &cudarc::driver::CudaView<f32>,
3276 t: usize,
3277 cache: &mut Cache,
3278 il: usize,
3279 pad_len: Option<&CudaSlice<i32>>,
3280 ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3281 let cfg = &self.cfg;
3282 let ssm = cfg.ssm.as_ref().unwrap();
3283 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;
3291 debug_assert!(
3292 t >= d_conv - 1,
3293 "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3294 );
3295
3296 let rl = cache.recur[il].as_mut().unwrap();
3301 let hk = Self::gdn_hk(e, t, num_v, num_k);
3302 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3303 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
3305 let mut k_g = e.uninit(d_state * hk * t)?;
3306 let mut v_g = e.uninit(d_state * num_v * t)?;
3307 if conv_fuse {
3308 e.ssm_conv1d_gdn_state_pad(
3309 qkv_mixed,
3310 &mut rl.conv_state,
3311 la.ssm_conv1d.float_data(),
3312 &mut q_g,
3313 &mut k_g,
3314 &mut v_g,
3315 conv_dim,
3316 t,
3317 d_conv,
3318 d_state,
3319 num_v,
3320 num_k,
3321 key_dim,
3322 hk,
3323 pad_len,
3324 )?;
3325 } else {
3326 let mut conv_out = e.uninit(conv_dim * t)?; e.ssm_conv1d_tm_state_pad_v(
3328 qkv_mixed,
3329 &mut rl.conv_state,
3330 la.ssm_conv1d.float_data(),
3331 &mut conv_out,
3332 conv_dim,
3333 t,
3334 d_conv,
3335 pad_len,
3336 )?;
3337 e.qkv_to_gdn_repack(
3338 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3339 )?;
3340 }
3341 let mut q_l2 = e.uninit(d_state * hk * t)?;
3342 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3346 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3347 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3348 Some(qb)
3349 } else {
3350 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3351 None
3352 };
3353 let mut k_l2 = e.uninit(d_state * hk * t)?;
3354 let kb16 = if Engine::l2_v2_on(d_state) {
3356 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3357 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3358 Some(kb)
3359 } else {
3360 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3361 None
3362 };
3363 let mut beta = e.uninit(t * num_v)?;
3364 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3365 let mut g_log = e.uninit(t * num_v)?;
3366 e.gdn_glog_v(
3367 alpha,
3368 la.ssm_dt.float_data(),
3369 la.ssm_a.float_data(),
3370 &mut g_log,
3371 num_v,
3372 t,
3373 )?;
3374 if let Some(len_d) = pad_len {
3375 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3376 }
3377 Ok(GdnPrep {
3378 hk,
3379 q_l2,
3380 k_l2,
3381 v_g,
3382 beta,
3383 g_log,
3384 kb16,
3385 qb16,
3386 })
3387 }
3388
3389 #[allow(clippy::too_many_arguments)]
3394 fn linear_attn_prime_core_batch(
3395 &self,
3396 e: &Engine,
3397 la: &LinearAttnLayer,
3398 g4: &[CudaSlice<f32>],
3399 offs: &[usize],
3400 ts: &[usize],
3401 caches: &mut [&mut Cache],
3402 il: usize,
3403 ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3404 let ssm = self.cfg.ssm.as_ref().unwrap();
3405 let d_state = ssm.state_size as usize;
3406 let num_k = ssm.group_count as usize;
3407 let num_v = ssm.time_step_rank as usize;
3408 let key_dim = d_state * num_k;
3409 let value_dim = d_state * num_v;
3410 let conv_dim = key_dim * 2 + value_dim;
3411 let eps = self.cfg.rms_eps;
3412 let scale = 1.0 / (d_state as f32).sqrt();
3413 let b = ts.len();
3414 let c = Engine::gdn_chunk_size();
3415 let carried = caches.iter().any(|c| c.pos > 0);
3418 let use_vl = !carried
3419 && (2..=8).contains(&b)
3420 && Engine::gdn_chunked_enabled()
3421 && ts.iter().all(|&t| t >= 16)
3422 && e.gdn_mma_enabled(c)
3423 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3424 if !use_vl {
3425 return (0..b)
3426 .map(|s| {
3427 let (o, t) = (offs[s], ts[s]);
3428 self.linear_attn_prime_core_pad_view(
3429 e,
3430 la,
3431 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3432 &g4[1].slice(o * value_dim..(o + t) * value_dim),
3433 &g4[2].slice(o * num_v..(o + t) * num_v),
3434 &g4[3].slice(o * num_v..(o + t) * num_v),
3435 t,
3436 caches[s],
3437 il,
3438 None,
3439 )
3440 })
3441 .collect();
3442 }
3443 struct SeqBufs {
3447 conv_out: CudaSlice<f32>,
3448 q_g: CudaSlice<f32>,
3449 k_g: CudaSlice<f32>,
3450 v_g: CudaSlice<f32>,
3451 q_l2: CudaSlice<f32>,
3452 k_l2: CudaSlice<f32>,
3453 beta: CudaSlice<f32>,
3454 g_log: CudaSlice<f32>,
3455 gn: CudaSlice<f32>,
3456 gn16: CudaSlice<u8>,
3457 }
3458 let d_conv = ssm.conv_kernel as usize;
3459 let f16o = Self::f16out_on(e, 16);
3460 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
3462 let mut pres = Vec::with_capacity(b);
3463 for &t in ts.iter().take(b) {
3464 sb.push(SeqBufs {
3465 conv_out: e.uninit(conv_dim * t)?,
3466 q_g: e.uninit(d_state * hk * t)?,
3467 k_g: e.uninit(d_state * hk * t)?,
3468 v_g: e.uninit(d_state * num_v * t)?,
3469 q_l2: e.uninit(d_state * hk * t)?,
3470 k_l2: e.uninit(d_state * hk * t)?,
3471 beta: e.uninit(t * num_v)?,
3472 g_log: e.uninit(t * num_v)?,
3473 gn: e.uninit(d_state * num_v * t)?,
3474 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3475 });
3476 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3477 }
3478 let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3479 .map(|s| {
3480 let (o, t) = (offs[s], ts[s]);
3481 let rl = caches[s].recur[il].as_ref().unwrap();
3482 crate::GdnPrepVl {
3483 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3484 conv_state: e.addr_f32(&rl.conv_state),
3485 conv_out: e.addr_f32(&sb[s].conv_out),
3486 q_g: e.addr_f32(&sb[s].q_g),
3487 k_g: e.addr_f32(&sb[s].k_g),
3488 v_g: e.addr_f32(&sb[s].v_g),
3489 q_l2: e.addr_f32(&sb[s].q_l2),
3490 k_l2: e.addr_f32(&sb[s].k_l2),
3491 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3492 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3493 beta: e.addr_f32(&sb[s].beta),
3494 g_log: e.addr_f32(&sb[s].g_log),
3495 o: e.addr_f32(&pres[s].o),
3496 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3497 gn: e.addr_f32(&sb[s].gn),
3498 gn16: e.addr_u8(&sb[s].gn16),
3499 kb16: if Engine::l2_v2_on(d_state) {
3500 e.addr_u8(&pres[s].kb16)
3501 } else {
3502 0
3503 },
3504 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3505 e.addr_u8(&pres[s].qb16)
3506 } else {
3507 0
3508 },
3509 t: t as i32,
3510 pad: 0,
3511 }
3512 })
3513 .collect();
3514 let args: Vec<crate::GdnSeqVl> = (0..b)
3515 .map(|s| {
3516 let rl = caches[s].recur[il].as_ref().unwrap();
3517 crate::GdnSeqVl {
3518 kb16: e.addr_u8(&pres[s].kb16),
3519 gcum: e.addr_f32(&pres[s].gcum),
3520 beta: e.addr_f32(&sb[s].beta),
3521 u: e.addr_f32(&pres[s].u),
3522 wb16: e.addr_u8(&pres[s].wb16),
3523 y: e.addr_u8(&pres[s].y16),
3524 ssnap: e.addr_u8(&pres[s].ssnap16),
3525 state_in: e.addr_f32(&rl.ssm_state),
3526 state_out: e.addr_f32(&rl.ssm_state_alt),
3527 q: e.addr_f32(&sb[s].q_l2),
3528 p: e.addr_f32(&pres[s].p),
3529 o: e.addr_f32(&pres[s].o),
3530 k: e.addr_f32(&sb[s].k_l2),
3531 v: e.addr_f32(&sb[s].v_g),
3532 g: e.addr_f32(&sb[s].g_log),
3533 a: e.addr_f32(&pres[s].a),
3534 w: e.addr_f32(&pres[s].w),
3535 t: ts[s] as i32,
3536 nc: pres[s].nc as i32,
3537 }
3538 })
3539 .collect();
3540 e.gdn_prep_vl8(
3541 &prep_args,
3542 la.ssm_conv1d.float_data(),
3543 la.ssm_dt.float_data(),
3544 la.ssm_a.float_data(),
3545 conv_dim,
3546 d_conv,
3547 d_state,
3548 num_v,
3549 num_k,
3550 key_dim,
3551 hk,
3552 eps,
3553 )?;
3554 if !Engine::l2_v2_on(d_state) {
3557 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3558 }
3559 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3561 if !Engine::l2_v2_on(d_state) {
3563 for s in 0..b {
3564 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3565 }
3566 }
3567 let mut wa = [crate::GdnWVl::default(); 8];
3568 for s in 0..b {
3569 wa[s] = crate::GdnWVl {
3570 qb16: e.addr_u8(&pres[s].qb16),
3571 pb16: e.addr_u8(&pres[s].pb16),
3572 };
3573 }
3574 Some(crate::GdnWVl8(wa))
3575 } else {
3576 None
3577 };
3578 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3579 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3580 if f16o {
3581 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3582 }
3583 let mut out = Vec::with_capacity(b);
3585 for (s, bufs) in sb.into_iter().enumerate() {
3586 let rl = caches[s].recur[il].as_mut().unwrap();
3587 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3588 let (o, t) = (offs[s], ts[s]);
3589 let SeqBufs { mut gn, gn16, .. } = bufs;
3590 if f16o {
3591 out.push((gn, Some(gn16)));
3592 } else {
3593 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3594 e.gated_rmsnorm_zv(
3595 &pres[s].o,
3596 la.ssm_norm.float_data(),
3597 &z_v,
3598 &mut gn,
3599 d_state,
3600 num_v * t,
3601 eps,
3602 )?;
3603 out.push((gn, None));
3604 }
3605 }
3606 Ok(out)
3607 }
3608
3609 #[allow(clippy::too_many_arguments)]
3613 fn linear_attn_prime_core_pad_view(
3614 &self,
3615 e: &Engine,
3616 la: &LinearAttnLayer,
3617 qkv_mixed: &cudarc::driver::CudaView<f32>,
3618 z: &cudarc::driver::CudaView<f32>,
3619 beta_raw: &cudarc::driver::CudaView<f32>,
3620 alpha: &cudarc::driver::CudaView<f32>,
3621 t: usize,
3622 cache: &mut Cache,
3623 il: usize,
3624 pad_len: Option<&CudaSlice<i32>>,
3625 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3626 let cfg = &self.cfg;
3627 let ssm = cfg.ssm.as_ref().unwrap();
3628 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
3631 let scale = 1.0 / (d_state as f32).sqrt();
3632
3633 let prep =
3634 self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3635
3636 let mut o = e.uninit(d_state * num_v * t)?;
3642 let rl = cache.recur[il].as_mut().unwrap();
3643 {
3644 let crate::cache::RecurLayer {
3645 ssm_state,
3646 ssm_state_alt,
3647 ..
3648 } = rl;
3649 e.gdn_scan_prefill(
3650 &prep.q_l2,
3651 &prep.k_l2,
3652 &prep.v_g,
3653 &prep.g_log,
3654 &prep.beta,
3655 prep.kb16.as_ref(),
3656 prep.qb16.as_ref(),
3657 ssm_state,
3658 ssm_state_alt,
3659 &mut o,
3660 num_v,
3661 t,
3662 scale,
3663 prep.hk,
3664 )?;
3665 }
3666 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3667
3668 let mut gn = e.uninit(d_state * num_v * t)?;
3671 let gn16 = if Self::f16out_on(e, t) {
3672 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3673 e.gated_rmsnorm_f16out_zv(
3674 &o,
3675 la.ssm_norm.float_data(),
3676 z,
3677 &mut gn,
3678 &mut g16,
3679 d_state,
3680 num_v * t,
3681 eps,
3682 )?;
3683 Some(g16)
3684 } else {
3685 e.gated_rmsnorm_zv(
3686 &o,
3687 la.ssm_norm.float_data(),
3688 z,
3689 &mut gn,
3690 d_state,
3691 num_v * t,
3692 eps,
3693 )?;
3694 None
3695 };
3696 Ok((gn, gn16))
3697 }
3698
3699 #[allow(clippy::too_many_arguments)]
3701 fn linear_attn_prime_core_pad(
3702 &self,
3703 e: &Engine,
3704 la: &LinearAttnLayer,
3705 g4: Vec<CudaSlice<f32>>,
3706 t: usize,
3707 cache: &mut Cache,
3708 il: usize,
3709 pad_len: Option<&CudaSlice<i32>>,
3710 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3711 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3712 if let Some(xh) = &gn16 {
3713 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3714 return Ok(y);
3715 }
3716 }
3717 Ok(e.matmul(&la.ssm_out, &gn, t)?)
3718 }
3719
3720 pub fn full_attn(
3725 &self,
3726 e: &Engine,
3727 fa: &FullAttnLayer,
3728 h: &CudaSlice<f32>,
3729 pos_d: &CudaSlice<i32>,
3730 t: usize,
3731 il: usize,
3732 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3733 if self.cfg.step35.is_some() {
3734 return self.step35_attn(e, fa, h, pos_d, t, il);
3735 }
3736 let cfg = &self.cfg;
3737 let _n_embd = cfg.n_embd as usize;
3738 let geometry = cfg.full_attention_geometry_at(il as u32);
3739 let n_head = geometry.n_head as usize;
3740 let n_head_kv = geometry.n_head_kv as usize;
3741 let head_dim = geometry.head_dim_k as usize;
3742 let eps = cfg.rms_eps;
3743 let scale = geometry.attention_scale();
3744
3745 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3748 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3750 let v = g3.pop().unwrap();
3751 let mut k = g3.pop().unwrap();
3752 let qf = g3.pop().unwrap();
3753 let (mut q, gate) = if gated {
3754 let mut q = e.uninit(t * n_head * head_dim)?;
3755 let mut gate = e.uninit(t * n_head * head_dim)?;
3756 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3757 (q, Some(gate))
3758 } else {
3759 (qf, None)
3760 };
3761
3762 let mut qn = e.uninit(t * n_head * head_dim)?;
3764 e.rms_norm(
3765 &q,
3766 fa.q_norm.float_data(),
3767 &mut qn,
3768 head_dim,
3769 n_head * t,
3770 eps,
3771 )?;
3772 q = qn;
3773 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3774 e.rms_norm(
3775 &k,
3776 fa.k_norm.float_data(),
3777 &mut kn,
3778 head_dim,
3779 n_head_kv * t,
3780 eps,
3781 )?;
3782 k = kn;
3783 let rope_dims = geometry.n_rot as usize;
3784 e.rope_neox(
3785 &mut q,
3786 pos_d,
3787 head_dim,
3788 rope_dims,
3789 n_head,
3790 t,
3791 geometry.rope_base,
3792 1.0,
3793 )?;
3794 e.rope_neox(
3795 &mut k,
3796 pos_d,
3797 head_dim,
3798 rope_dims,
3799 n_head_kv,
3800 t,
3801 geometry.rope_base,
3802 1.0,
3803 )?;
3804
3805 let mut attn = e.uninit(t * n_head * head_dim)?;
3807 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3810 e.sdpa_naive(
3812 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3813 )?;
3814 } else {
3815 e.fa_prefill(
3816 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3817 )?;
3818 }
3819
3820 let attn_g = match &gate {
3822 Some(gate) => {
3823 let mut gsig = e.uninit(t * n_head * head_dim)?;
3824 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3825 let mut ag = e.uninit(t * n_head * head_dim)?;
3826 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3827 ag
3828 }
3829 None => attn,
3830 };
3831
3832 let o = e.matmul(&fa.wo, &attn_g, t)?;
3834 Ok(o)
3835 }
3836
3837 pub fn linear_attn(
3839 &self,
3840 e: &Engine,
3841 la: &LinearAttnLayer,
3842 h: &CudaSlice<f32>,
3843 t: usize,
3844 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3845 let cfg = &self.cfg;
3846 let _n_embd = cfg.n_embd as usize;
3847 let ssm = cfg.ssm.as_ref().unwrap();
3848 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;
3853 let head_v = d_state;
3854 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;
3858 let scale = 1.0 / (d_state as f32).sqrt();
3859
3860 let mut g4 = e.matmul_group(
3863 &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
3864 h,
3865 t,
3866 )?;
3867 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);
3879 let mut q_g = e.uninit(d_state * num_v * t)?;
3880 let mut k_g = e.uninit(d_state * num_v * t)?;
3881 let mut v_g = e.uninit(d_state * num_v * t)?;
3882 e.ssm_conv1d_gdn(
3883 &qkv_mixed,
3884 la.ssm_conv1d.float_data(),
3885 &mut q_g,
3886 &mut k_g,
3887 &mut v_g,
3888 conv_dim,
3889 t,
3890 d_conv,
3891 d_state,
3892 num_v,
3893 num_k,
3894 key_dim,
3895 )?;
3896 let mut q_l2 = e.uninit(d_state * num_v * t)?;
3898 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3899 let mut k_l2 = e.uninit(d_state * num_v * t)?;
3900 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3901 let v_gd = v_g;
3902
3903 let mut beta = e.uninit(t * num_v)?;
3906 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3907 let mut g_log = e.uninit(t * num_v)?;
3909 e.gdn_glog(
3910 &alpha,
3911 la.ssm_dt.float_data(),
3912 la.ssm_a.float_data(),
3913 &mut g_log,
3914 num_v,
3915 t,
3916 )?;
3917
3918 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
3921 let mut o = e.uninit(d_state * num_v * t)?;
3922 e.gdn_scan_prefill(
3923 &q_l2,
3924 &k_l2,
3925 &v_gd,
3926 &g_log,
3927 &beta,
3928 None,
3929 None,
3930 &state_in,
3931 &mut state_out,
3932 &mut o,
3933 num_v,
3934 t,
3935 scale,
3936 num_v,
3937 )?;
3938
3939 let mut gn = e.uninit(d_state * num_v * t)?;
3944 e.gated_rmsnorm(
3945 &o,
3946 la.ssm_norm.float_data(),
3947 &z,
3948 &mut gn,
3949 d_state,
3950 num_v * t,
3951 eps,
3952 )?;
3953
3954 let out = e.matmul(&la.ssm_out, &gn, t)?;
3958 Ok(out)
3959 }
3960}
3961
3962impl HybridModel {
3963 pub fn moe_ffn_il(
3974 &self,
3975 e: &Engine,
3976 m: &MoeWeights,
3977 z: &CudaSlice<f32>,
3978 t: usize,
3979 il: u16,
3980 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3981 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
3982 }
3983
3984 pub fn moe_ffn_il_prefill(
3987 &self,
3988 e: &Engine,
3989 m: &MoeWeights,
3990 z: &CudaSlice<f32>,
3991 t: usize,
3992 il: u16,
3993 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3994 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
3995 }
3996
3997 pub fn moe_ffn_il_zq8(
4001 &self,
4002 e: &Engine,
4003 m: &MoeWeights,
4004 z: &CudaSlice<f32>,
4005 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4006 t: usize,
4007 il: u16,
4008 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4009 Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4010 }
4011
4012 pub(crate) fn moe_ffn(
4020 e: &Engine,
4021 m: &MoeWeights,
4022 z: &CudaSlice<f32>,
4023 t: usize,
4024 cfg: &ModelConfig,
4025 il: u16,
4026 max_block: usize,
4027 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4028 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4029 }
4030
4031 #[allow(clippy::too_many_arguments)]
4032 pub(crate) fn moe_ffn_inner(
4033 e: &Engine,
4034 m: &MoeWeights,
4035 z: &CudaSlice<f32>,
4036 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4037 t: usize,
4038 cfg: &ModelConfig,
4039 il: u16,
4040 max_block: usize,
4041 prefill: bool,
4042 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4043 let worker_io = crate::spill_pread::worker_enabled();
4044 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4045 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4046 e.with_moe_cache(max_block, |cache, _| {
4047 cache.begin_forward_epoch(il, t);
4048 if worker_io {
4049 cache.begin_worker_scope();
4050 }
4051 Ok(())
4052 })?;
4053 }
4054 if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4055 let moe = cfg.moe.as_ref().unwrap();
4056 let n_expert = moe.expert_count as usize;
4057 let n_used = moe.expert_used_count as usize;
4058 let sigmoid = cfg.sigmoid_router().unwrap();
4059 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4060 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4061 return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4062 }
4063 if t > 1 && moe_grouped_enabled(cfg, prefill) {
4066 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4067 if std::env::var("MEMRA_MOE_GATE").is_ok() {
4072 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4073 let g_host = e.dtoh(&grouped_out)?;
4074 let s_host = e.dtoh(&seq_out)?;
4075 let g_bytes: &[u8] = unsafe {
4076 std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4077 };
4078 let s_bytes: &[u8] = unsafe {
4079 std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4080 };
4081 if g_bytes == s_bytes {
4082 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4083 } else {
4084 let diffs = g_host
4085 .iter()
4086 .zip(s_host.iter())
4087 .enumerate()
4088 .filter(|(_, (a, b))| a != b)
4089 .count();
4090 let maxdiff = g_host
4091 .iter()
4092 .zip(s_host.iter())
4093 .map(|(a, b)| (a - b).abs())
4094 .fold(0.0f32, f32::max);
4095 panic!(
4096 "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4097 g_host.len()
4098 );
4099 }
4100 }
4101 return Ok(grouped_out);
4102 }
4103 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4104 }
4105
4106 fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4107 let Some(moe) = cfg.moe.as_ref() else {
4108 return false;
4109 };
4110 static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4113 let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4114 std::env::var("MEMRA_MOE_STATS").is_ok()
4115 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4116 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4117 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4118 || std::env::var("MEMRA_MOE_GATE").is_ok()
4119 });
4120 cfg.step35.is_some()
4121 && sigmoid_router_enabled()
4122 && moe_dev_enabled()
4123 && moe_slab_enabled()
4124 && !observation_mode
4125 && moe.expert_used_count <= 8
4126 && m.has_uniform_expert_layout()
4127 && m.gate_exps.macros.is_none()
4128 && m.up_exps.macros.is_none()
4129 && m.down_exps.macros.is_none()
4130 && !m.has_macros
4131 && moe_q8_enabled()
4132 && q8_expert_supported(m.gate_exps.qtype)
4133 && q8_expert_supported(m.up_exps.qtype)
4134 && q8_expert_supported(m.down_exps.qtype)
4135 && m.dev_exps
4136 .as_ref()
4137 .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4138 }
4139
4140 pub(crate) fn moe_ffn_sequential(
4142 e: &Engine,
4143 m: &MoeWeights,
4144 z: &CudaSlice<f32>,
4145 t: usize,
4146 cfg: &ModelConfig,
4147 il: u16,
4148 max_block: usize,
4149 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4150 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4151 }
4152
4153 fn moe_router_logits(
4157 e: &Engine,
4158 m: &MoeWeights,
4159 z: &CudaSlice<f32>,
4160 t: usize,
4161 cfg: &ModelConfig,
4162 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4163 if t < PRIME_MIN_T {
4164 if crate::router_kernel_on() {
4166 e.router_gemv(
4167 m.gate_inp.float_data(),
4168 z,
4169 cfg.n_embd as usize,
4170 m.gate_exps.n_expert,
4171 t,
4172 )
4173 } else {
4174 e.matmul_decode_exact(&m.gate_inp, z, t)
4175 }
4176 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4177 e.router_gemv(
4178 m.gate_inp.float_data(),
4179 z,
4180 cfg.n_embd as usize,
4181 m.gate_exps.n_expert,
4182 t,
4183 )
4184 } else {
4185 e.matmul(&m.gate_inp, z, t)
4186 }
4187 }
4188
4189 fn trace_moe_routes(
4193 il: u16,
4194 t: usize,
4195 sel_all: &[u32],
4196 weights: &[f32],
4197 ) -> Result<(), Box<dyn std::error::Error>> {
4198 use std::io::Write as _;
4199 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4200 let mut f = std::fs::OpenOptions::new()
4201 .create(true)
4202 .append(true)
4203 .open(path)?;
4204 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4205 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4206 }
4207 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4208 let mut f = std::fs::OpenOptions::new()
4209 .create(true)
4210 .append(true)
4211 .open(path)?;
4212 let pairs: Vec<String> = sel_all
4213 .iter()
4214 .zip(weights)
4215 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4216 .collect();
4217 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4218 }
4219 Ok(())
4220 }
4221
4222 #[allow(clippy::too_many_arguments)]
4223 fn trace_sigmoid_router_logits(
4224 e: &Engine,
4225 il: u16,
4226 t: usize,
4227 n_expert: usize,
4228 n_used: usize,
4229 logits: &CudaSlice<f32>,
4230 m: &MoeWeights,
4231 (scaling_factor, route_norm): (f32, bool),
4232 ) -> Result<(), Box<dyn std::error::Error>> {
4233 if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4234 return Ok(());
4235 }
4236 let logits = e.dtoh(logits)?;
4237 let active: Vec<u8> = m
4238 .active_experts
4239 .as_ref()
4240 .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4241 .unwrap_or_else(|| vec![1; n_expert]);
4242 let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4243 crate::sigrouter_contract::capture_served_logits(
4244 il as u32,
4245 t,
4246 n_expert,
4247 n_used,
4248 scaling_factor,
4249 route_norm,
4250 &active,
4251 &bias,
4252 &logits,
4253 )?;
4254 Ok(())
4255 }
4256
4257 fn trace_moe_input(
4262 e: &Engine,
4263 il: u16,
4264 t: usize,
4265 n_embd: usize,
4266 z: &CudaSlice<f32>,
4267 ) -> Result<(), Box<dyn std::error::Error>> {
4268 use std::io::Write as _;
4269 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4270 return Ok(());
4271 };
4272 let host = e.dtoh(z)?;
4273 if host.len() != t * n_embd {
4274 return Err(format!(
4275 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4276 host.len(),
4277 t,
4278 n_embd
4279 )
4280 .into());
4281 }
4282 let bytes = unsafe {
4283 std::slice::from_raw_parts(
4284 host.as_ptr().cast::<u8>(),
4285 host.len() * std::mem::size_of::<f32>(),
4286 )
4287 };
4288 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4289 let mut state = state
4290 .lock()
4291 .map_err(|_| "MoE input trace writer lock is poisoned")?;
4292 if state.is_none() {
4293 let dir = std::path::PathBuf::from(&dir);
4294 std::fs::create_dir_all(&dir)?;
4295 let index = std::fs::OpenOptions::new()
4296 .create(true)
4297 .append(true)
4298 .open(dir.join("index.jsonl"))?;
4299 *state = Some(MoeInputTraceWriter {
4300 dir,
4301 index,
4302 payloads: std::collections::HashMap::new(),
4303 });
4304 }
4305 let writer = state.as_mut().unwrap();
4306 if writer.dir != std::path::Path::new(&dir) {
4307 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4308 }
4309 let file_name = format!("layer-{il:03}.f32");
4310 if !writer.payloads.contains_key(&il) {
4311 let payload = std::fs::OpenOptions::new()
4312 .create(true)
4313 .append(true)
4314 .open(writer.dir.join(&file_name))?;
4315 let offset = payload.metadata()?.len();
4316 writer.payloads.insert(il, (payload, offset));
4317 }
4318 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4319 let row_offset = *offset;
4320 payload.write_all(bytes)?;
4321 *offset += bytes.len() as u64;
4322 writeln!(
4323 writer.index,
4324 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4325 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4326 \"payload_bytes\":{}}}",
4327 bytes.len()
4328 )?;
4329 Ok(())
4330 }
4331
4332 #[allow(clippy::too_many_arguments)]
4333 pub(crate) fn moe_ffn_sequential_zq8(
4334 e: &Engine,
4335 m: &MoeWeights,
4336 z: &CudaSlice<f32>,
4337 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4338 t: usize,
4339 cfg: &ModelConfig,
4340 il: u16,
4341 max_block: usize,
4342 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4343 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4344 let moe = cfg.moe.as_ref().unwrap();
4345 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);
4352 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4353 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);
4356
4357 let lim_exp = cfg.clamp_exp_at(il as u32);
4360 let lim_shexp = cfg.clamp_shexp_at(il as u32);
4361 let use_cache = Engine::moe_cache_enabled();
4362 let uniform_experts = m.has_uniform_expert_layout();
4363 let moe_q8 = uniform_experts
4364 && moe_q8_enabled()
4365 && q8_expert_supported(m.gate_exps.qtype)
4366 && q8_expert_supported(m.up_exps.qtype)
4367 && q8_expert_supported(m.down_exps.qtype);
4368 let cpu_expert_requested = crate::cpu_experts::configured();
4375 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4376 return Err(std::io::Error::other(
4377 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4378 )
4379 .into());
4380 }
4381 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4382 let freeze_cpu_residency = cpu_expert_requested
4388 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4389 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4390 .ok()
4391 .and_then(|value| value.parse::<usize>().ok())
4392 .is_some_and(|tokens| tokens > 0);
4393 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4394 e.freeze_moe_cache();
4395 }
4396 let cache_frozen = use_cache && e.moe_cache_frozen();
4397 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4398
4399 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4402 if let Some(sig) = cfg.sigmoid_router() {
4403 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4404 }
4405
4406 let no_exp_macros = m.gate_exps.macros.is_none()
4445 && m.up_exps.macros.is_none()
4446 && m.down_exps.macros.is_none();
4447 if cfg.sigmoid_router().is_none()
4451 && cfg.m3.is_none()
4452 && cfg.hy3.is_none()
4453 && !cfg.swiglu_clamped_at(il as u32)
4454 && no_exp_macros
4455 && t >= PRIME_MIN_T
4456 && m.dev_exps.is_some()
4457 && moe_q8_enabled()
4458 && q8_expert_supported(m.gate_exps.qtype)
4459 && q8_expert_supported(m.up_exps.qtype)
4460 && q8_expert_supported(m.down_exps.qtype)
4461 && std::env::var("MEMRA_MOE_PAIRS")
4462 .map(|v| v != "0")
4463 .unwrap_or(true)
4464 && std::env::var("MEMRA_MOE_STATS").is_err()
4465 {
4466 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4467 }
4468
4469 let dev_ok = uniform_experts
4487 && cfg.sigmoid_router().is_none()
4488 && cfg.m3.is_none()
4489 && cfg.hy3.is_none()
4490 && !cfg.swiglu_clamped_at(il as u32);
4491 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4495 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4496 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4497 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4498 if dev_ok
4499 && t < PRIME_MIN_T
4500 && m.dev_exps.is_some()
4501 && n_used <= 8
4502 && moe_dev_enabled()
4503 && !observe_routes
4504 {
4505 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4506 }
4507 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4508 let row_ok = e.with_moe_cache(max_block, |c, eng| {
4509 if moe_prewarm_enabled() {
4510 c.prewarm_layer(il, m, eng)?;
4511 }
4512 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4513 })?;
4514 if row_ok {
4515 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4516 }
4517 }
4518
4519 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4521 if cpu_hybrid {
4522 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4523 e,
4524 &logits,
4525 z,
4526 t,
4527 n_expert,
4528 n_used,
4529 m.exp_probs_b.as_deref(),
4530 sig,
4531 m.active_experts.as_deref(),
4532 )?;
4533 (sel, w, Some(input))
4534 } else {
4535 let (sel, w) =
4536 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4537 (sel, w, None)
4538 }
4539 } else {
4540 let (sel, w) =
4541 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4542 (sel, w, None)
4543 };
4544 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4545
4546 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4550 Self::trace_moe_input(e, il, t, n_embd, z)?;
4551
4552 let worker_disk_prefetch =
4564 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4565 let promote_worker_h2d =
4566 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4567 if promote_worker_h2d {
4568 let mut selected_blocks = Vec::with_capacity(n_used * 3);
4569 for &ex in sel_all.iter().take(n_used) {
4570 let ex = ex as u16;
4571 selected_blocks.extend([
4572 BlockId::new(il, PROJ_GATE, ex),
4573 BlockId::new(il, PROJ_UP, ex),
4574 BlockId::new(il, PROJ_DOWN, ex),
4575 ]);
4576 }
4577 for &ex in sel_all.iter().take(n_used) {
4578 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4579 }
4580 e.with_moe_cache(max_block, |cache, eng| {
4581 cache.promote_worker_reads_at_safe_boundary(
4582 &selected_blocks,
4583 &selected_blocks,
4584 eng,
4585 )?;
4586 Ok(())
4587 })?;
4588 }
4589
4590 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4593 let mut cnt = vec![0u32; n_expert];
4594 for &s in sel_all.iter() {
4595 cnt[s as usize] += 1;
4596 }
4597 let total = sel_all.len() as f64;
4598 let mut h = 0.0f64;
4599 let mut active = 0usize;
4600 for &c in &cnt {
4601 if c > 0 {
4602 active += 1;
4603 let p = c as f64 / total;
4604 h -= p * p.log2();
4605 }
4606 }
4607 let maxc = cnt.iter().copied().max().unwrap_or(0);
4608 println!(
4609 "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4610 il,
4611 t,
4612 sel_all.len(),
4613 active,
4614 n_expert,
4615 h,
4616 (n_expert as f64).log2(),
4617 total / active.max(1) as f64,
4618 maxc
4619 );
4620 }
4621
4622 let gdec_may_fire = uniform_experts
4635 && use_cache
4636 && n_used <= 8
4637 && gdec_enabled()
4638 && !cfg.swiglu_clamped_at(il as u32);
4639 let slab_local = m
4655 .dev_exps
4656 .as_ref()
4657 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4658 let slab_bases = slab_local.map(|d| {
4659 use cudarc::driver::DevicePtr;
4660 let s = e.stream();
4661 let (pg, _g0) = d.gate.device_ptr(&s);
4662 let (pu, _g1) = d.up.device_ptr(&s);
4663 let (pd, _g2) = d.down.device_ptr(&s);
4664 (pg as u64, pu as u64, pd as u64)
4665 });
4666 let slab_fused_may_fire = slab_bases.is_some()
4676 && n_used <= 8
4677 && gdec_enabled()
4678 && !cfg.swiglu_clamped_at(il as u32)
4679 && cfg.m3.is_none()
4680 && no_exp_macros
4681 && moe_q8;
4682 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4685 e.uninit(t * n_embd)?
4686 } else {
4687 e.zeros(t * n_embd)?
4688 };
4689 let cpu_input = if cpu_hybrid {
4692 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4693 } else {
4694 None
4695 };
4696
4697 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;
4705 let mut scratch_u: Option<CudaSlice<u8>> = None;
4706 let mut scratch_d: Option<CudaSlice<u8>> = None;
4707 let page_window = moe_page_prefetch_window();
4715
4716 for tok in 0..t {
4719 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4720 let w = &w_all[tok * n_used..(tok + 1) * n_used];
4721 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4723
4724 let no_macros = m.gate_exps.macros.is_none()
4738 && m.up_exps.macros.is_none()
4739 && m.down_exps.macros.is_none();
4740 if slab_fused_may_fire {
4750 let (pg, pu, pd) = slab_bases.unwrap();
4751 let mut gp = [0u64; 8];
4752 let mut up = [0u64; 8];
4753 let mut dp = [0u64; 8];
4754 for (j, &ex) in sel.iter().enumerate() {
4755 let ex = ex as usize;
4756 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4757 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4758 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4759 }
4760 let mut wv = [0f32; 8];
4761 wv[..n_used].copy_from_slice(w);
4762 if tok_q8.is_none() {
4763 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4764 }
4765 let (zq, zd) = tok_q8.as_ref().unwrap();
4766 let act = e.moe_gate_up_silu8_q8(
4767 crate::WPtr8(gp),
4768 crate::WPtr8(up),
4769 zq,
4770 zd,
4771 n_embd,
4772 n_ff_exp,
4773 n_used,
4774 m.gate_exps.qtype,
4775 m.up_exps.qtype,
4776 m.gate_exps.row_bytes,
4777 m.up_exps.row_bytes,
4778 )?;
4779 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4780 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4781 e.moe_down8_fma_q8(
4782 crate::WPtr8(dp),
4783 crate::F32x8(wv),
4784 &aq2,
4785 &ad2,
4786 &mut dst,
4787 n_ff_exp,
4788 n_embd,
4789 n_used,
4790 m.down_exps.qtype,
4791 m.down_exps.row_bytes,
4792 )?;
4793 continue;
4794 }
4795 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4796 if tok_q8.is_none() {
4797 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4798 }
4799 let (zq, zd) = tok_q8.as_ref().unwrap();
4800 if Self::moe_gdec_token_q8(
4801 e,
4802 m,
4803 il,
4804 max_block,
4805 zq,
4806 zd,
4807 sel,
4808 w,
4809 &mut moe_out,
4810 tok,
4811 n_embd,
4812 n_ff_exp,
4813 n_used,
4814 )? {
4815 continue;
4816 }
4817 } else if gdec_may_fire
4818 && cfg.m3.is_none()
4819 && no_macros
4820 && Self::moe_gdec_token(
4821 e,
4822 m,
4823 il,
4824 max_block,
4825 &zt,
4826 sel,
4827 w,
4828 &mut moe_out,
4829 tok,
4830 n_embd,
4831 n_ff_exp,
4832 n_used,
4833 )?
4834 {
4835 continue;
4836 }
4837
4838 if gdec_may_fire || slab_fused_may_fire {
4844 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4845 e.memset_zeros_view(&mut row)?;
4846 }
4847
4848 let mut cpu_mask = vec![false; sel.len()];
4854 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
4855 let gpu_resident = if use_cache {
4856 e.with_moe_cache(max_block, |cache, _| {
4857 Ok(sel
4858 .iter()
4859 .map(|&expert| {
4860 let expert = expert as u16;
4861 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
4862 .into_iter()
4863 .filter(|&projection| {
4864 cache
4865 .resident(BlockId::new(il, projection, expert))
4866 .is_some()
4867 })
4868 .count()
4869 })
4870 .collect::<Vec<_>>())
4871 })?
4872 } else {
4873 vec![0; sel.len()]
4874 };
4875 let mut cpu_selected = Vec::new();
4876 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
4877 if gpu_resident[index] != 3 {
4878 cpu_mask[index] = true;
4879 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
4880 let expert = expert as usize;
4881 cpu_selected.push((expert, route_weight));
4882 }
4883 }
4884 if crate::cpu_experts::predictor_enabled() {
4885 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4889 crate::cpu_experts::predictor_submit(il, row);
4890 }
4891 if cpu_selected.is_empty() {
4892 None
4893 } else {
4894 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4895 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
4896 .map_err(std::io::Error::other)?;
4897 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
4898 }
4899 } else {
4900 None
4901 };
4902
4903 let worker_window = worker_disk_prefetch
4904 .then(worker_prefetch_window)
4905 .unwrap_or(0);
4906 for (j, &ex) in sel.iter().enumerate() {
4907 if cpu_mask[j] {
4908 continue;
4909 }
4910 let ex = ex as usize;
4911 if let Some(d) = slab_local {
4918 let gl = m.gate_exps.expert_layout(ex);
4919 let ul = m.up_exps.expert_layout(ex);
4920 let dl = m.down_exps.expert_layout(ex);
4921 let (g0, u0, d0) = (
4922 ex * m.gate_exps.expert_stride,
4923 ex * m.up_exps.expert_stride,
4924 ex * m.down_exps.expert_stride,
4925 );
4926 let (gate, up) = if moe_q8 {
4927 if tok_q8.is_none() {
4928 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4929 }
4930 let (zq, zd) = tok_q8.as_ref().unwrap();
4931 (
4932 e.qmatvec_expert_q8(
4933 &d.gate,
4934 g0..g0 + gl.len,
4935 zq,
4936 zd,
4937 1,
4938 m.gate_exps.in_f,
4939 m.gate_exps.out_f,
4940 gl.qtype,
4941 gl.row_bytes,
4942 )?,
4943 e.qmatvec_expert_q8(
4944 &d.up,
4945 u0..u0 + ul.len,
4946 zq,
4947 zd,
4948 1,
4949 m.up_exps.in_f,
4950 m.up_exps.out_f,
4951 ul.qtype,
4952 ul.row_bytes,
4953 )?,
4954 )
4955 } else {
4956 (
4957 e.qmatvec_view(
4958 &d.gate,
4959 g0..g0 + gl.len,
4960 &zt,
4961 1,
4962 m.gate_exps.in_f,
4963 m.gate_exps.out_f,
4964 gl.qtype,
4965 gl.row_bytes,
4966 )?,
4967 e.qmatvec_view(
4968 &d.up,
4969 u0..u0 + ul.len,
4970 &zt,
4971 1,
4972 m.up_exps.in_f,
4973 m.up_exps.out_f,
4974 ul.qtype,
4975 ul.row_bytes,
4976 )?,
4977 )
4978 };
4979 let mut act = e.uninit(n_ff_exp)?;
4980 Self::ffn_act_lim(
4981 e,
4982 cfg,
4983 &gate,
4984 &up,
4985 m.gate_exps.macro_scale(ex),
4986 m.up_exps.macro_scale(ex),
4987 lim_exp,
4988 &mut act,
4989 n_ff_exp,
4990 )?;
4991 let y = if moe_q8 {
4992 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
4993 e.qmatvec_expert_q8(
4994 &d.down,
4995 d0..d0 + dl.len,
4996 &aq2,
4997 &ad2,
4998 1,
4999 m.down_exps.in_f,
5000 m.down_exps.out_f,
5001 dl.qtype,
5002 dl.row_bytes,
5003 )?
5004 } else {
5005 let actv = act.slice(0..n_ff_exp);
5006 e.qmatvec_view(
5007 &d.down,
5008 d0..d0 + dl.len,
5009 &actv,
5010 1,
5011 m.down_exps.in_f,
5012 m.down_exps.out_f,
5013 dl.qtype,
5014 dl.row_bytes,
5015 )?
5016 };
5017 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5018 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5019 continue;
5020 }
5021 for next in page_prefetch_positions(j, sel.len(), page_window) {
5022 Self::moe_prefetch_host_expert(sel[next] as usize, m);
5023 }
5024 let keep = [
5025 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5026 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5027 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5028 ];
5029 if worker_disk_prefetch && worker_window > 0 {
5030 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5031 Self::moe_prefetch_disk_expert(
5032 e,
5033 il,
5034 sel[next] as usize,
5035 m,
5036 max_block,
5037 &keep,
5038 )?;
5039 }
5040 } else if cache_dispatch
5041 && !cpu_hybrid
5042 && moe_prefetch_enabled()
5043 && j + 1 < sel.len()
5044 {
5045 let next = sel[j + 1] as usize;
5046 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5047 }
5048 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5049 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5050 if (gate_q8 || up_q8) && tok_q8.is_none() {
5053 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5054 }
5055 let gate = if gate_q8 {
5056 let (zq, zd) = tok_q8.as_ref().unwrap();
5057 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5058 } else {
5059 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5060 };
5061 let up = if up_q8 {
5062 let (zq, zd) = tok_q8.as_ref().unwrap();
5063 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5064 } else {
5065 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5066 };
5067 let mut act = e.uninit(n_ff_exp)?;
5068 Self::ffn_act_lim(
5069 e,
5070 cfg,
5071 &gate,
5072 &up,
5073 m.gate_exps.macro_scale(ex),
5074 m.up_exps.macro_scale(ex),
5075 lim_exp,
5076 &mut act,
5077 n_ff_exp,
5078 )?;
5079 let y = if down_q8 {
5080 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5081 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5082 } else {
5083 let actv = act.slice(0..n_ff_exp);
5084 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5085 };
5086 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5087 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5089 } else if cache_dispatch {
5090 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5095 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5096 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5098 e,
5099 cfg,
5100 &gate,
5101 &up,
5102 m.gate_exps.macro_scale(ex),
5103 m.up_exps.macro_scale(ex),
5104 lim_exp,
5105 &mut act,
5106 n_ff_exp,
5107 )?;
5108 let actv = act.slice(0..n_ff_exp);
5109 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5110 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5111 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5113 } else if cache_frozen {
5114 let gate = Self::moe_frozen_gemm(
5119 e,
5120 il,
5121 PROJ_GATE,
5122 ex,
5123 m,
5124 max_block,
5125 &zt,
5126 &mut scratch_g,
5127 g_len,
5128 )?;
5129 let up = Self::moe_frozen_gemm(
5130 e,
5131 il,
5132 PROJ_UP,
5133 ex,
5134 m,
5135 max_block,
5136 &zt,
5137 &mut scratch_u,
5138 u_len,
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 actv = act.slice(0..n_ff_exp);
5153 let y = Self::moe_frozen_gemm(
5154 e,
5155 il,
5156 PROJ_DOWN,
5157 ex,
5158 m,
5159 max_block,
5160 &actv,
5161 &mut scratch_d,
5162 d_len,
5163 )?;
5164 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5165 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5166 } else {
5167 if scratch_g.is_none() {
5171 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5172 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5173 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5174 }
5175 let (sg, su, sd) = (
5176 scratch_g.as_mut().unwrap(),
5177 scratch_u.as_mut().unwrap(),
5178 scratch_d.as_mut().unwrap(),
5179 );
5180 let gl = m.gate_exps.expert_layout(ex);
5181 let ul = m.up_exps.expert_layout(ex);
5182 let dl = m.down_exps.expert_layout(ex);
5183 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5184 let gate = e.qmatvec_view(
5185 sg,
5186 0..gl.len,
5187 &zt,
5188 1,
5189 m.gate_exps.in_f,
5190 m.gate_exps.out_f,
5191 gl.qtype,
5192 gl.row_bytes,
5193 )?;
5194
5195 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5196 let up = e.qmatvec_view(
5197 su,
5198 0..ul.len,
5199 &zt,
5200 1,
5201 m.up_exps.in_f,
5202 m.up_exps.out_f,
5203 ul.qtype,
5204 ul.row_bytes,
5205 )?;
5206
5207 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5209 e,
5210 cfg,
5211 &gate,
5212 &up,
5213 m.gate_exps.macro_scale(ex),
5214 m.up_exps.macro_scale(ex),
5215 lim_exp,
5216 &mut act,
5217 n_ff_exp,
5218 )?;
5219
5220 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5221 let actv = act.slice(0..n_ff_exp);
5222 let y = e.qmatvec_view(
5223 sd,
5224 0..dl.len,
5225 &actv,
5226 1,
5227 m.down_exps.in_f,
5228 m.down_exps.out_f,
5229 dl.qtype,
5230 dl.row_bytes,
5231 )?;
5232
5233 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5234 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5235 }
5236 }
5237 if let Some(worker) = cpu_worker {
5238 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5239 let cpu_output = e.htod(&cpu_output)?;
5240 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5241 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5242 }
5243 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5244 for (j, &ex) in sel.iter().enumerate() {
5245 if cpu_mask[j] {
5246 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5247 }
5248 }
5249 }
5250 }
5251
5252 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5257 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5258 {
5259 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
5268 let (sg_gate, sg_up) = if t == 1 {
5269 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5270 Some(pair) => pair,
5271 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5272 }
5273 } else if verify_t {
5274 (
5275 e.matmul_decode_exact(gate_shexp, z, t)?,
5276 e.matmul_decode_exact(up_shexp, z, t)?,
5277 )
5278 } else {
5279 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
5281 let mut sa = e.uninit(t * n_ff_sh)?; Self::ffn_act_lim(
5283 e,
5284 cfg,
5285 &sg_gate,
5286 &sg_up,
5287 1.0,
5288 1.0,
5289 lim_shexp,
5290 &mut sa,
5291 t * n_ff_sh,
5292 )?;
5293 let sh = if verify_t {
5294 e.matmul_decode_exact(down_shexp, &sa, t)?
5295 } else {
5296 e.matmul(down_shexp, &sa, t)?
5297 }; let g = match &m.gate_inp_shexp {
5311 Some(gate_inp_shexp) => {
5312 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5313 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5314 } else {
5315 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5316 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
5318 g
5319 }
5320 }
5321 None => e.htod(&vec![1.0f32; t])?,
5322 };
5323 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5325 }
5326
5327 Ok(moe_out)
5328 }
5329
5330 pub fn stage1_h2d_per_token(&self) -> u64 {
5333 use crate::hybrid::Ffn;
5334 let n_used = self
5335 .cfg
5336 .moe
5337 .as_ref()
5338 .map(|m| m.expert_used_count as u64)
5339 .unwrap_or(0);
5340 let mut bytes = 0u64;
5341 for l in self.layers.iter() {
5342 if let Ffn::Moe(m) = &l.ffn {
5343 bytes += n_used
5344 * (m.gate_exps.max_expert_bytes()
5345 + m.up_exps.max_expert_bytes()
5346 + m.down_exps.max_expert_bytes()) as u64;
5347 }
5348 }
5349 bytes
5350 }
5351
5352 pub(crate) fn max_moe_block(&self) -> usize {
5356 use crate::hybrid::Ffn;
5357 let mut mx = 0usize;
5358 let mut scan = |ffn: &Ffn| {
5359 if let Ffn::Moe(m) = ffn {
5360 mx = mx
5361 .max(m.gate_exps.max_expert_bytes())
5362 .max(m.up_exps.max_expert_bytes())
5363 .max(m.down_exps.max_expert_bytes());
5364 }
5365 };
5366 for l in self.layers.iter() {
5367 scan(&l.ffn);
5368 }
5369 if let Some(mtp) = self.mtp.as_ref() {
5370 scan(&mtp.ffn);
5371 }
5372 mx
5373 }
5374
5375 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5378 use crate::hybrid::Ffn;
5379 let mut sizes = Vec::new();
5380 let mut scan = |ffn: &Ffn| {
5381 let Ffn::Moe(m) = ffn else { return };
5382 for ex in 0..m.gate_exps.n_expert {
5383 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5384 continue;
5385 }
5386 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5387 let len = exps.expert_layout(ex).len;
5388 if len > 0 {
5389 sizes.push(len);
5390 }
5391 }
5392 }
5393 };
5394 for layer in &self.layers {
5395 scan(&layer.ffn);
5396 }
5397 if let Some(mtp) = &self.mtp {
5398 scan(&mtp.ffn);
5399 }
5400 sizes
5401 }
5402
5403 pub fn save_cpu_expert_residency_profile(
5409 &self,
5410 e: &Engine,
5411 path: &std::path::Path,
5412 ) -> Result<(), Box<dyn std::error::Error>> {
5413 let Some(ids) = e.export_moe_residency() else {
5414 return Err("no MoE residency cache to persist".into());
5415 };
5416 let mut body = format!(
5417 "memra-freeze-profile v1 max_block={} blocks={}\n",
5418 self.max_moe_block(),
5419 ids.len()
5420 );
5421 for (layer, proj, ex) in &ids {
5422 body.push_str(&format!("{layer} {proj} {ex}\n"));
5423 }
5424 let tmp = path.with_extension("tmp");
5425 std::fs::write(&tmp, body)?;
5426 std::fs::rename(&tmp, path)?;
5427 println!(
5428 "[moe-cache] freeze profile saved: {} blocks -> {}",
5429 ids.len(),
5430 path.display()
5431 );
5432 Ok(())
5433 }
5434
5435 pub fn restore_cpu_expert_residency_profile(
5439 &self,
5440 e: &Engine,
5441 path: &std::path::Path,
5442 ) -> Result<bool, Box<dyn std::error::Error>> {
5443 use crate::hybrid::Ffn;
5444 use crate::moe_cache::BlockId;
5445 let Ok(content) = std::fs::read_to_string(path) else {
5446 return Ok(false);
5447 };
5448 let mut lines = content.lines();
5449 let Some(header) = lines.next() else {
5450 return Ok(false);
5451 };
5452 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5453 if !header.starts_with(&expected) {
5454 println!(
5455 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5456 path.display()
5457 );
5458 return Ok(false);
5459 }
5460 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5461 std::collections::HashMap::new();
5462 for line in lines {
5463 let mut fields = line.split_whitespace();
5464 let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5465 else {
5466 continue;
5467 };
5468 let (Ok(layer), Ok(proj), Ok(ex)) =
5469 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5470 else {
5471 continue;
5472 };
5473 by_layer
5474 .entry(layer)
5475 .or_default()
5476 .push(BlockId::new(layer, proj, ex));
5477 }
5478 let requested: usize = by_layer.values().map(Vec::len).sum();
5479 if requested == 0 {
5480 return Ok(false);
5481 }
5482 let max_block = self.max_moe_block();
5483 let mut restaged = 0usize;
5484 let mut stage_layer =
5485 |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5486 let Ffn::Moe(m) = ffn else { return Ok(()) };
5487 let Some(ids) = by_layer.get(&layer_index) else {
5488 return Ok(());
5489 };
5490 e.with_moe_cache(max_block, |cache, eng| {
5491 for id in ids {
5492 if cache.restage_block(*id, m, eng)? {
5493 restaged += 1;
5494 }
5495 }
5496 Ok(())
5497 })
5498 };
5499 for (index, layer) in self.layers.iter().enumerate() {
5500 stage_layer(index as u16, &layer.ffn)?;
5501 }
5502 if let Some(mtp) = self.mtp.as_ref() {
5503 stage_layer(u16::MAX, &mtp.ffn)?;
5504 }
5505 e.freeze_moe_cache();
5506 println!(
5507 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5508 path.display()
5509 );
5510 Ok(true)
5511 }
5512
5513 pub fn freeze_cpu_expert_residency(
5515 &self,
5516 e: &Engine,
5517 ) -> Result<(), Box<dyn std::error::Error>> {
5518 e.freeze_moe_cache();
5519 Ok(())
5520 }
5521
5522 pub fn ffn_act(
5530 e: &Engine,
5531 cfg: &ModelConfig,
5532 gate: &CudaSlice<f32>,
5533 up: &CudaSlice<f32>,
5534 act: &mut CudaSlice<f32>,
5535 n: usize,
5536 ) -> Result<(), Box<dyn std::error::Error>> {
5537 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5538 }
5539
5540 #[allow(clippy::too_many_arguments)]
5544 pub(crate) fn ffn_act_scaled(
5545 e: &Engine,
5546 cfg: &ModelConfig,
5547 gate: &CudaSlice<f32>,
5548 up: &CudaSlice<f32>,
5549 gs: f32,
5550 us: f32,
5551 act: &mut CudaSlice<f32>,
5552 n: usize,
5553 ) -> Result<(), Box<dyn std::error::Error>> {
5554 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5555 }
5556
5557 #[allow(clippy::too_many_arguments)]
5566 pub(crate) fn ffn_act_lim(
5567 e: &Engine,
5568 cfg: &ModelConfig,
5569 gate: &CudaSlice<f32>,
5570 up: &CudaSlice<f32>,
5571 gs: f32,
5572 us: f32,
5573 limit: Option<f32>,
5574 act: &mut CudaSlice<f32>,
5575 n: usize,
5576 ) -> Result<(), Box<dyn std::error::Error>> {
5577 if let Some(m3) = cfg.m3.as_ref() {
5578 debug_assert!(
5579 limit.is_none(),
5580 "m3 swigluoai and step35 clamp are different archs"
5581 );
5582 return e.swigluoai_mul_scaled(
5583 gate,
5584 up,
5585 gs,
5586 us,
5587 m3.swiglu_alpha,
5588 m3.swiglu_limit,
5589 act,
5590 n,
5591 );
5592 }
5593 if let Some(l) = limit {
5594 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5595 }
5596 if gs == 1.0 && us == 1.0 {
5597 return e.silu_mul(gate, up, act, n);
5598 }
5599 e.silu_mul_scaled(gate, up, gs, us, act, n)
5600 }
5601
5602 fn moe_route(
5608 e: &Engine,
5609 logits: &CudaSlice<f32>,
5610 t: usize,
5611 n_expert: usize,
5612 n_used: usize,
5613 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5614 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5615 }
5616
5617 #[allow(clippy::too_many_arguments)]
5625 fn moe_route_sigmoid_cfg(
5626 e: &Engine,
5627 logits: &CudaSlice<f32>,
5628 t: usize,
5629 n_expert: usize,
5630 n_used: usize,
5631 m: &MoeWeights,
5632 (sf, route_norm): (f32, bool),
5633 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5634 if sigmoid_router_enabled() {
5635 return e.moe_router_sigmoid_topk_host(
5636 logits,
5637 t,
5638 n_expert,
5639 n_used,
5640 m.active_count(),
5641 &m.exp_probs_b_dev,
5642 &m.active_experts_dev,
5643 sf,
5644 route_norm,
5645 );
5646 }
5647 let lg = e.dtoh(logits)?;
5648 Self::moe_route_sigmoid_host(
5649 &lg,
5650 t,
5651 n_expert,
5652 n_used,
5653 m.exp_probs_b.as_deref(),
5654 sf,
5655 route_norm,
5656 m.active_experts.as_deref(),
5657 )
5658 }
5659
5660 fn moe_route_cfg(
5663 e: &Engine,
5664 logits: &CudaSlice<f32>,
5665 t: usize,
5666 n_expert: usize,
5667 n_used: usize,
5668 active: Option<&[bool]>,
5669 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5670 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5673 return e.moe_router_topk_host(logits, t, n_expert, n_used);
5674 }
5675 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
5678 let mut w_out = vec![0f32; t * n_used];
5679 for tok in 0..t {
5680 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5681 let maxl = row
5683 .iter()
5684 .enumerate()
5685 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5686 .map(|(_, &x)| x)
5687 .fold(f32::NEG_INFINITY, f32::max);
5688 let mut probs = vec![0f32; n_expert];
5689 let mut den = 0f32;
5690 for i in 0..n_expert {
5691 if active.is_some_and(|mask| !mask[i]) {
5692 continue;
5693 }
5694 let x = (row[i] - maxl).exp();
5695 probs[i] = x;
5696 den += x;
5697 }
5698 for p in probs.iter_mut() {
5699 *p /= den;
5700 }
5701 let mut idx: Vec<usize> = (0..n_expert)
5703 .filter(|&i| active.is_none_or(|mask| mask[i]))
5704 .collect();
5705 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5706 let sl = &idx[..n_used];
5707 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5708 let mut ws: f32 = wv.iter().sum();
5709 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() {
5711 *x /= ws;
5712 }
5713 for j in 0..n_used {
5714 sel[tok * n_used + j] = sl[j] as u32;
5715 w_out[tok * n_used + j] = wv[j];
5716 }
5717 }
5718 Ok((sel, w_out))
5719 }
5720
5721 #[allow(clippy::too_many_arguments)]
5722 fn moe_route_sigmoid_with_input(
5723 e: &Engine,
5724 logits: &CudaSlice<f32>,
5725 input: &CudaSlice<f32>,
5726 t: usize,
5727 n_expert: usize,
5728 n_used: usize,
5729 bias: Option<&[f32]>,
5730 (sf, route_norm): (f32, bool),
5731 active: Option<&[bool]>,
5732 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5733 let (lg, input) = e.dtoh_pair(logits, input)?;
5734 let (sel, w) =
5735 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5736 Ok((sel, w, input))
5737 }
5738
5739 pub fn start_moe_prefetch_predictor(
5744 &self,
5745 e: &Engine,
5746 cfg: &ModelConfig,
5747 ) -> Result<(), Box<dyn std::error::Error>> {
5748 use crate::hybrid::Ffn;
5749 let Some(sig) = cfg.sigmoid_router() else {
5750 return Err("prefetch predictor requires a sigmoid-router arch".into());
5751 };
5752 let resident: std::collections::HashSet<(u16, u8, u16)> = e
5753 .export_moe_residency()
5754 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5755 .into_iter()
5756 .collect();
5757 let mut layers = Vec::new();
5758 for (index, layer) in self.layers.iter().enumerate() {
5759 let Ffn::Moe(m) = &layer.ffn else { continue };
5760 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5761 continue;
5762 };
5763 let router = e.dtoh(data)?;
5764 let n_expert = m.gate_exps.n_expert;
5765 let n_embd = m.gate_exps.in_f;
5766 if router.len() != n_embd * n_expert {
5767 continue;
5768 }
5769 let build = |exps: &crate::model::HostExps| {
5770 (0..n_expert)
5771 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5772 .collect::<Vec<_>>()
5773 };
5774 layers.push((
5775 index as u16,
5776 crate::cpu_experts::PredictLayerInit {
5777 router,
5778 bias: m.exp_probs_b.clone(),
5779 active: m.active_experts.clone(),
5780 n_embd,
5781 n_used: cfg
5782 .moe
5783 .as_ref()
5784 .map(|moe| moe.expert_used_count as usize)
5785 .ok_or("prefetch predictor requires MoE config")?,
5786 sig,
5787 weights_n_expert: n_expert,
5788 gate: build(&m.gate_exps),
5789 up: build(&m.up_exps),
5790 down: build(&m.down_exps),
5791 },
5792 ));
5793 }
5794 crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5795 }
5796
5797 #[allow(clippy::too_many_arguments)]
5800 pub fn moe_route_sigmoid_host_public(
5801 logits: &[f32],
5802 t: usize,
5803 n_expert: usize,
5804 n_used: usize,
5805 bias: Option<&[f32]>,
5806 sf: f32,
5807 route_norm: bool,
5808 active: Option<&[bool]>,
5809 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5810 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
5811 }
5812
5813 #[allow(clippy::too_many_arguments)]
5814 fn moe_route_sigmoid_host(
5815 lg: &[f32],
5816 t: usize,
5817 n_expert: usize,
5818 n_used: usize,
5819 bias: Option<&[f32]>,
5820 sf: f32,
5821 route_norm: bool,
5822 active: Option<&[bool]>,
5823 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5824 let active_count = active
5825 .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
5826 .unwrap_or(n_expert);
5827 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5828 if lg.len() != t * n_expert {
5829 return Err(format!(
5830 "sigmoid router logits length mismatch: got {}, expected {}",
5831 lg.len(),
5832 t * n_expert,
5833 )
5834 .into());
5835 }
5836 let mut sel = vec![0u32; t * n_used];
5837 let mut w_out = vec![0f32; t * n_used];
5838 for tok in 0..t {
5839 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5840 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
5841 let selsc: Vec<f32> = match bias {
5843 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
5844 None => scores.clone(),
5845 };
5846 let mut idx: Vec<usize> = (0..n_expert)
5847 .filter(|&i| active.is_none_or(|mask| mask[i]))
5848 .collect();
5849 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
5850 let sl = &idx[..n_used];
5851 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
5852 if route_norm {
5853 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
5854 for x in wv.iter_mut() {
5855 *x = *x / ws * sf;
5856 }
5857 } else {
5858 for x in wv.iter_mut() {
5859 *x *= sf;
5860 }
5861 }
5862 for j in 0..n_used {
5863 sel[tok * n_used + j] = sl[j] as u32;
5864 w_out[tok * n_used + j] = wv[j];
5865 }
5866 }
5867 Ok((sel, w_out))
5868 }
5869
5870 #[allow(clippy::too_many_arguments)]
5874 fn moe_ffn_sigmoid_dev(
5875 e: &Engine,
5876 m: &MoeWeights,
5877 z: &CudaSlice<f32>,
5878 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
5879 logits: &CudaSlice<f32>,
5880 t: usize,
5881 cfg: &ModelConfig,
5882 il: u16,
5883 (scaling_factor, route_norm): (f32, bool),
5884 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5885 let moe = cfg.moe.as_ref().unwrap();
5886 let n_embd = cfg.n_embd as usize;
5887 let n_expert = moe.expert_count as usize;
5888 let n_used = moe.expert_used_count as usize;
5889 let n_ff_exp = moe.expert_ff_length as usize;
5890 let dev = m.dev_exps.as_ref().unwrap();
5891 debug_assert!(cfg.step35.is_some());
5892 debug_assert_eq!(dev.dev, e.ctx().ordinal());
5893 debug_assert!(m.has_uniform_expert_layout());
5894 debug_assert!(!m.has_macros);
5895
5896 let (sel_d, w_d) = e.moe_router_sigmoid_topk(
5897 logits,
5898 t,
5899 n_expert,
5900 n_used,
5901 m.active_count(),
5902 &m.exp_probs_b_dev,
5903 &m.active_experts_dev,
5904 scaling_factor,
5905 route_norm,
5906 )?;
5907 crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
5908 let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
5909 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5910 (combined, combined)
5911 } else {
5912 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5913 };
5914 let (zq, zd) = match (t, zq8) {
5915 (1, Some((q, d))) => (q.clone(), d.clone()),
5916 _ => e.quantize_q8_1(z, t, n_embd)?,
5917 };
5918 let n_pairs = t * n_used;
5919 let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
5920 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5924 let pair_tok_d = e.htod_i32(&pair_tok)?;
5925 let gate = e.moe_pairs_matvec_q8(
5926 &dev.ptr_row,
5927 0,
5928 &pair_tok_d,
5929 &sel_d,
5930 &zq,
5931 &zd,
5932 n_embd,
5933 n_ff_exp,
5934 n_expert,
5935 n_pairs,
5936 m.gate_exps.qtype,
5937 gate_row_bytes,
5938 )?;
5939 let up = e.moe_pairs_matvec_q8(
5940 &dev.ptr_row,
5941 1,
5942 &pair_tok_d,
5943 &sel_d,
5944 &zq,
5945 &zd,
5946 n_embd,
5947 n_ff_exp,
5948 n_expert,
5949 n_pairs,
5950 m.up_exps.qtype,
5951 up_row_bytes,
5952 )?;
5953 let mut act = e.uninit(n_pairs * n_ff_exp)?;
5954 Self::ffn_act_lim(
5955 e,
5956 cfg,
5957 &gate,
5958 &up,
5959 1.0,
5960 1.0,
5961 cfg.clamp_exp_at(il as u32),
5962 &mut act,
5963 n_pairs * n_ff_exp,
5964 )?;
5965 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5966 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5967 let pair_self_d = e.htod_i32(&pair_self)?;
5968 let down = e.moe_pairs_matvec_q8(
5969 &dev.ptr_row,
5970 2,
5971 &pair_self_d,
5972 &sel_d,
5973 &aq2,
5974 &ad2,
5975 n_ff_exp,
5976 n_embd,
5977 n_expert,
5978 n_pairs,
5979 m.down_exps.qtype,
5980 m.down_exps.row_bytes,
5981 )?;
5982 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5983 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5984 let tok_off_d = e.htod_i32(&tok_off)?;
5985 let tok_ids_d = e.htod_i32(&tok_ids)?;
5986 let mut output = e.uninit(t * n_embd)?;
5987 e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
5988 output
5989 } else {
5990 let act = e.moe_gate_up_silu8_dev_q8_rows(
5991 &dev.ptr_row,
5992 &sel_d,
5993 &zq,
5994 &zd,
5995 t,
5996 n_embd,
5997 n_ff_exp,
5998 n_used,
5999 n_expert,
6000 m.gate_exps.qtype,
6001 m.up_exps.qtype,
6002 gate_row_bytes,
6003 up_row_bytes,
6004 &m.dev_macros,
6005 )?;
6006 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6007 let mut output = e.uninit(t * n_embd)?;
6008 e.moe_down8_fma_dev_q8_rows_g(
6009 &dev.ptr_row,
6010 &sel_d,
6011 &w_d,
6012 &aq2,
6013 &ad2,
6014 &mut output,
6015 t,
6016 n_ff_exp,
6017 n_embd,
6018 n_used,
6019 n_expert,
6020 m.down_exps.qtype,
6021 m.down_exps.row_bytes,
6022 )?;
6023 output
6024 };
6025
6026 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6027 eprintln!(
6028 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6029 cfg.clamp_exp_at(il as u32).is_some(),
6030 dev.gu_il,
6031 );
6032 }
6033 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6034 Ok(moe_out)
6035 }
6036
6037 fn moe_ffn_pairs(
6046 e: &Engine,
6047 m: &MoeWeights,
6048 z: &CudaSlice<f32>,
6049 logits: &CudaSlice<f32>,
6050 t: usize,
6051 cfg: &ModelConfig,
6052 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6053 let moe = cfg.moe.as_ref().unwrap();
6054 let n_embd = cfg.n_embd as usize;
6055 let n_expert = moe.expert_count as usize;
6056 let n_used = moe.expert_used_count as usize;
6057 let n_ff_exp = moe.expert_ff_length as usize;
6058 debug_assert!(
6063 !cfg.swiglu_clamped_anywhere(),
6064 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6065 );
6066 let dev = m.dev_exps.as_ref().unwrap();
6067 let (rbg_d, rbu_d) = if dev.gu_il {
6069 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6070 (sxx, sxx)
6071 } else {
6072 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6073 };
6074
6075 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6076 let n_pairs = t * n_used;
6077 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6080 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6081 let pair_w: Vec<f32> = w_all.clone();
6082 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6083 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6084 let pt = e.htod_i32(&pair_tok)?;
6085 let px = e.htod_i32(&pair_ex)?;
6086 let pw = e.htod(&pair_w)?;
6087 let toff = e.htod_i32(&tok_off)?;
6088 let tids = e.htod_i32(&tok_ids)?;
6089
6090 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6094 for p in 0..n_pairs {
6095 by_ex[pair_ex[p] as usize].push(p as i32);
6096 }
6097 let mut ex_ids: Vec<i32> = Vec::new();
6098 let mut ex_off: Vec<i32> = vec![0];
6099 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6100 for (ex, list) in by_ex.iter().enumerate() {
6101 if list.is_empty() {
6102 continue;
6103 }
6104 ex_ids.push(ex as i32);
6105 ex_pairs.extend_from_slice(list);
6106 ex_off.push(ex_pairs.len() as i32);
6107 }
6108 let n_active = ex_ids.len();
6109 let exi = e.htod_i32(&ex_ids)?;
6110 let exo = e.htod_i32(&ex_off)?;
6111 let exp_d = e.htod_i32(&ex_pairs)?;
6112 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6133 let mma_t = *MMA_T.get_or_init(|| {
6134 std::env::var("MEMRA_MOE_MMA_T")
6135 .ok()
6136 .and_then(|v| v.parse().ok())
6137 .unwrap_or(16)
6138 });
6139 let use_mma = std::env::var("MEMRA_MOE_MMA")
6140 .map(|v| v != "0")
6141 .unwrap_or(true)
6142 && t >= mma_t
6143 && q8_expert_dec_supported(m.gate_exps.qtype)
6144 && q8_expert_dec_supported(m.up_exps.qtype)
6145 && q8_expert_dec_supported(m.down_exps.qtype)
6146 && n_embd % 256 == 0
6147 && n_ff_exp % 256 == 0;
6148 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6164 && q8_expert_dec_supported(m.up_exps.qtype)
6165 && q8_expert_dec_supported(m.down_exps.qtype)
6166 && n_embd % 256 == 0
6167 && n_ff_exp % 256 == 0;
6168 let f16g_mode = crate::moe_f16g_mode();
6169 let f16g = f16g_mode != 0
6170 && t >= mma_t
6171 && (f16g_mode != 3 || !mma_capable)
6172 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6173 && f16g_proj_ok(m.up_exps.qtype, n_embd)
6174 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6175 if use_mma || f16g {
6176 let y_down = if f16g {
6184 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6188 let csr_tok_d = e.htod_i32(&csr_tok)?;
6189 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6190 let g_csr = e.moe_f16_grouped(
6191 &dev.ptr_row,
6192 0,
6193 n_expert,
6194 &exi,
6195 &ex_off,
6196 &exo,
6197 &z_f16,
6198 &z_s,
6199 n_embd,
6200 n_ff_exp,
6201 n_active,
6202 n_pairs,
6203 m.gate_exps.qtype,
6204 rbg_d,
6205 )?;
6206 let u_csr = e.moe_f16_grouped(
6207 &dev.ptr_row,
6208 1,
6209 n_expert,
6210 &exi,
6211 &ex_off,
6212 &exo,
6213 &z_f16,
6214 &z_s,
6215 n_embd,
6216 n_ff_exp,
6217 n_active,
6218 n_pairs,
6219 m.up_exps.qtype,
6220 rbu_d,
6221 )?;
6222 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6223 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6224 let d_csr = e.moe_f16_grouped(
6225 &dev.ptr_row,
6226 2,
6227 n_expert,
6228 &exi,
6229 &ex_off,
6230 &exo,
6231 &a_f16,
6232 &a_s,
6233 n_ff_exp,
6234 n_embd,
6235 n_active,
6236 n_pairs,
6237 m.down_exps.qtype,
6238 m.down_exps.row_bytes,
6239 )?;
6240 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6241 } else {
6242 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6244 let gate = e.mmq_iq_experts(
6245 &dev.ptr_row,
6246 0,
6247 n_expert,
6248 &exi,
6249 &exo,
6250 &exp_d,
6251 &pt,
6252 &z_scr,
6253 n_embd,
6254 n_ff_exp,
6255 n_active,
6256 n_pairs,
6257 t,
6258 m.gate_exps.qtype,
6259 rbg_d,
6260 )?;
6261 let up = e.mmq_iq_experts(
6262 &dev.ptr_row,
6263 1,
6264 n_expert,
6265 &exi,
6266 &exo,
6267 &exp_d,
6268 &pt,
6269 &z_scr,
6270 n_embd,
6271 n_ff_exp,
6272 n_active,
6273 n_pairs,
6274 t,
6275 m.up_exps.qtype,
6276 rbu_d,
6277 )?;
6278 let a_scr = if crate::moe_fuse_actq_on() {
6284 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6285 } else {
6286 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6287 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6288 };
6289 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6290 let pself = e.htod_i32(&pair_self)?;
6291 e.mmq_iq_experts(
6292 &dev.ptr_row,
6293 2,
6294 n_expert,
6295 &exi,
6296 &exo,
6297 &exp_d,
6298 &pself,
6299 &a_scr,
6300 n_ff_exp,
6301 n_embd,
6302 n_active,
6303 n_pairs,
6304 n_pairs,
6305 m.down_exps.qtype,
6306 m.down_exps.row_bytes,
6307 )?
6308 };
6309 let mut moe_out = e.uninit(t * n_embd)?;
6310 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6311 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6312 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6313 {
6314 let n_ff_sh = gate_shexp.out_features();
6315 let sg_gate = e.matmul(gate_shexp, z, t)?;
6316 let sg_up = e.matmul(up_shexp, z, t)?;
6317 let mut sa = e.uninit(t * n_ff_sh)?;
6318 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6319 let sh = e.matmul(down_shexp, &sa, t)?;
6320 let g = match &m.gate_inp_shexp {
6326 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6327 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6328 }
6329 Some(gate_inp_shexp) => {
6330 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6331 let mut g = e.uninit(t)?;
6332 e.sigmoid(&gs, &mut g, t)?;
6333 g
6334 }
6335 None => e.htod(&vec![1.0f32; t])?,
6336 };
6337 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6338 }
6339 return Ok(moe_out);
6340 }
6341
6342 let dec = std::env::var("MEMRA_MOE_DEC")
6345 .map(|v| v != "0")
6346 .unwrap_or(true);
6347 let matvec = |proj,
6348 exi: &_,
6349 exo: &_,
6350 exp_d: &_,
6351 pt: &_,
6352 aq: &_,
6353 ad: &_,
6354 inf,
6355 outf,
6356 qtype,
6357 rb|
6358 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6359 let dec = dec && q8_expert_dec_supported(qtype);
6361 if dec {
6362 e.moe_pairs_matvec_q8_dec(
6363 &dev.ptr_row,
6364 proj,
6365 exi,
6366 exo,
6367 exp_d,
6368 pt,
6369 aq,
6370 ad,
6371 inf,
6372 outf,
6373 n_expert,
6374 n_active,
6375 n_pairs,
6376 qtype,
6377 rb,
6378 )
6379 } else {
6380 e.moe_pairs_matvec_q8_em(
6381 &dev.ptr_row,
6382 proj,
6383 exi,
6384 exo,
6385 exp_d,
6386 pt,
6387 aq,
6388 ad,
6389 inf,
6390 outf,
6391 n_expert,
6392 n_active,
6393 n_pairs,
6394 qtype,
6395 rb,
6396 )
6397 }
6398 };
6399 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6400 let gate = matvec(
6401 0,
6402 &exi,
6403 &exo,
6404 &exp_d,
6405 &pt,
6406 &zq,
6407 &zd,
6408 n_embd,
6409 n_ff_exp,
6410 m.gate_exps.qtype,
6411 rbg_d,
6412 )?;
6413 let up = matvec(
6414 1,
6415 &exi,
6416 &exo,
6417 &exp_d,
6418 &pt,
6419 &zq,
6420 &zd,
6421 n_embd,
6422 n_ff_exp,
6423 m.up_exps.qtype,
6424 rbu_d,
6425 )?;
6426 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6427 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6428 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6430 let pself = e.htod_i32(&pair_self)?;
6431 let y_down = matvec(
6432 2,
6433 &exi,
6434 &exo,
6435 &exp_d,
6436 &pself,
6437 &aq2,
6438 &ad2,
6439 n_ff_exp,
6440 n_embd,
6441 m.down_exps.qtype,
6442 m.down_exps.row_bytes,
6443 )?;
6444 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6446
6447 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6451 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6452 {
6453 let n_ff_sh = gate_shexp.out_features();
6454 let step_exact = cfg.step35.is_some();
6458 let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6459 let (sg_gate, sg_up) = if step_exact && t == 1 {
6460 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6461 Some(pair) => pair,
6462 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6463 }
6464 } else if verify_t {
6465 let mut fused = None;
6466 if crate::spec::spec_fused_t()
6467 && (2..=4).contains(&t)
6468 && e.uses_q8_1_fast(gate_shexp)
6469 && e.uses_q8_1_fast(up_shexp)
6470 {
6471 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6472 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6473 }
6474 match fused {
6475 Some(pair) => pair,
6476 None => (
6477 e.matmul_decode_exact(gate_shexp, z, t)?,
6478 e.matmul_decode_exact(up_shexp, z, t)?,
6479 ),
6480 }
6481 } else {
6482 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6483 };
6484 let mut sa = e.uninit(t * n_ff_sh)?;
6485 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6486 let sh = if verify_t {
6487 e.matmul_decode_exact(down_shexp, &sa, t)?
6488 } else {
6489 e.matmul(down_shexp, &sa, t)?
6490 };
6491 let g = match &m.gate_inp_shexp {
6496 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6497 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6498 }
6499 Some(gate_inp_shexp) => {
6500 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6501 let mut g = e.uninit(t)?;
6502 e.sigmoid(&gs, &mut g, t)?;
6503 g
6504 }
6505 None => e.htod(&vec![1.0f32; t])?,
6506 };
6507 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6508 }
6509 Ok(moe_out)
6510 }
6511
6512 #[allow(clippy::too_many_arguments)]
6514 #[allow(clippy::too_many_arguments)]
6515 fn moe_ffn_dev(
6516 e: &Engine,
6517 m: &MoeWeights,
6518 z: &CudaSlice<f32>,
6519 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6520 logits: &CudaSlice<f32>,
6521 t: usize,
6522 cfg: &ModelConfig,
6523 il: u16,
6524 max_block: usize,
6525 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6526 let moe = cfg.moe.as_ref().unwrap();
6527 let n_embd = cfg.n_embd as usize;
6528 let n_expert = moe.expert_count as usize;
6529 let n_used = moe.expert_used_count as usize;
6530 let n_ff_exp = moe.expert_ff_length as usize;
6531 debug_assert!(
6535 cfg.sigmoid_router().is_none(),
6536 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6537 );
6538 debug_assert!(
6539 !cfg.swiglu_clamped_at(il as u32),
6540 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6541 );
6542
6543 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6545 if m.has_macros {
6548 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6549 }
6550
6551 let mut moe_out = e.uninit(t * n_embd)?;
6553
6554 if let Some(dev) = m.dev_exps.as_ref() {
6557 let (rbg_d, rbu_d) = if dev.gu_il {
6560 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6561 (sxx, sxx)
6562 } else {
6563 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6564 };
6565 let q8 = moe_q8_enabled()
6566 && q8_expert_supported(m.gate_exps.qtype)
6567 && q8_expert_supported(m.up_exps.qtype)
6568 && q8_expert_supported(m.down_exps.qtype);
6569 let rows_arm = q8
6578 && t > 1
6579 && crate::spec::spec_m2()
6580 && n_ff_exp == 512
6581 && n_used <= 8
6582 && std::env::var("MEMRA_MOE_DEVQ8_GU")
6583 .map(|v| v.is_empty() || v == "v")
6584 .unwrap_or(true)
6585 && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6586 .map(|v| v.is_empty() || v == "w8h2v")
6587 .unwrap_or(true);
6588 let csr_mode = std::env::var("MEMRA_MOE_CSR")
6597 .ok()
6598 .and_then(|v| v.parse::<i32>().ok())
6599 .unwrap_or(1);
6600 let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6601 let csr_arm = rows_arm
6602 && csr_mode > 0
6603 && t <= 10
6604 && csr_qt(m.gate_exps.qtype)
6605 && csr_qt(m.up_exps.qtype)
6606 && csr_qt(m.down_exps.qtype);
6607 if csr_arm {
6608 if csr_mode == 2 {
6609 static ENGAGED: std::sync::Once = std::sync::Once::new();
6610 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6611 }
6612 let n_pairs = t * n_used;
6613 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6614 let act = e.moe_gate_up_silu8_dev_q8_csr(
6615 &dev.ptr_row,
6616 &sel_d,
6617 &zq,
6618 &zd,
6619 n_pairs,
6620 n_embd,
6621 n_ff_exp,
6622 n_used,
6623 n_expert,
6624 m.gate_exps.qtype,
6625 m.up_exps.qtype,
6626 rbg_d,
6627 rbu_d,
6628 )?;
6629 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6630 e.moe_down8_fma_dev_q8_rows(
6634 &dev.ptr_row,
6635 &sel_d,
6636 &w_d,
6637 &aq2,
6638 &ad2,
6639 &mut moe_out,
6640 t,
6641 n_ff_exp,
6642 n_embd,
6643 n_used,
6644 n_expert,
6645 m.down_exps.qtype,
6646 m.down_exps.row_bytes,
6647 )?;
6648 if csr_mode == 2 {
6649 let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6651 &dev.ptr_row,
6652 &sel_d,
6653 &zq,
6654 &zd,
6655 t,
6656 n_embd,
6657 n_ff_exp,
6658 n_used,
6659 n_expert,
6660 m.gate_exps.qtype,
6661 m.up_exps.qtype,
6662 rbg_d,
6663 rbu_d,
6664 &m.dev_macros,
6665 )?;
6666 let mut out_r = e.uninit(t * n_embd)?;
6667 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6668 e.moe_down8_fma_dev_q8_rows(
6669 &dev.ptr_row,
6670 &sel_d,
6671 &w_d,
6672 &aq2r,
6673 &ad2r,
6674 &mut out_r,
6675 t,
6676 n_ff_exp,
6677 n_embd,
6678 n_used,
6679 n_expert,
6680 m.down_exps.qtype,
6681 m.down_exps.row_bytes,
6682 )?;
6683 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6684 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6685 let ba = a1
6686 .iter()
6687 .zip(&a2)
6688 .filter(|(x, y)| x.to_bits() != y.to_bits())
6689 .count();
6690 let bo = o1
6691 .iter()
6692 .zip(&o2)
6693 .filter(|(x, y)| x.to_bits() != y.to_bits())
6694 .count();
6695 if ba + bo > 0 {
6696 eprintln!(
6697 "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6698 a1.len(),
6699 o1.len()
6700 );
6701 let sel_h = e.dtoh_i32(&sel_d)?;
6703 let mut shown = 0;
6704 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6705 if x.to_bits() != y.to_bits() && shown < 4 {
6706 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6707 let ex = sel_h[p];
6708 let npx = sel_h.iter().filter(|&&v| v == ex).count();
6709 eprintln!(
6710 " ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6711 );
6712 shown += 1;
6713 }
6714 }
6715 std::process::exit(3);
6716 }
6717 }
6718 } else if rows_arm {
6719 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6722 use std::sync::atomic::{AtomicU64, Ordering};
6723 static PAIRS: AtomicU64 = AtomicU64::new(0);
6724 static UNIQ: AtomicU64 = AtomicU64::new(0);
6725 static CALLS: AtomicU64 = AtomicU64::new(0);
6726 let sel_h = e.dtoh_i32(&sel_d)?;
6727 let mut u: Vec<i32> = sel_h.clone();
6728 u.sort_unstable();
6729 u.dedup();
6730 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6731 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6732 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6733 if c % 480 == 0 {
6734 let p = PAIRS.load(Ordering::Relaxed);
6735 let q = UNIQ.load(Ordering::Relaxed);
6736 eprintln!(
6737 "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6738 q as f64 / p as f64
6739 );
6740 }
6741 }
6742 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6743 let act = e.moe_gate_up_silu8_dev_q8_rows(
6744 &dev.ptr_row,
6745 &sel_d,
6746 &zq,
6747 &zd,
6748 t,
6749 n_embd,
6750 n_ff_exp,
6751 n_used,
6752 n_expert,
6753 m.gate_exps.qtype,
6754 m.up_exps.qtype,
6755 rbg_d,
6756 rbu_d,
6757 &m.dev_macros,
6758 )?;
6759 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6760 e.moe_down8_fma_dev_q8_rows(
6761 &dev.ptr_row,
6762 &sel_d,
6763 &w_d,
6764 &aq2,
6765 &ad2,
6766 &mut moe_out,
6767 t,
6768 n_ff_exp,
6769 n_embd,
6770 n_used,
6771 n_expert,
6772 m.down_exps.qtype,
6773 m.down_exps.row_bytes,
6774 )?;
6775 } else {
6776 for tok in 0..t {
6777 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6778 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6779 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6780 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6781 if q8 {
6782 let (zq, zd) = match (t, zq8) {
6783 (1, Some((q, d))) => (q.clone(), d.clone()),
6784 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6785 };
6786 let act = e.moe_gate_up_silu8_dev_q8(
6787 &dev.ptr_row,
6788 &selt,
6789 &zq,
6790 &zd,
6791 n_embd,
6792 n_ff_exp,
6793 n_used,
6794 n_expert,
6795 m.gate_exps.qtype,
6796 m.up_exps.qtype,
6797 rbg_d,
6798 rbu_d,
6799 &m.dev_macros,
6800 )?;
6801 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6802 e.moe_down8_fma_dev_q8(
6803 &dev.ptr_row,
6804 &selt,
6805 &wt,
6806 &aq2,
6807 &ad2,
6808 &mut dst,
6809 n_ff_exp,
6810 n_embd,
6811 n_used,
6812 n_expert,
6813 m.down_exps.qtype,
6814 m.down_exps.row_bytes,
6815 )?;
6816 } else {
6817 let act = e.moe_gate_up_silu8_dev(
6818 &dev.ptr_row,
6819 &selt,
6820 &zt,
6821 n_embd,
6822 n_ff_exp,
6823 n_used,
6824 n_expert,
6825 m.gate_exps.qtype,
6826 m.up_exps.qtype,
6827 rbg_d,
6828 rbu_d,
6829 &m.dev_macros,
6830 )?;
6831 e.moe_down8_fma_dev(
6832 &dev.ptr_row,
6833 &selt,
6834 &wt,
6835 &act,
6836 &mut dst,
6837 n_ff_exp,
6838 n_embd,
6839 n_used,
6840 n_expert,
6841 m.down_exps.qtype,
6842 m.down_exps.row_bytes,
6843 )?;
6844 }
6845 }
6846 }
6847 } else {
6848 let q8 = moe_q8_enabled()
6855 && q8_expert_supported(m.gate_exps.qtype)
6856 && q8_expert_supported(m.up_exps.qtype)
6857 && q8_expert_supported(m.down_exps.qtype);
6858 e.with_moe_cache(max_block, |c, eng| {
6859 let row = c
6860 .layer_dev_row(il, n_expert, eng)?
6861 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
6862 for tok in 0..t {
6863 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6864 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6865 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6866 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6867 if q8 {
6868 let (zq, zd) = match (t, zq8) {
6869 (1, Some((q, d))) => (q.clone(), d.clone()),
6870 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
6871 };
6872 let act = eng.moe_gate_up_silu8_dev_q8(
6873 row,
6874 &selt,
6875 &zq,
6876 &zd,
6877 n_embd,
6878 n_ff_exp,
6879 n_used,
6880 n_expert,
6881 m.gate_exps.qtype,
6882 m.up_exps.qtype,
6883 m.gate_exps.row_bytes,
6884 m.up_exps.row_bytes,
6885 &m.dev_macros,
6886 )?;
6887 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
6888 eng.moe_down8_fma_dev_q8(
6889 row,
6890 &selt,
6891 &wt,
6892 &aq2,
6893 &ad2,
6894 &mut dst,
6895 n_ff_exp,
6896 n_embd,
6897 n_used,
6898 n_expert,
6899 m.down_exps.qtype,
6900 m.down_exps.row_bytes,
6901 )?;
6902 } else {
6903 let act = eng.moe_gate_up_silu8_dev(
6904 row,
6905 &selt,
6906 &zt,
6907 n_embd,
6908 n_ff_exp,
6909 n_used,
6910 n_expert,
6911 m.gate_exps.qtype,
6912 m.up_exps.qtype,
6913 m.gate_exps.row_bytes,
6914 m.up_exps.row_bytes,
6915 &m.dev_macros,
6916 )?;
6917 eng.moe_down8_fma_dev(
6918 row,
6919 &selt,
6920 &wt,
6921 &act,
6922 &mut dst,
6923 n_ff_exp,
6924 n_embd,
6925 n_used,
6926 n_expert,
6927 m.down_exps.qtype,
6928 m.down_exps.row_bytes,
6929 )?;
6930 }
6931 }
6932 c.hits += (t * 3 * n_used) as u64;
6934 Ok(())
6935 })?;
6936 }
6937
6938 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6943 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6944 {
6945 let n_ff_sh = gate_shexp.out_features();
6946 let verify_t = t > 1 && t < PRIME_MIN_T;
6949 let (sg_gate, sg_up) = if t == 1 {
6950 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6951 Some(pair) => pair,
6952 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6953 }
6954 } else if verify_t {
6955 let mut fused = None;
6959 if crate::spec::spec_fused_t()
6960 && (2..=4).contains(&t)
6961 && e.uses_q8_1_fast(gate_shexp)
6962 && e.uses_q8_1_fast(up_shexp)
6963 {
6964 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6965 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6966 }
6967 match fused {
6968 Some(pair) => pair,
6969 None => (
6970 e.matmul_decode_exact(gate_shexp, z, t)?,
6971 e.matmul_decode_exact(up_shexp, z, t)?,
6972 ),
6973 }
6974 } else {
6975 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6976 };
6977 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6979 let sh = if verify_t {
6980 e.matmul_decode_exact(down_shexp, &sa, t)?
6981 } else {
6982 e.matmul(down_shexp, &sa, t)?
6983 };
6984 let g = match &m.gate_inp_shexp {
6988 Some(gate_inp_shexp) => {
6989 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
6992 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6993 } else {
6994 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6995 let mut g = e.uninit(t)?;
6996 e.sigmoid(&gs, &mut g, t)?;
6997 g
6998 }
6999 }
7000 None => e.htod(&vec![1.0f32; t])?,
7001 };
7002 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7003 }
7004
7005 Ok(moe_out)
7006 }
7007
7008 #[allow(clippy::too_many_arguments)]
7018 #[allow(clippy::too_many_arguments)]
7021 fn moe_gdec_token_q8(
7022 e: &Engine,
7023 m: &MoeWeights,
7024 il: u16,
7025 max_block: usize,
7026 zq: &CudaSlice<i8>,
7027 zd: &CudaSlice<f32>,
7028 sel: &[u32],
7029 w: &[f32],
7030 moe_out: &mut CudaSlice<f32>,
7031 tok: usize,
7032 n_embd: usize,
7033 n_ff_exp: usize,
7034 n_used: usize,
7035 ) -> Result<bool, Box<dyn std::error::Error>> {
7036 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7037 use cudarc::driver::DevicePtr;
7038 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7039 let mut g = [0u64; 8];
7040 let mut u = [0u64; 8];
7041 let mut d = [0u64; 8];
7042 for (j, &ex) in sel.iter().enumerate() {
7043 let ex = ex as u16;
7044 let (Some(sg), Some(su), Some(sd)) = (
7045 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7046 c.resident(BlockId::new(il, PROJ_UP, ex)),
7047 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7048 ) else {
7049 return Ok(None);
7050 };
7051 let __s = eng.stream();
7052 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7053 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7054 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7055 g[j] = pg as u64;
7056 u[j] = pu as u64;
7057 d[j] = pd as u64;
7058 }
7059 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7060 for &ex in sel {
7061 let ex = ex as u16;
7062 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7063 c.note_profile_hit(BlockId::new(il, proj, ex));
7064 }
7065 }
7066 }
7067 c.hits += (3 * n_used) as u64;
7068 Ok(Some((g, u, d)))
7069 })?;
7070 let Some((g, u, d)) = ptrs else {
7071 return Ok(false);
7072 };
7073 let mut wv = [0f32; 8];
7074 wv[..n_used].copy_from_slice(w);
7075 let act = e.moe_gate_up_silu8_q8(
7076 crate::WPtr8(g),
7077 crate::WPtr8(u),
7078 zq,
7079 zd,
7080 n_embd,
7081 n_ff_exp,
7082 n_used,
7083 m.gate_exps.qtype,
7084 m.up_exps.qtype,
7085 m.gate_exps.row_bytes,
7086 m.up_exps.row_bytes,
7087 )?;
7088 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7090 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7091 e.moe_down8_fma_q8(
7092 crate::WPtr8(d),
7093 crate::F32x8(wv),
7094 &aq2,
7095 &ad2,
7096 &mut dst,
7097 n_ff_exp,
7098 n_embd,
7099 n_used,
7100 m.down_exps.qtype,
7101 m.down_exps.row_bytes,
7102 )?;
7103 Ok(true)
7104 }
7105
7106 fn moe_gdec_token(
7107 e: &Engine,
7108 m: &MoeWeights,
7109 il: u16,
7110 max_block: usize,
7111 zt: &cudarc::driver::CudaView<f32>,
7112 sel: &[u32],
7113 w: &[f32],
7114 moe_out: &mut CudaSlice<f32>,
7115 tok: usize,
7116 n_embd: usize,
7117 n_ff_exp: usize,
7118 n_used: usize,
7119 ) -> Result<bool, Box<dyn std::error::Error>> {
7120 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7121 use cudarc::driver::DevicePtr;
7122 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7124 let mut g = [0u64; 8];
7125 let mut u = [0u64; 8];
7126 let mut d = [0u64; 8];
7127 for (j, &ex) in sel.iter().enumerate() {
7128 let ex = ex as u16;
7129 let (Some(sg), Some(su), Some(sd)) = (
7130 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7131 c.resident(BlockId::new(il, PROJ_UP, ex)),
7132 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7133 ) else {
7134 return Ok(None);
7135 };
7136 let __s = eng.stream();
7137 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7138 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7139 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7140 g[j] = pg as u64;
7141 u[j] = pu as u64;
7142 d[j] = pd as u64;
7143 }
7144 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7145 for &ex in sel {
7146 let ex = ex as u16;
7147 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7148 c.note_profile_hit(BlockId::new(il, proj, ex));
7149 }
7150 }
7151 }
7152 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
7154 })?;
7155 let Some((g, u, d)) = ptrs else {
7156 return Ok(false);
7157 };
7158 let mut wv = [0f32; 8];
7159 wv[..n_used].copy_from_slice(w);
7160 let act = e.moe_gate_up_silu8(
7162 crate::WPtr8(g),
7163 crate::WPtr8(u),
7164 zt,
7165 n_embd,
7166 n_ff_exp,
7167 n_used,
7168 m.gate_exps.qtype,
7169 m.up_exps.qtype,
7170 m.gate_exps.row_bytes,
7171 m.up_exps.row_bytes,
7172 )?;
7173 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7174 e.moe_down8_fma_into(
7175 crate::WPtr8(d),
7176 crate::F32x8(wv),
7177 &act,
7178 &mut dst,
7179 n_ff_exp,
7180 n_embd,
7181 n_used,
7182 m.down_exps.qtype,
7183 m.down_exps.row_bytes,
7184 )?;
7185 Ok(true)
7186 }
7187
7188 fn moe_cached_gemm_q8(
7193 e: &Engine,
7194 il: u16,
7195 proj: u8,
7196 ex: usize,
7197 m: &MoeWeights,
7198 max_block: usize,
7199 aq: &CudaSlice<i8>,
7200 ad: &CudaSlice<f32>,
7201 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7202 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7203 let exps = match proj {
7204 PROJ_GATE => &m.gate_exps,
7205 PROJ_UP => &m.up_exps,
7206 _ => &m.down_exps,
7207 };
7208 let layout = exps.expert_layout(ex);
7209 let id = BlockId::new(il, proj, ex as u16);
7210 let source = exps.expert_source(ex);
7211 e.with_moe_cache(max_block, |c, eng| {
7212 let slot = c.dispatch_source(id, source, eng)?;
7213 let DispatchSlot::Resident(sl) = slot;
7214 let buf = c.slot(sl);
7215 eng.qmatvec_expert_q8(
7216 buf,
7217 0..layout.len,
7218 aq,
7219 ad,
7220 1,
7221 exps.in_f,
7222 exps.out_f,
7223 layout.qtype,
7224 layout.row_bytes,
7225 )
7226 })
7227 }
7228
7229 fn moe_cached_gemm(
7230 e: &Engine,
7231 il: u16,
7232 proj: u8,
7233 ex: usize,
7234 m: &MoeWeights,
7235 max_block: usize,
7236 x: &cudarc::driver::CudaView<f32>,
7237 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7238 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7239 let exps = match proj {
7240 PROJ_GATE => &m.gate_exps,
7241 PROJ_UP => &m.up_exps,
7242 _ => &m.down_exps,
7243 };
7244 let layout = exps.expert_layout(ex);
7245 let id = BlockId::new(il, proj, ex as u16);
7246 let source = exps.expert_source(ex);
7247 e.with_moe_cache(max_block, |c, eng| {
7249 let slot = c.dispatch_source(id, source, eng)?;
7250 let DispatchSlot::Resident(sl) = slot;
7253 let buf = c.slot(sl);
7254 eng.qmatvec_view(
7255 buf,
7256 0..layout.len,
7257 x,
7258 1,
7259 exps.in_f,
7260 exps.out_f,
7261 layout.qtype,
7262 layout.row_bytes,
7263 )
7264 })
7265 }
7266
7267 fn moe_profile_admit_expert(
7271 e: &Engine,
7272 il: u16,
7273 ex: usize,
7274 m: &MoeWeights,
7275 max_block: usize,
7276 ) -> Result<(), Box<dyn std::error::Error>> {
7277 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7278 e.with_moe_cache(max_block, |cache, eng| {
7279 for (proj, exps) in [
7280 (PROJ_GATE, &m.gate_exps),
7281 (PROJ_UP, &m.up_exps),
7282 (PROJ_DOWN, &m.down_exps),
7283 ] {
7284 let id = BlockId::new(il, proj, ex as u16);
7285 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7286 }
7287 Ok(())
7288 })
7289 }
7290
7291 #[allow(clippy::too_many_arguments)]
7294 fn moe_frozen_gemm(
7295 e: &Engine,
7296 il: u16,
7297 proj: u8,
7298 ex: usize,
7299 m: &MoeWeights,
7300 max_block: usize,
7301 x: &cudarc::driver::CudaView<f32>,
7302 scratch: &mut Option<CudaSlice<u8>>,
7303 scratch_len: usize,
7304 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7305 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7306 let exps = match proj {
7307 PROJ_GATE => &m.gate_exps,
7308 PROJ_UP => &m.up_exps,
7309 _ => &m.down_exps,
7310 };
7311 let layout = exps.expert_layout(ex);
7312 let id = BlockId::new(il, proj, ex as u16);
7313 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7314 let Some(slot) = cache.resident(id) else {
7315 return Ok(None);
7316 };
7317 let buf = cache.slot(slot);
7318 Ok(Some(eng.qmatvec_view(
7319 buf,
7320 0..layout.len,
7321 x,
7322 1,
7323 exps.in_f,
7324 exps.out_f,
7325 layout.qtype,
7326 layout.row_bytes,
7327 )?))
7328 })? {
7329 return Ok(output);
7330 }
7331 if scratch.is_none() {
7332 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7333 }
7334 let scratch = scratch.as_mut().unwrap();
7335 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7336 e.qmatvec_view(
7337 scratch,
7338 0..layout.len,
7339 x,
7340 1,
7341 exps.in_f,
7342 exps.out_f,
7343 layout.qtype,
7344 layout.row_bytes,
7345 )
7346 }
7347
7348 fn moe_prefetch_expert(
7349 e: &Engine,
7350 il: u16,
7351 ex: usize,
7352 m: &MoeWeights,
7353 max_block: usize,
7354 keep: &[crate::moe_cache::BlockId],
7355 ) -> Result<(), Box<dyn std::error::Error>> {
7356 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7357 e.with_moe_cache(max_block, |c, eng| {
7358 for (proj, exps) in [
7359 (PROJ_GATE, &m.gate_exps),
7360 (PROJ_UP, &m.up_exps),
7361 (PROJ_DOWN, &m.down_exps),
7362 ] {
7363 let id = BlockId::new(il, proj, ex as u16);
7364 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7365 }
7366 Ok(())
7367 })
7368 }
7369
7370 fn moe_prefetch_disk_expert(
7373 e: &Engine,
7374 il: u16,
7375 ex: usize,
7376 m: &MoeWeights,
7377 max_block: usize,
7378 keep: &[crate::moe_cache::BlockId],
7379 ) -> Result<(), Box<dyn std::error::Error>> {
7380 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7381 e.with_moe_cache(max_block, |c, eng| {
7382 for (proj, exps) in [
7383 (PROJ_GATE, &m.gate_exps),
7384 (PROJ_UP, &m.up_exps),
7385 (PROJ_DOWN, &m.down_exps),
7386 ] {
7387 let source = exps.expert_source(ex);
7388 if let crate::model::ExpertSource::Disk { .. } = &source {
7389 let id = BlockId::new(il, proj, ex as u16);
7390 let _ = c.prefetch_source(id, source, keep, eng)?;
7391 }
7392 }
7393 Ok(())
7394 })
7395 }
7396
7397 #[inline]
7398 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7399 let _ = m.gate_exps.prefetch_expert_pages(ex);
7400 let _ = m.up_exps.prefetch_expert_pages(ex);
7401 let _ = m.down_exps.prefetch_expert_pages(ex);
7402 }
7403}
7404
7405impl HybridModel {
7422 #[allow(clippy::too_many_arguments)]
7426 fn moe_ffn_grouped_resident_q8(
7427 e: &Engine,
7428 m: &MoeWeights,
7429 z: &CudaSlice<f32>,
7430 t: usize,
7431 cfg: &ModelConfig,
7432 il: u16,
7433 sel_all: &[u32],
7434 w_all: &[f32],
7435 table: &CudaSlice<u64>,
7436 gu_il: bool,
7437 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7438 let moe = cfg.moe.as_ref().unwrap();
7439 let n_embd = cfg.n_embd as usize;
7440 let n_expert = moe.expert_count as usize;
7441 let n_used = moe.expert_used_count as usize;
7442 let n_ff_exp = moe.expert_ff_length as usize;
7443 let n_pairs = t * n_used;
7444 debug_assert_eq!(sel_all.len(), n_pairs);
7445 debug_assert_eq!(w_all.len(), n_pairs);
7446 debug_assert!(
7447 m.gate_exps.macros.is_none()
7448 && m.up_exps.macros.is_none()
7449 && m.down_exps.macros.is_none(),
7450 "resident grouped q8 does not fold per-expert macro scales",
7451 );
7452
7453 if !cfg.swiglu_clamped_at(il as u32) {
7459 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7460 let sel_d = e.htod_i32(&sel)?;
7461 let w_d = e.htod(w_all)?;
7462 let (gate_row_bytes, up_row_bytes) = if gu_il {
7463 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7464 (combined, combined)
7465 } else {
7466 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7467 };
7468 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7469 let act = e.moe_gate_up_silu8_dev_q8_rows(
7470 table,
7471 &sel_d,
7472 &zq,
7473 &zd,
7474 t,
7475 n_embd,
7476 n_ff_exp,
7477 n_used,
7478 n_expert,
7479 m.gate_exps.qtype,
7480 m.up_exps.qtype,
7481 gate_row_bytes,
7482 up_row_bytes,
7483 &m.dev_macros,
7484 )?;
7485 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7486 let mut moe_out = e.uninit(t * n_embd)?;
7487 e.moe_down8_fma_dev_q8_rows_g(
7488 table,
7489 &sel_d,
7490 &w_d,
7491 &aq2,
7492 &ad2,
7493 &mut moe_out,
7494 t,
7495 n_ff_exp,
7496 n_embd,
7497 n_used,
7498 n_expert,
7499 m.down_exps.qtype,
7500 m.down_exps.row_bytes,
7501 )?;
7502
7503 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7504 let mut counts = vec![0usize; n_expert];
7505 for &expert in sel_all {
7506 counts[expert as usize] += 1;
7507 }
7508 let mut sizes: Vec<usize> =
7509 counts.into_iter().filter(|&count| count != 0).collect();
7510 sizes.sort_unstable();
7511 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7512 println!(
7513 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7514 m_e: min={} median={} mean={mean:.1} max={}",
7515 sizes.len(),
7516 n_expert,
7517 sizes.first().copied().unwrap_or(0),
7518 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7519 sizes.last().copied().unwrap_or(0),
7520 );
7521 }
7522 return Ok(moe_out);
7523 }
7524
7525 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7529 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7530 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7531 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7532
7533 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7534 for (pair, &expert) in pair_ex.iter().enumerate() {
7535 by_expert[expert as usize].push(pair as i32);
7536 }
7537
7538 let pair_tok_d = e.htod_i32(&pair_tok)?;
7539 let pair_ex_d = e.htod_i32(&pair_ex)?;
7540 let pair_w_d = e.htod(w_all)?;
7541 let tok_off_d = e.htod_i32(&tok_off)?;
7542 let tok_ids_d = e.htod_i32(&tok_ids)?;
7543
7544 let matvec = |proj: i32,
7545 pair_rows: &CudaSlice<i32>,
7546 aq: &CudaSlice<i8>,
7547 ad: &CudaSlice<f32>,
7548 in_f: usize,
7549 out_f: usize,
7550 qtype: i32,
7551 row_bytes: usize|
7552 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7553 e.moe_pairs_matvec_q8(
7554 table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7555 row_bytes,
7556 )
7557 };
7558
7559 let (gate_row_bytes, up_row_bytes) = if gu_il {
7560 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7561 (combined, combined)
7562 } else {
7563 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7564 };
7565 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7566 let gate = matvec(
7567 0,
7568 &pair_tok_d,
7569 &zq,
7570 &zd,
7571 n_embd,
7572 n_ff_exp,
7573 m.gate_exps.qtype,
7574 gate_row_bytes,
7575 )?;
7576 let up = matvec(
7577 1,
7578 &pair_tok_d,
7579 &zq,
7580 &zd,
7581 n_embd,
7582 n_ff_exp,
7583 m.up_exps.qtype,
7584 up_row_bytes,
7585 )?;
7586 let mut act = e.uninit(n_pairs * n_ff_exp)?;
7587 Self::ffn_act_lim(
7588 e,
7589 cfg,
7590 &gate,
7591 &up,
7592 1.0,
7593 1.0,
7594 cfg.clamp_exp_at(il as u32),
7595 &mut act,
7596 n_pairs * n_ff_exp,
7597 )?;
7598 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7599 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7600 let pair_self_d = e.htod_i32(&pair_self)?;
7601 let down = matvec(
7602 2,
7603 &pair_self_d,
7604 &aq2,
7605 &ad2,
7606 n_ff_exp,
7607 n_embd,
7608 m.down_exps.qtype,
7609 m.down_exps.row_bytes,
7610 )?;
7611 let mut moe_out = e.uninit(t * n_embd)?;
7612 e.moe_pairs_scatter(
7613 &down,
7614 &pair_w_d,
7615 &tok_off_d,
7616 &tok_ids_d,
7617 &mut moe_out,
7618 t,
7619 n_embd,
7620 )?;
7621
7622 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7623 let mut sizes: Vec<usize> = by_expert
7624 .iter()
7625 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7626 .collect();
7627 sizes.sort_unstable();
7628 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7629 println!(
7630 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7631 m_e: min={} median={} mean={mean:.1} max={}",
7632 sizes.len(),
7633 n_expert,
7634 sizes.first().copied().unwrap_or(0),
7635 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7636 sizes.last().copied().unwrap_or(0),
7637 );
7638 }
7639 Ok(moe_out)
7640 }
7641
7642 fn moe_ffn_grouped_add_shared(
7643 e: &Engine,
7644 m: &MoeWeights,
7645 z: &CudaSlice<f32>,
7646 t: usize,
7647 cfg: &ModelConfig,
7648 il: u16,
7649 moe_out: &mut CudaSlice<f32>,
7650 ) -> Result<(), Box<dyn std::error::Error>> {
7651 let n_embd = cfg.n_embd as usize;
7652 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7653 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7654 {
7655 let n_ff_sh = gate_shexp.out_features();
7656 let sg_gate = e.matmul(gate_shexp, z, t)?;
7657 let sg_up = e.matmul(up_shexp, z, t)?;
7658 let mut sa = e.uninit(t * n_ff_sh)?;
7659 Self::ffn_act_lim(
7660 e,
7661 cfg,
7662 &sg_gate,
7663 &sg_up,
7664 1.0,
7665 1.0,
7666 cfg.clamp_shexp_at(il as u32),
7667 &mut sa,
7668 t * n_ff_sh,
7669 )?;
7670 let sh = e.matmul(down_shexp, &sa, t)?;
7671 let gate = match &m.gate_inp_shexp {
7672 Some(gate_inp_shexp) => {
7673 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7674 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7675 } else {
7676 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7677 let mut gate = e.uninit(t)?;
7678 e.sigmoid(&raw, &mut gate, t)?;
7679 gate
7680 }
7681 }
7682 None => e.htod(&vec![1.0f32; t])?,
7683 };
7684 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7685 }
7686 Ok(())
7687 }
7688
7689 pub(crate) fn moe_ffn_grouped(
7692 e: &Engine,
7693 m: &MoeWeights,
7694 z: &CudaSlice<f32>,
7695 t: usize,
7696 cfg: &ModelConfig,
7697 il: u16,
7698 max_block: usize,
7699 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7700 let moe = cfg.moe.as_ref().unwrap();
7701 let n_embd = cfg.n_embd as usize;
7702 let n_expert = moe.expert_count as usize;
7703 let n_used = moe.expert_used_count as usize;
7704 let n_ff_exp = moe.expert_ff_length as usize;
7705 let lim_exp = cfg.clamp_exp_at(il as u32);
7707
7708 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7712 if let Some(sig) = cfg.sigmoid_router() {
7713 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7714 }
7715 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7716 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7717 } else {
7718 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7719 };
7720 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7721 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7722 Self::trace_moe_input(e, il, t, n_embd, z)?;
7723
7724 let no_exp_macros = m.gate_exps.macros.is_none()
7729 && m.up_exps.macros.is_none()
7730 && m.down_exps.macros.is_none();
7731 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7732 m.has_uniform_expert_layout()
7733 && no_exp_macros
7734 && moe_q8_enabled()
7735 && q8_expert_supported(m.gate_exps.qtype)
7736 && q8_expert_supported(m.up_exps.qtype)
7737 && q8_expert_supported(m.down_exps.qtype)
7738 && moe_slab_enabled()
7739 && dev.dev == e.ctx().ordinal()
7740 });
7741 if let Some(dev) = resident_q8 {
7742 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7743 e,
7744 m,
7745 z,
7746 t,
7747 cfg,
7748 il,
7749 &sel_all,
7750 &w_all,
7751 &dev.ptr_row,
7752 dev.gu_il,
7753 )?;
7754 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7755 return Ok(moe_out);
7756 }
7757
7758 struct ExpertGroup {
7762 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
7766 let mut groups: Vec<ExpertGroup> = (0..n_expert)
7767 .map(|_| ExpertGroup {
7768 tok_indices: Vec::new(),
7769 slot_indices: Vec::new(),
7770 weights: Vec::new(),
7771 })
7772 .collect();
7773
7774 for tok in 0..t {
7775 for j in 0..n_used {
7776 let ex = sel_all[tok * n_used + j] as usize;
7777 let w = w_all[tok * n_used + j];
7778 groups[ex].tok_indices.push(tok as i32);
7779 groups[ex].slot_indices.push(j as i32);
7780 groups[ex].weights.push(w);
7781 }
7782 }
7783
7784 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7787 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
7791 let u_len = m.up_exps.max_expert_bytes();
7792 let d_len = m.down_exps.max_expert_bytes();
7793 let moe_q8 = m.has_uniform_expert_layout()
7794 && moe_q8_enabled()
7795 && q8_expert_supported(m.gate_exps.qtype)
7796 && q8_expert_supported(m.up_exps.qtype)
7797 && q8_expert_supported(m.down_exps.qtype);
7798 let slab_local = m
7801 .dev_exps
7802 .as_ref()
7803 .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
7804 let use_cache =
7805 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
7806 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
7809
7810 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
7812 (
7813 Some(e.alloc_u8(g_len)?),
7814 Some(e.alloc_u8(u_len)?),
7815 Some(e.alloc_u8(d_len)?),
7816 )
7817 } else {
7818 (None, None, None)
7819 };
7820
7821 let mut order: Vec<usize> = (0..n_expert)
7832 .filter(|&ex| !groups[ex].tok_indices.is_empty())
7833 .collect();
7834 order.sort_by(|&a, &b| {
7835 groups[b]
7836 .tok_indices
7837 .len()
7838 .cmp(&groups[a].tok_indices.len())
7839 .then(a.cmp(&b))
7840 });
7841 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
7843 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
7844 if worker_disk_prefetch {
7845 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
7846 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
7847 }
7848 }
7849 for (order_pos, &ex) in order.iter().enumerate() {
7850 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
7851 Self::moe_prefetch_host_expert(order[next], m);
7852 }
7853 if worker_disk_prefetch {
7854 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
7855 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7856 let keep = [
7857 BlockId::new(il, PROJ_GATE, ex as u16),
7858 BlockId::new(il, PROJ_UP, ex as u16),
7859 BlockId::new(il, PROJ_DOWN, ex as u16),
7860 ];
7861 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
7862 }
7863 }
7864 let grp = &groups[ex];
7865 let m_e = grp.tok_indices.len();
7866 m_dist.push(m_e);
7867 let gl = m.gate_exps.expert_layout(ex);
7868 let ul = m.up_exps.expert_layout(ex);
7869 let dl = m.down_exps.expert_layout(ex);
7870
7871 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
7875 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
7876 let dmac = m.down_exps.macro_scale(ex);
7877 let weight_d = if dmac == 1.0 {
7878 e.htod(&grp.weights)?
7879 } else {
7880 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
7881 e.htod(&scaled)?
7882 };
7883
7884 let mut gathered = e.zeros(m_e * n_embd)?;
7886 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
7887 let gv = gathered.slice(0..m_e * n_embd);
7888
7889 let y = if let Some(dev) = slab_local {
7892 let gate_start = ex * m.gate_exps.expert_stride;
7893 let up_start = ex * m.up_exps.expert_stride;
7894 let down_start = ex * m.down_exps.expert_stride;
7895 if grouped_q8 {
7896 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7897 let gate = e.qmatvec_expert_q8(
7898 &dev.gate,
7899 gate_start..gate_start + gl.len,
7900 &zq,
7901 &zd,
7902 m_e,
7903 m.gate_exps.in_f,
7904 m.gate_exps.out_f,
7905 gl.qtype,
7906 gl.row_bytes,
7907 )?;
7908 let up = e.qmatvec_expert_q8(
7909 &dev.up,
7910 up_start..up_start + ul.len,
7911 &zq,
7912 &zd,
7913 m_e,
7914 m.up_exps.in_f,
7915 m.up_exps.out_f,
7916 ul.qtype,
7917 ul.row_bytes,
7918 )?;
7919 let mut act = e.uninit(m_e * n_ff_exp)?;
7920 Self::ffn_act_lim(
7921 e,
7922 cfg,
7923 &gate,
7924 &up,
7925 m.gate_exps.macro_scale(ex),
7926 m.up_exps.macro_scale(ex),
7927 lim_exp,
7928 &mut act,
7929 m_e * n_ff_exp,
7930 )?;
7931 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
7932 e.qmatvec_expert_q8(
7933 &dev.down,
7934 down_start..down_start + dl.len,
7935 &aq2,
7936 &ad2,
7937 m_e,
7938 m.down_exps.in_f,
7939 m.down_exps.out_f,
7940 dl.qtype,
7941 dl.row_bytes,
7942 )?
7943 } else {
7944 let gate = e.qmatvec_view(
7945 &dev.gate,
7946 gate_start..gate_start + gl.len,
7947 &gv,
7948 m_e,
7949 m.gate_exps.in_f,
7950 m.gate_exps.out_f,
7951 gl.qtype,
7952 gl.row_bytes,
7953 )?;
7954 let up = e.qmatvec_view(
7955 &dev.up,
7956 up_start..up_start + ul.len,
7957 &gv,
7958 m_e,
7959 m.up_exps.in_f,
7960 m.up_exps.out_f,
7961 ul.qtype,
7962 ul.row_bytes,
7963 )?;
7964 let mut act = e.uninit(m_e * n_ff_exp)?;
7965 Self::ffn_act_lim(
7966 e,
7967 cfg,
7968 &gate,
7969 &up,
7970 m.gate_exps.macro_scale(ex),
7971 m.up_exps.macro_scale(ex),
7972 lim_exp,
7973 &mut act,
7974 m_e * n_ff_exp,
7975 )?;
7976 let actv = act.slice(0..m_e * n_ff_exp);
7977 e.qmatvec_view(
7978 &dev.down,
7979 down_start..down_start + dl.len,
7980 &actv,
7981 m_e,
7982 m.down_exps.in_f,
7983 m.down_exps.out_f,
7984 dl.qtype,
7985 dl.row_bytes,
7986 )?
7987 }
7988 } else if use_cache {
7989 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7990 if grouped_q8 {
7991 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7992 let gate = e.with_moe_cache(max_block, |cache, eng| {
7993 let id = BlockId::new(il, PROJ_GATE, ex as u16);
7994 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
7995 eng.qmatvec_expert_q8(
7996 cache.buf(slot),
7997 0..gl.len,
7998 &zq,
7999 &zd,
8000 m_e,
8001 m.gate_exps.in_f,
8002 m.gate_exps.out_f,
8003 gl.qtype,
8004 gl.row_bytes,
8005 )
8006 })?;
8007 let up = e.with_moe_cache(max_block, |cache, eng| {
8008 let id = BlockId::new(il, PROJ_UP, ex as u16);
8009 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8010 eng.qmatvec_expert_q8(
8011 cache.buf(slot),
8012 0..ul.len,
8013 &zq,
8014 &zd,
8015 m_e,
8016 m.up_exps.in_f,
8017 m.up_exps.out_f,
8018 ul.qtype,
8019 ul.row_bytes,
8020 )
8021 })?;
8022 let mut act = e.uninit(m_e * n_ff_exp)?;
8023 Self::ffn_act_lim(
8024 e,
8025 cfg,
8026 &gate,
8027 &up,
8028 m.gate_exps.macro_scale(ex),
8029 m.up_exps.macro_scale(ex),
8030 lim_exp,
8031 &mut act,
8032 m_e * n_ff_exp,
8033 )?;
8034 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8035 e.with_moe_cache(max_block, |cache, eng| {
8036 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8037 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8038 eng.qmatvec_expert_q8(
8039 cache.buf(slot),
8040 0..dl.len,
8041 &aq2,
8042 &ad2,
8043 m_e,
8044 m.down_exps.in_f,
8045 m.down_exps.out_f,
8046 dl.qtype,
8047 dl.row_bytes,
8048 )
8049 })?
8050 } else {
8051 let gate = e.with_moe_cache(max_block, |cache, eng| {
8052 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8053 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8054 eng.qmatvec_view(
8055 cache.buf(slot),
8056 0..gl.len,
8057 &gv,
8058 m_e,
8059 m.gate_exps.in_f,
8060 m.gate_exps.out_f,
8061 gl.qtype,
8062 gl.row_bytes,
8063 )
8064 })?;
8065 let up = e.with_moe_cache(max_block, |cache, eng| {
8066 let id = BlockId::new(il, PROJ_UP, ex as u16);
8067 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8068 eng.qmatvec_view(
8069 cache.buf(slot),
8070 0..ul.len,
8071 &gv,
8072 m_e,
8073 m.up_exps.in_f,
8074 m.up_exps.out_f,
8075 ul.qtype,
8076 ul.row_bytes,
8077 )
8078 })?;
8079 let mut act = e.uninit(m_e * n_ff_exp)?;
8080 Self::ffn_act_lim(
8081 e,
8082 cfg,
8083 &gate,
8084 &up,
8085 m.gate_exps.macro_scale(ex),
8086 m.up_exps.macro_scale(ex),
8087 lim_exp,
8088 &mut act,
8089 m_e * n_ff_exp,
8090 )?;
8091 let actv = act.slice(0..m_e * n_ff_exp);
8092 e.with_moe_cache(max_block, |cache, eng| {
8093 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8094 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8095 eng.qmatvec_view(
8096 cache.buf(slot),
8097 0..dl.len,
8098 &actv,
8099 m_e,
8100 m.down_exps.in_f,
8101 m.down_exps.out_f,
8102 dl.qtype,
8103 dl.row_bytes,
8104 )
8105 })?
8106 }
8107 } else {
8108 let sg = scratch_g.as_mut().unwrap();
8109 let su = scratch_u.as_mut().unwrap();
8110 let sd = scratch_d.as_mut().unwrap();
8111 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8112 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8113 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8114 if grouped_q8 {
8115 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8116 let gate = e.qmatvec_expert_q8(
8117 sg,
8118 0..gl.len,
8119 &zq,
8120 &zd,
8121 m_e,
8122 m.gate_exps.in_f,
8123 m.gate_exps.out_f,
8124 gl.qtype,
8125 gl.row_bytes,
8126 )?;
8127 let up = e.qmatvec_expert_q8(
8128 su,
8129 0..ul.len,
8130 &zq,
8131 &zd,
8132 m_e,
8133 m.up_exps.in_f,
8134 m.up_exps.out_f,
8135 ul.qtype,
8136 ul.row_bytes,
8137 )?;
8138 let mut act = e.uninit(m_e * n_ff_exp)?;
8139 Self::ffn_act_lim(
8140 e,
8141 cfg,
8142 &gate,
8143 &up,
8144 m.gate_exps.macro_scale(ex),
8145 m.up_exps.macro_scale(ex),
8146 lim_exp,
8147 &mut act,
8148 m_e * n_ff_exp,
8149 )?;
8150 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8151 e.qmatvec_expert_q8(
8152 sd,
8153 0..dl.len,
8154 &aq2,
8155 &ad2,
8156 m_e,
8157 m.down_exps.in_f,
8158 m.down_exps.out_f,
8159 dl.qtype,
8160 dl.row_bytes,
8161 )?
8162 } else {
8163 let gate = e.qmatvec_view(
8164 sg,
8165 0..gl.len,
8166 &gv,
8167 m_e,
8168 m.gate_exps.in_f,
8169 m.gate_exps.out_f,
8170 gl.qtype,
8171 gl.row_bytes,
8172 )?;
8173 let up = e.qmatvec_view(
8174 su,
8175 0..ul.len,
8176 &gv,
8177 m_e,
8178 m.up_exps.in_f,
8179 m.up_exps.out_f,
8180 ul.qtype,
8181 ul.row_bytes,
8182 )?;
8183 let mut act = e.uninit(m_e * n_ff_exp)?;
8184 Self::ffn_act_lim(
8185 e,
8186 cfg,
8187 &gate,
8188 &up,
8189 m.gate_exps.macro_scale(ex),
8190 m.up_exps.macro_scale(ex),
8191 lim_exp,
8192 &mut act,
8193 m_e * n_ff_exp,
8194 )?;
8195 let actv = act.slice(0..m_e * n_ff_exp);
8196 e.qmatvec_view(
8197 sd,
8198 0..dl.len,
8199 &actv,
8200 m_e,
8201 m.down_exps.in_f,
8202 m.down_exps.out_f,
8203 dl.qtype,
8204 dl.row_bytes,
8205 )?
8206 }
8207 };
8208
8209 e.scatter_slot(
8211 &y,
8212 &tok_idx_d,
8213 &slot_idx_d,
8214 &weight_d,
8215 &mut slot_buf,
8216 &mut wbuf,
8217 n_embd,
8218 n_used,
8219 m_e,
8220 )?;
8221 }
8222
8223 let mut moe_out = e.zeros(t * n_embd)?;
8225 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8226
8227 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8229 m_dist.sort_unstable();
8230 let active = m_dist.len();
8231 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8232 let median = m_dist[active / 2];
8233 let max_m = *m_dist.last().unwrap();
8234 let min_m = m_dist[0];
8235 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8236 println!(
8237 "moe-grouped il={il} t={t} active={active}/{n_expert} \
8238 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8239 above_gemm_threshold(>=16)={above16}/{active}"
8240 );
8241 }
8242
8243 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8244 Ok(moe_out)
8245 }
8246
8247 pub(crate) fn moe_ffn_lockstep(
8254 &self,
8255 e: &Engine,
8256 m: &MoeWeights,
8257 zbatch: &CudaSlice<f32>,
8258 mrows: usize,
8259 il: u16,
8260 max_block: usize,
8261 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8262 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8263 let cfg = &self.cfg;
8264 let moe = cfg.moe.as_ref().unwrap();
8265 let n_embd = cfg.n_embd as usize;
8266 let n_expert = moe.expert_count as usize;
8267 let n_used = moe.expert_used_count as usize;
8268 let n_ff_exp = moe.expert_ff_length as usize;
8269 let lim_exp = cfg.clamp_exp_at(il as u32);
8271 let lim_shexp = cfg.clamp_shexp_at(il as u32);
8272
8273 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8274 if let Some(sig) = cfg.sigmoid_router() {
8275 Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8276 }
8277 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8278 Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8279 } else {
8280 Self::moe_route_cfg(
8281 e,
8282 &logits,
8283 mrows,
8284 n_expert,
8285 n_used,
8286 m.active_experts.as_deref(),
8287 )?
8288 };
8289 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8290
8291 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8293 Ok((0..n_expert)
8294 .map(|ex| {
8295 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8296 .into_iter()
8297 .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8298 })
8299 .collect())
8300 })?;
8301
8302 struct Group {
8303 rows: Vec<i32>,
8304 slots: Vec<i32>,
8305 weights: Vec<f32>,
8306 }
8307 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8308 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8309 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8310 Default::default();
8311 for row in 0..mrows {
8312 for j in 0..n_used {
8313 let ex = sel_all[row * n_used + j] as usize;
8314 let w = w_all[row * n_used + j];
8315 if resident_expert[ex] {
8316 let group = groups.entry(ex).or_insert_with(|| Group {
8317 rows: Vec::new(),
8318 slots: Vec::new(),
8319 weights: Vec::new(),
8320 });
8321 group.rows.push(row as i32);
8322 group.slots.push(j as i32);
8323 group.weights.push(w);
8324 } else {
8325 crate::cpu_experts::record_incomplete_gpu_residency(0);
8326 cpu_rows[row].push((ex, w));
8327 cpu_by_expert.entry(ex).or_default().push((row, w));
8328 }
8329 }
8330 }
8331
8332 let host_rows = e.dtoh(zbatch)?;
8338 let rows_ok = crate::cpu_experts::rows_supported();
8339 enum CpuPart {
8340 Single { row: usize },
8341 Rows { rows: Vec<usize> },
8342 }
8343 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8344 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8345 if rows_ok {
8346 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8347 .into_iter()
8348 .filter(|(_, rows)| rows.len() >= 2)
8349 .collect();
8350 shared.sort_by_key(|(ex, _)| *ex);
8351 for (ex, mut row_weights) in shared {
8352 row_weights.sort_by_key(|(row, _)| *row);
8353 let inputs: Vec<(&[f32], f32)> = row_weights
8354 .iter()
8355 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8356 .collect();
8357 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8358 .map_err(std::io::Error::other)?;
8359 for &(row, _) in &row_weights {
8360 rows_served.insert((row, ex));
8361 }
8362 tickets.push((
8363 CpuPart::Rows {
8364 rows: row_weights.iter().map(|&(row, _)| row).collect(),
8365 },
8366 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8367 ));
8368 }
8369 }
8370 for (row, selected) in cpu_rows.iter().enumerate() {
8371 let leftover: Vec<(usize, f32)> = selected
8372 .iter()
8373 .copied()
8374 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8375 .collect();
8376 if leftover.is_empty() {
8377 continue;
8378 }
8379 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8380 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8381 .map_err(std::io::Error::other)?;
8382 tickets.push((
8383 CpuPart::Single { row },
8384 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8385 ));
8386 }
8387
8388 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8389 let mut wbuf = e.zeros(mrows * n_used)?;
8390 let mut order: Vec<usize> = groups.keys().copied().collect();
8391 order.sort_by(|&a, &b| {
8392 groups[&b]
8393 .rows
8394 .len()
8395 .cmp(&groups[&a].rows.len())
8396 .then(a.cmp(&b))
8397 });
8398 for &ex in &order {
8399 let group = &groups[&ex];
8400 let m_e = group.rows.len();
8401 let gl = m.gate_exps.expert_layout(ex);
8402 let ul = m.up_exps.expert_layout(ex);
8403 let dl = m.down_exps.expert_layout(ex);
8404 let row_idx_d = e.htod_i32(&group.rows)?;
8405 let slot_idx_d = e.htod_i32(&group.slots)?;
8406 let dmac = m.down_exps.macro_scale(ex);
8407 let weight_d = if dmac == 1.0 {
8408 e.htod(&group.weights)?
8409 } else {
8410 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8411 e.htod(&scaled)?
8412 };
8413 let mut gathered = e.zeros(m_e * n_embd)?;
8414 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8415 let gv = gathered.slice(0..m_e * n_embd);
8416 let gate = e.with_moe_cache(max_block, |c, eng| {
8417 let slot = c
8418 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8419 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8420 eng.qmatvec_view(
8421 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8422 0..gl.len,
8423 &gv,
8424 m_e,
8425 m.gate_exps.in_f,
8426 m.gate_exps.out_f,
8427 gl.qtype,
8428 gl.row_bytes,
8429 )
8430 })?;
8431 let up = e.with_moe_cache(max_block, |c, eng| {
8432 let slot = c
8433 .resident(BlockId::new(il, PROJ_UP, ex as u16))
8434 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8435 eng.qmatvec_view(
8436 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8437 0..ul.len,
8438 &gv,
8439 m_e,
8440 m.up_exps.in_f,
8441 m.up_exps.out_f,
8442 ul.qtype,
8443 ul.row_bytes,
8444 )
8445 })?;
8446 let mut act = e.zeros(m_e * n_ff_exp)?;
8447 Self::ffn_act_lim(
8448 e,
8449 cfg,
8450 &gate,
8451 &up,
8452 m.gate_exps.macro_scale(ex),
8453 m.up_exps.macro_scale(ex),
8454 lim_exp,
8455 &mut act,
8456 m_e * n_ff_exp,
8457 )?;
8458 let actv = act.slice(0..m_e * n_ff_exp);
8459 let y = e.with_moe_cache(max_block, |c, eng| {
8460 let slot = c
8461 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8462 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8463 eng.qmatvec_view(
8464 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8465 0..dl.len,
8466 &actv,
8467 m_e,
8468 m.down_exps.in_f,
8469 m.down_exps.out_f,
8470 dl.qtype,
8471 dl.row_bytes,
8472 )
8473 })?;
8474 e.scatter_slot(
8475 &y,
8476 &row_idx_d,
8477 &slot_idx_d,
8478 &weight_d,
8479 &mut slot_buf,
8480 &mut wbuf,
8481 n_embd,
8482 n_used,
8483 m_e,
8484 )?;
8485 }
8486 let mut moe_out = e.zeros(mrows * n_embd)?;
8487 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8488
8489 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8491 for (part, ticket) in tickets {
8492 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8493 let mut add_row = |row: usize, chunk: &[f32]| {
8494 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8495 for (accumulator, value) in sum.iter_mut().zip(chunk) {
8496 *accumulator += value;
8497 }
8498 };
8499 match part {
8500 CpuPart::Single { row } => add_row(row, &cpu_output),
8501 CpuPart::Rows { rows } => {
8502 for (slot, row) in rows.into_iter().enumerate() {
8503 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8504 }
8505 }
8506 }
8507 }
8508 for (row, sum) in row_sums.into_iter().enumerate() {
8509 let Some(sum) = sum else { continue };
8510 let cpu_output = e.htod(&sum)?;
8511 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8512 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8513 }
8514
8515 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8516 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8517 {
8518 let n_ff_sh = gate_shexp.out_features();
8519 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8520 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8521 let mut sa = e.zeros(mrows * n_ff_sh)?;
8522 Self::ffn_act_lim(
8523 e,
8524 cfg,
8525 &sg_gate,
8526 &sg_up,
8527 1.0,
8528 1.0,
8529 lim_shexp,
8530 &mut sa,
8531 mrows * n_ff_sh,
8532 )?;
8533 let sh = e.matmul(down_shexp, &sa, mrows)?;
8534 let g = match &m.gate_inp_shexp {
8537 Some(gate_inp_shexp) => {
8538 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8539 }
8540 None => e.htod(&vec![1.0f32; mrows])?,
8541 };
8542 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8543 }
8544
8545 Ok(moe_out)
8546 }
8547}
8548
8549impl HybridModel {
8555 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8557 let g = self.cfg.gemma4.as_ref().unwrap();
8558 let swa = g.swa_pattern[il];
8559 let hd = if swa {
8560 g.key_length_swa
8561 } else {
8562 g.key_length_global
8563 } as usize;
8564 (
8568 hd,
8569 g.head_count_kv[il] as usize,
8570 self.cfg.n_head as usize,
8571 if swa {
8572 g.rope_base_swa
8573 } else {
8574 g.rope_base_global
8575 },
8576 1.0,
8577 swa,
8578 )
8579 }
8580
8581 fn gemma4_suppress(
8585 &self,
8586 e: &Engine,
8587 ld: &mut CudaSlice<f32>,
8588 t: usize,
8589 ) -> Result<(), Box<dyn std::error::Error>> {
8590 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8591 #[cfg(debug_assertions)]
8596 crate::debug_assert_tensor_stream_device(
8597 ids,
8598 &e.stream(),
8599 "gemma4_suppress.suppress_d",
8600 );
8601 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8602 }
8603 Ok(())
8604 }
8605
8606 fn gemma4_attn_prime(
8611 &self,
8612 e: &Engine,
8613 fa: &crate::hybrid::FullAttnLayer,
8614 il: usize,
8615 h: &CudaSlice<f32>,
8616 pos_d: &CudaSlice<i32>,
8617 t: usize,
8618 cache: Option<&mut Cache>,
8619 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8620 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8621 let eps = self.cfg.rms_eps;
8622 let aux = self.gemma4_aux.as_ref().unwrap();
8623 let ones = aux.ones(e);
8624 #[cfg(debug_assertions)]
8625 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8626
8627 e.mmq_act_begin();
8630 let q0 = e.matmul(&fa.wq, h, t)?; let k0 = e.matmul(&fa.wk, h, t)?; let v0 = if swa {
8635 e.matmul(&fa.wv, h, t)?
8636 } else {
8637 e.clone_dtod(&k0)?
8638 };
8639
8640 let mut q = e.uninit(t * nh * hd)?;
8641 let mut k = e.uninit(t * nkv * hd)?;
8642 let mut v = e.uninit(t * nkv * hd)?;
8644 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8648 let emit = t >= 16
8649 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8650 && *EMIT.get_or_init(|| {
8651 std::env::var("MEMRA_FA_EMIT")
8652 .map(|s| s != "0")
8653 .unwrap_or(true)
8654 });
8655 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8656 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8657 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8658 let v_f16 = emit
8661 && crate::fa_f16pv_on()
8662 && match hd {
8663 512 => true,
8664 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8665 _ => false,
8666 };
8667 if emit {
8668 e.rms_norm_qkv_w4b(
8669 &q0,
8670 &k0,
8671 &v0,
8672 fa.q_norm.float_data(),
8673 fa.k_norm.float_data(),
8674 ones,
8675 &mut q,
8676 &mut k,
8677 &mut v,
8678 &mut vb,
8679 hd,
8680 nh * t,
8681 nkv * t,
8682 eps,
8683 v_f16,
8684 )?;
8685 } else {
8686 e.rms_norm_qkv(
8687 &q0,
8688 &k0,
8689 &v0,
8690 fa.q_norm.float_data(),
8691 fa.k_norm.float_data(),
8692 ones,
8693 &mut q,
8694 &mut k,
8695 &mut v,
8696 hd,
8697 nh * t,
8698 nkv * t,
8699 eps,
8700 )?;
8701 }
8702
8703 let ff = if swa {
8704 None
8705 } else {
8706 Some(
8707 aux.rope_freqs(e)
8708 .expect("gemma4 global rope needs rope_freqs.weight"),
8709 )
8710 };
8711 #[cfg(debug_assertions)]
8712 if let Some(ff) = ff {
8713 crate::debug_assert_tensor_stream_device(
8714 ff,
8715 &e.stream(),
8716 "gemma4_attn_prime.rope_freqs",
8717 );
8718 }
8719 if emit {
8720 e.rope_neox2_bf16e(
8721 &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8722 )?;
8723 } else {
8724 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8725 }
8726
8727 if let Some(cache) = cache {
8728 let kvl = cache.kv[il].as_mut().unwrap();
8729 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8730 e.append_kv_quantized_rows(
8731 &k,
8732 &v,
8733 &mut kvl.k,
8734 &mut kvl.v,
8735 kvl.len,
8736 t,
8737 kvl.kv_dim_k,
8738 kvl.kv_dim_v,
8739 kvl.k_tok_bytes,
8740 kvl.v_tok_bytes,
8741 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
8742 )?;
8743 kvl.len += t;
8744 }
8745 let mut attn = e.zeros(t * nh * hd)?;
8746 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8750 if swa && t > win {
8751 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8752 if emit {
8753 e.fa_prefill_w_pre(
8754 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
8755 )?;
8756 } else {
8757 e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8758 }
8759 } else {
8760 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8761 }
8762 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8763 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8764 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
8765 if emit {
8766 e.fa_prefill_hd512_pre(
8767 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
8768 )?;
8769 } else {
8770 e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8771 }
8772 } else {
8773 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8774 }
8775 Ok(e.matmul(&fa.wo, &attn, t)?)
8776 }
8777
8778 fn gemma4_attn(
8780 &self,
8781 e: &Engine,
8782 fa: &crate::hybrid::FullAttnLayer,
8783 il: usize,
8784 h: &CudaSlice<f32>,
8785 pos_d: &CudaSlice<i32>,
8786 t: usize,
8787 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8788 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
8789 }
8790
8791 fn gemma4_moe_q8(
8796 &self,
8797 e: &Engine,
8798 m: &crate::hybrid::MoeWeights,
8799 bits: &crate::hybrid::Gemma4MoeBits,
8800 mq: &(CudaSlice<i8>, CudaSlice<f32>),
8801 router_in: &CudaSlice<f32>,
8802 t: usize,
8803 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8804 let cfg = &self.cfg;
8805 let moe = cfg.moe.as_ref().unwrap();
8806 let n_embd = cfg.n_embd as usize;
8807 let n_expert = moe.expert_count as usize;
8808 let n_used = moe.expert_used_count as usize;
8809 let n_ff_exp = moe.expert_ff_length as usize;
8810 let logits = if crate::router_kernel_on() {
8814 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8815 } else {
8816 e.matmul(&m.gate_inp, router_in, t)?
8817 };
8818 let dev = m.dev_exps.as_ref().unwrap();
8819 let (sel_d, w_d) =
8820 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8821 let (zq, zd) = mq;
8822 if t == 1 {
8823 let selv = sel_d.slice(0..n_used);
8824 let wv = w_d.slice(0..n_used);
8825 let act = e.moe_gate_up_gelu8_dev_q8(
8826 &dev.ptr_row,
8827 &selv,
8828 zq,
8829 zd,
8830 n_embd,
8831 n_ff_exp,
8832 n_used,
8833 n_expert,
8834 m.gate_exps.qtype,
8835 m.up_exps.qtype,
8836 m.gate_exps.row_bytes,
8837 m.up_exps.row_bytes,
8838 )?;
8839 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8840 let mut moe_out = e.uninit(n_embd)?;
8841 e.moe_down8_fma_dev_q8(
8842 &dev.ptr_row,
8843 &selv,
8844 &wv,
8845 &aq2,
8846 &ad2,
8847 &mut moe_out.slice_mut(0..n_embd),
8848 n_ff_exp,
8849 n_embd,
8850 n_used,
8851 n_expert,
8852 m.down_exps.qtype,
8853 m.down_exps.row_bytes,
8854 )?;
8855 return Ok(moe_out);
8856 }
8857 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
8858 let act = if csr {
8859 e.moe_gate_up_gelu8_dev_q8_csr(
8860 &dev.ptr_row,
8861 &sel_d,
8862 zq,
8863 zd,
8864 t * n_used,
8865 n_embd,
8866 n_ff_exp,
8867 n_used,
8868 n_expert,
8869 m.gate_exps.qtype,
8870 m.up_exps.qtype,
8871 m.gate_exps.row_bytes,
8872 m.up_exps.row_bytes,
8873 )?
8874 } else {
8875 e.moe_gate_up_gelu8_dev_q8_rows(
8876 &dev.ptr_row,
8877 &sel_d,
8878 zq,
8879 zd,
8880 t,
8881 n_embd,
8882 n_ff_exp,
8883 n_used,
8884 n_expert,
8885 m.gate_exps.qtype,
8886 m.up_exps.qtype,
8887 m.gate_exps.row_bytes,
8888 m.up_exps.row_bytes,
8889 )?
8890 };
8891 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
8892 let mut moe_out = e.uninit(t * n_embd)?;
8893 e.moe_down8_fma_dev_q8_rows_g(
8896 &dev.ptr_row,
8897 &sel_d,
8898 &w_d,
8899 &aq2,
8900 &ad2,
8901 &mut moe_out,
8902 t,
8903 n_ff_exp,
8904 n_embd,
8905 n_used,
8906 n_expert,
8907 m.down_exps.qtype,
8908 m.down_exps.row_bytes,
8909 )?;
8910 Ok(moe_out)
8911 }
8912
8913 fn gemma4_moe(
8917 &self,
8918 e: &Engine,
8919 m: &crate::hybrid::MoeWeights,
8920 bits: &crate::hybrid::Gemma4MoeBits,
8921 moe_in: &CudaSlice<f32>,
8922 router_in: &CudaSlice<f32>,
8923 t: usize,
8924 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8925 let cfg = &self.cfg;
8926 let moe = cfg.moe.as_ref().unwrap();
8927 let n_embd = cfg.n_embd as usize;
8928 let n_expert = moe.expert_count as usize;
8929 let n_used = moe.expert_used_count as usize;
8930 let n_ff_exp = moe.expert_ff_length as usize;
8931
8932 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
8936 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8937 } else {
8938 e.matmul(&m.gate_inp, router_in, t)?
8939 };
8940
8941 if t < PRIME_MIN_T
8946 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
8947 && expert_dp4a_supported(m.gate_exps.qtype)
8948 && expert_dp4a_supported(m.up_exps.qtype)
8949 && expert_dp4a_supported(m.down_exps.qtype)
8950 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
8951 {
8952 let dev = m.dev_exps.as_ref().unwrap();
8953 let (sel_d, w_d) =
8954 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8955 if t == 1 {
8956 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
8957 let selv = sel_d.slice(0..n_used);
8958 let wv = w_d.slice(0..n_used);
8959 let act = e.moe_gate_up_gelu8_dev_q8(
8960 &dev.ptr_row,
8961 &selv,
8962 &zq,
8963 &zd,
8964 n_embd,
8965 n_ff_exp,
8966 n_used,
8967 n_expert,
8968 m.gate_exps.qtype,
8969 m.up_exps.qtype,
8970 m.gate_exps.row_bytes,
8971 m.up_exps.row_bytes,
8972 )?;
8973 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8974 let mut moe_out = e.uninit(n_embd)?;
8975 e.moe_down8_fma_dev_q8(
8976 &dev.ptr_row,
8977 &selv,
8978 &wv,
8979 &aq2,
8980 &ad2,
8981 &mut moe_out.slice_mut(0..n_embd),
8982 n_ff_exp,
8983 n_embd,
8984 n_used,
8985 n_expert,
8986 m.down_exps.qtype,
8987 m.down_exps.row_bytes,
8988 )?;
8989 return Ok(moe_out);
8990 }
8991 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
8996 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
8997 let act = if csr {
8998 e.moe_gate_up_gelu8_dev_q8_csr(
8999 &dev.ptr_row,
9000 &sel_d,
9001 &zq,
9002 &zd,
9003 t * n_used,
9004 n_embd,
9005 n_ff_exp,
9006 n_used,
9007 n_expert,
9008 m.gate_exps.qtype,
9009 m.up_exps.qtype,
9010 m.gate_exps.row_bytes,
9011 m.up_exps.row_bytes,
9012 )?
9013 } else {
9014 e.moe_gate_up_gelu8_dev_q8_rows(
9015 &dev.ptr_row,
9016 &sel_d,
9017 &zq,
9018 &zd,
9019 t,
9020 n_embd,
9021 n_ff_exp,
9022 n_used,
9023 n_expert,
9024 m.gate_exps.qtype,
9025 m.up_exps.qtype,
9026 m.gate_exps.row_bytes,
9027 m.up_exps.row_bytes,
9028 )?
9029 };
9030 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9031 let mut moe_out = e.uninit(t * n_embd)?;
9032 e.moe_down8_fma_dev_q8_rows_g(
9033 &dev.ptr_row,
9034 &sel_d,
9035 &w_d,
9036 &aq2,
9037 &ad2,
9038 &mut moe_out,
9039 t,
9040 n_ff_exp,
9041 n_embd,
9042 n_used,
9043 n_expert,
9044 m.down_exps.qtype,
9045 m.down_exps.row_bytes,
9046 )?;
9047 return Ok(moe_out);
9048 }
9049
9050 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9051 for (i, &sx) in sel_all.iter().enumerate() {
9052 w_all[i] *= bits.per_expert_scale[sx as usize];
9053 }
9054
9055 if t >= PRIME_MIN_T
9059 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9060 && expert_dp4a_supported(m.gate_exps.qtype)
9061 && expert_dp4a_supported(m.up_exps.qtype)
9062 && expert_dp4a_supported(m.down_exps.qtype)
9063 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9064 {
9065 let dev = m.dev_exps.as_ref().unwrap();
9066 let n_pairs = t * n_used;
9067 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9068 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9069 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9070 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9071 let pt = e.htod_i32(&pair_tok)?;
9072 let pw = e.htod(&w_all)?;
9073 let toff = e.htod_i32(&tok_off)?;
9074 let tids = e.htod_i32(&tok_ids)?;
9075 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9076 for p in 0..n_pairs {
9077 by_ex[pair_ex[p] as usize].push(p as i32);
9078 }
9079 let mut ex_ids: Vec<i32> = Vec::new();
9080 let mut ex_off: Vec<i32> = vec![0];
9081 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9082 for (ex, list) in by_ex.iter().enumerate() {
9083 if list.is_empty() {
9084 continue;
9085 }
9086 ex_ids.push(ex as i32);
9087 ex_pairs.extend_from_slice(list);
9088 ex_off.push(ex_pairs.len() as i32);
9089 }
9090 let n_active = ex_ids.len();
9091 let exi = e.htod_i32(&ex_ids)?;
9092 let exo = e.htod_i32(&ex_off)?;
9093 let exp_d = e.htod_i32(&ex_pairs)?;
9094 if crate::moe_f16g_gemma_on()
9102 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9103 && f16g_proj_ok(m.up_exps.qtype, n_embd)
9104 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9105 {
9106 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9107 let csr_tok_d = e.htod_i32(&csr_tok)?;
9108 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9109 let g_csr = e.moe_f16_grouped(
9110 &dev.ptr_row,
9111 0,
9112 n_expert,
9113 &exi,
9114 &ex_off,
9115 &exo,
9116 &z_f16,
9117 &z_s,
9118 n_embd,
9119 n_ff_exp,
9120 n_active,
9121 n_pairs,
9122 m.gate_exps.qtype,
9123 m.gate_exps.row_bytes,
9124 )?;
9125 let u_csr = e.moe_f16_grouped(
9126 &dev.ptr_row,
9127 1,
9128 n_expert,
9129 &exi,
9130 &ex_off,
9131 &exo,
9132 &z_f16,
9133 &z_s,
9134 n_embd,
9135 n_ff_exp,
9136 n_active,
9137 n_pairs,
9138 m.up_exps.qtype,
9139 m.up_exps.row_bytes,
9140 )?;
9141 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9142 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9143 let d_csr = e.moe_f16_grouped(
9144 &dev.ptr_row,
9145 2,
9146 n_expert,
9147 &exi,
9148 &ex_off,
9149 &exo,
9150 &a_f16,
9151 &a_s,
9152 n_ff_exp,
9153 n_embd,
9154 n_active,
9155 n_pairs,
9156 m.down_exps.qtype,
9157 m.down_exps.row_bytes,
9158 )?;
9159 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9160 let mut moe_out = e.uninit(t * n_embd)?;
9161 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9162 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9163 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9164 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9165 eprintln!(
9166 "[f16g-debug] post-permute bad={} post-scatter bad={}",
9167 scan(&yd),
9168 scan(&mo)
9169 );
9170 }
9171 return Ok(moe_out);
9172 }
9173 let mma =
9176 n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9177 let (gate, up) = if mma {
9178 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9179 (
9180 e.mmq_iq_experts(
9181 &dev.ptr_row,
9182 0,
9183 n_expert,
9184 &exi,
9185 &exo,
9186 &exp_d,
9187 &pt,
9188 &z_scr,
9189 n_embd,
9190 n_ff_exp,
9191 n_active,
9192 n_pairs,
9193 t,
9194 m.gate_exps.qtype,
9195 m.gate_exps.row_bytes,
9196 )?,
9197 e.mmq_iq_experts(
9198 &dev.ptr_row,
9199 1,
9200 n_expert,
9201 &exi,
9202 &exo,
9203 &exp_d,
9204 &pt,
9205 &z_scr,
9206 n_embd,
9207 n_ff_exp,
9208 n_active,
9209 n_pairs,
9210 t,
9211 m.up_exps.qtype,
9212 m.up_exps.row_bytes,
9213 )?,
9214 )
9215 } else {
9216 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9217 (
9218 e.moe_pairs_matvec_q8_dec(
9219 &dev.ptr_row,
9220 0,
9221 &exi,
9222 &exo,
9223 &exp_d,
9224 &pt,
9225 &zq,
9226 &zd,
9227 n_embd,
9228 n_ff_exp,
9229 n_expert,
9230 n_active,
9231 n_pairs,
9232 m.gate_exps.qtype,
9233 m.gate_exps.row_bytes,
9234 )?,
9235 e.moe_pairs_matvec_q8_dec(
9236 &dev.ptr_row,
9237 1,
9238 &exi,
9239 &exo,
9240 &exp_d,
9241 &pt,
9242 &zq,
9243 &zd,
9244 n_embd,
9245 n_ff_exp,
9246 n_expert,
9247 n_active,
9248 n_pairs,
9249 m.up_exps.qtype,
9250 m.up_exps.row_bytes,
9251 )?,
9252 )
9253 };
9254 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9255 let pself = e.htod_i32(&pair_self)?;
9256 let y_down = if mma {
9268 let in_pad = n_ff_exp.div_ceil(256) * 256;
9269 let a_scr = if crate::moe_fuse_actq_on() {
9270 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9271 } else {
9272 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9273 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9274 };
9275 e.mmq_iq_experts(
9276 &dev.ptr_row,
9277 2,
9278 n_expert,
9279 &exi,
9280 &exo,
9281 &exp_d,
9282 &pself,
9283 &a_scr,
9284 in_pad,
9285 n_embd,
9286 n_active,
9287 n_pairs,
9288 n_pairs,
9289 m.down_exps.qtype,
9290 m.down_exps.row_bytes,
9291 )?
9292 } else {
9293 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9294 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9295 e.moe_pairs_matvec_q8_dec(
9296 &dev.ptr_row,
9297 2,
9298 &exi,
9299 &exo,
9300 &exp_d,
9301 &pself,
9302 &aq2,
9303 &ad2,
9304 n_ff_exp,
9305 n_embd,
9306 n_expert,
9307 n_active,
9308 n_pairs,
9309 m.down_exps.qtype,
9310 m.down_exps.row_bytes,
9311 )?
9312 };
9313 let mut moe_out = e.uninit(t * n_embd)?;
9314 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9315 return Ok(moe_out);
9316 }
9317
9318 let g_len = m.gate_exps.expert_stride;
9319 let u_len = m.up_exps.expert_stride;
9320 let d_len = m.down_exps.expert_stride;
9321 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9325 let (mut sg, mut su, mut sd) = if dev.is_some() {
9326 (None, None, None)
9327 } else {
9328 (
9329 Some(e.alloc_u8_uninit(g_len)?),
9330 Some(e.alloc_u8_uninit(u_len)?),
9331 Some(e.alloc_u8_uninit(d_len)?),
9332 )
9333 };
9334 let mut moe_out = e.zeros(t * n_embd)?;
9335 for tok in 0..t {
9336 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9337 let w = &w_all[tok * n_used..(tok + 1) * n_used];
9338 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9339 for (j, &ex) in sel.iter().enumerate() {
9340 let ex = ex as usize;
9341 let gate = match dev {
9342 Some(d) => e.qmatvec_view(
9343 &d.gate,
9344 ex * g_len..(ex + 1) * g_len,
9345 &zt,
9346 1,
9347 m.gate_exps.in_f,
9348 m.gate_exps.out_f,
9349 m.gate_exps.qtype,
9350 m.gate_exps.row_bytes,
9351 )?,
9352 None => {
9353 let sg = sg.as_mut().unwrap();
9354 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9355 e.qmatvec_view(
9356 sg,
9357 0..g_len,
9358 &zt,
9359 1,
9360 m.gate_exps.in_f,
9361 m.gate_exps.out_f,
9362 m.gate_exps.qtype,
9363 m.gate_exps.row_bytes,
9364 )?
9365 }
9366 };
9367 let up = match dev {
9368 Some(d) => e.qmatvec_view(
9369 &d.up,
9370 ex * u_len..(ex + 1) * u_len,
9371 &zt,
9372 1,
9373 m.up_exps.in_f,
9374 m.up_exps.out_f,
9375 m.up_exps.qtype,
9376 m.up_exps.row_bytes,
9377 )?,
9378 None => {
9379 let su = su.as_mut().unwrap();
9380 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9381 e.qmatvec_view(
9382 su,
9383 0..u_len,
9384 &zt,
9385 1,
9386 m.up_exps.in_f,
9387 m.up_exps.out_f,
9388 m.up_exps.qtype,
9389 m.up_exps.row_bytes,
9390 )?
9391 }
9392 };
9393 let mut act = e.uninit(n_ff_exp)?;
9394 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9395 let actv = act.slice(0..n_ff_exp);
9396 let y = match dev {
9397 Some(d) => e.qmatvec_view(
9398 &d.down,
9399 ex * d_len..(ex + 1) * d_len,
9400 &actv,
9401 1,
9402 m.down_exps.in_f,
9403 m.down_exps.out_f,
9404 m.down_exps.qtype,
9405 m.down_exps.row_bytes,
9406 )?,
9407 None => {
9408 let sd = sd.as_mut().unwrap();
9409 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9410 e.qmatvec_view(
9411 sd,
9412 0..d_len,
9413 &actv,
9414 1,
9415 m.down_exps.in_f,
9416 m.down_exps.out_f,
9417 m.down_exps.qtype,
9418 m.down_exps.row_bytes,
9419 )?
9420 }
9421 };
9422 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9423 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9424 }
9425 }
9426 Ok(moe_out)
9427 }
9428
9429 fn gemma4_layer(
9431 &self,
9432 e: &Engine,
9433 il: usize,
9434 layer: &crate::hybrid::HybridLayer,
9435 x: &CudaSlice<f32>,
9436 pos_d: &CudaSlice<i32>,
9437 t: usize,
9438 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9439 let n_embd = self.cfg.n_embd as usize;
9440 let eps = self.cfg.rms_eps;
9441
9442 let mut h = e.zeros(t * n_embd)?;
9443 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9444 let Mixer::Full(fa) = &layer.mixer else {
9445 panic!("gemma4 layer {il} not full-attn")
9446 };
9447 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9448 let mut cur = e.zeros(t * n_embd)?;
9450 e.rms_norm(
9451 &o,
9452 layer.post_attn_norm.float_data(),
9453 &mut cur,
9454 n_embd,
9455 t,
9456 eps,
9457 )?;
9458 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9459 }
9460
9461 fn gemma4_layer_tail_add(
9465 &self,
9466 e: &Engine,
9467 layer: &crate::hybrid::HybridLayer,
9468 cur: &CudaSlice<f32>,
9469 x: &CudaSlice<f32>,
9470 t: usize,
9471 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9472 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9473 }
9474
9475 fn gemma4_layer_tail_add_n(
9478 &self,
9479 e: &Engine,
9480 layer: &crate::hybrid::HybridLayer,
9481 cur: &CudaSlice<f32>,
9482 x: &CudaSlice<f32>,
9483 t: usize,
9484 next_norm: Option<&CudaSlice<f32>>,
9485 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9486 let n_embd = self.cfg.n_embd as usize;
9487 let bits = layer.gemma4.as_ref().unwrap();
9488 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9489 let mut xn = e.uninit(t * n_embd)?;
9490 match next_norm {
9491 Some(w) => {
9492 let mut hn = e.uninit(t * n_embd)?;
9493 e.add_scale_rms_norm(
9494 &sn,
9495 &attn_out,
9496 bits.layer_scale,
9497 w,
9498 &mut xn,
9499 &mut hn,
9500 n_embd,
9501 t,
9502 self.cfg.rms_eps,
9503 )?;
9504 Ok((xn, Some(hn)))
9505 }
9506 None => {
9507 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9508 Ok((xn, None))
9509 }
9510 }
9511 }
9512
9513 fn gemma4_layer_tail_core(
9516 &self,
9517 e: &Engine,
9518 layer: &crate::hybrid::HybridLayer,
9519 cur: &CudaSlice<f32>,
9520 x: &CudaSlice<f32>,
9521 t: usize,
9522 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9523 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9524 }
9525
9526 fn gemma4_layer_tail_core_pn(
9533 &self,
9534 e: &Engine,
9535 layer: &crate::hybrid::HybridLayer,
9536 cur: &CudaSlice<f32>,
9537 x: &CudaSlice<f32>,
9538 t: usize,
9539 pre_norm: Option<&CudaSlice<f32>>,
9540 defer_post_norm: bool,
9541 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9542 let n_embd = self.cfg.n_embd as usize;
9543 let eps = self.cfg.rms_eps;
9544 let bits = layer.gemma4.as_ref().unwrap();
9545
9546 let Some(mbits) = bits.moe_bits.as_ref() else {
9549 let crate::hybrid::Ffn::Dense {
9550 ffn_gate,
9551 ffn_up,
9552 ffn_down,
9553 } = &layer.ffn
9554 else {
9555 panic!("gemma4 dense layer without Dense ffn")
9556 };
9557 let mut attn_out = e.uninit(t * n_embd)?;
9558 let mut zsh = e.uninit(t * n_embd)?;
9559 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9562 match pre_norm {
9563 Some(wa) if t == 1 => {
9564 zpair = Some(e.rms_pre_add_rms_norm_q8z(
9565 cur,
9566 wa,
9567 x,
9568 bits.ffn_norm.float_data(),
9569 &mut attn_out,
9570 &mut zsh,
9571 n_embd,
9572 t,
9573 eps,
9574 )?);
9575 }
9576 Some(wa) => e.rms_pre_add_rms_norm(
9577 cur,
9578 wa,
9579 x,
9580 bits.ffn_norm.float_data(),
9581 &mut attn_out,
9582 &mut zsh,
9583 n_embd,
9584 t,
9585 eps,
9586 )?,
9587 None => e.add_rms_norm(
9588 cur,
9589 x,
9590 bits.ffn_norm.float_data(),
9591 &mut attn_out,
9592 &mut zsh,
9593 n_embd,
9594 t,
9595 eps,
9596 )?,
9597 }
9598 let n_ff = ffn_gate.out_features();
9599 let (gate, up) = if t == 1 {
9605 let (zq, zd) = match zpair {
9606 Some(p) => p,
9607 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9608 };
9609 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9610 Some(p) => p,
9611 None => (
9612 e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9613 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9614 ),
9615 }
9616 } else {
9617 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9622 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9623 let fused = if f2b {
9624 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9625 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9626 } else {
9627 None
9628 };
9629 match fused {
9630 Some(p) => p,
9631 None => {
9632 e.mmq_act_begin();
9634 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9635 }
9636 }
9637 };
9638 let mut act = e.uninit(t * n_ff)?;
9639 let f0 = if e.uses_q8_1_fast(ffn_down) {
9642 let upv = e.view(&up, t * n_ff);
9643 let up_all = upv.slice(0..t * n_ff);
9644 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9645 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9646 } else {
9647 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9648 e.matmul(ffn_down, &act, t)?
9649 };
9650 if defer_post_norm {
9651 return Ok((f0, attn_out));
9652 }
9653 let mut sn = e.uninit(t * n_embd)?;
9654 e.rms_norm(
9655 &f0,
9656 bits.post_ffw_norm.float_data(),
9657 &mut sn,
9658 n_embd,
9659 t,
9660 eps,
9661 )?;
9662 return Ok((sn, attn_out));
9663 };
9664
9665 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9666 let mut attn_out = e.uninit(t * n_embd)?;
9671 let mut router_in = e.uninit(t * n_embd)?;
9672 let fast_moe = match &layer.ffn {
9673 crate::hybrid::Ffn::Moe(m) => {
9674 m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9675 && expert_dp4a_supported(m.gate_exps.qtype)
9676 && expert_dp4a_supported(m.up_exps.qtype)
9677 && expert_dp4a_supported(m.down_exps.qtype)
9678 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9679 }
9680 _ => false,
9681 };
9682 let q8z = t < PRIME_MIN_T && fast_moe;
9683 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9684 let (z0, m2) = e.add_rms_norm3_q8z(
9685 cur,
9686 x,
9687 bits.ffn_norm.float_data(),
9688 &mbits.router_scale_pre,
9689 mbits.pre_ffw_norm_2.float_data(),
9690 &mut attn_out,
9691 &mut router_in,
9692 n_embd,
9693 t,
9694 eps,
9695 )?;
9696 (None, Some(z0), Some(m2))
9697 } else {
9698 let mut zsh = e.uninit(t * n_embd)?;
9699 let mut moe_in = e.uninit(t * n_embd)?;
9700 e.add_rms_norm3(
9701 cur,
9702 x,
9703 bits.ffn_norm.float_data(),
9704 &mbits.router_scale_pre,
9705 mbits.pre_ffw_norm_2.float_data(),
9706 &mut attn_out,
9707 &mut zsh,
9708 &mut router_in,
9709 &mut moe_in,
9710 n_embd,
9711 t,
9712 eps,
9713 )?;
9714 (Some((zsh, moe_in)), None, None)
9715 };
9716 let attn_out2 = attn_out;
9717 #[allow(unused_variables)]
9718 let attn_out = &attn_out2;
9719 let n_ff = mbits.shared_gate.out_features();
9720 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9721 if t == 1 {
9722 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9723 Some(p) => p,
9724 None => {
9725 let h0 = e.zeros(0)?;
9726 (
9727 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
9728 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
9729 )
9730 }
9731 }
9732 } else {
9733 let h0 = e.zeros(0)?;
9735 (
9736 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
9737 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
9738 )
9739 }
9740 } else {
9741 let (zsh, _) = zsh_f32.as_ref().unwrap();
9742 (
9743 e.matmul(&mbits.shared_gate, zsh, t)?,
9744 e.matmul(&mbits.shared_up, zsh, t)?,
9745 )
9746 };
9747 let mut act = e.uninit(t * n_ff)?;
9748 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9749 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
9750 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
9751 panic!("gemma4 layer not MoE")
9752 };
9753 let moe0 = match (&moe_q8, &zsh_f32) {
9754 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
9755 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
9756 _ => unreachable!(),
9757 };
9758 let mut mlp = e.uninit(t * n_embd)?;
9760 let mut moe = e.uninit(t * n_embd)?;
9761 e.rms_norm2x(
9762 &mlp0,
9763 &moe0,
9764 mbits.post_ffw_norm_1.float_data(),
9765 mbits.post_ffw_norm_2.float_data(),
9766 &mut mlp,
9767 &mut moe,
9768 n_embd,
9769 t,
9770 eps,
9771 )?;
9772
9773 let mut sum = e.uninit(t * n_embd)?;
9776 let mut sn = e.uninit(t * n_embd)?;
9777 e.add_rms_norm(
9778 &mlp,
9779 &moe,
9780 bits.post_ffw_norm.float_data(),
9781 &mut sum,
9782 &mut sn,
9783 n_embd,
9784 t,
9785 eps,
9786 )?;
9787 Ok((sn, attn_out2))
9788 }
9789
9790 fn gemma4_layer_tail_add_nq(
9792 &self,
9793 e: &Engine,
9794 layer: &crate::hybrid::HybridLayer,
9795 cur: &CudaSlice<f32>,
9796 x: &CudaSlice<f32>,
9797 t: usize,
9798 next_norm: Option<&CudaSlice<f32>>,
9799 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9800 {
9801 let n_embd = self.cfg.n_embd as usize;
9802 let bits = layer.gemma4.as_ref().unwrap();
9803 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9804 let mut xn = e.uninit(t * n_embd)?;
9805 match next_norm {
9806 Some(w) => {
9807 let pair = e.add_scale_rms_norm_q8_1(
9808 &sn,
9809 &attn_out,
9810 bits.layer_scale,
9811 w,
9812 &mut xn,
9813 n_embd,
9814 t,
9815 self.cfg.rms_eps,
9816 )?;
9817 Ok((xn, Some(pair)))
9818 }
9819 None => {
9820 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9821 Ok((xn, None))
9822 }
9823 }
9824 }
9825
9826 fn gemma4_forward(
9829 &self,
9830 e: &Engine,
9831 tokens: &[u32],
9832 last_only: bool,
9833 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
9834 if self.is_gemma4_e4b() {
9837 return self.gemma4_e4b_forward(e, tokens, last_only);
9838 }
9839 let n_embd = self.cfg.n_embd as usize;
9840 let t = tokens.len();
9841 let pos: Vec<i32> = (0..t as i32).collect();
9842 let pos_d = e.htod_i32(&pos)?;
9843
9844 let mut x = self.embed(e, tokens)?;
9845 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9846 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
9849 let stat =
9850 |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
9851 let h = e.dtoh(x)?;
9852 let bad = h.iter().filter(|v| !v.is_finite()).count();
9853 let mx = h
9854 .iter()
9855 .filter(|v| v.is_finite())
9856 .fold(0.0f32, |m, v| m.max(v.abs()));
9857 eprintln!(
9858 "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
9859 &h[..3]
9860 );
9861 Ok(())
9862 };
9863 if probe {
9864 stat(e, &x, "embed")?;
9865 }
9866 for (il, layer) in self.layers.iter().enumerate() {
9867 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
9868 if probe {
9869 stat(e, &x, &format!("L{il}"))?;
9870 }
9871 }
9872 let mut hn = e.zeros(t * n_embd)?;
9873 e.rms_norm(
9874 &x,
9875 self.output_norm.float_data(),
9876 &mut hn,
9877 n_embd,
9878 t,
9879 self.cfg.rms_eps,
9880 )?;
9881 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9882 let n_vocab = self.output.out_features();
9883 let logits = if last_only {
9884 let hv = e.view(&hn, t * n_embd);
9885 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
9886 let mut hlast = e.zeros(n_embd)?;
9887 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
9888 let mut ld = e.matmul(&self.output, &hlast, 1)?;
9889 e.softcap(&mut ld, cap, n_vocab)?;
9890 self.gemma4_suppress(e, &mut ld, 1)?;
9891 e.dtoh(&ld)?
9892 } else {
9893 let mut ld = e.matmul(&self.output, &hn, t)?;
9894 e.softcap(&mut ld, cap, t * n_vocab)?;
9895 self.gemma4_suppress(e, &mut ld, t)?;
9896 e.dtoh(&ld)?
9897 };
9898 Ok(logits)
9899 }
9900
9901 pub(crate) fn gemma4_prime(
9906 &self,
9907 e: &Engine,
9908 tokens: &[u32],
9909 cache: &mut Cache,
9910 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9911 if cache.pos != 0 {
9916 return Err(
9917 "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
9918 — prime the full prompt in one call or decode tokenwise"
9919 .into(),
9920 );
9921 }
9922 let n_embd = self.cfg.n_embd as usize;
9923 let eps = self.cfg.rms_eps;
9924 let t = tokens.len();
9925 let pos: Vec<i32> = (0..t as i32).collect();
9926 let pos_d = e.htod_i32(&pos)?;
9927 let mut x = self.embed(e, tokens)?;
9928 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9929 for (il, layer) in self.layers.iter().enumerate() {
9930 let mut h = e.zeros(t * n_embd)?;
9931 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9932 let Mixer::Full(fa) = &layer.mixer else {
9933 panic!("gemma4 layer not full-attn")
9934 };
9935 let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
9936 let mut cur = e.zeros(t * n_embd)?;
9937 e.rms_norm(
9938 &o,
9939 layer.post_attn_norm.float_data(),
9940 &mut cur,
9941 n_embd,
9942 t,
9943 eps,
9944 )?;
9945 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
9946 self.dflash_tap(e, cache, il, &x, t)?;
9947 }
9948 cache.pos += t;
9949 let hiddens = e.clone_dtod(&x)?;
9950 let xv = e.view(&x, t * n_embd);
9951 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
9952 let mut h_seed = e.zeros(n_embd)?;
9953 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
9954 let mut hn = e.uninit(n_embd)?;
9955 e.rms_norm(
9956 &h_seed,
9957 self.output_norm.float_data(),
9958 &mut hn,
9959 n_embd,
9960 1,
9961 eps,
9962 )?;
9963 let mut ld = e.matmul(&self.output, &hn, 1)?;
9964 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9965 e.softcap(&mut ld, cap, self.output.out_features())?;
9966 self.gemma4_suppress(e, &mut ld, 1)?;
9967 let logits = e.dtoh(&ld)?;
9968 Ok((logits, h_seed, hiddens))
9969 }
9970
9971 fn gemma4_decode_attn(
9976 &self,
9977 e: &Engine,
9978 fa: &crate::hybrid::FullAttnLayer,
9979 il: usize,
9980 hq: &CudaSlice<i8>,
9981 hdq: &CudaSlice<f32>,
9982 pos_d: &CudaSlice<i32>,
9983 cache: &mut Cache,
9984 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9985 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
9986 let eps = self.cfg.rms_eps;
9987 let aux = self.gemma4_aux.as_ref().unwrap();
9988 let ones = aux.ones(e);
9989 #[cfg(debug_assertions)]
9990 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
9991 let (hq, hdq) = (hq, hdq);
9992 let h0 = e.zeros(0)?;
9993 let h = &h0;
9994 let (q0, k0, v0) = if swa {
9995 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
9996 Some(t3) => t3,
9997 None => (
9998 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
9999 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10000 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10001 ),
10002 }
10003 } else {
10004 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10005 Some(p) => p,
10006 None => (
10007 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10008 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10009 ),
10010 };
10011 let v0 = e.clone_dtod(&k0)?;
10012 (q0, k0, v0)
10013 };
10014 let mut q = e.uninit(nh * hd)?;
10015 let mut k = e.uninit(nkv * hd)?;
10016 let mut v = e.uninit(nkv * hd)?;
10017 let ff = if swa {
10020 None
10021 } else {
10022 Some(
10023 aux.rope_freqs(e)
10024 .expect("gemma4 global rope needs rope_freqs.weight"),
10025 )
10026 };
10027 #[cfg(debug_assertions)]
10028 if let Some(ff) = ff {
10029 crate::debug_assert_tensor_stream_device(
10030 ff,
10031 &e.stream(),
10032 "gemma4_decode_attn.rope_freqs",
10033 );
10034 }
10035 e.rms_norm_qkv_rope(
10036 &q0,
10037 &k0,
10038 &v0,
10039 fa.q_norm.float_data(),
10040 fa.k_norm.float_data(),
10041 ones,
10042 &mut q,
10043 &mut k,
10044 &mut v,
10045 hd,
10046 nh,
10047 nkv,
10048 pos_d,
10049 nh,
10050 nkv,
10051 base,
10052 1.0,
10053 ff,
10054 eps,
10055 )?;
10056 let kvl = cache.kv[il].as_mut().unwrap();
10057 e.append_kv_quantized(
10058 &k,
10059 &v,
10060 &mut kvl.k,
10061 &mut kvl.v,
10062 kvl.len,
10063 kvl.kv_dim_k,
10064 kvl.kv_dim_v,
10065 kvl.k_tok_bytes,
10066 kvl.v_tok_bytes,
10067 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
10068 )?;
10069 kvl.len += 1;
10070 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10074 let mut attn = e.uninit(nh * hd)?;
10075 if !swa
10077 && hd == 512
10078 && kvl.len >= crate::fa512_min_tkv()
10079 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10080 {
10081 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10082 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10083 let base = kvl.len as i32;
10085 e.i32_set_k(&mut kvl.len_d, base)?;
10086 e.fa_decode_rows(
10087 &q,
10088 &kp,
10089 &vp,
10090 &mut attn,
10091 hd,
10092 nh,
10093 nkv,
10094 kvl.len - 1,
10095 1,
10096 scale,
10097 kvl.k_tok_bytes,
10098 kvl.v_tok_bytes,
10099 Some((&kvl.len_d, -1)),
10100 false,
10101 false,
10102 None,
10103 )?;
10104 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10105 }
10106 if swa
10108 && kvl.len > win
10109 && hd == 256
10110 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10111 {
10112 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10113 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10114 let base = kvl.len as i32;
10115 e.i32_set_k(&mut kvl.len_d, base)?;
10116 e.fa_decode_rows_w(
10117 &q,
10118 &kp,
10119 &vp,
10120 &mut attn,
10121 hd,
10122 nh,
10123 nkv,
10124 &kvl.len_d,
10125 -1,
10126 1,
10127 scale,
10128 win,
10129 kvl.k_tok_bytes,
10130 kvl.v_tok_bytes,
10131 None,
10132 )?;
10133 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10134 }
10135 let (off_tok, t_kv) = if swa && kvl.len > win {
10136 (kvl.len - win, win)
10137 } else {
10138 (0, kvl.len)
10139 };
10140 let k_view = e.view_u8_range(
10141 &kvl.k,
10142 off_tok * kvl.k_tok_bytes,
10143 (off_tok + t_kv) * kvl.k_tok_bytes,
10144 );
10145 let v_view = e.view_u8_range(
10146 &kvl.v,
10147 off_tok * kvl.v_tok_bytes,
10148 (off_tok + t_kv) * kvl.v_tok_bytes,
10149 );
10150 e.fa_decode_kvmod(
10151 &q,
10152 &k_view,
10153 &v_view,
10154 &mut attn,
10155 hd,
10156 nh,
10157 nkv,
10158 t_kv,
10159 scale,
10160 kvl.k_tok_bytes,
10161 kvl.v_tok_bytes,
10162 swa && crate::Engine::wkv_on(),
10163 )?;
10164 Ok(e.matmul(&fa.wo, &attn, 1)?)
10165 }
10166
10167 #[allow(clippy::too_many_arguments)]
10174 pub fn gemma4_decode_step_dc(
10175 &self,
10176 e: &Engine,
10177 token_d: &CudaSlice<u32>,
10178 pos_d: &mut CudaSlice<i32>,
10179 embd_gpu: &CudaSlice<u8>,
10180 embd_qt: i32,
10181 embd_rb: usize,
10182 cache: &mut Cache,
10183 n_vocab: usize,
10184 cap_bucket_max: Option<(usize, usize)>,
10185 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10186 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10187 self.gemma4_decode_step_dc_into(
10188 e,
10189 token_d,
10190 pos_d,
10191 embd_gpu,
10192 embd_qt,
10193 embd_rb,
10194 cache,
10195 n_vocab,
10196 cap_bucket_max,
10197 &mut tok_out,
10198 )?;
10199 Ok(tok_out)
10200 }
10201
10202 #[allow(clippy::too_many_arguments)]
10205 pub fn gemma4_decode_step_dc_into(
10206 &self,
10207 e: &Engine,
10208 token_d: &CudaSlice<u32>,
10209 pos_d: &mut CudaSlice<i32>,
10210 embd_gpu: &CudaSlice<u8>,
10211 embd_qt: i32,
10212 embd_rb: usize,
10213 cache: &mut Cache,
10214 n_vocab: usize,
10215 cap_bucket_max: Option<(usize, usize)>,
10216 tok_out: &mut CudaSlice<u32>,
10217 ) -> Result<(), Box<dyn std::error::Error>> {
10218 let n_embd = self.cfg.n_embd as usize;
10219 let eps = self.cfg.rms_eps;
10220 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10221 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10222 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10223 let n_layers = self.layers.len();
10224 for (il, layer) in self.layers.iter().enumerate() {
10225 let (hq, hdq) = match h_carry.take() {
10226 Some(p) => p,
10227 None => {
10228 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10229 }
10230 };
10231 let Mixer::Full(fa) = &layer.mixer else {
10232 panic!("gemma4 layer {il} not full-attn")
10233 };
10234 let o =
10235 self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10236 let mut cur = e.uninit(n_embd)?;
10237 e.rms_norm(
10238 &o,
10239 layer.post_attn_norm.float_data(),
10240 &mut cur,
10241 n_embd,
10242 1,
10243 eps,
10244 )?;
10245 let next_norm = if il + 1 < n_layers {
10246 Some(self.layers[il + 1].attn_norm.float_data())
10247 } else {
10248 None
10249 };
10250 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
10251 x = xn;
10252 h_carry = hn;
10253 }
10254 let mut hn = e.uninit(n_embd)?;
10255 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10256 let mut logits = e.matmul(&self.output, &hn, 1)?;
10257 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10259 e.inc_seqlen(pos_d)?;
10260 if cap_bucket_max.is_none() {
10261 cache.pos += 1;
10262 }
10263 Ok(())
10264 }
10265
10266 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10273 let n_embd = self.cfg.n_embd as usize;
10274 let n_vocab = self.output.out_features();
10275 let n_layers = self.layers.len();
10276 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10277 for il in 0..n_layers {
10278 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10279 qmax = qmax.max(nh * hd);
10280 kvmax = kvmax.max(nkv * hd);
10281 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10282 ffmax = ffmax.max(ffn_gate.out_features());
10283 }
10284 }
10285 Ok(G4DcSlots {
10286 x: e.uninit(n_embd)?,
10287 xn: e.uninit(n_embd)?,
10288 cur: e.uninit(n_embd)?,
10289 hq: e.alloc_i8_uninit(n_embd)?,
10290 hd_: e.uninit(n_embd / 32)?,
10291 q0: e.uninit(qmax)?,
10292 k0: e.uninit(kvmax)?,
10293 v0: e.uninit(kvmax)?,
10294 q: e.uninit(qmax)?,
10295 k: e.uninit(kvmax)?,
10296 v: e.uninit(kvmax)?,
10297 attn: e.uninit(qmax)?,
10298 o: e.uninit(n_embd)?,
10299 attn_out: e.uninit(n_embd)?,
10300 zsh: e.uninit(n_embd)?,
10301 zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10304 zd: e.uninit(n_embd.max(qmax) / 32)?,
10305 gate: e.uninit(ffmax)?,
10306 up: e.uninit(ffmax)?,
10307 act: e.uninit(ffmax)?,
10308 actq: e.alloc_i8_uninit(ffmax)?,
10309 actd: e.uninit(ffmax / 32)?,
10310 f0: e.uninit(n_embd)?,
10311 sn: e.uninit(n_embd)?,
10312 hn: e.uninit(n_embd)?,
10313 logits: e.uninit(n_vocab)?,
10314 })
10315 }
10316
10317 fn g4_matvec_m1_into(
10320 &self,
10321 e: &Engine,
10322 w: &crate::model::GpuTensor,
10323 aq: &CudaSlice<i8>,
10324 ad: &CudaSlice<f32>,
10325 y: &mut CudaSlice<f32>,
10326 ) -> Result<(), Box<dyn std::error::Error>> {
10327 use crate::model::GpuTensor;
10328 let (bytes, qtype, row_bytes, scale, rp) = match w {
10329 GpuTensor::Quant {
10330 bytes,
10331 qtype,
10332 row_bytes,
10333 scale,
10334 rp,
10335 ..
10336 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10337 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10338 };
10339 let (mbytes, mrp) = match w {
10340 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10341 _ => (bytes, rp),
10342 };
10343 e.qmatvec_mmvq_into(
10344 mbytes,
10345 aq,
10346 ad,
10347 1,
10348 w.in_features(),
10349 w.out_features(),
10350 qtype,
10351 row_bytes,
10352 scale,
10353 mrp,
10354 y,
10355 )
10356 }
10357
10358 #[allow(clippy::too_many_arguments)]
10362 pub fn gemma4_decode_step_dc_slotted(
10363 &self,
10364 e: &Engine,
10365 token_d: &CudaSlice<u32>,
10366 pos_d: &mut CudaSlice<i32>,
10367 embd_gpu: &CudaSlice<u8>,
10368 embd_qt: i32,
10369 embd_rb: usize,
10370 cache: &mut Cache,
10371 n_vocab: usize,
10372 cap_bucket_max: Option<(usize, usize)>,
10373 sl: &mut G4DcSlots,
10374 tok_out: &mut CudaSlice<u32>,
10375 ring: Option<(&mut CudaSlice<u32>, usize)>,
10376 ) -> Result<(), Box<dyn std::error::Error>> {
10377 let n_embd = self.cfg.n_embd as usize;
10378 let eps = self.cfg.rms_eps;
10379 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10380 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10381 let n_layers = self.layers.len();
10382 let mut has_carry = false;
10383 for il in 0..n_layers {
10384 if !has_carry {
10385 e.rms_norm_q8_1_into(
10386 &sl.x,
10387 self.layers[il].attn_norm.float_data(),
10388 n_embd,
10389 1,
10390 eps,
10391 &mut sl.hq,
10392 &mut sl.hd_,
10393 )?;
10394 }
10395 has_carry = true;
10396 let layer = &self.layers[il];
10397 let Mixer::Full(fa) = &layer.mixer else {
10398 panic!("gemma4 layer {il} not full-attn")
10399 };
10400 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10401 e.rms_norm(
10402 &sl.o,
10403 layer.post_attn_norm.float_data(),
10404 &mut sl.cur,
10405 n_embd,
10406 1,
10407 eps,
10408 )?;
10409 let next_norm = if il + 1 < n_layers {
10410 Some(self.layers[il + 1].attn_norm.float_data())
10411 } else {
10412 None
10413 };
10414 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10415 std::mem::swap(&mut sl.x, &mut sl.xn);
10416 }
10417 e.rms_norm(
10418 &sl.x,
10419 self.output_norm.float_data(),
10420 &mut sl.hn,
10421 n_embd,
10422 1,
10423 eps,
10424 )?;
10425 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10426 {
10428 let (zq, zd) = (&sl.zq, &sl.zd);
10429 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10430 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10431 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10432 }
10433 self.gemma4_suppress(e, &mut sl.logits, 1)?;
10434 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10435 if let Some((ring, base)) = ring {
10436 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10440 }
10441 e.inc_seqlen(pos_d)?;
10442 if cap_bucket_max.is_none() {
10443 cache.pos += 1;
10444 }
10445 Ok(())
10446 }
10447
10448 #[allow(clippy::too_many_arguments)]
10450 fn gemma4_decode_attn_dc_slotted(
10451 &self,
10452 e: &Engine,
10453 fa: &crate::hybrid::FullAttnLayer,
10454 il: usize,
10455 pos_d: &CudaSlice<i32>,
10456 cache: &mut Cache,
10457 cap_bucket_max: Option<(usize, usize)>,
10458 sl: &mut G4DcSlots,
10459 ) -> Result<(), Box<dyn std::error::Error>> {
10460 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10461 let eps = self.cfg.rms_eps;
10462 let aux = self.gemma4_aux.as_ref().unwrap();
10463 let ones = aux.ones(e);
10464 #[cfg(debug_assertions)]
10465 crate::debug_assert_tensor_stream_device(
10466 ones,
10467 &e.stream(),
10468 "gemma4_decode_attn_dc_slotted.ones",
10469 );
10470 {
10471 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10472 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10473 if swa {
10474 if !e.matmul_q4_fused3_into(
10475 &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10476 )? {
10477 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10478 }
10479 } else {
10480 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
10481 return Err("slotted step: fused2 unavailable".into());
10482 }
10483 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10484 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10485 }
10486 }
10487 let ff = if swa {
10490 None
10491 } else {
10492 Some(
10493 aux.rope_freqs(e)
10494 .expect("gemma4 global rope needs rope_freqs.weight"),
10495 )
10496 };
10497 #[cfg(debug_assertions)]
10498 if let Some(ff) = ff {
10499 crate::debug_assert_tensor_stream_device(
10500 ff,
10501 &e.stream(),
10502 "gemma4_decode_attn_dc_slotted.rope_freqs",
10503 );
10504 }
10505 let kvl = cache.kv[il].as_mut().unwrap();
10506 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10507 if crate::Engine::qkv_append_on() {
10508 e.rms_norm_qkv_rope_append_dc(
10510 &sl.q0,
10511 &sl.k0,
10512 &sl.v0,
10513 fa.q_norm.float_data(),
10514 fa.k_norm.float_data(),
10515 ones,
10516 &mut sl.q,
10517 &mut sl.k,
10518 &mut sl.v,
10519 hd,
10520 nh,
10521 nkv,
10522 pos_d,
10523 nh,
10524 nkv,
10525 base,
10526 1.0,
10527 ff,
10528 eps,
10529 &mut kvl.k,
10530 &mut kvl.v,
10531 &kvl.len_d,
10532 kvl.k_tok_bytes,
10533 kvl.v_tok_bytes,
10534 kv_fp8,
10535 )?;
10536 } else {
10537 e.rms_norm_qkv_rope(
10538 &sl.q0,
10539 &sl.k0,
10540 &sl.v0,
10541 fa.q_norm.float_data(),
10542 fa.k_norm.float_data(),
10543 ones,
10544 &mut sl.q,
10545 &mut sl.k,
10546 &mut sl.v,
10547 hd,
10548 nh,
10549 nkv,
10550 pos_d,
10551 nh,
10552 nkv,
10553 base,
10554 1.0,
10555 ff,
10556 eps,
10557 )?;
10558 e.append_kv_quantized_dc(
10559 &sl.k,
10560 &sl.v,
10561 &mut kvl.k,
10562 &mut kvl.v,
10563 &kvl.len_d,
10564 kvl.kv_dim_k,
10565 kvl.kv_dim_v,
10566 kvl.k_tok_bytes,
10567 kvl.v_tok_bytes,
10568 kv_fp8,
10569 )?;
10570 }
10571 e.inc_seqlen(&mut kvl.len_d)?;
10572 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
10573 let k_view = e.view_u8(&kvl.k, kvl.k.len());
10574 let v_view = e.view_u8(&kvl.v, kvl.v.len());
10575 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10576 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10577 let mut fa_q8 = false;
10581 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10582 e.fa_decode_rows(
10583 &sl.q,
10584 &k_view,
10585 &v_view,
10586 &mut sl.attn,
10587 hd,
10588 nh,
10589 nkv,
10590 b_glob - 1,
10591 1,
10592 scale,
10593 kvl.k_tok_bytes,
10594 kvl.v_tok_bytes,
10595 Some((&kvl.len_d, -1)),
10596 false,
10597 false,
10598 Some((&mut sl.zq, &mut sl.zd)),
10599 )?;
10600 fa_q8 = true;
10601 } else if swa && b_swa > win && hd == 256 && rows_on {
10602 e.fa_decode_rows_w(
10603 &sl.q,
10604 &k_view,
10605 &v_view,
10606 &mut sl.attn,
10607 hd,
10608 nh,
10609 nkv,
10610 &kvl.len_d,
10611 -1,
10612 1,
10613 scale,
10614 win,
10615 kvl.k_tok_bytes,
10616 kvl.v_tok_bytes,
10617 Some((&mut sl.zq, &mut sl.zd)),
10618 )?;
10619 fa_q8 = true;
10620 } else {
10621 let b = if swa { b_swa } else { b_glob };
10622 e.fa_decode_dc(
10623 &sl.q,
10624 &k_view,
10625 &v_view,
10626 &mut sl.attn,
10627 hd,
10628 nh,
10629 nkv,
10630 &kvl.len_d,
10631 b,
10632 scale,
10633 kvl.k_tok_bytes,
10634 kvl.v_tok_bytes,
10635 swa && crate::Engine::wkv_on(),
10636 )?;
10637 }
10638 if !fa_q8 {
10639 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
10640 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
10641 }
10642 {
10643 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10644 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10645 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
10646 }
10647 Ok(())
10648 }
10649
10650 fn gemma4_layer_tail_slotted(
10653 &self,
10654 e: &Engine,
10655 layer: &crate::hybrid::HybridLayer,
10656 next_norm: Option<&CudaSlice<f32>>,
10657 sl: &mut G4DcSlots,
10658 ) -> Result<(), Box<dyn std::error::Error>> {
10659 let n_embd = self.cfg.n_embd as usize;
10660 let eps = self.cfg.rms_eps;
10661 let bits = layer.gemma4.as_ref().unwrap();
10662 let crate::hybrid::Ffn::Dense {
10663 ffn_gate,
10664 ffn_up,
10665 ffn_down,
10666 } = &layer.ffn
10667 else {
10668 return Err("slotted tail: dense ffn only".into());
10669 };
10670 e.add_rms_norm(
10671 &sl.cur,
10672 &sl.x,
10673 bits.ffn_norm.float_data(),
10674 &mut sl.attn_out,
10675 &mut sl.zsh,
10676 n_embd,
10677 1,
10678 eps,
10679 )?;
10680 let n_ff = ffn_gate.out_features();
10681 {
10682 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
10683 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10684 }
10685 {
10686 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10687 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10688 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
10689 return Err("slotted tail: ffn fused2 unavailable".into());
10690 }
10691 }
10692 debug_assert!(e.uses_q8_1_fast(ffn_down));
10693 {
10694 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
10695 let upv = e.view(upr, n_ff);
10696 let up_all = upv.slice(0..n_ff);
10697 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
10698 e.gelu_tanh_mul_q8_1_into(
10699 gr,
10700 &up_all,
10701 &mut sl.act,
10702 n_ff,
10703 1,
10704 &mut sl.actq,
10705 &mut sl.actd,
10706 )?;
10707 }
10708 {
10709 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
10710 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
10711 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
10712 }
10713 e.rms_norm(
10714 &sl.f0,
10715 bits.post_ffw_norm.float_data(),
10716 &mut sl.sn,
10717 n_embd,
10718 1,
10719 eps,
10720 )?;
10721 match next_norm {
10722 Some(w) => {
10723 e.add_scale_rms_norm_q8_1_into(
10724 &sl.sn,
10725 &sl.attn_out,
10726 bits.layer_scale,
10727 w,
10728 &mut sl.xn,
10729 n_embd,
10730 1,
10731 eps,
10732 &mut sl.hq,
10733 &mut sl.hd_,
10734 )?;
10735 }
10736 None => {
10737 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
10738 }
10739 }
10740 Ok(())
10741 }
10742
10743 #[allow(clippy::too_many_arguments)]
10745 fn gemma4_decode_attn_dc(
10746 &self,
10747 e: &Engine,
10748 fa: &crate::hybrid::FullAttnLayer,
10749 il: usize,
10750 hq: &CudaSlice<i8>,
10751 hdq: &CudaSlice<f32>,
10752 pos_d: &CudaSlice<i32>,
10753 cache: &mut Cache,
10754 cap_bucket_max: Option<(usize, usize)>,
10755 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10756 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10757 let eps = self.cfg.rms_eps;
10758 let aux = self.gemma4_aux.as_ref().unwrap();
10759 let ones = aux.ones(e);
10760 #[cfg(debug_assertions)]
10761 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
10762 let (q0, k0, v0) = if swa {
10763 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
10764 Some(t3) => t3,
10765 None => {
10766 let h0 = e.zeros(0)?;
10767 (
10768 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
10769 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
10770 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
10771 )
10772 }
10773 }
10774 } else {
10775 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
10776 Some(p) => p,
10777 None => {
10778 let h0 = e.zeros(0)?;
10779 (
10780 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
10781 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
10782 )
10783 }
10784 };
10785 let v0 = e.clone_dtod(&k0)?;
10786 (q0, k0, v0)
10787 };
10788 let mut q = e.uninit(nh * hd)?;
10789 let mut k = e.uninit(nkv * hd)?;
10790 let mut v = e.uninit(nkv * hd)?;
10791 let ff = if swa {
10793 None
10794 } else {
10795 Some(
10796 aux.rope_freqs(e)
10797 .expect("gemma4 global rope needs rope_freqs.weight"),
10798 )
10799 };
10800 #[cfg(debug_assertions)]
10801 if let Some(ff) = ff {
10802 crate::debug_assert_tensor_stream_device(
10803 ff,
10804 &e.stream(),
10805 "gemma4_decode_attn_dc.rope_freqs",
10806 );
10807 }
10808 let kvl = cache.kv[il].as_mut().unwrap();
10809 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10810 if crate::Engine::qkv_append_on() {
10811 e.rms_norm_qkv_rope_append_dc(
10813 &q0,
10814 &k0,
10815 &v0,
10816 fa.q_norm.float_data(),
10817 fa.k_norm.float_data(),
10818 ones,
10819 &mut q,
10820 &mut k,
10821 &mut v,
10822 hd,
10823 nh,
10824 nkv,
10825 pos_d,
10826 nh,
10827 nkv,
10828 base,
10829 1.0,
10830 ff,
10831 eps,
10832 &mut kvl.k,
10833 &mut kvl.v,
10834 &kvl.len_d,
10835 kvl.k_tok_bytes,
10836 kvl.v_tok_bytes,
10837 kv_fp8,
10838 )?;
10839 } else {
10840 e.rms_norm_qkv_rope(
10841 &q0,
10842 &k0,
10843 &v0,
10844 fa.q_norm.float_data(),
10845 fa.k_norm.float_data(),
10846 ones,
10847 &mut q,
10848 &mut k,
10849 &mut v,
10850 hd,
10851 nh,
10852 nkv,
10853 pos_d,
10854 nh,
10855 nkv,
10856 base,
10857 1.0,
10858 ff,
10859 eps,
10860 )?;
10861 e.append_kv_quantized_dc(
10862 &k,
10863 &v,
10864 &mut kvl.k,
10865 &mut kvl.v,
10866 &kvl.len_d,
10867 kvl.kv_dim_k,
10868 kvl.kv_dim_v,
10869 kvl.k_tok_bytes,
10870 kvl.v_tok_bytes,
10871 kv_fp8,
10872 )?;
10873 }
10874 e.inc_seqlen(&mut kvl.len_d)?;
10875 let mut attn = e.uninit(nh * hd)?;
10876 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10879 match cap_bucket_max {
10884 None => {
10885 kvl.len += 1;
10889 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10890 if !swa
10891 && hd == 512
10892 && kvl.len >= crate::fa512_min_tkv()
10893 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10894 {
10895 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10898 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10899 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
10900 e.fa_decode_rows(
10901 &q,
10902 &kp,
10903 &vp,
10904 &mut attn,
10905 hd,
10906 nh,
10907 nkv,
10908 kvl.len - 1,
10909 1,
10910 scale,
10911 kvl.k_tok_bytes,
10912 kvl.v_tok_bytes,
10913 Some((&kvl.len_d, -1)),
10914 false,
10915 false,
10916 Some((&mut aq8, &mut ad8)),
10917 )?;
10918 fa_q8 = Some((aq8, ad8));
10919 } else if swa
10920 && kvl.len > win
10921 && hd == 256
10922 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10923 {
10924 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10926 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10927 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
10928 e.fa_decode_rows_w(
10929 &q,
10930 &kp,
10931 &vp,
10932 &mut attn,
10933 hd,
10934 nh,
10935 nkv,
10936 &kvl.len_d,
10937 -1,
10938 1,
10939 scale,
10940 win,
10941 kvl.k_tok_bytes,
10942 kvl.v_tok_bytes,
10943 Some((&mut aq8, &mut ad8)),
10944 )?;
10945 fa_q8 = Some((aq8, ad8));
10946 } else {
10947 let (off_tok, t_kv) = if swa && kvl.len > win {
10948 (kvl.len - win, win)
10949 } else {
10950 (0, kvl.len)
10951 };
10952 let k_view = e.view_u8_range(
10953 &kvl.k,
10954 off_tok * kvl.k_tok_bytes,
10955 (off_tok + t_kv) * kvl.k_tok_bytes,
10956 );
10957 let v_view = e.view_u8_range(
10958 &kvl.v,
10959 off_tok * kvl.v_tok_bytes,
10960 (off_tok + t_kv) * kvl.v_tok_bytes,
10961 );
10962 e.fa_decode_kvmod(
10963 &q,
10964 &k_view,
10965 &v_view,
10966 &mut attn,
10967 hd,
10968 nh,
10969 nkv,
10970 t_kv,
10971 scale,
10972 kvl.k_tok_bytes,
10973 kvl.v_tok_bytes,
10974 swa && crate::Engine::wkv_on(),
10975 )?;
10976 }
10977 }
10978 Some((b_swa, b_glob)) => {
10979 let k_view = e.view_u8(&kvl.k, kvl.k.len());
10985 let v_view = e.view_u8(&kvl.v, kvl.v.len());
10986 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10987 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10988 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10989 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
10990 e.fa_decode_rows(
10991 &q,
10992 &k_view,
10993 &v_view,
10994 &mut attn,
10995 hd,
10996 nh,
10997 nkv,
10998 b_glob - 1,
10999 1,
11000 scale,
11001 kvl.k_tok_bytes,
11002 kvl.v_tok_bytes,
11003 Some((&kvl.len_d, -1)),
11004 false,
11005 false,
11006 Some((&mut aq8, &mut ad8)),
11007 )?;
11008 fa_q8 = Some((aq8, ad8));
11009 } else if swa && b_swa > win && hd == 256 && rows_on {
11010 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11011 e.fa_decode_rows_w(
11012 &q,
11013 &k_view,
11014 &v_view,
11015 &mut attn,
11016 hd,
11017 nh,
11018 nkv,
11019 &kvl.len_d,
11020 -1,
11021 1,
11022 scale,
11023 win,
11024 kvl.k_tok_bytes,
11025 kvl.v_tok_bytes,
11026 Some((&mut aq8, &mut ad8)),
11027 )?;
11028 fa_q8 = Some((aq8, ad8));
11029 } else {
11030 let b = if swa { b_swa } else { b_glob };
11031 e.fa_decode_dc(
11032 &q,
11033 &k_view,
11034 &v_view,
11035 &mut attn,
11036 hd,
11037 nh,
11038 nkv,
11039 &kvl.len_d,
11040 b,
11041 scale,
11042 kvl.k_tok_bytes,
11043 kvl.v_tok_bytes,
11044 swa && crate::Engine::wkv_on(),
11045 )?;
11046 }
11047 }
11048 }
11049 if let Some((aq8, ad8)) = fa_q8 {
11052 let mut y = e.uninit(fa.wo.out_features())?;
11053 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11054 return Ok(y);
11055 }
11056 Ok(e.matmul(&fa.wo, &attn, 1)?)
11057 }
11058
11059 pub fn gemma4_generate_graph(
11064 &self,
11065 e: &Engine,
11066 prompt_pos: usize,
11067 first_token: u32,
11068 cache: &mut Cache,
11069 max_new: usize,
11070 eos: &[u32],
11071 mut on_token: impl FnMut(u32) -> bool,
11072 ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11073 if self.is_gemma4_e4b() {
11074 return Err(
11075 "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11076 .into(),
11077 );
11078 }
11079 use crate::decode::StopReason;
11080 let n_vocab = self.output.out_features();
11081 let n_embd = self.cfg.n_embd as usize;
11082 let embd_gpu = self
11083 .embd_gpu
11084 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11085 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11086 for kvl in cache.kv.iter_mut().flatten() {
11087 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11088 }
11089 let mut token_d = e.stream().clone_htod(&[first_token])?;
11090 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11091 let g4 = self.cfg.gemma4.as_ref().unwrap();
11092 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11093 let nkv_s = g4
11095 .head_count_kv
11096 .iter()
11097 .zip(g4.swa_pattern.iter())
11098 .find(|p| *p.1)
11099 .map(|p| *p.0 as usize)
11100 .unwrap_or(8);
11101 let nkv_g = g4
11102 .head_count_kv
11103 .iter()
11104 .zip(g4.swa_pattern.iter())
11105 .find(|p| !*p.1)
11106 .map(|p| *p.0 as usize)
11107 .unwrap_or(2);
11108 let mut graphs: std::collections::HashMap<
11109 ((bool, usize), (bool, usize), bool, bool),
11110 (
11111 cudarc::driver::CudaGraph,
11112 Vec<Box<dyn std::any::Any + Send>>,
11113 ),
11114 > = Default::default();
11115 let mut slots = self.g4_dc_slots(e)?;
11118 const RING: usize = 64;
11121 const DRAIN: usize = 1;
11127 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11128 let ring_base = prompt_pos;
11129 let mut out = Vec::with_capacity(max_new);
11130 let mut reason = StopReason::MaxNew;
11131 let mut next = first_token;
11132 let mut captures = 0usize;
11133 for _ in 0..max_new {
11134 out.push(next);
11135 if eos.contains(&next) {
11136 reason = StopReason::Eos;
11137 break;
11138 }
11139 if !on_token(next) {
11140 reason = StopReason::Callback;
11141 break;
11142 }
11143 let t_kv = cache.pos + 1;
11144 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11152 let f512 = crate::fa512_min_tkv();
11153 let key_s = if t_kv > win {
11154 (true, usize::MAX)
11155 } else {
11156 e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11157 };
11158 let (key_g, rung_end) = if t_kv >= f512 {
11159 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11162 ((true, end), end)
11163 } else {
11164 (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11165 };
11166 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11167 if !graphs.contains_key(&key) {
11168 let bucket_max = (t_kv, rung_end);
11169 let snap = cache.snapshot(e)?;
11171 let pos_save = e.dtoh_i32_one(&pos_d)?;
11172 let len_save: Vec<Option<i32>> = cache
11173 .kv
11174 .iter()
11175 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11176 .collect();
11177 let tok_save = e.dtoh_u32_one(&token_d)?;
11178 let graph = {
11183 let tok_ref = &mut token_d;
11184 let pos_ref = &mut pos_d;
11185 let cache_ref = &mut *cache;
11186 let slots_ref = &mut slots;
11187 let ring_ref = &mut ring;
11188 e.capture_graph_retained_flags(
11189 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11190 |e| {
11191 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11193 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11194 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11195 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11196 cache_ref, n_vocab, Some(bucket_max),
11197 sl, tok_ref, Some((rg, ring_base)))
11198 })?
11199 };
11200 cache.rollback(e, &snap, 0)?;
11201 e.set_i32_one(&mut pos_d, pos_save)?;
11202 for (il, ls) in len_save.iter().enumerate() {
11203 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11204 e.set_i32_one(&mut kvl.len_d, *v)?;
11205 }
11206 }
11207 e.set_u32_one(&mut token_d, tok_save)?;
11208 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11209 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11210 eprintln!("[graph-census] {c:?}");
11211 }
11212 }
11213 graphs.insert(key, graph);
11214 captures += 1;
11215 }
11216 let mut chunk = 1usize;
11221 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11222 .ok()
11223 .and_then(|v| v.parse().ok())
11224 .unwrap_or(DRAIN);
11225 while chunk < drain_cap && out.len() + chunk < max_new {
11226 let t_next = cache.pos + 1 + chunk;
11227 let key_s2 = if t_next > win {
11228 (true, usize::MAX)
11229 } else {
11230 e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11231 };
11232 let key_g2 = if t_next >= f512 {
11233 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11234 } else {
11235 e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11236 };
11237 if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11238 break;
11239 }
11240 chunk += 1;
11241 }
11242 let g = &graphs.get(&key).unwrap().0;
11243 for _ in 0..chunk {
11244 g.launch()?;
11245 }
11246 e.stream().synchronize()?;
11247 let ringh = e.dtoh_u32(&ring)?;
11248 for j in 0..chunk {
11249 let pos_j = cache.pos + j;
11250 let tok_j = ringh[(pos_j - ring_base) % RING];
11251 cache.pos += 0; if j + 1 == chunk {
11253 next = tok_j;
11254 } else {
11255 out.push(tok_j);
11256 if eos.contains(&tok_j) || !on_token(tok_j) {
11257 reason = if eos.contains(&tok_j) {
11258 StopReason::Eos
11259 } else {
11260 StopReason::Callback
11261 };
11262 let keep = cache.pos + j + 1;
11264 e.set_i32_one(&mut pos_d, keep as i32)?;
11265 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11266 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11267 kvl.len = keep;
11268 }
11269 cache.pos = keep;
11270 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11271 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11272 }
11273 return Ok((out, reason));
11274 }
11275 }
11276 }
11277 cache.pos += chunk;
11278 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11279 kvl.len += chunk;
11280 }
11281 }
11282 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11283 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11284 }
11285 Ok((out, reason))
11286 }
11287
11288 pub(crate) fn gemma4_decode_step_t(
11294 &self,
11295 e: &Engine,
11296 tokens: &[u32],
11297 pos0: usize,
11298 cache: &mut Cache,
11299 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11300 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11301 }
11302
11303 pub(crate) fn gemma4_decode_step_t_am(
11307 &self,
11308 e: &Engine,
11309 tokens: &[u32],
11310 pos0: usize,
11311 cache: &mut Cache,
11312 ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11313 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11314 let t = tokens.len();
11315 let n_vocab = self.output.out_features();
11316 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11317 for i in 0..t {
11318 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11319 }
11320 Ok((e.dtoh_u32(&toks)?, hn))
11321 }
11322
11323 pub(crate) fn gemma4_decode_step_t_am_dev(
11326 &self,
11327 e: &Engine,
11328 tok_d: &CudaSlice<u32>,
11329 t: usize,
11330 pos0: usize,
11331 cache: &mut Cache,
11332 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11333 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11334 let n_vocab = self.output.out_features();
11335 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11336 for i in 0..t {
11337 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11338 }
11339 Ok((vam, hn))
11340 }
11341
11342 pub(crate) fn gemma4_decode_step_t_h(
11345 &self,
11346 e: &Engine,
11347 tokens: &[u32],
11348 pos0: usize,
11349 cache: &mut Cache,
11350 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11351 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11352 let t = tokens.len();
11353 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11354 e.softcap(&mut ld, cap, t * self.output.out_features())?;
11355 Ok((e.dtoh(&ld)?, hn))
11356 }
11357
11358 pub(crate) fn verify_stream_scratch(
11361 &self,
11362 e: &Engine,
11363 cap: usize,
11364 ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11365 Ok(VerifyStreamScratch {
11366 pos_d: e.htod_i32(&vec![0i32; cap])?,
11367 row_ctrs: (0..cap)
11368 .map(|_| e.htod_i32(&[0]))
11369 .collect::<Result<_, _>>()?,
11370 })
11371 }
11372
11373 pub(crate) fn gemma4_verify_t_am_stream(
11381 &self,
11382 e: &Engine,
11383 tok_d: &CudaSlice<u32>,
11384 t: usize,
11385 ctr: &CudaSlice<i32>,
11386 hint: usize,
11387 cache: &mut Cache,
11388 scr: &mut VerifyStreamScratch,
11389 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11390 let n_embd = self.cfg.n_embd as usize;
11391 let eps = self.cfg.rms_eps;
11392 assert!(t <= scr.row_ctrs.len() && t <= 64);
11393 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11394 for i in 0..t {
11395 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11396 }
11397 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11398 let embd_gpu = self
11399 .embd_gpu
11400 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11401 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11402 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11403 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11404 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11405 let n_layers = self.layers.len();
11406 for (il, layer) in self.layers.iter().enumerate() {
11407 let (hq, hdq) = match h_carry.take() {
11408 Some(p) => p,
11409 None => {
11410 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11411 }
11412 };
11413 let Mixer::Full(fa) = &layer.mixer else {
11414 panic!("gemma4 layer {il} not full-attn")
11415 };
11416 let o = self
11417 .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11418 let mut cur = e.uninit(t * n_embd)?;
11419 e.rms_norm(
11420 &o,
11421 layer.post_attn_norm.float_data(),
11422 &mut cur,
11423 n_embd,
11424 t,
11425 eps,
11426 )?;
11427 let next_norm = if il + 1 < n_layers {
11428 Some(self.layers[il + 1].attn_norm.float_data())
11429 } else {
11430 None
11431 };
11432 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
11433 x = xn;
11434 h_carry = hn;
11435 self.dflash_tap(e, cache, il, &x, t)?;
11436 }
11437 let mut hn = e.uninit(t * n_embd)?;
11438 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11439 let ld = e.matmul(&self.output, &hn, t)?;
11440 let n_vocab = self.output.out_features();
11441 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11442 for i in 0..t {
11443 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11444 }
11445 Ok((vam, hn))
11446 }
11447
11448 fn dflash_tap(
11455 &self,
11456 e: &Engine,
11457 cache: &mut Cache,
11458 il: usize,
11459 x: &CudaSlice<f32>,
11460 t: usize,
11461 ) -> Result<(), Box<dyn std::error::Error>> {
11462 let Some(taps) = cache.dflash_taps.as_mut() else {
11463 return Ok(());
11464 };
11465 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11466 return Ok(());
11467 };
11468 let h = taps.hidden;
11469 let n_taps = taps.layer_ids.len();
11470 debug_assert_eq!(taps.t, t);
11471 let xv = e.view(x, t * h);
11472 for r in 0..t {
11473 let row = xv.slice(r * h..(r + 1) * h);
11474 e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
11475 }
11476 Ok(())
11477 }
11478
11479 fn gemma4_verify_trunk(
11480 &self,
11481 e: &Engine,
11482 tokens: &[u32],
11483 pos0: usize,
11484 cache: &mut Cache,
11485 tok_dev: Option<&CudaSlice<u32>>,
11486 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11487 let n_embd = self.cfg.n_embd as usize;
11488 let eps = self.cfg.rms_eps;
11489 let t = tokens.len();
11490 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
11491 let pos_d = e.htod_i32(&pos)?;
11492 let mut x = match tok_dev {
11493 Some(td) => {
11494 let embd_gpu = self
11495 .embd_gpu
11496 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11497 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11498 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
11499 }
11500 None => e.htod(&self.embd.gather(n_embd, tokens))?,
11501 };
11502 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11503 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11504 let n_layers = self.layers.len();
11505 for (il, layer) in self.layers.iter().enumerate() {
11506 let (hq, hdq) = match h_carry.take() {
11507 Some(p) => p,
11508 None => {
11509 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11510 }
11511 };
11512 let Mixer::Full(fa) = &layer.mixer else {
11513 panic!("gemma4 layer {il} not full-attn")
11514 };
11515 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
11516 let mut cur = e.uninit(t * n_embd)?;
11517 e.rms_norm(
11518 &o,
11519 layer.post_attn_norm.float_data(),
11520 &mut cur,
11521 n_embd,
11522 t,
11523 eps,
11524 )?;
11525 let next_norm = if il + 1 < n_layers {
11526 Some(self.layers[il + 1].attn_norm.float_data())
11527 } else {
11528 None
11529 };
11530 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
11531 x = xn;
11532 h_carry = hn;
11533 self.dflash_tap(e, cache, il, &x, t)?;
11534 }
11535 let mut hn = e.uninit(t * n_embd)?;
11536 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11537 let mut ld = e.matmul(&self.output, &hn, t)?;
11538 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
11540 Ok((ld, hn))
11541 }
11542
11543 #[allow(clippy::too_many_arguments)]
11551 fn gemma4_verify_attn_stream(
11552 &self,
11553 e: &Engine,
11554 fa: &crate::hybrid::FullAttnLayer,
11555 il: usize,
11556 hq: &CudaSlice<i8>,
11557 hdq: &CudaSlice<f32>,
11558 pos_d: &CudaSlice<i32>,
11559 t: usize,
11560 cache: &mut Cache,
11561 hint: usize,
11562 row_ctrs: &[CudaSlice<i32>],
11563 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11564 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11565 let eps = self.cfg.rms_eps;
11566 let aux = self.gemma4_aux.as_ref().unwrap();
11567 let ones = aux.ones(e);
11568 #[cfg(debug_assertions)]
11569 crate::debug_assert_tensor_stream_device(
11570 ones,
11571 &e.stream(),
11572 "gemma4_verify_attn_stream.ones",
11573 );
11574 let h0 = e.zeros(0)?;
11575 let h = &h0;
11576 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11579 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11580 let fused_qkv = if f2b {
11581 if swa {
11582 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11583 .map(|(a, b, c)| (a, b, Some(c)))
11584 } else {
11585 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11586 .map(|(a, b)| (a, b, None))
11587 }
11588 } else {
11589 None
11590 };
11591 let (q0, k0, v0) = match fused_qkv {
11592 Some((a, b, cv)) => {
11593 let v = match cv {
11594 Some(c) => c,
11595 None => e.clone_dtod(&b)?,
11596 };
11597 (a, b, v)
11598 }
11599 None => {
11600 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11601 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11602 let v0 = if swa {
11603 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
11604 } else {
11605 e.clone_dtod(&k0)?
11606 };
11607 (q0, k0, v0)
11608 }
11609 };
11610 let mut q = e.uninit(t * nh * hd)?;
11611 let mut k = e.uninit(t * nkv * hd)?;
11612 let mut v = e.uninit(t * nkv * hd)?;
11613 let ff = if swa {
11616 None
11617 } else {
11618 Some(
11619 aux.rope_freqs(e)
11620 .expect("gemma4 global rope needs rope_freqs.weight"),
11621 )
11622 };
11623 #[cfg(debug_assertions)]
11624 if let Some(ff) = ff {
11625 crate::debug_assert_tensor_stream_device(
11626 ff,
11627 &e.stream(),
11628 "gemma4_verify_attn_stream.rope_freqs",
11629 );
11630 }
11631 e.rms_norm_qkv_rope(
11632 &q0,
11633 &k0,
11634 &v0,
11635 fa.q_norm.float_data(),
11636 fa.k_norm.float_data(),
11637 ones,
11638 &mut q,
11639 &mut k,
11640 &mut v,
11641 hd,
11642 nh * t,
11643 nkv * t,
11644 pos_d,
11645 nh,
11646 nkv,
11647 base,
11648 1.0,
11649 ff,
11650 eps,
11651 )?;
11652 let kvl = cache.kv[il].as_mut().unwrap();
11653 e.append_kv_quantized_rows_dc(
11655 &k,
11656 &v,
11657 &mut kvl.k,
11658 &mut kvl.v,
11659 &kvl.len_d,
11660 t,
11661 kvl.kv_dim_k,
11662 kvl.kv_dim_v,
11663 kvl.k_tok_bytes,
11664 kvl.v_tok_bytes,
11665 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
11666 )?;
11667 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11670 let mut attn = e.uninit(t * nh * hd)?;
11671 let k_view = e.view_u8(&kvl.k, kvl.k.len());
11672 let v_view = e.view_u8(&kvl.v, kvl.v.len());
11673 if swa && hint + 1 >= win {
11676 e.fa_decode_rows_w(
11679 &q,
11680 &k_view,
11681 &v_view,
11682 &mut attn,
11683 hd,
11684 nh,
11685 nkv,
11686 &kvl.len_d,
11687 0,
11688 t,
11689 scale,
11690 win,
11691 kvl.k_tok_bytes,
11692 kvl.v_tok_bytes,
11693 None,
11694 )?;
11695 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
11696 let bucket = (hint + t + 2)
11709 .next_power_of_two()
11710 .min(crate::fa512_min_tkv().saturating_sub(1));
11711 let qv = e.view(&q, t * nh * hd);
11712 for i in 0..t {
11713 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
11714 let mut q_one = e.uninit(nh * hd)?;
11715 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
11716 let mut a_one = e.uninit(nh * hd)?;
11717 e.fa_decode_dc(
11718 &q_one,
11719 &k_view,
11720 &v_view,
11721 &mut a_one,
11722 hd,
11723 nh,
11724 nkv,
11725 &row_ctrs[i],
11726 bucket,
11727 scale,
11728 kvl.k_tok_bytes,
11729 kvl.v_tok_bytes,
11730 false,
11731 )?;
11732 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
11733 }
11734 } else if hd == 512 {
11735 e.fa_decode_rows(
11738 &q,
11739 &k_view,
11740 &v_view,
11741 &mut attn,
11742 hd,
11743 nh,
11744 nkv,
11745 hint,
11746 t,
11747 scale,
11748 kvl.k_tok_bytes,
11749 kvl.v_tok_bytes,
11750 Some((&kvl.len_d, 0)),
11751 false,
11752 false,
11753 None,
11754 )?;
11755 } else {
11756 e.fa_decode_rows_dc(
11758 &q,
11759 &k_view,
11760 &v_view,
11761 &mut attn,
11762 hd,
11763 nh,
11764 nkv,
11765 &kvl.len_d,
11766 hint + t,
11767 t,
11768 scale,
11769 kvl.k_tok_bytes,
11770 kvl.v_tok_bytes,
11771 0,
11772 swa && crate::Engine::wkv_on(),
11773 )?;
11774 }
11775 Ok(e.matmul(&fa.wo, &attn, t)?)
11776 }
11777
11778 fn gemma4_verify_attn(
11779 &self,
11780 e: &Engine,
11781 fa: &crate::hybrid::FullAttnLayer,
11782 il: usize,
11783 hq: &CudaSlice<i8>,
11784 hdq: &CudaSlice<f32>,
11785 pos_d: &CudaSlice<i32>,
11786 t: usize,
11787 cache: &mut Cache,
11788 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11789 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11790 let eps = self.cfg.rms_eps;
11791 let aux = self.gemma4_aux.as_ref().unwrap();
11792 let ones = aux.ones(e);
11793 #[cfg(debug_assertions)]
11794 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
11795 let n_embd = self.cfg.n_embd as usize;
11796 let _ = n_embd;
11797
11798 let h0 = e.zeros(0)?;
11799 let h = &h0;
11800 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11803 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11804 let fused_qkv = if f2b {
11805 if swa {
11806 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11807 .map(|(a, b, c)| (a, b, Some(c)))
11808 } else {
11809 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11810 .map(|(a, b)| (a, b, None))
11811 }
11812 } else {
11813 None
11814 };
11815 let (q0, k0, v0) = match fused_qkv {
11816 Some((a, b, cv)) => {
11817 let v = match cv {
11818 Some(c) => c,
11819 None => e.clone_dtod(&b)?,
11820 };
11821 (a, b, v)
11822 }
11823 None => {
11824 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11825 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11826 let v0 = if swa {
11827 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
11828 } else {
11829 e.clone_dtod(&k0)?
11830 };
11831 (q0, k0, v0)
11832 }
11833 };
11834 let mut q = e.uninit(t * nh * hd)?;
11835 let mut k = e.uninit(t * nkv * hd)?;
11836 let mut v = e.uninit(t * nkv * hd)?;
11837 let ff = if swa {
11840 None
11841 } else {
11842 Some(
11843 aux.rope_freqs(e)
11844 .expect("gemma4 global rope needs rope_freqs.weight"),
11845 )
11846 };
11847 #[cfg(debug_assertions)]
11848 if let Some(ff) = ff {
11849 crate::debug_assert_tensor_stream_device(
11850 ff,
11851 &e.stream(),
11852 "gemma4_verify_attn.rope_freqs",
11853 );
11854 }
11855 e.rms_norm_qkv_rope(
11856 &q0,
11857 &k0,
11858 &v0,
11859 fa.q_norm.float_data(),
11860 fa.k_norm.float_data(),
11861 ones,
11862 &mut q,
11863 &mut k,
11864 &mut v,
11865 hd,
11866 nh * t,
11867 nkv * t,
11868 pos_d,
11869 nh,
11870 nkv,
11871 base,
11872 1.0,
11873 ff,
11874 eps,
11875 )?;
11876 let kvl = cache.kv[il].as_mut().unwrap();
11877 let base_len = kvl.len;
11878 e.append_kv_quantized_rows(
11879 &k,
11880 &v,
11881 &mut kvl.k,
11882 &mut kvl.v,
11883 base_len,
11884 t,
11885 kvl.kv_dim_k,
11886 kvl.kv_dim_v,
11887 kvl.k_tok_bytes,
11888 kvl.v_tok_bytes,
11889 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
11890 )?;
11891 kvl.len += t;
11892 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11893 let mut attn = e.uninit(t * nh * hd)?;
11894 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
11897 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
11900 if rows_ok && (!swa || base_len + t <= win) {
11901 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
11902 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
11903 if hd == 512 {
11904 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
11906 e.fa_decode_rows(
11907 &q,
11908 &k_view,
11909 &v_view,
11910 &mut attn,
11911 hd,
11912 nh,
11913 nkv,
11914 base_len,
11915 t,
11916 scale,
11917 kvl.k_tok_bytes,
11918 kvl.v_tok_bytes,
11919 Some((&kvl.len_d, 0)),
11920 false,
11921 swa && crate::Engine::wkv_on(),
11922 None,
11923 )?;
11924 } else {
11925 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
11929 e.fa_decode_rows_dc(
11930 &q,
11931 &k_view,
11932 &v_view,
11933 &mut attn,
11934 hd,
11935 nh,
11936 nkv,
11937 &kvl.len_d,
11938 base_len + t,
11939 t,
11940 scale,
11941 kvl.k_tok_bytes,
11942 kvl.v_tok_bytes,
11943 0,
11944 swa && crate::Engine::wkv_on(),
11945 )?;
11946 }
11947 return Ok(e.matmul(&fa.wo, &attn, t)?);
11948 }
11949 if hd == 256
11957 && swa
11958 && base_len + 1 >= win
11959 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11960 {
11961 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
11962 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
11963 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
11964 e.fa_decode_rows_w(
11965 &q,
11966 &k_view,
11967 &v_view,
11968 &mut attn,
11969 hd,
11970 nh,
11971 nkv,
11972 &kvl.len_d,
11973 0,
11974 t,
11975 scale,
11976 win,
11977 kvl.k_tok_bytes,
11978 kvl.v_tok_bytes,
11979 None,
11980 )?;
11981 return Ok(e.matmul(&fa.wo, &attn, t)?);
11982 }
11983 for i in 0..t {
11984 let avail = base_len + i + 1;
11985 let (off_tok, t_kv) = if swa && avail > win {
11986 (avail - win, win)
11987 } else {
11988 (0, avail)
11989 };
11990 let k_view = e.view_u8_range(
11991 &kvl.k,
11992 off_tok * kvl.k_tok_bytes,
11993 (off_tok + t_kv) * kvl.k_tok_bytes,
11994 );
11995 let v_view = e.view_u8_range(
11996 &kvl.v,
11997 off_tok * kvl.v_tok_bytes,
11998 (off_tok + t_kv) * kvl.v_tok_bytes,
11999 );
12000 let qi = e.view(&q, t * nh * hd);
12001 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12002 let mut q_one = e.uninit(nh * hd)?;
12003 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12004 let mut a_one = e.uninit(nh * hd)?;
12005 if swa
12009 && avail > win
12010 && hd == 256
12011 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12012 {
12013 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12014 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12015 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12016 e.fa_decode_rows_w(
12017 &q_one,
12018 &kp,
12019 &vp,
12020 &mut a_one,
12021 hd,
12022 nh,
12023 nkv,
12024 &kvl.len_d,
12025 0,
12026 1,
12027 scale,
12028 win,
12029 kvl.k_tok_bytes,
12030 kvl.v_tok_bytes,
12031 None,
12032 )?;
12033 } else if !swa
12034 && hd == 512
12035 && avail >= crate::fa512_min_tkv()
12036 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12037 {
12038 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12039 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12040 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12041 e.fa_decode_rows(
12042 &q_one,
12043 &kp,
12044 &vp,
12045 &mut a_one,
12046 hd,
12047 nh,
12048 nkv,
12049 avail - 1,
12050 1,
12051 scale,
12052 kvl.k_tok_bytes,
12053 kvl.v_tok_bytes,
12054 Some((&kvl.len_d, 0)),
12055 false,
12056 false,
12057 None,
12058 )?;
12059 } else {
12060 e.fa_decode_kvmod(
12061 &q_one,
12062 &k_view,
12063 &v_view,
12064 &mut a_one,
12065 hd,
12066 nh,
12067 nkv,
12068 t_kv,
12069 scale,
12070 kvl.k_tok_bytes,
12071 kvl.v_tok_bytes,
12072 swa && crate::Engine::wkv_on(),
12073 )?;
12074 }
12075 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12076 }
12077 Ok(e.matmul(&fa.wo, &attn, t)?)
12078 }
12079
12080 pub(crate) fn gemma4_decode_step_h(
12083 &self,
12084 e: &Engine,
12085 token: u32,
12086 cache: &mut Cache,
12087 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12088 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12093 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12094 }
12095 if crate::pp::pp_cuts(self.layers.len()).is_some() {
12096 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12097 }
12098 let n_embd = self.cfg.n_embd as usize;
12099 let eps = self.cfg.rms_eps;
12100 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12101 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12102 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12103 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12106 let n_layers = self.layers.len();
12107 for (il, layer) in self.layers.iter().enumerate() {
12108 let (hq, hdq) = match h_carry.take() {
12109 Some(p) => p,
12110 None => {
12111 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12112 }
12113 };
12114 let Mixer::Full(fa) = &layer.mixer else {
12115 panic!("gemma4 layer {il} not full-attn")
12116 };
12117 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12118 let mut cur = e.uninit(n_embd)?;
12119 e.rms_norm(
12120 &o,
12121 layer.post_attn_norm.float_data(),
12122 &mut cur,
12123 n_embd,
12124 1,
12125 eps,
12126 )?;
12127 let next_norm = if il + 1 < n_layers {
12128 Some(self.layers[il + 1].attn_norm.float_data())
12129 } else {
12130 None
12131 };
12132 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
12133 x = xn;
12134 h_carry = hn;
12135 }
12136 let mut hn = e.uninit(n_embd)?;
12137 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12138 let h_seed = e.clone_dtod(&x)?;
12139 let mut ld = e.matmul(&self.output, &hn, 1)?;
12140 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12141 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
12143 let logits = e.dtoh(&ld)?;
12144 cache.pos += 1;
12145 Ok((logits, h_seed))
12146 }
12147
12148 fn gemma4_decode_layers(
12156 &self,
12157 e: &Engine,
12158 mut x: CudaSlice<f32>,
12159 lo: usize,
12160 hi: usize,
12161 pos_d: &CudaSlice<i32>,
12162 cache: &mut Cache,
12163 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12164 let n_embd = self.cfg.n_embd as usize;
12165 let eps = self.cfg.rms_eps;
12166 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12167 for il in lo..hi {
12168 let layer = &self.layers[il];
12169 let (hq, hdq) = match h_carry.take() {
12170 Some(p) => p,
12171 None => {
12173 e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12174 }
12175 };
12176 let Mixer::Full(fa) = &layer.mixer else {
12177 panic!("gemma4 layer {il} not full-attn")
12178 };
12179 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12180 let mut cur = e.uninit(n_embd)?;
12181 e.rms_norm(
12182 &o,
12183 layer.post_attn_norm.float_data(),
12184 &mut cur,
12185 n_embd,
12186 1,
12187 eps,
12188 )?;
12189 let next_norm = if il + 1 < hi {
12190 Some(self.layers[il + 1].attn_norm.float_data())
12191 } else {
12192 None
12193 };
12194 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
12195 x = xn;
12196 h_carry = hn;
12197 }
12198 Ok(x)
12199 }
12200
12201 fn gemma4_decode_step_h_pp2(
12209 &self,
12210 e: &Engine,
12211 token: u32,
12212 cache: &mut Cache,
12213 split: usize,
12214 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12215 if crate::pp::pp2_streams_off() {
12216 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12217 }
12218 let rt = crate::pp::Pp2Rt::get(e)?;
12219 let e0 = rt.engine(0, e);
12220 let e1 = rt.engine(1, e);
12221 let n_embd = self.cfg.n_embd as usize;
12222 let eps = self.cfg.rms_eps;
12223 let pos = cache.pos as i32;
12224
12225 let slot = {
12227 let _st0 = rt.enter(0);
12228 let pos_d = e0.htod_i32(&[pos])?;
12229 #[cfg(debug_assertions)]
12230 crate::debug_assert_tensor_stream_device(
12231 &pos_d,
12232 &e0.stream(),
12233 "gemma4_decode_step_h_pp2.stage0.pos_d",
12234 );
12235 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12236 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12237 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12238 rt.tx(0, &x, n_embd)?
12239 };
12240
12241 let _st1 = rt.enter(1);
12243 let pos_d = e1.htod_i32(&[pos])?;
12244 #[cfg(debug_assertions)]
12245 crate::debug_assert_tensor_stream_device(
12246 &pos_d,
12247 &e1.stream(),
12248 "gemma4_decode_step_h_pp2.stage1.pos_d",
12249 );
12250 let x = rt.rx(0, slot, n_embd)?;
12251 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12252
12253 let mut hn = e1.uninit(n_embd)?;
12254 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12255 let h_seed = e1.clone_dtod(&x)?;
12256 let mut ld = e1.matmul(&self.output, &hn, 1)?;
12257 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12258 e1.softcap(&mut ld, cap, self.output.out_features())?;
12259 self.gemma4_suppress(e1, &mut ld, 1)?;
12260 let logits = e1.dtoh(&ld)?;
12261 cache.pos += 1;
12262 Ok((logits, h_seed))
12263 }
12264
12265 fn gemma4_decode_step_h_pp2_samestream(
12268 &self,
12269 e: &Engine,
12270 token: u32,
12271 cache: &mut Cache,
12272 split: usize,
12273 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12274 let n_embd = self.cfg.n_embd as usize;
12275 let eps = self.cfg.rms_eps;
12276 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12277
12278 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12280 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12281 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12282
12283 let boundary_tx = e.clone_dtod(&x)?;
12285 let boundary_rx = e.clone_dtod(&boundary_tx)?;
12286
12287 let x =
12289 self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12290
12291 let mut hn = e.uninit(n_embd)?;
12292 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12293 let h_seed = e.clone_dtod(&x)?;
12294 let mut ld = e.matmul(&self.output, &hn, 1)?;
12295 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12296 e.softcap(&mut ld, cap, self.output.out_features())?;
12297 self.gemma4_suppress(e, &mut ld, 1)?;
12298 let logits = e.dtoh(&ld)?;
12299 cache.pos += 1;
12300 Ok((logits, h_seed))
12301 }
12302}
12303
12304impl HybridModel {
12323 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12326 let geometry = self
12327 .cfg
12328 .layer_geometry(il as u32)
12329 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12330 debug_assert_eq!(
12331 geometry.attention_gate,
12332 memra_gguf::config::AttentionGateKind::SeparateHead
12333 );
12334 geometry
12335 }
12336
12337 #[allow(clippy::too_many_arguments)]
12397 fn step35_attn_pre_wo(
12398 &self,
12399 e: &Engine,
12400 fa: &FullAttnLayer,
12401 mut g3: Vec<CudaSlice<f32>>,
12402 hg: Option<&CudaSlice<f32>>,
12403 gt_pre: Option<&CudaSlice<f32>>,
12404 pos_d: &CudaSlice<i32>,
12405 t: usize,
12406 cache: Option<&mut Cache>,
12407 il: usize,
12408 seq_end: usize,
12409 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12410 let geometry = self.step35_geom(il);
12411 let hd = geometry.head_dim_k as usize;
12412 let nkv = geometry.n_head_kv as usize;
12413 let nh = geometry.n_head as usize;
12414 let rbase = geometry.rope_base;
12415 let scale = geometry.attention_scale();
12416 let swa = geometry.window.is_some();
12417 let eps = self.cfg.rms_eps;
12418 let win = geometry.window.unwrap_or(0) as usize;
12419 let n_rot = geometry.n_rot as usize;
12420
12421 let v = g3.pop().unwrap();
12422 let k0 = g3.pop().unwrap();
12423 let q0 = g3.pop().unwrap();
12424
12425 let mut q = e.uninit(t * nh * hd)?;
12429 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12430 let mut k = e.uninit(t * nkv * hd)?;
12431 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12432 let ff = if geometry.rope_factors {
12433 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12434 } else {
12435 None
12436 };
12437 #[cfg(debug_assertions)]
12438 if let Some(ff) = ff {
12439 crate::debug_assert_tensor_stream_device(
12440 ff,
12441 &e.stream(),
12442 "step35_attn_pre_wo.rope_freqs",
12443 );
12444 }
12445 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12446
12447 let mut attn = e.uninit(t * nh * hd)?;
12448 match cache {
12449 Some(cache) => {
12450 let base_len = cache.kv[il].as_ref().unwrap().len;
12451 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12453 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12454 let off = if swa {
12455 let raw = base_len.saturating_sub(win - 1);
12456 if legacy_tkv || legacy_calllocal {
12457 raw
12458 } else {
12459 raw & !31usize
12460 }
12461 } else {
12462 0
12463 };
12464 {
12465 let kvl = cache.kv[il].as_mut().unwrap();
12466 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12467 let write_row = e.prepare_kv_append(kvl, off, t)?;
12468 e.append_kv_quantized_rows(
12469 &k,
12470 &v,
12471 &mut kvl.k,
12472 &mut kvl.v,
12473 write_row,
12474 t,
12475 kvl.kv_dim_k,
12476 kvl.kv_dim_v,
12477 kvl.k_tok_bytes,
12478 kvl.v_tok_bytes,
12479 crate::Engine::kv_fp8_on(),
12480 )?;
12481 kvl.len += t;
12482 let new_len = kvl.len as i32;
12483 e.set_i32_one(&mut kvl.len_d, new_len)?;
12484 }
12485 let kvl = cache.kv[il].as_ref().unwrap();
12486 let t_kv = base_len + t - off;
12509 let physical = kvl.physical_rows(off, off + t_kv)?;
12510 let k_view = e.view_u8_range(
12511 &kvl.k,
12512 physical.start * kvl.k_tok_bytes,
12513 physical.end * kvl.k_tok_bytes,
12514 );
12515 let v_view = e.view_u8_range(
12516 &kvl.v,
12517 physical.start * kvl.v_tok_bytes,
12518 physical.end * kvl.v_tok_bytes,
12519 );
12520 let swa_naive = if legacy_tkv {
12532 t_kv > win
12533 } else {
12534 seq_end > win
12535 };
12536 if swa && swa_naive {
12537 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
12550 e.sdpa_naive_w_quantized_view(
12551 &q,
12552 &k_view,
12553 &v_view,
12554 &mut attn,
12555 hd,
12556 nh,
12557 nkv,
12558 t,
12559 t_kv,
12560 scale,
12561 true,
12562 win,
12563 kvl.k_tok_bytes,
12564 kvl.v_tok_bytes,
12565 )?;
12566 } else {
12567 e.fa_prefill_view_ws_w_hd128(
12568 &q,
12569 &k_view,
12570 &v_view,
12571 &mut attn,
12572 hd,
12573 nh,
12574 nkv,
12575 t,
12576 t_kv,
12577 scale,
12578 true,
12579 win,
12580 kvl.k_tok_bytes,
12581 kvl.v_tok_bytes,
12582 )?;
12583 }
12584 } else if std::env::var("MEMRA_NOFA").is_ok() {
12585 e.sdpa_naive_quantized_view(
12586 &q,
12587 &k_view,
12588 &v_view,
12589 &mut attn,
12590 hd,
12591 nh,
12592 nkv,
12593 t,
12594 t_kv,
12595 scale,
12596 true,
12597 kvl.k_tok_bytes,
12598 kvl.v_tok_bytes,
12599 )?;
12600 } else {
12601 e.fa_prefill_view_ws(
12606 &q,
12607 &k_view,
12608 &v_view,
12609 &mut attn,
12610 hd,
12611 nh,
12612 nkv,
12613 t,
12614 t_kv,
12615 scale,
12616 true,
12617 kvl.k_tok_bytes,
12618 kvl.v_tok_bytes,
12619 crate::Engine::kv_fp8_on(),
12620 )?;
12621 }
12622 }
12623 None => {
12624 debug_assert_eq!(
12629 seq_end, t,
12630 "step35 cacheless prefill is monolithic (seq_end == t)"
12631 );
12632 if swa && seq_end > win {
12633 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
12634 } else if std::env::var("MEMRA_NOFA").is_ok() {
12635 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12636 } else {
12637 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12638 }
12639 }
12640 }
12641
12642 let gw = fa
12645 .attn_gate
12646 .as_ref()
12647 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
12648 let gt_owned = if gt_pre.is_none() {
12649 Some(e.matmul(
12650 gw,
12651 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
12652 t,
12653 )?)
12654 } else {
12655 None
12656 };
12657 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
12658 let mut ag = e.uninit(t * nh * hd)?;
12659 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
12660 Ok(ag)
12661 }
12662
12663 pub(crate) fn step35_attn(
12666 &self,
12667 e: &Engine,
12668 fa: &FullAttnLayer,
12669 h: &CudaSlice<f32>,
12670 pos_d: &CudaSlice<i32>,
12671 t: usize,
12672 il: usize,
12673 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12674 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
12675 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
12677 Ok(e.matmul(&fa.wo, &ag, t)?)
12678 }
12679
12680 #[allow(clippy::too_many_arguments)]
12687 pub(crate) fn step35_attn_prime(
12688 &self,
12689 e: &Engine,
12690 fa: &FullAttnLayer,
12691 h: &CudaSlice<f32>,
12692 hx: Option<&CudaSlice<u8>>,
12693 pos_d: &CudaSlice<i32>,
12694 t: usize,
12695 cache: &mut Cache,
12696 il: usize,
12697 seq_end: usize,
12698 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12699 let g3 = match hx {
12700 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
12701 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
12702 };
12703 let ag =
12704 self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
12705 Ok(e.matmul(&fa.wo, &ag, t)?)
12706 }
12707
12708 #[allow(clippy::too_many_arguments)]
12718 pub(crate) fn step35_decode_attn(
12719 &self,
12720 e: &Engine,
12721 fa: &FullAttnLayer,
12722 il: usize,
12723 h: &CudaSlice<f32>,
12724 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
12725 pos_d: &CudaSlice<i32>,
12726 cache: &mut Cache,
12727 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12728 let geometry = self.step35_geom(il);
12729 let hd = geometry.head_dim_k as usize;
12730 let nkv = geometry.n_head_kv as usize;
12731 let nh = geometry.n_head as usize;
12732 let rbase = geometry.rope_base;
12733 let scale = geometry.attention_scale();
12734 let swa = geometry.window.is_some();
12735 let eps = self.cfg.rms_eps;
12736 let win = geometry.window.unwrap_or(0) as usize;
12737 let n_rot = geometry.n_rot as usize;
12738 let n_embd = self.cfg.n_embd as usize;
12739 let gw = fa
12740 .attn_gate
12741 .as_ref()
12742 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
12743
12744 let (q0, k0, v0, gt) = match pre_q {
12745 Some((hq, hdq)) => {
12746 debug_assert!(
12747 e.uses_q8_1_fast(gw),
12748 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
12749 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
12750 );
12751 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
12752 Some(t3) => t3,
12753 None => (
12754 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
12755 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
12756 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
12757 ),
12758 };
12759 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
12760 (a, b, c, gt)
12761 }
12762 None => {
12763 if e.uses_q8_1_fast(&fa.wq)
12764 && e.uses_q8_1_fast(&fa.wk)
12765 && e.uses_q8_1_fast(&fa.wv)
12766 && e.uses_q8_1_fast(gw)
12767 {
12768 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
12769 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
12770 Some(t3) => t3,
12771 None => (
12772 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
12773 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
12774 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
12775 ),
12776 };
12777 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
12778 (a, b, c, gt)
12779 } else {
12780 (
12781 e.matmul(&fa.wq, h, 1)?,
12782 e.matmul(&fa.wk, h, 1)?,
12783 e.matmul(&fa.wv, h, 1)?,
12784 e.matmul(gw, h, 1)?,
12785 )
12786 }
12787 }
12788 };
12789
12790 let mut q = e.uninit(nh * hd)?;
12791 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
12792 let mut k = e.uninit(nkv * hd)?;
12793 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
12794 let ff = if swa {
12795 None
12796 } else {
12797 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12798 };
12799 #[cfg(debug_assertions)]
12800 if let Some(ff) = ff {
12801 crate::debug_assert_tensor_stream_device(
12802 ff,
12803 &e.stream(),
12804 "step35_decode_attn.rope_freqs",
12805 );
12806 }
12807 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
12808
12809 if std::env::var("MEMRA_NOFA").is_ok() {
12810 return Err(
12811 "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
12812 cache; unset MEMRA_NOFA to use fa_decode"
12813 .into(),
12814 );
12815 }
12816 let kvl = cache.kv[il].as_mut().unwrap();
12817 let next_len = kvl.len + 1;
12818 let (off, t_kv) = if swa && next_len > win {
12819 (next_len - win, win)
12820 } else {
12821 (0, next_len)
12822 };
12823 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
12824 e.append_kv_quantized(
12825 &k,
12826 &v0,
12827 &mut kvl.k,
12828 &mut kvl.v,
12829 write_row,
12830 kvl.kv_dim_k,
12831 kvl.kv_dim_v,
12832 kvl.k_tok_bytes,
12833 kvl.v_tok_bytes,
12834 crate::Engine::kv_fp8_on(),
12835 )?;
12836 kvl.len = next_len;
12837 let physical = kvl.physical_rows(off, off + t_kv)?;
12838 let k_view = e.view_u8_range(
12839 &kvl.k,
12840 physical.start * kvl.k_tok_bytes,
12841 physical.end * kvl.k_tok_bytes,
12842 );
12843 let v_view = e.view_u8_range(
12844 &kvl.v,
12845 physical.start * kvl.v_tok_bytes,
12846 physical.end * kvl.v_tok_bytes,
12847 );
12848 let mut attn = e.uninit(nh * hd)?;
12849 e.fa_decode_kvmod(
12850 &q,
12851 &k_view,
12852 &v_view,
12853 &mut attn,
12854 hd,
12855 nh,
12856 nkv,
12857 t_kv,
12858 scale,
12859 kvl.k_tok_bytes,
12860 kvl.v_tok_bytes,
12861 crate::Engine::kv_fp8_on(),
12862 )?;
12863
12864 let mut ag = e.uninit(nh * hd)?;
12865 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
12866 Ok(e.matmul(&fa.wo, &ag, 1)?)
12867 }
12868}
12869
12870impl HybridModel {
12879 pub fn is_gemma4_e4b(&self) -> bool {
12880 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
12881 }
12882
12883 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
12887 let g = self.cfg.gemma4.as_ref().unwrap();
12888 let swa = g.swa_pattern[il];
12889 let hd = if swa {
12890 g.key_length_swa
12891 } else {
12892 g.key_length_global
12893 } as usize;
12894 let Mixer::Full(fa) = &self.layers[il].mixer else {
12895 panic!("e4b layer {il} not full-attn")
12896 };
12897 let nh = fa.wq.out_features() / hd;
12898 let nkv = fa.wk.out_features() / hd;
12899 (
12900 hd,
12901 nkv,
12902 nh,
12903 if swa {
12904 g.rope_base_swa
12905 } else {
12906 g.rope_base_global
12907 },
12908 1.0,
12909 swa,
12910 )
12911 }
12912
12913 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
12915 self.layers[il]
12916 .gemma4
12917 .as_ref()
12918 .and_then(|b| b.e4b.as_ref())
12919 .and_then(|e4| e4.kv_share.map(|t| t as usize))
12920 }
12921
12922 fn gemma4_e4b_inp_pl(
12927 &self,
12928 e: &Engine,
12929 tokens: &[u32],
12930 x_scaled: &CudaSlice<f32>,
12931 t: usize,
12932 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12933 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
12934 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
12935 }
12936
12937 fn gemma4_e4b_inp_pl_dev(
12939 &self,
12940 e: &Engine,
12941 tok_d: &CudaSlice<u32>,
12942 x_scaled: &CudaSlice<f32>,
12943 t: usize,
12944 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12945 let aux = self.gemma4_aux.as_ref().unwrap();
12946 let m = aux.e4b.as_ref().unwrap();
12947 let n_embd = self.cfg.n_embd as usize;
12948 let n_layer = self.layers.len();
12949 let width = m.n_epl * n_layer;
12950 let tbl = m.tok_tbl_gpu.get_or_init(|| {
12951 e.upload_u8(&m.tok_embd_bytes)
12952 .expect("e4b per-layer token table upload")
12953 });
12954 let mut a =
12955 e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
12956 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
12957 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
12958 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
12959 let mut pn = e.uninit(t * width)?;
12960 e.rms_norm(
12961 &p,
12962 m.proj_norm.float_data(),
12963 &mut pn,
12964 m.n_epl,
12965 t * n_layer,
12966 self.cfg.rms_eps,
12967 )?;
12968 let mut out = e.uninit(t * width)?;
12969 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
12970 Ok(out)
12971 }
12972
12973 #[allow(clippy::too_many_arguments)]
12978 fn gemma4_e4b_attn(
12979 &self,
12980 e: &Engine,
12981 il: usize,
12982 hq: &CudaSlice<i8>,
12983 hdq: &CudaSlice<f32>,
12984 pos_d: &CudaSlice<i32>,
12985 t: usize,
12986 cache: &mut Cache,
12987 dc_bucket: Option<usize>,
12988 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12989 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
12990 let eps = self.cfg.rms_eps;
12991 let aux = self.gemma4_aux.as_ref().unwrap();
12992 let ones = aux.ones(e);
12993 #[cfg(debug_assertions)]
12994 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
12995 let Mixer::Full(fa) = &self.layers[il].mixer else {
12996 unreachable!()
12997 };
12998 let h0 = e.zeros(0)?;
13002 let h = &h0;
13003
13004 let ff = if swa {
13005 None
13006 } else {
13007 Some(
13008 aux.rope_freqs(e)
13009 .expect("e4b global rope needs rope_freqs.weight"),
13010 )
13011 };
13012 #[cfg(debug_assertions)]
13013 if let Some(ff) = ff {
13014 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13015 }
13016 let share = self.gemma4_e4b_kv_target(il);
13017 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13019 let mut q;
13020 if let Some(_tgt) = share {
13021 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13022 q = e.uninit(t * nh * hd)?;
13023 let mut kdummy = e.uninit(1)?;
13026 let mut vdummy = e.uninit(1)?;
13027 e.rms_norm_qkv_rope(
13028 &q0,
13029 &q0,
13030 &q0,
13031 fa.q_norm.float_data(),
13032 fa.q_norm.float_data(),
13033 ones,
13034 &mut q,
13035 &mut kdummy,
13036 &mut vdummy,
13037 hd,
13038 nh * t,
13039 0,
13040 pos_d,
13041 nh,
13042 1,
13043 base,
13044 1.0,
13045 ff,
13046 eps,
13047 )?;
13048 } else {
13049 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13053 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13054 q = e.uninit(t * nh * hd)?;
13055 let mut k = e.uninit(t * nkv * hd)?;
13056 let mut v = e.uninit(t * nkv * hd)?;
13057 if t == 1 && cat.is_some() {
13058 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13059 e.rms_norm_qkv_rope_cat(
13060 &qkv0,
13061 fa.q_norm.float_data(),
13062 fa.k_norm.float_data(),
13063 ones,
13064 &mut q,
13065 &mut k,
13066 &mut v,
13067 hd,
13068 nh,
13069 nkv,
13070 pos_d,
13071 nh,
13072 nkv,
13073 base,
13074 1.0,
13075 ff,
13076 eps,
13077 )?;
13078 } else {
13079 let (q0, k0, v0) = match if t == 1 {
13080 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13081 } else {
13082 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13085 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13086 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13087 } else {
13088 None
13089 }
13090 } {
13091 Some(triple) => triple,
13092 None => (
13093 e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13094 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13095 e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13096 ), };
13098 e.rms_norm_qkv_rope(
13101 &q0,
13102 &k0,
13103 &v0,
13104 fa.q_norm.float_data(),
13105 fa.k_norm.float_data(),
13106 ones,
13107 &mut q,
13108 &mut k,
13109 &mut v,
13110 hd,
13111 nh * t,
13112 nkv * t,
13113 pos_d,
13114 nh,
13115 nkv,
13116 base,
13117 1.0,
13118 ff,
13119 eps,
13120 )?;
13121 }
13122 let kvl = cache.kv[il].as_mut().unwrap();
13123 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13127 if dc_bucket.is_some() {
13128 debug_assert!(t == 1);
13133 e.append_kv_quantized_row_dc_inc(
13135 &k,
13136 &v,
13137 &mut kvl.k,
13138 &mut kvl.v,
13139 &mut kvl.len_d,
13140 kvl.kv_dim_k,
13141 kvl.kv_dim_v,
13142 kvl.k_tok_bytes,
13143 kvl.v_tok_bytes,
13144 cls,
13145 )?;
13146 } else {
13147 e.append_kv_quantized_rows(
13148 &k,
13149 &v,
13150 &mut kvl.k,
13151 &mut kvl.v,
13152 kvl.len,
13153 t,
13154 kvl.kv_dim_k,
13155 kvl.kv_dim_v,
13156 kvl.k_tok_bytes,
13157 kvl.v_tok_bytes,
13158 cls,
13159 )?;
13160 kvl.len += t;
13161 }
13162 kv_f32 = Some((k, v));
13163 }
13164 let kvl_idx = share.unwrap_or(il);
13167 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13168 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13170 let mut attn = e.uninit(t * nh * hd)?;
13171 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13183 if let Some((kf, vf)) = &kv_f32 {
13184 if hd == 256 && t <= win {
13185 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13186 return Ok(e.matmul(&fa.wo, &attn, t)?);
13187 }
13188 if hd == 256 && swa && t > win {
13189 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13190 return Ok(e.matmul(&fa.wo, &attn, t)?);
13191 }
13192 if hd == 512 && !swa {
13193 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13194 return Ok(e.matmul(&fa.wo, &attn, t)?);
13195 }
13196 } else if share.is_some() {
13197 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13198 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13199 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13200 if hd == 256 && (!swa || t <= win) {
13201 e.fa_prefill_view(
13203 &q,
13204 &k_view,
13205 &v_view,
13206 &mut attn,
13207 hd,
13208 nh,
13209 nkv,
13210 t,
13211 t,
13212 scale,
13213 true,
13214 kvl.k_tok_bytes,
13215 kvl.v_tok_bytes,
13216 g,
13217 )?;
13218 return Ok(e.matmul(&fa.wo, &attn, t)?);
13219 }
13220 let kv_dim = nkv * hd;
13223 let mut kf = e.uninit(t * kv_dim)?;
13224 let mut vf = e.uninit(t * kv_dim)?;
13225 e.fa_dequant_kv_view_f32(
13226 &k_view,
13227 &v_view,
13228 &mut kf,
13229 &mut vf,
13230 kv_dim,
13231 kv_dim,
13232 t,
13233 kvl.k_tok_bytes,
13234 kvl.v_tok_bytes,
13235 g,
13236 )?;
13237 if hd == 512 {
13238 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13239 } else {
13240 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13241 }
13242 return Ok(e.matmul(&fa.wo, &attn, t)?);
13243 }
13244 }
13245 if let Some(bucket) = dc_bucket {
13246 assert!(t == 1);
13251 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13257 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13258 } else {
13259 bucket
13260 };
13261 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13262 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13263 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13264 if crate::Engine::wpf_level() >= 1 {
13272 e.prefetch_weight_l2(&fa.wo)?;
13273 }
13274 if e.uses_q8_1_fast(&fa.wo) {
13277 let mut oq = e.alloc_i8_uninit(nh * hd)?;
13278 let mut od = e.zeros(nh * hd / 32)?;
13279 e.fa_decode_dc_q8(
13280 &q,
13281 &k_view,
13282 &v_view,
13283 &mut attn,
13284 hd,
13285 nh,
13286 nkv,
13287 &kvl.len_d,
13288 bucket,
13289 scale,
13290 kvl.k_tok_bytes,
13291 kvl.v_tok_bytes,
13292 g,
13293 Some((&mut oq, &mut od)),
13294 )?;
13295 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13296 }
13297 e.fa_decode_dc(
13298 &q,
13299 &k_view,
13300 &v_view,
13301 &mut attn,
13302 hd,
13303 nh,
13304 nkv,
13305 &kvl.len_d,
13306 bucket,
13307 scale,
13308 kvl.k_tok_bytes,
13309 kvl.v_tok_bytes,
13310 g,
13311 )?;
13312 return Ok(e.matmul(&fa.wo, &attn, t)?);
13313 }
13314 for i in 0..t {
13315 let avail = base_len + i + 1;
13316 let (off_tok, t_kv) = if swa && avail > win {
13317 (avail - win, win)
13318 } else {
13319 (0, avail)
13320 };
13321 let k_view = e.view_u8_range(
13322 &kvl.k,
13323 off_tok * kvl.k_tok_bytes,
13324 (off_tok + t_kv) * kvl.k_tok_bytes,
13325 );
13326 let v_view = e.view_u8_range(
13327 &kvl.v,
13328 off_tok * kvl.v_tok_bytes,
13329 (off_tok + t_kv) * kvl.v_tok_bytes,
13330 );
13331 let qv = e.view(&q, t * nh * hd);
13332 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13333 let mut q_one = e.uninit(nh * hd)?;
13334 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13335 let mut a_one = e.uninit(nh * hd)?;
13336 e.fa_decode_kvmod(
13340 &q_one,
13341 &k_view,
13342 &v_view,
13343 &mut a_one,
13344 hd,
13345 nh,
13346 nkv,
13347 t_kv,
13348 scale,
13349 kvl.k_tok_bytes,
13350 kvl.v_tok_bytes,
13351 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13352 )?;
13353 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13354 }
13355 Ok(e.matmul(&fa.wo, &attn, t)?)
13356 }
13357
13358 fn gemma4_e4b_trunk(
13363 &self,
13364 e: &Engine,
13365 tokens: &[u32],
13366 pos0: usize,
13367 cache: &mut Cache,
13368 head_last: bool,
13369 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13370 let n_embd = self.cfg.n_embd as usize;
13371 let t = tokens.len();
13372 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13373 let pos_d = e.htod_i32(&pos)?;
13374 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13375 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13376 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13377 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13378 }
13379
13380 fn gemma4_e4b_trunk_core(
13384 &self,
13385 e: &Engine,
13386 x_in: CudaSlice<f32>,
13387 inp_pl: CudaSlice<f32>,
13388 pos_d: &CudaSlice<i32>,
13389 t: usize,
13390 cache: &mut Cache,
13391 dc_bucket: Option<usize>,
13392 cap_logits: bool,
13393 head_last: bool,
13394 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13395 let n_embd = self.cfg.n_embd as usize;
13396 let eps = self.cfg.rms_eps;
13397 let n_layer = self.layers.len();
13398 let mut x = x_in;
13399 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13400 let n_epl = aux_e4b.n_epl;
13401
13402 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13408 for il in 0..n_layer {
13409 let layer = &self.layers[il];
13410 let (hq, hdq) = match h_carry.take() {
13411 Some(p) => p,
13412 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13413 };
13414 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13415 let bits = layer.gemma4.as_ref().unwrap();
13418 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13419 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13430 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13431 e,
13432 layer,
13433 &o,
13434 &x,
13435 t,
13436 Some(layer.post_attn_norm.float_data()),
13437 fuse_exit,
13438 )?;
13439 let mut resid = e.uninit(t * n_embd)?;
13440 let g = if fuse_exit {
13446 let (rq, rd) = e.rms_pre_add_q8_1(
13448 &sn,
13449 bits.post_ffw_norm.float_data(),
13450 &attn_out,
13451 &mut resid,
13452 n_embd,
13453 t,
13454 self.cfg.rms_eps,
13455 )?;
13456 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13457 } else {
13458 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13459 e.matmul(&e4b.inp_gate, &resid, t)?
13460 };
13461 let mut act = e.uninit(t * n_epl)?;
13462 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13463 let ipv = e.view(&inp_pl, n_epl * n_layer);
13464 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13465 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13466 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13467 } else {
13468 let mut inp_this = e.uninit(t * n_epl)?;
13469 e.copy_rows_strided(
13470 &inp_pl,
13471 &mut inp_this,
13472 n_epl,
13473 t,
13474 n_epl * n_layer,
13475 il * n_epl,
13476 )?;
13477 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13478 e.matmul(&e4b.proj, &act, t)?
13479 };
13480 let next_norm = if il + 1 < n_layer {
13483 self.layers[il + 1].attn_norm.float_data()
13484 } else {
13485 self.output_norm.float_data()
13486 };
13487 let mut xn = e.uninit(t * n_embd)?;
13488 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13489 &y,
13490 e4b.post_norm.float_data(),
13491 &resid,
13492 bits.layer_scale,
13493 next_norm,
13494 &mut xn,
13495 n_embd,
13496 t,
13497 eps,
13498 )?;
13499 h_carry = Some(pair);
13500 x = xn;
13501 }
13502 let (oq, odq) = h_carry.take().unwrap();
13506 let h0 = e.zeros(0)?;
13507 let hm = if head_last { 1 } else { t };
13508 let (hq, hd) = if head_last && t > 1 {
13509 let mut q1 = e.uninit_i8(n_embd)?;
13510 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
13511 let nb = n_embd / 32;
13512 let mut d1 = e.uninit(nb)?;
13513 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
13514 (q1, d1)
13515 } else {
13516 (oq, odq)
13517 };
13518 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
13519 if cap_logits {
13523 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
13524 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
13525 }
13526 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
13528 }
13529
13530 pub fn gemma4_e4b_decode_step_t_am_dev(
13537 &self,
13538 e: &Engine,
13539 tok_d: &CudaSlice<u32>,
13540 t: usize,
13541 pos0: usize,
13542 cache: &mut Cache,
13543 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13544 let n_embd = self.cfg.n_embd as usize;
13545 let eps = self.cfg.rms_eps;
13546 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13547 let pos_d = e.htod_i32(&pos)?;
13548 let embd_gpu = self
13549 .embd_gpu
13550 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
13551 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
13552 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
13553 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13554 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
13555 let (ld, xp) =
13556 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
13557 let n_vocab = self.output.out_features();
13560 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
13561 for i in 0..t {
13562 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
13563 }
13564 let mut hn = e.uninit(t * n_embd)?;
13565 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13566 cache.pos += t;
13567 Ok((vam, hn))
13568 }
13569
13570 pub(crate) fn gemma4_e4b_decode_step_t_h(
13573 &self,
13574 e: &Engine,
13575 tokens: &[u32],
13576 pos0: usize,
13577 cache: &mut Cache,
13578 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13579 let n_embd = self.cfg.n_embd as usize;
13580 let eps = self.cfg.rms_eps;
13581 let t = tokens.len();
13582 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
13583 let mut hn = e.uninit(t * n_embd)?;
13584 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13585 cache.pos += t;
13586 Ok((e.dtoh(&ld)?, hn))
13587 }
13588
13589 pub fn gemma4_e4b_decode_step_dcg(
13595 &self,
13596 e: &Engine,
13597 token_d: &mut CudaSlice<u32>,
13598 pos_d: &mut CudaSlice<i32>,
13599 embd_gpu: &CudaSlice<u8>,
13600 embd_qt: i32,
13601 embd_rb: usize,
13602 cache: &mut Cache,
13603 n_vocab: usize,
13604 bucket: usize,
13605 ) -> Result<(), Box<dyn std::error::Error>> {
13606 let n_embd = self.cfg.n_embd as usize;
13607 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13608 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13609 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13610 let (ld, _x) =
13611 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
13612 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
13613 e.inc_seqlen(pos_d)?;
13614 Ok(())
13615 }
13616
13617 #[allow(clippy::too_many_arguments)]
13625 pub fn gemma4_e4b_decode_step_dc(
13626 &self,
13627 e: &Engine,
13628 token_d: &CudaSlice<u32>,
13629 pos_d: &mut CudaSlice<i32>,
13630 embd_gpu: &CudaSlice<u8>,
13631 embd_qt: i32,
13632 embd_rb: usize,
13633 cache: &mut Cache,
13634 n_vocab: usize,
13635 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13636 let n_embd = self.cfg.n_embd as usize;
13637 let eps = self.cfg.rms_eps;
13638 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13639 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13640 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13641 let (ld, _x) =
13642 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
13643 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
13644 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
13645 e.inc_seqlen(pos_d)?;
13646 cache.pos += 1;
13647 let _ = eps;
13648 Ok(tok_out)
13649 }
13650
13651 pub(crate) fn gemma4_e4b_decode_step_h(
13654 &self,
13655 e: &Engine,
13656 token: u32,
13657 cache: &mut Cache,
13658 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13659 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
13660 let logits = e.dtoh(&ld)?;
13661 cache.pos += 1;
13662 Ok((logits, x))
13663 }
13664
13665 pub(crate) fn gemma4_e4b_prime(
13669 &self,
13670 e: &Engine,
13671 tokens: &[u32],
13672 cache: &mut Cache,
13673 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13674 if cache.pos != 0 {
13677 return Err(
13678 "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
13679 call or decode tokenwise"
13680 .into(),
13681 );
13682 }
13683 let n_embd = self.cfg.n_embd as usize;
13684 let t = tokens.len();
13685 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
13686 cache.pos += t;
13687 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
13689 let row = xv.slice((t - 1) * n_embd..t * n_embd);
13690 let mut h_seed = e.uninit(n_embd)?;
13691 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
13692 Ok((last, h_seed, x))
13693 }
13694
13695 pub(crate) fn gemma4_e4b_forward(
13697 &self,
13698 e: &Engine,
13699 tokens: &[u32],
13700 last_only: bool,
13701 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
13702 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
13703 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
13704 Ok(e.dtoh(&ld)?) }
13706}
13707
13708#[cfg(test)]
13709mod prime_chunk_schedule_tests {
13710 use super::{
13711 PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
13712 fixed_prime_chunk_ranges_for_ring,
13713 };
13714
13715 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
13716 ranges.iter().map(|(start, end)| end - start).collect()
13717 }
13718
13719 fn auto_chunk(t: usize) -> usize {
13720 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
13721 }
13722
13723 #[test]
13724 fn fixed_schedule_retains_measured_geometry() {
13725 assert_eq!(
13726 sizes(&fixed_prime_chunk_ranges(461, 128)),
13727 vec![128, 128, 128, 77]
13728 );
13729 assert_eq!(
13730 sizes(&fixed_prime_chunk_ranges(1833, 230)),
13731 vec![230, 230, 230, 230, 230, 230, 230, 223]
13732 );
13733 assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
13734 let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
13735 assert_eq!(capped, vec![4096, 4088, 16]);
13736 assert!(capped.iter().all(|&rows| rows <= 4096));
13737 assert_eq!(
13738 sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
13739 vec![4100],
13740 "flag-off schedule remains byte-for-byte the legacy monolithic tail",
13741 );
13742 }
13743
13744 #[test]
13745 fn dynamic_schedule_matches_registered_shapes() {
13746 let cases = [
13747 (461, vec![64, 141, 132, 124]),
13748 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
13749 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
13750 ];
13751 for (t, expected) in cases {
13752 let chunk = auto_chunk(t);
13753 let fixed = fixed_prime_chunk_ranges(t, chunk);
13754 assert_eq!(
13755 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
13756 expected
13757 );
13758 }
13759 }
13760
13761 #[test]
13762 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
13763 for t in 256..=8192 {
13764 let chunk = auto_chunk(t);
13765 let fixed = fixed_prime_chunk_ranges(t, chunk);
13766 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
13767 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
13768 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
13769 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
13770 for pair in dynamic.windows(2) {
13771 assert_eq!(pair[0].1, pair[1].0, "T={t}");
13772 }
13773 assert!(
13774 dynamic
13775 .iter()
13776 .all(|(start, end)| end - start >= PRIME_MIN_T),
13777 "T={t} sizes={:?}",
13778 sizes(&dynamic)
13779 );
13780 if dynamic.len() >= 3 {
13781 let chunk_sizes = sizes(&dynamic);
13782 assert!(
13783 chunk_sizes[0] < chunk_sizes[1],
13784 "T={t} sizes={chunk_sizes:?}"
13785 );
13786 assert!(
13787 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
13788 "T={t} sizes={chunk_sizes:?}"
13789 );
13790 }
13791 }
13792 }
13793}
13794
13795#[cfg(test)]
13796mod page_prefetch_tests {
13797 use super::{
13798 grouped_worker_prefetch_position, page_prefetch_positions,
13799 page_prefetch_window_from_values, worker_prefetch_positions,
13800 };
13801
13802 #[test]
13803 fn page_prefetch_window_keeps_existing_opt_in_default() {
13804 assert_eq!(page_prefetch_window_from_values(false, None), 0);
13805 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
13806 assert_eq!(page_prefetch_window_from_values(true, None), 1);
13807 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
13808 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
13809 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
13810 }
13811
13812 #[test]
13813 fn rolling_page_prefetch_advises_each_future_expert_once() {
13814 let advised: Vec<_> = (0..7)
13815 .flat_map(|position| page_prefetch_positions(position, 7, 3))
13816 .collect();
13817 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
13818
13819 let one_ahead: Vec<_> = (0..4)
13820 .flat_map(|position| page_prefetch_positions(position, 4, 1))
13821 .collect();
13822 assert_eq!(one_ahead, vec![1, 2, 3]);
13823 assert!(page_prefetch_positions(0, 4, 0).is_empty());
13824 }
13825
13826 #[test]
13827 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
13828 assert_eq!(grouped_worker_prefetch_position(0, None), None);
13829 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
13830 .chain(
13831 (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
13832 )
13833 .collect();
13834 assert_eq!(positions, vec![0, 1, 2, 3]);
13835 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
13836 }
13837
13838 #[test]
13839 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
13840 let queued: Vec<_> = (0..8)
13841 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
13842 .collect();
13843 assert_eq!(queued, (0..8).collect::<Vec<_>>());
13844
13845 let one_at_a_time: Vec<_> = (0..4)
13846 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
13847 .collect();
13848 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
13849 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
13850 }
13851}
13852
13853pub struct G4DcSlots {
13854 x: CudaSlice<f32>,
13855 xn: CudaSlice<f32>,
13856 cur: CudaSlice<f32>,
13857 hq: CudaSlice<i8>,
13858 hd_: CudaSlice<f32>,
13859 q0: CudaSlice<f32>,
13860 k0: CudaSlice<f32>,
13861 v0: CudaSlice<f32>,
13862 q: CudaSlice<f32>,
13863 k: CudaSlice<f32>,
13864 v: CudaSlice<f32>,
13865 attn: CudaSlice<f32>,
13866 o: CudaSlice<f32>,
13867 attn_out: CudaSlice<f32>,
13868 zsh: CudaSlice<f32>,
13869 zq: CudaSlice<i8>,
13870 zd: CudaSlice<f32>,
13871 gate: CudaSlice<f32>,
13872 up: CudaSlice<f32>,
13873 act: CudaSlice<f32>,
13874 actq: CudaSlice<i8>,
13875 actd: CudaSlice<f32>,
13876 f0: CudaSlice<f32>,
13877 sn: CudaSlice<f32>,
13878 hn: CudaSlice<f32>,
13879 logits: CudaSlice<f32>,
13880}