1use crate::Engine;
6use crate::cache::Cache;
7use cudarc::driver::CudaSlice;
8use memra_gguf::config::ModelConfig;
9
10pub struct PrimeSlabs {
13 pub t_cap: usize,
14 pub h: CudaSlice<f32>,
15 pub x1: CudaSlice<f32>,
16 pub z: CudaSlice<f32>,
17 pub act: CudaSlice<f32>,
18 pub xa: CudaSlice<f32>,
19 pub xb: CudaSlice<f32>,
20 pub h16: CudaSlice<u8>,
21 pub z16: CudaSlice<u8>,
22 pub gate: CudaSlice<f32>, pub up: CudaSlice<f32>, pub ffn_out: CudaSlice<f32>, pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
32 pub mixed: CudaSlice<f32>,
36 pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
37 pub seg_t: usize,
38}
39
40unsafe impl Send for PrimeSlabs {}
43
44fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
45 std::iter::repeat_with(|| None).take(n).collect()
46}
47
48struct PrimeCacheStages<'a> {
53 parent: &'a mut Cache,
54 cut: usize,
55 stage0: Cache,
56 stage1: Cache,
57}
58
59impl<'a> PrimeCacheStages<'a> {
60 fn new(parent: &'a mut Cache, cut: usize) -> Self {
61 let n = parent.kv.len();
62 assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
63 assert!(cut <= n, "PP-2 cache cut {cut} exceeds {n} layers");
64 let mut kv0 = empty_cache_layers(n);
65 let mut kv1 = empty_cache_layers(n);
66 let mut recur0 = empty_cache_layers(n);
67 let mut recur1 = empty_cache_layers(n);
68 for i in 0..cut {
69 kv0[i] = parent.kv[i].take();
70 recur0[i] = parent.recur[i].take();
71 }
72 for i in cut..n {
73 kv1[i] = parent.kv[i].take();
74 recur1[i] = parent.recur[i].take();
75 }
76 let pos = parent.pos;
77 let max_ctx = parent.max_ctx;
78 Self {
79 parent,
80 cut,
81 stage0: Cache {
82 kv: kv0,
83 recur: recur0,
84 pos,
85 max_ctx,
86 last_logits_dev: None,
87 dflash_taps: None,
88 },
89 stage1: Cache {
90 kv: kv1,
91 recur: recur1,
92 pos,
93 max_ctx,
94 last_logits_dev: None,
95 dflash_taps: None,
96 },
97 }
98 }
99
100 fn parts(&mut self) -> (&mut Cache, &mut Cache) {
101 (&mut self.stage0, &mut self.stage1)
102 }
103}
104
105impl Drop for PrimeCacheStages<'_> {
106 fn drop(&mut self) {
107 let n = self.parent.kv.len();
108 for i in 0..n {
109 let source = if i < self.cut {
110 &mut self.stage0
111 } else {
112 &mut self.stage1
113 };
114 debug_assert!(self.parent.kv[i].is_none());
115 debug_assert!(self.parent.recur[i].is_none());
116 self.parent.kv[i] = source.kv[i].take();
117 self.parent.recur[i] = source.recur[i].take();
118 }
119 self.parent.pos = self.stage0.pos.min(self.stage1.pos);
120 }
121}
122
123pub(crate) struct AttnPre {
125 pub q: cudarc::driver::CudaSlice<f32>,
126 pub k: cudarc::driver::CudaSlice<f32>,
127 pub v: cudarc::driver::CudaSlice<f32>,
128 pub gate: Option<cudarc::driver::CudaSlice<f32>>,
129}
130
131pub(crate) struct GdnPrep {
133 pub hk: usize,
134 pub q_l2: cudarc::driver::CudaSlice<f32>,
135 pub k_l2: cudarc::driver::CudaSlice<f32>,
136 pub v_g: cudarc::driver::CudaSlice<f32>,
137 pub beta: cudarc::driver::CudaSlice<f32>,
138 pub g_log: cudarc::driver::CudaSlice<f32>,
139 pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
140 pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
141}
142
143pub(crate) struct VerifyStreamScratch {
145 pub pos_d: CudaSlice<i32>,
146 pub row_ctrs: Vec<CudaSlice<i32>>,
147}
148use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
149
150struct MoeInputTraceWriter {
151 dir: std::path::PathBuf,
152 index: std::fs::File,
153 payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
154}
155
156static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
157 std::sync::OnceLock::new();
158
159fn gdec_enabled() -> bool {
162 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *E.get_or_init(|| {
164 std::env::var("MEMRA_MOE_GDEC")
165 .map(|v| v != "0")
166 .unwrap_or(true)
167 })
168}
169
170fn moe_slab_enabled() -> bool {
181 std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
182}
183
184fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
188 std::env::var("MEMRA_MOE_GROUPED")
189 .map(|value| value != "0")
190 .unwrap_or(false)
191}
192
193fn moe_prefetch_enabled() -> bool {
196 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
197 *E.get_or_init(|| {
198 std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
199 || crate::spill_pread::worker_enabled()
200 })
201}
202
203fn moe_page_prefetch_window() -> usize {
208 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
209 *W.get_or_init(|| {
210 page_prefetch_window_from_values(
211 std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
212 std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
213 .ok()
214 .as_deref(),
215 )
216 })
217}
218
219fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
220 if !enabled {
221 return 0;
222 }
223 raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
224}
225
226fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
230 if window == 0 || position >= len {
231 return len..len;
232 }
233 let (start, count) = if position == 0 {
234 (1, window)
235 } else {
236 (position.saturating_add(window), 1)
237 };
238 let start = start.min(len);
239 start..start.saturating_add(count).min(len)
240}
241
242fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
245 let position = current.map_or(0, |position| position.saturating_add(1));
246 (position < order_len).then_some(position)
247}
248
249fn worker_prefetch_window() -> usize {
254 static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
255 *WINDOW.get_or_init(|| {
256 let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
257 std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
258 .ok()
259 .and_then(|value| value.parse::<usize>().ok())
260 .unwrap_or(automatic.max(1))
261 })
262}
263
264fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
268 if window == 0 || position >= len {
269 return len..len;
270 }
271 let (start, count) = if position == 0 {
272 (0, window)
273 } else {
274 (position.saturating_add(window).saturating_sub(1), 1)
275 };
276 let start = start.min(len);
277 start..start.saturating_add(count).min(len)
278}
279
280fn moe_dev_enabled() -> bool {
285 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
286 *E.get_or_init(|| {
287 std::env::var("MEMRA_MOE_DEV")
288 .map(|v| v != "0")
289 .unwrap_or(true)
290 && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
291 })
292}
293
294fn sigmoid_router_enabled() -> bool {
297 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
298 *E.get_or_init(|| {
299 std::env::var("MEMRA_SIG_ROUTER")
300 .map(|v| v != "0")
301 .unwrap_or(true)
302 })
303}
304
305fn moe_q8_enabled() -> bool {
310 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
311 *E.get_or_init(|| {
312 std::env::var("MEMRA_MOE_Q8")
313 .map(|v| v != "0")
314 .unwrap_or(true)
315 })
316}
317
318fn expert_dp4a_supported(qt: i32) -> bool {
321 qt == crate::QT_Q4_0
322 || qt == crate::QT_IQ3_S
323 || qt == crate::QT_IQ4_XS
324 || qt == crate::QT_Q3_K
325 || qt == crate::QT_Q4_K
326 || qt == crate::QT_Q6_K
327}
328
329fn q8_expert_supported(qt: i32) -> bool {
330 static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
336 let kq = *KQ.get_or_init(|| {
337 std::env::var("MEMRA_MOE_Q8_KQ")
338 .map(|v| v != "0")
339 .unwrap_or(true)
340 });
341 let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
348 .map(|v| v != "0")
349 .unwrap_or(true);
350 qt == crate::QT_IQ3_S
351 || qt == crate::QT_IQ4_XS
352 || (nvfp4_q8 && qt == crate::QT_NVFP4)
353 || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
354}
355
356fn q8_expert_dec_supported(qt: i32) -> bool {
359 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
360}
361
362fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
368 match qt {
369 crate::QT_Q4_0 => in_f % 32 == 0,
370 crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
371 in_f % 256 == 0
372 }
373 _ => false,
374 }
375}
376
377fn moe_prewarm_enabled() -> bool {
380 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
381 *E.get_or_init(|| {
382 std::env::var("MEMRA_MOE_PREWARM")
383 .map(|v| v != "0")
384 .unwrap_or(true)
385 })
386}
387
388fn cpu_expert_profile_admit_enabled() -> bool {
392 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
393 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
394}
395
396pub const PRIME_MIN_T: usize = 16;
400const PRIME_PIPE_MICROBATCHES: usize = 8;
401const PRIME_PIPE_MIN_CHUNK: usize = 128;
402const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
403const PRIME_PIPE_LINEAR_WORK: usize = 8;
404
405fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
406 crate::pp::prime_pp_on()
407 && !crate::pp::pp2_streams_off()
408 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
409}
410
411pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
415 if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
416 let parsed = value
417 .parse::<usize>()
418 .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
419 return if crate::cache::swa_ring_on() {
420 if parsed == 0 {
421 crate::cache::PRIME_CHUNK_MAX_TOKENS
422 } else {
423 parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
424 }
425 } else {
426 parsed
427 };
428 }
429 let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
430 if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
431 chunk.min(
432 t.div_ceil(PRIME_PIPE_MICROBATCHES)
433 .max(PRIME_PIPE_MIN_CHUNK),
434 )
435 } else {
436 chunk
437 }
438}
439
440fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
441 fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
442}
443
444fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
445 if chunk == 0 || t <= chunk {
446 return vec![(0, t)];
447 }
448 let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
449 let mut start = 0usize;
450 while start < t {
451 let mut end = (start + chunk).min(t);
452 if t - end > 0 && t - end < PRIME_MIN_T {
453 if ring_on {
454 let shifted = t - PRIME_MIN_T;
455 end = if shifted > start { shifted } else { t };
456 } else {
457 end = t;
458 }
459 }
460 ranges.push((start, end));
461 start = end;
462 }
463 ranges
464}
465
466fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
467 let prefix = prefix as u128;
468 prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
469}
470
471fn dynamic_prime_chunk_ranges(
472 t: usize,
473 fixed_chunk: usize,
474 fixed: &[(usize, usize)],
475) -> Vec<(usize, usize)> {
476 let n = fixed.len();
477 if n < 3 {
478 return fixed.to_vec();
479 }
480
481 let max_first = t - (n - 1) * PRIME_MIN_T;
482 let first = fixed_chunk
483 .div_ceil(2)
484 .max(PRIME_PIPE_EDGE_MIN_CHUNK)
485 .min(max_first);
486 let mut ranges = Vec::with_capacity(n);
487 ranges.push((0, first));
488
489 let first_work = prime_chunk_work(first, t);
490 let work_span = prime_chunk_work(t, t) - first_work;
491 let denominator = (n - 1) as u128;
492 let mut previous = first;
493 for boundary in 1..n - 1 {
494 let target = first_work * denominator + work_span * (boundary as u128);
495 let remaining = n - 1 - boundary;
496 let mut low = previous + PRIME_MIN_T;
497 let mut high = t - remaining * PRIME_MIN_T;
498 while low < high {
499 let mid = low + (high - low) / 2;
500 if prime_chunk_work(mid, t) * denominator >= target {
501 high = mid;
502 } else {
503 low = mid + 1;
504 }
505 }
506 ranges.push((previous, low));
507 previous = low;
508 }
509 ranges.push((previous, t));
510 ranges
511}
512
513pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
517 let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
518 let chunk = prime_chunk_tokens(t, n_layers);
519 let fixed = fixed_prime_chunk_ranges(t, chunk);
520 let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
521 Ok(value) => value == "dynamic",
522 Err(_) => true,
523 };
524 if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
525 fixed
526 } else {
527 dynamic_prime_chunk_ranges(t, chunk, &fixed)
528 }
529}
530
531impl HybridModel {
532 fn prime_trace_path() -> Option<&'static str> {
537 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
538 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
539 .as_deref()
540 }
541
542 pub fn forward(
544 &self,
545 e: &Engine,
546 tokens: &[u32],
547 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
548 if self.is_gemma4_e4b() {
549 return self.gemma4_e4b_forward(e, tokens, false);
550 }
551 if self.cfg.gemma4.is_some() {
552 return self.gemma4_forward(e, tokens, false);
553 }
554 let cfg = &self.cfg;
555 let n_embd = cfg.n_embd as usize;
556 let t = tokens.len();
557 let eps = cfg.rms_eps;
558 let pos: Vec<i32> = (0..t as i32).collect();
559 let pos_d = e.htod_i32(&pos)?;
560
561 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
564 let mut h = e.uninit(t * n_embd)?;
566 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
567
568 let mixed = match &layer.mixer {
569 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
570 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
571 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
572 };
573
574 let mut x1 = e.uninit(t * n_embd)?;
576 e.add(&x, &mixed, &mut x1, t * n_embd)?;
577
578 let mut z = e.uninit(t * n_embd)?;
580 e.rms_norm(
581 &x1,
582 layer.post_attn_norm.float_data(),
583 &mut z,
584 n_embd,
585 t,
586 eps,
587 )?;
588 let ffn_out = match &layer.ffn {
589 crate::hybrid::Ffn::Dense {
590 ffn_gate,
591 ffn_up,
592 ffn_down,
593 } => {
594 let n_ff = ffn_gate.out_features();
595 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
596 let up = g2.pop().unwrap();
597 let gate = g2.pop().unwrap();
598 let mut act = e.uninit(t * n_ff)?;
599 Self::ffn_act_lim(
604 e,
605 &self.cfg,
606 &gate,
607 &up,
608 1.0,
609 1.0,
610 self.cfg.clamp_shexp_at(il as u32),
611 &mut act,
612 t * n_ff,
613 )?;
614 e.matmul(ffn_down, &act, t)?
615 }
616 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
617 };
618 let mut x2 = e.uninit(t * n_embd)?;
619 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
620 x = x2;
621 }
622
623 let mut hn = e.uninit(t * n_embd)?;
624 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
625 let logits = e.matmul(&self.output, &hn, t)?;
626 Ok(e.dtoh(&logits)?)
627 }
628
629 pub fn forward_last(
635 &self,
636 e: &Engine,
637 tokens: &[u32],
638 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
639 if self.cfg.gemma4.is_some() {
640 return self.gemma4_forward(e, tokens, true);
641 }
642 let cfg = &self.cfg;
643 let n_embd = cfg.n_embd as usize;
644 let t = tokens.len();
645 let eps = cfg.rms_eps;
646 let pos: Vec<i32> = (0..t as i32).collect();
647 let pos_d = e.htod_i32(&pos)?;
648
649 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
653 for (il, layer) in self.layers.iter().enumerate() {
654 let mut h = e.uninit(t * n_embd)?;
655 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
656 if probe {
657 e.stream().synchronize()?;
658 eprintln!("[probe] L{il} norm ok");
659 }
660 let mixed = match &layer.mixer {
661 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
662 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
663 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
664 };
665 if probe {
666 e.stream().synchronize()?;
667 eprintln!("[probe] L{il} mixer ok");
668 }
669 let mut x1 = e.uninit(t * n_embd)?;
670 e.add(&x, &mixed, &mut x1, t * n_embd)?;
671 let mut z = e.uninit(t * n_embd)?;
672 e.rms_norm(
673 &x1,
674 layer.post_attn_norm.float_data(),
675 &mut z,
676 n_embd,
677 t,
678 eps,
679 )?;
680 let ffn_out = match &layer.ffn {
681 crate::hybrid::Ffn::Dense {
682 ffn_gate,
683 ffn_up,
684 ffn_down,
685 } => {
686 let n_ff = ffn_gate.out_features();
687 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
688 let up = g2.pop().unwrap();
689 let gate = g2.pop().unwrap();
690 let mut act = e.uninit(t * n_ff)?;
691 Self::ffn_act_lim(
693 e,
694 &self.cfg,
695 &gate,
696 &up,
697 1.0,
698 1.0,
699 self.cfg.clamp_shexp_at(il as u32),
700 &mut act,
701 t * n_ff,
702 )?;
703 e.matmul(ffn_down, &act, t)?
704 }
705 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
706 };
707 if probe {
708 e.stream().synchronize()?;
709 eprintln!("[probe] L{il} ffn ok");
710 }
711 let mut x2 = e.uninit(t * n_embd)?;
712 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
713 x = x2;
714 }
715 let mut hn = e.uninit(t * n_embd)?;
717 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
718 let last = e.view(&hn, t * n_embd); let last_row = last.slice((t - 1) * n_embd..t * n_embd); let mut hlast = e.uninit(n_embd)?;
721 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
722 let logits = e.matmul(&self.output, &hlast, 1)?; Ok(e.dtoh(&logits)?)
724 }
725
726 pub fn prime_cache(
758 &self,
759 e: &Engine,
760 tokens: &[u32],
761 cache: &mut Cache,
762 queued_after: usize,
763 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
764 self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
765 }
766
767 pub fn prime_cache_overlaid(
773 &self,
774 e: &Engine,
775 tokens: &[u32],
776 cache: &mut Cache,
777 queued_after: usize,
778 overlay: Option<&crate::vision::EmbedOverlay>,
779 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
780 let n_embd = self.cfg.n_embd as usize;
781 let t = tokens.len();
782 assert!(
786 t >= PRIME_MIN_T,
787 "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
788 );
789 assert!(
790 cache.pos + t <= cache.max_ctx,
791 "prime_cache: prompt exceeds cache max_ctx"
792 );
793
794 if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
806 if self.is_gemma4_e4b() {
807 if overlay.is_some() {
808 return Err(
809 "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
810 );
811 }
812 return self.gemma4_e4b_prime(e, tokens, cache);
813 }
814 return self.gemma4_prime(e, tokens, cache, overlay);
819 }
820 let ranges = prime_chunk_ranges(t, self.layers.len());
821 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
857 let seq_end = if legacy_calllocal {
858 cache.pos + t
859 } else {
860 cache.pos + t + queued_after
861 };
862 if ranges.len() == 1 {
863 return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
864 }
865 if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
870 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
871 if overlay.is_some() {
872 return Err(
873 "vision embedding overlay + pipelined PP prime unsupported (v1); \
874 run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
875 .into(),
876 );
877 }
878 if crate::pp::pp_multi_stream_same_device() {
879 return Err(
880 "prime chunk pipeline refused with 2 stage streams on one device — \
881 that concurrent-stream placement remains quarantined by the deferred \
882 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
883 the serial split."
884 .into(),
885 );
886 }
887 return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
888 }
889 }
890 let mut hiddens = e.uninit(t * n_embd)?;
891 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
892 for &(start, end) in &ranges {
893 let (l, hs, x) =
894 self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
895 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
896 last = Some((l, hs));
897 }
898 let (logits, h_seed) = last.unwrap();
899 Ok((logits, h_seed, hiddens))
900 }
901
902 fn prime_cache_pp2_pipelined(
907 &self,
908 e: &Engine,
909 tokens: &[u32],
910 cache: &mut Cache,
911 seq_end: usize,
912 ranges: &[(usize, usize)],
913 fence: &[usize],
914 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
915 debug_assert_eq!(fence.len(), 3);
916 debug_assert!(ranges.len() >= 2);
917 let rt = crate::pp::PpNRt::get(e)?;
918 assert_eq!(
919 rt.n_stages(),
920 2,
921 "prime pipeline requires exactly two PP stages"
922 );
923 let n_embd = self.cfg.n_embd as usize;
924 let t = tokens.len();
925 let initial_base = cache.pos;
926 let caller_stream = e.stream();
927
928 rt.fence_stages_behind(&caller_stream)?;
933 let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
934 rt.prepare_overlap_slots(0, max_payload)?;
935
936 let mut hiddens = e.uninit(t * n_embd)?;
937 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
938 let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
939 let (cache0, cache1) = stage_caches.parts();
940 let (first_start, first_end) = ranges[0];
941 let mut slot = self.prime_pp2_stage0_enqueue(
942 e,
943 rt,
944 &tokens[first_start..first_end],
945 cache0,
946 seq_end,
947 fence,
948 initial_base + first_start,
949 true,
950 )?;
951 cache0.pos = initial_base + first_end;
952
953 for (i, &(start, end)) in ranges.iter().enumerate() {
954 let base = initial_base + start;
955 debug_assert_eq!(
956 cache1.pos, base,
957 "stage 1 must drain chunks in original position order"
958 );
959 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
960 let next_base = initial_base + next_start;
961 debug_assert_eq!(
962 cache0.pos, next_base,
963 "stage 0 must issue chunks in original position order"
964 );
965 let cache0_stage = &mut *cache0;
966 std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
971 let stage0 = scope.spawn(move || -> Result<usize, String> {
972 let next = self
973 .prime_pp2_stage0_enqueue(
974 e,
975 rt,
976 &tokens[next_start..next_end],
977 cache0_stage,
978 seq_end,
979 fence,
980 next_base,
981 true,
982 )
983 .map_err(|err| err.to_string())?;
984 cache0_stage.pos = initial_base + next_end;
985 Ok(next)
986 });
987 let x = self.prime_pp2_stage1_enqueue(
988 e,
989 rt,
990 slot,
991 end - start,
992 cache1,
993 seq_end,
994 fence,
995 base,
996 true,
997 )?;
998 let out = {
999 rt.bind_stage(1)?;
1000 let _st1 = rt.enter(1);
1001 let e1 = rt.engine(1, e);
1002 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1003 };
1004 let next = stage0
1005 .join()
1006 .map_err(|_| "pipeprime stage-0 host walker panicked")?
1007 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1008 Ok((out, Some(next)))
1009 })?
1010 } else {
1011 let x = self.prime_pp2_stage1_enqueue(
1012 e,
1013 rt,
1014 slot,
1015 end - start,
1016 cache1,
1017 seq_end,
1018 fence,
1019 base,
1020 true,
1021 )?;
1022 let out = {
1023 rt.bind_stage(1)?;
1024 let _st1 = rt.enter(1);
1025 let e1 = rt.engine(1, e);
1026 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1027 };
1028 (out, None)
1029 };
1030
1031 rt.publish_to(1, &caller_stream)?;
1032 e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1033 last = Some((out.0, out.1));
1034 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1035
1036 if let Some(next) = next_slot {
1037 rt.fence_stages_behind(&caller_stream)?;
1042 slot = next;
1043 }
1044 }
1045
1046 debug_assert_eq!(cache0.pos, initial_base + t);
1047 debug_assert_eq!(cache1.pos, initial_base + t);
1048 let (logits, h_seed) = last.unwrap();
1049 Ok((logits, h_seed, hiddens))
1050 }
1051
1052 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1059 if Engine::gdn_db_on()
1060 && Engine::gdn_chunked_enabled()
1061 && t >= 16
1062 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1063 && num_k * 2 == num_v
1064 {
1065 num_k
1066 } else {
1067 num_v
1068 }
1069 }
1070
1071 fn f16out_on(e: &Engine, t: usize) -> bool {
1076 crate::f16_ffi::pp_f16_enabled()
1077 && t >= 16
1078 && !e.verify_exact_on()
1079 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1080 }
1081
1082 pub fn prime_slabs_get(
1090 &self,
1091 e: &Engine,
1092 t: usize,
1093 n_embd: usize,
1094 n_ff_max: usize,
1095 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1096 let mut slabs = self.prime_slabs.lock().unwrap();
1097 let dev = e.ctx().ordinal();
1098 let need_new = match slabs.get(&dev) {
1099 None => true,
1100 Some(sl) => sl.lock().unwrap().t_cap < t,
1101 };
1102 if need_new {
1103 slabs.insert(
1104 dev,
1105 std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1106 t_cap: t,
1107 h: e.uninit(t * n_embd)?,
1108 x1: e.uninit(t * n_embd)?,
1109 z: e.uninit(t * n_embd)?,
1110 act: e.uninit(t * n_ff_max)?,
1111 xa: e.uninit(t * n_embd)?,
1112 xb: e.uninit(t * n_embd)?,
1113 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1114 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1115 gate: e.uninit(t * n_ff_max)?,
1116 up: e.uninit(t * n_ff_max)?,
1117 ffn_out: e.uninit(t * n_embd)?,
1118 seg_glue: Vec::new(),
1119 mixed: e.uninit(t * n_embd)?,
1120 seg_mid: Vec::new(),
1121 seg_t: 0,
1122 })),
1123 );
1124 }
1125 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1126 }
1127
1128 fn prime_chunk(
1132 &self,
1133 e: &Engine,
1134 tokens: &[u32],
1135 cache: &mut Cache,
1136 seq_end: usize,
1137 chunk_off: usize,
1138 overlay: Option<&crate::vision::EmbedOverlay>,
1139 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1140 if crate::pp::pp_host_bounce_active()
1141 && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1142 {
1143 return Err(
1144 "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1145 has no active prime stage split and would peer-read remote weights; keep \
1146 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1147 .into(),
1148 );
1149 }
1150 if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1159 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1160 if overlay.is_some() {
1161 return Err("vision embedding overlay + PP prime unsupported (v1); \
1162 run single-device or MEMRA_PRIME_PP=0"
1163 .into());
1164 }
1165 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1166 }
1167 }
1168 if crate::pp::pp_host_bounce_active() {
1169 return Err(
1170 "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1171 refusing an unsplit remote-weight walk"
1172 .into(),
1173 );
1174 }
1175 let t = tokens.len();
1176 let base = cache.pos;
1177 debug_assert!(
1178 seq_end >= base + t,
1179 "prime_chunk: seq_end must cover this chunk"
1180 );
1181 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1182 let pos_d = e.htod_i32(&pos)?;
1183
1184 let mut x_embed = self.embed(e, tokens)?; if let Some(ov) = overlay {
1186 let n_embd = self.cfg.n_embd as usize;
1190 for &(pos, row_off, n_rows) in &ov.spans {
1191 let lo = pos.max(chunk_off);
1192 let hi = (pos + n_rows).min(chunk_off + t);
1193 if lo < hi {
1194 let src_row = row_off + (lo - pos);
1195 let view = ov
1196 .rows
1197 .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
1198 e.copy_view_into(
1199 &mut x_embed,
1200 (lo - chunk_off) * n_embd,
1201 &view,
1202 (hi - lo) * n_embd,
1203 )?;
1204 }
1205 }
1206 }
1207 let x = self.prime_layers(
1208 e,
1209 x_embed,
1210 0,
1211 self.layers.len(),
1212 &pos_d,
1213 t,
1214 base,
1215 cache,
1216 seq_end,
1217 )?;
1218 self.prime_chunk_epilogue(e, x, t, cache)
1219 }
1220
1221 #[allow(clippy::too_many_arguments)]
1237 fn prime_layers(
1238 &self,
1239 e: &Engine,
1240 x_in: CudaSlice<f32>,
1241 lo: usize,
1242 hi: usize,
1243 pos_d: &CudaSlice<i32>,
1244 t: usize,
1245 base: usize,
1246 cache: &mut Cache,
1247 seq_end: usize,
1248 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1249 let cfg = &self.cfg;
1250 let n_embd = cfg.n_embd as usize;
1251 let eps = cfg.rms_eps;
1252 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1256 let n_ff_max = self
1262 .layers
1263 .iter()
1264 .map(|l| match &l.ffn {
1265 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1266 _ => n_embd,
1267 })
1268 .max()
1269 .unwrap_or(n_embd)
1270 .max(n_embd);
1271 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1272 let slab = if use_slabs {
1273 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1274 } else {
1275 None
1276 };
1277 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1278 let mut x_own; type SlabRefs<'a> = (
1280 &'a mut CudaSlice<f32>,
1281 &'a mut CudaSlice<f32>,
1282 &'a mut CudaSlice<f32>,
1283 &'a mut CudaSlice<f32>,
1284 &'a mut CudaSlice<u8>,
1285 &'a mut CudaSlice<u8>,
1286 &'a mut CudaSlice<f32>,
1287 &'a mut CudaSlice<f32>,
1288 &'a mut CudaSlice<f32>,
1289 );
1290 let (mut x_cur, mut x_nxt, sl): (
1291 &mut CudaSlice<f32>,
1292 &mut CudaSlice<f32>,
1293 Option<SlabRefs>,
1294 );
1295 let mut seg: Option<(
1296 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1297 &mut Vec<Option<cudarc::driver::CudaGraph>>,
1298 &mut CudaSlice<f32>,
1299 &mut usize,
1300 )> = None;
1301 let mut x_own2;
1302 match slab_guard.as_mut() {
1303 Some(g) => {
1304 let slabs = &mut **g;
1305 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1306 let PrimeSlabs {
1307 xa,
1308 xb,
1309 h,
1310 x1,
1311 z,
1312 act,
1313 h16,
1314 z16,
1315 gate,
1316 up,
1317 ffn_out,
1318 seg_glue,
1319 mixed,
1320 seg_mid,
1321 seg_t,
1322 ..
1323 } = slabs;
1324 x_cur = xa;
1325 x_nxt = xb;
1326 seg = Some((seg_glue, seg_mid, mixed, seg_t));
1327 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1328 }
1329 None => {
1330 x_own = x_in;
1331 x_own2 = e.uninit(t * n_embd)?;
1332 x_cur = &mut x_own;
1333 x_nxt = &mut x_own2;
1334 sl = None;
1335 }
1336 }
1337 let mut alloc_h;
1338 let mut alloc_x1;
1339 let mut alloc_z;
1340 let mut alloc_act;
1341 let mut alloc_h16;
1342 let mut alloc_z16;
1343 let mut alloc_gate;
1344 let mut alloc_up;
1345 let mut alloc_fo;
1346 let (h, x1, z, act): (
1347 &mut CudaSlice<f32>,
1348 &mut CudaSlice<f32>,
1349 &mut CudaSlice<f32>,
1350 &mut CudaSlice<f32>,
1351 );
1352 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1353 let (sl_gate, sl_up, sl_fo): (
1354 &mut CudaSlice<f32>,
1355 &mut CudaSlice<f32>,
1356 &mut CudaSlice<f32>,
1357 );
1358 match sl {
1359 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1360 h = a;
1361 x1 = b;
1362 z = c;
1363 act = d;
1364 h16 = e16;
1365 z16 = f16b;
1366 sl_gate = g;
1367 sl_up = u;
1368 sl_fo = fo;
1369 }
1370 None => {
1371 alloc_h = e.uninit(t * n_embd)?;
1372 alloc_x1 = e.uninit(t * n_embd)?;
1373 alloc_z = e.uninit(t * n_embd)?;
1374 alloc_act = e.uninit(t * n_ff_max)?;
1375 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1376 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1377 alloc_gate = e.uninit(t * n_ff_max)?;
1378 alloc_up = e.uninit(t * n_ff_max)?;
1379 alloc_fo = e.uninit(t * n_embd)?;
1380 h = &mut alloc_h;
1381 x1 = &mut alloc_x1;
1382 z = &mut alloc_z;
1383 act = &mut alloc_act;
1384 h16 = &mut alloc_h16;
1385 z16 = &mut alloc_z16;
1386 sl_gate = &mut alloc_gate;
1387 sl_up = &mut alloc_up;
1388 sl_fo = &mut alloc_fo;
1389 }
1390 }
1391 let n_layers = self.layers.len();
1396 let use_seg = f16fuse
1406 && seg.is_some()
1407 && self.cfg.step35.is_none()
1408 && lo == 0
1409 && hi == n_layers
1410 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1411 if let Some((sg, sm, _, st)) = seg.as_mut() {
1412 if **st != t {
1413 sg.clear();
1414 sg.extend((0..n_layers).map(|_| None));
1415 sm.clear();
1416 sm.extend((0..n_layers).map(|_| None));
1417 **st = t;
1418 }
1419 }
1420 {
1421 let layer_lo = &self.layers[lo];
1422 if f16fuse {
1423 e.rms_norm_f16out(
1424 x_cur,
1425 layer_lo.attn_norm.float_data(),
1426 h,
1427 h16,
1428 n_embd,
1429 t,
1430 eps,
1431 )?;
1432 } else {
1433 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1434 }
1435 }
1436 for il in lo..hi {
1437 let layer = &self.layers[il];
1438 let hx16 = if f16fuse { Some(&*h16) } else { None };
1439 if use_seg {
1440 let (pre, pre16, w_out) = match &layer.mixer {
1443 Mixer::Full(fa) => {
1444 let g3 = match hx16 {
1445 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1446 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1447 };
1448 let (pre, pre16) =
1449 self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1450 (pre, pre16, &fa.wo)
1451 }
1452 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1453 Mixer::Linear(la) => {
1454 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1455 let g4 = match hx16 {
1456 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1457 None => e.matmul_group(&ws, h, t)?,
1458 };
1459 let (pre, pre16) =
1460 self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1461 (pre, pre16, &la.ssm_out)
1462 }
1463 };
1464 {
1465 let (_, sm, mslab, _) = seg.as_mut().unwrap();
1466 let pre_n = pre.len() / t;
1467 let xh_pre = match pre16 {
1468 Some(x) => x,
1469 None => e.f16_act(&pre, t * pre_n, pre_n)?,
1470 };
1471 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1472 let y = e.matmul(w_out, &pre, t)?;
1473 e.copy_into(mslab, 0, &y, t * n_embd)?;
1474 }
1475 if sm[il].is_none() {
1476 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1477 let w_post = layer.post_attn_norm.float_data();
1478 e.stream().synchronize()?;
1479 e.stream()
1480 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1481 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1482 e.add(x_cur, mslab, x1, t * n_embd)?;
1483 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1484 Ok(())
1485 })();
1486 let g = e.stream().end_capture(
1487 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1488 r?;
1489 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1490 }
1491 sm[il].as_ref().unwrap().launch()?;
1492 }
1493 } else {
1494 let mixed = match &layer.mixer {
1495 Mixer::Full(fa) => {
1496 self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?
1497 }
1498 Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1499 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1500 };
1501 if f16fuse {
1502 e.add_rms_norm_f16out(
1505 x_cur,
1506 &mixed,
1507 layer.post_attn_norm.float_data(),
1508 x1,
1509 z,
1510 z16,
1511 n_embd,
1512 t,
1513 eps,
1514 )?;
1515 } else {
1516 e.add(x_cur, &mixed, x1, t * n_embd)?;
1517 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1518 }
1519 }
1520 let zx16 = if f16fuse { Some(&*z16) } else { None };
1521 match &layer.ffn {
1522 crate::hybrid::Ffn::Dense {
1523 ffn_gate,
1524 ffn_up,
1525 ffn_down,
1526 } => {
1527 let n_ff = ffn_gate.out_features();
1528 let mut into_ok = false;
1531 if let Some(xh) = zx16 {
1532 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1533 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1534 }
1535 if !into_ok {
1536 let mut g2 = match zx16 {
1537 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1538 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1539 };
1540 let up_y = g2.pop().unwrap();
1541 let gate_y = g2.pop().unwrap();
1542 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1543 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1544 }
1545 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1550 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1551 {
1552 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1553 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1554 Some(a16)
1555 } else {
1556 Self::ffn_act_lim(
1557 e,
1558 &self.cfg,
1559 sl_gate,
1560 sl_up,
1561 1.0,
1562 1.0,
1563 d_lim,
1564 act,
1565 t * n_ff,
1566 )?;
1567 None
1568 };
1569 let xh_act = match act16 {
1571 Some(x) => x,
1572 None => e.f16_act(act, t * n_ff, n_ff)?,
1573 };
1574 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1575 let y = e.matmul(ffn_down, &*act, t)?;
1576 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1577 }
1578 }
1579 crate::hybrid::Ffn::Moe(m) => {
1580 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1581 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1582 }
1583 }
1584 if use_seg && il + 1 < hi {
1585 let w_next = self.layers[il + 1].attn_norm.float_data();
1587 let (sg, _, _, _) = seg.as_mut().unwrap();
1588 if sg[il].is_none() {
1589 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1590 e.stream().synchronize()?;
1591 e.stream()
1592 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1593 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1594 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1595 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1596 Ok(())
1597 })();
1598 let g = e.stream().end_capture(
1599 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1600 );
1601 r?;
1602 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1603 }
1604 sg[il].as_ref().unwrap().launch()?;
1605 } else {
1606 if il + 1 < hi {
1607 let w_next = self.layers[il + 1].attn_norm.float_data();
1608 if f16fuse {
1609 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1610 } else {
1611 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1612 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1613 }
1614 } else {
1615 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1616 }
1617 }
1618 if let Some(path) = Self::prime_trace_path() {
1624 let row = (base + t - 1) as usize;
1625 let host = e.dtoh(x_nxt)?;
1626 let last = &host[(t - 1) * n_embd..t * n_embd];
1627 use std::io::Write as _;
1628 let mut f = std::fs::OpenOptions::new()
1629 .create(true)
1630 .append(true)
1631 .open(path)?;
1632 let mut h64: u64 = 0xcbf29ce484222325;
1633 for v in last {
1634 h64 ^= v.to_bits() as u64;
1635 h64 = h64.wrapping_mul(0x100000001b3);
1636 }
1637 writeln!(
1638 f,
1639 "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1640 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1641 last[0], last[1], last[2]
1642 )?;
1643 }
1644 std::mem::swap(&mut x_cur, &mut x_nxt);
1645 }
1646 let mut x = e.uninit(t * n_embd)?;
1648 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1649 drop(slab_guard);
1650 Ok(x)
1651 }
1652
1653 fn prime_chunk_epilogue(
1658 &self,
1659 e: &Engine,
1660 x: CudaSlice<f32>,
1661 t: usize,
1662 cache: &mut Cache,
1663 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1664 let n_embd = self.cfg.n_embd as usize;
1665 let eps = self.cfg.rms_eps;
1666 let mut h_seed = e.uninit(n_embd)?;
1670 if !crate::spec::spec_hpost() {
1671 e.copy_view_into(
1672 &mut h_seed,
1673 0,
1674 &x.slice((t - 1) * n_embd..t * n_embd),
1675 n_embd,
1676 )?;
1677 }
1678 let mut hn = e.uninit(t * n_embd)?;
1680 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1681 if crate::spec::spec_hpost() {
1682 e.copy_view_into(
1683 &mut h_seed,
1684 0,
1685 &hn.slice((t - 1) * n_embd..t * n_embd),
1686 n_embd,
1687 )?;
1688 }
1689 let last = e.view(&hn, t * n_embd);
1690 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1691 let mut hlast = e.uninit(n_embd)?;
1692 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1693 let logits = e.matmul(&self.output, &hlast, 1)?;
1694 cache.pos += t;
1695 Ok((
1698 e.dtoh(&logits)?,
1699 h_seed,
1700 if crate::spec::spec_hpost() { hn } else { x },
1701 ))
1702 }
1703
1704 fn prime_chunk_ppn(
1728 &self,
1729 e: &Engine,
1730 tokens: &[u32],
1731 cache: &mut Cache,
1732 seq_end: usize,
1733 fence: &[usize],
1734 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1735 let rt = crate::pp::PpNRt::get(e)?;
1736 let n_st = fence.len() - 1;
1737 assert_eq!(
1738 rt.n_stages(),
1739 n_st,
1740 "PpNRt stage count {} != fence stages {n_st}",
1741 rt.n_stages()
1742 );
1743 let n_embd = self.cfg.n_embd as usize;
1744 let t = tokens.len();
1745 let base = cache.pos;
1746 debug_assert!(
1747 seq_end >= base + t,
1748 "prime_chunk_ppn: seq_end must cover this chunk"
1749 );
1750 let payload = t * n_embd;
1751 let caller_stream = e.stream();
1755 rt.fence_stages_behind(&caller_stream)?;
1756
1757 if n_st == 2 {
1758 let slot =
1759 self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1760 let x =
1761 self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1762 let out = {
1763 rt.bind_stage(1)?;
1764 let _st1 = rt.enter(1);
1765 let e1 = rt.engine(1, e);
1766 self.prime_chunk_epilogue(e1, x, t, cache)?
1767 };
1768 rt.publish_to(1, &caller_stream)?;
1769 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1770 return Ok(out);
1771 }
1772
1773 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1774
1775 let mut slot = {
1777 let _st0 = rt.enter(0);
1778 let e0 = rt.engine(0, e);
1779 let pos_d = e0.htod_i32(&pos)?;
1780 let x = self.embed(e0, tokens)?;
1781 let x =
1782 self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1783 rt.tx(0, &x, payload)?
1784 };
1786
1787 for s in 1..n_st - 1 {
1789 let _st = rt.enter(s);
1790 let es = rt.engine(s, e);
1791 let pos_d = es.htod_i32(&pos)?;
1792 let x = rt.rx(s - 1, slot, payload)?;
1793 let x = self.prime_layers(
1794 es,
1795 x,
1796 fence[s],
1797 fence[s + 1],
1798 &pos_d,
1799 t,
1800 base,
1801 cache,
1802 seq_end,
1803 )?;
1804 slot = rt.tx(s, &x, payload)?;
1805 }
1806
1807 let _stl = rt.enter(n_st - 1);
1809 let el = rt.engine(n_st - 1, e);
1810 let pos_d = el.htod_i32(&pos)?;
1811 let x = rt.rx(n_st - 2, slot, payload)?;
1812 let x = self.prime_layers(
1813 el,
1814 x,
1815 fence[n_st - 1],
1816 fence[n_st],
1817 &pos_d,
1818 t,
1819 base,
1820 cache,
1821 seq_end,
1822 )?;
1823 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1824 rt.publish_to(n_st - 1, &caller_stream)?;
1830 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1831 Ok(out)
1832 }
1833
1834 fn prime_pp2_stage0_enqueue(
1835 &self,
1836 e: &Engine,
1837 rt: &crate::pp::PpNRt,
1838 tokens: &[u32],
1839 cache: &mut Cache,
1840 seq_end: usize,
1841 fence: &[usize],
1842 base: usize,
1843 pipelined: bool,
1844 ) -> Result<usize, Box<dyn std::error::Error>> {
1845 let t = tokens.len();
1846 let n_embd = self.cfg.n_embd as usize;
1847 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1848 rt.bind_stage(0)?;
1849 let _st0 = rt.enter(0);
1850 let e0 = rt.engine(0, e);
1851 let pos_d = e0.htod_i32(&pos)?;
1852 let x = self.embed(e0, tokens)?;
1853 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1854 let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1855 if pipelined {
1856 rt.tx_pipelined(0, &x, t * n_embd)
1857 } else {
1858 rt.tx(0, &x, t * n_embd)
1859 }
1860 }
1861
1862 fn prime_pp2_stage1_enqueue(
1863 &self,
1864 e: &Engine,
1865 rt: &crate::pp::PpNRt,
1866 slot: usize,
1867 t: usize,
1868 cache: &mut Cache,
1869 seq_end: usize,
1870 fence: &[usize],
1871 base: usize,
1872 pipelined: bool,
1873 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1874 let n_embd = self.cfg.n_embd as usize;
1875 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1876 rt.bind_stage(1)?;
1877 let _st1 = rt.enter(1);
1878 let e1 = rt.engine(1, e);
1879 let pos_d = e1.htod_i32(&pos)?;
1880 let x = rt.rx(0, slot, t * n_embd)?;
1881 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1882 self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
1883 }
1884
1885 pub fn prime_chunk_captured(
1901 &self,
1902 e: &Engine,
1903 x_in: &CudaSlice<f32>,
1904 pos_d: &CudaSlice<i32>,
1905 t: usize,
1906 cache: &mut Cache,
1907 len_d: &CudaSlice<i32>,
1908 logits_out: &mut CudaSlice<f32>,
1909 h_seed_out: &mut CudaSlice<f32>,
1910 ) -> Result<(), Box<dyn std::error::Error>> {
1911 let cfg = &self.cfg;
1912 let n_embd = cfg.n_embd as usize;
1913 let eps = cfg.rms_eps;
1914 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1915 let mut x = e.uninit(t * n_embd)?;
1916 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1917 for (il, layer) in self.layers.iter().enumerate() {
1918 let mut h = e.uninit(t * n_embd)?;
1919 let mut hx16: Option<CudaSlice<u8>> = None;
1920 if f16fuse {
1921 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1922 e.rms_norm_f16out(
1923 &x,
1924 layer.attn_norm.float_data(),
1925 &mut h,
1926 &mut b16,
1927 n_embd,
1928 t,
1929 eps,
1930 )?;
1931 hx16 = Some(b16);
1932 } else {
1933 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1934 }
1935 let mixed = match &layer.mixer {
1936 Mixer::Full(fa) => {
1940 self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
1941 }
1942 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1943 Mixer::Linear(la) => {
1944 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1945 let g4 = match hx16.as_ref() {
1946 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1947 None => e.matmul_group(&ws, &h, t)?,
1948 };
1949 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1950 }
1951 };
1952 let mut x1 = e.uninit(t * n_embd)?;
1953 e.add(&x, &mixed, &mut x1, t * n_embd)?;
1954 let mut z = e.uninit(t * n_embd)?;
1955 let mut zx16: Option<CudaSlice<u8>> = None;
1956 if f16fuse {
1957 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1958 e.rms_norm_f16out(
1959 &x1,
1960 layer.post_attn_norm.float_data(),
1961 &mut z,
1962 &mut b16,
1963 n_embd,
1964 t,
1965 eps,
1966 )?;
1967 zx16 = Some(b16);
1968 } else {
1969 e.rms_norm(
1970 &x1,
1971 layer.post_attn_norm.float_data(),
1972 &mut z,
1973 n_embd,
1974 t,
1975 eps,
1976 )?;
1977 }
1978 let ffn_out = match &layer.ffn {
1979 crate::hybrid::Ffn::Dense {
1980 ffn_gate,
1981 ffn_up,
1982 ffn_down,
1983 } => {
1984 let n_ff = ffn_gate.out_features();
1985 let mut g2 = match &zx16 {
1986 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1987 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1988 };
1989 let up = g2.pop().unwrap();
1990 let gate = g2.pop().unwrap();
1991 let mut act = e.uninit(t * n_ff)?;
1992 Self::ffn_act_lim(
1994 e,
1995 &self.cfg,
1996 &gate,
1997 &up,
1998 1.0,
1999 1.0,
2000 self.cfg.clamp_shexp_at(il as u32),
2001 &mut act,
2002 t * n_ff,
2003 )?;
2004 e.matmul(ffn_down, &act, t)?
2005 }
2006 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
2007 };
2008 let mut x2 = e.uninit(t * n_embd)?;
2009 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
2010 x = x2;
2011 }
2012 if !crate::spec::spec_hpost() {
2014 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
2015 }
2016 let mut hn = e.uninit(t * n_embd)?;
2017 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
2018 if crate::spec::spec_hpost() {
2019 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
2020 }
2021 let mut hlast = e.uninit(n_embd)?;
2022 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
2023 let logits = e.matmul(&self.output, &hlast, 1)?;
2024 let nv = logits.len();
2025 e.copy_into(logits_out, 0, &logits, nv)?;
2026 Ok(())
2027 }
2028
2029 fn step35_prime_batch_on() -> bool {
2030 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
2031 }
2032
2033 #[allow(clippy::too_many_arguments)]
2036 fn step35_prime_batch_layers(
2037 &self,
2038 e: &Engine,
2039 mut x: CudaSlice<f32>,
2040 lo: usize,
2041 hi: usize,
2042 ts: &[usize],
2043 offs: &[usize],
2044 pos_ds: &[CudaSlice<i32>],
2045 caches: &mut [&mut Cache],
2046 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2047 let cfg = &self.cfg;
2048 let n_embd = cfg.n_embd as usize;
2049 let eps = cfg.rms_eps;
2050 let b = ts.len();
2051 let total: usize = ts.iter().sum();
2052 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
2053
2054 let split = |e: &Engine,
2055 y: &CudaSlice<f32>,
2056 dim: usize|
2057 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2058 let mut out = Vec::with_capacity(b);
2059 for s in 0..b {
2060 let mut ys = e.uninit(ts[s] * dim)?;
2061 e.copy_view_into(
2062 &mut ys,
2063 0,
2064 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2065 ts[s] * dim,
2066 )?;
2067 out.push(ys);
2068 }
2069 Ok(out)
2070 };
2071
2072 for il in lo..hi {
2073 let layer = &self.layers[il];
2074 let Mixer::Full(fa) = &layer.mixer else {
2075 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2076 };
2077
2078 let mut h = e.uninit(total * n_embd)?;
2079 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2080 if f16fuse {
2081 e.rms_norm_f16out(
2082 &x,
2083 layer.attn_norm.float_data(),
2084 &mut h,
2085 &mut hx16,
2086 n_embd,
2087 total,
2088 eps,
2089 )?;
2090 } else {
2091 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2092 }
2093
2094 let gate_w = fa
2098 .attn_gate
2099 .as_ref()
2100 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2101 let mut g4 = if f16fuse {
2102 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2103 } else {
2104 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2105 };
2106 let gate = g4.pop().unwrap();
2107 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2108 (0..b).map(|_| Vec::with_capacity(3)).collect();
2109 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2110 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2111 parts[s].push(ys);
2112 }
2113 }
2114 let gates = split(e, &gate, gate_w.out_features())?;
2115 let geometry = self.step35_geom(il);
2116 let hd = geometry.head_dim_k as usize;
2117 let nh = geometry.n_head as usize;
2118 let mut ag_cat = e.uninit(total * nh * hd)?;
2119 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2120 let ag = self.step35_attn_pre_wo(
2121 e,
2122 fa,
2123 g3s,
2124 None,
2125 Some(&gate),
2126 &pos_ds[s],
2127 ts[s],
2128 Some(&mut *caches[s]),
2129 il,
2130 ts[s],
2131 )?;
2132 e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2133 }
2134 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2135
2136 let mut x1 = e.uninit(total * n_embd)?;
2137 let mut z = e.uninit(total * n_embd)?;
2138 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2139 if f16fuse {
2140 e.add_rms_norm_f16out(
2141 &x,
2142 &mixed,
2143 layer.post_attn_norm.float_data(),
2144 &mut x1,
2145 &mut z,
2146 &mut zx16,
2147 n_embd,
2148 total,
2149 eps,
2150 )?;
2151 } else {
2152 e.add(&x, &mixed, &mut x1, total * n_embd)?;
2153 e.rms_norm(
2154 &x1,
2155 layer.post_attn_norm.float_data(),
2156 &mut z,
2157 n_embd,
2158 total,
2159 eps,
2160 )?;
2161 }
2162
2163 let ffn_out = match &layer.ffn {
2164 crate::hybrid::Ffn::Dense {
2165 ffn_gate,
2166 ffn_up,
2167 ffn_down,
2168 } => {
2169 let n_ff = ffn_gate.out_features();
2170 let mut g2 = if f16fuse {
2171 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2172 } else {
2173 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2174 };
2175 let up = g2.pop().unwrap();
2176 let gate = g2.pop().unwrap();
2177 let mut act = e.uninit(total * n_ff)?;
2178 let d_lim = cfg.clamp_shexp_at(il as u32);
2179 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2180 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2181 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2182 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2183 Some(y) => y,
2184 None => e.matmul(ffn_down, &act, total)?,
2185 }
2186 } else {
2187 Self::ffn_act_lim(
2188 e,
2189 cfg,
2190 &gate,
2191 &up,
2192 1.0,
2193 1.0,
2194 d_lim,
2195 &mut act,
2196 total * n_ff,
2197 )?;
2198 e.matmul(ffn_down, &act, total)?
2199 }
2200 }
2201 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2202 };
2203 let mut x2 = e.uninit(total * n_embd)?;
2204 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2205 x = x2;
2206 }
2207 Ok(x)
2208 }
2209
2210 fn step35_prime_batch_epilogue(
2211 &self,
2212 e: &Engine,
2213 x: CudaSlice<f32>,
2214 ts: &[usize],
2215 offs: &[usize],
2216 caches: &mut [&mut Cache],
2217 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2218 let n_embd = self.cfg.n_embd as usize;
2219 let total: usize = ts.iter().sum();
2220 let mut hn = e.uninit(total * n_embd)?;
2221 e.rms_norm(
2222 &x,
2223 self.output_norm.float_data(),
2224 &mut hn,
2225 n_embd,
2226 total,
2227 self.cfg.rms_eps,
2228 )?;
2229
2230 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2231 let mut out = Vec::with_capacity(ts.len());
2232 for s in 0..ts.len() {
2233 let mut hidden = e.uninit(ts[s] * n_embd)?;
2234 e.copy_view_into(
2235 &mut hidden,
2236 0,
2237 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2238 ts[s] * n_embd,
2239 )?;
2240 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2241 let mut h_seed = e.uninit(n_embd)?;
2242 e.copy_view_into(
2243 &mut h_seed,
2244 0,
2245 &hidden_src.slice(last0..last0 + n_embd),
2246 n_embd,
2247 )?;
2248 let mut hlast = e.uninit(n_embd)?;
2250 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2251 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2252 caches[s].pos += ts[s];
2253 out.push((logits, h_seed, hidden));
2254 }
2255 Ok(out)
2256 }
2257
2258 fn step35_prime_cache_batch(
2259 &self,
2260 e: &Engine,
2261 prompts: &[&[u32]],
2262 caches: &mut [&mut Cache],
2263 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2264 if crate::pp::pp_host_bounce_active()
2265 && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2266 {
2267 return Err(
2268 "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2269 stage split; refusing an unsplit remote-weight walk"
2270 .into(),
2271 );
2272 }
2273 if !Self::step35_prime_batch_on() {
2274 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2275 }
2276 if caches.iter().any(|c| c.pos != 0) {
2277 return Err(
2278 "step35 batched prime currently supports complete fresh prompts only; \
2279 continuation/tick chunks require per-request queued_after"
2280 .into(),
2281 );
2282 }
2283
2284 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2285 for &t in &ts {
2286 assert!(
2287 t >= PRIME_MIN_T,
2288 "step35 batched prime needs T >= {PRIME_MIN_T}"
2289 );
2290 }
2291 for (s, c) in caches.iter().enumerate() {
2292 assert!(
2293 ts[s] <= c.max_ctx,
2294 "step35 batched prime exceeds cache max_ctx"
2295 );
2296 }
2297 let offs: Vec<usize> = ts
2298 .iter()
2299 .scan(0usize, |a, &t| {
2300 let o = *a;
2301 *a += t;
2302 Some(o)
2303 })
2304 .collect();
2305 let total: usize = ts.iter().sum();
2306 let payload = total * self.cfg.n_embd as usize;
2307 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2308 let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2309 let upload_positions =
2310 |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2311 positions
2312 .iter()
2313 .map(|p| e.htod_i32(p))
2314 .collect::<Result<_, _>>()
2315 };
2316
2317 static ONCE: std::sync::Once = std::sync::Once::new();
2318 ONCE.call_once(|| {
2319 eprintln!(
2320 "[step35-prime-batch] first concat prime: B={} tokens={total}",
2321 prompts.len()
2322 );
2323 });
2324
2325 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2326 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2327 let rt = crate::pp::PpNRt::get(e)?;
2328 let n_st = fence.len() - 1;
2329 assert_eq!(
2330 rt.n_stages(),
2331 n_st,
2332 "step35 prime batch stage count mismatch"
2333 );
2334 let caller_stream = e.stream();
2335 rt.fence_stages_behind(&caller_stream)?;
2336
2337 let mut slot = {
2338 let _st0 = rt.enter(0);
2339 let e0 = rt.engine(0, e);
2340 let pos_ds = upload_positions(e0)?;
2341 let x = self.embed(e0, &cat_tokens)?;
2342 let x = self.step35_prime_batch_layers(
2343 e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2344 )?;
2345 rt.tx(0, &x, payload)?
2346 };
2347 for s in 1..n_st - 1 {
2348 let _st = rt.enter(s);
2349 let es = rt.engine(s, e);
2350 let pos_ds = upload_positions(es)?;
2351 let x = rt.rx(s - 1, slot, payload)?;
2352 let x = self.step35_prime_batch_layers(
2353 es,
2354 x,
2355 fence[s],
2356 fence[s + 1],
2357 &ts,
2358 &offs,
2359 &pos_ds,
2360 caches,
2361 )?;
2362 slot = rt.tx(s, &x, payload)?;
2363 }
2364
2365 let _stl = rt.enter(n_st - 1);
2366 let el = rt.engine(n_st - 1, e);
2367 let pos_ds = upload_positions(el)?;
2368 let x = rt.rx(n_st - 2, slot, payload)?;
2369 let x = self.step35_prime_batch_layers(
2370 el,
2371 x,
2372 fence[n_st - 1],
2373 fence[n_st],
2374 &ts,
2375 &offs,
2376 &pos_ds,
2377 caches,
2378 )?;
2379 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2380 rt.publish_to(n_st - 1, &caller_stream)?;
2381 crate::pp::STEP35_PRIME_BATCH_SPLITS
2382 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2383 out
2384 } else {
2385 let pos_ds = upload_positions(e)?;
2386 let x = self.embed(e, &cat_tokens)?;
2387 let x = self.step35_prime_batch_layers(
2388 e,
2389 x,
2390 0,
2391 self.layers.len(),
2392 &ts,
2393 &offs,
2394 &pos_ds,
2395 caches,
2396 )?;
2397 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2398 }
2399 } else {
2400 let pos_ds = upload_positions(e)?;
2401 let x = self.embed(e, &cat_tokens)?;
2402 let x = self.step35_prime_batch_layers(
2403 e,
2404 x,
2405 0,
2406 self.layers.len(),
2407 &ts,
2408 &offs,
2409 &pos_ds,
2410 caches,
2411 )?;
2412 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2413 };
2414 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2415 Ok(out)
2416 }
2417
2418 pub fn prime_cache_batch(
2435 &self,
2436 e: &Engine,
2437 prompts: &[&[u32]],
2438 caches: &mut [&mut Cache],
2439 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2440 let cfg = &self.cfg;
2441 let n_embd = cfg.n_embd as usize;
2442 let eps = cfg.rms_eps;
2443 let b = prompts.len();
2444 assert!(b >= 1 && b == caches.len());
2445 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2446 let carried = pos0s.iter().any(|&p| p > 0);
2447 if cfg.gemma4.is_some() {
2453 return Err(
2454 "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2455 swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2456 .into(),
2457 );
2458 }
2459 if cfg.step35.is_some() {
2462 return self.step35_prime_cache_batch(e, prompts, caches);
2463 }
2464 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2465 for &t in &ts {
2466 assert!(
2467 t >= PRIME_MIN_T,
2468 "prime_cache_batch needs T >= {PRIME_MIN_T}"
2469 );
2470 }
2471 for (s, c) in caches.iter().enumerate() {
2472 assert!(
2473 c.pos + ts[s] <= c.max_ctx,
2474 "prime_cache_batch: prompt exceeds cache max_ctx"
2475 );
2476 }
2477 let total: usize = ts.iter().sum();
2478 let offs: Vec<usize> = ts
2479 .iter()
2480 .scan(0usize, |a, &t| {
2481 let o = *a;
2482 *a += t;
2483 Some(o)
2484 })
2485 .collect();
2486 let pos_ds: Vec<CudaSlice<i32>> = ts
2488 .iter()
2489 .zip(&pos0s)
2490 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2491 .collect::<Result<_, _>>()?;
2492 let split = |e: &Engine,
2494 y: &CudaSlice<f32>,
2495 dim: usize|
2496 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2497 let mut out = Vec::with_capacity(b);
2498 for s in 0..b {
2499 let mut ys = e.uninit(ts[s] * dim)?;
2500 e.copy_view_into(
2501 &mut ys,
2502 0,
2503 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2504 ts[s] * dim,
2505 )?;
2506 out.push(ys);
2507 }
2508 Ok(out)
2509 };
2510
2511 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2512 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
2514 let mut h = e.uninit(total * n_embd)?;
2515 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2516 e.rms_norm_f16out(
2517 &x,
2518 layer.attn_norm.float_data(),
2519 &mut h,
2520 &mut hx16,
2521 n_embd,
2522 total,
2523 eps,
2524 )?;
2525 let mut mixed = e.uninit(total * n_embd)?;
2527 match &layer.mixer {
2528 Mixer::Full(fa) => {
2529 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2530 let geometry = self.cfg.full_attention_geometry_at(il as u32);
2536 let (n_head, n_head_kv, head_dim) = (
2537 geometry.n_head as usize,
2538 geometry.n_head_kv as usize,
2539 geometry.head_dim_k as usize,
2540 );
2541 let fa_scale = geometry.attention_scale();
2542 let use_favl = !carried
2543 && (2..=8).contains(&b)
2544 && (head_dim == 256 || head_dim == 128)
2545 && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2546 && std::env::var("MEMRA_NOFA").is_err()
2547 && std::env::var("MEMRA_FA_FLOOR").is_err()
2548 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2549 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2550 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2551 if use_favl {
2552 let (qf_w, kf_w, vf_w) = (
2553 fa.wq.out_features(),
2554 fa.wk.out_features(),
2555 fa.wv.out_features(),
2556 );
2557 struct APre {
2558 q: CudaSlice<f32>,
2559 gate: Option<CudaSlice<f32>>,
2560 qn: CudaSlice<f32>,
2561 kn: CudaSlice<f32>,
2562 }
2563 let mut aps = Vec::with_capacity(b);
2564 for &t in ts.iter().take(b) {
2565 aps.push(APre {
2566 q: e.uninit(t * n_head * head_dim)?,
2567 gate: Some(e.uninit(t * n_head * head_dim)?),
2568 qn: e.uninit(t * n_head * head_dim)?,
2569 kn: e.uninit(t * n_head_kv * head_dim)?,
2570 });
2571 }
2572 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2573 let kvl = caches[0].kv[il].as_ref().unwrap();
2574 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2575 };
2576 let pargs: Vec<crate::AttnPreVl> = (0..b)
2577 .map(|s| {
2578 let (o, t) = (offs[s], ts[s]);
2579 let kvl = caches[s].kv[il].as_ref().unwrap();
2580 assert!(
2581 kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2582 "prime_cache_batch attn vl: fresh + capacity"
2583 );
2584 crate::AttnPreVl {
2585 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2586 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2587 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2588 q: e.addr_f32(&aps[s].q),
2589 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2590 qn: e.addr_f32(&aps[s].qn),
2591 kn: e.addr_f32(&aps[s].kn),
2592 kc: e.addr_u8(&kvl.k),
2593 vc: e.addr_u8(&kvl.v),
2594 t: t as i32,
2595 pad: 0,
2596 }
2597 })
2598 .collect();
2599 e.attn_pre_vl8(
2600 &pargs,
2601 fa.q_norm.float_data(),
2602 fa.k_norm.float_data(),
2603 head_dim,
2604 geometry.n_rot as usize,
2605 n_head,
2606 n_head_kv,
2607 self.cfg.rms_eps,
2608 geometry.rope_base,
2609 1.0,
2610 kv_dim_k,
2611 kv_dim_v,
2612 ktb,
2613 vtb,
2614 )?;
2615 for s in 0..b {
2616 let kvl = caches[s].kv[il].as_mut().unwrap();
2617 kvl.len += ts[s];
2618 let new_len = kvl.len as i32;
2619 e.set_i32_one(&mut kvl.len_d, new_len)?;
2620 }
2621 let mut attns = Vec::with_capacity(b);
2622 let mut mirrors = Vec::with_capacity(b);
2623 for &t in ts.iter().take(b) {
2624 attns.push(e.uninit(t * n_head * head_dim)?);
2625 let n = t * n_head_kv * head_dim;
2626 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2627 }
2628 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2631 Ok("0") => false,
2632 Ok("1") => true,
2633 _ => cfg!(memra_hopper_mma),
2634 };
2635 if fa3_on {
2636 let mut q16s = Vec::with_capacity(b);
2637 let mut v16s = Vec::with_capacity(b);
2638 for s in 0..b {
2639 let t = ts[s];
2640 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2641 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2642 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2643 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2644 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2645 e.f32_to_bf16_v(
2646 &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2647 &mut v16,
2648 t * n_head_kv * head_dim,
2649 )?;
2650 q16s.push(q16);
2651 v16s.push((k16, v16));
2652 }
2653 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2654 let mut kp = qp;
2655 let mut vp = qp;
2656 let mut op = [core::ptr::null_mut::<f32>(); 8];
2657 let mut tsv = [0i32; 8];
2658 for s in 0..b {
2659 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2660 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2661 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2662 op[s] = e.addr_f32(&attns[s]) as *mut f32;
2663 tsv[s] = ts[s] as i32;
2664 }
2665 let rc = unsafe {
2666 crate::fa3_vl_raw(
2667 qp.as_ptr(),
2668 kp.as_ptr(),
2669 vp.as_ptr(),
2670 op.as_ptr(),
2671 tsv.as_ptr(),
2672 b as i32,
2673 n_head as i32,
2674 n_head_kv as i32,
2675 head_dim as i32,
2676 fa_scale,
2677 e.stream().cu_stream() as *mut core::ffi::c_void,
2678 )
2679 };
2680 if rc != 0 {
2681 return Err(format!("memra_fa3_vl rc={rc}").into());
2682 }
2683 } else {
2684 let fargs: Vec<crate::FaSeqVl> = (0..b)
2685 .map(|s| crate::FaSeqVl {
2686 q: e.addr_f32(&aps[s].qn),
2687 k16: e.addr_u8(&mirrors[s].0),
2688 v16: e.addr_u8(&mirrors[s].1),
2689 o: e.addr_f32(&attns[s]),
2690 kf: e.addr_f32(&aps[s].kn),
2691 vf: e.addr_f32v(
2692 &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2693 ),
2694 t: ts[s] as i32,
2695 pad: 0,
2696 })
2697 .collect();
2698 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2699 }
2700 for (s, attn) in attns.into_iter().enumerate() {
2701 let (attn_g, ag16) = self.full_attn_prime_post_fa(
2702 e,
2703 attn,
2704 &aps[s].gate,
2705 ts[s],
2706 n_head,
2707 head_dim,
2708 )?;
2709 let mut done = false;
2710 if let Some(xh) = &ag16 {
2711 done = e.try_f16_gemm_pre_into_off(
2712 &fa.wo,
2713 xh,
2714 ts[s],
2715 &mut mixed,
2716 offs[s] * n_embd,
2717 )?;
2718 }
2719 if !done {
2720 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2721 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2722 }
2723 }
2724 } else {
2725 let mut parts: Vec<Vec<CudaSlice<f32>>> =
2726 (0..b).map(|_| Vec::new()).collect();
2727 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2728 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2729 parts[s].push(ys);
2730 }
2731 }
2732 for (s, g3s) in parts.into_iter().enumerate() {
2733 let (attn_g, ag16) = self.full_attn_prime_core_inner(
2735 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2736 )?;
2737 let mut done = false;
2738 if let Some(xh) = &ag16 {
2739 done = e.try_f16_gemm_pre_into_off(
2740 &fa.wo,
2741 xh,
2742 ts[s],
2743 &mut mixed,
2744 offs[s] * n_embd,
2745 )?;
2746 }
2747 if !done {
2748 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2749 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2750 }
2751 }
2752 }
2753 }
2754 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2755 Mixer::Linear(la) => {
2756 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2761 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2762 let outs =
2763 self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2764 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2765 let (o, t) = (offs[s], ts[s]);
2766 let mut done = false;
2767 if let Some(xh) = &gn16 {
2768 done = e.try_f16_gemm_pre_into_off(
2769 &la.ssm_out,
2770 xh,
2771 t,
2772 &mut mixed,
2773 o * n_embd,
2774 )?;
2775 }
2776 if !done {
2777 let m = e.matmul(&la.ssm_out, &gn, t)?;
2778 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2779 }
2780 }
2781 }
2782 }
2783 let mut x1 = e.uninit(total * n_embd)?;
2784 let mut z = e.uninit(total * n_embd)?;
2785 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2786 e.add_rms_norm_f16out(
2787 &x,
2788 &mixed,
2789 layer.post_attn_norm.float_data(),
2790 &mut x1,
2791 &mut z,
2792 &mut zx16,
2793 n_embd,
2794 total,
2795 eps,
2796 )?;
2797 let ffn_out = match &layer.ffn {
2798 crate::hybrid::Ffn::Dense {
2799 ffn_gate,
2800 ffn_up,
2801 ffn_down,
2802 } => {
2803 let n_ff = ffn_gate.out_features();
2804 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2805 let up = g2.pop().unwrap();
2806 let gate = g2.pop().unwrap();
2807 let mut act = e.uninit(total * n_ff)?;
2808 let d_lim = self.cfg.clamp_shexp_at(il as u32);
2812 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2813 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2814 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2815 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2816 Some(y) => y,
2817 None => e.matmul(ffn_down, &act, total)?,
2818 }
2819 } else {
2820 Self::ffn_act_lim(
2821 e,
2822 &self.cfg,
2823 &gate,
2824 &up,
2825 1.0,
2826 1.0,
2827 d_lim,
2828 &mut act,
2829 total * n_ff,
2830 )?;
2831 e.matmul(ffn_down, &act, total)?
2832 }
2833 }
2834 crate::hybrid::Ffn::Moe(m) => {
2835 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2836 }
2837 };
2838 let mut x2 = e.uninit(total * n_embd)?;
2839 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2840 x = x2;
2841 }
2842 let mut hn = e.uninit(total * n_embd)?;
2844 e.rms_norm(
2845 &x,
2846 self.output_norm.float_data(),
2847 &mut hn,
2848 n_embd,
2849 total,
2850 eps,
2851 )?;
2852 let mut hcat = e.uninit(b * n_embd)?;
2858 for s in 0..b {
2859 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2860 e.copy_view_into(
2861 &mut hcat,
2862 s * n_embd,
2863 &hn.slice(last0..last0 + n_embd),
2864 n_embd,
2865 )?;
2866 }
2867 let logits_cat = if b >= 2 {
2868 e.try_f16_gemm(&self.output, &hcat, b)?
2869 } else {
2870 None
2871 };
2872 let logits_host: Option<Vec<f32>> = match &logits_cat {
2873 Some(lc) => Some(e.dtoh(lc)?),
2874 None => None,
2875 };
2876 let n_vocab = self.output.out_features();
2877 let mut hidden_all = if crate::spec::spec_hpost() {
2878 split(e, &hn, n_embd)?
2879 } else {
2880 split(e, &x, n_embd)?
2881 };
2882 let mut out = Vec::with_capacity(b);
2883 for s in 0..b {
2884 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2885 let mut h_seed = e.uninit(n_embd)?;
2886 if !crate::spec::spec_hpost() {
2887 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2888 } else {
2889 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2890 }
2891 let logits = match &logits_host {
2892 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2893 None => {
2894 let mut hlast = e.uninit(n_embd)?;
2895 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2896 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2897 }
2898 };
2899 caches[s].pos += ts[s];
2900 out.push((logits, h_seed, hidden_all.remove(0)));
2901 }
2902 Ok(out)
2903 }
2904
2905 #[allow(clippy::too_many_arguments)]
2916 fn full_attn_prime(
2917 &self,
2918 e: &Engine,
2919 fa: &FullAttnLayer,
2920 h: &CudaSlice<f32>,
2921 hx: Option<&CudaSlice<u8>>,
2922 pos_d: &CudaSlice<i32>,
2923 t: usize,
2924 cache: &mut Cache,
2925 il: usize,
2926 seq_end: usize,
2927 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2928 if self.cfg.step35.is_some() {
2929 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2930 }
2931 let g3 = match hx {
2936 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2937 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2938 };
2939 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2940 }
2941
2942 fn full_attn_prime_core(
2946 &self,
2947 e: &Engine,
2948 fa: &FullAttnLayer,
2949 g3: Vec<CudaSlice<f32>>,
2950 pos_d: &CudaSlice<i32>,
2951 t: usize,
2952 cache: &mut Cache,
2953 il: usize,
2954 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2955 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2956 if let Some(xh) = &ag16 {
2957 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2958 return Ok(y);
2959 }
2960 }
2961 Ok(e.matmul(&fa.wo, &attn_g, t)?)
2962 }
2963
2964 fn full_attn_prime_core_inner(
2965 &self,
2966 e: &Engine,
2967 fa: &FullAttnLayer,
2968 g3: Vec<CudaSlice<f32>>,
2969 pos_d: &CudaSlice<i32>,
2970 t: usize,
2971 cache: &mut Cache,
2972 il: usize,
2973 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2974 let cfg = &self.cfg;
2975 let geometry = cfg.full_attention_geometry_at(il as u32);
2976 let n_head = geometry.n_head as usize;
2977 let n_head_kv = geometry.n_head_kv as usize;
2978 let head_dim = geometry.head_dim_k as usize;
2979 let scale = geometry.attention_scale();
2980 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2981 let AttnPre { q, k, v, gate } = pre;
2982 let mut attn = e.uninit(t * n_head * head_dim)?;
2983 self.full_attn_prime_fa_dispatch(
2984 e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
2985 )?;
2986 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2987 }
2988
2989 #[allow(clippy::type_complexity)]
2993 fn full_attn_prime_pre_fa(
2994 &self,
2995 e: &Engine,
2996 fa: &FullAttnLayer,
2997 mut g3: Vec<CudaSlice<f32>>,
2998 pos_d: &CudaSlice<i32>,
2999 t: usize,
3000 cache: &mut Cache,
3001 il: usize,
3002 ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
3003 let cfg = &self.cfg;
3004 let geometry = cfg.full_attention_geometry_at(il as u32);
3005 let n_head = geometry.n_head as usize;
3006 let n_head_kv = geometry.n_head_kv as usize;
3007 let head_dim = geometry.head_dim_k as usize;
3008 let eps = cfg.rms_eps;
3009
3010 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3014 let v = g3.pop().unwrap();
3015 let mut k = g3.pop().unwrap();
3016 let qf = g3.pop().unwrap();
3017 let (mut q, gate) = if gated {
3018 let mut q = e.uninit(t * n_head * head_dim)?;
3019 let mut gate = e.uninit(t * n_head * head_dim)?;
3020 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3021 (q, Some(gate))
3022 } else {
3023 (qf, None)
3024 };
3025
3026 let mut qn = e.uninit(t * n_head * head_dim)?;
3027 e.rms_norm(
3028 &q,
3029 fa.q_norm.float_data(),
3030 &mut qn,
3031 head_dim,
3032 n_head * t,
3033 eps,
3034 )?;
3035 q = qn;
3036 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3037 e.rms_norm(
3038 &k,
3039 fa.k_norm.float_data(),
3040 &mut kn,
3041 head_dim,
3042 n_head_kv * t,
3043 eps,
3044 )?;
3045 k = kn;
3046 let rope_dims = geometry.n_rot as usize;
3047 e.rope_neox(
3048 &mut q,
3049 pos_d,
3050 head_dim,
3051 rope_dims,
3052 n_head,
3053 t,
3054 geometry.rope_base,
3055 1.0,
3056 )?;
3057 e.rope_neox(
3058 &mut k,
3059 pos_d,
3060 head_dim,
3061 rope_dims,
3062 n_head_kv,
3063 t,
3064 geometry.rope_base,
3065 1.0,
3066 )?;
3067
3068 {
3071 let kvl = cache.kv[il].as_mut().unwrap();
3072 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3073 e.append_kv_quantized_rows(
3074 &k,
3075 &v,
3076 &mut kvl.k,
3077 &mut kvl.v,
3078 kvl.len,
3079 t,
3080 kvl.kv_dim_k,
3081 kvl.kv_dim_v,
3082 kvl.k_tok_bytes,
3083 kvl.v_tok_bytes,
3084 crate::Engine::kv_fp8_on(),
3085 )?;
3086 kvl.len += t;
3087 let new_len = kvl.len as i32;
3088 e.set_i32_one(&mut kvl.len_d, new_len)?;
3089 }
3090
3091 let base_len = {
3092 let kvl = cache.kv[il].as_ref().unwrap();
3093 kvl.len - t };
3095 Ok((AttnPre { q, k, v, gate }, base_len))
3096 }
3097
3098 #[allow(clippy::too_many_arguments)]
3105 fn full_attn_prime_fa_dispatch(
3106 &self,
3107 e: &Engine,
3108 q: &CudaSlice<f32>,
3109 k: &CudaSlice<f32>,
3110 v: &CudaSlice<f32>,
3111 attn: &mut CudaSlice<f32>,
3112 base_len: usize,
3113 t: usize,
3114 cache: &mut Cache,
3115 il: usize,
3116 head_dim: usize,
3117 n_head: usize,
3118 n_head_kv: usize,
3119 scale: f32,
3120 ) -> Result<(), Box<dyn std::error::Error>> {
3121 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3134 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3135 e.sdpa_naive(
3136 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3137 )?;
3138 } else {
3139 e.fa_prefill(
3140 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3141 )?;
3142 }
3143 return Ok(());
3144 }
3145 let kvl = cache.kv[il].as_ref().unwrap();
3146 let t_kv = base_len + t;
3147 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3148 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3149 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3153 e.sdpa_naive_quantized_view(
3154 q,
3155 &k_view,
3156 &v_view,
3157 attn,
3158 head_dim,
3159 n_head,
3160 n_head_kv,
3161 t,
3162 t_kv,
3163 scale,
3164 true,
3165 kvl.k_tok_bytes,
3166 kvl.v_tok_bytes,
3167 )?;
3168 return Ok(());
3169 }
3170 let deqw = std::env::var("MEMRA_PRIME_DEQW")
3178 .map(|v| v != "0")
3179 .unwrap_or(true);
3180 if deqw {
3181 e.fa_prefill_view_ws(
3182 q,
3183 &k_view,
3184 &v_view,
3185 attn,
3186 head_dim,
3187 n_head,
3188 n_head_kv,
3189 t,
3190 t_kv,
3191 scale,
3192 true,
3193 kvl.k_tok_bytes,
3194 kvl.v_tok_bytes,
3195 crate::Engine::kv_fp8_on(),
3196 )?;
3197 } else {
3198 e.fa_prefill_view(
3199 q,
3200 &k_view,
3201 &v_view,
3202 attn,
3203 head_dim,
3204 n_head,
3205 n_head_kv,
3206 t,
3207 t_kv,
3208 scale,
3209 true,
3210 kvl.k_tok_bytes,
3211 kvl.v_tok_bytes,
3212 crate::Engine::kv_fp8_on(),
3213 )?;
3214 }
3215 Ok(())
3216 }
3217
3218 fn full_attn_prime_post_fa(
3221 &self,
3222 e: &Engine,
3223 attn: CudaSlice<f32>,
3224 gate: &Option<CudaSlice<f32>>,
3225 t: usize,
3226 n_head: usize,
3227 head_dim: usize,
3228 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3229 let (attn_g, ag16) = match gate {
3230 Some(gate) => {
3231 let n = t * n_head * head_dim;
3232 let mut ag = e.uninit(n)?;
3233 if Self::f16out_on(e, t) {
3234 let mut a16 = e.alloc_u8_uninit(n * 2)?;
3235 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3236 (ag, Some(a16))
3237 } else {
3238 let mut gsig = e.uninit(n)?;
3239 e.sigmoid(gate, &mut gsig, n)?;
3240 e.mul(&attn, &gsig, &mut ag, n)?;
3241 (ag, None)
3242 }
3243 }
3244 None => (attn, None),
3245 };
3246 Ok((attn_g, ag16))
3247 }
3248
3249 fn linear_attn_prime(
3256 &self,
3257 e: &Engine,
3258 la: &LinearAttnLayer,
3259 h: &CudaSlice<f32>,
3260 hx: Option<&CudaSlice<u8>>,
3261 t: usize,
3262 cache: &mut Cache,
3263 il: usize,
3264 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3265 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3267 let g4 = match hx {
3268 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3269 None => e.matmul_group(&ws, h, t)?,
3270 };
3271 self.linear_attn_prime_core(e, la, g4, t, cache, il)
3272 }
3273
3274 fn linear_attn_prime_core(
3276 &self,
3277 e: &Engine,
3278 la: &LinearAttnLayer,
3279 mut g4: Vec<CudaSlice<f32>>,
3280 t: usize,
3281 cache: &mut Cache,
3282 il: usize,
3283 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3284 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3285 }
3286
3287 #[allow(clippy::too_many_arguments)]
3291 fn linear_attn_prime_core_pad_inner(
3292 &self,
3293 e: &Engine,
3294 la: &LinearAttnLayer,
3295 mut g4: Vec<CudaSlice<f32>>,
3296 t: usize,
3297 cache: &mut Cache,
3298 il: usize,
3299 pad_len: Option<&CudaSlice<i32>>,
3300 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3301 let ssm = self.cfg.ssm.as_ref().unwrap();
3303 let d_state = ssm.state_size as usize;
3304 let num_k = ssm.group_count as usize;
3305 let num_v = ssm.time_step_rank as usize;
3306 let key_dim = d_state * num_k;
3307 let value_dim = d_state * num_v;
3308 let conv_dim = key_dim * 2 + value_dim;
3309 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(
3314 e,
3315 la,
3316 &qkv_mixed.slice(0..t * conv_dim),
3317 &z.slice(0..t * value_dim),
3318 &beta_raw.slice(0..t * num_v),
3319 &alpha.slice(0..t * num_v),
3320 t,
3321 cache,
3322 il,
3323 pad_len,
3324 )
3325 }
3326
3327 #[allow(clippy::too_many_arguments)]
3330 fn linear_attn_gdn_prep(
3331 &self,
3332 e: &Engine,
3333 la: &LinearAttnLayer,
3334 qkv_mixed: &cudarc::driver::CudaView<f32>,
3335 beta_raw: &cudarc::driver::CudaView<f32>,
3336 alpha: &cudarc::driver::CudaView<f32>,
3337 t: usize,
3338 cache: &mut Cache,
3339 il: usize,
3340 pad_len: Option<&CudaSlice<i32>>,
3341 ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3342 let cfg = &self.cfg;
3343 let ssm = cfg.ssm.as_ref().unwrap();
3344 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;
3352 debug_assert!(
3353 t >= d_conv - 1,
3354 "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3355 );
3356
3357 let rl = cache.recur[il].as_mut().unwrap();
3362 let hk = Self::gdn_hk(e, t, num_v, num_k);
3363 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3364 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
3366 let mut k_g = e.uninit(d_state * hk * t)?;
3367 let mut v_g = e.uninit(d_state * num_v * t)?;
3368 if conv_fuse {
3369 e.ssm_conv1d_gdn_state_pad(
3370 qkv_mixed,
3371 &mut rl.conv_state,
3372 la.ssm_conv1d.float_data(),
3373 &mut q_g,
3374 &mut k_g,
3375 &mut v_g,
3376 conv_dim,
3377 t,
3378 d_conv,
3379 d_state,
3380 num_v,
3381 num_k,
3382 key_dim,
3383 hk,
3384 pad_len,
3385 )?;
3386 } else {
3387 let mut conv_out = e.uninit(conv_dim * t)?; e.ssm_conv1d_tm_state_pad_v(
3389 qkv_mixed,
3390 &mut rl.conv_state,
3391 la.ssm_conv1d.float_data(),
3392 &mut conv_out,
3393 conv_dim,
3394 t,
3395 d_conv,
3396 pad_len,
3397 )?;
3398 e.qkv_to_gdn_repack(
3399 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3400 )?;
3401 }
3402 let mut q_l2 = e.uninit(d_state * hk * t)?;
3403 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3407 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3408 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3409 Some(qb)
3410 } else {
3411 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3412 None
3413 };
3414 let mut k_l2 = e.uninit(d_state * hk * t)?;
3415 let kb16 = if Engine::l2_v2_on(d_state) {
3417 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3418 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3419 Some(kb)
3420 } else {
3421 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3422 None
3423 };
3424 let mut beta = e.uninit(t * num_v)?;
3425 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3426 let mut g_log = e.uninit(t * num_v)?;
3427 e.gdn_glog_v(
3428 alpha,
3429 la.ssm_dt.float_data(),
3430 la.ssm_a.float_data(),
3431 &mut g_log,
3432 num_v,
3433 t,
3434 )?;
3435 if let Some(len_d) = pad_len {
3436 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3437 }
3438 Ok(GdnPrep {
3439 hk,
3440 q_l2,
3441 k_l2,
3442 v_g,
3443 beta,
3444 g_log,
3445 kb16,
3446 qb16,
3447 })
3448 }
3449
3450 #[allow(clippy::too_many_arguments)]
3455 fn linear_attn_prime_core_batch(
3456 &self,
3457 e: &Engine,
3458 la: &LinearAttnLayer,
3459 g4: &[CudaSlice<f32>],
3460 offs: &[usize],
3461 ts: &[usize],
3462 caches: &mut [&mut Cache],
3463 il: usize,
3464 ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3465 let ssm = self.cfg.ssm.as_ref().unwrap();
3466 let d_state = ssm.state_size as usize;
3467 let num_k = ssm.group_count as usize;
3468 let num_v = ssm.time_step_rank as usize;
3469 let key_dim = d_state * num_k;
3470 let value_dim = d_state * num_v;
3471 let conv_dim = key_dim * 2 + value_dim;
3472 let eps = self.cfg.rms_eps;
3473 let scale = 1.0 / (d_state as f32).sqrt();
3474 let b = ts.len();
3475 let c = Engine::gdn_chunk_size();
3476 let carried = caches.iter().any(|c| c.pos > 0);
3479 let use_vl = !carried
3480 && (2..=8).contains(&b)
3481 && Engine::gdn_chunked_enabled()
3482 && ts.iter().all(|&t| t >= 16)
3483 && e.gdn_mma_enabled(c)
3484 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3485 if !use_vl {
3486 return (0..b)
3487 .map(|s| {
3488 let (o, t) = (offs[s], ts[s]);
3489 self.linear_attn_prime_core_pad_view(
3490 e,
3491 la,
3492 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3493 &g4[1].slice(o * value_dim..(o + t) * value_dim),
3494 &g4[2].slice(o * num_v..(o + t) * num_v),
3495 &g4[3].slice(o * num_v..(o + t) * num_v),
3496 t,
3497 caches[s],
3498 il,
3499 None,
3500 )
3501 })
3502 .collect();
3503 }
3504 struct SeqBufs {
3508 conv_out: CudaSlice<f32>,
3509 q_g: CudaSlice<f32>,
3510 k_g: CudaSlice<f32>,
3511 v_g: CudaSlice<f32>,
3512 q_l2: CudaSlice<f32>,
3513 k_l2: CudaSlice<f32>,
3514 beta: CudaSlice<f32>,
3515 g_log: CudaSlice<f32>,
3516 gn: CudaSlice<f32>,
3517 gn16: CudaSlice<u8>,
3518 }
3519 let d_conv = ssm.conv_kernel as usize;
3520 let f16o = Self::f16out_on(e, 16);
3521 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
3523 let mut pres = Vec::with_capacity(b);
3524 for &t in ts.iter().take(b) {
3525 sb.push(SeqBufs {
3526 conv_out: e.uninit(conv_dim * t)?,
3527 q_g: e.uninit(d_state * hk * t)?,
3528 k_g: e.uninit(d_state * hk * t)?,
3529 v_g: e.uninit(d_state * num_v * t)?,
3530 q_l2: e.uninit(d_state * hk * t)?,
3531 k_l2: e.uninit(d_state * hk * t)?,
3532 beta: e.uninit(t * num_v)?,
3533 g_log: e.uninit(t * num_v)?,
3534 gn: e.uninit(d_state * num_v * t)?,
3535 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3536 });
3537 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3538 }
3539 let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3540 .map(|s| {
3541 let (o, t) = (offs[s], ts[s]);
3542 let rl = caches[s].recur[il].as_ref().unwrap();
3543 crate::GdnPrepVl {
3544 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3545 conv_state: e.addr_f32(&rl.conv_state),
3546 conv_out: e.addr_f32(&sb[s].conv_out),
3547 q_g: e.addr_f32(&sb[s].q_g),
3548 k_g: e.addr_f32(&sb[s].k_g),
3549 v_g: e.addr_f32(&sb[s].v_g),
3550 q_l2: e.addr_f32(&sb[s].q_l2),
3551 k_l2: e.addr_f32(&sb[s].k_l2),
3552 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3553 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3554 beta: e.addr_f32(&sb[s].beta),
3555 g_log: e.addr_f32(&sb[s].g_log),
3556 o: e.addr_f32(&pres[s].o),
3557 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3558 gn: e.addr_f32(&sb[s].gn),
3559 gn16: e.addr_u8(&sb[s].gn16),
3560 kb16: if Engine::l2_v2_on(d_state) {
3561 e.addr_u8(&pres[s].kb16)
3562 } else {
3563 0
3564 },
3565 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3566 e.addr_u8(&pres[s].qb16)
3567 } else {
3568 0
3569 },
3570 t: t as i32,
3571 pad: 0,
3572 }
3573 })
3574 .collect();
3575 let args: Vec<crate::GdnSeqVl> = (0..b)
3576 .map(|s| {
3577 let rl = caches[s].recur[il].as_ref().unwrap();
3578 crate::GdnSeqVl {
3579 kb16: e.addr_u8(&pres[s].kb16),
3580 gcum: e.addr_f32(&pres[s].gcum),
3581 beta: e.addr_f32(&sb[s].beta),
3582 u: e.addr_f32(&pres[s].u),
3583 wb16: e.addr_u8(&pres[s].wb16),
3584 y: e.addr_u8(&pres[s].y16),
3585 ssnap: e.addr_u8(&pres[s].ssnap16),
3586 state_in: e.addr_f32(&rl.ssm_state),
3587 state_out: e.addr_f32(&rl.ssm_state_alt),
3588 q: e.addr_f32(&sb[s].q_l2),
3589 p: e.addr_f32(&pres[s].p),
3590 o: e.addr_f32(&pres[s].o),
3591 k: e.addr_f32(&sb[s].k_l2),
3592 v: e.addr_f32(&sb[s].v_g),
3593 g: e.addr_f32(&sb[s].g_log),
3594 a: e.addr_f32(&pres[s].a),
3595 w: e.addr_f32(&pres[s].w),
3596 t: ts[s] as i32,
3597 nc: pres[s].nc as i32,
3598 }
3599 })
3600 .collect();
3601 e.gdn_prep_vl8(
3602 &prep_args,
3603 la.ssm_conv1d.float_data(),
3604 la.ssm_dt.float_data(),
3605 la.ssm_a.float_data(),
3606 conv_dim,
3607 d_conv,
3608 d_state,
3609 num_v,
3610 num_k,
3611 key_dim,
3612 hk,
3613 eps,
3614 )?;
3615 if !Engine::l2_v2_on(d_state) {
3618 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3619 }
3620 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3622 if !Engine::l2_v2_on(d_state) {
3624 for s in 0..b {
3625 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3626 }
3627 }
3628 let mut wa = [crate::GdnWVl::default(); 8];
3629 for s in 0..b {
3630 wa[s] = crate::GdnWVl {
3631 qb16: e.addr_u8(&pres[s].qb16),
3632 pb16: e.addr_u8(&pres[s].pb16),
3633 };
3634 }
3635 Some(crate::GdnWVl8(wa))
3636 } else {
3637 None
3638 };
3639 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3640 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3641 if f16o {
3642 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3643 }
3644 let mut out = Vec::with_capacity(b);
3646 for (s, bufs) in sb.into_iter().enumerate() {
3647 let rl = caches[s].recur[il].as_mut().unwrap();
3648 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3649 let (o, t) = (offs[s], ts[s]);
3650 let SeqBufs { mut gn, gn16, .. } = bufs;
3651 if f16o {
3652 out.push((gn, Some(gn16)));
3653 } else {
3654 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3655 e.gated_rmsnorm_zv(
3656 &pres[s].o,
3657 la.ssm_norm.float_data(),
3658 &z_v,
3659 &mut gn,
3660 d_state,
3661 num_v * t,
3662 eps,
3663 )?;
3664 out.push((gn, None));
3665 }
3666 }
3667 Ok(out)
3668 }
3669
3670 #[allow(clippy::too_many_arguments)]
3674 fn linear_attn_prime_core_pad_view(
3675 &self,
3676 e: &Engine,
3677 la: &LinearAttnLayer,
3678 qkv_mixed: &cudarc::driver::CudaView<f32>,
3679 z: &cudarc::driver::CudaView<f32>,
3680 beta_raw: &cudarc::driver::CudaView<f32>,
3681 alpha: &cudarc::driver::CudaView<f32>,
3682 t: usize,
3683 cache: &mut Cache,
3684 il: usize,
3685 pad_len: Option<&CudaSlice<i32>>,
3686 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3687 let cfg = &self.cfg;
3688 let ssm = cfg.ssm.as_ref().unwrap();
3689 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
3692 let scale = 1.0 / (d_state as f32).sqrt();
3693
3694 let prep =
3695 self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3696
3697 let mut o = e.uninit(d_state * num_v * t)?;
3703 let rl = cache.recur[il].as_mut().unwrap();
3704 {
3705 let crate::cache::RecurLayer {
3706 ssm_state,
3707 ssm_state_alt,
3708 ..
3709 } = rl;
3710 e.gdn_scan_prefill(
3711 &prep.q_l2,
3712 &prep.k_l2,
3713 &prep.v_g,
3714 &prep.g_log,
3715 &prep.beta,
3716 prep.kb16.as_ref(),
3717 prep.qb16.as_ref(),
3718 ssm_state,
3719 ssm_state_alt,
3720 &mut o,
3721 num_v,
3722 t,
3723 scale,
3724 prep.hk,
3725 )?;
3726 }
3727 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3728
3729 let mut gn = e.uninit(d_state * num_v * t)?;
3732 let gn16 = if Self::f16out_on(e, t) {
3733 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3734 e.gated_rmsnorm_f16out_zv(
3735 &o,
3736 la.ssm_norm.float_data(),
3737 z,
3738 &mut gn,
3739 &mut g16,
3740 d_state,
3741 num_v * t,
3742 eps,
3743 )?;
3744 Some(g16)
3745 } else {
3746 e.gated_rmsnorm_zv(
3747 &o,
3748 la.ssm_norm.float_data(),
3749 z,
3750 &mut gn,
3751 d_state,
3752 num_v * t,
3753 eps,
3754 )?;
3755 None
3756 };
3757 Ok((gn, gn16))
3758 }
3759
3760 #[allow(clippy::too_many_arguments)]
3762 fn linear_attn_prime_core_pad(
3763 &self,
3764 e: &Engine,
3765 la: &LinearAttnLayer,
3766 g4: Vec<CudaSlice<f32>>,
3767 t: usize,
3768 cache: &mut Cache,
3769 il: usize,
3770 pad_len: Option<&CudaSlice<i32>>,
3771 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3772 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3773 if let Some(xh) = &gn16 {
3774 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3775 return Ok(y);
3776 }
3777 }
3778 Ok(e.matmul(&la.ssm_out, &gn, t)?)
3779 }
3780
3781 pub fn full_attn(
3786 &self,
3787 e: &Engine,
3788 fa: &FullAttnLayer,
3789 h: &CudaSlice<f32>,
3790 pos_d: &CudaSlice<i32>,
3791 t: usize,
3792 il: usize,
3793 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3794 if self.cfg.step35.is_some() {
3795 return self.step35_attn(e, fa, h, pos_d, t, il);
3796 }
3797 let cfg = &self.cfg;
3798 let _n_embd = cfg.n_embd as usize;
3799 let geometry = cfg.full_attention_geometry_at(il as u32);
3800 let n_head = geometry.n_head as usize;
3801 let n_head_kv = geometry.n_head_kv as usize;
3802 let head_dim = geometry.head_dim_k as usize;
3803 let eps = cfg.rms_eps;
3804 let scale = geometry.attention_scale();
3805
3806 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3809 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3811 let v = g3.pop().unwrap();
3812 let mut k = g3.pop().unwrap();
3813 let qf = g3.pop().unwrap();
3814 let (mut q, gate) = if gated {
3815 let mut q = e.uninit(t * n_head * head_dim)?;
3816 let mut gate = e.uninit(t * n_head * head_dim)?;
3817 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3818 (q, Some(gate))
3819 } else {
3820 (qf, None)
3821 };
3822
3823 let mut qn = e.uninit(t * n_head * head_dim)?;
3825 e.rms_norm(
3826 &q,
3827 fa.q_norm.float_data(),
3828 &mut qn,
3829 head_dim,
3830 n_head * t,
3831 eps,
3832 )?;
3833 q = qn;
3834 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3835 e.rms_norm(
3836 &k,
3837 fa.k_norm.float_data(),
3838 &mut kn,
3839 head_dim,
3840 n_head_kv * t,
3841 eps,
3842 )?;
3843 k = kn;
3844 let rope_dims = geometry.n_rot as usize;
3845 e.rope_neox(
3846 &mut q,
3847 pos_d,
3848 head_dim,
3849 rope_dims,
3850 n_head,
3851 t,
3852 geometry.rope_base,
3853 1.0,
3854 )?;
3855 e.rope_neox(
3856 &mut k,
3857 pos_d,
3858 head_dim,
3859 rope_dims,
3860 n_head_kv,
3861 t,
3862 geometry.rope_base,
3863 1.0,
3864 )?;
3865
3866 let mut attn = e.uninit(t * n_head * head_dim)?;
3868 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3871 e.sdpa_naive(
3873 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3874 )?;
3875 } else {
3876 e.fa_prefill(
3877 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3878 )?;
3879 }
3880
3881 let attn_g = match &gate {
3883 Some(gate) => {
3884 let mut gsig = e.uninit(t * n_head * head_dim)?;
3885 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3886 let mut ag = e.uninit(t * n_head * head_dim)?;
3887 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3888 ag
3889 }
3890 None => attn,
3891 };
3892
3893 let o = e.matmul(&fa.wo, &attn_g, t)?;
3895 Ok(o)
3896 }
3897
3898 pub fn linear_attn(
3900 &self,
3901 e: &Engine,
3902 la: &LinearAttnLayer,
3903 h: &CudaSlice<f32>,
3904 t: usize,
3905 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3906 let cfg = &self.cfg;
3907 let _n_embd = cfg.n_embd as usize;
3908 let ssm = cfg.ssm.as_ref().unwrap();
3909 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;
3914 let head_v = d_state;
3915 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;
3919 let scale = 1.0 / (d_state as f32).sqrt();
3920
3921 let mut g4 = e.matmul_group(
3924 &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
3925 h,
3926 t,
3927 )?;
3928 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);
3940 let mut q_g = e.uninit(d_state * num_v * t)?;
3941 let mut k_g = e.uninit(d_state * num_v * t)?;
3942 let mut v_g = e.uninit(d_state * num_v * t)?;
3943 e.ssm_conv1d_gdn(
3944 &qkv_mixed,
3945 la.ssm_conv1d.float_data(),
3946 &mut q_g,
3947 &mut k_g,
3948 &mut v_g,
3949 conv_dim,
3950 t,
3951 d_conv,
3952 d_state,
3953 num_v,
3954 num_k,
3955 key_dim,
3956 )?;
3957 let mut q_l2 = e.uninit(d_state * num_v * t)?;
3959 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3960 let mut k_l2 = e.uninit(d_state * num_v * t)?;
3961 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3962 let v_gd = v_g;
3963
3964 let mut beta = e.uninit(t * num_v)?;
3967 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3968 let mut g_log = e.uninit(t * num_v)?;
3970 e.gdn_glog(
3971 &alpha,
3972 la.ssm_dt.float_data(),
3973 la.ssm_a.float_data(),
3974 &mut g_log,
3975 num_v,
3976 t,
3977 )?;
3978
3979 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
3982 let mut o = e.uninit(d_state * num_v * t)?;
3983 e.gdn_scan_prefill(
3984 &q_l2,
3985 &k_l2,
3986 &v_gd,
3987 &g_log,
3988 &beta,
3989 None,
3990 None,
3991 &state_in,
3992 &mut state_out,
3993 &mut o,
3994 num_v,
3995 t,
3996 scale,
3997 num_v,
3998 )?;
3999
4000 let mut gn = e.uninit(d_state * num_v * t)?;
4005 e.gated_rmsnorm(
4006 &o,
4007 la.ssm_norm.float_data(),
4008 &z,
4009 &mut gn,
4010 d_state,
4011 num_v * t,
4012 eps,
4013 )?;
4014
4015 let out = e.matmul(&la.ssm_out, &gn, t)?;
4019 Ok(out)
4020 }
4021}
4022
4023impl HybridModel {
4024 pub fn moe_ffn_il(
4035 &self,
4036 e: &Engine,
4037 m: &MoeWeights,
4038 z: &CudaSlice<f32>,
4039 t: usize,
4040 il: u16,
4041 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4042 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
4043 }
4044
4045 pub fn moe_ffn_il_prefill(
4048 &self,
4049 e: &Engine,
4050 m: &MoeWeights,
4051 z: &CudaSlice<f32>,
4052 t: usize,
4053 il: u16,
4054 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4055 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
4056 }
4057
4058 pub fn moe_ffn_il_zq8(
4062 &self,
4063 e: &Engine,
4064 m: &MoeWeights,
4065 z: &CudaSlice<f32>,
4066 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4067 t: usize,
4068 il: u16,
4069 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4070 Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4071 }
4072
4073 pub(crate) fn moe_ffn(
4081 e: &Engine,
4082 m: &MoeWeights,
4083 z: &CudaSlice<f32>,
4084 t: usize,
4085 cfg: &ModelConfig,
4086 il: u16,
4087 max_block: usize,
4088 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4089 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4090 }
4091
4092 #[allow(clippy::too_many_arguments)]
4093 pub(crate) fn moe_ffn_inner(
4094 e: &Engine,
4095 m: &MoeWeights,
4096 z: &CudaSlice<f32>,
4097 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4098 t: usize,
4099 cfg: &ModelConfig,
4100 il: u16,
4101 max_block: usize,
4102 prefill: bool,
4103 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4104 let worker_io = crate::spill_pread::worker_enabled();
4105 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4106 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4107 e.with_moe_cache(max_block, |cache, _| {
4108 cache.begin_forward_epoch(il, t);
4109 if worker_io {
4110 cache.begin_worker_scope();
4111 }
4112 Ok(())
4113 })?;
4114 }
4115 if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4116 let moe = cfg.moe.as_ref().unwrap();
4117 let n_expert = moe.expert_count as usize;
4118 let n_used = moe.expert_used_count as usize;
4119 let sigmoid = cfg.sigmoid_router().unwrap();
4120 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4121 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4122 return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4123 }
4124 if t > 1 && moe_grouped_enabled(cfg, prefill) {
4127 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4128 if std::env::var("MEMRA_MOE_GATE").is_ok() {
4133 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4134 let g_host = e.dtoh(&grouped_out)?;
4135 let s_host = e.dtoh(&seq_out)?;
4136 let g_bytes: &[u8] = unsafe {
4137 std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4138 };
4139 let s_bytes: &[u8] = unsafe {
4140 std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4141 };
4142 if g_bytes == s_bytes {
4143 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4144 } else {
4145 let diffs = g_host
4146 .iter()
4147 .zip(s_host.iter())
4148 .enumerate()
4149 .filter(|(_, (a, b))| a != b)
4150 .count();
4151 let maxdiff = g_host
4152 .iter()
4153 .zip(s_host.iter())
4154 .map(|(a, b)| (a - b).abs())
4155 .fold(0.0f32, f32::max);
4156 panic!(
4157 "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4158 g_host.len()
4159 );
4160 }
4161 }
4162 return Ok(grouped_out);
4163 }
4164 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4165 }
4166
4167 fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4168 let Some(moe) = cfg.moe.as_ref() else {
4169 return false;
4170 };
4171 static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4174 let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4175 std::env::var("MEMRA_MOE_STATS").is_ok()
4176 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4177 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4178 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4179 || std::env::var("MEMRA_MOE_GATE").is_ok()
4180 });
4181 cfg.step35.is_some()
4182 && sigmoid_router_enabled()
4183 && moe_dev_enabled()
4184 && moe_slab_enabled()
4185 && !observation_mode
4186 && moe.expert_used_count <= 8
4187 && m.has_uniform_expert_layout()
4188 && m.gate_exps.macros.is_none()
4189 && m.up_exps.macros.is_none()
4190 && m.down_exps.macros.is_none()
4191 && !m.has_macros
4192 && moe_q8_enabled()
4193 && q8_expert_supported(m.gate_exps.qtype)
4194 && q8_expert_supported(m.up_exps.qtype)
4195 && q8_expert_supported(m.down_exps.qtype)
4196 && m.dev_exps
4197 .as_ref()
4198 .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4199 }
4200
4201 pub(crate) fn moe_ffn_sequential(
4203 e: &Engine,
4204 m: &MoeWeights,
4205 z: &CudaSlice<f32>,
4206 t: usize,
4207 cfg: &ModelConfig,
4208 il: u16,
4209 max_block: usize,
4210 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4211 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4212 }
4213
4214 fn moe_router_logits(
4218 e: &Engine,
4219 m: &MoeWeights,
4220 z: &CudaSlice<f32>,
4221 t: usize,
4222 cfg: &ModelConfig,
4223 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4224 if t < PRIME_MIN_T {
4225 if crate::router_kernel_on() {
4227 e.router_gemv(
4228 m.gate_inp.float_data(),
4229 z,
4230 cfg.n_embd as usize,
4231 m.gate_exps.n_expert,
4232 t,
4233 )
4234 } else {
4235 e.matmul_decode_exact(&m.gate_inp, z, t)
4236 }
4237 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4238 e.router_gemv(
4239 m.gate_inp.float_data(),
4240 z,
4241 cfg.n_embd as usize,
4242 m.gate_exps.n_expert,
4243 t,
4244 )
4245 } else {
4246 e.matmul(&m.gate_inp, z, t)
4247 }
4248 }
4249
4250 fn trace_moe_routes(
4254 il: u16,
4255 t: usize,
4256 sel_all: &[u32],
4257 weights: &[f32],
4258 ) -> Result<(), Box<dyn std::error::Error>> {
4259 use std::io::Write as _;
4260 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4261 let mut f = std::fs::OpenOptions::new()
4262 .create(true)
4263 .append(true)
4264 .open(path)?;
4265 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4266 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4267 }
4268 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4269 let mut f = std::fs::OpenOptions::new()
4270 .create(true)
4271 .append(true)
4272 .open(path)?;
4273 let pairs: Vec<String> = sel_all
4274 .iter()
4275 .zip(weights)
4276 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4277 .collect();
4278 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4279 }
4280 Ok(())
4281 }
4282
4283 #[allow(clippy::too_many_arguments)]
4284 fn trace_sigmoid_router_logits(
4285 e: &Engine,
4286 il: u16,
4287 t: usize,
4288 n_expert: usize,
4289 n_used: usize,
4290 logits: &CudaSlice<f32>,
4291 m: &MoeWeights,
4292 (scaling_factor, route_norm): (f32, bool),
4293 ) -> Result<(), Box<dyn std::error::Error>> {
4294 if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4295 return Ok(());
4296 }
4297 let logits = e.dtoh(logits)?;
4298 let active: Vec<u8> = m
4299 .active_experts
4300 .as_ref()
4301 .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4302 .unwrap_or_else(|| vec![1; n_expert]);
4303 let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4304 crate::sigrouter_contract::capture_served_logits(
4305 il as u32,
4306 t,
4307 n_expert,
4308 n_used,
4309 scaling_factor,
4310 route_norm,
4311 &active,
4312 &bias,
4313 &logits,
4314 )?;
4315 Ok(())
4316 }
4317
4318 fn trace_moe_input(
4323 e: &Engine,
4324 il: u16,
4325 t: usize,
4326 n_embd: usize,
4327 z: &CudaSlice<f32>,
4328 ) -> Result<(), Box<dyn std::error::Error>> {
4329 use std::io::Write as _;
4330 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4331 return Ok(());
4332 };
4333 let host = e.dtoh(z)?;
4334 if host.len() != t * n_embd {
4335 return Err(format!(
4336 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4337 host.len(),
4338 t,
4339 n_embd
4340 )
4341 .into());
4342 }
4343 let bytes = unsafe {
4344 std::slice::from_raw_parts(
4345 host.as_ptr().cast::<u8>(),
4346 host.len() * std::mem::size_of::<f32>(),
4347 )
4348 };
4349 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4350 let mut state = state
4351 .lock()
4352 .map_err(|_| "MoE input trace writer lock is poisoned")?;
4353 if state.is_none() {
4354 let dir = std::path::PathBuf::from(&dir);
4355 std::fs::create_dir_all(&dir)?;
4356 let index = std::fs::OpenOptions::new()
4357 .create(true)
4358 .append(true)
4359 .open(dir.join("index.jsonl"))?;
4360 *state = Some(MoeInputTraceWriter {
4361 dir,
4362 index,
4363 payloads: std::collections::HashMap::new(),
4364 });
4365 }
4366 let writer = state.as_mut().unwrap();
4367 if writer.dir != std::path::Path::new(&dir) {
4368 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4369 }
4370 let file_name = format!("layer-{il:03}.f32");
4371 if !writer.payloads.contains_key(&il) {
4372 let payload = std::fs::OpenOptions::new()
4373 .create(true)
4374 .append(true)
4375 .open(writer.dir.join(&file_name))?;
4376 let offset = payload.metadata()?.len();
4377 writer.payloads.insert(il, (payload, offset));
4378 }
4379 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4380 let row_offset = *offset;
4381 payload.write_all(bytes)?;
4382 *offset += bytes.len() as u64;
4383 writeln!(
4384 writer.index,
4385 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4386 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4387 \"payload_bytes\":{}}}",
4388 bytes.len()
4389 )?;
4390 Ok(())
4391 }
4392
4393 #[allow(clippy::too_many_arguments)]
4394 pub(crate) fn moe_ffn_sequential_zq8(
4395 e: &Engine,
4396 m: &MoeWeights,
4397 z: &CudaSlice<f32>,
4398 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4399 t: usize,
4400 cfg: &ModelConfig,
4401 il: u16,
4402 max_block: usize,
4403 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4404 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4405 let moe = cfg.moe.as_ref().unwrap();
4406 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);
4413 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4414 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);
4417
4418 let lim_exp = cfg.clamp_exp_at(il as u32);
4421 let lim_shexp = cfg.clamp_shexp_at(il as u32);
4422 let use_cache = Engine::moe_cache_enabled();
4423 let uniform_experts = m.has_uniform_expert_layout();
4424 let moe_q8 = uniform_experts
4425 && moe_q8_enabled()
4426 && q8_expert_supported(m.gate_exps.qtype)
4427 && q8_expert_supported(m.up_exps.qtype)
4428 && q8_expert_supported(m.down_exps.qtype);
4429 let cpu_expert_requested = crate::cpu_experts::configured();
4436 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4437 return Err(std::io::Error::other(
4438 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4439 )
4440 .into());
4441 }
4442 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4443 let freeze_cpu_residency = cpu_expert_requested
4449 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4450 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4451 .ok()
4452 .and_then(|value| value.parse::<usize>().ok())
4453 .is_some_and(|tokens| tokens > 0);
4454 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4455 e.freeze_moe_cache();
4456 }
4457 let cache_frozen = use_cache && e.moe_cache_frozen();
4458 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4459
4460 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4463 if let Some(sig) = cfg.sigmoid_router() {
4464 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4465 }
4466
4467 let no_exp_macros = m.gate_exps.macros.is_none()
4506 && m.up_exps.macros.is_none()
4507 && m.down_exps.macros.is_none();
4508 if cfg.sigmoid_router().is_none()
4512 && cfg.m3.is_none()
4513 && cfg.hy3.is_none()
4514 && !cfg.swiglu_clamped_at(il as u32)
4515 && no_exp_macros
4516 && t >= PRIME_MIN_T
4517 && m.dev_exps.is_some()
4518 && moe_q8_enabled()
4519 && q8_expert_supported(m.gate_exps.qtype)
4520 && q8_expert_supported(m.up_exps.qtype)
4521 && q8_expert_supported(m.down_exps.qtype)
4522 && std::env::var("MEMRA_MOE_PAIRS")
4523 .map(|v| v != "0")
4524 .unwrap_or(true)
4525 && std::env::var("MEMRA_MOE_STATS").is_err()
4526 {
4527 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4528 }
4529
4530 let dev_ok = uniform_experts
4548 && cfg.sigmoid_router().is_none()
4549 && cfg.m3.is_none()
4550 && cfg.hy3.is_none()
4551 && !cfg.swiglu_clamped_at(il as u32);
4552 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4556 || std::env::var("MEMRA_MOE_TRACE").is_ok()
4557 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4558 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4559 if dev_ok
4560 && t < PRIME_MIN_T
4561 && m.dev_exps.is_some()
4562 && n_used <= 8
4563 && moe_dev_enabled()
4564 && !observe_routes
4565 {
4566 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4567 }
4568 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4569 let row_ok = e.with_moe_cache(max_block, |c, eng| {
4570 if moe_prewarm_enabled() {
4571 c.prewarm_layer(il, m, eng)?;
4572 }
4573 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4574 })?;
4575 if row_ok {
4576 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4577 }
4578 }
4579
4580 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4582 if cpu_hybrid {
4583 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4584 e,
4585 &logits,
4586 z,
4587 t,
4588 n_expert,
4589 n_used,
4590 m.exp_probs_b.as_deref(),
4591 sig,
4592 m.active_experts.as_deref(),
4593 )?;
4594 (sel, w, Some(input))
4595 } else {
4596 let (sel, w) =
4597 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4598 (sel, w, None)
4599 }
4600 } else {
4601 let (sel, w) =
4602 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4603 (sel, w, None)
4604 };
4605 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4606
4607 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4611 Self::trace_moe_input(e, il, t, n_embd, z)?;
4612
4613 let worker_disk_prefetch =
4625 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4626 let promote_worker_h2d =
4627 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4628 if promote_worker_h2d {
4629 let mut selected_blocks = Vec::with_capacity(n_used * 3);
4630 for &ex in sel_all.iter().take(n_used) {
4631 let ex = ex as u16;
4632 selected_blocks.extend([
4633 BlockId::new(il, PROJ_GATE, ex),
4634 BlockId::new(il, PROJ_UP, ex),
4635 BlockId::new(il, PROJ_DOWN, ex),
4636 ]);
4637 }
4638 for &ex in sel_all.iter().take(n_used) {
4639 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4640 }
4641 e.with_moe_cache(max_block, |cache, eng| {
4642 cache.promote_worker_reads_at_safe_boundary(
4643 &selected_blocks,
4644 &selected_blocks,
4645 eng,
4646 )?;
4647 Ok(())
4648 })?;
4649 }
4650
4651 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4654 let mut cnt = vec![0u32; n_expert];
4655 for &s in sel_all.iter() {
4656 cnt[s as usize] += 1;
4657 }
4658 let total = sel_all.len() as f64;
4659 let mut h = 0.0f64;
4660 let mut active = 0usize;
4661 for &c in &cnt {
4662 if c > 0 {
4663 active += 1;
4664 let p = c as f64 / total;
4665 h -= p * p.log2();
4666 }
4667 }
4668 let maxc = cnt.iter().copied().max().unwrap_or(0);
4669 println!(
4670 "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4671 il,
4672 t,
4673 sel_all.len(),
4674 active,
4675 n_expert,
4676 h,
4677 (n_expert as f64).log2(),
4678 total / active.max(1) as f64,
4679 maxc
4680 );
4681 }
4682
4683 let gdec_may_fire = uniform_experts
4696 && use_cache
4697 && n_used <= 8
4698 && gdec_enabled()
4699 && !cfg.swiglu_clamped_at(il as u32);
4700 let slab_local = m
4716 .dev_exps
4717 .as_ref()
4718 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4719 let slab_bases = slab_local.map(|d| {
4720 use cudarc::driver::DevicePtr;
4721 let s = e.stream();
4722 let (pg, _g0) = d.gate.device_ptr(&s);
4723 let (pu, _g1) = d.up.device_ptr(&s);
4724 let (pd, _g2) = d.down.device_ptr(&s);
4725 (pg as u64, pu as u64, pd as u64)
4726 });
4727 let slab_fused_may_fire = slab_bases.is_some()
4737 && n_used <= 8
4738 && gdec_enabled()
4739 && !cfg.swiglu_clamped_at(il as u32)
4740 && cfg.m3.is_none()
4741 && no_exp_macros
4742 && moe_q8;
4743 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4746 e.uninit(t * n_embd)?
4747 } else {
4748 e.zeros(t * n_embd)?
4749 };
4750 let cpu_input = if cpu_hybrid {
4753 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4754 } else {
4755 None
4756 };
4757
4758 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;
4766 let mut scratch_u: Option<CudaSlice<u8>> = None;
4767 let mut scratch_d: Option<CudaSlice<u8>> = None;
4768 let page_window = moe_page_prefetch_window();
4776
4777 for tok in 0..t {
4780 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4781 let w = &w_all[tok * n_used..(tok + 1) * n_used];
4782 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4784
4785 let no_macros = m.gate_exps.macros.is_none()
4799 && m.up_exps.macros.is_none()
4800 && m.down_exps.macros.is_none();
4801 if slab_fused_may_fire {
4811 let (pg, pu, pd) = slab_bases.unwrap();
4812 let mut gp = [0u64; 8];
4813 let mut up = [0u64; 8];
4814 let mut dp = [0u64; 8];
4815 for (j, &ex) in sel.iter().enumerate() {
4816 let ex = ex as usize;
4817 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4818 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4819 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4820 }
4821 let mut wv = [0f32; 8];
4822 wv[..n_used].copy_from_slice(w);
4823 if tok_q8.is_none() {
4824 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4825 }
4826 let (zq, zd) = tok_q8.as_ref().unwrap();
4827 let act = e.moe_gate_up_silu8_q8(
4828 crate::WPtr8(gp),
4829 crate::WPtr8(up),
4830 zq,
4831 zd,
4832 n_embd,
4833 n_ff_exp,
4834 n_used,
4835 m.gate_exps.qtype,
4836 m.up_exps.qtype,
4837 m.gate_exps.row_bytes,
4838 m.up_exps.row_bytes,
4839 )?;
4840 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4841 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4842 e.moe_down8_fma_q8(
4843 crate::WPtr8(dp),
4844 crate::F32x8(wv),
4845 &aq2,
4846 &ad2,
4847 &mut dst,
4848 n_ff_exp,
4849 n_embd,
4850 n_used,
4851 m.down_exps.qtype,
4852 m.down_exps.row_bytes,
4853 )?;
4854 continue;
4855 }
4856 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4857 if tok_q8.is_none() {
4858 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4859 }
4860 let (zq, zd) = tok_q8.as_ref().unwrap();
4861 if Self::moe_gdec_token_q8(
4862 e,
4863 m,
4864 il,
4865 max_block,
4866 zq,
4867 zd,
4868 sel,
4869 w,
4870 &mut moe_out,
4871 tok,
4872 n_embd,
4873 n_ff_exp,
4874 n_used,
4875 )? {
4876 continue;
4877 }
4878 } else if gdec_may_fire
4879 && cfg.m3.is_none()
4880 && no_macros
4881 && Self::moe_gdec_token(
4882 e,
4883 m,
4884 il,
4885 max_block,
4886 &zt,
4887 sel,
4888 w,
4889 &mut moe_out,
4890 tok,
4891 n_embd,
4892 n_ff_exp,
4893 n_used,
4894 )?
4895 {
4896 continue;
4897 }
4898
4899 if gdec_may_fire || slab_fused_may_fire {
4905 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4906 e.memset_zeros_view(&mut row)?;
4907 }
4908
4909 let mut cpu_mask = vec![false; sel.len()];
4915 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
4916 let gpu_resident = if use_cache {
4917 e.with_moe_cache(max_block, |cache, _| {
4918 Ok(sel
4919 .iter()
4920 .map(|&expert| {
4921 let expert = expert as u16;
4922 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
4923 .into_iter()
4924 .filter(|&projection| {
4925 cache
4926 .resident(BlockId::new(il, projection, expert))
4927 .is_some()
4928 })
4929 .count()
4930 })
4931 .collect::<Vec<_>>())
4932 })?
4933 } else {
4934 vec![0; sel.len()]
4935 };
4936 let mut cpu_selected = Vec::new();
4937 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
4938 if gpu_resident[index] != 3 {
4939 cpu_mask[index] = true;
4940 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
4941 let expert = expert as usize;
4942 cpu_selected.push((expert, route_weight));
4943 }
4944 }
4945 if crate::cpu_experts::predictor_enabled() {
4946 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4950 crate::cpu_experts::predictor_submit(il, row);
4951 }
4952 if cpu_selected.is_empty() {
4953 None
4954 } else {
4955 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4956 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
4957 .map_err(std::io::Error::other)?;
4958 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
4959 }
4960 } else {
4961 None
4962 };
4963
4964 let worker_window = worker_disk_prefetch
4965 .then(worker_prefetch_window)
4966 .unwrap_or(0);
4967 for (j, &ex) in sel.iter().enumerate() {
4968 if cpu_mask[j] {
4969 continue;
4970 }
4971 let ex = ex as usize;
4972 if let Some(d) = slab_local {
4979 let gl = m.gate_exps.expert_layout(ex);
4980 let ul = m.up_exps.expert_layout(ex);
4981 let dl = m.down_exps.expert_layout(ex);
4982 let (g0, u0, d0) = (
4983 ex * m.gate_exps.expert_stride,
4984 ex * m.up_exps.expert_stride,
4985 ex * m.down_exps.expert_stride,
4986 );
4987 let (gate, up) = if moe_q8 {
4988 if tok_q8.is_none() {
4989 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4990 }
4991 let (zq, zd) = tok_q8.as_ref().unwrap();
4992 (
4993 e.qmatvec_expert_q8(
4994 &d.gate,
4995 g0..g0 + gl.len,
4996 zq,
4997 zd,
4998 1,
4999 m.gate_exps.in_f,
5000 m.gate_exps.out_f,
5001 gl.qtype,
5002 gl.row_bytes,
5003 )?,
5004 e.qmatvec_expert_q8(
5005 &d.up,
5006 u0..u0 + ul.len,
5007 zq,
5008 zd,
5009 1,
5010 m.up_exps.in_f,
5011 m.up_exps.out_f,
5012 ul.qtype,
5013 ul.row_bytes,
5014 )?,
5015 )
5016 } else {
5017 (
5018 e.qmatvec_view(
5019 &d.gate,
5020 g0..g0 + gl.len,
5021 &zt,
5022 1,
5023 m.gate_exps.in_f,
5024 m.gate_exps.out_f,
5025 gl.qtype,
5026 gl.row_bytes,
5027 )?,
5028 e.qmatvec_view(
5029 &d.up,
5030 u0..u0 + ul.len,
5031 &zt,
5032 1,
5033 m.up_exps.in_f,
5034 m.up_exps.out_f,
5035 ul.qtype,
5036 ul.row_bytes,
5037 )?,
5038 )
5039 };
5040 let mut act = e.uninit(n_ff_exp)?;
5041 Self::ffn_act_lim(
5042 e,
5043 cfg,
5044 &gate,
5045 &up,
5046 m.gate_exps.macro_scale(ex),
5047 m.up_exps.macro_scale(ex),
5048 lim_exp,
5049 &mut act,
5050 n_ff_exp,
5051 )?;
5052 let y = if moe_q8 {
5053 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5054 e.qmatvec_expert_q8(
5055 &d.down,
5056 d0..d0 + dl.len,
5057 &aq2,
5058 &ad2,
5059 1,
5060 m.down_exps.in_f,
5061 m.down_exps.out_f,
5062 dl.qtype,
5063 dl.row_bytes,
5064 )?
5065 } else {
5066 let actv = act.slice(0..n_ff_exp);
5067 e.qmatvec_view(
5068 &d.down,
5069 d0..d0 + dl.len,
5070 &actv,
5071 1,
5072 m.down_exps.in_f,
5073 m.down_exps.out_f,
5074 dl.qtype,
5075 dl.row_bytes,
5076 )?
5077 };
5078 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5079 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5080 continue;
5081 }
5082 for next in page_prefetch_positions(j, sel.len(), page_window) {
5083 Self::moe_prefetch_host_expert(sel[next] as usize, m);
5084 }
5085 let keep = [
5086 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5087 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5088 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5089 ];
5090 if worker_disk_prefetch && worker_window > 0 {
5091 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5092 Self::moe_prefetch_disk_expert(
5093 e,
5094 il,
5095 sel[next] as usize,
5096 m,
5097 max_block,
5098 &keep,
5099 )?;
5100 }
5101 } else if cache_dispatch
5102 && !cpu_hybrid
5103 && moe_prefetch_enabled()
5104 && j + 1 < sel.len()
5105 {
5106 let next = sel[j + 1] as usize;
5107 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5108 }
5109 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5110 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5111 if (gate_q8 || up_q8) && tok_q8.is_none() {
5114 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5115 }
5116 let gate = if gate_q8 {
5117 let (zq, zd) = tok_q8.as_ref().unwrap();
5118 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5119 } else {
5120 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5121 };
5122 let up = if up_q8 {
5123 let (zq, zd) = tok_q8.as_ref().unwrap();
5124 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5125 } else {
5126 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5127 };
5128 let mut act = e.uninit(n_ff_exp)?;
5129 Self::ffn_act_lim(
5130 e,
5131 cfg,
5132 &gate,
5133 &up,
5134 m.gate_exps.macro_scale(ex),
5135 m.up_exps.macro_scale(ex),
5136 lim_exp,
5137 &mut act,
5138 n_ff_exp,
5139 )?;
5140 let y = if down_q8 {
5141 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5142 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5143 } else {
5144 let actv = act.slice(0..n_ff_exp);
5145 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5146 };
5147 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5148 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5150 } else if cache_dispatch {
5151 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5156 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5157 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5159 e,
5160 cfg,
5161 &gate,
5162 &up,
5163 m.gate_exps.macro_scale(ex),
5164 m.up_exps.macro_scale(ex),
5165 lim_exp,
5166 &mut act,
5167 n_ff_exp,
5168 )?;
5169 let actv = act.slice(0..n_ff_exp);
5170 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5171 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5172 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5174 } else if cache_frozen {
5175 let gate = Self::moe_frozen_gemm(
5180 e,
5181 il,
5182 PROJ_GATE,
5183 ex,
5184 m,
5185 max_block,
5186 &zt,
5187 &mut scratch_g,
5188 g_len,
5189 )?;
5190 let up = Self::moe_frozen_gemm(
5191 e,
5192 il,
5193 PROJ_UP,
5194 ex,
5195 m,
5196 max_block,
5197 &zt,
5198 &mut scratch_u,
5199 u_len,
5200 )?;
5201 let mut act = e.uninit(n_ff_exp)?;
5202 Self::ffn_act_lim(
5203 e,
5204 cfg,
5205 &gate,
5206 &up,
5207 m.gate_exps.macro_scale(ex),
5208 m.up_exps.macro_scale(ex),
5209 lim_exp,
5210 &mut act,
5211 n_ff_exp,
5212 )?;
5213 let actv = act.slice(0..n_ff_exp);
5214 let y = Self::moe_frozen_gemm(
5215 e,
5216 il,
5217 PROJ_DOWN,
5218 ex,
5219 m,
5220 max_block,
5221 &actv,
5222 &mut scratch_d,
5223 d_len,
5224 )?;
5225 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5226 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5227 } else {
5228 if scratch_g.is_none() {
5232 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5233 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5234 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5235 }
5236 let (sg, su, sd) = (
5237 scratch_g.as_mut().unwrap(),
5238 scratch_u.as_mut().unwrap(),
5239 scratch_d.as_mut().unwrap(),
5240 );
5241 let gl = m.gate_exps.expert_layout(ex);
5242 let ul = m.up_exps.expert_layout(ex);
5243 let dl = m.down_exps.expert_layout(ex);
5244 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5245 let gate = e.qmatvec_view(
5246 sg,
5247 0..gl.len,
5248 &zt,
5249 1,
5250 m.gate_exps.in_f,
5251 m.gate_exps.out_f,
5252 gl.qtype,
5253 gl.row_bytes,
5254 )?;
5255
5256 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5257 let up = e.qmatvec_view(
5258 su,
5259 0..ul.len,
5260 &zt,
5261 1,
5262 m.up_exps.in_f,
5263 m.up_exps.out_f,
5264 ul.qtype,
5265 ul.row_bytes,
5266 )?;
5267
5268 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
5270 e,
5271 cfg,
5272 &gate,
5273 &up,
5274 m.gate_exps.macro_scale(ex),
5275 m.up_exps.macro_scale(ex),
5276 lim_exp,
5277 &mut act,
5278 n_ff_exp,
5279 )?;
5280
5281 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5282 let actv = act.slice(0..n_ff_exp);
5283 let y = e.qmatvec_view(
5284 sd,
5285 0..dl.len,
5286 &actv,
5287 1,
5288 m.down_exps.in_f,
5289 m.down_exps.out_f,
5290 dl.qtype,
5291 dl.row_bytes,
5292 )?;
5293
5294 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5295 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5296 }
5297 }
5298 if let Some(worker) = cpu_worker {
5299 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5300 let cpu_output = e.htod(&cpu_output)?;
5301 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5302 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5303 }
5304 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5305 for (j, &ex) in sel.iter().enumerate() {
5306 if cpu_mask[j] {
5307 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5308 }
5309 }
5310 }
5311 }
5312
5313 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5318 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5319 {
5320 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
5329 let (sg_gate, sg_up) = if t == 1 {
5330 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5331 Some(pair) => pair,
5332 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5333 }
5334 } else if verify_t {
5335 (
5336 e.matmul_decode_exact(gate_shexp, z, t)?,
5337 e.matmul_decode_exact(up_shexp, z, t)?,
5338 )
5339 } else {
5340 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
5342 let mut sa = e.uninit(t * n_ff_sh)?; Self::ffn_act_lim(
5344 e,
5345 cfg,
5346 &sg_gate,
5347 &sg_up,
5348 1.0,
5349 1.0,
5350 lim_shexp,
5351 &mut sa,
5352 t * n_ff_sh,
5353 )?;
5354 let sh = if verify_t {
5355 e.matmul_decode_exact(down_shexp, &sa, t)?
5356 } else {
5357 e.matmul(down_shexp, &sa, t)?
5358 }; let g = match &m.gate_inp_shexp {
5372 Some(gate_inp_shexp) => {
5373 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5374 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5375 } else {
5376 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5377 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
5379 g
5380 }
5381 }
5382 None => e.htod(&vec![1.0f32; t])?,
5383 };
5384 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5386 }
5387
5388 Ok(moe_out)
5389 }
5390
5391 pub fn stage1_h2d_per_token(&self) -> u64 {
5394 use crate::hybrid::Ffn;
5395 let n_used = self
5396 .cfg
5397 .moe
5398 .as_ref()
5399 .map(|m| m.expert_used_count as u64)
5400 .unwrap_or(0);
5401 let mut bytes = 0u64;
5402 for l in self.layers.iter() {
5403 if let Ffn::Moe(m) = &l.ffn {
5404 bytes += n_used
5405 * (m.gate_exps.max_expert_bytes()
5406 + m.up_exps.max_expert_bytes()
5407 + m.down_exps.max_expert_bytes()) as u64;
5408 }
5409 }
5410 bytes
5411 }
5412
5413 pub(crate) fn max_moe_block(&self) -> usize {
5417 use crate::hybrid::Ffn;
5418 let mut mx = 0usize;
5419 let mut scan = |ffn: &Ffn| {
5420 if let Ffn::Moe(m) = ffn {
5421 mx = mx
5422 .max(m.gate_exps.max_expert_bytes())
5423 .max(m.up_exps.max_expert_bytes())
5424 .max(m.down_exps.max_expert_bytes());
5425 }
5426 };
5427 for l in self.layers.iter() {
5428 scan(&l.ffn);
5429 }
5430 if let Some(mtp) = self.mtp.as_ref() {
5431 scan(&mtp.ffn);
5432 }
5433 mx
5434 }
5435
5436 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5439 use crate::hybrid::Ffn;
5440 let mut sizes = Vec::new();
5441 let mut scan = |ffn: &Ffn| {
5442 let Ffn::Moe(m) = ffn else { return };
5443 for ex in 0..m.gate_exps.n_expert {
5444 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5445 continue;
5446 }
5447 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5448 let len = exps.expert_layout(ex).len;
5449 if len > 0 {
5450 sizes.push(len);
5451 }
5452 }
5453 }
5454 };
5455 for layer in &self.layers {
5456 scan(&layer.ffn);
5457 }
5458 if let Some(mtp) = &self.mtp {
5459 scan(&mtp.ffn);
5460 }
5461 sizes
5462 }
5463
5464 pub fn save_cpu_expert_residency_profile(
5470 &self,
5471 e: &Engine,
5472 path: &std::path::Path,
5473 ) -> Result<(), Box<dyn std::error::Error>> {
5474 let Some(ids) = e.export_moe_residency() else {
5475 return Err("no MoE residency cache to persist".into());
5476 };
5477 let mut body = format!(
5478 "memra-freeze-profile v1 max_block={} blocks={}\n",
5479 self.max_moe_block(),
5480 ids.len()
5481 );
5482 for (layer, proj, ex) in &ids {
5483 body.push_str(&format!("{layer} {proj} {ex}\n"));
5484 }
5485 let tmp = path.with_extension("tmp");
5486 std::fs::write(&tmp, body)?;
5487 std::fs::rename(&tmp, path)?;
5488 println!(
5489 "[moe-cache] freeze profile saved: {} blocks -> {}",
5490 ids.len(),
5491 path.display()
5492 );
5493 Ok(())
5494 }
5495
5496 pub fn restore_cpu_expert_residency_profile(
5500 &self,
5501 e: &Engine,
5502 path: &std::path::Path,
5503 ) -> Result<bool, Box<dyn std::error::Error>> {
5504 use crate::hybrid::Ffn;
5505 use crate::moe_cache::BlockId;
5506 let Ok(content) = std::fs::read_to_string(path) else {
5507 return Ok(false);
5508 };
5509 let mut lines = content.lines();
5510 let Some(header) = lines.next() else {
5511 return Ok(false);
5512 };
5513 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5514 if !header.starts_with(&expected) {
5515 println!(
5516 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5517 path.display()
5518 );
5519 return Ok(false);
5520 }
5521 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5522 std::collections::HashMap::new();
5523 for line in lines {
5524 let mut fields = line.split_whitespace();
5525 let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5526 else {
5527 continue;
5528 };
5529 let (Ok(layer), Ok(proj), Ok(ex)) =
5530 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5531 else {
5532 continue;
5533 };
5534 by_layer
5535 .entry(layer)
5536 .or_default()
5537 .push(BlockId::new(layer, proj, ex));
5538 }
5539 let requested: usize = by_layer.values().map(Vec::len).sum();
5540 if requested == 0 {
5541 return Ok(false);
5542 }
5543 let max_block = self.max_moe_block();
5544 let mut restaged = 0usize;
5545 let mut stage_layer =
5546 |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5547 let Ffn::Moe(m) = ffn else { return Ok(()) };
5548 let Some(ids) = by_layer.get(&layer_index) else {
5549 return Ok(());
5550 };
5551 e.with_moe_cache(max_block, |cache, eng| {
5552 for id in ids {
5553 if cache.restage_block(*id, m, eng)? {
5554 restaged += 1;
5555 }
5556 }
5557 Ok(())
5558 })
5559 };
5560 for (index, layer) in self.layers.iter().enumerate() {
5561 stage_layer(index as u16, &layer.ffn)?;
5562 }
5563 if let Some(mtp) = self.mtp.as_ref() {
5564 stage_layer(u16::MAX, &mtp.ffn)?;
5565 }
5566 e.freeze_moe_cache();
5567 println!(
5568 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5569 path.display()
5570 );
5571 Ok(true)
5572 }
5573
5574 pub fn freeze_cpu_expert_residency(
5576 &self,
5577 e: &Engine,
5578 ) -> Result<(), Box<dyn std::error::Error>> {
5579 e.freeze_moe_cache();
5580 Ok(())
5581 }
5582
5583 pub fn ffn_act(
5591 e: &Engine,
5592 cfg: &ModelConfig,
5593 gate: &CudaSlice<f32>,
5594 up: &CudaSlice<f32>,
5595 act: &mut CudaSlice<f32>,
5596 n: usize,
5597 ) -> Result<(), Box<dyn std::error::Error>> {
5598 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5599 }
5600
5601 #[allow(clippy::too_many_arguments)]
5605 pub(crate) fn ffn_act_scaled(
5606 e: &Engine,
5607 cfg: &ModelConfig,
5608 gate: &CudaSlice<f32>,
5609 up: &CudaSlice<f32>,
5610 gs: f32,
5611 us: f32,
5612 act: &mut CudaSlice<f32>,
5613 n: usize,
5614 ) -> Result<(), Box<dyn std::error::Error>> {
5615 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5616 }
5617
5618 #[allow(clippy::too_many_arguments)]
5627 pub(crate) fn ffn_act_lim(
5628 e: &Engine,
5629 cfg: &ModelConfig,
5630 gate: &CudaSlice<f32>,
5631 up: &CudaSlice<f32>,
5632 gs: f32,
5633 us: f32,
5634 limit: Option<f32>,
5635 act: &mut CudaSlice<f32>,
5636 n: usize,
5637 ) -> Result<(), Box<dyn std::error::Error>> {
5638 if let Some(m3) = cfg.m3.as_ref() {
5639 debug_assert!(
5640 limit.is_none(),
5641 "m3 swigluoai and step35 clamp are different archs"
5642 );
5643 return e.swigluoai_mul_scaled(
5644 gate,
5645 up,
5646 gs,
5647 us,
5648 m3.swiglu_alpha,
5649 m3.swiglu_limit,
5650 act,
5651 n,
5652 );
5653 }
5654 if let Some(l) = limit {
5655 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5656 }
5657 if gs == 1.0 && us == 1.0 {
5658 return e.silu_mul(gate, up, act, n);
5659 }
5660 e.silu_mul_scaled(gate, up, gs, us, act, n)
5661 }
5662
5663 fn moe_route(
5669 e: &Engine,
5670 logits: &CudaSlice<f32>,
5671 t: usize,
5672 n_expert: usize,
5673 n_used: usize,
5674 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5675 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5676 }
5677
5678 #[allow(clippy::too_many_arguments)]
5686 fn moe_route_sigmoid_cfg(
5687 e: &Engine,
5688 logits: &CudaSlice<f32>,
5689 t: usize,
5690 n_expert: usize,
5691 n_used: usize,
5692 m: &MoeWeights,
5693 (sf, route_norm): (f32, bool),
5694 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5695 if sigmoid_router_enabled() {
5696 return e.moe_router_sigmoid_topk_host(
5697 logits,
5698 t,
5699 n_expert,
5700 n_used,
5701 m.active_count(),
5702 &m.exp_probs_b_dev,
5703 &m.active_experts_dev,
5704 sf,
5705 route_norm,
5706 );
5707 }
5708 let lg = e.dtoh(logits)?;
5709 Self::moe_route_sigmoid_host(
5710 &lg,
5711 t,
5712 n_expert,
5713 n_used,
5714 m.exp_probs_b.as_deref(),
5715 sf,
5716 route_norm,
5717 m.active_experts.as_deref(),
5718 )
5719 }
5720
5721 fn moe_route_cfg(
5724 e: &Engine,
5725 logits: &CudaSlice<f32>,
5726 t: usize,
5727 n_expert: usize,
5728 n_used: usize,
5729 active: Option<&[bool]>,
5730 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5731 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5734 return e.moe_router_topk_host(logits, t, n_expert, n_used);
5735 }
5736 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
5739 let mut w_out = vec![0f32; t * n_used];
5740 for tok in 0..t {
5741 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5742 let maxl = row
5744 .iter()
5745 .enumerate()
5746 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5747 .map(|(_, &x)| x)
5748 .fold(f32::NEG_INFINITY, f32::max);
5749 let mut probs = vec![0f32; n_expert];
5750 let mut den = 0f32;
5751 for i in 0..n_expert {
5752 if active.is_some_and(|mask| !mask[i]) {
5753 continue;
5754 }
5755 let x = (row[i] - maxl).exp();
5756 probs[i] = x;
5757 den += x;
5758 }
5759 for p in probs.iter_mut() {
5760 *p /= den;
5761 }
5762 let mut idx: Vec<usize> = (0..n_expert)
5764 .filter(|&i| active.is_none_or(|mask| mask[i]))
5765 .collect();
5766 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5767 let sl = &idx[..n_used];
5768 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5769 let mut ws: f32 = wv.iter().sum();
5770 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() {
5772 *x /= ws;
5773 }
5774 for j in 0..n_used {
5775 sel[tok * n_used + j] = sl[j] as u32;
5776 w_out[tok * n_used + j] = wv[j];
5777 }
5778 }
5779 Ok((sel, w_out))
5780 }
5781
5782 #[allow(clippy::too_many_arguments)]
5783 fn moe_route_sigmoid_with_input(
5784 e: &Engine,
5785 logits: &CudaSlice<f32>,
5786 input: &CudaSlice<f32>,
5787 t: usize,
5788 n_expert: usize,
5789 n_used: usize,
5790 bias: Option<&[f32]>,
5791 (sf, route_norm): (f32, bool),
5792 active: Option<&[bool]>,
5793 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5794 let (lg, input) = e.dtoh_pair(logits, input)?;
5795 let (sel, w) =
5796 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5797 Ok((sel, w, input))
5798 }
5799
5800 pub fn start_moe_prefetch_predictor(
5805 &self,
5806 e: &Engine,
5807 cfg: &ModelConfig,
5808 ) -> Result<(), Box<dyn std::error::Error>> {
5809 use crate::hybrid::Ffn;
5810 let Some(sig) = cfg.sigmoid_router() else {
5811 return Err("prefetch predictor requires a sigmoid-router arch".into());
5812 };
5813 let resident: std::collections::HashSet<(u16, u8, u16)> = e
5814 .export_moe_residency()
5815 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5816 .into_iter()
5817 .collect();
5818 let mut layers = Vec::new();
5819 for (index, layer) in self.layers.iter().enumerate() {
5820 let Ffn::Moe(m) = &layer.ffn else { continue };
5821 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5822 continue;
5823 };
5824 let router = e.dtoh(data)?;
5825 let n_expert = m.gate_exps.n_expert;
5826 let n_embd = m.gate_exps.in_f;
5827 if router.len() != n_embd * n_expert {
5828 continue;
5829 }
5830 let build = |exps: &crate::model::HostExps| {
5831 (0..n_expert)
5832 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5833 .collect::<Vec<_>>()
5834 };
5835 layers.push((
5836 index as u16,
5837 crate::cpu_experts::PredictLayerInit {
5838 router,
5839 bias: m.exp_probs_b.clone(),
5840 active: m.active_experts.clone(),
5841 n_embd,
5842 n_used: cfg
5843 .moe
5844 .as_ref()
5845 .map(|moe| moe.expert_used_count as usize)
5846 .ok_or("prefetch predictor requires MoE config")?,
5847 sig,
5848 weights_n_expert: n_expert,
5849 gate: build(&m.gate_exps),
5850 up: build(&m.up_exps),
5851 down: build(&m.down_exps),
5852 },
5853 ));
5854 }
5855 crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5856 }
5857
5858 #[allow(clippy::too_many_arguments)]
5861 pub fn moe_route_sigmoid_host_public(
5862 logits: &[f32],
5863 t: usize,
5864 n_expert: usize,
5865 n_used: usize,
5866 bias: Option<&[f32]>,
5867 sf: f32,
5868 route_norm: bool,
5869 active: Option<&[bool]>,
5870 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5871 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
5872 }
5873
5874 #[allow(clippy::too_many_arguments)]
5875 fn moe_route_sigmoid_host(
5876 lg: &[f32],
5877 t: usize,
5878 n_expert: usize,
5879 n_used: usize,
5880 bias: Option<&[f32]>,
5881 sf: f32,
5882 route_norm: bool,
5883 active: Option<&[bool]>,
5884 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5885 let active_count = active
5886 .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
5887 .unwrap_or(n_expert);
5888 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5889 if lg.len() != t * n_expert {
5890 return Err(format!(
5891 "sigmoid router logits length mismatch: got {}, expected {}",
5892 lg.len(),
5893 t * n_expert,
5894 )
5895 .into());
5896 }
5897 let mut sel = vec![0u32; t * n_used];
5898 let mut w_out = vec![0f32; t * n_used];
5899 for tok in 0..t {
5900 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5901 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
5902 let selsc: Vec<f32> = match bias {
5904 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
5905 None => scores.clone(),
5906 };
5907 let mut idx: Vec<usize> = (0..n_expert)
5908 .filter(|&i| active.is_none_or(|mask| mask[i]))
5909 .collect();
5910 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
5911 let sl = &idx[..n_used];
5912 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
5913 if route_norm {
5914 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
5915 for x in wv.iter_mut() {
5916 *x = *x / ws * sf;
5917 }
5918 } else {
5919 for x in wv.iter_mut() {
5920 *x *= sf;
5921 }
5922 }
5923 for j in 0..n_used {
5924 sel[tok * n_used + j] = sl[j] as u32;
5925 w_out[tok * n_used + j] = wv[j];
5926 }
5927 }
5928 Ok((sel, w_out))
5929 }
5930
5931 #[allow(clippy::too_many_arguments)]
5935 fn moe_ffn_sigmoid_dev(
5936 e: &Engine,
5937 m: &MoeWeights,
5938 z: &CudaSlice<f32>,
5939 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
5940 logits: &CudaSlice<f32>,
5941 t: usize,
5942 cfg: &ModelConfig,
5943 il: u16,
5944 (scaling_factor, route_norm): (f32, bool),
5945 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5946 let moe = cfg.moe.as_ref().unwrap();
5947 let n_embd = cfg.n_embd as usize;
5948 let n_expert = moe.expert_count as usize;
5949 let n_used = moe.expert_used_count as usize;
5950 let n_ff_exp = moe.expert_ff_length as usize;
5951 let dev = m.dev_exps.as_ref().unwrap();
5952 debug_assert!(cfg.step35.is_some());
5953 debug_assert_eq!(dev.dev, e.ctx().ordinal());
5954 debug_assert!(m.has_uniform_expert_layout());
5955 debug_assert!(!m.has_macros);
5956
5957 let (sel_d, w_d) = e.moe_router_sigmoid_topk(
5958 logits,
5959 t,
5960 n_expert,
5961 n_used,
5962 m.active_count(),
5963 &m.exp_probs_b_dev,
5964 &m.active_experts_dev,
5965 scaling_factor,
5966 route_norm,
5967 )?;
5968 crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
5969 let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
5970 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5971 (combined, combined)
5972 } else {
5973 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5974 };
5975 let (zq, zd) = match (t, zq8) {
5976 (1, Some((q, d))) => (q.clone(), d.clone()),
5977 _ => e.quantize_q8_1(z, t, n_embd)?,
5978 };
5979 let n_pairs = t * n_used;
5980 let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
5981 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5985 let pair_tok_d = e.htod_i32(&pair_tok)?;
5986 let gate = e.moe_pairs_matvec_q8(
5987 &dev.ptr_row,
5988 0,
5989 &pair_tok_d,
5990 &sel_d,
5991 &zq,
5992 &zd,
5993 n_embd,
5994 n_ff_exp,
5995 n_expert,
5996 n_pairs,
5997 m.gate_exps.qtype,
5998 gate_row_bytes,
5999 )?;
6000 let up = e.moe_pairs_matvec_q8(
6001 &dev.ptr_row,
6002 1,
6003 &pair_tok_d,
6004 &sel_d,
6005 &zq,
6006 &zd,
6007 n_embd,
6008 n_ff_exp,
6009 n_expert,
6010 n_pairs,
6011 m.up_exps.qtype,
6012 up_row_bytes,
6013 )?;
6014 let mut act = e.uninit(n_pairs * n_ff_exp)?;
6015 Self::ffn_act_lim(
6016 e,
6017 cfg,
6018 &gate,
6019 &up,
6020 1.0,
6021 1.0,
6022 cfg.clamp_exp_at(il as u32),
6023 &mut act,
6024 n_pairs * n_ff_exp,
6025 )?;
6026 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6027 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6028 let pair_self_d = e.htod_i32(&pair_self)?;
6029 let down = e.moe_pairs_matvec_q8(
6030 &dev.ptr_row,
6031 2,
6032 &pair_self_d,
6033 &sel_d,
6034 &aq2,
6035 &ad2,
6036 n_ff_exp,
6037 n_embd,
6038 n_expert,
6039 n_pairs,
6040 m.down_exps.qtype,
6041 m.down_exps.row_bytes,
6042 )?;
6043 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6044 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6045 let tok_off_d = e.htod_i32(&tok_off)?;
6046 let tok_ids_d = e.htod_i32(&tok_ids)?;
6047 let mut output = e.uninit(t * n_embd)?;
6048 e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
6049 output
6050 } else {
6051 let act = e.moe_gate_up_silu8_dev_q8_rows(
6052 &dev.ptr_row,
6053 &sel_d,
6054 &zq,
6055 &zd,
6056 t,
6057 n_embd,
6058 n_ff_exp,
6059 n_used,
6060 n_expert,
6061 m.gate_exps.qtype,
6062 m.up_exps.qtype,
6063 gate_row_bytes,
6064 up_row_bytes,
6065 &m.dev_macros,
6066 )?;
6067 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6068 let mut output = e.uninit(t * n_embd)?;
6069 e.moe_down8_fma_dev_q8_rows_g(
6070 &dev.ptr_row,
6071 &sel_d,
6072 &w_d,
6073 &aq2,
6074 &ad2,
6075 &mut output,
6076 t,
6077 n_ff_exp,
6078 n_embd,
6079 n_used,
6080 n_expert,
6081 m.down_exps.qtype,
6082 m.down_exps.row_bytes,
6083 )?;
6084 output
6085 };
6086
6087 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6088 eprintln!(
6089 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6090 cfg.clamp_exp_at(il as u32).is_some(),
6091 dev.gu_il,
6092 );
6093 }
6094 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6095 Ok(moe_out)
6096 }
6097
6098 fn moe_ffn_pairs(
6107 e: &Engine,
6108 m: &MoeWeights,
6109 z: &CudaSlice<f32>,
6110 logits: &CudaSlice<f32>,
6111 t: usize,
6112 cfg: &ModelConfig,
6113 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6114 let moe = cfg.moe.as_ref().unwrap();
6115 let n_embd = cfg.n_embd as usize;
6116 let n_expert = moe.expert_count as usize;
6117 let n_used = moe.expert_used_count as usize;
6118 let n_ff_exp = moe.expert_ff_length as usize;
6119 debug_assert!(
6124 !cfg.swiglu_clamped_anywhere(),
6125 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6126 );
6127 let dev = m.dev_exps.as_ref().unwrap();
6128 let (rbg_d, rbu_d) = if dev.gu_il {
6130 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6131 (sxx, sxx)
6132 } else {
6133 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6134 };
6135
6136 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6137 let n_pairs = t * n_used;
6138 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6141 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6142 let pair_w: Vec<f32> = w_all.clone();
6143 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6144 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6145 let pt = e.htod_i32(&pair_tok)?;
6146 let px = e.htod_i32(&pair_ex)?;
6147 let pw = e.htod(&pair_w)?;
6148 let toff = e.htod_i32(&tok_off)?;
6149 let tids = e.htod_i32(&tok_ids)?;
6150
6151 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6155 for p in 0..n_pairs {
6156 by_ex[pair_ex[p] as usize].push(p as i32);
6157 }
6158 let mut ex_ids: Vec<i32> = Vec::new();
6159 let mut ex_off: Vec<i32> = vec![0];
6160 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6161 for (ex, list) in by_ex.iter().enumerate() {
6162 if list.is_empty() {
6163 continue;
6164 }
6165 ex_ids.push(ex as i32);
6166 ex_pairs.extend_from_slice(list);
6167 ex_off.push(ex_pairs.len() as i32);
6168 }
6169 let n_active = ex_ids.len();
6170 let exi = e.htod_i32(&ex_ids)?;
6171 let exo = e.htod_i32(&ex_off)?;
6172 let exp_d = e.htod_i32(&ex_pairs)?;
6173 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6194 let mma_t = *MMA_T.get_or_init(|| {
6195 std::env::var("MEMRA_MOE_MMA_T")
6196 .ok()
6197 .and_then(|v| v.parse().ok())
6198 .unwrap_or(16)
6199 });
6200 let use_mma = std::env::var("MEMRA_MOE_MMA")
6201 .map(|v| v != "0")
6202 .unwrap_or(true)
6203 && t >= mma_t
6204 && q8_expert_dec_supported(m.gate_exps.qtype)
6205 && q8_expert_dec_supported(m.up_exps.qtype)
6206 && q8_expert_dec_supported(m.down_exps.qtype)
6207 && n_embd % 256 == 0
6208 && n_ff_exp % 256 == 0;
6209 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6225 && q8_expert_dec_supported(m.up_exps.qtype)
6226 && q8_expert_dec_supported(m.down_exps.qtype)
6227 && n_embd % 256 == 0
6228 && n_ff_exp % 256 == 0;
6229 let f16g_mode = crate::moe_f16g_mode();
6230 let f16g = f16g_mode != 0
6231 && t >= mma_t
6232 && (f16g_mode != 3 || !mma_capable)
6233 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6234 && f16g_proj_ok(m.up_exps.qtype, n_embd)
6235 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6236 if use_mma || f16g {
6237 let y_down = if f16g {
6245 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6249 let csr_tok_d = e.htod_i32(&csr_tok)?;
6250 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6251 let g_csr = e.moe_f16_grouped(
6252 &dev.ptr_row,
6253 0,
6254 n_expert,
6255 &exi,
6256 &ex_off,
6257 &exo,
6258 &z_f16,
6259 &z_s,
6260 n_embd,
6261 n_ff_exp,
6262 n_active,
6263 n_pairs,
6264 m.gate_exps.qtype,
6265 rbg_d,
6266 )?;
6267 let u_csr = e.moe_f16_grouped(
6268 &dev.ptr_row,
6269 1,
6270 n_expert,
6271 &exi,
6272 &ex_off,
6273 &exo,
6274 &z_f16,
6275 &z_s,
6276 n_embd,
6277 n_ff_exp,
6278 n_active,
6279 n_pairs,
6280 m.up_exps.qtype,
6281 rbu_d,
6282 )?;
6283 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6284 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6285 let d_csr = e.moe_f16_grouped(
6286 &dev.ptr_row,
6287 2,
6288 n_expert,
6289 &exi,
6290 &ex_off,
6291 &exo,
6292 &a_f16,
6293 &a_s,
6294 n_ff_exp,
6295 n_embd,
6296 n_active,
6297 n_pairs,
6298 m.down_exps.qtype,
6299 m.down_exps.row_bytes,
6300 )?;
6301 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6302 } else {
6303 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6305 let gate = e.mmq_iq_experts(
6306 &dev.ptr_row,
6307 0,
6308 n_expert,
6309 &exi,
6310 &exo,
6311 &exp_d,
6312 &pt,
6313 &z_scr,
6314 n_embd,
6315 n_ff_exp,
6316 n_active,
6317 n_pairs,
6318 t,
6319 m.gate_exps.qtype,
6320 rbg_d,
6321 )?;
6322 let up = e.mmq_iq_experts(
6323 &dev.ptr_row,
6324 1,
6325 n_expert,
6326 &exi,
6327 &exo,
6328 &exp_d,
6329 &pt,
6330 &z_scr,
6331 n_embd,
6332 n_ff_exp,
6333 n_active,
6334 n_pairs,
6335 t,
6336 m.up_exps.qtype,
6337 rbu_d,
6338 )?;
6339 let a_scr = if crate::moe_fuse_actq_on() {
6345 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6346 } else {
6347 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6348 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6349 };
6350 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6351 let pself = e.htod_i32(&pair_self)?;
6352 e.mmq_iq_experts(
6353 &dev.ptr_row,
6354 2,
6355 n_expert,
6356 &exi,
6357 &exo,
6358 &exp_d,
6359 &pself,
6360 &a_scr,
6361 n_ff_exp,
6362 n_embd,
6363 n_active,
6364 n_pairs,
6365 n_pairs,
6366 m.down_exps.qtype,
6367 m.down_exps.row_bytes,
6368 )?
6369 };
6370 let mut moe_out = e.uninit(t * n_embd)?;
6371 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6372 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6373 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6374 {
6375 let n_ff_sh = gate_shexp.out_features();
6376 let sg_gate = e.matmul(gate_shexp, z, t)?;
6377 let sg_up = e.matmul(up_shexp, z, t)?;
6378 let mut sa = e.uninit(t * n_ff_sh)?;
6379 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6380 let sh = e.matmul(down_shexp, &sa, t)?;
6381 let g = match &m.gate_inp_shexp {
6387 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6388 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6389 }
6390 Some(gate_inp_shexp) => {
6391 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6392 let mut g = e.uninit(t)?;
6393 e.sigmoid(&gs, &mut g, t)?;
6394 g
6395 }
6396 None => e.htod(&vec![1.0f32; t])?,
6397 };
6398 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6399 }
6400 return Ok(moe_out);
6401 }
6402
6403 let dec = std::env::var("MEMRA_MOE_DEC")
6406 .map(|v| v != "0")
6407 .unwrap_or(true);
6408 let matvec = |proj,
6409 exi: &_,
6410 exo: &_,
6411 exp_d: &_,
6412 pt: &_,
6413 aq: &_,
6414 ad: &_,
6415 inf,
6416 outf,
6417 qtype,
6418 rb|
6419 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6420 let dec = dec && q8_expert_dec_supported(qtype);
6422 if dec {
6423 e.moe_pairs_matvec_q8_dec(
6424 &dev.ptr_row,
6425 proj,
6426 exi,
6427 exo,
6428 exp_d,
6429 pt,
6430 aq,
6431 ad,
6432 inf,
6433 outf,
6434 n_expert,
6435 n_active,
6436 n_pairs,
6437 qtype,
6438 rb,
6439 )
6440 } else {
6441 e.moe_pairs_matvec_q8_em(
6442 &dev.ptr_row,
6443 proj,
6444 exi,
6445 exo,
6446 exp_d,
6447 pt,
6448 aq,
6449 ad,
6450 inf,
6451 outf,
6452 n_expert,
6453 n_active,
6454 n_pairs,
6455 qtype,
6456 rb,
6457 )
6458 }
6459 };
6460 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6461 let gate = matvec(
6462 0,
6463 &exi,
6464 &exo,
6465 &exp_d,
6466 &pt,
6467 &zq,
6468 &zd,
6469 n_embd,
6470 n_ff_exp,
6471 m.gate_exps.qtype,
6472 rbg_d,
6473 )?;
6474 let up = matvec(
6475 1,
6476 &exi,
6477 &exo,
6478 &exp_d,
6479 &pt,
6480 &zq,
6481 &zd,
6482 n_embd,
6483 n_ff_exp,
6484 m.up_exps.qtype,
6485 rbu_d,
6486 )?;
6487 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6488 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6489 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6491 let pself = e.htod_i32(&pair_self)?;
6492 let y_down = matvec(
6493 2,
6494 &exi,
6495 &exo,
6496 &exp_d,
6497 &pself,
6498 &aq2,
6499 &ad2,
6500 n_ff_exp,
6501 n_embd,
6502 m.down_exps.qtype,
6503 m.down_exps.row_bytes,
6504 )?;
6505 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6507
6508 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6512 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6513 {
6514 let n_ff_sh = gate_shexp.out_features();
6515 let step_exact = cfg.step35.is_some();
6519 let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6520 let (sg_gate, sg_up) = if step_exact && t == 1 {
6521 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6522 Some(pair) => pair,
6523 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6524 }
6525 } else if verify_t {
6526 let mut fused = None;
6527 if crate::spec::spec_fused_t()
6528 && (2..=4).contains(&t)
6529 && e.uses_q8_1_fast(gate_shexp)
6530 && e.uses_q8_1_fast(up_shexp)
6531 {
6532 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6533 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6534 }
6535 match fused {
6536 Some(pair) => pair,
6537 None => (
6538 e.matmul_decode_exact(gate_shexp, z, t)?,
6539 e.matmul_decode_exact(up_shexp, z, t)?,
6540 ),
6541 }
6542 } else {
6543 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6544 };
6545 let mut sa = e.uninit(t * n_ff_sh)?;
6546 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6547 let sh = if verify_t {
6548 e.matmul_decode_exact(down_shexp, &sa, t)?
6549 } else {
6550 e.matmul(down_shexp, &sa, t)?
6551 };
6552 let g = match &m.gate_inp_shexp {
6557 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6558 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6559 }
6560 Some(gate_inp_shexp) => {
6561 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6562 let mut g = e.uninit(t)?;
6563 e.sigmoid(&gs, &mut g, t)?;
6564 g
6565 }
6566 None => e.htod(&vec![1.0f32; t])?,
6567 };
6568 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6569 }
6570 Ok(moe_out)
6571 }
6572
6573 #[allow(clippy::too_many_arguments)]
6575 #[allow(clippy::too_many_arguments)]
6576 fn moe_ffn_dev(
6577 e: &Engine,
6578 m: &MoeWeights,
6579 z: &CudaSlice<f32>,
6580 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6581 logits: &CudaSlice<f32>,
6582 t: usize,
6583 cfg: &ModelConfig,
6584 il: u16,
6585 max_block: usize,
6586 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6587 let moe = cfg.moe.as_ref().unwrap();
6588 let n_embd = cfg.n_embd as usize;
6589 let n_expert = moe.expert_count as usize;
6590 let n_used = moe.expert_used_count as usize;
6591 let n_ff_exp = moe.expert_ff_length as usize;
6592 debug_assert!(
6596 cfg.sigmoid_router().is_none(),
6597 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6598 );
6599 debug_assert!(
6600 !cfg.swiglu_clamped_at(il as u32),
6601 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6602 );
6603
6604 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6606 if m.has_macros {
6609 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6610 }
6611
6612 let mut moe_out = e.uninit(t * n_embd)?;
6614
6615 if let Some(dev) = m.dev_exps.as_ref() {
6618 let (rbg_d, rbu_d) = if dev.gu_il {
6621 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6622 (sxx, sxx)
6623 } else {
6624 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6625 };
6626 let q8 = moe_q8_enabled()
6627 && q8_expert_supported(m.gate_exps.qtype)
6628 && q8_expert_supported(m.up_exps.qtype)
6629 && q8_expert_supported(m.down_exps.qtype);
6630 let rows_arm = q8
6639 && t > 1
6640 && crate::spec::spec_m2()
6641 && n_ff_exp == 512
6642 && n_used <= 8
6643 && std::env::var("MEMRA_MOE_DEVQ8_GU")
6644 .map(|v| v.is_empty() || v == "v")
6645 .unwrap_or(true)
6646 && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6647 .map(|v| v.is_empty() || v == "w8h2v")
6648 .unwrap_or(true);
6649 let csr_mode = std::env::var("MEMRA_MOE_CSR")
6658 .ok()
6659 .and_then(|v| v.parse::<i32>().ok())
6660 .unwrap_or(1);
6661 let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6662 let csr_arm = rows_arm
6663 && csr_mode > 0
6664 && t <= 10
6665 && csr_qt(m.gate_exps.qtype)
6666 && csr_qt(m.up_exps.qtype)
6667 && csr_qt(m.down_exps.qtype);
6668 if csr_arm {
6669 if csr_mode == 2 {
6670 static ENGAGED: std::sync::Once = std::sync::Once::new();
6671 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6672 }
6673 let n_pairs = t * n_used;
6674 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6675 let act = e.moe_gate_up_silu8_dev_q8_csr(
6676 &dev.ptr_row,
6677 &sel_d,
6678 &zq,
6679 &zd,
6680 n_pairs,
6681 n_embd,
6682 n_ff_exp,
6683 n_used,
6684 n_expert,
6685 m.gate_exps.qtype,
6686 m.up_exps.qtype,
6687 rbg_d,
6688 rbu_d,
6689 )?;
6690 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6691 e.moe_down8_fma_dev_q8_rows(
6695 &dev.ptr_row,
6696 &sel_d,
6697 &w_d,
6698 &aq2,
6699 &ad2,
6700 &mut moe_out,
6701 t,
6702 n_ff_exp,
6703 n_embd,
6704 n_used,
6705 n_expert,
6706 m.down_exps.qtype,
6707 m.down_exps.row_bytes,
6708 )?;
6709 if csr_mode == 2 {
6710 let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6712 &dev.ptr_row,
6713 &sel_d,
6714 &zq,
6715 &zd,
6716 t,
6717 n_embd,
6718 n_ff_exp,
6719 n_used,
6720 n_expert,
6721 m.gate_exps.qtype,
6722 m.up_exps.qtype,
6723 rbg_d,
6724 rbu_d,
6725 &m.dev_macros,
6726 )?;
6727 let mut out_r = e.uninit(t * n_embd)?;
6728 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6729 e.moe_down8_fma_dev_q8_rows(
6730 &dev.ptr_row,
6731 &sel_d,
6732 &w_d,
6733 &aq2r,
6734 &ad2r,
6735 &mut out_r,
6736 t,
6737 n_ff_exp,
6738 n_embd,
6739 n_used,
6740 n_expert,
6741 m.down_exps.qtype,
6742 m.down_exps.row_bytes,
6743 )?;
6744 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6745 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6746 let ba = a1
6747 .iter()
6748 .zip(&a2)
6749 .filter(|(x, y)| x.to_bits() != y.to_bits())
6750 .count();
6751 let bo = o1
6752 .iter()
6753 .zip(&o2)
6754 .filter(|(x, y)| x.to_bits() != y.to_bits())
6755 .count();
6756 if ba + bo > 0 {
6757 eprintln!(
6758 "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6759 a1.len(),
6760 o1.len()
6761 );
6762 let sel_h = e.dtoh_i32(&sel_d)?;
6764 let mut shown = 0;
6765 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6766 if x.to_bits() != y.to_bits() && shown < 4 {
6767 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6768 let ex = sel_h[p];
6769 let npx = sel_h.iter().filter(|&&v| v == ex).count();
6770 eprintln!(
6771 " ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6772 );
6773 shown += 1;
6774 }
6775 }
6776 std::process::exit(3);
6777 }
6778 }
6779 } else if rows_arm {
6780 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6783 use std::sync::atomic::{AtomicU64, Ordering};
6784 static PAIRS: AtomicU64 = AtomicU64::new(0);
6785 static UNIQ: AtomicU64 = AtomicU64::new(0);
6786 static CALLS: AtomicU64 = AtomicU64::new(0);
6787 let sel_h = e.dtoh_i32(&sel_d)?;
6788 let mut u: Vec<i32> = sel_h.clone();
6789 u.sort_unstable();
6790 u.dedup();
6791 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6792 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6793 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6794 if c % 480 == 0 {
6795 let p = PAIRS.load(Ordering::Relaxed);
6796 let q = UNIQ.load(Ordering::Relaxed);
6797 eprintln!(
6798 "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6799 q as f64 / p as f64
6800 );
6801 }
6802 }
6803 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6804 let act = e.moe_gate_up_silu8_dev_q8_rows(
6805 &dev.ptr_row,
6806 &sel_d,
6807 &zq,
6808 &zd,
6809 t,
6810 n_embd,
6811 n_ff_exp,
6812 n_used,
6813 n_expert,
6814 m.gate_exps.qtype,
6815 m.up_exps.qtype,
6816 rbg_d,
6817 rbu_d,
6818 &m.dev_macros,
6819 )?;
6820 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6821 e.moe_down8_fma_dev_q8_rows(
6822 &dev.ptr_row,
6823 &sel_d,
6824 &w_d,
6825 &aq2,
6826 &ad2,
6827 &mut moe_out,
6828 t,
6829 n_ff_exp,
6830 n_embd,
6831 n_used,
6832 n_expert,
6833 m.down_exps.qtype,
6834 m.down_exps.row_bytes,
6835 )?;
6836 } else {
6837 for tok in 0..t {
6838 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6839 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6840 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6841 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6842 if q8 {
6843 let (zq, zd) = match (t, zq8) {
6844 (1, Some((q, d))) => (q.clone(), d.clone()),
6845 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6846 };
6847 let act = e.moe_gate_up_silu8_dev_q8(
6848 &dev.ptr_row,
6849 &selt,
6850 &zq,
6851 &zd,
6852 n_embd,
6853 n_ff_exp,
6854 n_used,
6855 n_expert,
6856 m.gate_exps.qtype,
6857 m.up_exps.qtype,
6858 rbg_d,
6859 rbu_d,
6860 &m.dev_macros,
6861 )?;
6862 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6863 e.moe_down8_fma_dev_q8(
6864 &dev.ptr_row,
6865 &selt,
6866 &wt,
6867 &aq2,
6868 &ad2,
6869 &mut dst,
6870 n_ff_exp,
6871 n_embd,
6872 n_used,
6873 n_expert,
6874 m.down_exps.qtype,
6875 m.down_exps.row_bytes,
6876 )?;
6877 } else {
6878 let act = e.moe_gate_up_silu8_dev(
6879 &dev.ptr_row,
6880 &selt,
6881 &zt,
6882 n_embd,
6883 n_ff_exp,
6884 n_used,
6885 n_expert,
6886 m.gate_exps.qtype,
6887 m.up_exps.qtype,
6888 rbg_d,
6889 rbu_d,
6890 &m.dev_macros,
6891 )?;
6892 e.moe_down8_fma_dev(
6893 &dev.ptr_row,
6894 &selt,
6895 &wt,
6896 &act,
6897 &mut dst,
6898 n_ff_exp,
6899 n_embd,
6900 n_used,
6901 n_expert,
6902 m.down_exps.qtype,
6903 m.down_exps.row_bytes,
6904 )?;
6905 }
6906 }
6907 }
6908 } else {
6909 let q8 = moe_q8_enabled()
6916 && q8_expert_supported(m.gate_exps.qtype)
6917 && q8_expert_supported(m.up_exps.qtype)
6918 && q8_expert_supported(m.down_exps.qtype);
6919 e.with_moe_cache(max_block, |c, eng| {
6920 let row = c
6921 .layer_dev_row(il, n_expert, eng)?
6922 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
6923 for tok in 0..t {
6924 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6925 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6926 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6927 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6928 if q8 {
6929 let (zq, zd) = match (t, zq8) {
6930 (1, Some((q, d))) => (q.clone(), d.clone()),
6931 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
6932 };
6933 let act = eng.moe_gate_up_silu8_dev_q8(
6934 row,
6935 &selt,
6936 &zq,
6937 &zd,
6938 n_embd,
6939 n_ff_exp,
6940 n_used,
6941 n_expert,
6942 m.gate_exps.qtype,
6943 m.up_exps.qtype,
6944 m.gate_exps.row_bytes,
6945 m.up_exps.row_bytes,
6946 &m.dev_macros,
6947 )?;
6948 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
6949 eng.moe_down8_fma_dev_q8(
6950 row,
6951 &selt,
6952 &wt,
6953 &aq2,
6954 &ad2,
6955 &mut dst,
6956 n_ff_exp,
6957 n_embd,
6958 n_used,
6959 n_expert,
6960 m.down_exps.qtype,
6961 m.down_exps.row_bytes,
6962 )?;
6963 } else {
6964 let act = eng.moe_gate_up_silu8_dev(
6965 row,
6966 &selt,
6967 &zt,
6968 n_embd,
6969 n_ff_exp,
6970 n_used,
6971 n_expert,
6972 m.gate_exps.qtype,
6973 m.up_exps.qtype,
6974 m.gate_exps.row_bytes,
6975 m.up_exps.row_bytes,
6976 &m.dev_macros,
6977 )?;
6978 eng.moe_down8_fma_dev(
6979 row,
6980 &selt,
6981 &wt,
6982 &act,
6983 &mut dst,
6984 n_ff_exp,
6985 n_embd,
6986 n_used,
6987 n_expert,
6988 m.down_exps.qtype,
6989 m.down_exps.row_bytes,
6990 )?;
6991 }
6992 }
6993 c.hits += (t * 3 * n_used) as u64;
6995 Ok(())
6996 })?;
6997 }
6998
6999 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7004 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7005 {
7006 let n_ff_sh = gate_shexp.out_features();
7007 let verify_t = t > 1 && t < PRIME_MIN_T;
7010 let (sg_gate, sg_up) = if t == 1 {
7011 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
7012 Some(pair) => pair,
7013 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
7014 }
7015 } else if verify_t {
7016 let mut fused = None;
7020 if crate::spec::spec_fused_t()
7021 && (2..=4).contains(&t)
7022 && e.uses_q8_1_fast(gate_shexp)
7023 && e.uses_q8_1_fast(up_shexp)
7024 {
7025 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7026 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
7027 }
7028 match fused {
7029 Some(pair) => pair,
7030 None => (
7031 e.matmul_decode_exact(gate_shexp, z, t)?,
7032 e.matmul_decode_exact(up_shexp, z, t)?,
7033 ),
7034 }
7035 } else {
7036 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
7037 };
7038 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
7040 let sh = if verify_t {
7041 e.matmul_decode_exact(down_shexp, &sa, t)?
7042 } else {
7043 e.matmul(down_shexp, &sa, t)?
7044 };
7045 let g = match &m.gate_inp_shexp {
7049 Some(gate_inp_shexp) => {
7050 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7053 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7054 } else {
7055 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7056 let mut g = e.uninit(t)?;
7057 e.sigmoid(&gs, &mut g, t)?;
7058 g
7059 }
7060 }
7061 None => e.htod(&vec![1.0f32; t])?,
7062 };
7063 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7064 }
7065
7066 Ok(moe_out)
7067 }
7068
7069 #[allow(clippy::too_many_arguments)]
7079 #[allow(clippy::too_many_arguments)]
7082 fn moe_gdec_token_q8(
7083 e: &Engine,
7084 m: &MoeWeights,
7085 il: u16,
7086 max_block: usize,
7087 zq: &CudaSlice<i8>,
7088 zd: &CudaSlice<f32>,
7089 sel: &[u32],
7090 w: &[f32],
7091 moe_out: &mut CudaSlice<f32>,
7092 tok: usize,
7093 n_embd: usize,
7094 n_ff_exp: usize,
7095 n_used: usize,
7096 ) -> Result<bool, Box<dyn std::error::Error>> {
7097 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7098 use cudarc::driver::DevicePtr;
7099 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7100 let mut g = [0u64; 8];
7101 let mut u = [0u64; 8];
7102 let mut d = [0u64; 8];
7103 for (j, &ex) in sel.iter().enumerate() {
7104 let ex = ex as u16;
7105 let (Some(sg), Some(su), Some(sd)) = (
7106 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7107 c.resident(BlockId::new(il, PROJ_UP, ex)),
7108 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7109 ) else {
7110 return Ok(None);
7111 };
7112 let __s = eng.stream();
7113 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7114 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7115 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7116 g[j] = pg as u64;
7117 u[j] = pu as u64;
7118 d[j] = pd as u64;
7119 }
7120 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7121 for &ex in sel {
7122 let ex = ex as u16;
7123 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7124 c.note_profile_hit(BlockId::new(il, proj, ex));
7125 }
7126 }
7127 }
7128 c.hits += (3 * n_used) as u64;
7129 Ok(Some((g, u, d)))
7130 })?;
7131 let Some((g, u, d)) = ptrs else {
7132 return Ok(false);
7133 };
7134 let mut wv = [0f32; 8];
7135 wv[..n_used].copy_from_slice(w);
7136 let act = e.moe_gate_up_silu8_q8(
7137 crate::WPtr8(g),
7138 crate::WPtr8(u),
7139 zq,
7140 zd,
7141 n_embd,
7142 n_ff_exp,
7143 n_used,
7144 m.gate_exps.qtype,
7145 m.up_exps.qtype,
7146 m.gate_exps.row_bytes,
7147 m.up_exps.row_bytes,
7148 )?;
7149 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7151 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7152 e.moe_down8_fma_q8(
7153 crate::WPtr8(d),
7154 crate::F32x8(wv),
7155 &aq2,
7156 &ad2,
7157 &mut dst,
7158 n_ff_exp,
7159 n_embd,
7160 n_used,
7161 m.down_exps.qtype,
7162 m.down_exps.row_bytes,
7163 )?;
7164 Ok(true)
7165 }
7166
7167 fn moe_gdec_token(
7168 e: &Engine,
7169 m: &MoeWeights,
7170 il: u16,
7171 max_block: usize,
7172 zt: &cudarc::driver::CudaView<f32>,
7173 sel: &[u32],
7174 w: &[f32],
7175 moe_out: &mut CudaSlice<f32>,
7176 tok: usize,
7177 n_embd: usize,
7178 n_ff_exp: usize,
7179 n_used: usize,
7180 ) -> Result<bool, Box<dyn std::error::Error>> {
7181 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7182 use cudarc::driver::DevicePtr;
7183 let ptrs = e.with_moe_cache(max_block, |c, eng| {
7185 let mut g = [0u64; 8];
7186 let mut u = [0u64; 8];
7187 let mut d = [0u64; 8];
7188 for (j, &ex) in sel.iter().enumerate() {
7189 let ex = ex as u16;
7190 let (Some(sg), Some(su), Some(sd)) = (
7191 c.resident(BlockId::new(il, PROJ_GATE, ex)),
7192 c.resident(BlockId::new(il, PROJ_UP, ex)),
7193 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7194 ) else {
7195 return Ok(None);
7196 };
7197 let __s = eng.stream();
7198 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7199 let (pu, _e1) = c.slot(su).device_ptr(&__s);
7200 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7201 g[j] = pg as u64;
7202 u[j] = pu as u64;
7203 d[j] = pd as u64;
7204 }
7205 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7206 for &ex in sel {
7207 let ex = ex as u16;
7208 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7209 c.note_profile_hit(BlockId::new(il, proj, ex));
7210 }
7211 }
7212 }
7213 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
7215 })?;
7216 let Some((g, u, d)) = ptrs else {
7217 return Ok(false);
7218 };
7219 let mut wv = [0f32; 8];
7220 wv[..n_used].copy_from_slice(w);
7221 let act = e.moe_gate_up_silu8(
7223 crate::WPtr8(g),
7224 crate::WPtr8(u),
7225 zt,
7226 n_embd,
7227 n_ff_exp,
7228 n_used,
7229 m.gate_exps.qtype,
7230 m.up_exps.qtype,
7231 m.gate_exps.row_bytes,
7232 m.up_exps.row_bytes,
7233 )?;
7234 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7235 e.moe_down8_fma_into(
7236 crate::WPtr8(d),
7237 crate::F32x8(wv),
7238 &act,
7239 &mut dst,
7240 n_ff_exp,
7241 n_embd,
7242 n_used,
7243 m.down_exps.qtype,
7244 m.down_exps.row_bytes,
7245 )?;
7246 Ok(true)
7247 }
7248
7249 fn moe_cached_gemm_q8(
7254 e: &Engine,
7255 il: u16,
7256 proj: u8,
7257 ex: usize,
7258 m: &MoeWeights,
7259 max_block: usize,
7260 aq: &CudaSlice<i8>,
7261 ad: &CudaSlice<f32>,
7262 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7263 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7264 let exps = match proj {
7265 PROJ_GATE => &m.gate_exps,
7266 PROJ_UP => &m.up_exps,
7267 _ => &m.down_exps,
7268 };
7269 let layout = exps.expert_layout(ex);
7270 let id = BlockId::new(il, proj, ex as u16);
7271 let source = exps.expert_source(ex);
7272 e.with_moe_cache(max_block, |c, eng| {
7273 let slot = c.dispatch_source(id, source, eng)?;
7274 let DispatchSlot::Resident(sl) = slot;
7275 let buf = c.slot(sl);
7276 eng.qmatvec_expert_q8(
7277 buf,
7278 0..layout.len,
7279 aq,
7280 ad,
7281 1,
7282 exps.in_f,
7283 exps.out_f,
7284 layout.qtype,
7285 layout.row_bytes,
7286 )
7287 })
7288 }
7289
7290 fn moe_cached_gemm(
7291 e: &Engine,
7292 il: u16,
7293 proj: u8,
7294 ex: usize,
7295 m: &MoeWeights,
7296 max_block: usize,
7297 x: &cudarc::driver::CudaView<f32>,
7298 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7299 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7300 let exps = match proj {
7301 PROJ_GATE => &m.gate_exps,
7302 PROJ_UP => &m.up_exps,
7303 _ => &m.down_exps,
7304 };
7305 let layout = exps.expert_layout(ex);
7306 let id = BlockId::new(il, proj, ex as u16);
7307 let source = exps.expert_source(ex);
7308 e.with_moe_cache(max_block, |c, eng| {
7310 let slot = c.dispatch_source(id, source, eng)?;
7311 let DispatchSlot::Resident(sl) = slot;
7314 let buf = c.slot(sl);
7315 eng.qmatvec_view(
7316 buf,
7317 0..layout.len,
7318 x,
7319 1,
7320 exps.in_f,
7321 exps.out_f,
7322 layout.qtype,
7323 layout.row_bytes,
7324 )
7325 })
7326 }
7327
7328 fn moe_profile_admit_expert(
7332 e: &Engine,
7333 il: u16,
7334 ex: usize,
7335 m: &MoeWeights,
7336 max_block: usize,
7337 ) -> Result<(), Box<dyn std::error::Error>> {
7338 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7339 e.with_moe_cache(max_block, |cache, eng| {
7340 for (proj, exps) in [
7341 (PROJ_GATE, &m.gate_exps),
7342 (PROJ_UP, &m.up_exps),
7343 (PROJ_DOWN, &m.down_exps),
7344 ] {
7345 let id = BlockId::new(il, proj, ex as u16);
7346 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7347 }
7348 Ok(())
7349 })
7350 }
7351
7352 #[allow(clippy::too_many_arguments)]
7355 fn moe_frozen_gemm(
7356 e: &Engine,
7357 il: u16,
7358 proj: u8,
7359 ex: usize,
7360 m: &MoeWeights,
7361 max_block: usize,
7362 x: &cudarc::driver::CudaView<f32>,
7363 scratch: &mut Option<CudaSlice<u8>>,
7364 scratch_len: usize,
7365 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7366 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7367 let exps = match proj {
7368 PROJ_GATE => &m.gate_exps,
7369 PROJ_UP => &m.up_exps,
7370 _ => &m.down_exps,
7371 };
7372 let layout = exps.expert_layout(ex);
7373 let id = BlockId::new(il, proj, ex as u16);
7374 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7375 let Some(slot) = cache.resident(id) else {
7376 return Ok(None);
7377 };
7378 let buf = cache.slot(slot);
7379 Ok(Some(eng.qmatvec_view(
7380 buf,
7381 0..layout.len,
7382 x,
7383 1,
7384 exps.in_f,
7385 exps.out_f,
7386 layout.qtype,
7387 layout.row_bytes,
7388 )?))
7389 })? {
7390 return Ok(output);
7391 }
7392 if scratch.is_none() {
7393 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7394 }
7395 let scratch = scratch.as_mut().unwrap();
7396 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7397 e.qmatvec_view(
7398 scratch,
7399 0..layout.len,
7400 x,
7401 1,
7402 exps.in_f,
7403 exps.out_f,
7404 layout.qtype,
7405 layout.row_bytes,
7406 )
7407 }
7408
7409 fn moe_prefetch_expert(
7410 e: &Engine,
7411 il: u16,
7412 ex: usize,
7413 m: &MoeWeights,
7414 max_block: usize,
7415 keep: &[crate::moe_cache::BlockId],
7416 ) -> Result<(), Box<dyn std::error::Error>> {
7417 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7418 e.with_moe_cache(max_block, |c, eng| {
7419 for (proj, exps) in [
7420 (PROJ_GATE, &m.gate_exps),
7421 (PROJ_UP, &m.up_exps),
7422 (PROJ_DOWN, &m.down_exps),
7423 ] {
7424 let id = BlockId::new(il, proj, ex as u16);
7425 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7426 }
7427 Ok(())
7428 })
7429 }
7430
7431 fn moe_prefetch_disk_expert(
7434 e: &Engine,
7435 il: u16,
7436 ex: usize,
7437 m: &MoeWeights,
7438 max_block: usize,
7439 keep: &[crate::moe_cache::BlockId],
7440 ) -> Result<(), Box<dyn std::error::Error>> {
7441 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7442 e.with_moe_cache(max_block, |c, eng| {
7443 for (proj, exps) in [
7444 (PROJ_GATE, &m.gate_exps),
7445 (PROJ_UP, &m.up_exps),
7446 (PROJ_DOWN, &m.down_exps),
7447 ] {
7448 let source = exps.expert_source(ex);
7449 if let crate::model::ExpertSource::Disk { .. } = &source {
7450 let id = BlockId::new(il, proj, ex as u16);
7451 let _ = c.prefetch_source(id, source, keep, eng)?;
7452 }
7453 }
7454 Ok(())
7455 })
7456 }
7457
7458 #[inline]
7459 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7460 let _ = m.gate_exps.prefetch_expert_pages(ex);
7461 let _ = m.up_exps.prefetch_expert_pages(ex);
7462 let _ = m.down_exps.prefetch_expert_pages(ex);
7463 }
7464}
7465
7466impl HybridModel {
7483 #[allow(clippy::too_many_arguments)]
7487 fn moe_ffn_grouped_resident_q8(
7488 e: &Engine,
7489 m: &MoeWeights,
7490 z: &CudaSlice<f32>,
7491 t: usize,
7492 cfg: &ModelConfig,
7493 il: u16,
7494 sel_all: &[u32],
7495 w_all: &[f32],
7496 table: &CudaSlice<u64>,
7497 gu_il: bool,
7498 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7499 let moe = cfg.moe.as_ref().unwrap();
7500 let n_embd = cfg.n_embd as usize;
7501 let n_expert = moe.expert_count as usize;
7502 let n_used = moe.expert_used_count as usize;
7503 let n_ff_exp = moe.expert_ff_length as usize;
7504 let n_pairs = t * n_used;
7505 debug_assert_eq!(sel_all.len(), n_pairs);
7506 debug_assert_eq!(w_all.len(), n_pairs);
7507 debug_assert!(
7508 m.gate_exps.macros.is_none()
7509 && m.up_exps.macros.is_none()
7510 && m.down_exps.macros.is_none(),
7511 "resident grouped q8 does not fold per-expert macro scales",
7512 );
7513
7514 if !cfg.swiglu_clamped_at(il as u32) {
7520 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7521 let sel_d = e.htod_i32(&sel)?;
7522 let w_d = e.htod(w_all)?;
7523 let (gate_row_bytes, up_row_bytes) = if gu_il {
7524 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7525 (combined, combined)
7526 } else {
7527 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7528 };
7529 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7530 let act = e.moe_gate_up_silu8_dev_q8_rows(
7531 table,
7532 &sel_d,
7533 &zq,
7534 &zd,
7535 t,
7536 n_embd,
7537 n_ff_exp,
7538 n_used,
7539 n_expert,
7540 m.gate_exps.qtype,
7541 m.up_exps.qtype,
7542 gate_row_bytes,
7543 up_row_bytes,
7544 &m.dev_macros,
7545 )?;
7546 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7547 let mut moe_out = e.uninit(t * n_embd)?;
7548 e.moe_down8_fma_dev_q8_rows_g(
7549 table,
7550 &sel_d,
7551 &w_d,
7552 &aq2,
7553 &ad2,
7554 &mut moe_out,
7555 t,
7556 n_ff_exp,
7557 n_embd,
7558 n_used,
7559 n_expert,
7560 m.down_exps.qtype,
7561 m.down_exps.row_bytes,
7562 )?;
7563
7564 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7565 let mut counts = vec![0usize; n_expert];
7566 for &expert in sel_all {
7567 counts[expert as usize] += 1;
7568 }
7569 let mut sizes: Vec<usize> =
7570 counts.into_iter().filter(|&count| count != 0).collect();
7571 sizes.sort_unstable();
7572 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7573 println!(
7574 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7575 m_e: min={} median={} mean={mean:.1} max={}",
7576 sizes.len(),
7577 n_expert,
7578 sizes.first().copied().unwrap_or(0),
7579 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7580 sizes.last().copied().unwrap_or(0),
7581 );
7582 }
7583 return Ok(moe_out);
7584 }
7585
7586 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7590 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7591 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7592 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7593
7594 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7595 for (pair, &expert) in pair_ex.iter().enumerate() {
7596 by_expert[expert as usize].push(pair as i32);
7597 }
7598
7599 let pair_tok_d = e.htod_i32(&pair_tok)?;
7600 let pair_ex_d = e.htod_i32(&pair_ex)?;
7601 let pair_w_d = e.htod(w_all)?;
7602 let tok_off_d = e.htod_i32(&tok_off)?;
7603 let tok_ids_d = e.htod_i32(&tok_ids)?;
7604
7605 let matvec = |proj: i32,
7606 pair_rows: &CudaSlice<i32>,
7607 aq: &CudaSlice<i8>,
7608 ad: &CudaSlice<f32>,
7609 in_f: usize,
7610 out_f: usize,
7611 qtype: i32,
7612 row_bytes: usize|
7613 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7614 e.moe_pairs_matvec_q8(
7615 table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7616 row_bytes,
7617 )
7618 };
7619
7620 let (gate_row_bytes, up_row_bytes) = if gu_il {
7621 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7622 (combined, combined)
7623 } else {
7624 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7625 };
7626 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7627 let gate = matvec(
7628 0,
7629 &pair_tok_d,
7630 &zq,
7631 &zd,
7632 n_embd,
7633 n_ff_exp,
7634 m.gate_exps.qtype,
7635 gate_row_bytes,
7636 )?;
7637 let up = matvec(
7638 1,
7639 &pair_tok_d,
7640 &zq,
7641 &zd,
7642 n_embd,
7643 n_ff_exp,
7644 m.up_exps.qtype,
7645 up_row_bytes,
7646 )?;
7647 let mut act = e.uninit(n_pairs * n_ff_exp)?;
7648 Self::ffn_act_lim(
7649 e,
7650 cfg,
7651 &gate,
7652 &up,
7653 1.0,
7654 1.0,
7655 cfg.clamp_exp_at(il as u32),
7656 &mut act,
7657 n_pairs * n_ff_exp,
7658 )?;
7659 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7660 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7661 let pair_self_d = e.htod_i32(&pair_self)?;
7662 let down = matvec(
7663 2,
7664 &pair_self_d,
7665 &aq2,
7666 &ad2,
7667 n_ff_exp,
7668 n_embd,
7669 m.down_exps.qtype,
7670 m.down_exps.row_bytes,
7671 )?;
7672 let mut moe_out = e.uninit(t * n_embd)?;
7673 e.moe_pairs_scatter(
7674 &down,
7675 &pair_w_d,
7676 &tok_off_d,
7677 &tok_ids_d,
7678 &mut moe_out,
7679 t,
7680 n_embd,
7681 )?;
7682
7683 if std::env::var("MEMRA_MOE_STATS").is_ok() {
7684 let mut sizes: Vec<usize> = by_expert
7685 .iter()
7686 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7687 .collect();
7688 sizes.sort_unstable();
7689 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7690 println!(
7691 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7692 m_e: min={} median={} mean={mean:.1} max={}",
7693 sizes.len(),
7694 n_expert,
7695 sizes.first().copied().unwrap_or(0),
7696 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7697 sizes.last().copied().unwrap_or(0),
7698 );
7699 }
7700 Ok(moe_out)
7701 }
7702
7703 fn moe_ffn_grouped_add_shared(
7704 e: &Engine,
7705 m: &MoeWeights,
7706 z: &CudaSlice<f32>,
7707 t: usize,
7708 cfg: &ModelConfig,
7709 il: u16,
7710 moe_out: &mut CudaSlice<f32>,
7711 ) -> Result<(), Box<dyn std::error::Error>> {
7712 let n_embd = cfg.n_embd as usize;
7713 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7714 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7715 {
7716 let n_ff_sh = gate_shexp.out_features();
7717 let sg_gate = e.matmul(gate_shexp, z, t)?;
7718 let sg_up = e.matmul(up_shexp, z, t)?;
7719 let mut sa = e.uninit(t * n_ff_sh)?;
7720 Self::ffn_act_lim(
7721 e,
7722 cfg,
7723 &sg_gate,
7724 &sg_up,
7725 1.0,
7726 1.0,
7727 cfg.clamp_shexp_at(il as u32),
7728 &mut sa,
7729 t * n_ff_sh,
7730 )?;
7731 let sh = e.matmul(down_shexp, &sa, t)?;
7732 let gate = match &m.gate_inp_shexp {
7733 Some(gate_inp_shexp) => {
7734 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7735 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7736 } else {
7737 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7738 let mut gate = e.uninit(t)?;
7739 e.sigmoid(&raw, &mut gate, t)?;
7740 gate
7741 }
7742 }
7743 None => e.htod(&vec![1.0f32; t])?,
7744 };
7745 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7746 }
7747 Ok(())
7748 }
7749
7750 pub(crate) fn moe_ffn_grouped(
7753 e: &Engine,
7754 m: &MoeWeights,
7755 z: &CudaSlice<f32>,
7756 t: usize,
7757 cfg: &ModelConfig,
7758 il: u16,
7759 max_block: usize,
7760 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7761 let moe = cfg.moe.as_ref().unwrap();
7762 let n_embd = cfg.n_embd as usize;
7763 let n_expert = moe.expert_count as usize;
7764 let n_used = moe.expert_used_count as usize;
7765 let n_ff_exp = moe.expert_ff_length as usize;
7766 let lim_exp = cfg.clamp_exp_at(il as u32);
7768
7769 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7773 if let Some(sig) = cfg.sigmoid_router() {
7774 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7775 }
7776 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7777 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7778 } else {
7779 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7780 };
7781 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7782 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7783 Self::trace_moe_input(e, il, t, n_embd, z)?;
7784
7785 let no_exp_macros = m.gate_exps.macros.is_none()
7790 && m.up_exps.macros.is_none()
7791 && m.down_exps.macros.is_none();
7792 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7793 m.has_uniform_expert_layout()
7794 && no_exp_macros
7795 && moe_q8_enabled()
7796 && q8_expert_supported(m.gate_exps.qtype)
7797 && q8_expert_supported(m.up_exps.qtype)
7798 && q8_expert_supported(m.down_exps.qtype)
7799 && moe_slab_enabled()
7800 && dev.dev == e.ctx().ordinal()
7801 });
7802 if let Some(dev) = resident_q8 {
7803 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7804 e,
7805 m,
7806 z,
7807 t,
7808 cfg,
7809 il,
7810 &sel_all,
7811 &w_all,
7812 &dev.ptr_row,
7813 dev.gu_il,
7814 )?;
7815 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7816 return Ok(moe_out);
7817 }
7818
7819 struct ExpertGroup {
7823 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
7827 let mut groups: Vec<ExpertGroup> = (0..n_expert)
7828 .map(|_| ExpertGroup {
7829 tok_indices: Vec::new(),
7830 slot_indices: Vec::new(),
7831 weights: Vec::new(),
7832 })
7833 .collect();
7834
7835 for tok in 0..t {
7836 for j in 0..n_used {
7837 let ex = sel_all[tok * n_used + j] as usize;
7838 let w = w_all[tok * n_used + j];
7839 groups[ex].tok_indices.push(tok as i32);
7840 groups[ex].slot_indices.push(j as i32);
7841 groups[ex].weights.push(w);
7842 }
7843 }
7844
7845 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7848 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
7852 let u_len = m.up_exps.max_expert_bytes();
7853 let d_len = m.down_exps.max_expert_bytes();
7854 let moe_q8 = m.has_uniform_expert_layout()
7855 && moe_q8_enabled()
7856 && q8_expert_supported(m.gate_exps.qtype)
7857 && q8_expert_supported(m.up_exps.qtype)
7858 && q8_expert_supported(m.down_exps.qtype);
7859 let slab_local = m
7862 .dev_exps
7863 .as_ref()
7864 .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
7865 let use_cache =
7866 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
7867 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
7870
7871 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
7873 (
7874 Some(e.alloc_u8(g_len)?),
7875 Some(e.alloc_u8(u_len)?),
7876 Some(e.alloc_u8(d_len)?),
7877 )
7878 } else {
7879 (None, None, None)
7880 };
7881
7882 let mut order: Vec<usize> = (0..n_expert)
7893 .filter(|&ex| !groups[ex].tok_indices.is_empty())
7894 .collect();
7895 order.sort_by(|&a, &b| {
7896 groups[b]
7897 .tok_indices
7898 .len()
7899 .cmp(&groups[a].tok_indices.len())
7900 .then(a.cmp(&b))
7901 });
7902 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
7904 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
7905 if worker_disk_prefetch {
7906 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
7907 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
7908 }
7909 }
7910 for (order_pos, &ex) in order.iter().enumerate() {
7911 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
7912 Self::moe_prefetch_host_expert(order[next], m);
7913 }
7914 if worker_disk_prefetch {
7915 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
7916 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7917 let keep = [
7918 BlockId::new(il, PROJ_GATE, ex as u16),
7919 BlockId::new(il, PROJ_UP, ex as u16),
7920 BlockId::new(il, PROJ_DOWN, ex as u16),
7921 ];
7922 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
7923 }
7924 }
7925 let grp = &groups[ex];
7926 let m_e = grp.tok_indices.len();
7927 m_dist.push(m_e);
7928 let gl = m.gate_exps.expert_layout(ex);
7929 let ul = m.up_exps.expert_layout(ex);
7930 let dl = m.down_exps.expert_layout(ex);
7931
7932 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
7936 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
7937 let dmac = m.down_exps.macro_scale(ex);
7938 let weight_d = if dmac == 1.0 {
7939 e.htod(&grp.weights)?
7940 } else {
7941 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
7942 e.htod(&scaled)?
7943 };
7944
7945 let mut gathered = e.zeros(m_e * n_embd)?;
7947 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
7948 let gv = gathered.slice(0..m_e * n_embd);
7949
7950 let y = if let Some(dev) = slab_local {
7953 let gate_start = ex * m.gate_exps.expert_stride;
7954 let up_start = ex * m.up_exps.expert_stride;
7955 let down_start = ex * m.down_exps.expert_stride;
7956 if grouped_q8 {
7957 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7958 let gate = e.qmatvec_expert_q8(
7959 &dev.gate,
7960 gate_start..gate_start + gl.len,
7961 &zq,
7962 &zd,
7963 m_e,
7964 m.gate_exps.in_f,
7965 m.gate_exps.out_f,
7966 gl.qtype,
7967 gl.row_bytes,
7968 )?;
7969 let up = e.qmatvec_expert_q8(
7970 &dev.up,
7971 up_start..up_start + ul.len,
7972 &zq,
7973 &zd,
7974 m_e,
7975 m.up_exps.in_f,
7976 m.up_exps.out_f,
7977 ul.qtype,
7978 ul.row_bytes,
7979 )?;
7980 let mut act = e.uninit(m_e * n_ff_exp)?;
7981 Self::ffn_act_lim(
7982 e,
7983 cfg,
7984 &gate,
7985 &up,
7986 m.gate_exps.macro_scale(ex),
7987 m.up_exps.macro_scale(ex),
7988 lim_exp,
7989 &mut act,
7990 m_e * n_ff_exp,
7991 )?;
7992 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
7993 e.qmatvec_expert_q8(
7994 &dev.down,
7995 down_start..down_start + dl.len,
7996 &aq2,
7997 &ad2,
7998 m_e,
7999 m.down_exps.in_f,
8000 m.down_exps.out_f,
8001 dl.qtype,
8002 dl.row_bytes,
8003 )?
8004 } else {
8005 let gate = e.qmatvec_view(
8006 &dev.gate,
8007 gate_start..gate_start + gl.len,
8008 &gv,
8009 m_e,
8010 m.gate_exps.in_f,
8011 m.gate_exps.out_f,
8012 gl.qtype,
8013 gl.row_bytes,
8014 )?;
8015 let up = e.qmatvec_view(
8016 &dev.up,
8017 up_start..up_start + ul.len,
8018 &gv,
8019 m_e,
8020 m.up_exps.in_f,
8021 m.up_exps.out_f,
8022 ul.qtype,
8023 ul.row_bytes,
8024 )?;
8025 let mut act = e.uninit(m_e * n_ff_exp)?;
8026 Self::ffn_act_lim(
8027 e,
8028 cfg,
8029 &gate,
8030 &up,
8031 m.gate_exps.macro_scale(ex),
8032 m.up_exps.macro_scale(ex),
8033 lim_exp,
8034 &mut act,
8035 m_e * n_ff_exp,
8036 )?;
8037 let actv = act.slice(0..m_e * n_ff_exp);
8038 e.qmatvec_view(
8039 &dev.down,
8040 down_start..down_start + dl.len,
8041 &actv,
8042 m_e,
8043 m.down_exps.in_f,
8044 m.down_exps.out_f,
8045 dl.qtype,
8046 dl.row_bytes,
8047 )?
8048 }
8049 } else if use_cache {
8050 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8051 if grouped_q8 {
8052 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8053 let gate = e.with_moe_cache(max_block, |cache, eng| {
8054 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8055 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8056 eng.qmatvec_expert_q8(
8057 cache.buf(slot),
8058 0..gl.len,
8059 &zq,
8060 &zd,
8061 m_e,
8062 m.gate_exps.in_f,
8063 m.gate_exps.out_f,
8064 gl.qtype,
8065 gl.row_bytes,
8066 )
8067 })?;
8068 let up = e.with_moe_cache(max_block, |cache, eng| {
8069 let id = BlockId::new(il, PROJ_UP, ex as u16);
8070 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8071 eng.qmatvec_expert_q8(
8072 cache.buf(slot),
8073 0..ul.len,
8074 &zq,
8075 &zd,
8076 m_e,
8077 m.up_exps.in_f,
8078 m.up_exps.out_f,
8079 ul.qtype,
8080 ul.row_bytes,
8081 )
8082 })?;
8083 let mut act = e.uninit(m_e * n_ff_exp)?;
8084 Self::ffn_act_lim(
8085 e,
8086 cfg,
8087 &gate,
8088 &up,
8089 m.gate_exps.macro_scale(ex),
8090 m.up_exps.macro_scale(ex),
8091 lim_exp,
8092 &mut act,
8093 m_e * n_ff_exp,
8094 )?;
8095 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8096 e.with_moe_cache(max_block, |cache, eng| {
8097 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8098 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8099 eng.qmatvec_expert_q8(
8100 cache.buf(slot),
8101 0..dl.len,
8102 &aq2,
8103 &ad2,
8104 m_e,
8105 m.down_exps.in_f,
8106 m.down_exps.out_f,
8107 dl.qtype,
8108 dl.row_bytes,
8109 )
8110 })?
8111 } else {
8112 let gate = e.with_moe_cache(max_block, |cache, eng| {
8113 let id = BlockId::new(il, PROJ_GATE, ex as u16);
8114 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8115 eng.qmatvec_view(
8116 cache.buf(slot),
8117 0..gl.len,
8118 &gv,
8119 m_e,
8120 m.gate_exps.in_f,
8121 m.gate_exps.out_f,
8122 gl.qtype,
8123 gl.row_bytes,
8124 )
8125 })?;
8126 let up = e.with_moe_cache(max_block, |cache, eng| {
8127 let id = BlockId::new(il, PROJ_UP, ex as u16);
8128 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8129 eng.qmatvec_view(
8130 cache.buf(slot),
8131 0..ul.len,
8132 &gv,
8133 m_e,
8134 m.up_exps.in_f,
8135 m.up_exps.out_f,
8136 ul.qtype,
8137 ul.row_bytes,
8138 )
8139 })?;
8140 let mut act = e.uninit(m_e * n_ff_exp)?;
8141 Self::ffn_act_lim(
8142 e,
8143 cfg,
8144 &gate,
8145 &up,
8146 m.gate_exps.macro_scale(ex),
8147 m.up_exps.macro_scale(ex),
8148 lim_exp,
8149 &mut act,
8150 m_e * n_ff_exp,
8151 )?;
8152 let actv = act.slice(0..m_e * n_ff_exp);
8153 e.with_moe_cache(max_block, |cache, eng| {
8154 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8155 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8156 eng.qmatvec_view(
8157 cache.buf(slot),
8158 0..dl.len,
8159 &actv,
8160 m_e,
8161 m.down_exps.in_f,
8162 m.down_exps.out_f,
8163 dl.qtype,
8164 dl.row_bytes,
8165 )
8166 })?
8167 }
8168 } else {
8169 let sg = scratch_g.as_mut().unwrap();
8170 let su = scratch_u.as_mut().unwrap();
8171 let sd = scratch_d.as_mut().unwrap();
8172 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8173 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8174 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8175 if grouped_q8 {
8176 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8177 let gate = e.qmatvec_expert_q8(
8178 sg,
8179 0..gl.len,
8180 &zq,
8181 &zd,
8182 m_e,
8183 m.gate_exps.in_f,
8184 m.gate_exps.out_f,
8185 gl.qtype,
8186 gl.row_bytes,
8187 )?;
8188 let up = e.qmatvec_expert_q8(
8189 su,
8190 0..ul.len,
8191 &zq,
8192 &zd,
8193 m_e,
8194 m.up_exps.in_f,
8195 m.up_exps.out_f,
8196 ul.qtype,
8197 ul.row_bytes,
8198 )?;
8199 let mut act = e.uninit(m_e * n_ff_exp)?;
8200 Self::ffn_act_lim(
8201 e,
8202 cfg,
8203 &gate,
8204 &up,
8205 m.gate_exps.macro_scale(ex),
8206 m.up_exps.macro_scale(ex),
8207 lim_exp,
8208 &mut act,
8209 m_e * n_ff_exp,
8210 )?;
8211 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8212 e.qmatvec_expert_q8(
8213 sd,
8214 0..dl.len,
8215 &aq2,
8216 &ad2,
8217 m_e,
8218 m.down_exps.in_f,
8219 m.down_exps.out_f,
8220 dl.qtype,
8221 dl.row_bytes,
8222 )?
8223 } else {
8224 let gate = e.qmatvec_view(
8225 sg,
8226 0..gl.len,
8227 &gv,
8228 m_e,
8229 m.gate_exps.in_f,
8230 m.gate_exps.out_f,
8231 gl.qtype,
8232 gl.row_bytes,
8233 )?;
8234 let up = e.qmatvec_view(
8235 su,
8236 0..ul.len,
8237 &gv,
8238 m_e,
8239 m.up_exps.in_f,
8240 m.up_exps.out_f,
8241 ul.qtype,
8242 ul.row_bytes,
8243 )?;
8244 let mut act = e.uninit(m_e * n_ff_exp)?;
8245 Self::ffn_act_lim(
8246 e,
8247 cfg,
8248 &gate,
8249 &up,
8250 m.gate_exps.macro_scale(ex),
8251 m.up_exps.macro_scale(ex),
8252 lim_exp,
8253 &mut act,
8254 m_e * n_ff_exp,
8255 )?;
8256 let actv = act.slice(0..m_e * n_ff_exp);
8257 e.qmatvec_view(
8258 sd,
8259 0..dl.len,
8260 &actv,
8261 m_e,
8262 m.down_exps.in_f,
8263 m.down_exps.out_f,
8264 dl.qtype,
8265 dl.row_bytes,
8266 )?
8267 }
8268 };
8269
8270 e.scatter_slot(
8272 &y,
8273 &tok_idx_d,
8274 &slot_idx_d,
8275 &weight_d,
8276 &mut slot_buf,
8277 &mut wbuf,
8278 n_embd,
8279 n_used,
8280 m_e,
8281 )?;
8282 }
8283
8284 let mut moe_out = e.zeros(t * n_embd)?;
8286 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8287
8288 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8290 m_dist.sort_unstable();
8291 let active = m_dist.len();
8292 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8293 let median = m_dist[active / 2];
8294 let max_m = *m_dist.last().unwrap();
8295 let min_m = m_dist[0];
8296 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8297 println!(
8298 "moe-grouped il={il} t={t} active={active}/{n_expert} \
8299 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8300 above_gemm_threshold(>=16)={above16}/{active}"
8301 );
8302 }
8303
8304 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8305 Ok(moe_out)
8306 }
8307
8308 pub(crate) fn moe_ffn_lockstep(
8315 &self,
8316 e: &Engine,
8317 m: &MoeWeights,
8318 zbatch: &CudaSlice<f32>,
8319 mrows: usize,
8320 il: u16,
8321 max_block: usize,
8322 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8323 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8324 let cfg = &self.cfg;
8325 let moe = cfg.moe.as_ref().unwrap();
8326 let n_embd = cfg.n_embd as usize;
8327 let n_expert = moe.expert_count as usize;
8328 let n_used = moe.expert_used_count as usize;
8329 let n_ff_exp = moe.expert_ff_length as usize;
8330 let lim_exp = cfg.clamp_exp_at(il as u32);
8332 let lim_shexp = cfg.clamp_shexp_at(il as u32);
8333
8334 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8335 if let Some(sig) = cfg.sigmoid_router() {
8336 Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8337 }
8338 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8339 Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8340 } else {
8341 Self::moe_route_cfg(
8342 e,
8343 &logits,
8344 mrows,
8345 n_expert,
8346 n_used,
8347 m.active_experts.as_deref(),
8348 )?
8349 };
8350 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8351
8352 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8354 Ok((0..n_expert)
8355 .map(|ex| {
8356 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8357 .into_iter()
8358 .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8359 })
8360 .collect())
8361 })?;
8362
8363 struct Group {
8364 rows: Vec<i32>,
8365 slots: Vec<i32>,
8366 weights: Vec<f32>,
8367 }
8368 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8369 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8370 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8371 Default::default();
8372 for row in 0..mrows {
8373 for j in 0..n_used {
8374 let ex = sel_all[row * n_used + j] as usize;
8375 let w = w_all[row * n_used + j];
8376 if resident_expert[ex] {
8377 let group = groups.entry(ex).or_insert_with(|| Group {
8378 rows: Vec::new(),
8379 slots: Vec::new(),
8380 weights: Vec::new(),
8381 });
8382 group.rows.push(row as i32);
8383 group.slots.push(j as i32);
8384 group.weights.push(w);
8385 } else {
8386 crate::cpu_experts::record_incomplete_gpu_residency(0);
8387 cpu_rows[row].push((ex, w));
8388 cpu_by_expert.entry(ex).or_default().push((row, w));
8389 }
8390 }
8391 }
8392
8393 let host_rows = e.dtoh(zbatch)?;
8399 let rows_ok = crate::cpu_experts::rows_supported();
8400 enum CpuPart {
8401 Single { row: usize },
8402 Rows { rows: Vec<usize> },
8403 }
8404 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8405 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8406 if rows_ok {
8407 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8408 .into_iter()
8409 .filter(|(_, rows)| rows.len() >= 2)
8410 .collect();
8411 shared.sort_by_key(|(ex, _)| *ex);
8412 for (ex, mut row_weights) in shared {
8413 row_weights.sort_by_key(|(row, _)| *row);
8414 let inputs: Vec<(&[f32], f32)> = row_weights
8415 .iter()
8416 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8417 .collect();
8418 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8419 .map_err(std::io::Error::other)?;
8420 for &(row, _) in &row_weights {
8421 rows_served.insert((row, ex));
8422 }
8423 tickets.push((
8424 CpuPart::Rows {
8425 rows: row_weights.iter().map(|&(row, _)| row).collect(),
8426 },
8427 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8428 ));
8429 }
8430 }
8431 for (row, selected) in cpu_rows.iter().enumerate() {
8432 let leftover: Vec<(usize, f32)> = selected
8433 .iter()
8434 .copied()
8435 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8436 .collect();
8437 if leftover.is_empty() {
8438 continue;
8439 }
8440 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8441 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8442 .map_err(std::io::Error::other)?;
8443 tickets.push((
8444 CpuPart::Single { row },
8445 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8446 ));
8447 }
8448
8449 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8450 let mut wbuf = e.zeros(mrows * n_used)?;
8451 let mut order: Vec<usize> = groups.keys().copied().collect();
8452 order.sort_by(|&a, &b| {
8453 groups[&b]
8454 .rows
8455 .len()
8456 .cmp(&groups[&a].rows.len())
8457 .then(a.cmp(&b))
8458 });
8459 for &ex in &order {
8460 let group = &groups[&ex];
8461 let m_e = group.rows.len();
8462 let gl = m.gate_exps.expert_layout(ex);
8463 let ul = m.up_exps.expert_layout(ex);
8464 let dl = m.down_exps.expert_layout(ex);
8465 let row_idx_d = e.htod_i32(&group.rows)?;
8466 let slot_idx_d = e.htod_i32(&group.slots)?;
8467 let dmac = m.down_exps.macro_scale(ex);
8468 let weight_d = if dmac == 1.0 {
8469 e.htod(&group.weights)?
8470 } else {
8471 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8472 e.htod(&scaled)?
8473 };
8474 let mut gathered = e.zeros(m_e * n_embd)?;
8475 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8476 let gv = gathered.slice(0..m_e * n_embd);
8477 let gate = e.with_moe_cache(max_block, |c, eng| {
8478 let slot = c
8479 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8480 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8481 eng.qmatvec_view(
8482 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8483 0..gl.len,
8484 &gv,
8485 m_e,
8486 m.gate_exps.in_f,
8487 m.gate_exps.out_f,
8488 gl.qtype,
8489 gl.row_bytes,
8490 )
8491 })?;
8492 let up = e.with_moe_cache(max_block, |c, eng| {
8493 let slot = c
8494 .resident(BlockId::new(il, PROJ_UP, ex as u16))
8495 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8496 eng.qmatvec_view(
8497 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8498 0..ul.len,
8499 &gv,
8500 m_e,
8501 m.up_exps.in_f,
8502 m.up_exps.out_f,
8503 ul.qtype,
8504 ul.row_bytes,
8505 )
8506 })?;
8507 let mut act = e.zeros(m_e * n_ff_exp)?;
8508 Self::ffn_act_lim(
8509 e,
8510 cfg,
8511 &gate,
8512 &up,
8513 m.gate_exps.macro_scale(ex),
8514 m.up_exps.macro_scale(ex),
8515 lim_exp,
8516 &mut act,
8517 m_e * n_ff_exp,
8518 )?;
8519 let actv = act.slice(0..m_e * n_ff_exp);
8520 let y = e.with_moe_cache(max_block, |c, eng| {
8521 let slot = c
8522 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8523 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8524 eng.qmatvec_view(
8525 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8526 0..dl.len,
8527 &actv,
8528 m_e,
8529 m.down_exps.in_f,
8530 m.down_exps.out_f,
8531 dl.qtype,
8532 dl.row_bytes,
8533 )
8534 })?;
8535 e.scatter_slot(
8536 &y,
8537 &row_idx_d,
8538 &slot_idx_d,
8539 &weight_d,
8540 &mut slot_buf,
8541 &mut wbuf,
8542 n_embd,
8543 n_used,
8544 m_e,
8545 )?;
8546 }
8547 let mut moe_out = e.zeros(mrows * n_embd)?;
8548 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8549
8550 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8552 for (part, ticket) in tickets {
8553 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8554 let mut add_row = |row: usize, chunk: &[f32]| {
8555 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8556 for (accumulator, value) in sum.iter_mut().zip(chunk) {
8557 *accumulator += value;
8558 }
8559 };
8560 match part {
8561 CpuPart::Single { row } => add_row(row, &cpu_output),
8562 CpuPart::Rows { rows } => {
8563 for (slot, row) in rows.into_iter().enumerate() {
8564 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8565 }
8566 }
8567 }
8568 }
8569 for (row, sum) in row_sums.into_iter().enumerate() {
8570 let Some(sum) = sum else { continue };
8571 let cpu_output = e.htod(&sum)?;
8572 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8573 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8574 }
8575
8576 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8577 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8578 {
8579 let n_ff_sh = gate_shexp.out_features();
8580 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8581 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8582 let mut sa = e.zeros(mrows * n_ff_sh)?;
8583 Self::ffn_act_lim(
8584 e,
8585 cfg,
8586 &sg_gate,
8587 &sg_up,
8588 1.0,
8589 1.0,
8590 lim_shexp,
8591 &mut sa,
8592 mrows * n_ff_sh,
8593 )?;
8594 let sh = e.matmul(down_shexp, &sa, mrows)?;
8595 let g = match &m.gate_inp_shexp {
8598 Some(gate_inp_shexp) => {
8599 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8600 }
8601 None => e.htod(&vec![1.0f32; mrows])?,
8602 };
8603 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8604 }
8605
8606 Ok(moe_out)
8607 }
8608}
8609
8610impl HybridModel {
8616 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8618 let g = self.cfg.gemma4.as_ref().unwrap();
8619 let swa = g.swa_pattern[il];
8620 let hd = if swa {
8621 g.key_length_swa
8622 } else {
8623 g.key_length_global
8624 } as usize;
8625 (
8629 hd,
8630 g.head_count_kv[il] as usize,
8631 self.cfg.n_head as usize,
8632 if swa {
8633 g.rope_base_swa
8634 } else {
8635 g.rope_base_global
8636 },
8637 1.0,
8638 swa,
8639 )
8640 }
8641
8642 pub(crate) fn gemma4_suppress(
8646 &self,
8647 e: &Engine,
8648 ld: &mut CudaSlice<f32>,
8649 t: usize,
8650 ) -> Result<(), Box<dyn std::error::Error>> {
8651 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8652 #[cfg(debug_assertions)]
8657 crate::debug_assert_tensor_stream_device(
8658 ids,
8659 &e.stream(),
8660 "gemma4_suppress.suppress_d",
8661 );
8662 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8663 }
8664 Ok(())
8665 }
8666
8667 #[allow(clippy::too_many_arguments)]
8672 fn gemma4_attn_prime(
8673 &self,
8674 e: &Engine,
8675 fa: &crate::hybrid::FullAttnLayer,
8676 il: usize,
8677 h: &CudaSlice<f32>,
8678 pos_d: &CudaSlice<i32>,
8679 t: usize,
8680 cache: Option<&mut Cache>,
8681 island: Option<&CudaSlice<i32>>,
8682 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8683 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8684 let eps = self.cfg.rms_eps;
8685 let aux = self.gemma4_aux.as_ref().unwrap();
8686 let ones = aux.ones(e);
8687 #[cfg(debug_assertions)]
8688 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8689
8690 e.mmq_act_begin();
8693 let q0 = e.matmul(&fa.wq, h, t)?; if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8695 let v = e.dtoh(&q0)?;
8696 let nan = v.iter().filter(|x| x.is_nan()).count();
8697 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8698 eprintln!(
8699 "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
8700 v.len()
8701 );
8702 }
8703 let k0 = e.matmul(&fa.wk, h, t)?; let v0 = if swa {
8707 e.matmul(&fa.wv, h, t)?
8708 } else {
8709 e.clone_dtod(&k0)?
8710 };
8711 if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8712 for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
8713 let v = e.dtoh(buf)?;
8714 let nan = v.iter().filter(|x| x.is_nan()).count();
8715 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8716 eprintln!(
8717 "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
8718 v.len()
8719 );
8720 }
8721 }
8722
8723 let mut q = e.uninit(t * nh * hd)?;
8724 let mut k = e.uninit(t * nkv * hd)?;
8725 let mut v = e.uninit(t * nkv * hd)?;
8727 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8731 let emit = island.is_none()
8734 && t >= 16
8735 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8736 && *EMIT.get_or_init(|| {
8737 std::env::var("MEMRA_FA_EMIT")
8738 .map(|s| s != "0")
8739 .unwrap_or(true)
8740 });
8741 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8742 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8743 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8744 let v_f16 = emit
8747 && crate::fa_f16pv_on()
8748 && match hd {
8749 512 => true,
8750 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8751 _ => false,
8752 };
8753 if emit {
8754 e.rms_norm_qkv_w4b(
8755 &q0,
8756 &k0,
8757 &v0,
8758 fa.q_norm.float_data(),
8759 fa.k_norm.float_data(),
8760 ones,
8761 &mut q,
8762 &mut k,
8763 &mut v,
8764 &mut vb,
8765 hd,
8766 nh * t,
8767 nkv * t,
8768 eps,
8769 v_f16,
8770 )?;
8771 } else {
8772 e.rms_norm_qkv(
8773 &q0,
8774 &k0,
8775 &v0,
8776 fa.q_norm.float_data(),
8777 fa.k_norm.float_data(),
8778 ones,
8779 &mut q,
8780 &mut k,
8781 &mut v,
8782 hd,
8783 nh * t,
8784 nkv * t,
8785 eps,
8786 )?;
8787 }
8788
8789 let ff = if swa {
8790 None
8791 } else {
8792 Some(
8793 aux.rope_freqs(e)
8794 .expect("gemma4 global rope needs rope_freqs.weight"),
8795 )
8796 };
8797 #[cfg(debug_assertions)]
8798 if let Some(ff) = ff {
8799 crate::debug_assert_tensor_stream_device(
8800 ff,
8801 &e.stream(),
8802 "gemma4_attn_prime.rope_freqs",
8803 );
8804 }
8805 if emit {
8806 e.rope_neox2_bf16e(
8807 &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8808 )?;
8809 } else {
8810 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8811 }
8812
8813 if let Some(cache) = cache {
8814 let kvl = cache.kv[il].as_mut().unwrap();
8815 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8816 e.append_kv_quantized_rows(
8817 &k,
8818 &v,
8819 &mut kvl.k,
8820 &mut kvl.v,
8821 kvl.len,
8822 t,
8823 kvl.kv_dim_k,
8824 kvl.kv_dim_v,
8825 kvl.k_tok_bytes,
8826 kvl.v_tok_bytes,
8827 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
8828 )?;
8829 kvl.len += t;
8830 }
8831 let mut attn = e.zeros(t * nh * hd)?;
8832 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8836 if let Some(span) = island {
8837 let w = if swa && t > win { win } else { 0 };
8842 e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
8843 } else if swa && t > win {
8844 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8845 if emit {
8846 e.fa_prefill_w_pre(
8847 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
8848 )?;
8849 } else {
8850 e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8851 }
8852 } else {
8853 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8854 }
8855 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8856 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8857 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
8858 if emit {
8859 e.fa_prefill_hd512_pre(
8860 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
8861 )?;
8862 } else {
8863 e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8864 }
8865 } else {
8866 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8867 }
8868 Ok(e.matmul(&fa.wo, &attn, t)?)
8869 }
8870
8871 fn gemma4_attn(
8873 &self,
8874 e: &Engine,
8875 fa: &crate::hybrid::FullAttnLayer,
8876 il: usize,
8877 h: &CudaSlice<f32>,
8878 pos_d: &CudaSlice<i32>,
8879 t: usize,
8880 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8881 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
8882 }
8883
8884 fn gemma4_moe_q8(
8889 &self,
8890 e: &Engine,
8891 m: &crate::hybrid::MoeWeights,
8892 bits: &crate::hybrid::Gemma4MoeBits,
8893 mq: &(CudaSlice<i8>, CudaSlice<f32>),
8894 router_in: &CudaSlice<f32>,
8895 t: usize,
8896 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8897 let cfg = &self.cfg;
8898 let moe = cfg.moe.as_ref().unwrap();
8899 let n_embd = cfg.n_embd as usize;
8900 let n_expert = moe.expert_count as usize;
8901 let n_used = moe.expert_used_count as usize;
8902 let n_ff_exp = moe.expert_ff_length as usize;
8903 let logits = if crate::router_kernel_on() {
8907 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8908 } else {
8909 e.matmul(&m.gate_inp, router_in, t)?
8910 };
8911 let dev = m.dev_exps.as_ref().unwrap();
8912 let (sel_d, w_d) =
8913 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8914 let (zq, zd) = mq;
8915 if t == 1 {
8916 let selv = sel_d.slice(0..n_used);
8917 let wv = w_d.slice(0..n_used);
8918 let act = e.moe_gate_up_gelu8_dev_q8(
8919 &dev.ptr_row,
8920 &selv,
8921 zq,
8922 zd,
8923 n_embd,
8924 n_ff_exp,
8925 n_used,
8926 n_expert,
8927 m.gate_exps.qtype,
8928 m.up_exps.qtype,
8929 m.gate_exps.row_bytes,
8930 m.up_exps.row_bytes,
8931 )?;
8932 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8933 let mut moe_out = e.uninit(n_embd)?;
8934 e.moe_down8_fma_dev_q8(
8935 &dev.ptr_row,
8936 &selv,
8937 &wv,
8938 &aq2,
8939 &ad2,
8940 &mut moe_out.slice_mut(0..n_embd),
8941 n_ff_exp,
8942 n_embd,
8943 n_used,
8944 n_expert,
8945 m.down_exps.qtype,
8946 m.down_exps.row_bytes,
8947 )?;
8948 return Ok(moe_out);
8949 }
8950 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
8951 let act = if csr {
8952 e.moe_gate_up_gelu8_dev_q8_csr(
8953 &dev.ptr_row,
8954 &sel_d,
8955 zq,
8956 zd,
8957 t * n_used,
8958 n_embd,
8959 n_ff_exp,
8960 n_used,
8961 n_expert,
8962 m.gate_exps.qtype,
8963 m.up_exps.qtype,
8964 m.gate_exps.row_bytes,
8965 m.up_exps.row_bytes,
8966 )?
8967 } else {
8968 e.moe_gate_up_gelu8_dev_q8_rows(
8969 &dev.ptr_row,
8970 &sel_d,
8971 zq,
8972 zd,
8973 t,
8974 n_embd,
8975 n_ff_exp,
8976 n_used,
8977 n_expert,
8978 m.gate_exps.qtype,
8979 m.up_exps.qtype,
8980 m.gate_exps.row_bytes,
8981 m.up_exps.row_bytes,
8982 )?
8983 };
8984 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
8985 let mut moe_out = e.uninit(t * n_embd)?;
8986 e.moe_down8_fma_dev_q8_rows_g(
8989 &dev.ptr_row,
8990 &sel_d,
8991 &w_d,
8992 &aq2,
8993 &ad2,
8994 &mut moe_out,
8995 t,
8996 n_ff_exp,
8997 n_embd,
8998 n_used,
8999 n_expert,
9000 m.down_exps.qtype,
9001 m.down_exps.row_bytes,
9002 )?;
9003 Ok(moe_out)
9004 }
9005
9006 fn gemma4_moe(
9010 &self,
9011 e: &Engine,
9012 m: &crate::hybrid::MoeWeights,
9013 bits: &crate::hybrid::Gemma4MoeBits,
9014 moe_in: &CudaSlice<f32>,
9015 router_in: &CudaSlice<f32>,
9016 t: usize,
9017 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9018 let cfg = &self.cfg;
9019 let moe = cfg.moe.as_ref().unwrap();
9020 let n_embd = cfg.n_embd as usize;
9021 let n_expert = moe.expert_count as usize;
9022 let n_used = moe.expert_used_count as usize;
9023 let n_ff_exp = moe.expert_ff_length as usize;
9024
9025 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
9029 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9030 } else {
9031 e.matmul(&m.gate_inp, router_in, t)?
9032 };
9033
9034 if t < PRIME_MIN_T
9039 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9040 && expert_dp4a_supported(m.gate_exps.qtype)
9041 && expert_dp4a_supported(m.up_exps.qtype)
9042 && expert_dp4a_supported(m.down_exps.qtype)
9043 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9044 {
9045 let dev = m.dev_exps.as_ref().unwrap();
9046 let (sel_d, w_d) =
9047 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9048 if t == 1 {
9049 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
9050 let selv = sel_d.slice(0..n_used);
9051 let wv = w_d.slice(0..n_used);
9052 let act = e.moe_gate_up_gelu8_dev_q8(
9053 &dev.ptr_row,
9054 &selv,
9055 &zq,
9056 &zd,
9057 n_embd,
9058 n_ff_exp,
9059 n_used,
9060 n_expert,
9061 m.gate_exps.qtype,
9062 m.up_exps.qtype,
9063 m.gate_exps.row_bytes,
9064 m.up_exps.row_bytes,
9065 )?;
9066 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9067 let mut moe_out = e.uninit(n_embd)?;
9068 e.moe_down8_fma_dev_q8(
9069 &dev.ptr_row,
9070 &selv,
9071 &wv,
9072 &aq2,
9073 &ad2,
9074 &mut moe_out.slice_mut(0..n_embd),
9075 n_ff_exp,
9076 n_embd,
9077 n_used,
9078 n_expert,
9079 m.down_exps.qtype,
9080 m.down_exps.row_bytes,
9081 )?;
9082 return Ok(moe_out);
9083 }
9084 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9089 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9090 let act = if csr {
9091 e.moe_gate_up_gelu8_dev_q8_csr(
9092 &dev.ptr_row,
9093 &sel_d,
9094 &zq,
9095 &zd,
9096 t * n_used,
9097 n_embd,
9098 n_ff_exp,
9099 n_used,
9100 n_expert,
9101 m.gate_exps.qtype,
9102 m.up_exps.qtype,
9103 m.gate_exps.row_bytes,
9104 m.up_exps.row_bytes,
9105 )?
9106 } else {
9107 e.moe_gate_up_gelu8_dev_q8_rows(
9108 &dev.ptr_row,
9109 &sel_d,
9110 &zq,
9111 &zd,
9112 t,
9113 n_embd,
9114 n_ff_exp,
9115 n_used,
9116 n_expert,
9117 m.gate_exps.qtype,
9118 m.up_exps.qtype,
9119 m.gate_exps.row_bytes,
9120 m.up_exps.row_bytes,
9121 )?
9122 };
9123 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9124 let mut moe_out = e.uninit(t * n_embd)?;
9125 e.moe_down8_fma_dev_q8_rows_g(
9126 &dev.ptr_row,
9127 &sel_d,
9128 &w_d,
9129 &aq2,
9130 &ad2,
9131 &mut moe_out,
9132 t,
9133 n_ff_exp,
9134 n_embd,
9135 n_used,
9136 n_expert,
9137 m.down_exps.qtype,
9138 m.down_exps.row_bytes,
9139 )?;
9140 return Ok(moe_out);
9141 }
9142
9143 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9144 for (i, &sx) in sel_all.iter().enumerate() {
9145 w_all[i] *= bits.per_expert_scale[sx as usize];
9146 }
9147
9148 if t >= PRIME_MIN_T
9152 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9153 && expert_dp4a_supported(m.gate_exps.qtype)
9154 && expert_dp4a_supported(m.up_exps.qtype)
9155 && expert_dp4a_supported(m.down_exps.qtype)
9156 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9157 {
9158 let dev = m.dev_exps.as_ref().unwrap();
9159 let n_pairs = t * n_used;
9160 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9161 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9162 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9163 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9164 let pt = e.htod_i32(&pair_tok)?;
9165 let pw = e.htod(&w_all)?;
9166 let toff = e.htod_i32(&tok_off)?;
9167 let tids = e.htod_i32(&tok_ids)?;
9168 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9169 for p in 0..n_pairs {
9170 by_ex[pair_ex[p] as usize].push(p as i32);
9171 }
9172 let mut ex_ids: Vec<i32> = Vec::new();
9173 let mut ex_off: Vec<i32> = vec![0];
9174 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9175 for (ex, list) in by_ex.iter().enumerate() {
9176 if list.is_empty() {
9177 continue;
9178 }
9179 ex_ids.push(ex as i32);
9180 ex_pairs.extend_from_slice(list);
9181 ex_off.push(ex_pairs.len() as i32);
9182 }
9183 let n_active = ex_ids.len();
9184 let exi = e.htod_i32(&ex_ids)?;
9185 let exo = e.htod_i32(&ex_off)?;
9186 let exp_d = e.htod_i32(&ex_pairs)?;
9187 if crate::moe_f16g_gemma_on()
9195 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9196 && f16g_proj_ok(m.up_exps.qtype, n_embd)
9197 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9198 {
9199 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9200 let csr_tok_d = e.htod_i32(&csr_tok)?;
9201 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9202 let g_csr = e.moe_f16_grouped(
9203 &dev.ptr_row,
9204 0,
9205 n_expert,
9206 &exi,
9207 &ex_off,
9208 &exo,
9209 &z_f16,
9210 &z_s,
9211 n_embd,
9212 n_ff_exp,
9213 n_active,
9214 n_pairs,
9215 m.gate_exps.qtype,
9216 m.gate_exps.row_bytes,
9217 )?;
9218 let u_csr = e.moe_f16_grouped(
9219 &dev.ptr_row,
9220 1,
9221 n_expert,
9222 &exi,
9223 &ex_off,
9224 &exo,
9225 &z_f16,
9226 &z_s,
9227 n_embd,
9228 n_ff_exp,
9229 n_active,
9230 n_pairs,
9231 m.up_exps.qtype,
9232 m.up_exps.row_bytes,
9233 )?;
9234 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9235 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9236 let d_csr = e.moe_f16_grouped(
9237 &dev.ptr_row,
9238 2,
9239 n_expert,
9240 &exi,
9241 &ex_off,
9242 &exo,
9243 &a_f16,
9244 &a_s,
9245 n_ff_exp,
9246 n_embd,
9247 n_active,
9248 n_pairs,
9249 m.down_exps.qtype,
9250 m.down_exps.row_bytes,
9251 )?;
9252 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9253 let mut moe_out = e.uninit(t * n_embd)?;
9254 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9255 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9256 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9257 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9258 eprintln!(
9259 "[f16g-debug] post-permute bad={} post-scatter bad={}",
9260 scan(&yd),
9261 scan(&mo)
9262 );
9263 }
9264 return Ok(moe_out);
9265 }
9266 let mma =
9269 n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9270 let (gate, up) = if mma {
9271 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9272 (
9273 e.mmq_iq_experts(
9274 &dev.ptr_row,
9275 0,
9276 n_expert,
9277 &exi,
9278 &exo,
9279 &exp_d,
9280 &pt,
9281 &z_scr,
9282 n_embd,
9283 n_ff_exp,
9284 n_active,
9285 n_pairs,
9286 t,
9287 m.gate_exps.qtype,
9288 m.gate_exps.row_bytes,
9289 )?,
9290 e.mmq_iq_experts(
9291 &dev.ptr_row,
9292 1,
9293 n_expert,
9294 &exi,
9295 &exo,
9296 &exp_d,
9297 &pt,
9298 &z_scr,
9299 n_embd,
9300 n_ff_exp,
9301 n_active,
9302 n_pairs,
9303 t,
9304 m.up_exps.qtype,
9305 m.up_exps.row_bytes,
9306 )?,
9307 )
9308 } else {
9309 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9310 (
9311 e.moe_pairs_matvec_q8_dec(
9312 &dev.ptr_row,
9313 0,
9314 &exi,
9315 &exo,
9316 &exp_d,
9317 &pt,
9318 &zq,
9319 &zd,
9320 n_embd,
9321 n_ff_exp,
9322 n_expert,
9323 n_active,
9324 n_pairs,
9325 m.gate_exps.qtype,
9326 m.gate_exps.row_bytes,
9327 )?,
9328 e.moe_pairs_matvec_q8_dec(
9329 &dev.ptr_row,
9330 1,
9331 &exi,
9332 &exo,
9333 &exp_d,
9334 &pt,
9335 &zq,
9336 &zd,
9337 n_embd,
9338 n_ff_exp,
9339 n_expert,
9340 n_active,
9341 n_pairs,
9342 m.up_exps.qtype,
9343 m.up_exps.row_bytes,
9344 )?,
9345 )
9346 };
9347 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9348 let pself = e.htod_i32(&pair_self)?;
9349 let y_down = if mma {
9361 let in_pad = n_ff_exp.div_ceil(256) * 256;
9362 let a_scr = if crate::moe_fuse_actq_on() {
9363 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9364 } else {
9365 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9366 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9367 };
9368 e.mmq_iq_experts(
9369 &dev.ptr_row,
9370 2,
9371 n_expert,
9372 &exi,
9373 &exo,
9374 &exp_d,
9375 &pself,
9376 &a_scr,
9377 in_pad,
9378 n_embd,
9379 n_active,
9380 n_pairs,
9381 n_pairs,
9382 m.down_exps.qtype,
9383 m.down_exps.row_bytes,
9384 )?
9385 } else {
9386 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9387 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9388 e.moe_pairs_matvec_q8_dec(
9389 &dev.ptr_row,
9390 2,
9391 &exi,
9392 &exo,
9393 &exp_d,
9394 &pself,
9395 &aq2,
9396 &ad2,
9397 n_ff_exp,
9398 n_embd,
9399 n_expert,
9400 n_active,
9401 n_pairs,
9402 m.down_exps.qtype,
9403 m.down_exps.row_bytes,
9404 )?
9405 };
9406 let mut moe_out = e.uninit(t * n_embd)?;
9407 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9408 return Ok(moe_out);
9409 }
9410
9411 let g_len = m.gate_exps.expert_stride;
9412 let u_len = m.up_exps.expert_stride;
9413 let d_len = m.down_exps.expert_stride;
9414 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9418 let (mut sg, mut su, mut sd) = if dev.is_some() {
9419 (None, None, None)
9420 } else {
9421 (
9422 Some(e.alloc_u8_uninit(g_len)?),
9423 Some(e.alloc_u8_uninit(u_len)?),
9424 Some(e.alloc_u8_uninit(d_len)?),
9425 )
9426 };
9427 let mut moe_out = e.zeros(t * n_embd)?;
9428 for tok in 0..t {
9429 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9430 let w = &w_all[tok * n_used..(tok + 1) * n_used];
9431 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9432 for (j, &ex) in sel.iter().enumerate() {
9433 let ex = ex as usize;
9434 let gate = match dev {
9435 Some(d) => e.qmatvec_view(
9436 &d.gate,
9437 ex * g_len..(ex + 1) * g_len,
9438 &zt,
9439 1,
9440 m.gate_exps.in_f,
9441 m.gate_exps.out_f,
9442 m.gate_exps.qtype,
9443 m.gate_exps.row_bytes,
9444 )?,
9445 None => {
9446 let sg = sg.as_mut().unwrap();
9447 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9448 e.qmatvec_view(
9449 sg,
9450 0..g_len,
9451 &zt,
9452 1,
9453 m.gate_exps.in_f,
9454 m.gate_exps.out_f,
9455 m.gate_exps.qtype,
9456 m.gate_exps.row_bytes,
9457 )?
9458 }
9459 };
9460 let up = match dev {
9461 Some(d) => e.qmatvec_view(
9462 &d.up,
9463 ex * u_len..(ex + 1) * u_len,
9464 &zt,
9465 1,
9466 m.up_exps.in_f,
9467 m.up_exps.out_f,
9468 m.up_exps.qtype,
9469 m.up_exps.row_bytes,
9470 )?,
9471 None => {
9472 let su = su.as_mut().unwrap();
9473 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9474 e.qmatvec_view(
9475 su,
9476 0..u_len,
9477 &zt,
9478 1,
9479 m.up_exps.in_f,
9480 m.up_exps.out_f,
9481 m.up_exps.qtype,
9482 m.up_exps.row_bytes,
9483 )?
9484 }
9485 };
9486 let mut act = e.uninit(n_ff_exp)?;
9487 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9488 let actv = act.slice(0..n_ff_exp);
9489 let y = match dev {
9490 Some(d) => e.qmatvec_view(
9491 &d.down,
9492 ex * d_len..(ex + 1) * d_len,
9493 &actv,
9494 1,
9495 m.down_exps.in_f,
9496 m.down_exps.out_f,
9497 m.down_exps.qtype,
9498 m.down_exps.row_bytes,
9499 )?,
9500 None => {
9501 let sd = sd.as_mut().unwrap();
9502 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9503 e.qmatvec_view(
9504 sd,
9505 0..d_len,
9506 &actv,
9507 1,
9508 m.down_exps.in_f,
9509 m.down_exps.out_f,
9510 m.down_exps.qtype,
9511 m.down_exps.row_bytes,
9512 )?
9513 }
9514 };
9515 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9516 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9517 }
9518 }
9519 Ok(moe_out)
9520 }
9521
9522 fn gemma4_layer(
9524 &self,
9525 e: &Engine,
9526 il: usize,
9527 layer: &crate::hybrid::HybridLayer,
9528 x: &CudaSlice<f32>,
9529 pos_d: &CudaSlice<i32>,
9530 t: usize,
9531 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9532 let n_embd = self.cfg.n_embd as usize;
9533 let eps = self.cfg.rms_eps;
9534
9535 let mut h = e.zeros(t * n_embd)?;
9536 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9537 let Mixer::Full(fa) = &layer.mixer else {
9538 panic!("gemma4 layer {il} not full-attn")
9539 };
9540 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9541 let mut cur = e.zeros(t * n_embd)?;
9543 e.rms_norm(
9544 &o,
9545 layer.post_attn_norm.float_data(),
9546 &mut cur,
9547 n_embd,
9548 t,
9549 eps,
9550 )?;
9551 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9552 }
9553
9554 fn gemma4_layer_tail_add(
9558 &self,
9559 e: &Engine,
9560 layer: &crate::hybrid::HybridLayer,
9561 cur: &CudaSlice<f32>,
9562 x: &CudaSlice<f32>,
9563 t: usize,
9564 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9565 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9566 }
9567
9568 fn gemma4_layer_tail_add_n(
9571 &self,
9572 e: &Engine,
9573 layer: &crate::hybrid::HybridLayer,
9574 cur: &CudaSlice<f32>,
9575 x: &CudaSlice<f32>,
9576 t: usize,
9577 next_norm: Option<&CudaSlice<f32>>,
9578 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9579 let n_embd = self.cfg.n_embd as usize;
9580 let bits = layer.gemma4.as_ref().unwrap();
9581 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9582 let mut xn = e.uninit(t * n_embd)?;
9583 match next_norm {
9584 Some(w) => {
9585 let mut hn = e.uninit(t * n_embd)?;
9586 e.add_scale_rms_norm(
9587 &sn,
9588 &attn_out,
9589 bits.layer_scale,
9590 w,
9591 &mut xn,
9592 &mut hn,
9593 n_embd,
9594 t,
9595 self.cfg.rms_eps,
9596 )?;
9597 Ok((xn, Some(hn)))
9598 }
9599 None => {
9600 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9601 Ok((xn, None))
9602 }
9603 }
9604 }
9605
9606 fn gemma4_layer_tail_core(
9609 &self,
9610 e: &Engine,
9611 layer: &crate::hybrid::HybridLayer,
9612 cur: &CudaSlice<f32>,
9613 x: &CudaSlice<f32>,
9614 t: usize,
9615 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9616 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9617 }
9618
9619 fn gemma4_layer_tail_core_pn(
9626 &self,
9627 e: &Engine,
9628 layer: &crate::hybrid::HybridLayer,
9629 cur: &CudaSlice<f32>,
9630 x: &CudaSlice<f32>,
9631 t: usize,
9632 pre_norm: Option<&CudaSlice<f32>>,
9633 defer_post_norm: bool,
9634 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9635 let n_embd = self.cfg.n_embd as usize;
9636 let eps = self.cfg.rms_eps;
9637 let bits = layer.gemma4.as_ref().unwrap();
9638
9639 let Some(mbits) = bits.moe_bits.as_ref() else {
9642 let crate::hybrid::Ffn::Dense {
9643 ffn_gate,
9644 ffn_up,
9645 ffn_down,
9646 } = &layer.ffn
9647 else {
9648 panic!("gemma4 dense layer without Dense ffn")
9649 };
9650 let mut attn_out = e.uninit(t * n_embd)?;
9651 let mut zsh = e.uninit(t * n_embd)?;
9652 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9655 match pre_norm {
9656 Some(wa) if t == 1 => {
9657 zpair = Some(e.rms_pre_add_rms_norm_q8z(
9658 cur,
9659 wa,
9660 x,
9661 bits.ffn_norm.float_data(),
9662 &mut attn_out,
9663 &mut zsh,
9664 n_embd,
9665 t,
9666 eps,
9667 )?);
9668 }
9669 Some(wa) => e.rms_pre_add_rms_norm(
9670 cur,
9671 wa,
9672 x,
9673 bits.ffn_norm.float_data(),
9674 &mut attn_out,
9675 &mut zsh,
9676 n_embd,
9677 t,
9678 eps,
9679 )?,
9680 None => e.add_rms_norm(
9681 cur,
9682 x,
9683 bits.ffn_norm.float_data(),
9684 &mut attn_out,
9685 &mut zsh,
9686 n_embd,
9687 t,
9688 eps,
9689 )?,
9690 }
9691 let n_ff = ffn_gate.out_features();
9692 let (gate, up) = if t == 1 {
9698 let (zq, zd) = match zpair {
9699 Some(p) => p,
9700 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9701 };
9702 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9703 Some(p) => p,
9704 None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
9706 Some(p) => p,
9707 None => (
9708 e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9709 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9710 ),
9711 },
9712 }
9713 } else {
9714 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9719 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9720 let fused = if f2b {
9721 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9722 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9723 } else {
9724 None
9725 };
9726 match fused {
9727 Some(p) => p,
9728 None => {
9729 e.mmq_act_begin();
9731 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9732 }
9733 }
9734 };
9735 let mut act = e.uninit(t * n_ff)?;
9736 let f0 = if e.uses_q8_1_fast(ffn_down) {
9739 let upv = e.view(&up, t * n_ff);
9740 let up_all = upv.slice(0..t * n_ff);
9741 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9742 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9743 } else {
9744 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9745 e.matmul(ffn_down, &act, t)?
9746 };
9747 if defer_post_norm {
9748 return Ok((f0, attn_out));
9749 }
9750 let mut sn = e.uninit(t * n_embd)?;
9751 e.rms_norm(
9752 &f0,
9753 bits.post_ffw_norm.float_data(),
9754 &mut sn,
9755 n_embd,
9756 t,
9757 eps,
9758 )?;
9759 return Ok((sn, attn_out));
9760 };
9761
9762 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9763 let mut attn_out = e.uninit(t * n_embd)?;
9768 let mut router_in = e.uninit(t * n_embd)?;
9769 let fast_moe = match &layer.ffn {
9770 crate::hybrid::Ffn::Moe(m) => {
9771 m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9772 && expert_dp4a_supported(m.gate_exps.qtype)
9773 && expert_dp4a_supported(m.up_exps.qtype)
9774 && expert_dp4a_supported(m.down_exps.qtype)
9775 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9776 }
9777 _ => false,
9778 };
9779 let q8z = t < PRIME_MIN_T && fast_moe;
9780 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9781 let (z0, m2) = e.add_rms_norm3_q8z(
9782 cur,
9783 x,
9784 bits.ffn_norm.float_data(),
9785 &mbits.router_scale_pre,
9786 mbits.pre_ffw_norm_2.float_data(),
9787 &mut attn_out,
9788 &mut router_in,
9789 n_embd,
9790 t,
9791 eps,
9792 )?;
9793 (None, Some(z0), Some(m2))
9794 } else {
9795 let mut zsh = e.uninit(t * n_embd)?;
9796 let mut moe_in = e.uninit(t * n_embd)?;
9797 e.add_rms_norm3(
9798 cur,
9799 x,
9800 bits.ffn_norm.float_data(),
9801 &mbits.router_scale_pre,
9802 mbits.pre_ffw_norm_2.float_data(),
9803 &mut attn_out,
9804 &mut zsh,
9805 &mut router_in,
9806 &mut moe_in,
9807 n_embd,
9808 t,
9809 eps,
9810 )?;
9811 (Some((zsh, moe_in)), None, None)
9812 };
9813 let attn_out2 = attn_out;
9814 #[allow(unused_variables)]
9815 let attn_out = &attn_out2;
9816 let n_ff = mbits.shared_gate.out_features();
9817 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9818 if t == 1 {
9819 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9820 Some(p) => p,
9821 None => match e.matmul_nvfp4_fused2(
9822 &mbits.shared_gate,
9823 &mbits.shared_up,
9824 zq,
9825 zd,
9826 1,
9827 )? {
9828 Some(p) => p,
9829 None => {
9830 let h0 = e.zeros(0)?;
9831 (
9832 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
9833 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
9834 )
9835 }
9836 },
9837 }
9838 } else {
9839 let h0 = e.zeros(0)?;
9841 (
9842 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
9843 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
9844 )
9845 }
9846 } else {
9847 let (zsh, _) = zsh_f32.as_ref().unwrap();
9848 (
9849 e.matmul(&mbits.shared_gate, zsh, t)?,
9850 e.matmul(&mbits.shared_up, zsh, t)?,
9851 )
9852 };
9853 let mut act = e.uninit(t * n_ff)?;
9854 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9855 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
9856 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
9857 panic!("gemma4 layer not MoE")
9858 };
9859 let moe0 = match (&moe_q8, &zsh_f32) {
9860 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
9861 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
9862 _ => unreachable!(),
9863 };
9864 let mut mlp = e.uninit(t * n_embd)?;
9866 let mut moe = e.uninit(t * n_embd)?;
9867 e.rms_norm2x(
9868 &mlp0,
9869 &moe0,
9870 mbits.post_ffw_norm_1.float_data(),
9871 mbits.post_ffw_norm_2.float_data(),
9872 &mut mlp,
9873 &mut moe,
9874 n_embd,
9875 t,
9876 eps,
9877 )?;
9878
9879 let mut sum = e.uninit(t * n_embd)?;
9882 let mut sn = e.uninit(t * n_embd)?;
9883 e.add_rms_norm(
9884 &mlp,
9885 &moe,
9886 bits.post_ffw_norm.float_data(),
9887 &mut sum,
9888 &mut sn,
9889 n_embd,
9890 t,
9891 eps,
9892 )?;
9893 Ok((sn, attn_out2))
9894 }
9895
9896 pub(crate) fn gemma4_layer_tail_add_nq_pn(
9906 &self,
9907 e: &Engine,
9908 layer: &crate::hybrid::HybridLayer,
9909 o: &CudaSlice<f32>,
9910 x: &CudaSlice<f32>,
9911 t: usize,
9912 next_norm: Option<&CudaSlice<f32>>,
9913 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9914 {
9915 let n_embd = self.cfg.n_embd as usize;
9916 let eps = self.cfg.rms_eps;
9917 let bits = layer.gemma4.as_ref().unwrap();
9918 if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
9919 let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
9920 e,
9921 layer,
9922 o,
9923 x,
9924 t,
9925 Some(layer.post_attn_norm.float_data()),
9926 true,
9927 )?;
9928 let mut xn = e.uninit(t * n_embd)?;
9929 return match next_norm {
9930 Some(w) => {
9931 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
9932 &f0,
9933 bits.post_ffw_norm.float_data(),
9934 &attn_out,
9935 bits.layer_scale,
9936 w,
9937 &mut xn,
9938 n_embd,
9939 t,
9940 eps,
9941 )?;
9942 Ok((xn, Some(pair)))
9943 }
9944 None => {
9945 let mut sn = e.uninit(t * n_embd)?;
9946 e.rms_norm(
9947 &f0,
9948 bits.post_ffw_norm.float_data(),
9949 &mut sn,
9950 n_embd,
9951 t,
9952 eps,
9953 )?;
9954 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9955 Ok((xn, None))
9956 }
9957 };
9958 }
9959 let mut cur = e.uninit(t * n_embd)?;
9960 e.rms_norm(
9961 o,
9962 layer.post_attn_norm.float_data(),
9963 &mut cur,
9964 n_embd,
9965 t,
9966 eps,
9967 )?;
9968 self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
9969 }
9970
9971 pub(crate) fn gemma4_layer_tail_add_nq(
9972 &self,
9973 e: &Engine,
9974 layer: &crate::hybrid::HybridLayer,
9975 cur: &CudaSlice<f32>,
9976 x: &CudaSlice<f32>,
9977 t: usize,
9978 next_norm: Option<&CudaSlice<f32>>,
9979 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9980 {
9981 let n_embd = self.cfg.n_embd as usize;
9982 let bits = layer.gemma4.as_ref().unwrap();
9983 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9984 let mut xn = e.uninit(t * n_embd)?;
9985 match next_norm {
9986 Some(w) => {
9987 let pair = e.add_scale_rms_norm_q8_1(
9988 &sn,
9989 &attn_out,
9990 bits.layer_scale,
9991 w,
9992 &mut xn,
9993 n_embd,
9994 t,
9995 self.cfg.rms_eps,
9996 )?;
9997 Ok((xn, Some(pair)))
9998 }
9999 None => {
10000 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10001 Ok((xn, None))
10002 }
10003 }
10004 }
10005
10006 fn gemma4_forward(
10009 &self,
10010 e: &Engine,
10011 tokens: &[u32],
10012 last_only: bool,
10013 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
10014 if self.is_gemma4_e4b() {
10017 return self.gemma4_e4b_forward(e, tokens, last_only);
10018 }
10019 let n_embd = self.cfg.n_embd as usize;
10020 let t = tokens.len();
10021 let pos: Vec<i32> = (0..t as i32).collect();
10022 let pos_d = e.htod_i32(&pos)?;
10023
10024 let mut x = self.embed(e, tokens)?;
10025 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10026 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
10029 let stat =
10030 |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
10031 let h = e.dtoh(x)?;
10032 let bad = h.iter().filter(|v| !v.is_finite()).count();
10033 let mx = h
10034 .iter()
10035 .filter(|v| v.is_finite())
10036 .fold(0.0f32, |m, v| m.max(v.abs()));
10037 eprintln!(
10038 "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
10039 &h[..3]
10040 );
10041 Ok(())
10042 };
10043 if probe {
10044 stat(e, &x, "embed")?;
10045 }
10046 for (il, layer) in self.layers.iter().enumerate() {
10047 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
10048 if probe {
10049 stat(e, &x, &format!("L{il}"))?;
10050 }
10051 }
10052 let mut hn = e.zeros(t * n_embd)?;
10053 e.rms_norm(
10054 &x,
10055 self.output_norm.float_data(),
10056 &mut hn,
10057 n_embd,
10058 t,
10059 self.cfg.rms_eps,
10060 )?;
10061 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10062 let n_vocab = self.output.out_features();
10063 let logits = if last_only {
10064 let hv = e.view(&hn, t * n_embd);
10065 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
10066 let mut hlast = e.zeros(n_embd)?;
10067 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
10068 let mut ld = e.matmul(&self.output, &hlast, 1)?;
10069 e.softcap(&mut ld, cap, n_vocab)?;
10070 self.gemma4_suppress(e, &mut ld, 1)?;
10071 e.dtoh(&ld)?
10072 } else {
10073 let mut ld = e.matmul(&self.output, &hn, t)?;
10074 e.softcap(&mut ld, cap, t * n_vocab)?;
10075 self.gemma4_suppress(e, &mut ld, t)?;
10076 e.dtoh(&ld)?
10077 };
10078 Ok(logits)
10079 }
10080
10081 pub(crate) fn gemma4_prime(
10086 &self,
10087 e: &Engine,
10088 tokens: &[u32],
10089 cache: &mut Cache,
10090 overlay: Option<&crate::vision::EmbedOverlay>,
10091 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10092 if cache.pos != 0 {
10097 return Err(
10098 "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
10099 — prime the full prompt in one call or decode tokenwise"
10100 .into(),
10101 );
10102 }
10103 let n_embd = self.cfg.n_embd as usize;
10104 let eps = self.cfg.rms_eps;
10105 let t = tokens.len();
10106 let pos: Vec<i32> = (0..t as i32).collect();
10107 let pos_d = e.htod_i32(&pos)?;
10108 let mut x = self.embed(e, tokens)?;
10109 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10110 let island: Option<CudaSlice<i32>> = match overlay {
10117 Some(ov) => {
10118 let mut span_id = vec![-1i32; t];
10119 for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
10120 if pos + n_rows > t {
10121 return Err(format!(
10122 "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
10123 pos + n_rows
10124 )
10125 .into());
10126 }
10127 let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
10128 e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
10129 for s in span_id.iter_mut().skip(pos).take(n_rows) {
10130 *s = i as i32;
10131 }
10132 }
10133 if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
10137 None
10138 } else {
10139 Some(e.htod_i32(&span_id)?)
10140 }
10141 }
10142 None => None,
10143 };
10144 for (il, layer) in self.layers.iter().enumerate() {
10145 let mut h = e.zeros(t * n_embd)?;
10146 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
10147 let Mixer::Full(fa) = &layer.mixer else {
10148 panic!("gemma4 layer not full-attn")
10149 };
10150 let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
10151 if trace {
10152 let v = e.dtoh(&h)?;
10153 let nan = v.iter().filter(|x| x.is_nan()).count();
10154 eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
10155 }
10156 let o =
10157 self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
10158 if trace {
10159 let v = e.dtoh(&o)?;
10160 let nan = v.iter().filter(|x| x.is_nan()).count();
10161 eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
10162 }
10163 let mut cur = e.zeros(t * n_embd)?;
10164 e.rms_norm(
10165 &o,
10166 layer.post_attn_norm.float_data(),
10167 &mut cur,
10168 n_embd,
10169 t,
10170 eps,
10171 )?;
10172 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
10173 self.dflash_tap(e, cache, il, &x, t)?;
10174 if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
10176 let h = e.dtoh(&x)?;
10177 let nan = h.iter().filter(|v| v.is_nan()).count();
10178 let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
10179 eprintln!(
10180 "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
10181 h.len()
10182 );
10183 if nan > 0 {
10184 return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
10185 }
10186 }
10187 }
10188 cache.pos += t;
10189 let hiddens = e.clone_dtod(&x)?;
10190 let xv = e.view(&x, t * n_embd);
10191 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
10192 let mut h_seed = e.zeros(n_embd)?;
10193 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
10194 let mut hn = e.uninit(n_embd)?;
10195 e.rms_norm(
10196 &h_seed,
10197 self.output_norm.float_data(),
10198 &mut hn,
10199 n_embd,
10200 1,
10201 eps,
10202 )?;
10203 let mut ld = e.matmul(&self.output, &hn, 1)?;
10204 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10205 e.softcap(&mut ld, cap, self.output.out_features())?;
10206 self.gemma4_suppress(e, &mut ld, 1)?;
10207 let logits = e.dtoh(&ld)?;
10208 Ok((logits, h_seed, hiddens))
10209 }
10210
10211 fn gemma4_decode_attn(
10216 &self,
10217 e: &Engine,
10218 fa: &crate::hybrid::FullAttnLayer,
10219 il: usize,
10220 hq: &CudaSlice<i8>,
10221 hdq: &CudaSlice<f32>,
10222 pos_d: &CudaSlice<i32>,
10223 cache: &mut Cache,
10224 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10225 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10226 let eps = self.cfg.rms_eps;
10227 let aux = self.gemma4_aux.as_ref().unwrap();
10228 let ones = aux.ones(e);
10229 #[cfg(debug_assertions)]
10230 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
10231 let (hq, hdq) = (hq, hdq);
10232 let h0 = e.zeros(0)?;
10233 let h = &h0;
10234 let (q0, k0, v0) = if swa {
10235 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
10236 Some(t3) => t3,
10237 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10240 Some((q0, k0)) => {
10241 let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?;
10242 (q0, k0, v0)
10243 }
10244 None => (
10245 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10246 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10247 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10248 ),
10249 },
10250 }
10251 } else {
10252 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10253 Some(p) => p,
10254 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10255 Some(p) => p,
10256 None => (
10257 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10258 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10259 ),
10260 },
10261 };
10262 let v0 = e.clone_dtod(&k0)?;
10263 (q0, k0, v0)
10264 };
10265 let mut q = e.uninit(nh * hd)?;
10266 let mut k = e.uninit(nkv * hd)?;
10267 let mut v = e.uninit(nkv * hd)?;
10268 let ff = if swa {
10271 None
10272 } else {
10273 Some(
10274 aux.rope_freqs(e)
10275 .expect("gemma4 global rope needs rope_freqs.weight"),
10276 )
10277 };
10278 #[cfg(debug_assertions)]
10279 if let Some(ff) = ff {
10280 crate::debug_assert_tensor_stream_device(
10281 ff,
10282 &e.stream(),
10283 "gemma4_decode_attn.rope_freqs",
10284 );
10285 }
10286 let kvl = cache.kv[il].as_mut().unwrap();
10287 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10288 if crate::Engine::qkv_append_on() {
10289 e.rms_norm_qkv_rope_append(
10293 &q0,
10294 &k0,
10295 &v0,
10296 fa.q_norm.float_data(),
10297 fa.k_norm.float_data(),
10298 ones,
10299 &mut q,
10300 &mut k,
10301 &mut v,
10302 hd,
10303 nh,
10304 nkv,
10305 pos_d,
10306 nh,
10307 nkv,
10308 base,
10309 1.0,
10310 ff,
10311 eps,
10312 &mut kvl.k,
10313 &mut kvl.v,
10314 kvl.len,
10315 kvl.k_tok_bytes,
10316 kvl.v_tok_bytes,
10317 kv_fp8,
10318 )?;
10319 } else {
10320 e.rms_norm_qkv_rope(
10321 &q0,
10322 &k0,
10323 &v0,
10324 fa.q_norm.float_data(),
10325 fa.k_norm.float_data(),
10326 ones,
10327 &mut q,
10328 &mut k,
10329 &mut v,
10330 hd,
10331 nh,
10332 nkv,
10333 pos_d,
10334 nh,
10335 nkv,
10336 base,
10337 1.0,
10338 ff,
10339 eps,
10340 )?;
10341 e.append_kv_quantized(
10342 &k,
10343 &v,
10344 &mut kvl.k,
10345 &mut kvl.v,
10346 kvl.len,
10347 kvl.kv_dim_k,
10348 kvl.kv_dim_v,
10349 kvl.k_tok_bytes,
10350 kvl.v_tok_bytes,
10351 kv_fp8,
10352 )?;
10353 }
10354 kvl.len += 1;
10355 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10359 let mut attn = e.uninit(nh * hd)?;
10360 if !swa
10362 && hd == 512
10363 && kvl.len >= crate::fa512_min_tkv()
10364 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10365 {
10366 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10367 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10368 let base = kvl.len as i32;
10370 e.i32_set_k(&mut kvl.len_d, base)?;
10371 e.fa_decode_rows(
10372 &q,
10373 &kp,
10374 &vp,
10375 &mut attn,
10376 hd,
10377 nh,
10378 nkv,
10379 kvl.len - 1,
10380 1,
10381 scale,
10382 kvl.k_tok_bytes,
10383 kvl.v_tok_bytes,
10384 Some((&kvl.len_d, -1)),
10385 false,
10386 false,
10387 None,
10388 )?;
10389 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10390 }
10391 if swa
10393 && kvl.len > win
10394 && hd == 256
10395 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10396 {
10397 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10398 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10399 let base = kvl.len as i32;
10400 e.i32_set_k(&mut kvl.len_d, base)?;
10401 e.fa_decode_rows_w(
10402 &q,
10403 &kp,
10404 &vp,
10405 &mut attn,
10406 hd,
10407 nh,
10408 nkv,
10409 &kvl.len_d,
10410 -1,
10411 1,
10412 scale,
10413 win,
10414 kvl.k_tok_bytes,
10415 kvl.v_tok_bytes,
10416 None,
10417 )?;
10418 return Ok(e.matmul(&fa.wo, &attn, 1)?);
10419 }
10420 let (off_tok, t_kv) = if swa && kvl.len > win {
10421 (kvl.len - win, win)
10422 } else {
10423 (0, kvl.len)
10424 };
10425 let k_view = e.view_u8_range(
10426 &kvl.k,
10427 off_tok * kvl.k_tok_bytes,
10428 (off_tok + t_kv) * kvl.k_tok_bytes,
10429 );
10430 let v_view = e.view_u8_range(
10431 &kvl.v,
10432 off_tok * kvl.v_tok_bytes,
10433 (off_tok + t_kv) * kvl.v_tok_bytes,
10434 );
10435 e.fa_decode_kvmod(
10436 &q,
10437 &k_view,
10438 &v_view,
10439 &mut attn,
10440 hd,
10441 nh,
10442 nkv,
10443 t_kv,
10444 scale,
10445 kvl.k_tok_bytes,
10446 kvl.v_tok_bytes,
10447 swa && crate::Engine::wkv_on(),
10448 )?;
10449 Ok(e.matmul(&fa.wo, &attn, 1)?)
10450 }
10451
10452 #[allow(clippy::too_many_arguments)]
10459 pub fn gemma4_decode_step_dc(
10460 &self,
10461 e: &Engine,
10462 token_d: &CudaSlice<u32>,
10463 pos_d: &mut CudaSlice<i32>,
10464 embd_gpu: &CudaSlice<u8>,
10465 embd_qt: i32,
10466 embd_rb: usize,
10467 cache: &mut Cache,
10468 n_vocab: usize,
10469 cap_bucket_max: Option<(usize, usize)>,
10470 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10471 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10472 self.gemma4_decode_step_dc_into(
10473 e,
10474 token_d,
10475 pos_d,
10476 embd_gpu,
10477 embd_qt,
10478 embd_rb,
10479 cache,
10480 n_vocab,
10481 cap_bucket_max,
10482 &mut tok_out,
10483 )?;
10484 Ok(tok_out)
10485 }
10486
10487 #[allow(clippy::too_many_arguments)]
10490 pub fn gemma4_decode_step_dc_into(
10491 &self,
10492 e: &Engine,
10493 token_d: &CudaSlice<u32>,
10494 pos_d: &mut CudaSlice<i32>,
10495 embd_gpu: &CudaSlice<u8>,
10496 embd_qt: i32,
10497 embd_rb: usize,
10498 cache: &mut Cache,
10499 n_vocab: usize,
10500 cap_bucket_max: Option<(usize, usize)>,
10501 tok_out: &mut CudaSlice<u32>,
10502 ) -> Result<(), Box<dyn std::error::Error>> {
10503 let n_embd = self.cfg.n_embd as usize;
10504 let eps = self.cfg.rms_eps;
10505 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10506 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10507 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10508 let n_layers = self.layers.len();
10509 for (il, layer) in self.layers.iter().enumerate() {
10510 let (hq, hdq) = match h_carry.take() {
10511 Some(p) => p,
10512 None => {
10513 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10514 }
10515 };
10516 let Mixer::Full(fa) = &layer.mixer else {
10517 panic!("gemma4 layer {il} not full-attn")
10518 };
10519 let o =
10520 self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10521 let next_norm = if il + 1 < n_layers {
10522 Some(self.layers[il + 1].attn_norm.float_data())
10523 } else {
10524 None
10525 };
10526 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
10527 x = xn;
10528 h_carry = hn;
10529 }
10530 let mut hn = e.uninit(n_embd)?;
10531 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10532 let mut logits = e.matmul(&self.output, &hn, 1)?;
10533 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10535 e.inc_seqlen(pos_d)?;
10536 if cap_bucket_max.is_none() {
10537 cache.pos += 1;
10538 }
10539 Ok(())
10540 }
10541
10542 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10549 let n_embd = self.cfg.n_embd as usize;
10550 let n_vocab = self.output.out_features();
10551 let n_layers = self.layers.len();
10552 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10553 for il in 0..n_layers {
10554 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10555 qmax = qmax.max(nh * hd);
10556 kvmax = kvmax.max(nkv * hd);
10557 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10558 ffmax = ffmax.max(ffn_gate.out_features());
10559 }
10560 }
10561 Ok(G4DcSlots {
10562 x: e.uninit(n_embd)?,
10563 xn: e.uninit(n_embd)?,
10564 cur: e.uninit(n_embd)?,
10565 hq: e.alloc_i8_uninit(n_embd)?,
10566 hd_: e.uninit(n_embd / 32)?,
10567 q0: e.uninit(qmax)?,
10568 k0: e.uninit(kvmax)?,
10569 v0: e.uninit(kvmax)?,
10570 q: e.uninit(qmax)?,
10571 k: e.uninit(kvmax)?,
10572 v: e.uninit(kvmax)?,
10573 attn: e.uninit(qmax)?,
10574 o: e.uninit(n_embd)?,
10575 attn_out: e.uninit(n_embd)?,
10576 zsh: e.uninit(n_embd)?,
10577 zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10580 zd: e.uninit(n_embd.max(qmax) / 32)?,
10581 gate: e.uninit(ffmax)?,
10582 up: e.uninit(ffmax)?,
10583 act: e.uninit(ffmax)?,
10584 actq: e.alloc_i8_uninit(ffmax)?,
10585 actd: e.uninit(ffmax / 32)?,
10586 f0: e.uninit(n_embd)?,
10587 sn: e.uninit(n_embd)?,
10588 hn: e.uninit(n_embd)?,
10589 logits: e.uninit(n_vocab)?,
10590 })
10591 }
10592
10593 fn g4_matvec_m1_into(
10596 &self,
10597 e: &Engine,
10598 w: &crate::model::GpuTensor,
10599 aq: &CudaSlice<i8>,
10600 ad: &CudaSlice<f32>,
10601 y: &mut CudaSlice<f32>,
10602 ) -> Result<(), Box<dyn std::error::Error>> {
10603 use crate::model::GpuTensor;
10604 let (bytes, qtype, row_bytes, scale, rp) = match w {
10605 GpuTensor::Quant {
10606 bytes,
10607 qtype,
10608 row_bytes,
10609 scale,
10610 rp,
10611 ..
10612 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10613 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10614 };
10615 let (mbytes, mrp) = match w {
10616 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10617 _ => (bytes, rp),
10618 };
10619 e.qmatvec_mmvq_into(
10620 mbytes,
10621 aq,
10622 ad,
10623 1,
10624 w.in_features(),
10625 w.out_features(),
10626 qtype,
10627 row_bytes,
10628 scale,
10629 mrp,
10630 y,
10631 )
10632 }
10633
10634 #[allow(clippy::too_many_arguments)]
10638 pub fn gemma4_decode_step_dc_slotted(
10639 &self,
10640 e: &Engine,
10641 token_d: &CudaSlice<u32>,
10642 pos_d: &mut CudaSlice<i32>,
10643 embd_gpu: &CudaSlice<u8>,
10644 embd_qt: i32,
10645 embd_rb: usize,
10646 cache: &mut Cache,
10647 n_vocab: usize,
10648 cap_bucket_max: Option<(usize, usize)>,
10649 sl: &mut G4DcSlots,
10650 tok_out: &mut CudaSlice<u32>,
10651 ring: Option<(&mut CudaSlice<u32>, usize)>,
10652 ) -> Result<(), Box<dyn std::error::Error>> {
10653 let n_embd = self.cfg.n_embd as usize;
10654 let eps = self.cfg.rms_eps;
10655 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10656 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10657 let n_layers = self.layers.len();
10658 let mut has_carry = false;
10659 for il in 0..n_layers {
10660 if !has_carry {
10661 e.rms_norm_q8_1_into(
10662 &sl.x,
10663 self.layers[il].attn_norm.float_data(),
10664 n_embd,
10665 1,
10666 eps,
10667 &mut sl.hq,
10668 &mut sl.hd_,
10669 )?;
10670 }
10671 has_carry = true;
10672 let layer = &self.layers[il];
10673 let Mixer::Full(fa) = &layer.mixer else {
10674 panic!("gemma4 layer {il} not full-attn")
10675 };
10676 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10677 if !Engine::g4_pnfold_on() {
10680 e.rms_norm(
10681 &sl.o,
10682 layer.post_attn_norm.float_data(),
10683 &mut sl.cur,
10684 n_embd,
10685 1,
10686 eps,
10687 )?;
10688 }
10689 let next_norm = if il + 1 < n_layers {
10690 Some(self.layers[il + 1].attn_norm.float_data())
10691 } else {
10692 None
10693 };
10694 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10695 std::mem::swap(&mut sl.x, &mut sl.xn);
10696 }
10697 e.rms_norm(
10698 &sl.x,
10699 self.output_norm.float_data(),
10700 &mut sl.hn,
10701 n_embd,
10702 1,
10703 eps,
10704 )?;
10705 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10706 {
10708 let (zq, zd) = (&sl.zq, &sl.zd);
10709 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10710 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10711 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10712 }
10713 self.gemma4_suppress(e, &mut sl.logits, 1)?;
10714 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10715 if let Some((ring, base)) = ring {
10716 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10720 }
10721 e.inc_seqlen(pos_d)?;
10722 if cap_bucket_max.is_none() {
10723 cache.pos += 1;
10724 }
10725 Ok(())
10726 }
10727
10728 #[allow(clippy::too_many_arguments)]
10730 fn gemma4_decode_attn_dc_slotted(
10731 &self,
10732 e: &Engine,
10733 fa: &crate::hybrid::FullAttnLayer,
10734 il: usize,
10735 pos_d: &CudaSlice<i32>,
10736 cache: &mut Cache,
10737 cap_bucket_max: Option<(usize, usize)>,
10738 sl: &mut G4DcSlots,
10739 ) -> Result<(), Box<dyn std::error::Error>> {
10740 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10741 let eps = self.cfg.rms_eps;
10742 let aux = self.gemma4_aux.as_ref().unwrap();
10743 let ones = aux.ones(e);
10744 #[cfg(debug_assertions)]
10745 crate::debug_assert_tensor_stream_device(
10746 ones,
10747 &e.stream(),
10748 "gemma4_decode_attn_dc_slotted.ones",
10749 );
10750 {
10751 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10752 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10753 if swa {
10754 if !e.matmul_q4_fused3_into(
10755 &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10756 )? {
10757 if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10761 {
10762 self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
10763 } else {
10764 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10765 }
10766 }
10767 } else {
10768 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10769 && !e
10770 .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10771 {
10772 return Err("slotted step: fused2 unavailable".into());
10773 }
10774 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10775 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10776 }
10777 }
10778 let ff = if swa {
10781 None
10782 } else {
10783 Some(
10784 aux.rope_freqs(e)
10785 .expect("gemma4 global rope needs rope_freqs.weight"),
10786 )
10787 };
10788 #[cfg(debug_assertions)]
10789 if let Some(ff) = ff {
10790 crate::debug_assert_tensor_stream_device(
10791 ff,
10792 &e.stream(),
10793 "gemma4_decode_attn_dc_slotted.rope_freqs",
10794 );
10795 }
10796 let kvl = cache.kv[il].as_mut().unwrap();
10797 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10798 if crate::Engine::qkv_append_on() {
10799 e.rms_norm_qkv_rope_append_dc(
10801 &sl.q0,
10802 &sl.k0,
10803 &sl.v0,
10804 fa.q_norm.float_data(),
10805 fa.k_norm.float_data(),
10806 ones,
10807 &mut sl.q,
10808 &mut sl.k,
10809 &mut sl.v,
10810 hd,
10811 nh,
10812 nkv,
10813 pos_d,
10814 nh,
10815 nkv,
10816 base,
10817 1.0,
10818 ff,
10819 eps,
10820 &mut kvl.k,
10821 &mut kvl.v,
10822 &kvl.len_d,
10823 kvl.k_tok_bytes,
10824 kvl.v_tok_bytes,
10825 kv_fp8,
10826 )?;
10827 } else {
10828 e.rms_norm_qkv_rope(
10829 &sl.q0,
10830 &sl.k0,
10831 &sl.v0,
10832 fa.q_norm.float_data(),
10833 fa.k_norm.float_data(),
10834 ones,
10835 &mut sl.q,
10836 &mut sl.k,
10837 &mut sl.v,
10838 hd,
10839 nh,
10840 nkv,
10841 pos_d,
10842 nh,
10843 nkv,
10844 base,
10845 1.0,
10846 ff,
10847 eps,
10848 )?;
10849 e.append_kv_quantized_dc(
10850 &sl.k,
10851 &sl.v,
10852 &mut kvl.k,
10853 &mut kvl.v,
10854 &kvl.len_d,
10855 kvl.kv_dim_k,
10856 kvl.kv_dim_v,
10857 kvl.k_tok_bytes,
10858 kvl.v_tok_bytes,
10859 kv_fp8,
10860 )?;
10861 }
10862 e.inc_seqlen(&mut kvl.len_d)?;
10863 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
10864 let k_view = e.view_u8(&kvl.k, kvl.k.len());
10865 let v_view = e.view_u8(&kvl.v, kvl.v.len());
10866 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10867 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10868 let mut fa_q8 = false;
10872 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10873 e.fa_decode_rows(
10874 &sl.q,
10875 &k_view,
10876 &v_view,
10877 &mut sl.attn,
10878 hd,
10879 nh,
10880 nkv,
10881 b_glob - 1,
10882 1,
10883 scale,
10884 kvl.k_tok_bytes,
10885 kvl.v_tok_bytes,
10886 Some((&kvl.len_d, -1)),
10887 false,
10888 false,
10889 Some((&mut sl.zq, &mut sl.zd)),
10890 )?;
10891 fa_q8 = true;
10892 } else if swa && b_swa > win && hd == 256 && rows_on {
10893 e.fa_decode_rows_w(
10894 &sl.q,
10895 &k_view,
10896 &v_view,
10897 &mut sl.attn,
10898 hd,
10899 nh,
10900 nkv,
10901 &kvl.len_d,
10902 -1,
10903 1,
10904 scale,
10905 win,
10906 kvl.k_tok_bytes,
10907 kvl.v_tok_bytes,
10908 Some((&mut sl.zq, &mut sl.zd)),
10909 )?;
10910 fa_q8 = true;
10911 } else {
10912 let b = if swa { b_swa } else { b_glob };
10913 e.fa_decode_dc(
10914 &sl.q,
10915 &k_view,
10916 &v_view,
10917 &mut sl.attn,
10918 hd,
10919 nh,
10920 nkv,
10921 &kvl.len_d,
10922 b,
10923 scale,
10924 kvl.k_tok_bytes,
10925 kvl.v_tok_bytes,
10926 swa && crate::Engine::wkv_on(),
10927 )?;
10928 }
10929 if !fa_q8 {
10930 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
10931 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
10932 }
10933 {
10934 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10935 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10936 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
10937 }
10938 Ok(())
10939 }
10940
10941 fn gemma4_layer_tail_slotted(
10944 &self,
10945 e: &Engine,
10946 layer: &crate::hybrid::HybridLayer,
10947 next_norm: Option<&CudaSlice<f32>>,
10948 sl: &mut G4DcSlots,
10949 ) -> Result<(), Box<dyn std::error::Error>> {
10950 let n_embd = self.cfg.n_embd as usize;
10951 let eps = self.cfg.rms_eps;
10952 let bits = layer.gemma4.as_ref().unwrap();
10953 let crate::hybrid::Ffn::Dense {
10954 ffn_gate,
10955 ffn_up,
10956 ffn_down,
10957 } = &layer.ffn
10958 else {
10959 return Err("slotted tail: dense ffn only".into());
10960 };
10961 let pnfold = Engine::g4_pnfold_on();
10962 if pnfold {
10963 let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
10966 let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
10967 e.rms_pre_add_rms_norm_q8z_into(
10968 or,
10969 layer.post_attn_norm.float_data(),
10970 xr,
10971 bits.ffn_norm.float_data(),
10972 &mut sl.attn_out,
10973 &mut sl.zsh,
10974 n_embd,
10975 1,
10976 eps,
10977 &mut sl.zq,
10978 &mut sl.zd,
10979 )?;
10980 } else {
10981 e.add_rms_norm(
10982 &sl.cur,
10983 &sl.x,
10984 bits.ffn_norm.float_data(),
10985 &mut sl.attn_out,
10986 &mut sl.zsh,
10987 n_embd,
10988 1,
10989 eps,
10990 )?;
10991 }
10992 let n_ff = ffn_gate.out_features();
10993 if !pnfold {
10994 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
10995 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10996 }
10997 {
10998 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10999 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11000 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
11001 && !e.matmul_nvfp4_fused2_into(
11002 ffn_gate,
11003 ffn_up,
11004 zq,
11005 zd,
11006 &mut sl.gate,
11007 &mut sl.up,
11008 )?
11009 {
11010 return Err("slotted tail: ffn fused2 unavailable".into());
11011 }
11012 }
11013 debug_assert!(e.uses_q8_1_fast(ffn_down));
11014 {
11015 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
11016 let upv = e.view(upr, n_ff);
11017 let up_all = upv.slice(0..n_ff);
11018 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
11019 e.gelu_tanh_mul_q8_1_into(
11020 gr,
11021 &up_all,
11022 &mut sl.act,
11023 n_ff,
11024 1,
11025 &mut sl.actq,
11026 &mut sl.actd,
11027 )?;
11028 }
11029 {
11030 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
11031 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
11032 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
11033 }
11034 if pnfold {
11035 if let Some(w) = next_norm {
11038 let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
11039 let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
11040 e.rms_pre_add_scale_rms_norm_q8_1_into(
11041 f0r,
11042 bits.post_ffw_norm.float_data(),
11043 aor,
11044 bits.layer_scale,
11045 w,
11046 &mut sl.xn,
11047 n_embd,
11048 1,
11049 eps,
11050 &mut sl.hq,
11051 &mut sl.hd_,
11052 )?;
11053 return Ok(());
11054 }
11055 }
11056 e.rms_norm(
11057 &sl.f0,
11058 bits.post_ffw_norm.float_data(),
11059 &mut sl.sn,
11060 n_embd,
11061 1,
11062 eps,
11063 )?;
11064 match next_norm {
11065 Some(w) => {
11066 e.add_scale_rms_norm_q8_1_into(
11067 &sl.sn,
11068 &sl.attn_out,
11069 bits.layer_scale,
11070 w,
11071 &mut sl.xn,
11072 n_embd,
11073 1,
11074 eps,
11075 &mut sl.hq,
11076 &mut sl.hd_,
11077 )?;
11078 }
11079 None => {
11080 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
11081 }
11082 }
11083 Ok(())
11084 }
11085
11086 #[allow(clippy::too_many_arguments)]
11088 fn gemma4_decode_attn_dc(
11089 &self,
11090 e: &Engine,
11091 fa: &crate::hybrid::FullAttnLayer,
11092 il: usize,
11093 hq: &CudaSlice<i8>,
11094 hdq: &CudaSlice<f32>,
11095 pos_d: &CudaSlice<i32>,
11096 cache: &mut Cache,
11097 cap_bucket_max: Option<(usize, usize)>,
11098 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11099 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11100 let eps = self.cfg.rms_eps;
11101 let aux = self.gemma4_aux.as_ref().unwrap();
11102 let ones = aux.ones(e);
11103 #[cfg(debug_assertions)]
11104 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
11105 let (q0, k0, v0) = if swa {
11106 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
11107 Some(t3) => t3,
11108 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11110 Some((q0, k0)) => {
11111 let h0 = e.zeros(0)?;
11112 let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
11113 (q0, k0, v0)
11114 }
11115 None => {
11116 let h0 = e.zeros(0)?;
11117 (
11118 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11119 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11120 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
11121 )
11122 }
11123 },
11124 }
11125 } else {
11126 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
11127 Some(p) => p,
11128 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11129 Some(p) => p,
11130 None => {
11131 let h0 = e.zeros(0)?;
11132 (
11133 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11134 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11135 )
11136 }
11137 },
11138 };
11139 let v0 = e.clone_dtod(&k0)?;
11140 (q0, k0, v0)
11141 };
11142 let mut q = e.uninit(nh * hd)?;
11143 let mut k = e.uninit(nkv * hd)?;
11144 let mut v = e.uninit(nkv * hd)?;
11145 let ff = if swa {
11147 None
11148 } else {
11149 Some(
11150 aux.rope_freqs(e)
11151 .expect("gemma4 global rope needs rope_freqs.weight"),
11152 )
11153 };
11154 #[cfg(debug_assertions)]
11155 if let Some(ff) = ff {
11156 crate::debug_assert_tensor_stream_device(
11157 ff,
11158 &e.stream(),
11159 "gemma4_decode_attn_dc.rope_freqs",
11160 );
11161 }
11162 let kvl = cache.kv[il].as_mut().unwrap();
11163 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
11164 if crate::Engine::qkv_append_on() {
11165 e.rms_norm_qkv_rope_append_dc(
11167 &q0,
11168 &k0,
11169 &v0,
11170 fa.q_norm.float_data(),
11171 fa.k_norm.float_data(),
11172 ones,
11173 &mut q,
11174 &mut k,
11175 &mut v,
11176 hd,
11177 nh,
11178 nkv,
11179 pos_d,
11180 nh,
11181 nkv,
11182 base,
11183 1.0,
11184 ff,
11185 eps,
11186 &mut kvl.k,
11187 &mut kvl.v,
11188 &kvl.len_d,
11189 kvl.k_tok_bytes,
11190 kvl.v_tok_bytes,
11191 kv_fp8,
11192 )?;
11193 } else {
11194 e.rms_norm_qkv_rope(
11195 &q0,
11196 &k0,
11197 &v0,
11198 fa.q_norm.float_data(),
11199 fa.k_norm.float_data(),
11200 ones,
11201 &mut q,
11202 &mut k,
11203 &mut v,
11204 hd,
11205 nh,
11206 nkv,
11207 pos_d,
11208 nh,
11209 nkv,
11210 base,
11211 1.0,
11212 ff,
11213 eps,
11214 )?;
11215 e.append_kv_quantized_dc(
11216 &k,
11217 &v,
11218 &mut kvl.k,
11219 &mut kvl.v,
11220 &kvl.len_d,
11221 kvl.kv_dim_k,
11222 kvl.kv_dim_v,
11223 kvl.k_tok_bytes,
11224 kvl.v_tok_bytes,
11225 kv_fp8,
11226 )?;
11227 }
11228 e.inc_seqlen(&mut kvl.len_d)?;
11229 let mut attn = e.uninit(nh * hd)?;
11230 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11233 match cap_bucket_max {
11238 None => {
11239 kvl.len += 1;
11243 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11244 if !swa
11245 && hd == 512
11246 && kvl.len >= crate::fa512_min_tkv()
11247 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11248 {
11249 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11252 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11253 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11254 e.fa_decode_rows(
11255 &q,
11256 &kp,
11257 &vp,
11258 &mut attn,
11259 hd,
11260 nh,
11261 nkv,
11262 kvl.len - 1,
11263 1,
11264 scale,
11265 kvl.k_tok_bytes,
11266 kvl.v_tok_bytes,
11267 Some((&kvl.len_d, -1)),
11268 false,
11269 false,
11270 Some((&mut aq8, &mut ad8)),
11271 )?;
11272 fa_q8 = Some((aq8, ad8));
11273 } else if swa
11274 && kvl.len > win
11275 && hd == 256
11276 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11277 {
11278 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11280 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11281 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11282 e.fa_decode_rows_w(
11283 &q,
11284 &kp,
11285 &vp,
11286 &mut attn,
11287 hd,
11288 nh,
11289 nkv,
11290 &kvl.len_d,
11291 -1,
11292 1,
11293 scale,
11294 win,
11295 kvl.k_tok_bytes,
11296 kvl.v_tok_bytes,
11297 Some((&mut aq8, &mut ad8)),
11298 )?;
11299 fa_q8 = Some((aq8, ad8));
11300 } else {
11301 let (off_tok, t_kv) = if swa && kvl.len > win {
11302 (kvl.len - win, win)
11303 } else {
11304 (0, kvl.len)
11305 };
11306 let k_view = e.view_u8_range(
11307 &kvl.k,
11308 off_tok * kvl.k_tok_bytes,
11309 (off_tok + t_kv) * kvl.k_tok_bytes,
11310 );
11311 let v_view = e.view_u8_range(
11312 &kvl.v,
11313 off_tok * kvl.v_tok_bytes,
11314 (off_tok + t_kv) * kvl.v_tok_bytes,
11315 );
11316 e.fa_decode_kvmod(
11317 &q,
11318 &k_view,
11319 &v_view,
11320 &mut attn,
11321 hd,
11322 nh,
11323 nkv,
11324 t_kv,
11325 scale,
11326 kvl.k_tok_bytes,
11327 kvl.v_tok_bytes,
11328 swa && crate::Engine::wkv_on(),
11329 )?;
11330 }
11331 }
11332 Some((b_swa, b_glob)) => {
11333 let k_view = e.view_u8(&kvl.k, kvl.k.len());
11339 let v_view = e.view_u8(&kvl.v, kvl.v.len());
11340 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11341 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11342 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11343 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11344 e.fa_decode_rows(
11345 &q,
11346 &k_view,
11347 &v_view,
11348 &mut attn,
11349 hd,
11350 nh,
11351 nkv,
11352 b_glob - 1,
11353 1,
11354 scale,
11355 kvl.k_tok_bytes,
11356 kvl.v_tok_bytes,
11357 Some((&kvl.len_d, -1)),
11358 false,
11359 false,
11360 Some((&mut aq8, &mut ad8)),
11361 )?;
11362 fa_q8 = Some((aq8, ad8));
11363 } else if swa && b_swa > win && hd == 256 && rows_on {
11364 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11365 e.fa_decode_rows_w(
11366 &q,
11367 &k_view,
11368 &v_view,
11369 &mut attn,
11370 hd,
11371 nh,
11372 nkv,
11373 &kvl.len_d,
11374 -1,
11375 1,
11376 scale,
11377 win,
11378 kvl.k_tok_bytes,
11379 kvl.v_tok_bytes,
11380 Some((&mut aq8, &mut ad8)),
11381 )?;
11382 fa_q8 = Some((aq8, ad8));
11383 } else {
11384 let b = if swa { b_swa } else { b_glob };
11385 e.fa_decode_dc(
11386 &q,
11387 &k_view,
11388 &v_view,
11389 &mut attn,
11390 hd,
11391 nh,
11392 nkv,
11393 &kvl.len_d,
11394 b,
11395 scale,
11396 kvl.k_tok_bytes,
11397 kvl.v_tok_bytes,
11398 swa && crate::Engine::wkv_on(),
11399 )?;
11400 }
11401 }
11402 }
11403 if let Some((aq8, ad8)) = fa_q8 {
11406 let mut y = e.uninit(fa.wo.out_features())?;
11407 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11408 return Ok(y);
11409 }
11410 Ok(e.matmul(&fa.wo, &attn, 1)?)
11411 }
11412
11413 pub fn gemma4_generate_graph(
11418 &self,
11419 e: &Engine,
11420 prompt_pos: usize,
11421 first_token: u32,
11422 cache: &mut Cache,
11423 max_new: usize,
11424 eos: &[u32],
11425 mut on_token: impl FnMut(u32) -> bool,
11426 ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11427 if self.is_gemma4_e4b() {
11428 return Err(
11429 "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11430 .into(),
11431 );
11432 }
11433 use crate::decode::StopReason;
11434 let n_vocab = self.output.out_features();
11435 let n_embd = self.cfg.n_embd as usize;
11436 let embd_gpu = self
11437 .embd_gpu
11438 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11439 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11440 for kvl in cache.kv.iter_mut().flatten() {
11441 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11442 }
11443 let mut token_d = e.stream().clone_htod(&[first_token])?;
11444 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11445 let g4 = self.cfg.gemma4.as_ref().unwrap();
11446 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11447 let nkv_s = g4
11449 .head_count_kv
11450 .iter()
11451 .zip(g4.swa_pattern.iter())
11452 .find(|p| *p.1)
11453 .map(|p| *p.0 as usize)
11454 .unwrap_or(8);
11455 let nkv_g = g4
11456 .head_count_kv
11457 .iter()
11458 .zip(g4.swa_pattern.iter())
11459 .find(|p| !*p.1)
11460 .map(|p| *p.0 as usize)
11461 .unwrap_or(2);
11462 let mut graphs: std::collections::HashMap<
11463 ((bool, usize), (bool, usize), bool, bool),
11464 (
11465 cudarc::driver::CudaGraph,
11466 Vec<Box<dyn std::any::Any + Send>>,
11467 ),
11468 > = Default::default();
11469 let mut slots = self.g4_dc_slots(e)?;
11472 const RING: usize = 64;
11475 const DRAIN: usize = 1;
11481 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11482 let ring_base = prompt_pos;
11483 let mut out = Vec::with_capacity(max_new);
11484 let mut reason = StopReason::MaxNew;
11485 let mut next = first_token;
11486 let mut captures = 0usize;
11487 for _ in 0..max_new {
11488 out.push(next);
11489 if eos.contains(&next) {
11490 reason = StopReason::Eos;
11491 break;
11492 }
11493 if !on_token(next) {
11494 reason = StopReason::Callback;
11495 break;
11496 }
11497 let t_kv = cache.pos + 1;
11498 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11506 let f512 = crate::fa512_min_tkv();
11507 let key_s = if t_kv > win {
11508 (true, usize::MAX)
11509 } else {
11510 e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11511 };
11512 let (key_g, rung_end) = if t_kv >= f512 {
11513 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11516 ((true, end), end)
11517 } else {
11518 (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11519 };
11520 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11521 if !graphs.contains_key(&key) {
11522 let bucket_max = (t_kv, rung_end);
11523 let snap = cache.snapshot(e)?;
11525 let pos_save = e.dtoh_i32_one(&pos_d)?;
11526 let len_save: Vec<Option<i32>> = cache
11527 .kv
11528 .iter()
11529 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11530 .collect();
11531 let tok_save = e.dtoh_u32_one(&token_d)?;
11532 let graph = {
11537 let tok_ref = &mut token_d;
11538 let pos_ref = &mut pos_d;
11539 let cache_ref = &mut *cache;
11540 let slots_ref = &mut slots;
11541 let ring_ref = &mut ring;
11542 e.capture_graph_retained_flags(
11543 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11544 |e| {
11545 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11547 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11548 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11549 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11550 cache_ref, n_vocab, Some(bucket_max),
11551 sl, tok_ref, Some((rg, ring_base)))
11552 })?
11553 };
11554 cache.rollback(e, &snap, 0)?;
11555 e.set_i32_one(&mut pos_d, pos_save)?;
11556 for (il, ls) in len_save.iter().enumerate() {
11557 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11558 e.set_i32_one(&mut kvl.len_d, *v)?;
11559 }
11560 }
11561 e.set_u32_one(&mut token_d, tok_save)?;
11562 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11563 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11564 eprintln!("[graph-census] {c:?}");
11565 }
11566 }
11567 graphs.insert(key, graph);
11568 captures += 1;
11569 }
11570 let mut chunk = 1usize;
11575 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11576 .ok()
11577 .and_then(|v| v.parse().ok())
11578 .unwrap_or(DRAIN);
11579 while chunk < drain_cap && out.len() + chunk < max_new {
11580 let t_next = cache.pos + 1 + chunk;
11581 let key_s2 = if t_next > win {
11582 (true, usize::MAX)
11583 } else {
11584 e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11585 };
11586 let key_g2 = if t_next >= f512 {
11587 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11588 } else {
11589 e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11590 };
11591 if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11592 break;
11593 }
11594 chunk += 1;
11595 }
11596 let g = &graphs.get(&key).unwrap().0;
11597 for _ in 0..chunk {
11598 g.launch()?;
11599 }
11600 e.stream().synchronize()?;
11601 let ringh = e.dtoh_u32(&ring)?;
11602 for j in 0..chunk {
11603 let pos_j = cache.pos + j;
11604 let tok_j = ringh[(pos_j - ring_base) % RING];
11605 cache.pos += 0; if j + 1 == chunk {
11607 next = tok_j;
11608 } else {
11609 out.push(tok_j);
11610 if eos.contains(&tok_j) || !on_token(tok_j) {
11611 reason = if eos.contains(&tok_j) {
11612 StopReason::Eos
11613 } else {
11614 StopReason::Callback
11615 };
11616 let keep = cache.pos + j + 1;
11618 e.set_i32_one(&mut pos_d, keep as i32)?;
11619 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11620 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11621 kvl.len = keep;
11622 }
11623 cache.pos = keep;
11624 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11625 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11626 }
11627 return Ok((out, reason));
11628 }
11629 }
11630 }
11631 cache.pos += chunk;
11632 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11633 kvl.len += chunk;
11634 }
11635 }
11636 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11637 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11638 }
11639 Ok((out, reason))
11640 }
11641
11642 pub(crate) fn gemma4_decode_step_t(
11648 &self,
11649 e: &Engine,
11650 tokens: &[u32],
11651 pos0: usize,
11652 cache: &mut Cache,
11653 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11654 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11655 }
11656
11657 pub(crate) fn gemma4_decode_step_t_am(
11661 &self,
11662 e: &Engine,
11663 tokens: &[u32],
11664 pos0: usize,
11665 cache: &mut Cache,
11666 ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11667 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11668 let t = tokens.len();
11669 let n_vocab = self.output.out_features();
11670 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11671 for i in 0..t {
11672 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11673 }
11674 Ok((e.dtoh_u32(&toks)?, hn))
11675 }
11676
11677 pub(crate) fn gemma4_decode_step_t_am_dev(
11680 &self,
11681 e: &Engine,
11682 tok_d: &CudaSlice<u32>,
11683 t: usize,
11684 pos0: usize,
11685 cache: &mut Cache,
11686 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11687 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11688 let n_vocab = self.output.out_features();
11689 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11690 for i in 0..t {
11691 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11692 }
11693 Ok((vam, hn))
11694 }
11695
11696 pub(crate) fn gemma4_decode_step_t_h(
11699 &self,
11700 e: &Engine,
11701 tokens: &[u32],
11702 pos0: usize,
11703 cache: &mut Cache,
11704 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11705 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11706 let t = tokens.len();
11707 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11708 e.softcap(&mut ld, cap, t * self.output.out_features())?;
11709 Ok((e.dtoh(&ld)?, hn))
11710 }
11711
11712 pub(crate) fn verify_stream_scratch(
11715 &self,
11716 e: &Engine,
11717 cap: usize,
11718 ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11719 Ok(VerifyStreamScratch {
11720 pos_d: e.htod_i32(&vec![0i32; cap])?,
11721 row_ctrs: (0..cap)
11722 .map(|_| e.htod_i32(&[0]))
11723 .collect::<Result<_, _>>()?,
11724 })
11725 }
11726
11727 pub(crate) fn gemma4_verify_t_am_stream(
11735 &self,
11736 e: &Engine,
11737 tok_d: &CudaSlice<u32>,
11738 t: usize,
11739 ctr: &CudaSlice<i32>,
11740 hint: usize,
11741 cache: &mut Cache,
11742 scr: &mut VerifyStreamScratch,
11743 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11744 let n_embd = self.cfg.n_embd as usize;
11745 let eps = self.cfg.rms_eps;
11746 assert!(t <= scr.row_ctrs.len() && t <= 64);
11747 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11748 for i in 0..t {
11749 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11750 }
11751 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11752 let embd_gpu = self
11753 .embd_gpu
11754 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11755 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11756 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11757 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11758 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11759 let n_layers = self.layers.len();
11760 for (il, layer) in self.layers.iter().enumerate() {
11761 let (hq, hdq) = match h_carry.take() {
11762 Some(p) => p,
11763 None => {
11764 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11765 }
11766 };
11767 let Mixer::Full(fa) = &layer.mixer else {
11768 panic!("gemma4 layer {il} not full-attn")
11769 };
11770 let o = self
11771 .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11772 let next_norm = if il + 1 < n_layers {
11773 Some(self.layers[il + 1].attn_norm.float_data())
11774 } else {
11775 None
11776 };
11777 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11778 x = xn;
11779 h_carry = hn;
11780 self.dflash_tap(e, cache, il, &x, t)?;
11781 }
11782 let mut hn = e.uninit(t * n_embd)?;
11783 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11784 let ld = e.matmul(&self.output, &hn, t)?;
11785 let n_vocab = self.output.out_features();
11786 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11787 for i in 0..t {
11788 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11789 }
11790 Ok((vam, hn))
11791 }
11792
11793 fn dflash_tap(
11800 &self,
11801 e: &Engine,
11802 cache: &mut Cache,
11803 il: usize,
11804 x: &CudaSlice<f32>,
11805 t: usize,
11806 ) -> Result<(), Box<dyn std::error::Error>> {
11807 let Some(taps) = cache.dflash_taps.as_mut() else {
11808 return Ok(());
11809 };
11810 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11811 return Ok(());
11812 };
11813 let h = taps.hidden;
11814 let n_taps = taps.layer_ids.len();
11815 debug_assert_eq!(taps.t, t);
11816 let xv = e.view(x, t * h);
11817 for r in 0..t {
11818 let row = xv.slice(r * h..(r + 1) * h);
11819 e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
11820 }
11821 Ok(())
11822 }
11823
11824 fn gemma4_verify_trunk(
11825 &self,
11826 e: &Engine,
11827 tokens: &[u32],
11828 pos0: usize,
11829 cache: &mut Cache,
11830 tok_dev: Option<&CudaSlice<u32>>,
11831 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11832 let n_embd = self.cfg.n_embd as usize;
11833 let eps = self.cfg.rms_eps;
11834 let t = tokens.len();
11835 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
11836 let pos_d = e.htod_i32(&pos)?;
11837 let mut x = match tok_dev {
11838 Some(td) => {
11839 let embd_gpu = self
11840 .embd_gpu
11841 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11842 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11843 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
11844 }
11845 None => e.htod(&self.embd.gather(n_embd, tokens))?,
11846 };
11847 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11848 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11849 let n_layers = self.layers.len();
11850 for (il, layer) in self.layers.iter().enumerate() {
11851 let (hq, hdq) = match h_carry.take() {
11852 Some(p) => p,
11853 None => {
11854 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11855 }
11856 };
11857 let Mixer::Full(fa) = &layer.mixer else {
11858 panic!("gemma4 layer {il} not full-attn")
11859 };
11860 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
11861 let next_norm = if il + 1 < n_layers {
11862 Some(self.layers[il + 1].attn_norm.float_data())
11863 } else {
11864 None
11865 };
11866 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11867 x = xn;
11868 h_carry = hn;
11869 self.dflash_tap(e, cache, il, &x, t)?;
11870 }
11871 let mut hn = e.uninit(t * n_embd)?;
11872 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11873 let mut ld = e.matmul(&self.output, &hn, t)?;
11874 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
11876 Ok((ld, hn))
11877 }
11878
11879 #[allow(clippy::too_many_arguments)]
11887 fn gemma4_verify_attn_stream(
11888 &self,
11889 e: &Engine,
11890 fa: &crate::hybrid::FullAttnLayer,
11891 il: usize,
11892 hq: &CudaSlice<i8>,
11893 hdq: &CudaSlice<f32>,
11894 pos_d: &CudaSlice<i32>,
11895 t: usize,
11896 cache: &mut Cache,
11897 hint: usize,
11898 row_ctrs: &[CudaSlice<i32>],
11899 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11900 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11901 let eps = self.cfg.rms_eps;
11902 let aux = self.gemma4_aux.as_ref().unwrap();
11903 let ones = aux.ones(e);
11904 #[cfg(debug_assertions)]
11905 crate::debug_assert_tensor_stream_device(
11906 ones,
11907 &e.stream(),
11908 "gemma4_verify_attn_stream.ones",
11909 );
11910 let h0 = e.zeros(0)?;
11911 let h = &h0;
11912 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11915 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11916 let fused_qkv = if f2b {
11917 if swa {
11918 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11919 .map(|(a, b, c)| (a, b, Some(c)))
11920 } else {
11921 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11922 .map(|(a, b)| (a, b, None))
11923 }
11924 } else {
11925 None
11926 };
11927 let (q0, k0, v0) = match fused_qkv {
11928 Some((a, b, cv)) => {
11929 let v = match cv {
11930 Some(c) => c,
11931 None => e.clone_dtod(&b)?,
11932 };
11933 (a, b, v)
11934 }
11935 None => {
11936 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11937 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11938 let v0 = if swa {
11939 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
11940 } else {
11941 e.clone_dtod(&k0)?
11942 };
11943 (q0, k0, v0)
11944 }
11945 };
11946 let mut q = e.uninit(t * nh * hd)?;
11947 let mut k = e.uninit(t * nkv * hd)?;
11948 let mut v = e.uninit(t * nkv * hd)?;
11949 let ff = if swa {
11952 None
11953 } else {
11954 Some(
11955 aux.rope_freqs(e)
11956 .expect("gemma4 global rope needs rope_freqs.weight"),
11957 )
11958 };
11959 #[cfg(debug_assertions)]
11960 if let Some(ff) = ff {
11961 crate::debug_assert_tensor_stream_device(
11962 ff,
11963 &e.stream(),
11964 "gemma4_verify_attn_stream.rope_freqs",
11965 );
11966 }
11967 e.rms_norm_qkv_rope(
11968 &q0,
11969 &k0,
11970 &v0,
11971 fa.q_norm.float_data(),
11972 fa.k_norm.float_data(),
11973 ones,
11974 &mut q,
11975 &mut k,
11976 &mut v,
11977 hd,
11978 nh * t,
11979 nkv * t,
11980 pos_d,
11981 nh,
11982 nkv,
11983 base,
11984 1.0,
11985 ff,
11986 eps,
11987 )?;
11988 let kvl = cache.kv[il].as_mut().unwrap();
11989 e.append_kv_quantized_rows_dc(
11991 &k,
11992 &v,
11993 &mut kvl.k,
11994 &mut kvl.v,
11995 &kvl.len_d,
11996 t,
11997 kvl.kv_dim_k,
11998 kvl.kv_dim_v,
11999 kvl.k_tok_bytes,
12000 kvl.v_tok_bytes,
12001 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12002 )?;
12003 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12006 let mut attn = e.uninit(t * nh * hd)?;
12007 let k_view = e.view_u8(&kvl.k, kvl.k.len());
12008 let v_view = e.view_u8(&kvl.v, kvl.v.len());
12009 if swa && hint + 1 >= win {
12012 e.fa_decode_rows_w(
12015 &q,
12016 &k_view,
12017 &v_view,
12018 &mut attn,
12019 hd,
12020 nh,
12021 nkv,
12022 &kvl.len_d,
12023 0,
12024 t,
12025 scale,
12026 win,
12027 kvl.k_tok_bytes,
12028 kvl.v_tok_bytes,
12029 None,
12030 )?;
12031 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
12032 let bucket = (hint + t + 2)
12045 .next_power_of_two()
12046 .min(crate::fa512_min_tkv().saturating_sub(1));
12047 let qv = e.view(&q, t * nh * hd);
12048 for i in 0..t {
12049 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
12050 let mut q_one = e.uninit(nh * hd)?;
12051 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12052 let mut a_one = e.uninit(nh * hd)?;
12053 e.fa_decode_dc(
12054 &q_one,
12055 &k_view,
12056 &v_view,
12057 &mut a_one,
12058 hd,
12059 nh,
12060 nkv,
12061 &row_ctrs[i],
12062 bucket,
12063 scale,
12064 kvl.k_tok_bytes,
12065 kvl.v_tok_bytes,
12066 false,
12067 )?;
12068 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12069 }
12070 } else if hd == 512 {
12071 e.fa_decode_rows(
12074 &q,
12075 &k_view,
12076 &v_view,
12077 &mut attn,
12078 hd,
12079 nh,
12080 nkv,
12081 hint,
12082 t,
12083 scale,
12084 kvl.k_tok_bytes,
12085 kvl.v_tok_bytes,
12086 Some((&kvl.len_d, 0)),
12087 false,
12088 false,
12089 None,
12090 )?;
12091 } else {
12092 e.fa_decode_rows_dc(
12094 &q,
12095 &k_view,
12096 &v_view,
12097 &mut attn,
12098 hd,
12099 nh,
12100 nkv,
12101 &kvl.len_d,
12102 hint + t,
12103 t,
12104 scale,
12105 kvl.k_tok_bytes,
12106 kvl.v_tok_bytes,
12107 0,
12108 swa && crate::Engine::wkv_on(),
12109 )?;
12110 }
12111 Ok(e.matmul(&fa.wo, &attn, t)?)
12112 }
12113
12114 fn gemma4_verify_attn(
12115 &self,
12116 e: &Engine,
12117 fa: &crate::hybrid::FullAttnLayer,
12118 il: usize,
12119 hq: &CudaSlice<i8>,
12120 hdq: &CudaSlice<f32>,
12121 pos_d: &CudaSlice<i32>,
12122 t: usize,
12123 cache: &mut Cache,
12124 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12125 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12126 let eps = self.cfg.rms_eps;
12127 let aux = self.gemma4_aux.as_ref().unwrap();
12128 let ones = aux.ones(e);
12129 #[cfg(debug_assertions)]
12130 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
12131 let n_embd = self.cfg.n_embd as usize;
12132 let _ = n_embd;
12133
12134 let h0 = e.zeros(0)?;
12135 let h = &h0;
12136 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12139 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12140 let fused_qkv = if f2b {
12141 if swa {
12142 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12143 .map(|(a, b, c)| (a, b, Some(c)))
12144 } else {
12145 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12146 .map(|(a, b)| (a, b, None))
12147 }
12148 } else {
12149 None
12150 };
12151 let (q0, k0, v0) = match fused_qkv {
12152 Some((a, b, cv)) => {
12153 let v = match cv {
12154 Some(c) => c,
12155 None => e.clone_dtod(&b)?,
12156 };
12157 (a, b, v)
12158 }
12159 None => {
12160 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12161 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12162 let v0 = if swa {
12163 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12164 } else {
12165 e.clone_dtod(&k0)?
12166 };
12167 (q0, k0, v0)
12168 }
12169 };
12170 let mut q = e.uninit(t * nh * hd)?;
12171 let mut k = e.uninit(t * nkv * hd)?;
12172 let mut v = e.uninit(t * nkv * hd)?;
12173 let ff = if swa {
12176 None
12177 } else {
12178 Some(
12179 aux.rope_freqs(e)
12180 .expect("gemma4 global rope needs rope_freqs.weight"),
12181 )
12182 };
12183 #[cfg(debug_assertions)]
12184 if let Some(ff) = ff {
12185 crate::debug_assert_tensor_stream_device(
12186 ff,
12187 &e.stream(),
12188 "gemma4_verify_attn.rope_freqs",
12189 );
12190 }
12191 e.rms_norm_qkv_rope(
12192 &q0,
12193 &k0,
12194 &v0,
12195 fa.q_norm.float_data(),
12196 fa.k_norm.float_data(),
12197 ones,
12198 &mut q,
12199 &mut k,
12200 &mut v,
12201 hd,
12202 nh * t,
12203 nkv * t,
12204 pos_d,
12205 nh,
12206 nkv,
12207 base,
12208 1.0,
12209 ff,
12210 eps,
12211 )?;
12212 let kvl = cache.kv[il].as_mut().unwrap();
12213 let base_len = kvl.len;
12214 e.append_kv_quantized_rows(
12215 &k,
12216 &v,
12217 &mut kvl.k,
12218 &mut kvl.v,
12219 base_len,
12220 t,
12221 kvl.kv_dim_k,
12222 kvl.kv_dim_v,
12223 kvl.k_tok_bytes,
12224 kvl.v_tok_bytes,
12225 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12226 )?;
12227 kvl.len += t;
12228 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12229 let mut attn = e.uninit(t * nh * hd)?;
12230 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
12233 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
12236 if rows_ok && (!swa || base_len + t <= win) {
12237 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12238 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12239 if hd == 512 {
12240 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12242 e.fa_decode_rows(
12243 &q,
12244 &k_view,
12245 &v_view,
12246 &mut attn,
12247 hd,
12248 nh,
12249 nkv,
12250 base_len,
12251 t,
12252 scale,
12253 kvl.k_tok_bytes,
12254 kvl.v_tok_bytes,
12255 Some((&kvl.len_d, 0)),
12256 false,
12257 swa && crate::Engine::wkv_on(),
12258 None,
12259 )?;
12260 } else {
12261 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12265 e.fa_decode_rows_dc(
12266 &q,
12267 &k_view,
12268 &v_view,
12269 &mut attn,
12270 hd,
12271 nh,
12272 nkv,
12273 &kvl.len_d,
12274 base_len + t,
12275 t,
12276 scale,
12277 kvl.k_tok_bytes,
12278 kvl.v_tok_bytes,
12279 0,
12280 swa && crate::Engine::wkv_on(),
12281 )?;
12282 }
12283 return Ok(e.matmul(&fa.wo, &attn, t)?);
12284 }
12285 if hd == 256
12293 && swa
12294 && base_len + 1 >= win
12295 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12296 {
12297 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12298 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12299 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12300 e.fa_decode_rows_w(
12301 &q,
12302 &k_view,
12303 &v_view,
12304 &mut attn,
12305 hd,
12306 nh,
12307 nkv,
12308 &kvl.len_d,
12309 0,
12310 t,
12311 scale,
12312 win,
12313 kvl.k_tok_bytes,
12314 kvl.v_tok_bytes,
12315 None,
12316 )?;
12317 return Ok(e.matmul(&fa.wo, &attn, t)?);
12318 }
12319 for i in 0..t {
12320 let avail = base_len + i + 1;
12321 let (off_tok, t_kv) = if swa && avail > win {
12322 (avail - win, win)
12323 } else {
12324 (0, avail)
12325 };
12326 let k_view = e.view_u8_range(
12327 &kvl.k,
12328 off_tok * kvl.k_tok_bytes,
12329 (off_tok + t_kv) * kvl.k_tok_bytes,
12330 );
12331 let v_view = e.view_u8_range(
12332 &kvl.v,
12333 off_tok * kvl.v_tok_bytes,
12334 (off_tok + t_kv) * kvl.v_tok_bytes,
12335 );
12336 let qi = e.view(&q, t * nh * hd);
12337 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12338 let mut q_one = e.uninit(nh * hd)?;
12339 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12340 let mut a_one = e.uninit(nh * hd)?;
12341 if swa
12345 && avail > win
12346 && hd == 256
12347 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12348 {
12349 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12350 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12351 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12352 e.fa_decode_rows_w(
12353 &q_one,
12354 &kp,
12355 &vp,
12356 &mut a_one,
12357 hd,
12358 nh,
12359 nkv,
12360 &kvl.len_d,
12361 0,
12362 1,
12363 scale,
12364 win,
12365 kvl.k_tok_bytes,
12366 kvl.v_tok_bytes,
12367 None,
12368 )?;
12369 } else if !swa
12370 && hd == 512
12371 && avail >= crate::fa512_min_tkv()
12372 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12373 {
12374 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12375 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12376 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12377 e.fa_decode_rows(
12378 &q_one,
12379 &kp,
12380 &vp,
12381 &mut a_one,
12382 hd,
12383 nh,
12384 nkv,
12385 avail - 1,
12386 1,
12387 scale,
12388 kvl.k_tok_bytes,
12389 kvl.v_tok_bytes,
12390 Some((&kvl.len_d, 0)),
12391 false,
12392 false,
12393 None,
12394 )?;
12395 } else {
12396 e.fa_decode_kvmod(
12397 &q_one,
12398 &k_view,
12399 &v_view,
12400 &mut a_one,
12401 hd,
12402 nh,
12403 nkv,
12404 t_kv,
12405 scale,
12406 kvl.k_tok_bytes,
12407 kvl.v_tok_bytes,
12408 swa && crate::Engine::wkv_on(),
12409 )?;
12410 }
12411 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12412 }
12413 Ok(e.matmul(&fa.wo, &attn, t)?)
12414 }
12415
12416 pub(crate) fn gemma4_decode_step_h(
12419 &self,
12420 e: &Engine,
12421 token: u32,
12422 cache: &mut Cache,
12423 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12424 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12429 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12430 }
12431 if crate::pp::pp_cuts(self.layers.len()).is_some() {
12432 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12433 }
12434 let n_embd = self.cfg.n_embd as usize;
12435 let eps = self.cfg.rms_eps;
12436 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12437 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12438 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12439 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12442 let n_layers = self.layers.len();
12443 for (il, layer) in self.layers.iter().enumerate() {
12444 let (hq, hdq) = match h_carry.take() {
12445 Some(p) => p,
12446 None => {
12447 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12448 }
12449 };
12450 let Mixer::Full(fa) = &layer.mixer else {
12451 panic!("gemma4 layer {il} not full-attn")
12452 };
12453 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12454 let next_norm = if il + 1 < n_layers {
12455 Some(self.layers[il + 1].attn_norm.float_data())
12456 } else {
12457 None
12458 };
12459 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12460 x = xn;
12461 h_carry = hn;
12462 }
12463 let mut hn = e.uninit(n_embd)?;
12464 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12465 let h_seed = e.clone_dtod(&x)?;
12466 let mut ld = e.matmul(&self.output, &hn, 1)?;
12467 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12468 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
12470 let logits = e.dtoh(&ld)?;
12471 cache.pos += 1;
12472 Ok((logits, h_seed))
12473 }
12474
12475 fn gemma4_decode_layers(
12483 &self,
12484 e: &Engine,
12485 mut x: CudaSlice<f32>,
12486 lo: usize,
12487 hi: usize,
12488 pos_d: &CudaSlice<i32>,
12489 cache: &mut Cache,
12490 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12491 let n_embd = self.cfg.n_embd as usize;
12492 let eps = self.cfg.rms_eps;
12493 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12494 for il in lo..hi {
12495 let layer = &self.layers[il];
12496 let (hq, hdq) = match h_carry.take() {
12497 Some(p) => p,
12498 None => {
12500 e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12501 }
12502 };
12503 let Mixer::Full(fa) = &layer.mixer else {
12504 panic!("gemma4 layer {il} not full-attn")
12505 };
12506 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12507 let next_norm = if il + 1 < hi {
12508 Some(self.layers[il + 1].attn_norm.float_data())
12509 } else {
12510 None
12511 };
12512 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12513 x = xn;
12514 h_carry = hn;
12515 }
12516 Ok(x)
12517 }
12518
12519 fn gemma4_decode_step_h_pp2(
12527 &self,
12528 e: &Engine,
12529 token: u32,
12530 cache: &mut Cache,
12531 split: usize,
12532 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12533 if crate::pp::pp2_streams_off() {
12534 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12535 }
12536 let rt = crate::pp::Pp2Rt::get(e)?;
12537 let e0 = rt.engine(0, e);
12538 let e1 = rt.engine(1, e);
12539 let n_embd = self.cfg.n_embd as usize;
12540 let eps = self.cfg.rms_eps;
12541 let pos = cache.pos as i32;
12542
12543 let slot = {
12545 let _st0 = rt.enter(0);
12546 let pos_d = e0.htod_i32(&[pos])?;
12547 #[cfg(debug_assertions)]
12548 crate::debug_assert_tensor_stream_device(
12549 &pos_d,
12550 &e0.stream(),
12551 "gemma4_decode_step_h_pp2.stage0.pos_d",
12552 );
12553 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12554 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12555 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12556 rt.tx(0, &x, n_embd)?
12557 };
12558
12559 let _st1 = rt.enter(1);
12561 let pos_d = e1.htod_i32(&[pos])?;
12562 #[cfg(debug_assertions)]
12563 crate::debug_assert_tensor_stream_device(
12564 &pos_d,
12565 &e1.stream(),
12566 "gemma4_decode_step_h_pp2.stage1.pos_d",
12567 );
12568 let x = rt.rx(0, slot, n_embd)?;
12569 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12570
12571 let mut hn = e1.uninit(n_embd)?;
12572 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12573 let h_seed = e1.clone_dtod(&x)?;
12574 let mut ld = e1.matmul(&self.output, &hn, 1)?;
12575 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12576 e1.softcap(&mut ld, cap, self.output.out_features())?;
12577 self.gemma4_suppress(e1, &mut ld, 1)?;
12578 let logits = e1.dtoh(&ld)?;
12579 cache.pos += 1;
12580 Ok((logits, h_seed))
12581 }
12582
12583 fn gemma4_decode_step_h_pp2_samestream(
12586 &self,
12587 e: &Engine,
12588 token: u32,
12589 cache: &mut Cache,
12590 split: usize,
12591 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12592 let n_embd = self.cfg.n_embd as usize;
12593 let eps = self.cfg.rms_eps;
12594 let pos_d = e.htod_i32(&[cache.pos as i32])?;
12595
12596 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12598 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12599 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12600
12601 let boundary_tx = e.clone_dtod(&x)?;
12603 let boundary_rx = e.clone_dtod(&boundary_tx)?;
12604
12605 let x =
12607 self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12608
12609 let mut hn = e.uninit(n_embd)?;
12610 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12611 let h_seed = e.clone_dtod(&x)?;
12612 let mut ld = e.matmul(&self.output, &hn, 1)?;
12613 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12614 e.softcap(&mut ld, cap, self.output.out_features())?;
12615 self.gemma4_suppress(e, &mut ld, 1)?;
12616 let logits = e.dtoh(&ld)?;
12617 cache.pos += 1;
12618 Ok((logits, h_seed))
12619 }
12620}
12621
12622impl HybridModel {
12641 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12644 let geometry = self
12645 .cfg
12646 .layer_geometry(il as u32)
12647 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12648 debug_assert_eq!(
12649 geometry.attention_gate,
12650 memra_gguf::config::AttentionGateKind::SeparateHead
12651 );
12652 geometry
12653 }
12654
12655 #[allow(clippy::too_many_arguments)]
12715 fn step35_attn_pre_wo(
12716 &self,
12717 e: &Engine,
12718 fa: &FullAttnLayer,
12719 mut g3: Vec<CudaSlice<f32>>,
12720 hg: Option<&CudaSlice<f32>>,
12721 gt_pre: Option<&CudaSlice<f32>>,
12722 pos_d: &CudaSlice<i32>,
12723 t: usize,
12724 cache: Option<&mut Cache>,
12725 il: usize,
12726 seq_end: usize,
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
12739 let v = g3.pop().unwrap();
12740 let k0 = g3.pop().unwrap();
12741 let q0 = g3.pop().unwrap();
12742
12743 let mut q = e.uninit(t * nh * hd)?;
12747 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12748 let mut k = e.uninit(t * nkv * hd)?;
12749 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12750 let ff = if geometry.rope_factors {
12751 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12752 } else {
12753 None
12754 };
12755 #[cfg(debug_assertions)]
12756 if let Some(ff) = ff {
12757 crate::debug_assert_tensor_stream_device(
12758 ff,
12759 &e.stream(),
12760 "step35_attn_pre_wo.rope_freqs",
12761 );
12762 }
12763 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12764
12765 let mut attn = e.uninit(t * nh * hd)?;
12766 match cache {
12767 Some(cache) => {
12768 let base_len = cache.kv[il].as_ref().unwrap().len;
12769 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12771 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12772 let off = if swa {
12773 let raw = base_len.saturating_sub(win - 1);
12774 if legacy_tkv || legacy_calllocal {
12775 raw
12776 } else {
12777 raw & !31usize
12778 }
12779 } else {
12780 0
12781 };
12782 {
12783 let kvl = cache.kv[il].as_mut().unwrap();
12784 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12785 let write_row = e.prepare_kv_append(kvl, off, t)?;
12786 e.append_kv_quantized_rows(
12787 &k,
12788 &v,
12789 &mut kvl.k,
12790 &mut kvl.v,
12791 write_row,
12792 t,
12793 kvl.kv_dim_k,
12794 kvl.kv_dim_v,
12795 kvl.k_tok_bytes,
12796 kvl.v_tok_bytes,
12797 crate::Engine::kv_fp8_on(),
12798 )?;
12799 kvl.len += t;
12800 let new_len = kvl.len as i32;
12801 e.set_i32_one(&mut kvl.len_d, new_len)?;
12802 }
12803 let kvl = cache.kv[il].as_ref().unwrap();
12804 let t_kv = base_len + t - off;
12827 let physical = kvl.physical_rows(off, off + t_kv)?;
12828 let k_view = e.view_u8_range(
12829 &kvl.k,
12830 physical.start * kvl.k_tok_bytes,
12831 physical.end * kvl.k_tok_bytes,
12832 );
12833 let v_view = e.view_u8_range(
12834 &kvl.v,
12835 physical.start * kvl.v_tok_bytes,
12836 physical.end * kvl.v_tok_bytes,
12837 );
12838 let swa_naive = if legacy_tkv {
12850 t_kv > win
12851 } else {
12852 seq_end > win
12853 };
12854 if swa && swa_naive {
12855 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
12868 e.sdpa_naive_w_quantized_view(
12869 &q,
12870 &k_view,
12871 &v_view,
12872 &mut attn,
12873 hd,
12874 nh,
12875 nkv,
12876 t,
12877 t_kv,
12878 scale,
12879 true,
12880 win,
12881 kvl.k_tok_bytes,
12882 kvl.v_tok_bytes,
12883 )?;
12884 } else {
12885 e.fa_prefill_view_ws_w_hd128(
12886 &q,
12887 &k_view,
12888 &v_view,
12889 &mut attn,
12890 hd,
12891 nh,
12892 nkv,
12893 t,
12894 t_kv,
12895 scale,
12896 true,
12897 win,
12898 kvl.k_tok_bytes,
12899 kvl.v_tok_bytes,
12900 )?;
12901 }
12902 } else if std::env::var("MEMRA_NOFA").is_ok() {
12903 e.sdpa_naive_quantized_view(
12904 &q,
12905 &k_view,
12906 &v_view,
12907 &mut attn,
12908 hd,
12909 nh,
12910 nkv,
12911 t,
12912 t_kv,
12913 scale,
12914 true,
12915 kvl.k_tok_bytes,
12916 kvl.v_tok_bytes,
12917 )?;
12918 } else {
12919 e.fa_prefill_view_ws(
12924 &q,
12925 &k_view,
12926 &v_view,
12927 &mut attn,
12928 hd,
12929 nh,
12930 nkv,
12931 t,
12932 t_kv,
12933 scale,
12934 true,
12935 kvl.k_tok_bytes,
12936 kvl.v_tok_bytes,
12937 crate::Engine::kv_fp8_on(),
12938 )?;
12939 }
12940 }
12941 None => {
12942 debug_assert_eq!(
12947 seq_end, t,
12948 "step35 cacheless prefill is monolithic (seq_end == t)"
12949 );
12950 if swa && seq_end > win {
12951 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
12952 } else if std::env::var("MEMRA_NOFA").is_ok() {
12953 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12954 } else {
12955 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12956 }
12957 }
12958 }
12959
12960 let gw = fa
12963 .attn_gate
12964 .as_ref()
12965 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
12966 let gt_owned = if gt_pre.is_none() {
12967 Some(e.matmul(
12968 gw,
12969 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
12970 t,
12971 )?)
12972 } else {
12973 None
12974 };
12975 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
12976 let mut ag = e.uninit(t * nh * hd)?;
12977 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
12978 Ok(ag)
12979 }
12980
12981 pub(crate) fn step35_attn(
12984 &self,
12985 e: &Engine,
12986 fa: &FullAttnLayer,
12987 h: &CudaSlice<f32>,
12988 pos_d: &CudaSlice<i32>,
12989 t: usize,
12990 il: usize,
12991 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12992 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
12993 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
12995 Ok(e.matmul(&fa.wo, &ag, t)?)
12996 }
12997
12998 #[allow(clippy::too_many_arguments)]
13005 pub(crate) fn step35_attn_prime(
13006 &self,
13007 e: &Engine,
13008 fa: &FullAttnLayer,
13009 h: &CudaSlice<f32>,
13010 hx: Option<&CudaSlice<u8>>,
13011 pos_d: &CudaSlice<i32>,
13012 t: usize,
13013 cache: &mut Cache,
13014 il: usize,
13015 seq_end: usize,
13016 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13017 let g3 = match hx {
13018 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
13019 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
13020 };
13021 let ag =
13022 self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
13023 Ok(e.matmul(&fa.wo, &ag, t)?)
13024 }
13025
13026 #[allow(clippy::too_many_arguments)]
13036 pub(crate) fn step35_decode_attn(
13037 &self,
13038 e: &Engine,
13039 fa: &FullAttnLayer,
13040 il: usize,
13041 h: &CudaSlice<f32>,
13042 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
13043 pos_d: &CudaSlice<i32>,
13044 cache: &mut Cache,
13045 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13046 let geometry = self.step35_geom(il);
13047 let hd = geometry.head_dim_k as usize;
13048 let nkv = geometry.n_head_kv as usize;
13049 let nh = geometry.n_head as usize;
13050 let rbase = geometry.rope_base;
13051 let scale = geometry.attention_scale();
13052 let swa = geometry.window.is_some();
13053 let eps = self.cfg.rms_eps;
13054 let win = geometry.window.unwrap_or(0) as usize;
13055 let n_rot = geometry.n_rot as usize;
13056 let n_embd = self.cfg.n_embd as usize;
13057 let gw = fa
13058 .attn_gate
13059 .as_ref()
13060 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13061
13062 let (q0, k0, v0, gt) = match pre_q {
13063 Some((hq, hdq)) => {
13064 debug_assert!(
13065 e.uses_q8_1_fast(gw),
13066 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
13067 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
13068 );
13069 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
13070 Some(t3) => t3,
13071 None => (
13072 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
13073 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
13074 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
13075 ),
13076 };
13077 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
13078 (a, b, c, gt)
13079 }
13080 None => {
13081 if e.uses_q8_1_fast(&fa.wq)
13082 && e.uses_q8_1_fast(&fa.wk)
13083 && e.uses_q8_1_fast(&fa.wv)
13084 && e.uses_q8_1_fast(gw)
13085 {
13086 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
13087 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
13088 Some(t3) => t3,
13089 None => (
13090 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
13091 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
13092 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
13093 ),
13094 };
13095 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
13096 (a, b, c, gt)
13097 } else {
13098 (
13099 e.matmul(&fa.wq, h, 1)?,
13100 e.matmul(&fa.wk, h, 1)?,
13101 e.matmul(&fa.wv, h, 1)?,
13102 e.matmul(gw, h, 1)?,
13103 )
13104 }
13105 }
13106 };
13107
13108 let mut q = e.uninit(nh * hd)?;
13109 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
13110 let mut k = e.uninit(nkv * hd)?;
13111 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
13112 let ff = if swa {
13113 None
13114 } else {
13115 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
13116 };
13117 #[cfg(debug_assertions)]
13118 if let Some(ff) = ff {
13119 crate::debug_assert_tensor_stream_device(
13120 ff,
13121 &e.stream(),
13122 "step35_decode_attn.rope_freqs",
13123 );
13124 }
13125 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
13126
13127 if std::env::var("MEMRA_NOFA").is_ok() {
13128 return Err(
13129 "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
13130 cache; unset MEMRA_NOFA to use fa_decode"
13131 .into(),
13132 );
13133 }
13134 let kvl = cache.kv[il].as_mut().unwrap();
13135 let next_len = kvl.len + 1;
13136 let (off, t_kv) = if swa && next_len > win {
13137 (next_len - win, win)
13138 } else {
13139 (0, next_len)
13140 };
13141 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
13142 e.append_kv_quantized(
13143 &k,
13144 &v0,
13145 &mut kvl.k,
13146 &mut kvl.v,
13147 write_row,
13148 kvl.kv_dim_k,
13149 kvl.kv_dim_v,
13150 kvl.k_tok_bytes,
13151 kvl.v_tok_bytes,
13152 crate::Engine::kv_fp8_on(),
13153 )?;
13154 kvl.len = next_len;
13155 let physical = kvl.physical_rows(off, off + t_kv)?;
13156 let k_view = e.view_u8_range(
13157 &kvl.k,
13158 physical.start * kvl.k_tok_bytes,
13159 physical.end * kvl.k_tok_bytes,
13160 );
13161 let v_view = e.view_u8_range(
13162 &kvl.v,
13163 physical.start * kvl.v_tok_bytes,
13164 physical.end * kvl.v_tok_bytes,
13165 );
13166 let mut attn = e.uninit(nh * hd)?;
13167 e.fa_decode_kvmod(
13168 &q,
13169 &k_view,
13170 &v_view,
13171 &mut attn,
13172 hd,
13173 nh,
13174 nkv,
13175 t_kv,
13176 scale,
13177 kvl.k_tok_bytes,
13178 kvl.v_tok_bytes,
13179 crate::Engine::kv_fp8_on(),
13180 )?;
13181
13182 let mut ag = e.uninit(nh * hd)?;
13183 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
13184 Ok(e.matmul(&fa.wo, &ag, 1)?)
13185 }
13186}
13187
13188impl HybridModel {
13197 pub fn is_gemma4_e4b(&self) -> bool {
13198 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
13199 }
13200
13201 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
13205 let g = self.cfg.gemma4.as_ref().unwrap();
13206 let swa = g.swa_pattern[il];
13207 let hd = if swa {
13208 g.key_length_swa
13209 } else {
13210 g.key_length_global
13211 } as usize;
13212 let Mixer::Full(fa) = &self.layers[il].mixer else {
13213 panic!("e4b layer {il} not full-attn")
13214 };
13215 let nh = fa.wq.out_features() / hd;
13216 let nkv = fa.wk.out_features() / hd;
13217 (
13218 hd,
13219 nkv,
13220 nh,
13221 if swa {
13222 g.rope_base_swa
13223 } else {
13224 g.rope_base_global
13225 },
13226 1.0,
13227 swa,
13228 )
13229 }
13230
13231 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
13233 self.layers[il]
13234 .gemma4
13235 .as_ref()
13236 .and_then(|b| b.e4b.as_ref())
13237 .and_then(|e4| e4.kv_share.map(|t| t as usize))
13238 }
13239
13240 fn gemma4_e4b_inp_pl(
13245 &self,
13246 e: &Engine,
13247 tokens: &[u32],
13248 x_scaled: &CudaSlice<f32>,
13249 t: usize,
13250 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13251 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
13252 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
13253 }
13254
13255 fn gemma4_e4b_inp_pl_dev(
13257 &self,
13258 e: &Engine,
13259 tok_d: &CudaSlice<u32>,
13260 x_scaled: &CudaSlice<f32>,
13261 t: usize,
13262 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13263 let aux = self.gemma4_aux.as_ref().unwrap();
13264 let m = aux.e4b.as_ref().unwrap();
13265 let n_embd = self.cfg.n_embd as usize;
13266 let n_layer = self.layers.len();
13267 let width = m.n_epl * n_layer;
13268 let tbl = m.tok_tbl_gpu.get_or_init(|| {
13269 e.upload_u8(&m.tok_embd_bytes)
13270 .expect("e4b per-layer token table upload")
13271 });
13272 let mut a =
13273 e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
13274 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
13275 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
13276 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
13277 let mut pn = e.uninit(t * width)?;
13278 e.rms_norm(
13279 &p,
13280 m.proj_norm.float_data(),
13281 &mut pn,
13282 m.n_epl,
13283 t * n_layer,
13284 self.cfg.rms_eps,
13285 )?;
13286 let mut out = e.uninit(t * width)?;
13287 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
13288 Ok(out)
13289 }
13290
13291 #[allow(clippy::too_many_arguments)]
13296 fn gemma4_e4b_attn(
13297 &self,
13298 e: &Engine,
13299 il: usize,
13300 hq: &CudaSlice<i8>,
13301 hdq: &CudaSlice<f32>,
13302 pos_d: &CudaSlice<i32>,
13303 t: usize,
13304 cache: &mut Cache,
13305 dc_bucket: Option<usize>,
13306 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13307 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
13308 let eps = self.cfg.rms_eps;
13309 let aux = self.gemma4_aux.as_ref().unwrap();
13310 let ones = aux.ones(e);
13311 #[cfg(debug_assertions)]
13312 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
13313 let Mixer::Full(fa) = &self.layers[il].mixer else {
13314 unreachable!()
13315 };
13316 let h0 = e.zeros(0)?;
13320 let h = &h0;
13321
13322 let ff = if swa {
13323 None
13324 } else {
13325 Some(
13326 aux.rope_freqs(e)
13327 .expect("e4b global rope needs rope_freqs.weight"),
13328 )
13329 };
13330 #[cfg(debug_assertions)]
13331 if let Some(ff) = ff {
13332 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13333 }
13334 let share = self.gemma4_e4b_kv_target(il);
13335 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13337 let mut q;
13338 if let Some(_tgt) = share {
13339 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13340 q = e.uninit(t * nh * hd)?;
13341 let mut kdummy = e.uninit(1)?;
13344 let mut vdummy = e.uninit(1)?;
13345 e.rms_norm_qkv_rope(
13346 &q0,
13347 &q0,
13348 &q0,
13349 fa.q_norm.float_data(),
13350 fa.q_norm.float_data(),
13351 ones,
13352 &mut q,
13353 &mut kdummy,
13354 &mut vdummy,
13355 hd,
13356 nh * t,
13357 0,
13358 pos_d,
13359 nh,
13360 1,
13361 base,
13362 1.0,
13363 ff,
13364 eps,
13365 )?;
13366 } else {
13367 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13371 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13372 q = e.uninit(t * nh * hd)?;
13373 let mut k = e.uninit(t * nkv * hd)?;
13374 let mut v = e.uninit(t * nkv * hd)?;
13375 if t == 1 && cat.is_some() {
13376 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13377 e.rms_norm_qkv_rope_cat(
13378 &qkv0,
13379 fa.q_norm.float_data(),
13380 fa.k_norm.float_data(),
13381 ones,
13382 &mut q,
13383 &mut k,
13384 &mut v,
13385 hd,
13386 nh,
13387 nkv,
13388 pos_d,
13389 nh,
13390 nkv,
13391 base,
13392 1.0,
13393 ff,
13394 eps,
13395 )?;
13396 } else {
13397 let (q0, k0, v0) = match if t == 1 {
13398 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13399 } else {
13400 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13403 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13404 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13405 } else {
13406 None
13407 }
13408 } {
13409 Some(triple) => triple,
13410 None => (
13411 e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13412 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13413 e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13414 ), };
13416 e.rms_norm_qkv_rope(
13419 &q0,
13420 &k0,
13421 &v0,
13422 fa.q_norm.float_data(),
13423 fa.k_norm.float_data(),
13424 ones,
13425 &mut q,
13426 &mut k,
13427 &mut v,
13428 hd,
13429 nh * t,
13430 nkv * t,
13431 pos_d,
13432 nh,
13433 nkv,
13434 base,
13435 1.0,
13436 ff,
13437 eps,
13438 )?;
13439 }
13440 let kvl = cache.kv[il].as_mut().unwrap();
13441 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13445 if dc_bucket.is_some() {
13446 debug_assert!(t == 1);
13451 e.append_kv_quantized_row_dc_inc(
13453 &k,
13454 &v,
13455 &mut kvl.k,
13456 &mut kvl.v,
13457 &mut kvl.len_d,
13458 kvl.kv_dim_k,
13459 kvl.kv_dim_v,
13460 kvl.k_tok_bytes,
13461 kvl.v_tok_bytes,
13462 cls,
13463 )?;
13464 } else {
13465 e.append_kv_quantized_rows(
13466 &k,
13467 &v,
13468 &mut kvl.k,
13469 &mut kvl.v,
13470 kvl.len,
13471 t,
13472 kvl.kv_dim_k,
13473 kvl.kv_dim_v,
13474 kvl.k_tok_bytes,
13475 kvl.v_tok_bytes,
13476 cls,
13477 )?;
13478 kvl.len += t;
13479 }
13480 kv_f32 = Some((k, v));
13481 }
13482 let kvl_idx = share.unwrap_or(il);
13485 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13486 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13488 let mut attn = e.uninit(t * nh * hd)?;
13489 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13501 if let Some((kf, vf)) = &kv_f32 {
13502 if hd == 256 && t <= win {
13503 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13504 return Ok(e.matmul(&fa.wo, &attn, t)?);
13505 }
13506 if hd == 256 && swa && t > win {
13507 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13508 return Ok(e.matmul(&fa.wo, &attn, t)?);
13509 }
13510 if hd == 512 && !swa {
13511 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13512 return Ok(e.matmul(&fa.wo, &attn, t)?);
13513 }
13514 } else if share.is_some() {
13515 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13516 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13517 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13518 if hd == 256 && (!swa || t <= win) {
13519 e.fa_prefill_view(
13521 &q,
13522 &k_view,
13523 &v_view,
13524 &mut attn,
13525 hd,
13526 nh,
13527 nkv,
13528 t,
13529 t,
13530 scale,
13531 true,
13532 kvl.k_tok_bytes,
13533 kvl.v_tok_bytes,
13534 g,
13535 )?;
13536 return Ok(e.matmul(&fa.wo, &attn, t)?);
13537 }
13538 let kv_dim = nkv * hd;
13541 let mut kf = e.uninit(t * kv_dim)?;
13542 let mut vf = e.uninit(t * kv_dim)?;
13543 e.fa_dequant_kv_view_f32(
13544 &k_view,
13545 &v_view,
13546 &mut kf,
13547 &mut vf,
13548 kv_dim,
13549 kv_dim,
13550 t,
13551 kvl.k_tok_bytes,
13552 kvl.v_tok_bytes,
13553 g,
13554 )?;
13555 if hd == 512 {
13556 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13557 } else {
13558 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13559 }
13560 return Ok(e.matmul(&fa.wo, &attn, t)?);
13561 }
13562 }
13563 if let Some(bucket) = dc_bucket {
13564 assert!(t == 1);
13569 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13575 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13576 } else {
13577 bucket
13578 };
13579 let k_view = e.view_u8(&kvl.k, kvl.k.len());
13580 let v_view = e.view_u8(&kvl.v, kvl.v.len());
13581 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13582 if crate::Engine::wpf_level() >= 1 {
13590 e.prefetch_weight_l2(&fa.wo)?;
13591 }
13592 if e.uses_q8_1_fast(&fa.wo) {
13595 let mut oq = e.alloc_i8_uninit(nh * hd)?;
13596 let mut od = e.zeros(nh * hd / 32)?;
13597 e.fa_decode_dc_q8(
13598 &q,
13599 &k_view,
13600 &v_view,
13601 &mut attn,
13602 hd,
13603 nh,
13604 nkv,
13605 &kvl.len_d,
13606 bucket,
13607 scale,
13608 kvl.k_tok_bytes,
13609 kvl.v_tok_bytes,
13610 g,
13611 Some((&mut oq, &mut od)),
13612 )?;
13613 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13614 }
13615 e.fa_decode_dc(
13616 &q,
13617 &k_view,
13618 &v_view,
13619 &mut attn,
13620 hd,
13621 nh,
13622 nkv,
13623 &kvl.len_d,
13624 bucket,
13625 scale,
13626 kvl.k_tok_bytes,
13627 kvl.v_tok_bytes,
13628 g,
13629 )?;
13630 return Ok(e.matmul(&fa.wo, &attn, t)?);
13631 }
13632 for i in 0..t {
13633 let avail = base_len + i + 1;
13634 let (off_tok, t_kv) = if swa && avail > win {
13635 (avail - win, win)
13636 } else {
13637 (0, avail)
13638 };
13639 let k_view = e.view_u8_range(
13640 &kvl.k,
13641 off_tok * kvl.k_tok_bytes,
13642 (off_tok + t_kv) * kvl.k_tok_bytes,
13643 );
13644 let v_view = e.view_u8_range(
13645 &kvl.v,
13646 off_tok * kvl.v_tok_bytes,
13647 (off_tok + t_kv) * kvl.v_tok_bytes,
13648 );
13649 let qv = e.view(&q, t * nh * hd);
13650 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13651 let mut q_one = e.uninit(nh * hd)?;
13652 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13653 let mut a_one = e.uninit(nh * hd)?;
13654 e.fa_decode_kvmod(
13658 &q_one,
13659 &k_view,
13660 &v_view,
13661 &mut a_one,
13662 hd,
13663 nh,
13664 nkv,
13665 t_kv,
13666 scale,
13667 kvl.k_tok_bytes,
13668 kvl.v_tok_bytes,
13669 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13670 )?;
13671 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13672 }
13673 Ok(e.matmul(&fa.wo, &attn, t)?)
13674 }
13675
13676 fn gemma4_e4b_trunk(
13681 &self,
13682 e: &Engine,
13683 tokens: &[u32],
13684 pos0: usize,
13685 cache: &mut Cache,
13686 head_last: bool,
13687 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13688 let n_embd = self.cfg.n_embd as usize;
13689 let t = tokens.len();
13690 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13691 let pos_d = e.htod_i32(&pos)?;
13692 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13693 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13694 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13695 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13696 }
13697
13698 fn gemma4_e4b_trunk_core(
13702 &self,
13703 e: &Engine,
13704 x_in: CudaSlice<f32>,
13705 inp_pl: CudaSlice<f32>,
13706 pos_d: &CudaSlice<i32>,
13707 t: usize,
13708 cache: &mut Cache,
13709 dc_bucket: Option<usize>,
13710 cap_logits: bool,
13711 head_last: bool,
13712 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13713 let n_embd = self.cfg.n_embd as usize;
13714 let eps = self.cfg.rms_eps;
13715 let n_layer = self.layers.len();
13716 let mut x = x_in;
13717 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13718 let n_epl = aux_e4b.n_epl;
13719
13720 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13726 for il in 0..n_layer {
13727 let layer = &self.layers[il];
13728 let (hq, hdq) = match h_carry.take() {
13729 Some(p) => p,
13730 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13731 };
13732 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13733 let bits = layer.gemma4.as_ref().unwrap();
13736 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13737 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13748 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13749 e,
13750 layer,
13751 &o,
13752 &x,
13753 t,
13754 Some(layer.post_attn_norm.float_data()),
13755 fuse_exit,
13756 )?;
13757 let mut resid = e.uninit(t * n_embd)?;
13758 let g = if fuse_exit {
13764 let (rq, rd) = e.rms_pre_add_q8_1(
13766 &sn,
13767 bits.post_ffw_norm.float_data(),
13768 &attn_out,
13769 &mut resid,
13770 n_embd,
13771 t,
13772 self.cfg.rms_eps,
13773 )?;
13774 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13775 } else {
13776 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13777 e.matmul(&e4b.inp_gate, &resid, t)?
13778 };
13779 let mut act = e.uninit(t * n_epl)?;
13780 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13781 let ipv = e.view(&inp_pl, n_epl * n_layer);
13782 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13783 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13784 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13785 } else {
13786 let mut inp_this = e.uninit(t * n_epl)?;
13787 e.copy_rows_strided(
13788 &inp_pl,
13789 &mut inp_this,
13790 n_epl,
13791 t,
13792 n_epl * n_layer,
13793 il * n_epl,
13794 )?;
13795 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13796 e.matmul(&e4b.proj, &act, t)?
13797 };
13798 let next_norm = if il + 1 < n_layer {
13801 self.layers[il + 1].attn_norm.float_data()
13802 } else {
13803 self.output_norm.float_data()
13804 };
13805 let mut xn = e.uninit(t * n_embd)?;
13806 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13807 &y,
13808 e4b.post_norm.float_data(),
13809 &resid,
13810 bits.layer_scale,
13811 next_norm,
13812 &mut xn,
13813 n_embd,
13814 t,
13815 eps,
13816 )?;
13817 h_carry = Some(pair);
13818 x = xn;
13819 }
13820 let (oq, odq) = h_carry.take().unwrap();
13824 let h0 = e.zeros(0)?;
13825 let hm = if head_last { 1 } else { t };
13826 let (hq, hd) = if head_last && t > 1 {
13827 let mut q1 = e.uninit_i8(n_embd)?;
13828 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
13829 let nb = n_embd / 32;
13830 let mut d1 = e.uninit(nb)?;
13831 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
13832 (q1, d1)
13833 } else {
13834 (oq, odq)
13835 };
13836 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
13837 if cap_logits {
13841 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
13842 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
13843 }
13844 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
13846 }
13847
13848 pub fn gemma4_e4b_decode_step_t_am_dev(
13855 &self,
13856 e: &Engine,
13857 tok_d: &CudaSlice<u32>,
13858 t: usize,
13859 pos0: usize,
13860 cache: &mut Cache,
13861 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13862 let n_embd = self.cfg.n_embd as usize;
13863 let eps = self.cfg.rms_eps;
13864 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13865 let pos_d = e.htod_i32(&pos)?;
13866 let embd_gpu = self
13867 .embd_gpu
13868 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
13869 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
13870 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
13871 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13872 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
13873 let (ld, xp) =
13874 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
13875 let n_vocab = self.output.out_features();
13878 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
13879 for i in 0..t {
13880 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
13881 }
13882 let mut hn = e.uninit(t * n_embd)?;
13883 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13884 cache.pos += t;
13885 Ok((vam, hn))
13886 }
13887
13888 pub(crate) fn gemma4_e4b_decode_step_t_h(
13891 &self,
13892 e: &Engine,
13893 tokens: &[u32],
13894 pos0: usize,
13895 cache: &mut Cache,
13896 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13897 let n_embd = self.cfg.n_embd as usize;
13898 let eps = self.cfg.rms_eps;
13899 let t = tokens.len();
13900 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
13901 let mut hn = e.uninit(t * n_embd)?;
13902 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13903 cache.pos += t;
13904 Ok((e.dtoh(&ld)?, hn))
13905 }
13906
13907 pub fn gemma4_e4b_decode_step_dcg(
13913 &self,
13914 e: &Engine,
13915 token_d: &mut CudaSlice<u32>,
13916 pos_d: &mut CudaSlice<i32>,
13917 embd_gpu: &CudaSlice<u8>,
13918 embd_qt: i32,
13919 embd_rb: usize,
13920 cache: &mut Cache,
13921 n_vocab: usize,
13922 bucket: usize,
13923 ) -> Result<(), Box<dyn std::error::Error>> {
13924 let n_embd = self.cfg.n_embd as usize;
13925 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13926 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13927 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13928 let (ld, _x) =
13929 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
13930 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
13931 e.inc_seqlen(pos_d)?;
13932 Ok(())
13933 }
13934
13935 #[allow(clippy::too_many_arguments)]
13943 pub fn gemma4_e4b_decode_step_dc(
13944 &self,
13945 e: &Engine,
13946 token_d: &CudaSlice<u32>,
13947 pos_d: &mut CudaSlice<i32>,
13948 embd_gpu: &CudaSlice<u8>,
13949 embd_qt: i32,
13950 embd_rb: usize,
13951 cache: &mut Cache,
13952 n_vocab: usize,
13953 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13954 let n_embd = self.cfg.n_embd as usize;
13955 let eps = self.cfg.rms_eps;
13956 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13957 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13958 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13959 let (ld, _x) =
13960 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
13961 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
13962 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
13963 e.inc_seqlen(pos_d)?;
13964 cache.pos += 1;
13965 let _ = eps;
13966 Ok(tok_out)
13967 }
13968
13969 pub(crate) fn gemma4_e4b_decode_step_h(
13972 &self,
13973 e: &Engine,
13974 token: u32,
13975 cache: &mut Cache,
13976 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13977 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
13978 let logits = e.dtoh(&ld)?;
13979 cache.pos += 1;
13980 Ok((logits, x))
13981 }
13982
13983 pub(crate) fn gemma4_e4b_prime(
13987 &self,
13988 e: &Engine,
13989 tokens: &[u32],
13990 cache: &mut Cache,
13991 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13992 if cache.pos != 0 {
13995 return Err(
13996 "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
13997 call or decode tokenwise"
13998 .into(),
13999 );
14000 }
14001 let n_embd = self.cfg.n_embd as usize;
14002 let t = tokens.len();
14003 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
14004 cache.pos += t;
14005 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
14007 let row = xv.slice((t - 1) * n_embd..t * n_embd);
14008 let mut h_seed = e.uninit(n_embd)?;
14009 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
14010 Ok((last, h_seed, x))
14011 }
14012
14013 pub(crate) fn gemma4_e4b_forward(
14015 &self,
14016 e: &Engine,
14017 tokens: &[u32],
14018 last_only: bool,
14019 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
14020 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
14021 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
14022 Ok(e.dtoh(&ld)?) }
14024}
14025
14026#[cfg(test)]
14027mod prime_chunk_schedule_tests {
14028 use super::{
14029 PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
14030 fixed_prime_chunk_ranges_for_ring,
14031 };
14032
14033 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
14034 ranges.iter().map(|(start, end)| end - start).collect()
14035 }
14036
14037 fn auto_chunk(t: usize) -> usize {
14038 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
14039 }
14040
14041 #[test]
14042 fn fixed_schedule_retains_measured_geometry() {
14043 assert_eq!(
14044 sizes(&fixed_prime_chunk_ranges(461, 128)),
14045 vec![128, 128, 128, 77]
14046 );
14047 assert_eq!(
14048 sizes(&fixed_prime_chunk_ranges(1833, 230)),
14049 vec![230, 230, 230, 230, 230, 230, 230, 223]
14050 );
14051 assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
14052 let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
14053 assert_eq!(capped, vec![4096, 4088, 16]);
14054 assert!(capped.iter().all(|&rows| rows <= 4096));
14055 assert_eq!(
14056 sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
14057 vec![4100],
14058 "flag-off schedule remains byte-for-byte the legacy monolithic tail",
14059 );
14060 }
14061
14062 #[test]
14063 fn dynamic_schedule_matches_registered_shapes() {
14064 let cases = [
14065 (461, vec![64, 141, 132, 124]),
14066 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
14067 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
14068 ];
14069 for (t, expected) in cases {
14070 let chunk = auto_chunk(t);
14071 let fixed = fixed_prime_chunk_ranges(t, chunk);
14072 assert_eq!(
14073 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
14074 expected
14075 );
14076 }
14077 }
14078
14079 #[test]
14080 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
14081 for t in 256..=8192 {
14082 let chunk = auto_chunk(t);
14083 let fixed = fixed_prime_chunk_ranges(t, chunk);
14084 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
14085 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
14086 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
14087 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
14088 for pair in dynamic.windows(2) {
14089 assert_eq!(pair[0].1, pair[1].0, "T={t}");
14090 }
14091 assert!(
14092 dynamic
14093 .iter()
14094 .all(|(start, end)| end - start >= PRIME_MIN_T),
14095 "T={t} sizes={:?}",
14096 sizes(&dynamic)
14097 );
14098 if dynamic.len() >= 3 {
14099 let chunk_sizes = sizes(&dynamic);
14100 assert!(
14101 chunk_sizes[0] < chunk_sizes[1],
14102 "T={t} sizes={chunk_sizes:?}"
14103 );
14104 assert!(
14105 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
14106 "T={t} sizes={chunk_sizes:?}"
14107 );
14108 }
14109 }
14110 }
14111}
14112
14113#[cfg(test)]
14114mod page_prefetch_tests {
14115 use super::{
14116 grouped_worker_prefetch_position, page_prefetch_positions,
14117 page_prefetch_window_from_values, worker_prefetch_positions,
14118 };
14119
14120 #[test]
14121 fn page_prefetch_window_keeps_existing_opt_in_default() {
14122 assert_eq!(page_prefetch_window_from_values(false, None), 0);
14123 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
14124 assert_eq!(page_prefetch_window_from_values(true, None), 1);
14125 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
14126 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
14127 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
14128 }
14129
14130 #[test]
14131 fn rolling_page_prefetch_advises_each_future_expert_once() {
14132 let advised: Vec<_> = (0..7)
14133 .flat_map(|position| page_prefetch_positions(position, 7, 3))
14134 .collect();
14135 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
14136
14137 let one_ahead: Vec<_> = (0..4)
14138 .flat_map(|position| page_prefetch_positions(position, 4, 1))
14139 .collect();
14140 assert_eq!(one_ahead, vec![1, 2, 3]);
14141 assert!(page_prefetch_positions(0, 4, 0).is_empty());
14142 }
14143
14144 #[test]
14145 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
14146 assert_eq!(grouped_worker_prefetch_position(0, None), None);
14147 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
14148 .chain(
14149 (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
14150 )
14151 .collect();
14152 assert_eq!(positions, vec![0, 1, 2, 3]);
14153 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
14154 }
14155
14156 #[test]
14157 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
14158 let queued: Vec<_> = (0..8)
14159 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
14160 .collect();
14161 assert_eq!(queued, (0..8).collect::<Vec<_>>());
14162
14163 let one_at_a_time: Vec<_> = (0..4)
14164 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
14165 .collect();
14166 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
14167 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
14168 }
14169}
14170
14171pub struct G4DcSlots {
14172 x: CudaSlice<f32>,
14173 xn: CudaSlice<f32>,
14174 cur: CudaSlice<f32>,
14175 hq: CudaSlice<i8>,
14176 hd_: CudaSlice<f32>,
14177 q0: CudaSlice<f32>,
14178 k0: CudaSlice<f32>,
14179 v0: CudaSlice<f32>,
14180 q: CudaSlice<f32>,
14181 k: CudaSlice<f32>,
14182 v: CudaSlice<f32>,
14183 attn: CudaSlice<f32>,
14184 o: CudaSlice<f32>,
14185 attn_out: CudaSlice<f32>,
14186 zsh: CudaSlice<f32>,
14187 zq: CudaSlice<i8>,
14188 zd: CudaSlice<f32>,
14189 gate: CudaSlice<f32>,
14190 up: CudaSlice<f32>,
14191 act: CudaSlice<f32>,
14192 actq: CudaSlice<i8>,
14193 actd: CudaSlice<f32>,
14194 f0: CudaSlice<f32>,
14195 sn: CudaSlice<f32>,
14196 hn: CudaSlice<f32>,
14197 logits: CudaSlice<f32>,
14198}