1use cudarc::driver::CudaSlice;
6use memra_gguf::config::ModelConfig;
7use crate::Engine;
8use crate::cache::Cache;
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
123
124pub(crate) struct AttnPre {
126 pub q: cudarc::driver::CudaSlice<f32>,
127 pub k: cudarc::driver::CudaSlice<f32>,
128 pub v: cudarc::driver::CudaSlice<f32>,
129 pub gate: Option<cudarc::driver::CudaSlice<f32>>,
130}
131
132pub(crate) struct GdnPrep {
134 pub hk: usize,
135 pub q_l2: cudarc::driver::CudaSlice<f32>,
136 pub k_l2: cudarc::driver::CudaSlice<f32>,
137 pub v_g: cudarc::driver::CudaSlice<f32>,
138 pub beta: cudarc::driver::CudaSlice<f32>,
139 pub g_log: cudarc::driver::CudaSlice<f32>,
140 pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
141 pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
142}
143
144pub(crate) struct VerifyStreamScratch {
146 pub pos_d: CudaSlice<i32>,
147 pub row_ctrs: Vec<CudaSlice<i32>>,
148}
149use crate::hybrid::{HybridModel, Mixer, FullAttnLayer, LinearAttnLayer, MoeWeights};
150
151struct MoeInputTraceWriter {
152 dir: std::path::PathBuf,
153 index: std::fs::File,
154 payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
155}
156
157static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<
158 std::sync::Mutex<Option<MoeInputTraceWriter>>,
159> = std::sync::OnceLock::new();
160
161fn gdec_enabled() -> bool {
164 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
165 *E.get_or_init(|| std::env::var("MEMRA_MOE_GDEC").map(|v| v != "0").unwrap_or(true))
166}
167
168fn moe_slab_enabled() -> bool {
179 std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
180}
181
182fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
186 std::env::var("MEMRA_MOE_GROUPED")
187 .map(|value| value != "0")
188 .unwrap_or(false)
189}
190
191fn moe_prefetch_enabled() -> bool {
194 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
195 *E.get_or_init(|| std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
196 || crate::spill_pread::worker_enabled())
197}
198
199fn moe_page_prefetch_window() -> usize {
204 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
205 *W.get_or_init(|| page_prefetch_window_from_values(
206 std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
207 std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW").ok().as_deref(),
208 ))
209}
210
211fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
212 if !enabled {
213 return 0;
214 }
215 raw_window
216 .and_then(|value| value.parse().ok())
217 .unwrap_or(1)
218}
219
220fn page_prefetch_positions(
224 position: usize,
225 len: usize,
226 window: usize,
227) -> std::ops::Range<usize> {
228 if window == 0 || position >= len {
229 return len..len;
230 }
231 let (start, count) = if position == 0 {
232 (1, window)
233 } else {
234 (position.saturating_add(window), 1)
235 };
236 let start = start.min(len);
237 start..start.saturating_add(count).min(len)
238}
239
240fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
243 let position = current.map_or(0, |position| position.saturating_add(1));
244 (position < order_len).then_some(position)
245}
246
247fn worker_prefetch_window() -> usize {
252 static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
253 *WINDOW.get_or_init(|| {
254 let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
255 std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
256 .ok()
257 .and_then(|value| value.parse::<usize>().ok())
258 .unwrap_or(automatic.max(1))
259 })
260}
261
262fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
266 if window == 0 || position >= len {
267 return len..len;
268 }
269 let (start, count) = if position == 0 {
270 (0, window)
271 } else {
272 (position.saturating_add(window).saturating_sub(1), 1)
273 };
274 let start = start.min(len);
275 start..start.saturating_add(count).min(len)
276}
277
278fn moe_dev_enabled() -> bool {
283 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
284 *E.get_or_init(|| std::env::var("MEMRA_MOE_DEV").map(|v| v != "0").unwrap_or(true)
285 && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")))
286}
287
288fn moe_q8_enabled() -> bool {
293 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
294 *E.get_or_init(|| std::env::var("MEMRA_MOE_Q8").map(|v| v != "0").unwrap_or(true))
295}
296
297fn expert_dp4a_supported(qt: i32) -> bool {
300 qt == crate::QT_Q4_0 || qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS
301 || qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K
302}
303
304fn q8_expert_supported(qt: i32) -> bool {
305 static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
311 let kq = *KQ.get_or_init(|| {
312 std::env::var("MEMRA_MOE_Q8_KQ").map(|v| v != "0").unwrap_or(true)
313 });
314 let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4").map(|v| v != "0").unwrap_or(true);
321 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || (nvfp4_q8 && qt == crate::QT_NVFP4)
322 || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
323}
324
325fn q8_expert_dec_supported(qt: i32) -> bool {
328 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
329}
330
331fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
337 match qt {
338 crate::QT_Q4_0 => in_f % 32 == 0,
339 crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K
340 | crate::QT_Q6_K => in_f % 256 == 0,
341 _ => false,
342 }
343}
344
345fn moe_prewarm_enabled() -> bool {
348 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
349 *E.get_or_init(|| std::env::var("MEMRA_MOE_PREWARM").map(|v| v != "0").unwrap_or(true))
350}
351
352fn cpu_expert_profile_admit_enabled() -> bool {
356 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
357 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
358}
359
360pub const PRIME_MIN_T: usize = 16;
364const PRIME_PIPE_MICROBATCHES: usize = 8;
365const PRIME_PIPE_MIN_CHUNK: usize = 128;
366const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
367const PRIME_PIPE_LINEAR_WORK: usize = 8;
368
369fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
370 crate::pp::prime_pp_on()
371 && !crate::pp::pp2_streams_off()
372 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
373}
374
375pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
379 if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
380 return value.parse().unwrap_or(4096);
381 }
382 let chunk = 4096usize;
383 if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
384 chunk.min(
385 t.div_ceil(PRIME_PIPE_MICROBATCHES)
386 .max(PRIME_PIPE_MIN_CHUNK),
387 )
388 } else {
389 chunk
390 }
391}
392
393fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
394 if chunk == 0 || t <= chunk {
395 return vec![(0, t)];
396 }
397 let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
398 let mut start = 0usize;
399 while start < t {
400 let mut end = (start + chunk).min(t);
401 if t - end > 0 && t - end < PRIME_MIN_T {
402 end = t;
403 }
404 ranges.push((start, end));
405 start = end;
406 }
407 ranges
408}
409
410fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
411 let prefix = prefix as u128;
412 prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
413}
414
415fn dynamic_prime_chunk_ranges(
416 t: usize,
417 fixed_chunk: usize,
418 fixed: &[(usize, usize)],
419) -> Vec<(usize, usize)> {
420 let n = fixed.len();
421 if n < 3 {
422 return fixed.to_vec();
423 }
424
425 let max_first = t - (n - 1) * PRIME_MIN_T;
426 let first = fixed_chunk
427 .div_ceil(2)
428 .max(PRIME_PIPE_EDGE_MIN_CHUNK)
429 .min(max_first);
430 let mut ranges = Vec::with_capacity(n);
431 ranges.push((0, first));
432
433 let first_work = prime_chunk_work(first, t);
434 let work_span = prime_chunk_work(t, t) - first_work;
435 let denominator = (n - 1) as u128;
436 let mut previous = first;
437 for boundary in 1..n - 1 {
438 let target = first_work * denominator + work_span * (boundary as u128);
439 let remaining = n - 1 - boundary;
440 let mut low = previous + PRIME_MIN_T;
441 let mut high = t - remaining * PRIME_MIN_T;
442 while low < high {
443 let mid = low + (high - low) / 2;
444 if prime_chunk_work(mid, t) * denominator >= target {
445 high = mid;
446 } else {
447 low = mid + 1;
448 }
449 }
450 ranges.push((previous, low));
451 previous = low;
452 }
453 ranges.push((previous, t));
454 ranges
455}
456
457pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
461 let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
462 let chunk = prime_chunk_tokens(t, n_layers);
463 let fixed = fixed_prime_chunk_ranges(t, chunk);
464 let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
465 Ok(value) => value == "dynamic",
466 Err(_) => true,
467 };
468 if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
469 fixed
470 } else {
471 dynamic_prime_chunk_ranges(t, chunk, &fixed)
472 }
473}
474
475impl HybridModel {
476 fn prime_trace_path() -> Option<&'static str> {
481 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
482 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
483 .as_deref()
484 }
485
486 pub fn forward(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
488 if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, false); }
489 if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, false); }
490 let cfg = &self.cfg;
491 let n_embd = cfg.n_embd as usize;
492 let t = tokens.len();
493 let eps = cfg.rms_eps;
494 let pos: Vec<i32> = (0..t as i32).collect();
495 let pos_d = e.htod_i32(&pos)?;
496
497 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
500 let mut h = e.uninit(t * n_embd)?;
502 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
503
504 let mixed = match &layer.mixer {
505 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
506 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
507 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
508 };
509
510 let mut x1 = e.uninit(t * n_embd)?;
512 e.add(&x, &mixed, &mut x1, t * n_embd)?;
513
514 let mut z = e.uninit(t * n_embd)?;
516 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
517 let ffn_out = match &layer.ffn {
518 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
519 let n_ff = ffn_gate.out_features();
520 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
521 let up = g2.pop().unwrap();
522 let gate = g2.pop().unwrap();
523 let mut act = e.uninit(t * n_ff)?;
524 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
529 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
530 e.matmul(ffn_down, &act, t)?
531 }
532 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
533 };
534 let mut x2 = e.uninit(t * n_embd)?;
535 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
536 x = x2;
537 }
538
539 let mut hn = e.uninit(t * n_embd)?;
540 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
541 let logits = e.matmul(&self.output, &hn, t)?;
542 Ok(e.dtoh(&logits)?)
543 }
544
545 pub fn forward_last(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
551 if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, true); }
552 let cfg = &self.cfg;
553 let n_embd = cfg.n_embd as usize;
554 let t = tokens.len();
555 let eps = cfg.rms_eps;
556 let pos: Vec<i32> = (0..t as i32).collect();
557 let pos_d = e.htod_i32(&pos)?;
558
559 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
563 for (il, layer) in self.layers.iter().enumerate() {
564 let mut h = e.uninit(t * n_embd)?;
565 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
566 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} norm ok"); }
567 let mixed = match &layer.mixer {
568 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
569 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
570 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
571 };
572 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} mixer ok"); }
573 let mut x1 = e.uninit(t * n_embd)?;
574 e.add(&x, &mixed, &mut x1, t * n_embd)?;
575 let mut z = e.uninit(t * n_embd)?;
576 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
577 let ffn_out = match &layer.ffn {
578 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
579 let n_ff = ffn_gate.out_features();
580 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
581 let up = g2.pop().unwrap();
582 let gate = g2.pop().unwrap();
583 let mut act = e.uninit(t * n_ff)?;
584 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
586 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
587 e.matmul(ffn_down, &act, t)?
588 }
589 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
590 };
591 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} ffn ok"); }
592 let mut x2 = e.uninit(t * n_embd)?;
593 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
594 x = x2;
595 }
596 let mut hn = e.uninit(t * n_embd)?;
598 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
599 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)?;
602 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
603 let logits = e.matmul(&self.output, &hlast, 1)?; Ok(e.dtoh(&logits)?)
605 }
606
607 pub fn prime_cache(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize)
639 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
640 let n_embd = self.cfg.n_embd as usize;
641 let t = tokens.len();
642 assert!(t >= PRIME_MIN_T, "prime_cache needs T >= {PRIME_MIN_T} (caller gates)");
646 assert!(cache.pos + t <= cache.max_ctx, "prime_cache: prompt exceeds cache max_ctx");
647
648 if self.is_gemma4_e4b() {
660 return self.gemma4_e4b_prime(e, tokens, cache);
661 }
662 if self.cfg.gemma4.is_some() {
663 return self.gemma4_prime(e, tokens, cache);
665 }
666 let ranges = prime_chunk_ranges(t, self.layers.len());
667 let legacy_calllocal =
703 std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
704 let seq_end = if legacy_calllocal {
705 cache.pos + t
706 } else {
707 cache.pos + t + queued_after
708 };
709 if ranges.len() == 1 {
710 return self.prime_chunk(e, tokens, cache, seq_end);
711 }
712 if crate::pp::prime_pipe_on()
717 && crate::pp::prime_pp_on()
718 && !crate::pp::pp2_streams_off()
719 {
720 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
721 if crate::pp::pp_multi_stream_same_device() {
722 return Err(
723 "prime chunk pipeline refused with 2 stage streams on one device — \
724 that concurrent-stream placement remains quarantined by the deferred \
725 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
726 the serial split."
727 .into(),
728 );
729 }
730 return self.prime_cache_pp2_pipelined(
731 e, tokens, cache, seq_end, &ranges, &fence,
732 );
733 }
734 }
735 let mut hiddens = e.uninit(t * n_embd)?;
736 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
737 for &(start, end) in &ranges {
738 let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
739 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
740 last = Some((l, hs));
741 }
742 let (logits, h_seed) = last.unwrap();
743 Ok((logits, h_seed, hiddens))
744 }
745
746 fn prime_cache_pp2_pipelined(
751 &self,
752 e: &Engine,
753 tokens: &[u32],
754 cache: &mut Cache,
755 seq_end: usize,
756 ranges: &[(usize, usize)],
757 fence: &[usize],
758 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
759 debug_assert_eq!(fence.len(), 3);
760 debug_assert!(ranges.len() >= 2);
761 let rt = crate::pp::PpNRt::get(e)?;
762 assert_eq!(rt.n_stages(), 2, "prime pipeline requires exactly two PP stages");
763 let n_embd = self.cfg.n_embd as usize;
764 let t = tokens.len();
765 let initial_base = cache.pos;
766 let caller_stream = e.stream();
767
768 rt.fence_stages_behind(&caller_stream)?;
773 let max_payload = ranges
774 .iter()
775 .map(|(s, e)| (e - s) * n_embd)
776 .max()
777 .unwrap();
778 rt.prepare_overlap_slots(0, max_payload)?;
779
780 let mut hiddens = e.uninit(t * n_embd)?;
781 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
782 let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
783 let (cache0, cache1) = stage_caches.parts();
784 let (first_start, first_end) = ranges[0];
785 let mut slot = self.prime_pp2_stage0_enqueue(
786 e,
787 rt,
788 &tokens[first_start..first_end],
789 cache0,
790 seq_end,
791 fence,
792 initial_base + first_start,
793 true,
794 )?;
795 cache0.pos = initial_base + first_end;
796
797 for (i, &(start, end)) in ranges.iter().enumerate() {
798 let base = initial_base + start;
799 debug_assert_eq!(
800 cache1.pos, base,
801 "stage 1 must drain chunks in original position order"
802 );
803 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
804 let next_base = initial_base + next_start;
805 debug_assert_eq!(
806 cache0.pos, next_base,
807 "stage 0 must issue chunks in original position order"
808 );
809 let cache0_stage = &mut *cache0;
810 std::thread::scope(
815 |scope| -> Result<_, Box<dyn std::error::Error>> {
816 let stage0 = scope.spawn(move || -> Result<usize, String> {
817 let next = self
818 .prime_pp2_stage0_enqueue(
819 e,
820 rt,
821 &tokens[next_start..next_end],
822 cache0_stage,
823 seq_end,
824 fence,
825 next_base,
826 true,
827 )
828 .map_err(|err| err.to_string())?;
829 cache0_stage.pos = initial_base + next_end;
830 Ok(next)
831 });
832 let x = self.prime_pp2_stage1_enqueue(
833 e,
834 rt,
835 slot,
836 end - start,
837 cache1,
838 seq_end,
839 fence,
840 base,
841 true,
842 )?;
843 let out = {
844 rt.bind_stage(1)?;
845 let _st1 = rt.enter(1);
846 let e1 = rt.engine(1, e);
847 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
848 };
849 let next = stage0
850 .join()
851 .map_err(|_| "pipeprime stage-0 host walker panicked")?
852 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
853 Ok((out, Some(next)))
854 },
855 )?
856 } else {
857 let x = self.prime_pp2_stage1_enqueue(
858 e,
859 rt,
860 slot,
861 end - start,
862 cache1,
863 seq_end,
864 fence,
865 base,
866 true,
867 )?;
868 let out = {
869 rt.bind_stage(1)?;
870 let _st1 = rt.enter(1);
871 let e1 = rt.engine(1, e);
872 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
873 };
874 (out, None)
875 };
876
877 rt.publish_to(1, &caller_stream)?;
878 e.copy_into(
879 &mut hiddens,
880 start * n_embd,
881 &out.2,
882 (end - start) * n_embd,
883 )?;
884 last = Some((out.0, out.1));
885 crate::pp::PRIME_SPLIT_CHUNKS
886 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
887
888 if let Some(next) = next_slot {
889 rt.fence_stages_behind(&caller_stream)?;
894 slot = next;
895 }
896 }
897
898 debug_assert_eq!(cache0.pos, initial_base + t);
899 debug_assert_eq!(cache1.pos, initial_base + t);
900 let (logits, h_seed) = last.unwrap();
901 Ok((logits, h_seed, hiddens))
902 }
903
904 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
911 if Engine::gdn_db_on()
912 && Engine::gdn_chunked_enabled() && t >= 16
913 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
914 && num_k * 2 == num_v
915 {
916 num_k
917 } else {
918 num_v
919 }
920 }
921
922 fn f16out_on(e: &Engine, t: usize) -> bool {
927 crate::f16_ffi::pp_f16_enabled() && t >= 16 && !e.verify_exact_on()
928 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
929 }
930
931 pub fn prime_slabs_get(
939 &self,
940 e: &Engine,
941 t: usize,
942 n_embd: usize,
943 n_ff_max: usize,
944 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
945 let mut slabs = self.prime_slabs.lock().unwrap();
946 let dev = e.ctx().ordinal();
947 let need_new = match slabs.get(&dev) {
948 None => true,
949 Some(sl) => sl.lock().unwrap().t_cap < t,
950 };
951 if need_new {
952 slabs.insert(dev, std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
953 t_cap: t,
954 h: e.uninit(t * n_embd)?,
955 x1: e.uninit(t * n_embd)?,
956 z: e.uninit(t * n_embd)?,
957 act: e.uninit(t * n_ff_max)?,
958 xa: e.uninit(t * n_embd)?,
959 xb: e.uninit(t * n_embd)?,
960 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
961 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
962 gate: e.uninit(t * n_ff_max)?,
963 up: e.uninit(t * n_ff_max)?,
964 ffn_out: e.uninit(t * n_embd)?,
965 seg_glue: Vec::new(),
966 mixed: e.uninit(t * n_embd)?,
967 seg_mid: Vec::new(),
968 seg_t: 0,
969 })));
970 }
971 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
972 }
973
974 fn prime_chunk(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize)
978 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
979 if self.cfg.gemma4.is_none()
988 && !crate::pp::pp2_streams_off()
989 && crate::pp::prime_pp_on()
990 {
991 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
992 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
993 }
994 }
995 let t = tokens.len();
996 let base = cache.pos;
997 debug_assert!(seq_end >= base + t, "prime_chunk: seq_end must cover this chunk");
998 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
999 let pos_d = e.htod_i32(&pos)?;
1000
1001 let x_embed = self.embed(e, tokens)?; let x = self.prime_layers(
1003 e, x_embed, 0, self.layers.len(), &pos_d, t, base, cache, seq_end,
1004 )?;
1005 self.prime_chunk_epilogue(e, x, t, cache)
1006 }
1007
1008 #[allow(clippy::too_many_arguments)]
1024 fn prime_layers(&self, e: &Engine, x_in: CudaSlice<f32>, lo: usize, hi: usize,
1025 pos_d: &CudaSlice<i32>, t: usize, base: usize, cache: &mut Cache,
1026 seq_end: usize)
1027 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1028 let cfg = &self.cfg;
1029 let n_embd = cfg.n_embd as usize;
1030 let eps = cfg.rms_eps;
1031 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1035 let n_ff_max = self.layers.iter().map(|l| match &l.ffn {
1041 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1042 _ => n_embd,
1043 }).max().unwrap_or(n_embd).max(n_embd);
1044 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1045 let slab = if use_slabs {
1046 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1047 } else {
1048 None
1049 };
1050 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1051 let mut x_own; type SlabRefs<'a> = (&'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<u8>, &'a mut CudaSlice<u8>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>);
1053 let (mut x_cur, mut x_nxt, sl): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, Option<SlabRefs>);
1054 let mut seg: Option<(&mut Vec<Option<cudarc::driver::CudaGraph>>, &mut Vec<Option<cudarc::driver::CudaGraph>>, &mut CudaSlice<f32>, &mut usize)> = None;
1055 let mut x_own2;
1056 match slab_guard.as_mut() {
1057 Some(g) => {
1058 let slabs = &mut **g;
1059 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1060 let PrimeSlabs { xa, xb, h, x1, z, act, h16, z16, gate, up, ffn_out, seg_glue, mixed, seg_mid, seg_t, .. } = slabs;
1061 x_cur = xa;
1062 x_nxt = xb;
1063 seg = Some((seg_glue, seg_mid, mixed, seg_t));
1064 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1065 }
1066 None => {
1067 x_own = x_in;
1068 x_own2 = e.uninit(t * n_embd)?;
1069 x_cur = &mut x_own;
1070 x_nxt = &mut x_own2;
1071 sl = None;
1072 }
1073 }
1074 let mut alloc_h; let mut alloc_x1; let mut alloc_z; let mut alloc_act;
1075 let mut alloc_h16; let mut alloc_z16;
1076 let mut alloc_gate; let mut alloc_up; let mut alloc_fo;
1077 let (h, x1, z, act): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
1078 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1079 let (sl_gate, sl_up, sl_fo): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
1080 match sl {
1081 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1082 h = a; x1 = b; z = c; act = d; h16 = e16; z16 = f16b;
1083 sl_gate = g; sl_up = u; sl_fo = fo;
1084 }
1085 None => {
1086 alloc_h = e.uninit(t * n_embd)?;
1087 alloc_x1 = e.uninit(t * n_embd)?;
1088 alloc_z = e.uninit(t * n_embd)?;
1089 alloc_act = e.uninit(t * n_ff_max)?;
1090 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1091 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1092 alloc_gate = e.uninit(t * n_ff_max)?;
1093 alloc_up = e.uninit(t * n_ff_max)?;
1094 alloc_fo = e.uninit(t * n_embd)?;
1095 h = &mut alloc_h; x1 = &mut alloc_x1; z = &mut alloc_z; act = &mut alloc_act;
1096 h16 = &mut alloc_h16; z16 = &mut alloc_z16;
1097 sl_gate = &mut alloc_gate; sl_up = &mut alloc_up; sl_fo = &mut alloc_fo;
1098 }
1099 }
1100 let n_layers = self.layers.len();
1105 let use_seg = f16fuse && seg.is_some() && self.cfg.step35.is_none()
1115 && lo == 0 && hi == n_layers
1116 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1117 if let Some((sg, sm, _, st)) = seg.as_mut() {
1118 if **st != t {
1119 sg.clear();
1120 sg.extend((0..n_layers).map(|_| None));
1121 sm.clear();
1122 sm.extend((0..n_layers).map(|_| None));
1123 **st = t;
1124 }
1125 }
1126 {
1127 let layer_lo = &self.layers[lo];
1128 if f16fuse {
1129 e.rms_norm_f16out(x_cur, layer_lo.attn_norm.float_data(), h, h16, n_embd, t, eps)?;
1130 } else {
1131 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1132 }
1133 }
1134 for il in lo..hi {
1135 let layer = &self.layers[il];
1136 let hx16 = if f16fuse { Some(&*h16) } else { None };
1137 if use_seg {
1138 let (pre, pre16, w_out) = match &layer.mixer {
1141 Mixer::Full(fa) => {
1142 let g3 = match hx16 {
1143 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1144 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1145 };
1146 let (pre, pre16) = self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1147 (pre, pre16, &fa.wo)
1148 }
1149 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1150 Mixer::Linear(la) => {
1151 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1152 let g4 = match hx16 {
1153 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1154 None => e.matmul_group(&ws, h, t)?,
1155 };
1156 let (pre, pre16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1157 (pre, pre16, &la.ssm_out)
1158 }
1159 };
1160 {
1161 let (_, sm, mslab, _) = seg.as_mut().unwrap();
1162 let pre_n = pre.len() / t;
1163 let xh_pre = match pre16 {
1164 Some(x) => x,
1165 None => e.f16_act(&pre, t * pre_n, pre_n)?,
1166 };
1167 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1168 let y = e.matmul(w_out, &pre, t)?;
1169 e.copy_into(mslab, 0, &y, t * n_embd)?;
1170 }
1171 if sm[il].is_none() {
1172 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1173 let w_post = layer.post_attn_norm.float_data();
1174 e.stream().synchronize()?;
1175 e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1176 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1177 e.add(x_cur, mslab, x1, t * n_embd)?;
1178 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1179 Ok(())
1180 })();
1181 let g = e.stream().end_capture(
1182 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1183 r?;
1184 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1185 }
1186 sm[il].as_ref().unwrap().launch()?;
1187 }
1188 } else {
1189 let mixed = match &layer.mixer {
1190 Mixer::Full(fa) => self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il,
1191 seq_end)?,
1192 Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1193 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1194 };
1195 if f16fuse {
1196 e.add_rms_norm_f16out(x_cur, &mixed, layer.post_attn_norm.float_data(),
1199 x1, z, z16, n_embd, t, eps)?;
1200 } else {
1201 e.add(x_cur, &mixed, x1, t * n_embd)?;
1202 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1203 }
1204 }
1205 let zx16 = if f16fuse { Some(&*z16) } else { None };
1206 match &layer.ffn {
1207 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1208 let n_ff = ffn_gate.out_features();
1209 let mut into_ok = false;
1212 if let Some(xh) = zx16 {
1213 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1214 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1215 }
1216 if !into_ok {
1217 let mut g2 = match zx16 {
1218 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1219 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1220 };
1221 let up_y = g2.pop().unwrap();
1222 let gate_y = g2.pop().unwrap();
1223 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1224 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1225 }
1226 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1231 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none()
1232 && d_lim.is_none() {
1233 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1234 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1235 Some(a16)
1236 } else {
1237 Self::ffn_act_lim(e, &self.cfg, sl_gate, sl_up, 1.0, 1.0, d_lim,
1238 act, t * n_ff)?;
1239 None
1240 };
1241 let xh_act = match act16 {
1243 Some(x) => x,
1244 None => e.f16_act(act, t * n_ff, n_ff)?,
1245 };
1246 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1247 let y = e.matmul(ffn_down, &*act, t)?;
1248 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1249 }
1250 }
1251 crate::hybrid::Ffn::Moe(m) => {
1252 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1253 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1254 }
1255 }
1256 if use_seg && il + 1 < hi {
1257 let w_next = self.layers[il + 1].attn_norm.float_data();
1259 let (sg, _, _, _) = seg.as_mut().unwrap();
1260 if sg[il].is_none() {
1261 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1262 e.stream().synchronize()?;
1263 e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1264 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1265 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1266 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1267 Ok(())
1268 })();
1269 let g = e.stream().end_capture(
1270 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1271 r?;
1272 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1273 }
1274 sg[il].as_ref().unwrap().launch()?;
1275 } else {
1276 if il + 1 < hi {
1277 let w_next = self.layers[il + 1].attn_norm.float_data();
1278 if f16fuse {
1279 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1280 } else {
1281 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1282 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1283 }
1284 } else {
1285 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1286 }
1287 }
1288 if let Some(path) = Self::prime_trace_path() {
1294 let row = (base + t - 1) as usize;
1295 let host = e.dtoh(x_nxt)?;
1296 let last = &host[(t - 1) * n_embd..t * n_embd];
1297 use std::io::Write as _;
1298 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
1299 let mut h64: u64 = 0xcbf29ce484222325;
1300 for v in last {
1301 h64 ^= v.to_bits() as u64;
1302 h64 = h64.wrapping_mul(0x100000001b3);
1303 }
1304 writeln!(f, "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1305 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1306 last[0], last[1], last[2])?;
1307 }
1308 std::mem::swap(&mut x_cur, &mut x_nxt);
1309 }
1310 let mut x = e.uninit(t * n_embd)?;
1312 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1313 drop(slab_guard);
1314 Ok(x)
1315 }
1316
1317 fn prime_chunk_epilogue(&self, e: &Engine, x: CudaSlice<f32>, t: usize, cache: &mut Cache)
1322 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1323 let n_embd = self.cfg.n_embd as usize;
1324 let eps = self.cfg.rms_eps;
1325 let mut h_seed = e.uninit(n_embd)?;
1329 if !crate::spec::spec_hpost() {
1330 e.copy_view_into(&mut h_seed, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
1331 }
1332 let mut hn = e.uninit(t * n_embd)?;
1334 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1335 if crate::spec::spec_hpost() {
1336 e.copy_view_into(&mut h_seed, 0, &hn.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
1337 }
1338 let last = e.view(&hn, t * n_embd);
1339 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1340 let mut hlast = e.uninit(n_embd)?;
1341 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1342 let logits = e.matmul(&self.output, &hlast, 1)?;
1343 cache.pos += t;
1344 Ok((e.dtoh(&logits)?, h_seed, if crate::spec::spec_hpost() { hn } else { x }))
1347 }
1348
1349 fn prime_chunk_ppn(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize,
1373 fence: &[usize])
1374 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1375 let rt = crate::pp::PpNRt::get(e)?;
1376 let n_st = fence.len() - 1;
1377 assert_eq!(
1378 rt.n_stages(), n_st,
1379 "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
1380 );
1381 let n_embd = self.cfg.n_embd as usize;
1382 let t = tokens.len();
1383 let base = cache.pos;
1384 debug_assert!(seq_end >= base + t, "prime_chunk_ppn: seq_end must cover this chunk");
1385 let payload = t * n_embd;
1386 let caller_stream = e.stream();
1390 rt.fence_stages_behind(&caller_stream)?;
1391
1392 if n_st == 2 {
1393 let slot = self.prime_pp2_stage0_enqueue(
1394 e, rt, tokens, cache, seq_end, fence, base, false,
1395 )?;
1396 let x = self.prime_pp2_stage1_enqueue(
1397 e, rt, slot, t, cache, seq_end, fence, base, false,
1398 )?;
1399 let out = {
1400 rt.bind_stage(1)?;
1401 let _st1 = rt.enter(1);
1402 let e1 = rt.engine(1, e);
1403 self.prime_chunk_epilogue(e1, x, t, cache)?
1404 };
1405 rt.publish_to(1, &caller_stream)?;
1406 crate::pp::PRIME_SPLIT_CHUNKS
1407 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1408 return Ok(out);
1409 }
1410
1411 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1412
1413 let mut slot = {
1415 let _st0 = rt.enter(0);
1416 let e0 = rt.engine(0, e);
1417 let pos_d = e0.htod_i32(&pos)?;
1418 let x = self.embed(e0, tokens)?;
1419 let x = self.prime_layers(
1420 e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end,
1421 )?;
1422 rt.tx(0, &x, payload)?
1423 };
1425
1426 for s in 1..n_st - 1 {
1428 let _st = rt.enter(s);
1429 let es = rt.engine(s, e);
1430 let pos_d = es.htod_i32(&pos)?;
1431 let x = rt.rx(s - 1, slot, payload)?;
1432 let x = self.prime_layers(
1433 es, x, fence[s], fence[s + 1], &pos_d, t, base, cache, seq_end,
1434 )?;
1435 slot = rt.tx(s, &x, payload)?;
1436 }
1437
1438 let _stl = rt.enter(n_st - 1);
1440 let el = rt.engine(n_st - 1, e);
1441 let pos_d = el.htod_i32(&pos)?;
1442 let x = rt.rx(n_st - 2, slot, payload)?;
1443 let x = self.prime_layers(
1444 el, x, fence[n_st - 1], fence[n_st], &pos_d, t, base, cache, seq_end,
1445 )?;
1446 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1447 rt.publish_to(n_st - 1, &caller_stream)?;
1453 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1454 Ok(out)
1455 }
1456
1457 fn prime_pp2_stage0_enqueue(
1458 &self,
1459 e: &Engine,
1460 rt: &crate::pp::PpNRt,
1461 tokens: &[u32],
1462 cache: &mut Cache,
1463 seq_end: usize,
1464 fence: &[usize],
1465 base: usize,
1466 pipelined: bool,
1467 ) -> Result<usize, Box<dyn std::error::Error>> {
1468 let t = tokens.len();
1469 let n_embd = self.cfg.n_embd as usize;
1470 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1471 rt.bind_stage(0)?;
1472 let _st0 = rt.enter(0);
1473 let e0 = rt.engine(0, e);
1474 let pos_d = e0.htod_i32(&pos)?;
1475 let x = self.embed(e0, tokens)?;
1476 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1477 let x = self.prime_layers(
1478 e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end,
1479 )?;
1480 if pipelined {
1481 rt.tx_pipelined(0, &x, t * n_embd)
1482 } else {
1483 rt.tx(0, &x, t * n_embd)
1484 }
1485 }
1486
1487 fn prime_pp2_stage1_enqueue(
1488 &self,
1489 e: &Engine,
1490 rt: &crate::pp::PpNRt,
1491 slot: usize,
1492 t: usize,
1493 cache: &mut Cache,
1494 seq_end: usize,
1495 fence: &[usize],
1496 base: usize,
1497 pipelined: bool,
1498 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1499 let n_embd = self.cfg.n_embd as usize;
1500 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1501 rt.bind_stage(1)?;
1502 let _st1 = rt.enter(1);
1503 let e1 = rt.engine(1, e);
1504 let pos_d = e1.htod_i32(&pos)?;
1505 let x = rt.rx(0, slot, t * n_embd)?;
1506 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1507 self.prime_layers(
1508 e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end,
1509 )
1510 }
1511
1512 pub fn prime_chunk_captured(&self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>,
1528 t: usize, cache: &mut Cache,
1529 len_d: &CudaSlice<i32>,
1530 logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>)
1531 -> Result<(), Box<dyn std::error::Error>> {
1532 let cfg = &self.cfg;
1533 let n_embd = cfg.n_embd as usize;
1534 let eps = cfg.rms_eps;
1535 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1536 let mut x = e.uninit(t * n_embd)?;
1537 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1538 for (il, layer) in self.layers.iter().enumerate() {
1539 let mut h = e.uninit(t * n_embd)?;
1540 let mut hx16: Option<CudaSlice<u8>> = None;
1541 if f16fuse {
1542 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1543 e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut b16, n_embd, t, eps)?;
1544 hx16 = Some(b16);
1545 } else {
1546 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1547 }
1548 let mixed = match &layer.mixer {
1549 Mixer::Full(fa) => self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache,
1553 il, t)?,
1554 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1555 Mixer::Linear(la) => {
1556 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1557 let g4 = match hx16.as_ref() {
1558 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1559 None => e.matmul_group(&ws, &h, t)?,
1560 };
1561 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1562 }
1563 };
1564 let mut x1 = e.uninit(t * n_embd)?;
1565 e.add(&x, &mixed, &mut x1, t * n_embd)?;
1566 let mut z = e.uninit(t * n_embd)?;
1567 let mut zx16: Option<CudaSlice<u8>> = None;
1568 if f16fuse {
1569 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1570 e.rms_norm_f16out(&x1, layer.post_attn_norm.float_data(), &mut z, &mut b16, n_embd, t, eps)?;
1571 zx16 = Some(b16);
1572 } else {
1573 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
1574 }
1575 let ffn_out = match &layer.ffn {
1576 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1577 let n_ff = ffn_gate.out_features();
1578 let mut g2 = match &zx16 {
1579 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1580 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1581 };
1582 let up = g2.pop().unwrap();
1583 let gate = g2.pop().unwrap();
1584 let mut act = e.uninit(t * n_ff)?;
1585 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
1587 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
1588 e.matmul(ffn_down, &act, t)?
1589 }
1590 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
1591 };
1592 let mut x2 = e.uninit(t * n_embd)?;
1593 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1594 x = x2;
1595 }
1596 if !crate::spec::spec_hpost() {
1598 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1599 }
1600 let mut hn = e.uninit(t * n_embd)?;
1601 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1602 if crate::spec::spec_hpost() {
1603 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1604 }
1605 let mut hlast = e.uninit(n_embd)?;
1606 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1607 let logits = e.matmul(&self.output, &hlast, 1)?;
1608 let nv = logits.len();
1609 e.copy_into(logits_out, 0, &logits, nv)?;
1610 Ok(())
1611 }
1612
1613 fn step35_prime_batch_on() -> bool {
1614 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
1615 }
1616
1617 #[allow(clippy::too_many_arguments)]
1620 fn step35_prime_batch_layers(
1621 &self,
1622 e: &Engine,
1623 mut x: CudaSlice<f32>,
1624 lo: usize,
1625 hi: usize,
1626 ts: &[usize],
1627 offs: &[usize],
1628 pos_ds: &[CudaSlice<i32>],
1629 caches: &mut [&mut Cache],
1630 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1631 let cfg = &self.cfg;
1632 let n_embd = cfg.n_embd as usize;
1633 let eps = cfg.rms_eps;
1634 let b = ts.len();
1635 let total: usize = ts.iter().sum();
1636 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
1637
1638 let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
1639 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1640 let mut out = Vec::with_capacity(b);
1641 for s in 0..b {
1642 let mut ys = e.uninit(ts[s] * dim)?;
1643 e.copy_view_into(
1644 &mut ys,
1645 0,
1646 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
1647 ts[s] * dim,
1648 )?;
1649 out.push(ys);
1650 }
1651 Ok(out)
1652 };
1653
1654 for il in lo..hi {
1655 let layer = &self.layers[il];
1656 let Mixer::Full(fa) = &layer.mixer else {
1657 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
1658 };
1659
1660 let mut h = e.uninit(total * n_embd)?;
1661 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1662 if f16fuse {
1663 e.rms_norm_f16out(
1664 &x,
1665 layer.attn_norm.float_data(),
1666 &mut h,
1667 &mut hx16,
1668 n_embd,
1669 total,
1670 eps,
1671 )?;
1672 } else {
1673 e.rms_norm(
1674 &x,
1675 layer.attn_norm.float_data(),
1676 &mut h,
1677 n_embd,
1678 total,
1679 eps,
1680 )?;
1681 }
1682
1683 let gate_w = fa
1687 .attn_gate
1688 .as_ref()
1689 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
1690 let mut g4 = if f16fuse {
1691 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
1692 } else {
1693 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
1694 };
1695 let gate = g4.pop().unwrap();
1696 let mut parts: Vec<Vec<CudaSlice<f32>>> =
1697 (0..b).map(|_| Vec::with_capacity(3)).collect();
1698 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
1699 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
1700 parts[s].push(ys);
1701 }
1702 }
1703 let gates = split(e, &gate, gate_w.out_features())?;
1704 let geometry = self.step35_geom(il);
1705 let hd = geometry.head_dim_k as usize;
1706 let nh = geometry.n_head as usize;
1707 let mut ag_cat = e.uninit(total * nh * hd)?;
1708 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
1709 let ag = self.step35_attn_pre_wo(
1710 e,
1711 fa,
1712 g3s,
1713 None,
1714 Some(&gate),
1715 &pos_ds[s],
1716 ts[s],
1717 Some(&mut *caches[s]),
1718 il,
1719 ts[s],
1720 )?;
1721 e.copy_into(
1722 &mut ag_cat,
1723 offs[s] * nh * hd,
1724 &ag,
1725 ts[s] * nh * hd,
1726 )?;
1727 }
1728 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
1729
1730 let mut x1 = e.uninit(total * n_embd)?;
1731 let mut z = e.uninit(total * n_embd)?;
1732 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1733 if f16fuse {
1734 e.add_rms_norm_f16out(
1735 &x,
1736 &mixed,
1737 layer.post_attn_norm.float_data(),
1738 &mut x1,
1739 &mut z,
1740 &mut zx16,
1741 n_embd,
1742 total,
1743 eps,
1744 )?;
1745 } else {
1746 e.add(&x, &mixed, &mut x1, total * n_embd)?;
1747 e.rms_norm(
1748 &x1,
1749 layer.post_attn_norm.float_data(),
1750 &mut z,
1751 n_embd,
1752 total,
1753 eps,
1754 )?;
1755 }
1756
1757 let ffn_out = match &layer.ffn {
1758 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1759 let n_ff = ffn_gate.out_features();
1760 let mut g2 = if f16fuse {
1761 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
1762 } else {
1763 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
1764 };
1765 let up = g2.pop().unwrap();
1766 let gate = g2.pop().unwrap();
1767 let mut act = e.uninit(total * n_ff)?;
1768 let d_lim = cfg.clamp_shexp_at(il as u32);
1769 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
1770 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
1771 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
1772 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
1773 Some(y) => y,
1774 None => e.matmul(ffn_down, &act, total)?,
1775 }
1776 } else {
1777 Self::ffn_act_lim(
1778 e,
1779 cfg,
1780 &gate,
1781 &up,
1782 1.0,
1783 1.0,
1784 d_lim,
1785 &mut act,
1786 total * n_ff,
1787 )?;
1788 e.matmul(ffn_down, &act, total)?
1789 }
1790 }
1791 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
1792 };
1793 let mut x2 = e.uninit(total * n_embd)?;
1794 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
1795 x = x2;
1796 }
1797 Ok(x)
1798 }
1799
1800 fn step35_prime_batch_epilogue(
1801 &self,
1802 e: &Engine,
1803 x: CudaSlice<f32>,
1804 ts: &[usize],
1805 offs: &[usize],
1806 caches: &mut [&mut Cache],
1807 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1808 let n_embd = self.cfg.n_embd as usize;
1809 let total: usize = ts.iter().sum();
1810 let mut hn = e.uninit(total * n_embd)?;
1811 e.rms_norm(
1812 &x,
1813 self.output_norm.float_data(),
1814 &mut hn,
1815 n_embd,
1816 total,
1817 self.cfg.rms_eps,
1818 )?;
1819
1820 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
1821 let mut out = Vec::with_capacity(ts.len());
1822 for s in 0..ts.len() {
1823 let mut hidden = e.uninit(ts[s] * n_embd)?;
1824 e.copy_view_into(
1825 &mut hidden,
1826 0,
1827 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
1828 ts[s] * n_embd,
1829 )?;
1830 let last0 = (offs[s] + ts[s] - 1) * n_embd;
1831 let mut h_seed = e.uninit(n_embd)?;
1832 e.copy_view_into(
1833 &mut h_seed,
1834 0,
1835 &hidden_src.slice(last0..last0 + n_embd),
1836 n_embd,
1837 )?;
1838 let mut hlast = e.uninit(n_embd)?;
1840 e.copy_view_into(
1841 &mut hlast,
1842 0,
1843 &hn.slice(last0..last0 + n_embd),
1844 n_embd,
1845 )?;
1846 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
1847 caches[s].pos += ts[s];
1848 out.push((logits, h_seed, hidden));
1849 }
1850 Ok(out)
1851 }
1852
1853 fn step35_prime_cache_batch(
1854 &self,
1855 e: &Engine,
1856 prompts: &[&[u32]],
1857 caches: &mut [&mut Cache],
1858 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1859 if !Self::step35_prime_batch_on() {
1860 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
1861 }
1862 if caches.iter().any(|c| c.pos != 0) {
1863 return Err(
1864 "step35 batched prime currently supports complete fresh prompts only; \
1865 continuation/tick chunks require per-request queued_after"
1866 .into(),
1867 );
1868 }
1869
1870 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
1871 for &t in &ts {
1872 assert!(t >= PRIME_MIN_T, "step35 batched prime needs T >= {PRIME_MIN_T}");
1873 }
1874 for (s, c) in caches.iter().enumerate() {
1875 assert!(ts[s] <= c.max_ctx, "step35 batched prime exceeds cache max_ctx");
1876 }
1877 let offs: Vec<usize> = ts
1878 .iter()
1879 .scan(0usize, |a, &t| {
1880 let o = *a;
1881 *a += t;
1882 Some(o)
1883 })
1884 .collect();
1885 let total: usize = ts.iter().sum();
1886 let payload = total * self.cfg.n_embd as usize;
1887 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
1888 let positions: Vec<Vec<i32>> = ts
1889 .iter()
1890 .map(|&t| (0..t as i32).collect())
1891 .collect();
1892 let upload_positions = |e: &Engine|
1893 -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
1894 positions
1895 .iter()
1896 .map(|p| e.htod_i32(p))
1897 .collect::<Result<_, _>>()
1898 };
1899
1900 static ONCE: std::sync::Once = std::sync::Once::new();
1901 ONCE.call_once(|| {
1902 eprintln!(
1903 "[step35-prime-batch] first concat prime: B={} tokens={total}",
1904 prompts.len()
1905 );
1906 });
1907
1908 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1909 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1910 let rt = crate::pp::PpNRt::get(e)?;
1911 let n_st = fence.len() - 1;
1912 assert_eq!(rt.n_stages(), n_st, "step35 prime batch stage count mismatch");
1913 let caller_stream = e.stream();
1914 rt.fence_stages_behind(&caller_stream)?;
1915
1916 let mut slot = {
1917 let _st0 = rt.enter(0);
1918 let e0 = rt.engine(0, e);
1919 let pos_ds = upload_positions(e0)?;
1920 let x = self.embed(e0, &cat_tokens)?;
1921 let x = self.step35_prime_batch_layers(
1922 e0,
1923 x,
1924 fence[0],
1925 fence[1],
1926 &ts,
1927 &offs,
1928 &pos_ds,
1929 caches,
1930 )?;
1931 rt.tx(0, &x, payload)?
1932 };
1933 for s in 1..n_st - 1 {
1934 let _st = rt.enter(s);
1935 let es = rt.engine(s, e);
1936 let pos_ds = upload_positions(es)?;
1937 let x = rt.rx(s - 1, slot, payload)?;
1938 let x = self.step35_prime_batch_layers(
1939 es,
1940 x,
1941 fence[s],
1942 fence[s + 1],
1943 &ts,
1944 &offs,
1945 &pos_ds,
1946 caches,
1947 )?;
1948 slot = rt.tx(s, &x, payload)?;
1949 }
1950
1951 let _stl = rt.enter(n_st - 1);
1952 let el = rt.engine(n_st - 1, e);
1953 let pos_ds = upload_positions(el)?;
1954 let x = rt.rx(n_st - 2, slot, payload)?;
1955 let x = self.step35_prime_batch_layers(
1956 el,
1957 x,
1958 fence[n_st - 1],
1959 fence[n_st],
1960 &ts,
1961 &offs,
1962 &pos_ds,
1963 caches,
1964 )?;
1965 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
1966 rt.publish_to(n_st - 1, &caller_stream)?;
1967 crate::pp::STEP35_PRIME_BATCH_SPLITS.fetch_add(
1968 1,
1969 std::sync::atomic::Ordering::Relaxed,
1970 );
1971 out
1972 } else {
1973 let pos_ds = upload_positions(e)?;
1974 let x = self.embed(e, &cat_tokens)?;
1975 let x = self.step35_prime_batch_layers(
1976 e,
1977 x,
1978 0,
1979 self.layers.len(),
1980 &ts,
1981 &offs,
1982 &pos_ds,
1983 caches,
1984 )?;
1985 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
1986 }
1987 } else {
1988 let pos_ds = upload_positions(e)?;
1989 let x = self.embed(e, &cat_tokens)?;
1990 let x = self.step35_prime_batch_layers(
1991 e,
1992 x,
1993 0,
1994 self.layers.len(),
1995 &ts,
1996 &offs,
1997 &pos_ds,
1998 caches,
1999 )?;
2000 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2001 };
2002 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2003 Ok(out)
2004 }
2005
2006 pub fn prime_cache_batch(&self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache])
2023 -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2024 let cfg = &self.cfg;
2025 let n_embd = cfg.n_embd as usize;
2026 let eps = cfg.rms_eps;
2027 let b = prompts.len();
2028 assert!(b >= 1 && b == caches.len());
2029 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2030 let carried = pos0s.iter().any(|&p| p > 0);
2031 if cfg.gemma4.is_some() {
2037 return Err("prime_cache_batch: gemma4 has no batched prime core (per-layer \
2038 swa/global geometry, softcapped head) — use gemma4_prime per sequence".into());
2039 }
2040 if cfg.step35.is_some() {
2043 return self.step35_prime_cache_batch(e, prompts, caches);
2044 }
2045 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2046 for &t in &ts { assert!(t >= PRIME_MIN_T, "prime_cache_batch needs T >= {PRIME_MIN_T}"); }
2047 for (s, c) in caches.iter().enumerate() {
2048 assert!(c.pos + ts[s] <= c.max_ctx, "prime_cache_batch: prompt exceeds cache max_ctx");
2049 }
2050 let total: usize = ts.iter().sum();
2051 let offs: Vec<usize> = ts.iter().scan(0usize, |a, &t| { let o = *a; *a += t; Some(o) }).collect();
2052 let pos_ds: Vec<CudaSlice<i32>> = ts.iter().zip(&pos0s)
2054 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2055 .collect::<Result<_, _>>()?;
2056 let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
2058 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2059 let mut out = Vec::with_capacity(b);
2060 for s in 0..b {
2061 let mut ys = e.uninit(ts[s] * dim)?;
2062 e.copy_view_into(&mut ys, 0, &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim), ts[s] * dim)?;
2063 out.push(ys);
2064 }
2065 Ok(out)
2066 };
2067
2068 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2069 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
2071 let mut h = e.uninit(total * n_embd)?;
2072 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2073 e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut hx16, n_embd, total, eps)?;
2074 let mut mixed = e.uninit(total * n_embd)?;
2076 match &layer.mixer {
2077 Mixer::Full(fa) => {
2078 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2079 let geometry = self.cfg.full_attention_geometry_at(il as u32);
2085 let (n_head, n_head_kv, head_dim) = (
2086 geometry.n_head as usize,
2087 geometry.n_head_kv as usize,
2088 geometry.head_dim_k as usize,
2089 );
2090 let fa_scale = geometry.attention_scale();
2091 let use_favl = !carried
2092 && (2..=8).contains(&b)
2093 && (head_dim == 256 || head_dim == 128)
2094 && geometry.attention_gate
2095 == memra_gguf::config::AttentionGateKind::FusedQ
2096 && std::env::var("MEMRA_NOFA").is_err()
2097 && std::env::var("MEMRA_FA_FLOOR").is_err()
2098 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2099 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2100 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2101 if use_favl {
2102 let (qf_w, kf_w, vf_w) =
2103 (fa.wq.out_features(), fa.wk.out_features(), fa.wv.out_features());
2104 struct APre {
2105 q: CudaSlice<f32>, gate: Option<CudaSlice<f32>>,
2106 qn: CudaSlice<f32>, kn: CudaSlice<f32>,
2107 }
2108 let mut aps = Vec::with_capacity(b);
2109 for &t in ts.iter().take(b) {
2110 aps.push(APre {
2111 q: e.uninit(t * n_head * head_dim)?,
2112 gate: Some(e.uninit(t * n_head * head_dim)?),
2113 qn: e.uninit(t * n_head * head_dim)?,
2114 kn: e.uninit(t * n_head_kv * head_dim)?,
2115 });
2116 }
2117 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2118 let kvl = caches[0].kv[il].as_ref().unwrap();
2119 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2120 };
2121 let pargs: Vec<crate::AttnPreVl> = (0..b).map(|s| {
2122 let (o, t) = (offs[s], ts[s]);
2123 let kvl = caches[s].kv[il].as_ref().unwrap();
2124 assert!(kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2125 "prime_cache_batch attn vl: fresh + capacity");
2126 crate::AttnPreVl {
2127 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2128 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2129 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2130 q: e.addr_f32(&aps[s].q),
2131 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2132 qn: e.addr_f32(&aps[s].qn), kn: e.addr_f32(&aps[s].kn),
2133 kc: e.addr_u8(&kvl.k), vc: e.addr_u8(&kvl.v),
2134 t: t as i32, pad: 0,
2135 }
2136 }).collect();
2137 e.attn_pre_vl8(&pargs, fa.q_norm.float_data(), fa.k_norm.float_data(),
2138 head_dim, geometry.n_rot as usize, n_head, n_head_kv,
2139 self.cfg.rms_eps, geometry.rope_base, 1.0,
2140 kv_dim_k, kv_dim_v, ktb, vtb)?;
2141 for s in 0..b {
2142 let kvl = caches[s].kv[il].as_mut().unwrap();
2143 kvl.len += ts[s];
2144 let new_len = kvl.len as i32;
2145 e.set_i32_one(&mut kvl.len_d, new_len)?;
2146 }
2147 let mut attns = Vec::with_capacity(b);
2148 let mut mirrors = Vec::with_capacity(b);
2149 for &t in ts.iter().take(b) {
2150 attns.push(e.uninit(t * n_head * head_dim)?);
2151 let n = t * n_head_kv * head_dim;
2152 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2153 }
2154 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2157 Ok("0") => false,
2158 Ok("1") => true,
2159 _ => cfg!(memra_hopper_mma),
2160 };
2161 if fa3_on {
2162 let mut q16s = Vec::with_capacity(b);
2163 let mut v16s = Vec::with_capacity(b);
2164 for s in 0..b {
2165 let t = ts[s];
2166 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2167 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2168 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2169 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2170 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2171 e.f32_to_bf16_v(&g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2172 &mut v16, t * n_head_kv * head_dim)?;
2173 q16s.push(q16);
2174 v16s.push((k16, v16));
2175 }
2176 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2177 let mut kp = qp;
2178 let mut vp = qp;
2179 let mut op = [core::ptr::null_mut::<f32>(); 8];
2180 let mut tsv = [0i32; 8];
2181 for s in 0..b {
2182 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2183 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2184 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2185 op[s] = e.addr_f32(&attns[s]) as *mut f32;
2186 tsv[s] = ts[s] as i32;
2187 }
2188 let rc = unsafe {
2189 crate::fa3_vl_raw(qp.as_ptr(), kp.as_ptr(), vp.as_ptr(), op.as_ptr(),
2190 tsv.as_ptr(), b as i32, n_head as i32,
2191 n_head_kv as i32, head_dim as i32, fa_scale,
2192 e.stream().cu_stream() as *mut core::ffi::c_void)
2193 };
2194 if rc != 0 {
2195 return Err(format!("memra_fa3_vl rc={rc}").into());
2196 }
2197 } else {
2198 let fargs: Vec<crate::FaSeqVl> = (0..b).map(|s| crate::FaSeqVl {
2199 q: e.addr_f32(&aps[s].qn), k16: e.addr_u8(&mirrors[s].0),
2200 v16: e.addr_u8(&mirrors[s].1), o: e.addr_f32(&attns[s]),
2201 kf: e.addr_f32(&aps[s].kn),
2202 vf: e.addr_f32v(&g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w)),
2203 t: ts[s] as i32, pad: 0,
2204 }).collect();
2205 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2206 }
2207 for (s, attn) in attns.into_iter().enumerate() {
2208 let (attn_g, ag16) = self.full_attn_prime_post_fa(
2209 e, attn, &aps[s].gate, ts[s], n_head, head_dim)?;
2210 let mut done = false;
2211 if let Some(xh) = &ag16 {
2212 done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
2213 }
2214 if !done {
2215 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2216 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2217 }
2218 }
2219 } else {
2220 let mut parts: Vec<Vec<CudaSlice<f32>>> = (0..b).map(|_| Vec::new()).collect();
2221 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2222 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2223 parts[s].push(ys);
2224 }
2225 }
2226 for (s, g3s) in parts.into_iter().enumerate() {
2227 let (attn_g, ag16) = self.full_attn_prime_core_inner(
2229 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il)?;
2230 let mut done = false;
2231 if let Some(xh) = &ag16 {
2232 done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
2233 }
2234 if !done {
2235 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2236 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2237 }
2238 }
2239 }
2240 }
2241 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2242 Mixer::Linear(la) => {
2243 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2248 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2249 let outs = self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2250 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2251 let (o, t) = (offs[s], ts[s]);
2252 let mut done = false;
2253 if let Some(xh) = &gn16 {
2254 done = e.try_f16_gemm_pre_into_off(&la.ssm_out, xh, t, &mut mixed, o * n_embd)?;
2255 }
2256 if !done {
2257 let m = e.matmul(&la.ssm_out, &gn, t)?;
2258 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2259 }
2260 }
2261 }
2262 }
2263 let mut x1 = e.uninit(total * n_embd)?;
2264 let mut z = e.uninit(total * n_embd)?;
2265 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2266 e.add_rms_norm_f16out(&x, &mixed, layer.post_attn_norm.float_data(),
2267 &mut x1, &mut z, &mut zx16, n_embd, total, eps)?;
2268 let ffn_out = match &layer.ffn {
2269 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
2270 let n_ff = ffn_gate.out_features();
2271 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2272 let up = g2.pop().unwrap();
2273 let gate = g2.pop().unwrap();
2274 let mut act = e.uninit(total * n_ff)?;
2275 let d_lim = self.cfg.clamp_shexp_at(il as u32);
2279 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2280 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2281 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2282 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2283 Some(y) => y,
2284 None => e.matmul(ffn_down, &act, total)?,
2285 }
2286 } else {
2287 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, d_lim,
2288 &mut act, total * n_ff)?;
2289 e.matmul(ffn_down, &act, total)?
2290 }
2291 }
2292 crate::hybrid::Ffn::Moe(m) => {
2293 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2294 }
2295 };
2296 let mut x2 = e.uninit(total * n_embd)?;
2297 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2298 x = x2;
2299 }
2300 let mut hn = e.uninit(total * n_embd)?;
2302 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, total, eps)?;
2303 let mut hcat = e.uninit(b * n_embd)?;
2309 for s in 0..b {
2310 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2311 e.copy_view_into(&mut hcat, s * n_embd, &hn.slice(last0..last0 + n_embd), n_embd)?;
2312 }
2313 let logits_cat = if b >= 2 { e.try_f16_gemm(&self.output, &hcat, b)? } else { None };
2314 let logits_host: Option<Vec<f32>> = match &logits_cat {
2315 Some(lc) => Some(e.dtoh(lc)?),
2316 None => None,
2317 };
2318 let n_vocab = self.output.out_features();
2319 let mut hidden_all = if crate::spec::spec_hpost() {
2320 split(e, &hn, n_embd)?
2321 } else {
2322 split(e, &x, n_embd)?
2323 };
2324 let mut out = Vec::with_capacity(b);
2325 for s in 0..b {
2326 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2327 let mut h_seed = e.uninit(n_embd)?;
2328 if !crate::spec::spec_hpost() {
2329 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2330 } else {
2331 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2332 }
2333 let logits = match &logits_host {
2334 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2335 None => {
2336 let mut hlast = e.uninit(n_embd)?;
2337 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2338 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2339 }
2340 };
2341 caches[s].pos += ts[s];
2342 out.push((logits, h_seed, hidden_all.remove(0)));
2343 }
2344 Ok(out)
2345 }
2346
2347 #[allow(clippy::too_many_arguments)]
2358 fn full_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
2359 hx: Option<&CudaSlice<u8>>,
2360 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize,
2361 seq_end: usize)
2362 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2363 if self.cfg.step35.is_some() {
2364 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2365 }
2366 let g3 = match hx {
2371 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2372 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2373 };
2374 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2375 }
2376
2377 fn full_attn_prime_core(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
2381 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2382 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2383 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2384 if let Some(xh) = &ag16 {
2385 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2386 return Ok(y);
2387 }
2388 }
2389 Ok(e.matmul(&fa.wo, &attn_g, t)?)
2390 }
2391
2392 fn full_attn_prime_core_inner(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
2393 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2394 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2395 let cfg = &self.cfg;
2396 let geometry = cfg.full_attention_geometry_at(il as u32);
2397 let n_head = geometry.n_head as usize;
2398 let n_head_kv = geometry.n_head_kv as usize;
2399 let head_dim = geometry.head_dim_k as usize;
2400 let scale = geometry.attention_scale();
2401 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2402 let AttnPre { q, k, v, gate } = pre;
2403 let mut attn = e.uninit(t * n_head * head_dim)?;
2404 self.full_attn_prime_fa_dispatch(e, &q, &k, &v, &mut attn, base_len, t, cache, il,
2405 head_dim, n_head, n_head_kv, scale)?;
2406 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2407 }
2408
2409 #[allow(clippy::type_complexity)]
2413 fn full_attn_prime_pre_fa(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
2414 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2415 -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
2416 let cfg = &self.cfg;
2417 let geometry = cfg.full_attention_geometry_at(il as u32);
2418 let n_head = geometry.n_head as usize;
2419 let n_head_kv = geometry.n_head_kv as usize;
2420 let head_dim = geometry.head_dim_k as usize;
2421 let eps = cfg.rms_eps;
2422
2423 let gated = geometry.attention_gate
2427 == memra_gguf::config::AttentionGateKind::FusedQ;
2428 let v = g3.pop().unwrap();
2429 let mut k = g3.pop().unwrap();
2430 let qf = g3.pop().unwrap();
2431 let (mut q, gate) = if gated {
2432 let mut q = e.uninit(t * n_head * head_dim)?;
2433 let mut gate = e.uninit(t * n_head * head_dim)?;
2434 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2435 (q, Some(gate))
2436 } else {
2437 (qf, None)
2438 };
2439
2440 let mut qn = e.uninit(t * n_head * head_dim)?;
2441 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2442 q = qn;
2443 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2444 e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2445 k = kn;
2446 let rope_dims = geometry.n_rot as usize;
2447 e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, geometry.rope_base, 1.0)?;
2448 e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, geometry.rope_base, 1.0)?;
2449
2450 {
2453 let kvl = cache.kv[il].as_mut().unwrap();
2454 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
2455 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
2456 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
2457 crate::Engine::kv_fp8_on())?;
2458 kvl.len += t;
2459 let new_len = kvl.len as i32;
2460 e.set_i32_one(&mut kvl.len_d, new_len)?;
2461 }
2462
2463 let base_len = {
2464 let kvl = cache.kv[il].as_ref().unwrap();
2465 kvl.len - t };
2467 Ok((AttnPre { q, k, v, gate }, base_len))
2468 }
2469
2470 #[allow(clippy::too_many_arguments)]
2477 fn full_attn_prime_fa_dispatch(&self, e: &Engine, q: &CudaSlice<f32>, k: &CudaSlice<f32>,
2478 v: &CudaSlice<f32>, attn: &mut CudaSlice<f32>, base_len: usize,
2479 t: usize, cache: &mut Cache, il: usize,
2480 head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32)
2481 -> Result<(), Box<dyn std::error::Error>> {
2482 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
2495 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2496 e.sdpa_naive(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2497 } else {
2498 e.fa_prefill(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2499 }
2500 return Ok(());
2501 }
2502 let kvl = cache.kv[il].as_ref().unwrap();
2503 let t_kv = base_len + t;
2504 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
2505 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
2506 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2510 e.sdpa_naive_quantized_view(q, &k_view, &v_view, attn, head_dim, n_head,
2511 n_head_kv, t, t_kv, scale, true,
2512 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
2513 return Ok(());
2514 }
2515 let deqw = std::env::var("MEMRA_PRIME_DEQW").map(|v| v != "0").unwrap_or(true);
2523 if deqw {
2524 e.fa_prefill_view_ws(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
2525 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
2526 crate::Engine::kv_fp8_on())?;
2527 } else {
2528 e.fa_prefill_view(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
2529 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
2530 crate::Engine::kv_fp8_on())?;
2531 }
2532 Ok(())
2533 }
2534
2535 fn full_attn_prime_post_fa(&self, e: &Engine, attn: CudaSlice<f32>,
2538 gate: &Option<CudaSlice<f32>>, t: usize,
2539 n_head: usize, head_dim: usize)
2540 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2541 let (attn_g, ag16) = match gate {
2542 Some(gate) => {
2543 let n = t * n_head * head_dim;
2544 let mut ag = e.uninit(n)?;
2545 if Self::f16out_on(e, t) {
2546 let mut a16 = e.alloc_u8_uninit(n * 2)?;
2547 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
2548 (ag, Some(a16))
2549 } else {
2550 let mut gsig = e.uninit(n)?;
2551 e.sigmoid(gate, &mut gsig, n)?;
2552 e.mul(&attn, &gsig, &mut ag, n)?;
2553 (ag, None)
2554 }
2555 }
2556 None => (attn, None),
2557 };
2558 Ok((attn_g, ag16))
2559 }
2560
2561 fn linear_attn_prime(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>,
2568 hx: Option<&CudaSlice<u8>>, t: usize,
2569 cache: &mut Cache, il: usize)
2570 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2571 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2573 let g4 = match hx {
2574 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
2575 None => e.matmul_group(&ws, h, t)?,
2576 };
2577 self.linear_attn_prime_core(e, la, g4, t, cache, il)
2578 }
2579
2580 fn linear_attn_prime_core(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
2582 t: usize, cache: &mut Cache, il: usize)
2583 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2584 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
2585 }
2586
2587 #[allow(clippy::too_many_arguments)]
2591 fn linear_attn_prime_core_pad_inner(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
2592 t: usize, cache: &mut Cache, il: usize,
2593 pad_len: Option<&CudaSlice<i32>>)
2594 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2595 let ssm = self.cfg.ssm.as_ref().unwrap();
2597 let d_state = ssm.state_size as usize;
2598 let num_k = ssm.group_count as usize;
2599 let num_v = ssm.time_step_rank as usize;
2600 let key_dim = d_state * num_k;
2601 let value_dim = d_state * num_v;
2602 let conv_dim = key_dim * 2 + value_dim;
2603 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(
2608 e, la,
2609 &qkv_mixed.slice(0..t * conv_dim), &z.slice(0..t * value_dim),
2610 &beta_raw.slice(0..t * num_v), &alpha.slice(0..t * num_v),
2611 t, cache, il, pad_len)
2612 }
2613
2614 #[allow(clippy::too_many_arguments)]
2617 fn linear_attn_gdn_prep(&self, e: &Engine, la: &LinearAttnLayer,
2618 qkv_mixed: &cudarc::driver::CudaView<f32>,
2619 beta_raw: &cudarc::driver::CudaView<f32>,
2620 alpha: &cudarc::driver::CudaView<f32>,
2621 t: usize, cache: &mut Cache, il: usize,
2622 pad_len: Option<&CudaSlice<i32>>)
2623 -> Result<GdnPrep, Box<dyn std::error::Error>> {
2624 let cfg = &self.cfg;
2625 let ssm = cfg.ssm.as_ref().unwrap();
2626 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;
2634 debug_assert!(t >= d_conv - 1, "stateful conv needs T >= pad (PRIME_MIN_T gates)");
2635
2636 let rl = cache.recur[il].as_mut().unwrap();
2641 let hk = Self::gdn_hk(e, t, num_v, num_k);
2642 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
2643 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
2645 let mut k_g = e.uninit(d_state * hk * t)?;
2646 let mut v_g = e.uninit(d_state * num_v * t)?;
2647 if conv_fuse {
2648 e.ssm_conv1d_gdn_state_pad(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
2649 &mut q_g, &mut k_g, &mut v_g,
2650 conv_dim, t, d_conv, d_state, num_v, num_k, key_dim, hk, pad_len)?;
2651 } else {
2652 let mut conv_out = e.uninit(conv_dim * t)?; e.ssm_conv1d_tm_state_pad_v(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
2654 &mut conv_out, conv_dim, t, d_conv, pad_len)?;
2655 e.qkv_to_gdn_repack(&conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t)?;
2656 }
2657 let mut q_l2 = e.uninit(d_state * hk * t)?;
2658 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
2662 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
2663 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
2664 Some(qb)
2665 } else {
2666 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
2667 None
2668 };
2669 let mut k_l2 = e.uninit(d_state * hk * t)?;
2670 let kb16 = if Engine::l2_v2_on(d_state) {
2672 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
2673 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
2674 Some(kb)
2675 } else {
2676 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
2677 None
2678 };
2679 let mut beta = e.uninit(t * num_v)?;
2680 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
2681 let mut g_log = e.uninit(t * num_v)?;
2682 e.gdn_glog_v(alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
2683 if let Some(len_d) = pad_len {
2684 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
2685 }
2686 Ok(GdnPrep { hk, q_l2, k_l2, v_g, beta, g_log, kb16, qb16 })
2687 }
2688
2689 #[allow(clippy::too_many_arguments)]
2694 fn linear_attn_prime_core_batch(&self, e: &Engine, la: &LinearAttnLayer,
2695 g4: &[CudaSlice<f32>], offs: &[usize], ts: &[usize],
2696 caches: &mut [&mut Cache], il: usize)
2697 -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
2698 let ssm = self.cfg.ssm.as_ref().unwrap();
2699 let d_state = ssm.state_size as usize;
2700 let num_k = ssm.group_count as usize;
2701 let num_v = ssm.time_step_rank as usize;
2702 let key_dim = d_state * num_k;
2703 let value_dim = d_state * num_v;
2704 let conv_dim = key_dim * 2 + value_dim;
2705 let eps = self.cfg.rms_eps;
2706 let scale = 1.0 / (d_state as f32).sqrt();
2707 let b = ts.len();
2708 let c = Engine::gdn_chunk_size();
2709 let carried = caches.iter().any(|c| c.pos > 0);
2712 let use_vl = !carried
2713 && (2..=8).contains(&b)
2714 && Engine::gdn_chunked_enabled() && ts.iter().all(|&t| t >= 16)
2715 && e.gdn_mma_enabled(c)
2716 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
2717 if !use_vl {
2718 return (0..b).map(|s| {
2719 let (o, t) = (offs[s], ts[s]);
2720 self.linear_attn_prime_core_pad_view(
2721 e, la,
2722 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
2723 &g4[1].slice(o * value_dim..(o + t) * value_dim),
2724 &g4[2].slice(o * num_v..(o + t) * num_v),
2725 &g4[3].slice(o * num_v..(o + t) * num_v),
2726 t, caches[s], il, None)
2727 }).collect();
2728 }
2729 struct SeqBufs {
2733 conv_out: CudaSlice<f32>, q_g: CudaSlice<f32>, k_g: CudaSlice<f32>, v_g: CudaSlice<f32>,
2734 q_l2: CudaSlice<f32>, k_l2: CudaSlice<f32>, beta: CudaSlice<f32>, g_log: CudaSlice<f32>,
2735 gn: CudaSlice<f32>, gn16: CudaSlice<u8>,
2736 }
2737 let d_conv = ssm.conv_kernel as usize;
2738 let f16o = Self::f16out_on(e, 16);
2739 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
2741 let mut pres = Vec::with_capacity(b);
2742 for &t in ts.iter().take(b) {
2743 sb.push(SeqBufs {
2744 conv_out: e.uninit(conv_dim * t)?,
2745 q_g: e.uninit(d_state * hk * t)?,
2746 k_g: e.uninit(d_state * hk * t)?,
2747 v_g: e.uninit(d_state * num_v * t)?,
2748 q_l2: e.uninit(d_state * hk * t)?,
2749 k_l2: e.uninit(d_state * hk * t)?,
2750 beta: e.uninit(t * num_v)?,
2751 g_log: e.uninit(t * num_v)?,
2752 gn: e.uninit(d_state * num_v * t)?,
2753 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
2754 });
2755 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
2756 }
2757 let prep_args: Vec<crate::GdnPrepVl> = (0..b).map(|s| {
2758 let (o, t) = (offs[s], ts[s]);
2759 let rl = caches[s].recur[il].as_ref().unwrap();
2760 crate::GdnPrepVl {
2761 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
2762 conv_state: e.addr_f32(&rl.conv_state),
2763 conv_out: e.addr_f32(&sb[s].conv_out),
2764 q_g: e.addr_f32(&sb[s].q_g), k_g: e.addr_f32(&sb[s].k_g), v_g: e.addr_f32(&sb[s].v_g),
2765 q_l2: e.addr_f32(&sb[s].q_l2), k_l2: e.addr_f32(&sb[s].k_l2),
2766 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
2767 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
2768 beta: e.addr_f32(&sb[s].beta), g_log: e.addr_f32(&sb[s].g_log),
2769 o: e.addr_f32(&pres[s].o),
2770 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
2771 gn: e.addr_f32(&sb[s].gn), gn16: e.addr_u8(&sb[s].gn16),
2772 kb16: if Engine::l2_v2_on(d_state) { e.addr_u8(&pres[s].kb16) } else { 0 },
2773 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) { e.addr_u8(&pres[s].qb16) } else { 0 },
2774 t: t as i32, pad: 0,
2775 }
2776 }).collect();
2777 let args: Vec<crate::GdnSeqVl> = (0..b).map(|s| {
2778 let rl = caches[s].recur[il].as_ref().unwrap();
2779 crate::GdnSeqVl {
2780 kb16: e.addr_u8(&pres[s].kb16), gcum: e.addr_f32(&pres[s].gcum),
2781 beta: e.addr_f32(&sb[s].beta), u: e.addr_f32(&pres[s].u),
2782 wb16: e.addr_u8(&pres[s].wb16), y: e.addr_u8(&pres[s].y16),
2783 ssnap: e.addr_u8(&pres[s].ssnap16),
2784 state_in: e.addr_f32(&rl.ssm_state), state_out: e.addr_f32(&rl.ssm_state_alt),
2785 q: e.addr_f32(&sb[s].q_l2), p: e.addr_f32(&pres[s].p),
2786 o: e.addr_f32(&pres[s].o),
2787 k: e.addr_f32(&sb[s].k_l2), v: e.addr_f32(&sb[s].v_g),
2788 g: e.addr_f32(&sb[s].g_log), a: e.addr_f32(&pres[s].a),
2789 w: e.addr_f32(&pres[s].w),
2790 t: ts[s] as i32, nc: pres[s].nc as i32,
2791 }
2792 }).collect();
2793 e.gdn_prep_vl8(&prep_args, la.ssm_conv1d.float_data(), la.ssm_dt.float_data(),
2794 la.ssm_a.float_data(), conv_dim, d_conv, d_state, num_v, num_k, key_dim, hk, eps)?;
2795 if !Engine::l2_v2_on(d_state) {
2798 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
2799 }
2800 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
2802 if !Engine::l2_v2_on(d_state) {
2804 for s in 0..b {
2805 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
2806 }
2807 }
2808 let mut wa = [crate::GdnWVl::default(); 8];
2809 for s in 0..b {
2810 wa[s] = crate::GdnWVl { qb16: e.addr_u8(&pres[s].qb16), pb16: e.addr_u8(&pres[s].pb16) };
2811 }
2812 Some(crate::GdnWVl8(wa))
2813 } else { None };
2814 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
2815 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
2816 if f16o {
2817 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
2818 }
2819 let mut out = Vec::with_capacity(b);
2821 for (s, bufs) in sb.into_iter().enumerate() {
2822 let rl = caches[s].recur[il].as_mut().unwrap();
2823 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2824 let (o, t) = (offs[s], ts[s]);
2825 let SeqBufs { mut gn, gn16, .. } = bufs;
2826 if f16o {
2827 out.push((gn, Some(gn16)));
2828 } else {
2829 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
2830 e.gated_rmsnorm_zv(&pres[s].o, la.ssm_norm.float_data(), &z_v, &mut gn,
2831 d_state, num_v * t, eps)?;
2832 out.push((gn, None));
2833 }
2834 }
2835 Ok(out)
2836 }
2837
2838 #[allow(clippy::too_many_arguments)]
2842 fn linear_attn_prime_core_pad_view(&self, e: &Engine, la: &LinearAttnLayer,
2843 qkv_mixed: &cudarc::driver::CudaView<f32>,
2844 z: &cudarc::driver::CudaView<f32>,
2845 beta_raw: &cudarc::driver::CudaView<f32>,
2846 alpha: &cudarc::driver::CudaView<f32>,
2847 t: usize, cache: &mut Cache, il: usize,
2848 pad_len: Option<&CudaSlice<i32>>)
2849 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2850 let cfg = &self.cfg;
2851 let ssm = cfg.ssm.as_ref().unwrap();
2852 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
2855 let scale = 1.0 / (d_state as f32).sqrt();
2856
2857 let prep = self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
2858
2859 let mut o = e.uninit(d_state * num_v * t)?;
2865 let rl = cache.recur[il].as_mut().unwrap();
2866 {
2867 let crate::cache::RecurLayer { ssm_state, ssm_state_alt, .. } = rl;
2868 e.gdn_scan_prefill(&prep.q_l2, &prep.k_l2, &prep.v_g, &prep.g_log, &prep.beta,
2869 prep.kb16.as_ref(), prep.qb16.as_ref(), ssm_state, ssm_state_alt, &mut o, num_v, t, scale,
2870 prep.hk)?;
2871 }
2872 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2873
2874 let mut gn = e.uninit(d_state * num_v * t)?;
2877 let gn16 = if Self::f16out_on(e, t) {
2878 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
2879 e.gated_rmsnorm_f16out_zv(&o, la.ssm_norm.float_data(), z, &mut gn, &mut g16,
2880 d_state, num_v * t, eps)?;
2881 Some(g16)
2882 } else {
2883 e.gated_rmsnorm_zv(&o, la.ssm_norm.float_data(), z, &mut gn, d_state, num_v * t, eps)?;
2884 None
2885 };
2886 Ok((gn, gn16))
2887 }
2888
2889 #[allow(clippy::too_many_arguments)]
2891 fn linear_attn_prime_core_pad(&self, e: &Engine, la: &LinearAttnLayer, g4: Vec<CudaSlice<f32>>,
2892 t: usize, cache: &mut Cache, il: usize,
2893 pad_len: Option<&CudaSlice<i32>>)
2894 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2895 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
2896 if let Some(xh) = &gn16 {
2897 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
2898 return Ok(y);
2899 }
2900 }
2901 Ok(e.matmul(&la.ssm_out, &gn, t)?)
2902 }
2903
2904 pub fn full_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, il: usize)
2909 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2910 if self.cfg.step35.is_some() {
2911 return self.step35_attn(e, fa, h, pos_d, t, il);
2912 }
2913 let cfg = &self.cfg;
2914 let _n_embd = cfg.n_embd as usize;
2915 let geometry = cfg.full_attention_geometry_at(il as u32);
2916 let n_head = geometry.n_head as usize;
2917 let n_head_kv = geometry.n_head_kv as usize;
2918 let head_dim = geometry.head_dim_k as usize;
2919 let eps = cfg.rms_eps;
2920 let scale = geometry.attention_scale();
2921
2922 let gated = geometry.attention_gate
2925 == memra_gguf::config::AttentionGateKind::FusedQ;
2926 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
2928 let v = g3.pop().unwrap();
2929 let mut k = g3.pop().unwrap();
2930 let qf = g3.pop().unwrap();
2931 let (mut q, gate) = if gated {
2932 let mut q = e.uninit(t * n_head * head_dim)?;
2933 let mut gate = e.uninit(t * n_head * head_dim)?;
2934 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2935 (q, Some(gate))
2936 } else {
2937 (qf, None)
2938 };
2939
2940 let mut qn = e.uninit(t * n_head * head_dim)?;
2942 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2943 q = qn;
2944 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2945 e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2946 k = kn;
2947 let rope_dims = geometry.n_rot as usize;
2948 e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, geometry.rope_base, 1.0)?;
2949 e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, geometry.rope_base, 1.0)?;
2950
2951 let mut attn = e.uninit(t * n_head * head_dim)?;
2953 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2956 e.sdpa_naive(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2958 } else {
2959 e.fa_prefill(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2960 }
2961
2962 let attn_g = match &gate {
2964 Some(gate) => {
2965 let mut gsig = e.uninit(t * n_head * head_dim)?;
2966 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
2967 let mut ag = e.uninit(t * n_head * head_dim)?;
2968 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
2969 ag
2970 }
2971 None => attn,
2972 };
2973
2974 let o = e.matmul(&fa.wo, &attn_g, t)?;
2976 Ok(o)
2977 }
2978
2979 pub fn linear_attn(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize)
2981 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2982 let cfg = &self.cfg;
2983 let _n_embd = cfg.n_embd as usize;
2984 let ssm = cfg.ssm.as_ref().unwrap();
2985 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; let head_v = d_state;
2990 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;
2994 let scale = 1.0 / (d_state as f32).sqrt();
2995
2996 let mut g4 = e.matmul_group(&[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha], h, t)?;
2999 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);
3011 let mut q_g = e.uninit(d_state * num_v * t)?;
3012 let mut k_g = e.uninit(d_state * num_v * t)?;
3013 let mut v_g = e.uninit(d_state * num_v * t)?;
3014 e.ssm_conv1d_gdn(&qkv_mixed, la.ssm_conv1d.float_data(), &mut q_g, &mut k_g, &mut v_g,
3015 conv_dim, t, d_conv, d_state, num_v, num_k, key_dim)?;
3016 let mut q_l2 = e.uninit(d_state * num_v * t)?;
3018 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3019 let mut k_l2 = e.uninit(d_state * num_v * t)?;
3020 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3021 let v_gd = v_g;
3022
3023 let mut beta = e.uninit(t * num_v)?;
3026 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3027 let mut g_log = e.uninit(t * num_v)?;
3029 e.gdn_glog(&alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
3030
3031 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
3034 let mut o = e.uninit(d_state * num_v * t)?;
3035 e.gdn_scan_prefill(&q_l2, &k_l2, &v_gd, &g_log, &beta, None, None, &state_in, &mut state_out, &mut o, num_v, t, scale, num_v)?;
3036
3037 let mut gn = e.uninit(d_state * num_v * t)?;
3042 e.gated_rmsnorm(&o, la.ssm_norm.float_data(), &z, &mut gn, d_state, num_v * t, eps)?;
3043
3044 let out = e.matmul(&la.ssm_out, &gn, t)?;
3048 Ok(out)
3049 }
3050}
3051
3052impl HybridModel {
3053 pub fn moe_ffn_il(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16)
3064 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3065 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
3066 }
3067
3068 pub fn moe_ffn_il_prefill(
3071 &self,
3072 e: &Engine,
3073 m: &MoeWeights,
3074 z: &CudaSlice<f32>,
3075 t: usize,
3076 il: u16,
3077 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3078 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
3079 }
3080
3081 pub fn moe_ffn_il_zq8(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
3085 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16)
3086 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3087 Self::moe_ffn_inner(
3088 e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false,
3089 )
3090 }
3091
3092 pub(crate) fn moe_ffn(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
3100 cfg: &ModelConfig, il: u16, max_block: usize)
3101 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3102 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
3103 }
3104
3105 #[allow(clippy::too_many_arguments)]
3106 pub(crate) fn moe_ffn_inner(
3107 e: &Engine,
3108 m: &MoeWeights,
3109 z: &CudaSlice<f32>,
3110 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
3111 t: usize,
3112 cfg: &ModelConfig,
3113 il: u16,
3114 max_block: usize,
3115 prefill: bool,
3116 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3117 let worker_io = crate::spill_pread::worker_enabled();
3118 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
3119 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
3120 e.with_moe_cache(max_block, |cache, _| {
3121 cache.begin_forward_epoch(il, t);
3122 if worker_io {
3123 cache.begin_worker_scope();
3124 }
3125 Ok(())
3126 })?;
3127 }
3128 if t > 1 && moe_grouped_enabled(cfg, prefill) {
3131 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
3132 if std::env::var("MEMRA_MOE_GATE").is_ok() {
3137 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
3138 let g_host = e.dtoh(&grouped_out)?;
3139 let s_host = e.dtoh(&seq_out)?;
3140 let g_bytes: &[u8] = unsafe { std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4) };
3141 let s_bytes: &[u8] = unsafe { std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4) };
3142 if g_bytes == s_bytes {
3143 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
3144 } else {
3145 let diffs = g_host.iter().zip(s_host.iter()).enumerate()
3146 .filter(|(_, (a, b))| a != b).count();
3147 let maxdiff = g_host.iter().zip(s_host.iter())
3148 .map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
3149 panic!("moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}", g_host.len());
3150 }
3151 }
3152 return Ok(grouped_out);
3153 }
3154 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
3155 }
3156
3157 pub(crate) fn moe_ffn_sequential(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
3159 cfg: &ModelConfig, il: u16, max_block: usize)
3160 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3161 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
3162 }
3163
3164 fn moe_router_logits(
3168 e: &Engine,
3169 m: &MoeWeights,
3170 z: &CudaSlice<f32>,
3171 t: usize,
3172 cfg: &ModelConfig,
3173 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3174 if t < PRIME_MIN_T {
3175 if crate::router_kernel_on() {
3177 e.router_gemv(
3178 m.gate_inp.float_data(),
3179 z,
3180 cfg.n_embd as usize,
3181 m.gate_exps.n_expert,
3182 t,
3183 )
3184 } else {
3185 e.matmul_decode_exact(&m.gate_inp, z, t)
3186 }
3187 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
3188 e.router_gemv(
3189 m.gate_inp.float_data(),
3190 z,
3191 cfg.n_embd as usize,
3192 m.gate_exps.n_expert,
3193 t,
3194 )
3195 } else {
3196 e.matmul(&m.gate_inp, z, t)
3197 }
3198 }
3199
3200 fn trace_moe_routes(il: u16, t: usize, sel_all: &[u32], weights: &[f32])
3204 -> Result<(), Box<dyn std::error::Error>> {
3205 use std::io::Write as _;
3206 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
3207 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
3208 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
3209 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
3210 }
3211 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
3212 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
3213 let pairs: Vec<String> = sel_all.iter().zip(weights)
3214 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
3215 .collect();
3216 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
3217 }
3218 Ok(())
3219 }
3220
3221 fn trace_moe_input(e: &Engine, il: u16, t: usize, n_embd: usize, z: &CudaSlice<f32>)
3226 -> Result<(), Box<dyn std::error::Error>> {
3227 use std::io::Write as _;
3228 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else { return Ok(()) };
3229 let host = e.dtoh(z)?;
3230 if host.len() != t * n_embd {
3231 return Err(format!(
3232 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
3233 host.len(), t, n_embd
3234 ).into());
3235 }
3236 let bytes = unsafe {
3237 std::slice::from_raw_parts(
3238 host.as_ptr().cast::<u8>(), host.len() * std::mem::size_of::<f32>()
3239 )
3240 };
3241 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
3242 let mut state = state.lock().map_err(|_| "MoE input trace writer lock is poisoned")?;
3243 if state.is_none() {
3244 let dir = std::path::PathBuf::from(&dir);
3245 std::fs::create_dir_all(&dir)?;
3246 let index = std::fs::OpenOptions::new().create(true).append(true)
3247 .open(dir.join("index.jsonl"))?;
3248 *state = Some(MoeInputTraceWriter {
3249 dir,
3250 index,
3251 payloads: std::collections::HashMap::new(),
3252 });
3253 }
3254 let writer = state.as_mut().unwrap();
3255 if writer.dir != std::path::Path::new(&dir) {
3256 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
3257 }
3258 let file_name = format!("layer-{il:03}.f32");
3259 if !writer.payloads.contains_key(&il) {
3260 let payload = std::fs::OpenOptions::new().create(true).append(true)
3261 .open(writer.dir.join(&file_name))?;
3262 let offset = payload.metadata()?.len();
3263 writer.payloads.insert(il, (payload, offset));
3264 }
3265 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
3266 let row_offset = *offset;
3267 payload.write_all(bytes)?;
3268 *offset += bytes.len() as u64;
3269 writeln!(
3270 writer.index,
3271 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
3272 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
3273 \"payload_bytes\":{}}}",
3274 bytes.len()
3275 )?;
3276 Ok(())
3277 }
3278
3279 #[allow(clippy::too_many_arguments)]
3280 pub(crate) fn moe_ffn_sequential_zq8(
3281 e: &Engine,
3282 m: &MoeWeights,
3283 z: &CudaSlice<f32>,
3284 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
3285 t: usize,
3286 cfg: &ModelConfig,
3287 il: u16,
3288 max_block: usize,
3289 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3290 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
3291 let moe = cfg.moe.as_ref().unwrap();
3292 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);
3299 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
3300 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);
3303
3304 let lim_exp = cfg.clamp_exp_at(il as u32);
3307 let lim_shexp = cfg.clamp_shexp_at(il as u32);
3308 let use_cache = Engine::moe_cache_enabled();
3309 let uniform_experts = m.has_uniform_expert_layout();
3310 let moe_q8 = uniform_experts && moe_q8_enabled()
3311 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3312 && q8_expert_supported(m.down_exps.qtype);
3313 let cpu_expert_requested = crate::cpu_experts::configured();
3320 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
3321 return Err(std::io::Error::other(
3322 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
3323 )
3324 .into());
3325 }
3326 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
3327 let freeze_cpu_residency = cpu_expert_requested
3333 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
3334 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
3335 .ok()
3336 .and_then(|value| value.parse::<usize>().ok())
3337 .is_some_and(|tokens| tokens > 0);
3338 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
3339 e.freeze_moe_cache();
3340 }
3341 let cache_frozen = use_cache && e.moe_cache_frozen();
3342 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
3343
3344 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
3347
3348 let no_exp_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
3386 && m.down_exps.macros.is_none();
3387 if cfg.sigmoid_router().is_none() && cfg.m3.is_none() && cfg.hy3.is_none()
3391 && !cfg.swiglu_clamped_at(il as u32)
3392 && no_exp_macros
3393 && t >= PRIME_MIN_T && m.dev_exps.is_some() && moe_q8_enabled()
3394 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3395 && q8_expert_supported(m.down_exps.qtype)
3396 && std::env::var("MEMRA_MOE_PAIRS").map(|v| v != "0").unwrap_or(true)
3397 && std::env::var("MEMRA_MOE_STATS").is_err() {
3398 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
3399 }
3400
3401 let dev_ok = uniform_experts && cfg.sigmoid_router().is_none()
3418 && cfg.m3.is_none() && cfg.hy3.is_none()
3419 && !cfg.swiglu_clamped_at(il as u32);
3420 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
3424 || std::env::var("MEMRA_MOE_TRACE").is_ok()
3425 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
3426 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
3427 if dev_ok && t < PRIME_MIN_T && m.dev_exps.is_some() && n_used <= 8 && moe_dev_enabled()
3428 && !observe_routes {
3429 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
3430 }
3431 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled()
3432 && !observe_routes {
3433 let row_ok = e.with_moe_cache(max_block, |c, eng| {
3434 if moe_prewarm_enabled() { c.prewarm_layer(il, m, eng)?; }
3435 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
3436 })?;
3437 if row_ok {
3438 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
3439 }
3440 }
3441
3442 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
3444 if cpu_hybrid {
3445 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
3446 e,
3447 &logits,
3448 z,
3449 t,
3450 n_expert,
3451 n_used,
3452 m.exp_probs_b.as_deref(),
3453 sig,
3454 m.active_experts.as_deref(),
3455 )?;
3456 (sel, w, Some(input))
3457 } else {
3458 let (sel, w) = Self::moe_route_cfg(
3459 e,
3460 &logits,
3461 t,
3462 n_expert,
3463 n_used,
3464 m.exp_probs_b.as_deref(),
3465 Some(sig),
3466 m.active_experts.as_deref(),
3467 )?;
3468 (sel, w, None)
3469 }
3470 } else {
3471 let (sel, w) = Self::moe_route_cfg(
3472 e,
3473 &logits,
3474 t,
3475 n_expert,
3476 n_used,
3477 None,
3478 None,
3479 m.active_experts.as_deref(),
3480 )?;
3481 (sel, w, None)
3482 };
3483
3484 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
3488 Self::trace_moe_input(e, il, t, n_embd, z)?;
3489
3490 let worker_disk_prefetch =
3502 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
3503 let promote_worker_h2d =
3504 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
3505 if promote_worker_h2d {
3506 let mut selected_blocks = Vec::with_capacity(n_used * 3);
3507 for &ex in sel_all.iter().take(n_used) {
3508 let ex = ex as u16;
3509 selected_blocks.extend([
3510 BlockId::new(il, PROJ_GATE, ex),
3511 BlockId::new(il, PROJ_UP, ex),
3512 BlockId::new(il, PROJ_DOWN, ex),
3513 ]);
3514 }
3515 for &ex in sel_all.iter().take(n_used) {
3516 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
3517 }
3518 e.with_moe_cache(max_block, |cache, eng| {
3519 cache.promote_worker_reads_at_safe_boundary(
3520 &selected_blocks,
3521 &selected_blocks,
3522 eng,
3523 )?;
3524 Ok(())
3525 })?;
3526 }
3527
3528 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
3531 let mut cnt = vec![0u32; n_expert];
3532 for &s in sel_all.iter() { cnt[s as usize] += 1; }
3533 let total = sel_all.len() as f64;
3534 let mut h = 0.0f64;
3535 let mut active = 0usize;
3536 for &c in &cnt { if c > 0 { active += 1; let p = c as f64 / total; h -= p * p.log2(); } }
3537 let maxc = cnt.iter().copied().max().unwrap_or(0);
3538 println!("moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
3539 il, t, sel_all.len(), active, n_expert, h, (n_expert as f64).log2(), total / active.max(1) as f64, maxc);
3540 }
3541
3542 let gdec_may_fire = uniform_experts && use_cache && n_used <= 8 && gdec_enabled()
3555 && !cfg.swiglu_clamped_at(il as u32);
3556 let slab_local = m.dev_exps.as_ref()
3572 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
3573 let slab_bases = slab_local.map(|d| {
3574 use cudarc::driver::DevicePtr;
3575 let s = e.stream();
3576 let (pg, _g0) = d.gate.device_ptr(&s);
3577 let (pu, _g1) = d.up.device_ptr(&s);
3578 let (pd, _g2) = d.down.device_ptr(&s);
3579 (pg as u64, pu as u64, pd as u64)
3580 });
3581 let slab_fused_may_fire = slab_bases.is_some() && n_used <= 8 && gdec_enabled()
3591 && !cfg.swiglu_clamped_at(il as u32) && cfg.m3.is_none()
3592 && no_exp_macros && moe_q8;
3593 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
3596 e.uninit(t * n_embd)?
3597 } else {
3598 e.zeros(t * n_embd)?
3599 };
3600 let cpu_input = if cpu_hybrid {
3603 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
3604 } else {
3605 None
3606 };
3607
3608 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;
3616 let mut scratch_u: Option<CudaSlice<u8>> = None;
3617 let mut scratch_d: Option<CudaSlice<u8>> = None;
3618 let page_window = moe_page_prefetch_window();
3626
3627 for tok in 0..t {
3630 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
3631 let w = &w_all[tok * n_used..(tok + 1) * n_used];
3632 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
3634
3635 let no_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
3649 && m.down_exps.macros.is_none();
3650 if slab_fused_may_fire {
3660 let (pg, pu, pd) = slab_bases.unwrap();
3661 let mut gp = [0u64; 8];
3662 let mut up = [0u64; 8];
3663 let mut dp = [0u64; 8];
3664 for (j, &ex) in sel.iter().enumerate() {
3665 let ex = ex as usize;
3666 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
3667 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
3668 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
3669 }
3670 let mut wv = [0f32; 8];
3671 wv[..n_used].copy_from_slice(w);
3672 if tok_q8.is_none() {
3673 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3674 }
3675 let (zq, zd) = tok_q8.as_ref().unwrap();
3676 let act = e.moe_gate_up_silu8_q8(crate::WPtr8(gp), crate::WPtr8(up), zq, zd,
3677 n_embd, n_ff_exp, n_used,
3678 m.gate_exps.qtype, m.up_exps.qtype,
3679 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
3680 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
3681 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3682 e.moe_down8_fma_q8(crate::WPtr8(dp), crate::F32x8(wv), &aq2, &ad2, &mut dst,
3683 n_ff_exp, n_embd, n_used,
3684 m.down_exps.qtype, m.down_exps.row_bytes)?;
3685 continue;
3686 }
3687 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
3688 if tok_q8.is_none() {
3689 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3690 }
3691 let (zq, zd) = tok_q8.as_ref().unwrap();
3692 if Self::moe_gdec_token_q8(e, m, il, max_block, zq, zd, sel, w,
3693 &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
3694 continue;
3695 }
3696 } else if gdec_may_fire && cfg.m3.is_none() && no_macros
3697 && Self::moe_gdec_token(e, m, il, max_block, &zt, sel, w,
3698 &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
3699 continue;
3700 }
3701
3702 if gdec_may_fire || slab_fused_may_fire {
3708 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3709 e.memset_zeros_view(&mut row)?;
3710 }
3711
3712 let mut cpu_mask = vec![false; sel.len()];
3718 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
3719 let gpu_resident = if use_cache {
3720 e.with_moe_cache(max_block, |cache, _| {
3721 Ok(sel
3722 .iter()
3723 .map(|&expert| {
3724 let expert = expert as u16;
3725 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
3726 .into_iter()
3727 .filter(|&projection| {
3728 cache
3729 .resident(BlockId::new(il, projection, expert))
3730 .is_some()
3731 })
3732 .count()
3733 })
3734 .collect::<Vec<_>>())
3735 })?
3736 } else {
3737 vec![0; sel.len()]
3738 };
3739 let mut cpu_selected = Vec::new();
3740 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
3741 if gpu_resident[index] != 3 {
3742 cpu_mask[index] = true;
3743 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
3744 let expert = expert as usize;
3745 cpu_selected.push((expert, route_weight));
3746 }
3747 }
3748 if crate::cpu_experts::predictor_enabled() {
3749 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
3753 crate::cpu_experts::predictor_submit(il, row);
3754 }
3755 if cpu_selected.is_empty() {
3756 None
3757 } else {
3758 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
3759 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
3760 .map_err(std::io::Error::other)?;
3761 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
3762 }
3763 } else {
3764 None
3765 };
3766
3767 let worker_window = worker_disk_prefetch
3768 .then(worker_prefetch_window)
3769 .unwrap_or(0);
3770 for (j, &ex) in sel.iter().enumerate() {
3771 if cpu_mask[j] {
3772 continue;
3773 }
3774 let ex = ex as usize;
3775 if let Some(d) = slab_local {
3782 let gl = m.gate_exps.expert_layout(ex);
3783 let ul = m.up_exps.expert_layout(ex);
3784 let dl = m.down_exps.expert_layout(ex);
3785 let (g0, u0, d0) = (ex * m.gate_exps.expert_stride,
3786 ex * m.up_exps.expert_stride,
3787 ex * m.down_exps.expert_stride);
3788 let (gate, up) = if moe_q8 {
3789 if tok_q8.is_none() {
3790 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3791 }
3792 let (zq, zd) = tok_q8.as_ref().unwrap();
3793 (e.qmatvec_expert_q8(&d.gate, g0..g0 + gl.len, zq, zd, 1,
3794 m.gate_exps.in_f, m.gate_exps.out_f,
3795 gl.qtype, gl.row_bytes)?,
3796 e.qmatvec_expert_q8(&d.up, u0..u0 + ul.len, zq, zd, 1,
3797 m.up_exps.in_f, m.up_exps.out_f,
3798 ul.qtype, ul.row_bytes)?)
3799 } else {
3800 (e.qmatvec_view(&d.gate, g0..g0 + gl.len, &zt, 1,
3801 m.gate_exps.in_f, m.gate_exps.out_f,
3802 gl.qtype, gl.row_bytes)?,
3803 e.qmatvec_view(&d.up, u0..u0 + ul.len, &zt, 1,
3804 m.up_exps.in_f, m.up_exps.out_f,
3805 ul.qtype, ul.row_bytes)?)
3806 };
3807 let mut act = e.uninit(n_ff_exp)?;
3808 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3809 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3810 let y = if moe_q8 {
3811 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
3812 e.qmatvec_expert_q8(&d.down, d0..d0 + dl.len, &aq2, &ad2, 1,
3813 m.down_exps.in_f, m.down_exps.out_f,
3814 dl.qtype, dl.row_bytes)?
3815 } else {
3816 let actv = act.slice(0..n_ff_exp);
3817 e.qmatvec_view(&d.down, d0..d0 + dl.len, &actv, 1,
3818 m.down_exps.in_f, m.down_exps.out_f,
3819 dl.qtype, dl.row_bytes)?
3820 };
3821 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3822 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3823 continue;
3824 }
3825 for next in page_prefetch_positions(j, sel.len(), page_window) {
3826 Self::moe_prefetch_host_expert(sel[next] as usize, m);
3827 }
3828 let keep = [
3829 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
3830 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
3831 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
3832 ];
3833 if worker_disk_prefetch && worker_window > 0 {
3834 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
3835 Self::moe_prefetch_disk_expert(
3836 e,
3837 il,
3838 sel[next] as usize,
3839 m,
3840 max_block,
3841 &keep,
3842 )?;
3843 }
3844 } else if cache_dispatch
3845 && !cpu_hybrid
3846 && moe_prefetch_enabled()
3847 && j + 1 < sel.len()
3848 {
3849 let next = sel[j + 1] as usize;
3850 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
3851 }
3852 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
3853 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
3854 if (gate_q8 || up_q8) && tok_q8.is_none() {
3857 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3858 }
3859 let gate = if gate_q8 {
3860 let (zq, zd) = tok_q8.as_ref().unwrap();
3861 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
3862 } else {
3863 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
3864 };
3865 let up = if up_q8 {
3866 let (zq, zd) = tok_q8.as_ref().unwrap();
3867 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
3868 } else {
3869 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
3870 };
3871 let mut act = e.uninit(n_ff_exp)?;
3872 Self::ffn_act_lim(
3873 e,
3874 cfg,
3875 &gate,
3876 &up,
3877 m.gate_exps.macro_scale(ex),
3878 m.up_exps.macro_scale(ex),
3879 lim_exp,
3880 &mut act,
3881 n_ff_exp,
3882 )?;
3883 let y = if down_q8 {
3884 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
3885 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
3886 } else {
3887 let actv = act.slice(0..n_ff_exp);
3888 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
3889 };
3890 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3891 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3893 } else if cache_dispatch {
3894 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
3899 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
3900 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3902 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3903 let actv = act.slice(0..n_ff_exp);
3904 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
3905 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3906 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3908 } else if cache_frozen {
3909 let gate = Self::moe_frozen_gemm(
3914 e,
3915 il,
3916 PROJ_GATE,
3917 ex,
3918 m,
3919 max_block,
3920 &zt,
3921 &mut scratch_g,
3922 g_len,
3923 )?;
3924 let up = Self::moe_frozen_gemm(
3925 e,
3926 il,
3927 PROJ_UP,
3928 ex,
3929 m,
3930 max_block,
3931 &zt,
3932 &mut scratch_u,
3933 u_len,
3934 )?;
3935 let mut act = e.uninit(n_ff_exp)?;
3936 Self::ffn_act_lim(
3937 e,
3938 cfg,
3939 &gate,
3940 &up,
3941 m.gate_exps.macro_scale(ex),
3942 m.up_exps.macro_scale(ex),
3943 lim_exp,
3944 &mut act,
3945 n_ff_exp,
3946 )?;
3947 let actv = act.slice(0..n_ff_exp);
3948 let y = Self::moe_frozen_gemm(
3949 e,
3950 il,
3951 PROJ_DOWN,
3952 ex,
3953 m,
3954 max_block,
3955 &actv,
3956 &mut scratch_d,
3957 d_len,
3958 )?;
3959 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3960 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3961 } else {
3962 if scratch_g.is_none() {
3966 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
3967 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
3968 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
3969 }
3970 let (sg, su, sd) = (scratch_g.as_mut().unwrap(), scratch_u.as_mut().unwrap(),
3971 scratch_d.as_mut().unwrap());
3972 let gl = m.gate_exps.expert_layout(ex);
3973 let ul = m.up_exps.expert_layout(ex);
3974 let dl = m.down_exps.expert_layout(ex);
3975 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
3976 let gate = e.qmatvec_view(sg, 0..gl.len, &zt, 1,
3977 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
3978
3979 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
3980 let up = e.qmatvec_view(su, 0..ul.len, &zt, 1,
3981 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
3982
3983 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3985 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3986
3987 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
3988 let actv = act.slice(0..n_ff_exp);
3989 let y = e.qmatvec_view(sd, 0..dl.len, &actv, 1,
3990 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?;
3991
3992 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3993 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3994 }
3995 }
3996 if let Some(worker) = cpu_worker {
3997 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
3998 let cpu_output = e.htod(&cpu_output)?;
3999 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4000 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
4001 }
4002 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
4003 for (j, &ex) in sel.iter().enumerate() {
4004 if cpu_mask[j] {
4005 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
4006 }
4007 }
4008 }
4009 }
4010
4011 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4016 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4017 {
4018 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
4027 let (sg_gate, sg_up) = if t == 1 {
4028 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
4029 Some(pair) => pair,
4030 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
4031 }
4032 } else if verify_t {
4033 (e.matmul_decode_exact(gate_shexp, z, t)?, e.matmul_decode_exact(up_shexp, z, t)?)
4034 } else {
4035 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
4037 let mut sa = e.uninit(t * n_ff_sh)?; Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp, &mut sa, t * n_ff_sh)?;
4039 let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
4040 else { e.matmul(down_shexp, &sa, t)? }; let g = match &m.gate_inp_shexp {
4054 Some(gate_inp_shexp) => {
4055 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4056 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4057 } else {
4058 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4059 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
4061 g
4062 }
4063 }
4064 None => e.htod(&vec![1.0f32; t])?,
4065 };
4066 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4068 }
4069
4070 Ok(moe_out)
4071 }
4072
4073 pub fn stage1_h2d_per_token(&self) -> u64 {
4076 use crate::hybrid::Ffn;
4077 let n_used = self.cfg.moe.as_ref().map(|m| m.expert_used_count as u64).unwrap_or(0);
4078 let mut bytes = 0u64;
4079 for l in self.layers.iter() {
4080 if let Ffn::Moe(m) = &l.ffn {
4081 bytes += n_used * (m.gate_exps.max_expert_bytes() + m.up_exps.max_expert_bytes()
4082 + m.down_exps.max_expert_bytes()) as u64;
4083 }
4084 }
4085 bytes
4086 }
4087
4088 pub(crate) fn max_moe_block(&self) -> usize {
4092 use crate::hybrid::Ffn;
4093 let mut mx = 0usize;
4094 let mut scan = |ffn: &Ffn| {
4095 if let Ffn::Moe(m) = ffn {
4096 mx = mx.max(m.gate_exps.max_expert_bytes())
4097 .max(m.up_exps.max_expert_bytes())
4098 .max(m.down_exps.max_expert_bytes());
4099 }
4100 };
4101 for l in self.layers.iter() { scan(&l.ffn); }
4102 if let Some(mtp) = self.mtp.as_ref() { scan(&mtp.ffn); }
4103 mx
4104 }
4105
4106 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
4109 use crate::hybrid::Ffn;
4110 let mut sizes = Vec::new();
4111 let mut scan = |ffn: &Ffn| {
4112 let Ffn::Moe(m) = ffn else { return };
4113 for ex in 0..m.gate_exps.n_expert {
4114 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
4115 continue;
4116 }
4117 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
4118 let len = exps.expert_layout(ex).len;
4119 if len > 0 {
4120 sizes.push(len);
4121 }
4122 }
4123 }
4124 };
4125 for layer in &self.layers {
4126 scan(&layer.ffn);
4127 }
4128 if let Some(mtp) = &self.mtp {
4129 scan(&mtp.ffn);
4130 }
4131 sizes
4132 }
4133
4134 pub fn save_cpu_expert_residency_profile(
4140 &self,
4141 e: &Engine,
4142 path: &std::path::Path,
4143 ) -> Result<(), Box<dyn std::error::Error>> {
4144 let Some(ids) = e.export_moe_residency() else {
4145 return Err("no MoE residency cache to persist".into());
4146 };
4147 let mut body = format!(
4148 "memra-freeze-profile v1 max_block={} blocks={}\n",
4149 self.max_moe_block(),
4150 ids.len()
4151 );
4152 for (layer, proj, ex) in &ids {
4153 body.push_str(&format!("{layer} {proj} {ex}\n"));
4154 }
4155 let tmp = path.with_extension("tmp");
4156 std::fs::write(&tmp, body)?;
4157 std::fs::rename(&tmp, path)?;
4158 println!(
4159 "[moe-cache] freeze profile saved: {} blocks -> {}",
4160 ids.len(),
4161 path.display()
4162 );
4163 Ok(())
4164 }
4165
4166 pub fn restore_cpu_expert_residency_profile(
4170 &self,
4171 e: &Engine,
4172 path: &std::path::Path,
4173 ) -> Result<bool, Box<dyn std::error::Error>> {
4174 use crate::hybrid::Ffn;
4175 use crate::moe_cache::BlockId;
4176 let Ok(content) = std::fs::read_to_string(path) else {
4177 return Ok(false);
4178 };
4179 let mut lines = content.lines();
4180 let Some(header) = lines.next() else { return Ok(false) };
4181 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
4182 if !header.starts_with(&expected) {
4183 println!(
4184 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
4185 path.display()
4186 );
4187 return Ok(false);
4188 }
4189 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
4190 std::collections::HashMap::new();
4191 for line in lines {
4192 let mut fields = line.split_whitespace();
4193 let (Some(layer), Some(proj), Some(ex)) =
4194 (fields.next(), fields.next(), fields.next())
4195 else {
4196 continue;
4197 };
4198 let (Ok(layer), Ok(proj), Ok(ex)) =
4199 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
4200 else {
4201 continue;
4202 };
4203 by_layer
4204 .entry(layer)
4205 .or_default()
4206 .push(BlockId::new(layer, proj, ex));
4207 }
4208 let requested: usize = by_layer.values().map(Vec::len).sum();
4209 if requested == 0 {
4210 return Ok(false);
4211 }
4212 let max_block = self.max_moe_block();
4213 let mut restaged = 0usize;
4214 let mut stage_layer = |layer_index: u16,
4215 ffn: &Ffn|
4216 -> Result<(), Box<dyn std::error::Error>> {
4217 let Ffn::Moe(m) = ffn else { return Ok(()) };
4218 let Some(ids) = by_layer.get(&layer_index) else {
4219 return Ok(());
4220 };
4221 e.with_moe_cache(max_block, |cache, eng| {
4222 for id in ids {
4223 if cache.restage_block(*id, m, eng)? {
4224 restaged += 1;
4225 }
4226 }
4227 Ok(())
4228 })
4229 };
4230 for (index, layer) in self.layers.iter().enumerate() {
4231 stage_layer(index as u16, &layer.ffn)?;
4232 }
4233 if let Some(mtp) = self.mtp.as_ref() {
4234 stage_layer(u16::MAX, &mtp.ffn)?;
4235 }
4236 e.freeze_moe_cache();
4237 println!(
4238 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
4239 path.display()
4240 );
4241 Ok(true)
4242 }
4243
4244 pub fn freeze_cpu_expert_residency(
4246 &self,
4247 e: &Engine,
4248 ) -> Result<(), Box<dyn std::error::Error>> {
4249 e.freeze_moe_cache();
4250 Ok(())
4251 }
4252
4253 pub fn ffn_act(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4261 act: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
4262 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
4263 }
4264
4265 #[allow(clippy::too_many_arguments)]
4269 pub(crate) fn ffn_act_scaled(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4270 gs: f32, us: f32, act: &mut CudaSlice<f32>, n: usize)
4271 -> Result<(), Box<dyn std::error::Error>> {
4272 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
4273 }
4274
4275 #[allow(clippy::too_many_arguments)]
4284 pub(crate) fn ffn_act_lim(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4285 gs: f32, us: f32, limit: Option<f32>, act: &mut CudaSlice<f32>, n: usize)
4286 -> Result<(), Box<dyn std::error::Error>> {
4287 if let Some(m3) = cfg.m3.as_ref() {
4288 debug_assert!(limit.is_none(), "m3 swigluoai and step35 clamp are different archs");
4289 return e.swigluoai_mul_scaled(gate, up, gs, us, m3.swiglu_alpha, m3.swiglu_limit, act, n);
4290 }
4291 if let Some(l) = limit {
4292 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
4293 }
4294 if gs == 1.0 && us == 1.0 { return e.silu_mul(gate, up, act, n); }
4295 e.silu_mul_scaled(gate, up, gs, us, act, n)
4296 }
4297
4298 fn moe_route(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
4304 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4305 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None, None, None)
4306 }
4307
4308 fn moe_route_cfg(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize,
4316 bias: Option<&[f32]>, sig: Option<(f32, bool)>, active: Option<&[bool]>)
4317 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4318 if let Some((sf, route_norm)) = sig {
4319 let lg = e.dtoh(logits)?;
4321 return Self::moe_route_sigmoid_host(
4322 &lg, t, n_expert, n_used, bias, sf, route_norm, active,
4323 );
4324 }
4325 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
4329 return e.moe_router_topk_host(logits, t, n_expert, n_used);
4330 }
4331 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
4334 let mut w_out = vec![0f32; t * n_used];
4335 for tok in 0..t {
4336 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
4337 let maxl = row.iter().enumerate()
4339 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
4340 .map(|(_, &x)| x).fold(f32::NEG_INFINITY, f32::max);
4341 let mut probs = vec![0f32; n_expert];
4342 let mut den = 0f32;
4343 for i in 0..n_expert {
4344 if active.is_some_and(|mask| !mask[i]) { continue; }
4345 let x = (row[i] - maxl).exp(); probs[i] = x; den += x;
4346 }
4347 for p in probs.iter_mut() { *p /= den; }
4348 let mut idx: Vec<usize> = (0..n_expert)
4350 .filter(|&i| active.is_none_or(|mask| mask[i])).collect();
4351 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
4352 let sl = &idx[..n_used];
4353 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
4354 let mut ws: f32 = wv.iter().sum();
4355 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() { *x /= ws; }
4357 for j in 0..n_used {
4358 sel[tok * n_used + j] = sl[j] as u32;
4359 w_out[tok * n_used + j] = wv[j];
4360 }
4361 }
4362 Ok((sel, w_out))
4363 }
4364
4365 #[allow(clippy::too_many_arguments)]
4366 fn moe_route_sigmoid_with_input(
4367 e: &Engine,
4368 logits: &CudaSlice<f32>,
4369 input: &CudaSlice<f32>,
4370 t: usize,
4371 n_expert: usize,
4372 n_used: usize,
4373 bias: Option<&[f32]>,
4374 (sf, route_norm): (f32, bool),
4375 active: Option<&[bool]>,
4376 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
4377 let (lg, input) = e.dtoh_pair(logits, input)?;
4378 let (sel, w) =
4379 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
4380 Ok((sel, w, input))
4381 }
4382
4383 pub fn start_moe_prefetch_predictor(
4388 &self,
4389 e: &Engine,
4390 cfg: &ModelConfig,
4391 ) -> Result<(), Box<dyn std::error::Error>> {
4392 use crate::hybrid::Ffn;
4393 let Some(sig) = cfg.sigmoid_router() else {
4394 return Err("prefetch predictor requires a sigmoid-router arch".into());
4395 };
4396 let resident: std::collections::HashSet<(u16, u8, u16)> = e
4397 .export_moe_residency()
4398 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
4399 .into_iter()
4400 .collect();
4401 let mut layers = Vec::new();
4402 for (index, layer) in self.layers.iter().enumerate() {
4403 let Ffn::Moe(m) = &layer.ffn else { continue };
4404 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else { continue };
4405 let router = e.dtoh(data)?;
4406 let n_expert = m.gate_exps.n_expert;
4407 let n_embd = m.gate_exps.in_f;
4408 if router.len() != n_embd * n_expert {
4409 continue;
4410 }
4411 let build = |exps: &crate::model::HostExps| {
4412 (0..n_expert)
4413 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
4414 .collect::<Vec<_>>()
4415 };
4416 layers.push((index as u16, crate::cpu_experts::PredictLayerInit {
4417 router,
4418 bias: m.exp_probs_b.clone(),
4419 active: m.active_experts.clone(),
4420 n_embd,
4421 n_used: cfg
4422 .moe
4423 .as_ref()
4424 .map(|moe| moe.expert_used_count as usize)
4425 .ok_or("prefetch predictor requires MoE config")?,
4426 sig,
4427 weights_n_expert: n_expert,
4428 gate: build(&m.gate_exps),
4429 up: build(&m.up_exps),
4430 down: build(&m.down_exps),
4431 }));
4432 }
4433 crate::cpu_experts::start_prefetch_predictor(layers, resident)
4434 .map_err(|error| error.into())
4435 }
4436
4437 #[allow(clippy::too_many_arguments)]
4440 pub(crate) fn moe_route_sigmoid_host_public(
4441 logits: &[f32],
4442 t: usize,
4443 n_expert: usize,
4444 n_used: usize,
4445 bias: Option<&[f32]>,
4446 sf: f32,
4447 route_norm: bool,
4448 active: Option<&[bool]>,
4449 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4450 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
4451 }
4452
4453 #[allow(clippy::too_many_arguments)]
4454 fn moe_route_sigmoid_host(
4455 lg: &[f32],
4456 t: usize,
4457 n_expert: usize,
4458 n_used: usize,
4459 bias: Option<&[f32]>,
4460 sf: f32,
4461 route_norm: bool,
4462 active: Option<&[bool]>,
4463 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4464 if lg.len() != t * n_expert {
4465 return Err(format!(
4466 "sigmoid router logits length mismatch: got {}, expected {}",
4467 lg.len(),
4468 t * n_expert,
4469 )
4470 .into());
4471 }
4472 let mut sel = vec![0u32; t * n_used];
4473 let mut w_out = vec![0f32; t * n_used];
4474 for tok in 0..t {
4475 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
4476 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
4477 let selsc: Vec<f32> = match bias {
4479 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
4480 None => scores.clone(),
4481 };
4482 let mut idx: Vec<usize> = (0..n_expert)
4483 .filter(|&i| active.is_none_or(|mask| mask[i]))
4484 .collect();
4485 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
4486 let sl = &idx[..n_used];
4487 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
4488 if route_norm {
4489 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
4490 for x in wv.iter_mut() {
4491 *x = *x / ws * sf;
4492 }
4493 } else {
4494 for x in wv.iter_mut() {
4495 *x *= sf;
4496 }
4497 }
4498 for j in 0..n_used {
4499 sel[tok * n_used + j] = sl[j] as u32;
4500 w_out[tok * n_used + j] = wv[j];
4501 }
4502 }
4503 Ok((sel, w_out))
4504 }
4505
4506 fn moe_ffn_pairs(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, logits: &CudaSlice<f32>,
4515 t: usize, cfg: &ModelConfig)
4516 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4517 let moe = cfg.moe.as_ref().unwrap();
4518 let n_embd = cfg.n_embd as usize;
4519 let n_expert = moe.expert_count as usize;
4520 let n_used = moe.expert_used_count as usize;
4521 let n_ff_exp = moe.expert_ff_length as usize;
4522 debug_assert!(!cfg.swiglu_clamped_anywhere(),
4527 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU");
4528 let dev = m.dev_exps.as_ref().unwrap();
4529 let (rbg_d, rbu_d) = if dev.gu_il {
4531 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
4532 } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
4533
4534 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
4535 let n_pairs = t * n_used;
4536 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
4539 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
4540 let pair_w: Vec<f32> = w_all.clone();
4541 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
4542 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
4543 let pt = e.htod_i32(&pair_tok)?;
4544 let px = e.htod_i32(&pair_ex)?;
4545 let pw = e.htod(&pair_w)?;
4546 let toff = e.htod_i32(&tok_off)?;
4547 let tids = e.htod_i32(&tok_ids)?;
4548
4549 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
4553 for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
4554 let mut ex_ids: Vec<i32> = Vec::new();
4555 let mut ex_off: Vec<i32> = vec![0];
4556 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
4557 for (ex, list) in by_ex.iter().enumerate() {
4558 if list.is_empty() { continue; }
4559 ex_ids.push(ex as i32);
4560 ex_pairs.extend_from_slice(list);
4561 ex_off.push(ex_pairs.len() as i32);
4562 }
4563 let n_active = ex_ids.len();
4564 let exi = e.htod_i32(&ex_ids)?;
4565 let exo = e.htod_i32(&ex_off)?;
4566 let exp_d = e.htod_i32(&ex_pairs)?;
4567 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4588 let mma_t = *MMA_T.get_or_init(|| {
4589 std::env::var("MEMRA_MOE_MMA_T").ok().and_then(|v| v.parse().ok()).unwrap_or(16)
4590 });
4591 let use_mma = std::env::var("MEMRA_MOE_MMA").map(|v| v != "0").unwrap_or(true)
4592 && t >= mma_t
4593 && q8_expert_dec_supported(m.gate_exps.qtype) && q8_expert_dec_supported(m.up_exps.qtype)
4594 && q8_expert_dec_supported(m.down_exps.qtype)
4595 && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
4596 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
4612 && q8_expert_dec_supported(m.up_exps.qtype)
4613 && q8_expert_dec_supported(m.down_exps.qtype)
4614 && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
4615 let f16g_mode = crate::moe_f16g_mode();
4616 let f16g = f16g_mode != 0 && t >= mma_t
4617 && (f16g_mode != 3 || !mma_capable)
4618 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
4619 && f16g_proj_ok(m.up_exps.qtype, n_embd)
4620 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
4621 if use_mma || f16g {
4622 let y_down = if f16g {
4630 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
4634 let csr_tok_d = e.htod_i32(&csr_tok)?;
4635 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
4636 let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
4637 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4638 m.gate_exps.qtype, rbg_d)?;
4639 let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
4640 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4641 m.up_exps.qtype, rbu_d)?;
4642 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
4643 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
4644 let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
4645 &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
4646 m.down_exps.qtype, m.down_exps.row_bytes)?;
4647 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
4648 } else {
4649 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
4651 let gate = e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4652 n_embd, n_ff_exp, n_active, n_pairs, t,
4653 m.gate_exps.qtype, rbg_d)?;
4654 let up = e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4655 n_embd, n_ff_exp, n_active, n_pairs, t,
4656 m.up_exps.qtype, rbu_d)?;
4657 let a_scr = if crate::moe_fuse_actq_on() {
4663 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
4664 } else {
4665 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
4666 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
4667 };
4668 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
4669 let pself = e.htod_i32(&pair_self)?;
4670 e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
4671 n_ff_exp, n_embd, n_active, n_pairs, n_pairs,
4672 m.down_exps.qtype, m.down_exps.row_bytes)?
4673 };
4674 let mut moe_out = e.uninit(t * n_embd)?;
4675 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
4676 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4677 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4678 {
4679 let n_ff_sh = gate_shexp.out_features();
4680 let sg_gate = e.matmul(gate_shexp, z, t)?;
4681 let sg_up = e.matmul(up_shexp, z, t)?;
4682 let mut sa = e.uninit(t * n_ff_sh)?;
4683 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
4684 let sh = e.matmul(down_shexp, &sa, t)?;
4685 let g = match &m.gate_inp_shexp {
4691 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
4692 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4693 }
4694 Some(gate_inp_shexp) => {
4695 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4696 let mut g = e.uninit(t)?;
4697 e.sigmoid(&gs, &mut g, t)?;
4698 g
4699 }
4700 None => e.htod(&vec![1.0f32; t])?,
4701 };
4702 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4703 }
4704 return Ok(moe_out);
4705 }
4706
4707 let dec = std::env::var("MEMRA_MOE_DEC").map(|v| v != "0").unwrap_or(true);
4710 let matvec = |proj, exi: &_, exo: &_, exp_d: &_, pt: &_, aq: &_, ad: &_,
4711 inf, outf, qtype, rb| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4712 let dec = dec && q8_expert_dec_supported(qtype);
4714 if dec { e.moe_pairs_matvec_q8_dec(&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
4715 inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
4716 else { e.moe_pairs_matvec_q8_em (&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
4717 inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
4718 };
4719 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4720 let gate = matvec(0, &exi, &exo, &exp_d, &pt, &zq, &zd,
4721 n_embd, n_ff_exp, m.gate_exps.qtype, rbg_d)?;
4722 let up = matvec(1, &exi, &exo, &exp_d, &pt, &zq, &zd,
4723 n_embd, n_ff_exp, m.up_exps.qtype, rbu_d)?;
4724 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
4725 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4726 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
4728 let pself = e.htod_i32(&pair_self)?;
4729 let y_down = matvec(2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
4730 n_ff_exp, n_embd, m.down_exps.qtype, m.down_exps.row_bytes)?;
4731 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
4733
4734 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4738 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4739 {
4740 let n_ff_sh = gate_shexp.out_features();
4741 let sg_gate = e.matmul(gate_shexp, z, t)?;
4742 let sg_up = e.matmul(up_shexp, z, t)?;
4743 let mut sa = e.uninit(t * n_ff_sh)?;
4744 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
4745 let sh = e.matmul(down_shexp, &sa, t)?;
4746 let g = match &m.gate_inp_shexp {
4751 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
4752 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4753 }
4754 Some(gate_inp_shexp) => {
4755 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4756 let mut g = e.uninit(t)?;
4757 e.sigmoid(&gs, &mut g, t)?;
4758 g
4759 }
4760 None => e.htod(&vec![1.0f32; t])?,
4761 };
4762 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4763 }
4764 Ok(moe_out)
4765 }
4766
4767 #[allow(clippy::too_many_arguments)]
4769 #[allow(clippy::too_many_arguments)]
4770 fn moe_ffn_dev(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
4771 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, logits: &CudaSlice<f32>,
4772 t: usize, cfg: &ModelConfig, il: u16, max_block: usize)
4773 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4774 let moe = cfg.moe.as_ref().unwrap();
4775 let n_embd = cfg.n_embd as usize;
4776 let n_expert = moe.expert_count as usize;
4777 let n_used = moe.expert_used_count as usize;
4778 let n_ff_exp = moe.expert_ff_length as usize;
4779 debug_assert!(cfg.sigmoid_router().is_none(),
4783 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts");
4784 debug_assert!(!cfg.swiglu_clamped_at(il as u32),
4785 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form");
4786
4787 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
4789 if m.has_macros {
4792 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
4793 }
4794
4795 let mut moe_out = e.uninit(t * n_embd)?;
4797
4798 if let Some(dev) = m.dev_exps.as_ref() {
4801 let (rbg_d, rbu_d) = if dev.gu_il {
4804 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
4805 } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
4806 let q8 = moe_q8_enabled()
4807 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
4808 && q8_expert_supported(m.down_exps.qtype);
4809 let rows_arm = q8 && t > 1 && crate::spec::spec_m2()
4818 && n_ff_exp == 512 && n_used <= 8
4819 && std::env::var("MEMRA_MOE_DEVQ8_GU").map(|v| v.is_empty() || v == "v").unwrap_or(true)
4820 && std::env::var("MEMRA_MOE_DEVQ8_DOWN").map(|v| v.is_empty() || v == "w8h2v").unwrap_or(true);
4821 let csr_mode = std::env::var("MEMRA_MOE_CSR").ok()
4830 .and_then(|v| v.parse::<i32>().ok()).unwrap_or(1);
4831 let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
4832 let csr_arm = rows_arm && csr_mode > 0 && t <= 10
4833 && csr_qt(m.gate_exps.qtype) && csr_qt(m.up_exps.qtype)
4834 && csr_qt(m.down_exps.qtype);
4835 if csr_arm {
4836 if csr_mode == 2 {
4837 static ENGAGED: std::sync::Once = std::sync::Once::new();
4838 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
4839 }
4840 let n_pairs = t * n_used;
4841 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4842 let act = e.moe_gate_up_silu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, n_pairs,
4843 n_embd, n_ff_exp, n_used, n_expert,
4844 m.gate_exps.qtype, m.up_exps.qtype,
4845 rbg_d, rbu_d)?;
4846 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4847 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
4851 t, n_ff_exp, n_embd, n_used, n_expert,
4852 m.down_exps.qtype, m.down_exps.row_bytes)?;
4853 if csr_mode == 2 {
4854 let act_r = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
4856 n_embd, n_ff_exp, n_used, n_expert,
4857 m.gate_exps.qtype, m.up_exps.qtype,
4858 rbg_d, rbu_d, &m.dev_macros)?;
4859 let mut out_r = e.uninit(t * n_embd)?;
4860 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
4861 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2r, &ad2r, &mut out_r,
4862 t, n_ff_exp, n_embd, n_used, n_expert,
4863 m.down_exps.qtype, m.down_exps.row_bytes)?;
4864 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
4865 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
4866 let ba = a1.iter().zip(&a2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
4867 let bo = o1.iter().zip(&o2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
4868 if ba + bo > 0 {
4869 eprintln!("[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
4870 a1.len(), o1.len());
4871 let sel_h = e.dtoh_i32(&sel_d)?;
4873 let mut shown = 0;
4874 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
4875 if x.to_bits() != y.to_bits() && shown < 4 {
4876 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
4877 let ex = sel_h[p];
4878 let npx = sel_h.iter().filter(|&&v| v == ex).count();
4879 eprintln!(" ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}");
4880 shown += 1;
4881 }
4882 }
4883 std::process::exit(3);
4884 }
4885 }
4886 } else if rows_arm {
4887 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
4890 use std::sync::atomic::{AtomicU64, Ordering};
4891 static PAIRS: AtomicU64 = AtomicU64::new(0);
4892 static UNIQ: AtomicU64 = AtomicU64::new(0);
4893 static CALLS: AtomicU64 = AtomicU64::new(0);
4894 let sel_h = e.dtoh_i32(&sel_d)?;
4895 let mut u: Vec<i32> = sel_h.clone(); u.sort_unstable(); u.dedup();
4896 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
4897 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
4898 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
4899 if c % 480 == 0 {
4900 let p = PAIRS.load(Ordering::Relaxed); let q = UNIQ.load(Ordering::Relaxed);
4901 eprintln!("[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
4902 q as f64 / p as f64);
4903 }
4904 }
4905 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4906 let act = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
4907 n_embd, n_ff_exp, n_used, n_expert,
4908 m.gate_exps.qtype, m.up_exps.qtype,
4909 rbg_d, rbu_d, &m.dev_macros)?;
4910 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
4911 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
4912 t, n_ff_exp, n_embd, n_used, n_expert,
4913 m.down_exps.qtype, m.down_exps.row_bytes)?;
4914 } else {
4915 for tok in 0..t {
4916 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
4917 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
4918 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
4919 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4920 if q8 {
4921 let (zq, zd) = match (t, zq8) {
4922 (1, Some((q, d))) => (q.clone(), d.clone()),
4923 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
4924 };
4925 let act = e.moe_gate_up_silu8_dev_q8(&dev.ptr_row, &selt, &zq, &zd,
4926 n_embd, n_ff_exp, n_used, n_expert,
4927 m.gate_exps.qtype, m.up_exps.qtype,
4928 rbg_d, rbu_d, &m.dev_macros)?;
4929 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4930 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selt, &wt, &aq2, &ad2, &mut dst,
4931 n_ff_exp, n_embd, n_used, n_expert,
4932 m.down_exps.qtype, m.down_exps.row_bytes)?;
4933 } else {
4934 let act = e.moe_gate_up_silu8_dev(&dev.ptr_row, &selt, &zt, n_embd, n_ff_exp,
4935 n_used, n_expert,
4936 m.gate_exps.qtype, m.up_exps.qtype,
4937 rbg_d, rbu_d, &m.dev_macros)?;
4938 e.moe_down8_fma_dev(&dev.ptr_row, &selt, &wt, &act, &mut dst,
4939 n_ff_exp, n_embd, n_used, n_expert,
4940 m.down_exps.qtype, m.down_exps.row_bytes)?;
4941 }
4942 }
4943 }
4944 } else {
4945 let q8 = moe_q8_enabled()
4952 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
4953 && q8_expert_supported(m.down_exps.qtype);
4954 e.with_moe_cache(max_block, |c, eng| {
4955 let row = c.layer_dev_row(il, n_expert, eng)?
4956 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
4957 for tok in 0..t {
4958 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
4959 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
4960 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
4961 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4962 if q8 {
4963 let (zq, zd) = match (t, zq8) {
4964 (1, Some((q, d))) => (q.clone(), d.clone()),
4965 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
4966 };
4967 let act = eng.moe_gate_up_silu8_dev_q8(row, &selt, &zq, &zd,
4968 n_embd, n_ff_exp, n_used, n_expert,
4969 m.gate_exps.qtype, m.up_exps.qtype,
4970 m.gate_exps.row_bytes, m.up_exps.row_bytes,
4971 &m.dev_macros)?;
4972 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
4973 eng.moe_down8_fma_dev_q8(row, &selt, &wt, &aq2, &ad2, &mut dst,
4974 n_ff_exp, n_embd, n_used, n_expert,
4975 m.down_exps.qtype, m.down_exps.row_bytes)?;
4976 } else {
4977 let act = eng.moe_gate_up_silu8_dev(row, &selt, &zt, n_embd, n_ff_exp,
4978 n_used, n_expert,
4979 m.gate_exps.qtype, m.up_exps.qtype,
4980 m.gate_exps.row_bytes, m.up_exps.row_bytes,
4981 &m.dev_macros)?;
4982 eng.moe_down8_fma_dev(row, &selt, &wt, &act, &mut dst,
4983 n_ff_exp, n_embd, n_used, n_expert,
4984 m.down_exps.qtype, m.down_exps.row_bytes)?;
4985 }
4986 }
4987 c.hits += (t * 3 * n_used) as u64;
4989 Ok(())
4990 })?;
4991 }
4992
4993 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4998 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4999 {
5000 let n_ff_sh = gate_shexp.out_features();
5001 let verify_t = t > 1 && t < PRIME_MIN_T;
5004 let (sg_gate, sg_up) = if t == 1 {
5005 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5006 Some(pair) => pair,
5007 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5008 }
5009 } else if verify_t {
5010 let mut fused = None;
5014 if crate::spec::spec_fused_t() && (2..=4).contains(&t)
5015 && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp) {
5016 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5017 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
5018 }
5019 match fused {
5020 Some(pair) => pair,
5021 None => (e.matmul_decode_exact(gate_shexp, z, t)?,
5022 e.matmul_decode_exact(up_shexp, z, t)?),
5023 }
5024 } else {
5025 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
5026 };
5027 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5029 let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
5030 else { e.matmul(down_shexp, &sa, t)? };
5031 let g = match &m.gate_inp_shexp {
5035 Some(gate_inp_shexp) => {
5036 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5039 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5040 } else {
5041 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5042 let mut g = e.uninit(t)?;
5043 e.sigmoid(&gs, &mut g, t)?;
5044 g
5045 }
5046 }
5047 None => e.htod(&vec![1.0f32; t])?,
5048 };
5049 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5050 }
5051
5052 Ok(moe_out)
5053 }
5054
5055 #[allow(clippy::too_many_arguments)]
5065 #[allow(clippy::too_many_arguments)]
5068 fn moe_gdec_token_q8(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
5069 zq: &CudaSlice<i8>, zd: &CudaSlice<f32>, sel: &[u32], w: &[f32],
5070 moe_out: &mut CudaSlice<f32>, tok: usize,
5071 n_embd: usize, n_ff_exp: usize, n_used: usize)
5072 -> Result<bool, Box<dyn std::error::Error>> {
5073 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
5074 use cudarc::driver::DevicePtr;
5075 let ptrs = e.with_moe_cache(max_block, |c, eng| {
5076 let mut g = [0u64; 8];
5077 let mut u = [0u64; 8];
5078 let mut d = [0u64; 8];
5079 for (j, &ex) in sel.iter().enumerate() {
5080 let ex = ex as u16;
5081 let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
5082 c.resident(BlockId::new(il, PROJ_UP, ex)),
5083 c.resident(BlockId::new(il, PROJ_DOWN, ex)))
5084 else { return Ok(None); };
5085 let __s = eng.stream();
5086 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
5087 let (pu, _e1) = c.slot(su).device_ptr(&__s);
5088 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
5089 g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
5090 }
5091 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
5092 for &ex in sel {
5093 let ex = ex as u16;
5094 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
5095 c.note_profile_hit(BlockId::new(il, proj, ex));
5096 }
5097 }
5098 }
5099 c.hits += (3 * n_used) as u64;
5100 Ok(Some((g, u, d)))
5101 })?;
5102 let Some((g, u, d)) = ptrs else { return Ok(false) };
5103 let mut wv = [0f32; 8];
5104 wv[..n_used].copy_from_slice(w);
5105 let act = e.moe_gate_up_silu8_q8(crate::WPtr8(g), crate::WPtr8(u), zq, zd,
5106 n_embd, n_ff_exp, n_used,
5107 m.gate_exps.qtype, m.up_exps.qtype,
5108 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5109 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5111 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5112 e.moe_down8_fma_q8(crate::WPtr8(d), crate::F32x8(wv), &aq2, &ad2, &mut dst,
5113 n_ff_exp, n_embd, n_used,
5114 m.down_exps.qtype, m.down_exps.row_bytes)?;
5115 Ok(true)
5116 }
5117
5118 fn moe_gdec_token(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
5119 zt: &cudarc::driver::CudaView<f32>, sel: &[u32], w: &[f32],
5120 moe_out: &mut CudaSlice<f32>, tok: usize,
5121 n_embd: usize, n_ff_exp: usize, n_used: usize)
5122 -> Result<bool, Box<dyn std::error::Error>> {
5123 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
5124 use cudarc::driver::DevicePtr;
5125 let ptrs = e.with_moe_cache(max_block, |c, eng| {
5127 let mut g = [0u64; 8];
5128 let mut u = [0u64; 8];
5129 let mut d = [0u64; 8];
5130 for (j, &ex) in sel.iter().enumerate() {
5131 let ex = ex as u16;
5132 let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
5133 c.resident(BlockId::new(il, PROJ_UP, ex)),
5134 c.resident(BlockId::new(il, PROJ_DOWN, ex)))
5135 else { return Ok(None); };
5136 let __s = eng.stream();
5137 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
5138 let (pu, _e1) = c.slot(su).device_ptr(&__s);
5139 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
5140 g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
5141 }
5142 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
5143 for &ex in sel {
5144 let ex = ex as u16;
5145 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
5146 c.note_profile_hit(BlockId::new(il, proj, ex));
5147 }
5148 }
5149 }
5150 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
5152 })?;
5153 let Some((g, u, d)) = ptrs else { return Ok(false) };
5154 let mut wv = [0f32; 8];
5155 wv[..n_used].copy_from_slice(w);
5156 let act = e.moe_gate_up_silu8(crate::WPtr8(g), crate::WPtr8(u), zt,
5158 n_embd, n_ff_exp, n_used,
5159 m.gate_exps.qtype, m.up_exps.qtype,
5160 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5161 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5162 e.moe_down8_fma_into(crate::WPtr8(d), crate::F32x8(wv), &act, &mut dst,
5163 n_ff_exp, n_embd, n_used,
5164 m.down_exps.qtype, m.down_exps.row_bytes)?;
5165 Ok(true)
5166 }
5167
5168 fn moe_cached_gemm_q8(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
5173 max_block: usize, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
5174 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5175 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
5176 let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
5177 let layout = exps.expert_layout(ex);
5178 let id = BlockId::new(il, proj, ex as u16);
5179 let source = exps.expert_source(ex);
5180 e.with_moe_cache(max_block, |c, eng| {
5181 let slot = c.dispatch_source(id, source, eng)?;
5182 let DispatchSlot::Resident(sl) = slot;
5183 let buf = c.slot(sl);
5184 eng.qmatvec_expert_q8(buf, 0..layout.len, aq, ad, 1, exps.in_f, exps.out_f,
5185 layout.qtype, layout.row_bytes)
5186 })
5187 }
5188
5189 fn moe_cached_gemm(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
5190 max_block: usize, x: &cudarc::driver::CudaView<f32>)
5191 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5192 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
5193 let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
5194 let layout = exps.expert_layout(ex);
5195 let id = BlockId::new(il, proj, ex as u16);
5196 let source = exps.expert_source(ex);
5197 e.with_moe_cache(max_block, |c, eng| {
5199 let slot = c.dispatch_source(id, source, eng)?;
5200 let DispatchSlot::Resident(sl) = slot;
5203 let buf = c.slot(sl);
5204 eng.qmatvec_view(buf, 0..layout.len, x, 1, exps.in_f, exps.out_f,
5205 layout.qtype, layout.row_bytes)
5206 })
5207 }
5208
5209 fn moe_profile_admit_expert(
5213 e: &Engine,
5214 il: u16,
5215 ex: usize,
5216 m: &MoeWeights,
5217 max_block: usize,
5218 ) -> Result<(), Box<dyn std::error::Error>> {
5219 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5220 e.with_moe_cache(max_block, |cache, eng| {
5221 for (proj, exps) in [
5222 (PROJ_GATE, &m.gate_exps),
5223 (PROJ_UP, &m.up_exps),
5224 (PROJ_DOWN, &m.down_exps),
5225 ] {
5226 let id = BlockId::new(il, proj, ex as u16);
5227 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
5228 }
5229 Ok(())
5230 })
5231 }
5232
5233 #[allow(clippy::too_many_arguments)]
5236 fn moe_frozen_gemm(
5237 e: &Engine,
5238 il: u16,
5239 proj: u8,
5240 ex: usize,
5241 m: &MoeWeights,
5242 max_block: usize,
5243 x: &cudarc::driver::CudaView<f32>,
5244 scratch: &mut Option<CudaSlice<u8>>,
5245 scratch_len: usize,
5246 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5247 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
5248 let exps = match proj {
5249 PROJ_GATE => &m.gate_exps,
5250 PROJ_UP => &m.up_exps,
5251 _ => &m.down_exps,
5252 };
5253 let layout = exps.expert_layout(ex);
5254 let id = BlockId::new(il, proj, ex as u16);
5255 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
5256 let Some(slot) = cache.resident(id) else {
5257 return Ok(None);
5258 };
5259 let buf = cache.slot(slot);
5260 Ok(Some(eng.qmatvec_view(
5261 buf,
5262 0..layout.len,
5263 x,
5264 1,
5265 exps.in_f,
5266 exps.out_f,
5267 layout.qtype,
5268 layout.row_bytes,
5269 )?))
5270 })? {
5271 return Ok(output);
5272 }
5273 if scratch.is_none() {
5274 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
5275 }
5276 let scratch = scratch.as_mut().unwrap();
5277 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
5278 e.qmatvec_view(
5279 scratch,
5280 0..layout.len,
5281 x,
5282 1,
5283 exps.in_f,
5284 exps.out_f,
5285 layout.qtype,
5286 layout.row_bytes,
5287 )
5288 }
5289
5290 fn moe_prefetch_expert(
5291 e: &Engine,
5292 il: u16,
5293 ex: usize,
5294 m: &MoeWeights,
5295 max_block: usize,
5296 keep: &[crate::moe_cache::BlockId],
5297 ) -> Result<(), Box<dyn std::error::Error>> {
5298 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5299 e.with_moe_cache(max_block, |c, eng| {
5300 for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
5301 (PROJ_DOWN, &m.down_exps)] {
5302 let id = BlockId::new(il, proj, ex as u16);
5303 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
5304 }
5305 Ok(())
5306 })
5307 }
5308
5309 fn moe_prefetch_disk_expert(e: &Engine, il: u16, ex: usize, m: &MoeWeights,
5312 max_block: usize, keep: &[crate::moe_cache::BlockId])
5313 -> Result<(), Box<dyn std::error::Error>> {
5314 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5315 e.with_moe_cache(max_block, |c, eng| {
5316 for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
5317 (PROJ_DOWN, &m.down_exps)] {
5318 let source = exps.expert_source(ex);
5319 if let crate::model::ExpertSource::Disk { .. } = &source {
5320 let id = BlockId::new(il, proj, ex as u16);
5321 let _ = c.prefetch_source(id, source, keep, eng)?;
5322 }
5323 }
5324 Ok(())
5325 })
5326 }
5327
5328 #[inline]
5329 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
5330 let _ = m.gate_exps.prefetch_expert_pages(ex);
5331 let _ = m.up_exps.prefetch_expert_pages(ex);
5332 let _ = m.down_exps.prefetch_expert_pages(ex);
5333 }
5334}
5335
5336impl HybridModel {
5353 #[allow(clippy::too_many_arguments)]
5357 fn moe_ffn_grouped_resident_q8(
5358 e: &Engine,
5359 m: &MoeWeights,
5360 z: &CudaSlice<f32>,
5361 t: usize,
5362 cfg: &ModelConfig,
5363 il: u16,
5364 sel_all: &[u32],
5365 w_all: &[f32],
5366 table: &CudaSlice<u64>,
5367 gu_il: bool,
5368 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5369 let moe = cfg.moe.as_ref().unwrap();
5370 let n_embd = cfg.n_embd as usize;
5371 let n_expert = moe.expert_count as usize;
5372 let n_used = moe.expert_used_count as usize;
5373 let n_ff_exp = moe.expert_ff_length as usize;
5374 let n_pairs = t * n_used;
5375 debug_assert_eq!(sel_all.len(), n_pairs);
5376 debug_assert_eq!(w_all.len(), n_pairs);
5377 debug_assert!(
5378 m.gate_exps.macros.is_none()
5379 && m.up_exps.macros.is_none()
5380 && m.down_exps.macros.is_none(),
5381 "resident grouped q8 does not fold per-expert macro scales",
5382 );
5383
5384 if !cfg.swiglu_clamped_at(il as u32) {
5389 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
5390 let sel_d = e.htod_i32(&sel)?;
5391 let w_d = e.htod(w_all)?;
5392 let (gate_row_bytes, up_row_bytes) = if gu_il {
5393 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5394 (combined, combined)
5395 } else {
5396 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5397 };
5398 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5399 let act = e.moe_gate_up_silu8_dev_q8_rows(
5400 table,
5401 &sel_d,
5402 &zq,
5403 &zd,
5404 t,
5405 n_embd,
5406 n_ff_exp,
5407 n_used,
5408 n_expert,
5409 m.gate_exps.qtype,
5410 m.up_exps.qtype,
5411 gate_row_bytes,
5412 up_row_bytes,
5413 &m.dev_macros,
5414 )?;
5415 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5416 let mut moe_out = e.uninit(t * n_embd)?;
5417 e.moe_down8_fma_dev_q8_rows_g(
5418 table,
5419 &sel_d,
5420 &w_d,
5421 &aq2,
5422 &ad2,
5423 &mut moe_out,
5424 t,
5425 n_ff_exp,
5426 n_embd,
5427 n_used,
5428 n_expert,
5429 m.down_exps.qtype,
5430 m.down_exps.row_bytes,
5431 )?;
5432
5433 if std::env::var("MEMRA_MOE_STATS").is_ok() {
5434 let mut counts = vec![0usize; n_expert];
5435 for &expert in sel_all {
5436 counts[expert as usize] += 1;
5437 }
5438 let mut sizes: Vec<usize> =
5439 counts.into_iter().filter(|&count| count != 0).collect();
5440 sizes.sort_unstable();
5441 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
5442 println!(
5443 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
5444 m_e: min={} median={} mean={mean:.1} max={}",
5445 sizes.len(),
5446 n_expert,
5447 sizes.first().copied().unwrap_or(0),
5448 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
5449 sizes.last().copied().unwrap_or(0),
5450 );
5451 }
5452 return Ok(moe_out);
5453 }
5454
5455 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5459 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
5460 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5461 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5462
5463 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
5464 for (pair, &expert) in pair_ex.iter().enumerate() {
5465 by_expert[expert as usize].push(pair as i32);
5466 }
5467
5468 let pair_tok_d = e.htod_i32(&pair_tok)?;
5469 let pair_ex_d = e.htod_i32(&pair_ex)?;
5470 let pair_w_d = e.htod(w_all)?;
5471 let tok_off_d = e.htod_i32(&tok_off)?;
5472 let tok_ids_d = e.htod_i32(&tok_ids)?;
5473
5474 let matvec = |
5475 proj: i32,
5476 pair_rows: &CudaSlice<i32>,
5477 aq: &CudaSlice<i8>,
5478 ad: &CudaSlice<f32>,
5479 in_f: usize,
5480 out_f: usize,
5481 qtype: i32,
5482 row_bytes: usize,
5483 | -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5484 e.moe_pairs_matvec_q8(
5485 table,
5486 proj,
5487 pair_rows,
5488 &pair_ex_d,
5489 aq,
5490 ad,
5491 in_f,
5492 out_f,
5493 n_expert,
5494 n_pairs,
5495 qtype,
5496 row_bytes,
5497 )
5498 };
5499
5500 let (gate_row_bytes, up_row_bytes) = if gu_il {
5501 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5502 (combined, combined)
5503 } else {
5504 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5505 };
5506 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5507 let gate = matvec(
5508 0,
5509 &pair_tok_d,
5510 &zq,
5511 &zd,
5512 n_embd,
5513 n_ff_exp,
5514 m.gate_exps.qtype,
5515 gate_row_bytes,
5516 )?;
5517 let up = matvec(
5518 1,
5519 &pair_tok_d,
5520 &zq,
5521 &zd,
5522 n_embd,
5523 n_ff_exp,
5524 m.up_exps.qtype,
5525 up_row_bytes,
5526 )?;
5527 let mut act = e.uninit(n_pairs * n_ff_exp)?;
5528 Self::ffn_act_lim(
5529 e,
5530 cfg,
5531 &gate,
5532 &up,
5533 1.0,
5534 1.0,
5535 cfg.clamp_exp_at(il as u32),
5536 &mut act,
5537 n_pairs * n_ff_exp,
5538 )?;
5539 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5540 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5541 let pair_self_d = e.htod_i32(&pair_self)?;
5542 let down = matvec(
5543 2,
5544 &pair_self_d,
5545 &aq2,
5546 &ad2,
5547 n_ff_exp,
5548 n_embd,
5549 m.down_exps.qtype,
5550 m.down_exps.row_bytes,
5551 )?;
5552 let mut moe_out = e.uninit(t * n_embd)?;
5553 e.moe_pairs_scatter(
5554 &down,
5555 &pair_w_d,
5556 &tok_off_d,
5557 &tok_ids_d,
5558 &mut moe_out,
5559 t,
5560 n_embd,
5561 )?;
5562
5563 if std::env::var("MEMRA_MOE_STATS").is_ok() {
5564 let mut sizes: Vec<usize> = by_expert
5565 .iter()
5566 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
5567 .collect();
5568 sizes.sort_unstable();
5569 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
5570 println!(
5571 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
5572 m_e: min={} median={} mean={mean:.1} max={}",
5573 sizes.len(),
5574 n_expert,
5575 sizes.first().copied().unwrap_or(0),
5576 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
5577 sizes.last().copied().unwrap_or(0),
5578 );
5579 }
5580 Ok(moe_out)
5581 }
5582
5583 fn moe_ffn_grouped_add_shared(
5584 e: &Engine,
5585 m: &MoeWeights,
5586 z: &CudaSlice<f32>,
5587 t: usize,
5588 cfg: &ModelConfig,
5589 il: u16,
5590 moe_out: &mut CudaSlice<f32>,
5591 ) -> Result<(), Box<dyn std::error::Error>> {
5592 let n_embd = cfg.n_embd as usize;
5593 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5594 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5595 {
5596 let n_ff_sh = gate_shexp.out_features();
5597 let sg_gate = e.matmul(gate_shexp, z, t)?;
5598 let sg_up = e.matmul(up_shexp, z, t)?;
5599 let mut sa = e.uninit(t * n_ff_sh)?;
5600 Self::ffn_act_lim(
5601 e,
5602 cfg,
5603 &sg_gate,
5604 &sg_up,
5605 1.0,
5606 1.0,
5607 cfg.clamp_shexp_at(il as u32),
5608 &mut sa,
5609 t * n_ff_sh,
5610 )?;
5611 let sh = e.matmul(down_shexp, &sa, t)?;
5612 let gate = match &m.gate_inp_shexp {
5613 Some(gate_inp_shexp) => {
5614 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5615 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5616 } else {
5617 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5618 let mut gate = e.uninit(t)?;
5619 e.sigmoid(&raw, &mut gate, t)?;
5620 gate
5621 }
5622 }
5623 None => e.htod(&vec![1.0f32; t])?,
5624 };
5625 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
5626 }
5627 Ok(())
5628 }
5629
5630 pub(crate) fn moe_ffn_grouped(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
5633 cfg: &ModelConfig, il: u16, max_block: usize)
5634 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5635 let moe = cfg.moe.as_ref().unwrap();
5636 let n_embd = cfg.n_embd as usize;
5637 let n_expert = moe.expert_count as usize;
5638 let n_used = moe.expert_used_count as usize;
5639 let n_ff_exp = moe.expert_ff_length as usize;
5640 let lim_exp = cfg.clamp_exp_at(il as u32);
5642
5643 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
5646 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
5647 Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
5648 m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
5649 } else {
5650 Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
5651 None, None, m.active_experts.as_deref())?
5652 };
5653 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
5654 Self::trace_moe_input(e, il, t, n_embd, z)?;
5655
5656 let no_exp_macros = m.gate_exps.macros.is_none()
5661 && m.up_exps.macros.is_none()
5662 && m.down_exps.macros.is_none();
5663 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
5664 m.has_uniform_expert_layout()
5665 && no_exp_macros
5666 && moe_q8_enabled()
5667 && q8_expert_supported(m.gate_exps.qtype)
5668 && q8_expert_supported(m.up_exps.qtype)
5669 && q8_expert_supported(m.down_exps.qtype)
5670 && moe_slab_enabled()
5671 && dev.dev == e.ctx().ordinal()
5672 });
5673 if let Some(dev) = resident_q8 {
5674 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
5675 e,
5676 m,
5677 z,
5678 t,
5679 cfg,
5680 il,
5681 &sel_all,
5682 &w_all,
5683 &dev.ptr_row,
5684 dev.gu_il,
5685 )?;
5686 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
5687 return Ok(moe_out);
5688 }
5689
5690 struct ExpertGroup {
5694 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
5698 let mut groups: Vec<ExpertGroup> = (0..n_expert).map(|_| ExpertGroup {
5699 tok_indices: Vec::new(), slot_indices: Vec::new(), weights: Vec::new(),
5700 }).collect();
5701
5702 for tok in 0..t {
5703 for j in 0..n_used {
5704 let ex = sel_all[tok * n_used + j] as usize;
5705 let w = w_all[tok * n_used + j];
5706 groups[ex].tok_indices.push(tok as i32);
5707 groups[ex].slot_indices.push(j as i32);
5708 groups[ex].weights.push(w);
5709 }
5710 }
5711
5712 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
5715 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
5719 let u_len = m.up_exps.max_expert_bytes();
5720 let d_len = m.down_exps.max_expert_bytes();
5721 let moe_q8 = m.has_uniform_expert_layout()
5722 && moe_q8_enabled()
5723 && q8_expert_supported(m.gate_exps.qtype)
5724 && q8_expert_supported(m.up_exps.qtype)
5725 && q8_expert_supported(m.down_exps.qtype);
5726 let slab_local = m.dev_exps.as_ref().filter(|dev| {
5729 !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal()
5730 });
5731 let use_cache =
5732 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
5733 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
5736
5737 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
5739 (Some(e.alloc_u8(g_len)?), Some(e.alloc_u8(u_len)?), Some(e.alloc_u8(d_len)?))
5740 } else {
5741 (None, None, None)
5742 };
5743
5744 let mut order: Vec<usize> =
5755 (0..n_expert).filter(|&ex| !groups[ex].tok_indices.is_empty()).collect();
5756 order.sort_by(|&a, &b| groups[b].tok_indices.len()
5757 .cmp(&groups[a].tok_indices.len()).then(a.cmp(&b)));
5758 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
5760 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
5761 if worker_disk_prefetch {
5762 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
5763 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
5764 }
5765 }
5766 for (order_pos, &ex) in order.iter().enumerate() {
5767 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
5768 Self::moe_prefetch_host_expert(order[next], m);
5769 }
5770 if worker_disk_prefetch {
5771 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
5772 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5773 let keep = [
5774 BlockId::new(il, PROJ_GATE, ex as u16),
5775 BlockId::new(il, PROJ_UP, ex as u16),
5776 BlockId::new(il, PROJ_DOWN, ex as u16),
5777 ];
5778 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
5779 }
5780 }
5781 let grp = &groups[ex];
5782 let m_e = grp.tok_indices.len();
5783 m_dist.push(m_e);
5784 let gl = m.gate_exps.expert_layout(ex);
5785 let ul = m.up_exps.expert_layout(ex);
5786 let dl = m.down_exps.expert_layout(ex);
5787
5788 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
5792 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
5793 let dmac = m.down_exps.macro_scale(ex);
5794 let weight_d = if dmac == 1.0 { e.htod(&grp.weights)? } else {
5795 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
5796 e.htod(&scaled)?
5797 };
5798
5799 let mut gathered = e.zeros(m_e * n_embd)?;
5801 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
5802 let gv = gathered.slice(0..m_e * n_embd);
5803
5804 let y = if let Some(dev) = slab_local {
5807 let gate_start = ex * m.gate_exps.expert_stride;
5808 let up_start = ex * m.up_exps.expert_stride;
5809 let down_start = ex * m.down_exps.expert_stride;
5810 if grouped_q8 {
5811 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
5812 let gate = e.qmatvec_expert_q8(
5813 &dev.gate,
5814 gate_start..gate_start + gl.len,
5815 &zq,
5816 &zd,
5817 m_e,
5818 m.gate_exps.in_f,
5819 m.gate_exps.out_f,
5820 gl.qtype,
5821 gl.row_bytes,
5822 )?;
5823 let up = e.qmatvec_expert_q8(
5824 &dev.up,
5825 up_start..up_start + ul.len,
5826 &zq,
5827 &zd,
5828 m_e,
5829 m.up_exps.in_f,
5830 m.up_exps.out_f,
5831 ul.qtype,
5832 ul.row_bytes,
5833 )?;
5834 let mut act = e.uninit(m_e * n_ff_exp)?;
5835 Self::ffn_act_lim(
5836 e,
5837 cfg,
5838 &gate,
5839 &up,
5840 m.gate_exps.macro_scale(ex),
5841 m.up_exps.macro_scale(ex),
5842 lim_exp,
5843 &mut act,
5844 m_e * n_ff_exp,
5845 )?;
5846 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
5847 e.qmatvec_expert_q8(
5848 &dev.down,
5849 down_start..down_start + dl.len,
5850 &aq2,
5851 &ad2,
5852 m_e,
5853 m.down_exps.in_f,
5854 m.down_exps.out_f,
5855 dl.qtype,
5856 dl.row_bytes,
5857 )?
5858 } else {
5859 let gate = e.qmatvec_view(
5860 &dev.gate,
5861 gate_start..gate_start + gl.len,
5862 &gv,
5863 m_e,
5864 m.gate_exps.in_f,
5865 m.gate_exps.out_f,
5866 gl.qtype,
5867 gl.row_bytes,
5868 )?;
5869 let up = e.qmatvec_view(
5870 &dev.up,
5871 up_start..up_start + ul.len,
5872 &gv,
5873 m_e,
5874 m.up_exps.in_f,
5875 m.up_exps.out_f,
5876 ul.qtype,
5877 ul.row_bytes,
5878 )?;
5879 let mut act = e.uninit(m_e * n_ff_exp)?;
5880 Self::ffn_act_lim(
5881 e,
5882 cfg,
5883 &gate,
5884 &up,
5885 m.gate_exps.macro_scale(ex),
5886 m.up_exps.macro_scale(ex),
5887 lim_exp,
5888 &mut act,
5889 m_e * n_ff_exp,
5890 )?;
5891 let actv = act.slice(0..m_e * n_ff_exp);
5892 e.qmatvec_view(
5893 &dev.down,
5894 down_start..down_start + dl.len,
5895 &actv,
5896 m_e,
5897 m.down_exps.in_f,
5898 m.down_exps.out_f,
5899 dl.qtype,
5900 dl.row_bytes,
5901 )?
5902 }
5903 } else if use_cache {
5904 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5905 if grouped_q8 {
5906 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
5907 let gate = e.with_moe_cache(max_block, |cache, eng| {
5908 let id = BlockId::new(il, PROJ_GATE, ex as u16);
5909 let slot =
5910 cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
5911 eng.qmatvec_expert_q8(
5912 cache.buf(slot),
5913 0..gl.len,
5914 &zq,
5915 &zd,
5916 m_e,
5917 m.gate_exps.in_f,
5918 m.gate_exps.out_f,
5919 gl.qtype,
5920 gl.row_bytes,
5921 )
5922 })?;
5923 let up = e.with_moe_cache(max_block, |cache, eng| {
5924 let id = BlockId::new(il, PROJ_UP, ex as u16);
5925 let slot =
5926 cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
5927 eng.qmatvec_expert_q8(
5928 cache.buf(slot),
5929 0..ul.len,
5930 &zq,
5931 &zd,
5932 m_e,
5933 m.up_exps.in_f,
5934 m.up_exps.out_f,
5935 ul.qtype,
5936 ul.row_bytes,
5937 )
5938 })?;
5939 let mut act = e.uninit(m_e * n_ff_exp)?;
5940 Self::ffn_act_lim(
5941 e,
5942 cfg,
5943 &gate,
5944 &up,
5945 m.gate_exps.macro_scale(ex),
5946 m.up_exps.macro_scale(ex),
5947 lim_exp,
5948 &mut act,
5949 m_e * n_ff_exp,
5950 )?;
5951 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
5952 e.with_moe_cache(max_block, |cache, eng| {
5953 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
5954 let slot =
5955 cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
5956 eng.qmatvec_expert_q8(
5957 cache.buf(slot),
5958 0..dl.len,
5959 &aq2,
5960 &ad2,
5961 m_e,
5962 m.down_exps.in_f,
5963 m.down_exps.out_f,
5964 dl.qtype,
5965 dl.row_bytes,
5966 )
5967 })?
5968 } else {
5969 let gate = e.with_moe_cache(max_block, |cache, eng| {
5970 let id = BlockId::new(il, PROJ_GATE, ex as u16);
5971 let slot =
5972 cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
5973 eng.qmatvec_view(
5974 cache.buf(slot),
5975 0..gl.len,
5976 &gv,
5977 m_e,
5978 m.gate_exps.in_f,
5979 m.gate_exps.out_f,
5980 gl.qtype,
5981 gl.row_bytes,
5982 )
5983 })?;
5984 let up = e.with_moe_cache(max_block, |cache, eng| {
5985 let id = BlockId::new(il, PROJ_UP, ex as u16);
5986 let slot =
5987 cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
5988 eng.qmatvec_view(
5989 cache.buf(slot),
5990 0..ul.len,
5991 &gv,
5992 m_e,
5993 m.up_exps.in_f,
5994 m.up_exps.out_f,
5995 ul.qtype,
5996 ul.row_bytes,
5997 )
5998 })?;
5999 let mut act = e.uninit(m_e * n_ff_exp)?;
6000 Self::ffn_act_lim(
6001 e,
6002 cfg,
6003 &gate,
6004 &up,
6005 m.gate_exps.macro_scale(ex),
6006 m.up_exps.macro_scale(ex),
6007 lim_exp,
6008 &mut act,
6009 m_e * n_ff_exp,
6010 )?;
6011 let actv = act.slice(0..m_e * n_ff_exp);
6012 e.with_moe_cache(max_block, |cache, eng| {
6013 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
6014 let slot =
6015 cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
6016 eng.qmatvec_view(
6017 cache.buf(slot),
6018 0..dl.len,
6019 &actv,
6020 m_e,
6021 m.down_exps.in_f,
6022 m.down_exps.out_f,
6023 dl.qtype,
6024 dl.row_bytes,
6025 )
6026 })?
6027 }
6028 } else {
6029 let sg = scratch_g.as_mut().unwrap();
6030 let su = scratch_u.as_mut().unwrap();
6031 let sd = scratch_d.as_mut().unwrap();
6032 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
6033 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
6034 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
6035 if grouped_q8 {
6036 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6037 let gate = e.qmatvec_expert_q8(
6038 sg,
6039 0..gl.len,
6040 &zq,
6041 &zd,
6042 m_e,
6043 m.gate_exps.in_f,
6044 m.gate_exps.out_f,
6045 gl.qtype,
6046 gl.row_bytes,
6047 )?;
6048 let up = e.qmatvec_expert_q8(
6049 su,
6050 0..ul.len,
6051 &zq,
6052 &zd,
6053 m_e,
6054 m.up_exps.in_f,
6055 m.up_exps.out_f,
6056 ul.qtype,
6057 ul.row_bytes,
6058 )?;
6059 let mut act = e.uninit(m_e * n_ff_exp)?;
6060 Self::ffn_act_lim(
6061 e,
6062 cfg,
6063 &gate,
6064 &up,
6065 m.gate_exps.macro_scale(ex),
6066 m.up_exps.macro_scale(ex),
6067 lim_exp,
6068 &mut act,
6069 m_e * n_ff_exp,
6070 )?;
6071 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6072 e.qmatvec_expert_q8(
6073 sd,
6074 0..dl.len,
6075 &aq2,
6076 &ad2,
6077 m_e,
6078 m.down_exps.in_f,
6079 m.down_exps.out_f,
6080 dl.qtype,
6081 dl.row_bytes,
6082 )?
6083 } else {
6084 let gate = e.qmatvec_view(
6085 sg,
6086 0..gl.len,
6087 &gv,
6088 m_e,
6089 m.gate_exps.in_f,
6090 m.gate_exps.out_f,
6091 gl.qtype,
6092 gl.row_bytes,
6093 )?;
6094 let up = e.qmatvec_view(
6095 su,
6096 0..ul.len,
6097 &gv,
6098 m_e,
6099 m.up_exps.in_f,
6100 m.up_exps.out_f,
6101 ul.qtype,
6102 ul.row_bytes,
6103 )?;
6104 let mut act = e.uninit(m_e * n_ff_exp)?;
6105 Self::ffn_act_lim(
6106 e,
6107 cfg,
6108 &gate,
6109 &up,
6110 m.gate_exps.macro_scale(ex),
6111 m.up_exps.macro_scale(ex),
6112 lim_exp,
6113 &mut act,
6114 m_e * n_ff_exp,
6115 )?;
6116 let actv = act.slice(0..m_e * n_ff_exp);
6117 e.qmatvec_view(
6118 sd,
6119 0..dl.len,
6120 &actv,
6121 m_e,
6122 m.down_exps.in_f,
6123 m.down_exps.out_f,
6124 dl.qtype,
6125 dl.row_bytes,
6126 )?
6127 }
6128 };
6129
6130 e.scatter_slot(&y, &tok_idx_d, &slot_idx_d, &weight_d,
6132 &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
6133 }
6134
6135 let mut moe_out = e.zeros(t * n_embd)?;
6137 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
6138
6139 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
6141 m_dist.sort_unstable();
6142 let active = m_dist.len();
6143 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
6144 let median = m_dist[active / 2];
6145 let max_m = *m_dist.last().unwrap();
6146 let min_m = m_dist[0];
6147 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
6148 println!("moe-grouped il={il} t={t} active={active}/{n_expert} \
6149 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
6150 above_gemm_threshold(>=16)={above16}/{active}");
6151 }
6152
6153 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6154 Ok(moe_out)
6155 }
6156
6157 pub(crate) fn moe_ffn_lockstep(
6164 &self,
6165 e: &Engine,
6166 m: &MoeWeights,
6167 zbatch: &CudaSlice<f32>,
6168 mrows: usize,
6169 il: u16,
6170 max_block: usize,
6171 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6172 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6173 let cfg = &self.cfg;
6174 let moe = cfg.moe.as_ref().unwrap();
6175 let n_embd = cfg.n_embd as usize;
6176 let n_expert = moe.expert_count as usize;
6177 let n_used = moe.expert_used_count as usize;
6178 let n_ff_exp = moe.expert_ff_length as usize;
6179 let lim_exp = cfg.clamp_exp_at(il as u32);
6181 let lim_shexp = cfg.clamp_shexp_at(il as u32);
6182
6183 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
6184 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
6185 Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
6186 m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
6187 } else {
6188 Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
6189 None, None, m.active_experts.as_deref())?
6190 };
6191 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
6192
6193 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
6195 Ok((0..n_expert)
6196 .map(|ex| {
6197 [PROJ_GATE, PROJ_UP, PROJ_DOWN].into_iter().all(|p| {
6198 c.resident(BlockId::new(il, p, ex as u16)).is_some()
6199 })
6200 })
6201 .collect())
6202 })?;
6203
6204 struct Group {
6205 rows: Vec<i32>,
6206 slots: Vec<i32>,
6207 weights: Vec<f32>,
6208 }
6209 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
6210 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
6211 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
6212 Default::default();
6213 for row in 0..mrows {
6214 for j in 0..n_used {
6215 let ex = sel_all[row * n_used + j] as usize;
6216 let w = w_all[row * n_used + j];
6217 if resident_expert[ex] {
6218 let group = groups.entry(ex).or_insert_with(|| Group {
6219 rows: Vec::new(),
6220 slots: Vec::new(),
6221 weights: Vec::new(),
6222 });
6223 group.rows.push(row as i32);
6224 group.slots.push(j as i32);
6225 group.weights.push(w);
6226 } else {
6227 crate::cpu_experts::record_incomplete_gpu_residency(0);
6228 cpu_rows[row].push((ex, w));
6229 cpu_by_expert.entry(ex).or_default().push((row, w));
6230 }
6231 }
6232 }
6233
6234 let host_rows = e.dtoh(zbatch)?;
6240 let rows_ok = crate::cpu_experts::rows_supported();
6241 enum CpuPart {
6242 Single { row: usize },
6243 Rows { rows: Vec<usize> },
6244 }
6245 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
6246 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
6247 if rows_ok {
6248 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
6249 .into_iter()
6250 .filter(|(_, rows)| rows.len() >= 2)
6251 .collect();
6252 shared.sort_by_key(|(ex, _)| *ex);
6253 for (ex, mut row_weights) in shared {
6254 row_weights.sort_by_key(|(row, _)| *row);
6255 let inputs: Vec<(&[f32], f32)> = row_weights
6256 .iter()
6257 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
6258 .collect();
6259 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
6260 .map_err(std::io::Error::other)?;
6261 for &(row, _) in &row_weights {
6262 rows_served.insert((row, ex));
6263 }
6264 tickets.push((
6265 CpuPart::Rows {
6266 rows: row_weights.iter().map(|&(row, _)| row).collect(),
6267 },
6268 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
6269 ));
6270 }
6271 }
6272 for (row, selected) in cpu_rows.iter().enumerate() {
6273 let leftover: Vec<(usize, f32)> = selected
6274 .iter()
6275 .copied()
6276 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
6277 .collect();
6278 if leftover.is_empty() {
6279 continue;
6280 }
6281 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
6282 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
6283 .map_err(std::io::Error::other)?;
6284 tickets.push((
6285 CpuPart::Single { row },
6286 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
6287 ));
6288 }
6289
6290 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
6291 let mut wbuf = e.zeros(mrows * n_used)?;
6292 let mut order: Vec<usize> = groups.keys().copied().collect();
6293 order.sort_by(|&a, &b| {
6294 groups[&b].rows.len().cmp(&groups[&a].rows.len()).then(a.cmp(&b))
6295 });
6296 for &ex in &order {
6297 let group = &groups[&ex];
6298 let m_e = group.rows.len();
6299 let gl = m.gate_exps.expert_layout(ex);
6300 let ul = m.up_exps.expert_layout(ex);
6301 let dl = m.down_exps.expert_layout(ex);
6302 let row_idx_d = e.htod_i32(&group.rows)?;
6303 let slot_idx_d = e.htod_i32(&group.slots)?;
6304 let dmac = m.down_exps.macro_scale(ex);
6305 let weight_d = if dmac == 1.0 {
6306 e.htod(&group.weights)?
6307 } else {
6308 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
6309 e.htod(&scaled)?
6310 };
6311 let mut gathered = e.zeros(m_e * n_embd)?;
6312 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
6313 let gv = gathered.slice(0..m_e * n_embd);
6314 let gate = e.with_moe_cache(max_block, |c, eng| {
6315 let slot = c
6316 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
6317 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6318 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..gl.len, &gv, m_e,
6319 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
6320 })?;
6321 let up = e.with_moe_cache(max_block, |c, eng| {
6322 let slot = c
6323 .resident(BlockId::new(il, PROJ_UP, ex as u16))
6324 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6325 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..ul.len, &gv, m_e,
6326 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
6327 })?;
6328 let mut act = e.zeros(m_e * n_ff_exp)?;
6329 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
6330 m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
6331 let actv = act.slice(0..m_e * n_ff_exp);
6332 let y = e.with_moe_cache(max_block, |c, eng| {
6333 let slot = c
6334 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
6335 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6336 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..dl.len, &actv, m_e,
6337 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
6338 })?;
6339 e.scatter_slot(&y, &row_idx_d, &slot_idx_d, &weight_d,
6340 &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
6341 }
6342 let mut moe_out = e.zeros(mrows * n_embd)?;
6343 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
6344
6345 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
6347 for (part, ticket) in tickets {
6348 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
6349 let mut add_row = |row: usize, chunk: &[f32]| {
6350 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
6351 for (accumulator, value) in sum.iter_mut().zip(chunk) {
6352 *accumulator += value;
6353 }
6354 };
6355 match part {
6356 CpuPart::Single { row } => add_row(row, &cpu_output),
6357 CpuPart::Rows { rows } => {
6358 for (slot, row) in rows.into_iter().enumerate() {
6359 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
6360 }
6361 }
6362 }
6363 }
6364 for (row, sum) in row_sums.into_iter().enumerate() {
6365 let Some(sum) = sum else { continue };
6366 let cpu_output = e.htod(&sum)?;
6367 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
6368 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
6369 }
6370
6371 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6372 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6373 {
6374 let n_ff_sh = gate_shexp.out_features();
6375 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
6376 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
6377 let mut sa = e.zeros(mrows * n_ff_sh)?;
6378 Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp,
6379 &mut sa, mrows * n_ff_sh)?;
6380 let sh = e.matmul(down_shexp, &sa, mrows)?;
6381 let g = match &m.gate_inp_shexp {
6384 Some(gate_inp_shexp) => {
6385 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
6386 }
6387 None => e.htod(&vec![1.0f32; mrows])?,
6388 };
6389 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
6390 }
6391
6392 Ok(moe_out)
6393 }
6394}
6395
6396impl HybridModel {
6402 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
6404 let g = self.cfg.gemma4.as_ref().unwrap();
6405 let swa = g.swa_pattern[il];
6406 let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
6407 (hd, g.head_count_kv[il] as usize, self.cfg.n_head as usize,
6411 if swa { g.rope_base_swa } else { g.rope_base_global },
6412 1.0, swa)
6413 }
6414
6415 fn gemma4_suppress(&self, e: &Engine, ld: &mut CudaSlice<f32>, t: usize)
6419 -> Result<(), Box<dyn std::error::Error>> {
6420 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
6421 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
6422 }
6423 Ok(())
6424 }
6425
6426 fn gemma4_attn_prime(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6431 h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
6432 cache: Option<&mut Cache>)
6433 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6434 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6435 let eps = self.cfg.rms_eps;
6436 let aux = self.gemma4_aux.as_ref().unwrap();
6437
6438 e.mmq_act_begin();
6441 let q0 = e.matmul(&fa.wq, h, t)?; let k0 = e.matmul(&fa.wk, h, t)?; let v0 = if swa { e.matmul(&fa.wv, h, t)? } else { e.clone_dtod(&k0)? };
6446
6447 let mut q = e.uninit(t * nh * hd)?;
6448 let mut k = e.uninit(t * nkv * hd)?;
6449 let mut v = e.uninit(t * nkv * hd)?;
6451 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6455 let emit = t >= 16 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
6456 && *EMIT.get_or_init(|| std::env::var("MEMRA_FA_EMIT").map(|s| s != "0").unwrap_or(true));
6457 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
6458 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
6459 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
6460 let v_f16 = emit && crate::fa_f16pv_on() && match hd {
6463 512 => true,
6464 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
6465 _ => false,
6466 };
6467 if emit {
6468 e.rms_norm_qkv_w4b(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6469 &aux.ones, &mut q, &mut k, &mut v, &mut vb,
6470 hd, nh * t, nkv * t, eps, v_f16)?;
6471 } else {
6472 e.rms_norm_qkv(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6473 &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t, eps)?;
6474 }
6475
6476 let ff = if swa { None } else {
6477 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6478 };
6479 if emit {
6480 e.rope_neox2_bf16e(&mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t,
6481 base, 1.0, ff)?;
6482 } else {
6483 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
6484 }
6485
6486 if let Some(cache) = cache {
6487 let kvl = cache.kv[il].as_mut().unwrap();
6488 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
6489 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
6490 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()))?;
6491 kvl.len += t;
6492 }
6493 let mut attn = e.zeros(t * nh * hd)?;
6494 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6498 if swa && t > win {
6499 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
6500 if emit { e.fa_prefill_w_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
6501 scale, true, win, v_f16)?; }
6502 else { e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true,
6503 win)?; }
6504 } else {
6505 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
6506 }
6507 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
6508 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6509 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
6510 if emit { e.fa_prefill_hd512_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
6511 scale, true, v_f16)?; }
6512 else { e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?; }
6513 } else {
6514 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6515 }
6516 Ok(e.matmul(&fa.wo, &attn, t)?)
6517 }
6518
6519 fn gemma4_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6521 h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
6522 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6523 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
6524 }
6525
6526 fn gemma4_moe_q8(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
6531 bits: &crate::hybrid::Gemma4MoeBits,
6532 mq: &(CudaSlice<i8>, CudaSlice<f32>),
6533 router_in: &CudaSlice<f32>, t: usize)
6534 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6535 let cfg = &self.cfg;
6536 let moe = cfg.moe.as_ref().unwrap();
6537 let n_embd = cfg.n_embd as usize;
6538 let n_expert = moe.expert_count as usize;
6539 let n_used = moe.expert_used_count as usize;
6540 let n_ff_exp = moe.expert_ff_length as usize;
6541 let logits = if crate::router_kernel_on() {
6545 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
6546 } else {
6547 e.matmul(&m.gate_inp, router_in, t)?
6548 };
6549 let dev = m.dev_exps.as_ref().unwrap();
6550 let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
6551 &bits.per_expert_scale_d)?;
6552 let (zq, zd) = mq;
6553 if t == 1 {
6554 let selv = sel_d.slice(0..n_used);
6555 let wv = w_d.slice(0..n_used);
6556 let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, zq, zd,
6557 n_embd, n_ff_exp, n_used, n_expert,
6558 m.gate_exps.qtype, m.up_exps.qtype,
6559 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
6560 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6561 let mut moe_out = e.uninit(n_embd)?;
6562 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
6563 &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
6564 n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
6565 return Ok(moe_out);
6566 }
6567 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
6568 let act = if csr {
6569 e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, zq, zd, t * n_used,
6570 n_embd, n_ff_exp, n_used, n_expert,
6571 m.gate_exps.qtype, m.up_exps.qtype,
6572 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6573 } else {
6574 e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, zq, zd, t,
6575 n_embd, n_ff_exp, n_used, n_expert,
6576 m.gate_exps.qtype, m.up_exps.qtype,
6577 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6578 };
6579 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6580 let mut moe_out = e.uninit(t * n_embd)?;
6581 e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
6584 n_ff_exp, n_embd, n_used, n_expert,
6585 m.down_exps.qtype, m.down_exps.row_bytes)?;
6586 Ok(moe_out)
6587 }
6588
6589 fn gemma4_moe(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
6593 bits: &crate::hybrid::Gemma4MoeBits, moe_in: &CudaSlice<f32>,
6594 router_in: &CudaSlice<f32>, t: usize)
6595 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6596 let cfg = &self.cfg;
6597 let moe = cfg.moe.as_ref().unwrap();
6598 let n_embd = cfg.n_embd as usize;
6599 let n_expert = moe.expert_count as usize;
6600 let n_used = moe.expert_used_count as usize;
6601 let n_ff_exp = moe.expert_ff_length as usize;
6602
6603 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
6607 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
6608 } else {
6609 e.matmul(&m.gate_inp, router_in, t)?
6610 };
6611
6612 if t < PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
6617 && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
6618 && expert_dp4a_supported(m.down_exps.qtype)
6619 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0") {
6620 let dev = m.dev_exps.as_ref().unwrap();
6621 let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
6622 &bits.per_expert_scale_d)?;
6623 if t == 1 {
6624 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
6625 let selv = sel_d.slice(0..n_used);
6626 let wv = w_d.slice(0..n_used);
6627 let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, &zq, &zd,
6628 n_embd, n_ff_exp, n_used, n_expert,
6629 m.gate_exps.qtype, m.up_exps.qtype,
6630 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
6631 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6632 let mut moe_out = e.uninit(n_embd)?;
6633 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
6634 &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
6635 n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
6636 return Ok(moe_out);
6637 }
6638 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
6643 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
6644 let act = if csr {
6645 e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, t * n_used,
6646 n_embd, n_ff_exp, n_used, n_expert,
6647 m.gate_exps.qtype, m.up_exps.qtype,
6648 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6649 } else {
6650 e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
6651 n_embd, n_ff_exp, n_used, n_expert,
6652 m.gate_exps.qtype, m.up_exps.qtype,
6653 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6654 };
6655 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6656 let mut moe_out = e.uninit(t * n_embd)?;
6657 e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
6658 n_ff_exp, n_embd, n_used, n_expert,
6659 m.down_exps.qtype, m.down_exps.row_bytes)?;
6660 return Ok(moe_out);
6661 }
6662
6663 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
6664 for (i, &sx) in sel_all.iter().enumerate() {
6665 w_all[i] *= bits.per_expert_scale[sx as usize];
6666 }
6667
6668 if t >= PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
6672 && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
6673 && expert_dp4a_supported(m.down_exps.qtype)
6674 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0") {
6675 let dev = m.dev_exps.as_ref().unwrap();
6676 let n_pairs = t * n_used;
6677 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6678 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6679 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6680 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6681 let pt = e.htod_i32(&pair_tok)?;
6682 let pw = e.htod(&w_all)?;
6683 let toff = e.htod_i32(&tok_off)?;
6684 let tids = e.htod_i32(&tok_ids)?;
6685 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6686 for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
6687 let mut ex_ids: Vec<i32> = Vec::new();
6688 let mut ex_off: Vec<i32> = vec![0];
6689 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6690 for (ex, list) in by_ex.iter().enumerate() {
6691 if list.is_empty() { continue; }
6692 ex_ids.push(ex as i32);
6693 ex_pairs.extend_from_slice(list);
6694 ex_off.push(ex_pairs.len() as i32);
6695 }
6696 let n_active = ex_ids.len();
6697 let exi = e.htod_i32(&ex_ids)?;
6698 let exo = e.htod_i32(&ex_off)?;
6699 let exp_d = e.htod_i32(&ex_pairs)?;
6700 if crate::moe_f16g_gemma_on()
6708 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6709 && f16g_proj_ok(m.up_exps.qtype, n_embd)
6710 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp) {
6711 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6712 let csr_tok_d = e.htod_i32(&csr_tok)?;
6713 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
6714 let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
6715 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
6716 m.gate_exps.qtype, m.gate_exps.row_bytes)?;
6717 let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
6718 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
6719 m.up_exps.qtype, m.up_exps.row_bytes)?;
6720 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6721 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6722 let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
6723 &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
6724 m.down_exps.qtype, m.down_exps.row_bytes)?;
6725 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
6726 let mut moe_out = e.uninit(t * n_embd)?;
6727 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6728 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
6729 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
6730 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
6731 eprintln!("[f16g-debug] post-permute bad={} post-scatter bad={}",
6732 scan(&yd), scan(&mo));
6733 }
6734 return Ok(moe_out);
6735 }
6736 let mma = n_embd % 256 == 0
6739 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
6740 let (gate, up) = if mma {
6741 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
6742 (e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
6743 n_embd, n_ff_exp, n_active, n_pairs, t,
6744 m.gate_exps.qtype, m.gate_exps.row_bytes)?,
6745 e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
6746 n_embd, n_ff_exp, n_active, n_pairs, t,
6747 m.up_exps.qtype, m.up_exps.row_bytes)?)
6748 } else {
6749 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
6750 (e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 0, &exi, &exo, &exp_d, &pt, &zq, &zd,
6751 n_embd, n_ff_exp, n_expert, n_active, n_pairs,
6752 m.gate_exps.qtype, m.gate_exps.row_bytes)?,
6753 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 1, &exi, &exo, &exp_d, &pt, &zq, &zd,
6754 n_embd, n_ff_exp, n_expert, n_active, n_pairs,
6755 m.up_exps.qtype, m.up_exps.row_bytes)?)
6756 };
6757 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6758 let pself = e.htod_i32(&pair_self)?;
6759 let y_down = if mma {
6771 let in_pad = n_ff_exp.div_ceil(256) * 256;
6772 let a_scr = if crate::moe_fuse_actq_on() {
6773 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
6774 } else {
6775 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6776 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6777 };
6778 e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
6779 in_pad, n_embd, n_active, n_pairs, n_pairs,
6780 m.down_exps.qtype, m.down_exps.row_bytes)?
6781 } else {
6782 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6783 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6784 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
6785 n_ff_exp, n_embd, n_expert, n_active, n_pairs,
6786 m.down_exps.qtype, m.down_exps.row_bytes)?
6787 };
6788 let mut moe_out = e.uninit(t * n_embd)?;
6789 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6790 return Ok(moe_out);
6791 }
6792
6793 let g_len = m.gate_exps.expert_stride;
6794 let u_len = m.up_exps.expert_stride;
6795 let d_len = m.down_exps.expert_stride;
6796 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
6800 let (mut sg, mut su, mut sd) = if dev.is_some() { (None, None, None) } else {
6801 (Some(e.alloc_u8_uninit(g_len)?), Some(e.alloc_u8_uninit(u_len)?), Some(e.alloc_u8_uninit(d_len)?))
6802 };
6803 let mut moe_out = e.zeros(t * n_embd)?;
6804 for tok in 0..t {
6805 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
6806 let w = &w_all[tok * n_used..(tok + 1) * n_used];
6807 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
6808 for (j, &ex) in sel.iter().enumerate() {
6809 let ex = ex as usize;
6810 let gate = match dev {
6811 Some(d) => e.qmatvec_view(&d.gate, ex * g_len..(ex + 1) * g_len, &zt, 1,
6812 m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?,
6813 None => {
6814 let sg = sg.as_mut().unwrap();
6815 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
6816 e.qmatvec_view(sg, 0..g_len, &zt, 1,
6817 m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?
6818 }
6819 };
6820 let up = match dev {
6821 Some(d) => e.qmatvec_view(&d.up, ex * u_len..(ex + 1) * u_len, &zt, 1,
6822 m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?,
6823 None => {
6824 let su = su.as_mut().unwrap();
6825 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
6826 e.qmatvec_view(su, 0..u_len, &zt, 1,
6827 m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?
6828 }
6829 };
6830 let mut act = e.uninit(n_ff_exp)?;
6831 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
6832 let actv = act.slice(0..n_ff_exp);
6833 let y = match dev {
6834 Some(d) => e.qmatvec_view(&d.down, ex * d_len..(ex + 1) * d_len, &actv, 1,
6835 m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?,
6836 None => {
6837 let sd = sd.as_mut().unwrap();
6838 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
6839 e.qmatvec_view(sd, 0..d_len, &actv, 1,
6840 m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?
6841 }
6842 };
6843 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6844 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
6845 }
6846 }
6847 Ok(moe_out)
6848 }
6849
6850 fn gemma4_layer(&self, e: &Engine, il: usize, layer: &crate::hybrid::HybridLayer,
6852 x: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
6853 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6854 let n_embd = self.cfg.n_embd as usize;
6855 let eps = self.cfg.rms_eps;
6856
6857 let mut h = e.zeros(t * n_embd)?;
6858 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6859 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6860 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
6861 let mut cur = e.zeros(t * n_embd)?;
6863 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6864 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
6865 }
6866
6867 fn gemma4_layer_tail_add(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
6871 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
6872 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6873 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
6874 }
6875
6876 fn gemma4_layer_tail_add_n(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
6879 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
6880 next_norm: Option<&CudaSlice<f32>>)
6881 -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
6882 let n_embd = self.cfg.n_embd as usize;
6883 let bits = layer.gemma4.as_ref().unwrap();
6884 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
6885 let mut xn = e.uninit(t * n_embd)?;
6886 match next_norm {
6887 Some(w) => {
6888 let mut hn = e.uninit(t * n_embd)?;
6889 e.add_scale_rms_norm(&sn, &attn_out, bits.layer_scale, w, &mut xn, &mut hn,
6890 n_embd, t, self.cfg.rms_eps)?;
6891 Ok((xn, Some(hn)))
6892 }
6893 None => {
6894 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
6895 Ok((xn, None))
6896 }
6897 }
6898 }
6899
6900 fn gemma4_layer_tail_core(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
6903 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
6904 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6905 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
6906 }
6907
6908 fn gemma4_layer_tail_core_pn(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
6915 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
6916 pre_norm: Option<&CudaSlice<f32>>, defer_post_norm: bool)
6917 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6918 let n_embd = self.cfg.n_embd as usize;
6919 let eps = self.cfg.rms_eps;
6920 let bits = layer.gemma4.as_ref().unwrap();
6921
6922 let Some(mbits) = bits.moe_bits.as_ref() else {
6925 let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
6926 else { panic!("gemma4 dense layer without Dense ffn") };
6927 let mut attn_out = e.uninit(t * n_embd)?;
6928 let mut zsh = e.uninit(t * n_embd)?;
6929 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6932 match pre_norm {
6933 Some(wa) if t == 1 => {
6934 zpair = Some(e.rms_pre_add_rms_norm_q8z(cur, wa, x,
6935 bits.ffn_norm.float_data(),
6936 &mut attn_out, &mut zsh,
6937 n_embd, t, eps)?);
6938 }
6939 Some(wa) => e.rms_pre_add_rms_norm(cur, wa, x, bits.ffn_norm.float_data(),
6940 &mut attn_out, &mut zsh, n_embd, t, eps)?,
6941 None => e.add_rms_norm(cur, x, bits.ffn_norm.float_data(), &mut attn_out,
6942 &mut zsh, n_embd, t, eps)?,
6943 }
6944 let n_ff = ffn_gate.out_features();
6945 let (gate, up) = if t == 1 {
6951 let (zq, zd) = match zpair {
6952 Some(p) => p,
6953 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
6954 };
6955 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
6956 Some(p) => p,
6957 None => (e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
6958 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?),
6959 }
6960 } else {
6961 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6966 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6967 let fused = if f2b {
6968 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
6969 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
6970 } else { None };
6971 match fused {
6972 Some(p) => p,
6973 None => {
6974 e.mmq_act_begin();
6976 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
6977 }
6978 }
6979 };
6980 let mut act = e.uninit(t * n_ff)?;
6981 let f0 = if e.uses_q8_1_fast(ffn_down) {
6984 let upv = e.view(&up, t * n_ff);
6985 let up_all = upv.slice(0..t * n_ff);
6986 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
6987 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
6988 } else {
6989 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
6990 e.matmul(ffn_down, &act, t)?
6991 };
6992 if defer_post_norm { return Ok((f0, attn_out)); }
6993 let mut sn = e.uninit(t * n_embd)?;
6994 e.rms_norm(&f0, bits.post_ffw_norm.float_data(), &mut sn, n_embd, t, eps)?;
6995 return Ok((sn, attn_out));
6996 };
6997
6998 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
6999 let mut attn_out = e.uninit(t * n_embd)?;
7004 let mut router_in = e.uninit(t * n_embd)?;
7005 let fast_moe = match &layer.ffn {
7006 crate::hybrid::Ffn::Moe(m) => m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7007 && expert_dp4a_supported(m.gate_exps.qtype)
7008 && expert_dp4a_supported(m.up_exps.qtype)
7009 && expert_dp4a_supported(m.down_exps.qtype)
7010 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0"),
7011 _ => false,
7012 };
7013 let q8z = t < PRIME_MIN_T && fast_moe;
7014 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
7015 let (z0, m2) = e.add_rms_norm3_q8z(cur, x, bits.ffn_norm.float_data(),
7016 &mbits.router_scale_pre,
7017 mbits.pre_ffw_norm_2.float_data(),
7018 &mut attn_out, &mut router_in, n_embd, t, eps)?;
7019 (None, Some(z0), Some(m2))
7020 } else {
7021 let mut zsh = e.uninit(t * n_embd)?;
7022 let mut moe_in = e.uninit(t * n_embd)?;
7023 e.add_rms_norm3(cur, x, bits.ffn_norm.float_data(), &mbits.router_scale_pre,
7024 mbits.pre_ffw_norm_2.float_data(), &mut attn_out, &mut zsh,
7025 &mut router_in, &mut moe_in, n_embd, t, eps)?;
7026 (Some((zsh, moe_in)), None, None)
7027 };
7028 let attn_out2 = attn_out;
7029 #[allow(unused_variables)]
7030 let attn_out = &attn_out2;
7031 let n_ff = mbits.shared_gate.out_features();
7032 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
7033 if t == 1 {
7034 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
7035 Some(p) => p,
7036 None => {
7037 let h0 = e.zeros(0)?;
7038 (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
7039 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?)
7040 }
7041 }
7042 } else {
7043 let h0 = e.zeros(0)?;
7045 (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
7046 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?)
7047 }
7048 } else {
7049 let (zsh, _) = zsh_f32.as_ref().unwrap();
7050 (e.matmul(&mbits.shared_gate, zsh, t)?, e.matmul(&mbits.shared_up, zsh, t)?)
7051 };
7052 let mut act = e.uninit(t * n_ff)?;
7053 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
7054 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
7055 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else { panic!("gemma4 layer not MoE") };
7056 let moe0 = match (&moe_q8, &zsh_f32) {
7057 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
7058 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
7059 _ => unreachable!(),
7060 };
7061 let mut mlp = e.uninit(t * n_embd)?;
7063 let mut moe = e.uninit(t * n_embd)?;
7064 e.rms_norm2x(&mlp0, &moe0, mbits.post_ffw_norm_1.float_data(),
7065 mbits.post_ffw_norm_2.float_data(), &mut mlp, &mut moe, n_embd, t, eps)?;
7066
7067 let mut sum = e.uninit(t * n_embd)?;
7070 let mut sn = e.uninit(t * n_embd)?;
7071 e.add_rms_norm(&mlp, &moe, bits.post_ffw_norm.float_data(), &mut sum, &mut sn,
7072 n_embd, t, eps)?;
7073 Ok((sn, attn_out2))
7074 }
7075
7076 fn gemma4_layer_tail_add_nq(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7078 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7079 next_norm: Option<&CudaSlice<f32>>)
7080 -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>> {
7081 let n_embd = self.cfg.n_embd as usize;
7082 let bits = layer.gemma4.as_ref().unwrap();
7083 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
7084 let mut xn = e.uninit(t * n_embd)?;
7085 match next_norm {
7086 Some(w) => {
7087 let pair = e.add_scale_rms_norm_q8_1(&sn, &attn_out, bits.layer_scale, w, &mut xn,
7088 n_embd, t, self.cfg.rms_eps)?;
7089 Ok((xn, Some(pair)))
7090 }
7091 None => {
7092 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
7093 Ok((xn, None))
7094 }
7095 }
7096 }
7097
7098 fn gemma4_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
7101 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7102 if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, last_only); }
7105 let n_embd = self.cfg.n_embd as usize;
7106 let t = tokens.len();
7107 let pos: Vec<i32> = (0..t as i32).collect();
7108 let pos_d = e.htod_i32(&pos)?;
7109
7110 let mut x = self.embed(e, tokens)?;
7111 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7112 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
7115 let stat = |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
7116 let h = e.dtoh(x)?;
7117 let bad = h.iter().filter(|v| !v.is_finite()).count();
7118 let mx = h.iter().filter(|v| v.is_finite()).fold(0.0f32, |m, v| m.max(v.abs()));
7119 eprintln!("[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}", &h[..3]);
7120 Ok(())
7121 };
7122 if probe { stat(e, &x, "embed")?; }
7123 for (il, layer) in self.layers.iter().enumerate() {
7124 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
7125 if probe { stat(e, &x, &format!("L{il}"))?; }
7126 }
7127 let mut hn = e.zeros(t * n_embd)?;
7128 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, self.cfg.rms_eps)?;
7129 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7130 let n_vocab = self.output.out_features();
7131 let logits = if last_only {
7132 let hv = e.view(&hn, t * n_embd);
7133 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
7134 let mut hlast = e.zeros(n_embd)?;
7135 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
7136 let mut ld = e.matmul(&self.output, &hlast, 1)?;
7137 e.softcap(&mut ld, cap, n_vocab)?;
7138 self.gemma4_suppress(e, &mut ld, 1)?;
7139 e.dtoh(&ld)?
7140 } else {
7141 let mut ld = e.matmul(&self.output, &hn, t)?;
7142 e.softcap(&mut ld, cap, t * n_vocab)?;
7143 self.gemma4_suppress(e, &mut ld, t)?;
7144 e.dtoh(&ld)?
7145 };
7146 Ok(logits)
7147 }
7148
7149 pub(crate) fn gemma4_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
7154 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7155 if cache.pos != 0 {
7160 return Err("gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
7161 — prime the full prompt in one call or decode tokenwise".into());
7162 }
7163 let n_embd = self.cfg.n_embd as usize;
7164 let eps = self.cfg.rms_eps;
7165 let t = tokens.len();
7166 let pos: Vec<i32> = (0..t as i32).collect();
7167 let pos_d = e.htod_i32(&pos)?;
7168 let mut x = self.embed(e, tokens)?;
7169 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7170 for (il, layer) in self.layers.iter().enumerate() {
7171 let mut h = e.zeros(t * n_embd)?;
7172 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7173 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer not full-attn") };
7174 let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
7175 let mut cur = e.zeros(t * n_embd)?;
7176 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
7177 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
7178 self.dflash_tap(e, cache, il, &x, t)?;
7179 }
7180 cache.pos += t;
7181 let hiddens = e.clone_dtod(&x)?;
7182 let xv = e.view(&x, t * n_embd);
7183 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
7184 let mut h_seed = e.zeros(n_embd)?;
7185 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
7186 let mut hn = e.uninit(n_embd)?;
7187 e.rms_norm(&h_seed, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7188 let mut ld = e.matmul(&self.output, &hn, 1)?;
7189 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7190 e.softcap(&mut ld, cap, self.output.out_features())?;
7191 self.gemma4_suppress(e, &mut ld, 1)?;
7192 let logits = e.dtoh(&ld)?;
7193 Ok((logits, h_seed, hiddens))
7194 }
7195
7196 fn gemma4_decode_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
7201 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
7202 pos_d: &CudaSlice<i32>, cache: &mut Cache)
7203 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7204 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7205 let eps = self.cfg.rms_eps;
7206 let aux = self.gemma4_aux.as_ref().unwrap();
7207 let (hq, hdq) = (hq, hdq);
7208 let h0 = e.zeros(0)?;
7209 let h = &h0;
7210 let (q0, k0, v0) = if swa {
7211 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
7212 Some(t3) => t3,
7213 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7214 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
7215 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
7216 }
7217 } else {
7218 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
7219 Some(p) => p,
7220 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7221 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?),
7222 };
7223 let v0 = e.clone_dtod(&k0)?;
7224 (q0, k0, v0)
7225 };
7226 let mut q = e.uninit(nh * hd)?;
7227 let mut k = e.uninit(nkv * hd)?;
7228 let mut v = e.uninit(nkv * hd)?;
7229 let ff = if swa { None } else {
7232 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
7233 };
7234 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7235 &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7236 pos_d, nh, nkv, base, 1.0, ff, eps)?;
7237 let kvl = cache.kv[il].as_mut().unwrap();
7238 e.append_kv_quantized(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len,
7239 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()))?;
7240 kvl.len += 1;
7241 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7245 let mut attn = e.uninit(nh * hd)?;
7246 if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
7248 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7249 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7250 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7251 let base = kvl.len as i32;
7253 e.i32_set_k(&mut kvl.len_d, base)?;
7254 e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1, scale,
7255 kvl.k_tok_bytes, kvl.v_tok_bytes, Some((&kvl.len_d, -1)), false,
7256 false, None)?;
7257 return Ok(e.matmul(&fa.wo, &attn, 1)?);
7258 }
7259 if swa && kvl.len > win && hd == 256
7261 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7262 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7263 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7264 let base = kvl.len as i32;
7265 e.i32_set_k(&mut kvl.len_d, base)?;
7266 e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1, 1, scale,
7267 win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
7268 return Ok(e.matmul(&fa.wo, &attn, 1)?);
7269 }
7270 let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
7271 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
7272 (off_tok + t_kv) * kvl.k_tok_bytes);
7273 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
7274 (off_tok + t_kv) * kvl.v_tok_bytes);
7275 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
7276 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
7277 Ok(e.matmul(&fa.wo, &attn, 1)?)
7278 }
7279
7280 #[allow(clippy::too_many_arguments)]
7287 pub fn gemma4_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
7288 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7289 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7290 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>)
7291 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7292 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
7293 self.gemma4_decode_step_dc_into(e, token_d, pos_d, embd_gpu, embd_qt, embd_rb, cache,
7294 n_vocab, cap_bucket_max, &mut tok_out)?;
7295 Ok(tok_out)
7296 }
7297
7298 #[allow(clippy::too_many_arguments)]
7301 pub fn gemma4_decode_step_dc_into(&self, e: &Engine, token_d: &CudaSlice<u32>,
7302 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7303 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7304 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
7305 tok_out: &mut CudaSlice<u32>)
7306 -> Result<(), Box<dyn std::error::Error>> {
7307 let n_embd = self.cfg.n_embd as usize;
7308 let eps = self.cfg.rms_eps;
7309 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7310 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7311 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7312 let n_layers = self.layers.len();
7313 for (il, layer) in self.layers.iter().enumerate() {
7314 let (hq, hdq) = match h_carry.take() {
7315 Some(p) => p,
7316 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
7317 };
7318 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7319 let o = self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
7320 let mut cur = e.uninit(n_embd)?;
7321 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
7322 let next_norm = if il + 1 < n_layers {
7323 Some(self.layers[il + 1].attn_norm.float_data())
7324 } else { None };
7325 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
7326 x = xn;
7327 h_carry = hn;
7328 }
7329 let mut hn = e.uninit(n_embd)?;
7330 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7331 let mut logits = e.matmul(&self.output, &hn, 1)?;
7332 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
7334 e.inc_seqlen(pos_d)?;
7335 if cap_bucket_max.is_none() { cache.pos += 1; }
7336 Ok(())
7337 }
7338
7339 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
7346 let n_embd = self.cfg.n_embd as usize;
7347 let n_vocab = self.output.out_features();
7348 let n_layers = self.layers.len();
7349 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
7350 for il in 0..n_layers {
7351 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
7352 qmax = qmax.max(nh * hd);
7353 kvmax = kvmax.max(nkv * hd);
7354 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
7355 ffmax = ffmax.max(ffn_gate.out_features());
7356 }
7357 }
7358 Ok(G4DcSlots {
7359 x: e.uninit(n_embd)?, xn: e.uninit(n_embd)?, cur: e.uninit(n_embd)?,
7360 hq: e.alloc_i8_uninit(n_embd)?, hd_: e.uninit(n_embd / 32)?,
7361 q0: e.uninit(qmax)?, k0: e.uninit(kvmax)?, v0: e.uninit(kvmax)?,
7362 q: e.uninit(qmax)?, k: e.uninit(kvmax)?, v: e.uninit(kvmax)?,
7363 attn: e.uninit(qmax)?, o: e.uninit(n_embd)?,
7364 attn_out: e.uninit(n_embd)?, zsh: e.uninit(n_embd)?,
7365 zq: e.alloc_i8_uninit(n_embd.max(qmax))?, zd: e.uninit(n_embd.max(qmax) / 32)?,
7368 gate: e.uninit(ffmax)?, up: e.uninit(ffmax)?,
7369 act: e.uninit(ffmax)?, actq: e.alloc_i8_uninit(ffmax)?, actd: e.uninit(ffmax / 32)?,
7370 f0: e.uninit(n_embd)?, sn: e.uninit(n_embd)?,
7371 hn: e.uninit(n_embd)?, logits: e.uninit(n_vocab)?,
7372 })
7373 }
7374
7375 fn g4_matvec_m1_into(&self, e: &Engine, w: &crate::model::GpuTensor,
7378 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y: &mut CudaSlice<f32>)
7379 -> Result<(), Box<dyn std::error::Error>> {
7380 use crate::model::GpuTensor;
7381 let (bytes, qtype, row_bytes, scale, rp) = match w {
7382 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
7383 (bytes, *qtype, *row_bytes, *scale, *rp),
7384 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
7385 };
7386 let (mbytes, mrp) = match w {
7387 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
7388 _ => (bytes, rp),
7389 };
7390 e.qmatvec_mmvq_into(mbytes, aq, ad, 1, w.in_features(), w.out_features(),
7391 qtype, row_bytes, scale, mrp, y)
7392 }
7393
7394 #[allow(clippy::too_many_arguments)]
7398 pub fn gemma4_decode_step_dc_slotted(&self, e: &Engine, token_d: &CudaSlice<u32>,
7399 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7400 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7401 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
7402 sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>,
7403 ring: Option<(&mut CudaSlice<u32>, usize)>)
7404 -> Result<(), Box<dyn std::error::Error>> {
7405 let n_embd = self.cfg.n_embd as usize;
7406 let eps = self.cfg.rms_eps;
7407 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
7408 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
7409 let n_layers = self.layers.len();
7410 let mut has_carry = false;
7411 for il in 0..n_layers {
7412 if !has_carry {
7413 e.rms_norm_q8_1_into(&sl.x, self.layers[il].attn_norm.float_data(), n_embd, 1,
7414 eps, &mut sl.hq, &mut sl.hd_)?;
7415 }
7416 has_carry = true;
7417 let layer = &self.layers[il];
7418 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7419 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
7420 e.rms_norm(&sl.o, layer.post_attn_norm.float_data(), &mut sl.cur, n_embd, 1, eps)?;
7421 let next_norm = if il + 1 < n_layers {
7422 Some(self.layers[il + 1].attn_norm.float_data())
7423 } else { None };
7424 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
7425 std::mem::swap(&mut sl.x, &mut sl.xn);
7426 }
7427 e.rms_norm(&sl.x, self.output_norm.float_data(), &mut sl.hn, n_embd, 1, eps)?;
7428 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
7429 {
7431 let (zq, zd) = (&sl.zq, &sl.zd);
7432 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
7433 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
7434 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
7435 }
7436 self.gemma4_suppress(e, &mut sl.logits, 1)?;
7437 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
7438 if let Some((ring, base)) = ring {
7439 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
7443 }
7444 e.inc_seqlen(pos_d)?;
7445 if cap_bucket_max.is_none() { cache.pos += 1; }
7446 Ok(())
7447 }
7448
7449 #[allow(clippy::too_many_arguments)]
7451 fn gemma4_decode_attn_dc_slotted(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer,
7452 il: usize, pos_d: &CudaSlice<i32>, cache: &mut Cache,
7453 cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots)
7454 -> Result<(), Box<dyn std::error::Error>> {
7455 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7456 let eps = self.cfg.rms_eps;
7457 let aux = self.gemma4_aux.as_ref().unwrap();
7458 {
7459 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
7460 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
7461 if swa {
7462 if !e.matmul_q4_fused3_into(&fa.wq, &fa.wk, &fa.wv, hq, hdq,
7463 &mut sl.q0, &mut sl.k0, &mut sl.v0)? {
7464 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
7465 }
7466 } else {
7467 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
7468 return Err("slotted step: fused2 unavailable".into());
7469 }
7470 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
7471 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
7472 }
7473 }
7474 let ff = if swa { None } else {
7477 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
7478 };
7479 let kvl = cache.kv[il].as_mut().unwrap();
7480 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7481 if crate::Engine::qkv_append_on() {
7482 e.rms_norm_qkv_rope_append_dc(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(),
7484 fa.k_norm.float_data(), &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
7485 pos_d, nh, nkv, base, 1.0, ff, eps,
7486 &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
7487 } else {
7488 e.rms_norm_qkv_rope(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7489 &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
7490 pos_d, nh, nkv, base, 1.0, ff, eps)?;
7491 e.append_kv_quantized_dc(&sl.k, &sl.v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
7492 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
7493 kv_fp8)?;
7494 }
7495 e.inc_seqlen(&mut kvl.len_d)?;
7496 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
7497 let k_view = e.view_u8(&kvl.k, kvl.k.len());
7498 let v_view = e.view_u8(&kvl.v, kvl.v.len());
7499 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
7500 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7501 let mut fa_q8 = false;
7505 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
7506 e.fa_decode_rows(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, b_glob - 1,
7507 1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7508 Some((&kvl.len_d, -1)), false, false,
7509 Some((&mut sl.zq, &mut sl.zd)))?;
7510 fa_q8 = true;
7511 } else if swa && b_swa > win && hd == 256 && rows_on {
7512 e.fa_decode_rows_w(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv,
7513 &kvl.len_d, -1, 1, scale, win,
7514 kvl.k_tok_bytes, kvl.v_tok_bytes,
7515 Some((&mut sl.zq, &mut sl.zd)))?;
7516 fa_q8 = true;
7517 } else {
7518 let b = if swa { b_swa } else { b_glob };
7519 e.fa_decode_dc(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, &kvl.len_d, b,
7520 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7521 swa && crate::Engine::wkv_on())?;
7522 }
7523 if !fa_q8 {
7524 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
7525 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
7526 }
7527 {
7528 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
7529 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
7530 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
7531 }
7532 Ok(())
7533 }
7534
7535 fn gemma4_layer_tail_slotted(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7538 next_norm: Option<&CudaSlice<f32>>, sl: &mut G4DcSlots)
7539 -> Result<(), Box<dyn std::error::Error>> {
7540 let n_embd = self.cfg.n_embd as usize;
7541 let eps = self.cfg.rms_eps;
7542 let bits = layer.gemma4.as_ref().unwrap();
7543 let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
7544 else { return Err("slotted tail: dense ffn only".into()) };
7545 e.add_rms_norm(&sl.cur, &sl.x, bits.ffn_norm.float_data(), &mut sl.attn_out,
7546 &mut sl.zsh, n_embd, 1, eps)?;
7547 let n_ff = ffn_gate.out_features();
7548 {
7549 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
7550 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
7551 }
7552 {
7553 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
7554 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
7555 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
7556 return Err("slotted tail: ffn fused2 unavailable".into());
7557 }
7558 }
7559 debug_assert!(e.uses_q8_1_fast(ffn_down));
7560 {
7561 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
7562 let upv = e.view(upr, n_ff);
7563 let up_all = upv.slice(0..n_ff);
7564 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
7565 e.gelu_tanh_mul_q8_1_into(gr, &up_all, &mut sl.act, n_ff, 1,
7566 &mut sl.actq, &mut sl.actd)?;
7567 }
7568 {
7569 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
7570 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
7571 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
7572 }
7573 e.rms_norm(&sl.f0, bits.post_ffw_norm.float_data(), &mut sl.sn, n_embd, 1, eps)?;
7574 match next_norm {
7575 Some(w) => {
7576 e.add_scale_rms_norm_q8_1_into(&sl.sn, &sl.attn_out, bits.layer_scale, w,
7577 &mut sl.xn, n_embd, 1, eps,
7578 &mut sl.hq, &mut sl.hd_)?;
7579 }
7580 None => {
7581 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
7582 }
7583 }
7584 Ok(())
7585 }
7586
7587 #[allow(clippy::too_many_arguments)]
7589 fn gemma4_decode_attn_dc(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
7590 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
7591 pos_d: &CudaSlice<i32>, cache: &mut Cache,
7592 cap_bucket_max: Option<(usize, usize)>)
7593 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7594 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7595 let eps = self.cfg.rms_eps;
7596 let aux = self.gemma4_aux.as_ref().unwrap();
7597 let (q0, k0, v0) = if swa {
7598 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
7599 Some(t3) => t3,
7600 None => {
7601 let h0 = e.zeros(0)?;
7602 (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
7603 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
7604 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?)
7605 }
7606 }
7607 } else {
7608 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
7609 Some(p) => p,
7610 None => {
7611 let h0 = e.zeros(0)?;
7612 (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
7613 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?)
7614 }
7615 };
7616 let v0 = e.clone_dtod(&k0)?;
7617 (q0, k0, v0)
7618 };
7619 let mut q = e.uninit(nh * hd)?;
7620 let mut k = e.uninit(nkv * hd)?;
7621 let mut v = e.uninit(nkv * hd)?;
7622 let ff = if swa { None } else {
7624 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
7625 };
7626 let kvl = cache.kv[il].as_mut().unwrap();
7627 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7628 if crate::Engine::qkv_append_on() {
7629 e.rms_norm_qkv_rope_append_dc(&q0, &k0, &v0, fa.q_norm.float_data(),
7631 fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7632 pos_d, nh, nkv, base, 1.0, ff, eps,
7633 &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
7634 } else {
7635 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7636 &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7637 pos_d, nh, nkv, base, 1.0, ff, eps)?;
7638 e.append_kv_quantized_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
7639 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
7640 }
7641 e.inc_seqlen(&mut kvl.len_d)?;
7642 let mut attn = e.uninit(nh * hd)?;
7643 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7646 match cap_bucket_max {
7651 None => {
7652 kvl.len += 1;
7656 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7657 if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
7658 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7659 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7662 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7663 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
7664 e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1,
7665 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7666 Some((&kvl.len_d, -1)), false, false,
7667 Some((&mut aq8, &mut ad8)))?;
7668 fa_q8 = Some((aq8, ad8));
7669 } else if swa && kvl.len > win && hd == 256
7670 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7671 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7673 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7674 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
7675 e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1,
7676 1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes,
7677 Some((&mut aq8, &mut ad8)))?;
7678 fa_q8 = Some((aq8, ad8));
7679 } else {
7680 let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) }
7681 else { (0, kvl.len) };
7682 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
7683 (off_tok + t_kv) * kvl.k_tok_bytes);
7684 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
7685 (off_tok + t_kv) * kvl.v_tok_bytes);
7686 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
7687 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
7688 }
7689 }
7690 Some((b_swa, b_glob)) => {
7691 let k_view = e.view_u8(&kvl.k, kvl.k.len());
7697 let v_view = e.view_u8(&kvl.v, kvl.v.len());
7698 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
7699 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7700 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
7701 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
7702 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, b_glob - 1,
7703 1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7704 Some((&kvl.len_d, -1)), false, false,
7705 Some((&mut aq8, &mut ad8)))?;
7706 fa_q8 = Some((aq8, ad8));
7707 } else if swa && b_swa > win && hd == 256 && rows_on {
7708 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
7709 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7710 &kvl.len_d, -1, 1, scale, win,
7711 kvl.k_tok_bytes, kvl.v_tok_bytes,
7712 Some((&mut aq8, &mut ad8)))?;
7713 fa_q8 = Some((aq8, ad8));
7714 } else {
7715 let b = if swa { b_swa } else { b_glob };
7716 e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, b,
7717 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7718 swa && crate::Engine::wkv_on())?;
7719 }
7720 }
7721 }
7722 if let Some((aq8, ad8)) = fa_q8 {
7725 let mut y = e.uninit(fa.wo.out_features())?;
7726 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
7727 return Ok(y);
7728 }
7729 Ok(e.matmul(&fa.wo, &attn, 1)?)
7730 }
7731
7732 pub fn gemma4_generate_graph(&self, e: &Engine, prompt_pos: usize, first_token: u32,
7737 cache: &mut Cache, max_new: usize, eos: &[u32],
7738 mut on_token: impl FnMut(u32) -> bool)
7739 -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
7740 if self.is_gemma4_e4b() {
7741 return Err("E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm".into());
7742 }
7743 use crate::decode::StopReason;
7744 let n_vocab = self.output.out_features();
7745 let n_embd = self.cfg.n_embd as usize;
7746 let embd_gpu = self.embd_gpu.get_or_init(|| {
7747 e.upload_u8(&self.embd.raw).expect("embed table upload")
7748 });
7749 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
7750 for kvl in cache.kv.iter_mut().flatten() {
7751 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7752 }
7753 let mut token_d = e.stream().clone_htod(&[first_token])?;
7754 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
7755 let g4 = self.cfg.gemma4.as_ref().unwrap();
7756 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
7757 let nkv_s = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
7759 .find(|p| *p.1).map(|p| *p.0 as usize).unwrap_or(8);
7760 let nkv_g = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
7761 .find(|p| !*p.1).map(|p| *p.0 as usize).unwrap_or(2);
7762 let mut graphs: std::collections::HashMap<((bool, usize), (bool, usize), bool, bool),
7763 (cudarc::driver::CudaGraph,
7764 Vec<Box<dyn std::any::Any + Send>>)> = Default::default();
7765 let mut slots = self.g4_dc_slots(e)?;
7768 const RING: usize = 64;
7771 const DRAIN: usize = 1;
7777 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
7778 let ring_base = prompt_pos;
7779 let mut out = Vec::with_capacity(max_new);
7780 let mut reason = StopReason::MaxNew;
7781 let mut next = first_token;
7782 let mut captures = 0usize;
7783 for _ in 0..max_new {
7784 out.push(next);
7785 if eos.contains(&next) { reason = StopReason::Eos; break; }
7786 if !on_token(next) { reason = StopReason::Callback; break; }
7787 let t_kv = cache.pos + 1;
7788 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7796 let f512 = crate::fa512_min_tkv();
7797 let key_s = if t_kv > win { (true, usize::MAX) }
7798 else { e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on()) };
7799 let (key_g, rung_end) = if t_kv >= f512 {
7800 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
7803 ((true, end), end)
7804 } else { (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv) };
7805 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
7806 if !graphs.contains_key(&key) {
7807 let bucket_max = (t_kv, rung_end);
7808 let snap = cache.snapshot(e)?;
7810 let pos_save = e.dtoh_i32_one(&pos_d)?;
7811 let len_save: Vec<Option<i32>> = cache.kv.iter()
7812 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
7813 let tok_save = e.dtoh_u32_one(&token_d)?;
7814 let graph = {
7819 let tok_ref = &mut token_d;
7820 let pos_ref = &mut pos_d;
7821 let cache_ref = &mut *cache;
7822 let slots_ref = &mut slots;
7823 let ring_ref = &mut ring;
7824 e.capture_graph_retained_flags(
7825 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
7826 |e| {
7827 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
7829 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
7830 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
7831 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
7832 cache_ref, n_vocab, Some(bucket_max),
7833 sl, tok_ref, Some((rg, ring_base)))
7834 })?
7835 };
7836 cache.rollback(e, &snap, 0)?;
7837 e.set_i32_one(&mut pos_d, pos_save)?;
7838 for (il, ls) in len_save.iter().enumerate() {
7839 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
7840 e.set_i32_one(&mut kvl.len_d, *v)?;
7841 }
7842 }
7843 e.set_u32_one(&mut token_d, tok_save)?;
7844 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
7845 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
7846 eprintln!("[graph-census] {c:?}");
7847 }
7848 }
7849 graphs.insert(key, graph);
7850 captures += 1;
7851 }
7852 let mut chunk = 1usize;
7857 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN").ok()
7858 .and_then(|v| v.parse().ok()).unwrap_or(DRAIN);
7859 while chunk < drain_cap && out.len() + chunk < max_new {
7860 let t_next = cache.pos + 1 + chunk;
7861 let key_s2 = if t_next > win { (true, usize::MAX) }
7862 else { e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on()) };
7863 let key_g2 = if t_next >= f512 {
7864 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
7865 } else { e.fa_bucket_key(t_next, hd_g, nkv_g, false) };
7866 if (key_s2, key_g2, t_next >= f512, t_next > win) != key { break; }
7867 chunk += 1;
7868 }
7869 let g = &graphs.get(&key).unwrap().0;
7870 for _ in 0..chunk { g.launch()?; }
7871 e.stream().synchronize()?;
7872 let ringh = e.dtoh_u32(&ring)?;
7873 for j in 0..chunk {
7874 let pos_j = cache.pos + j;
7875 let tok_j = ringh[(pos_j - ring_base) % RING];
7876 cache.pos += 0; if j + 1 == chunk { next = tok_j; }
7878 else {
7879 out.push(tok_j);
7880 if eos.contains(&tok_j) || !on_token(tok_j) {
7881 reason = if eos.contains(&tok_j) { StopReason::Eos }
7882 else { StopReason::Callback };
7883 let keep = cache.pos + j + 1;
7885 e.set_i32_one(&mut pos_d, keep as i32)?;
7886 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
7887 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
7888 kvl.len = keep;
7889 }
7890 cache.pos = keep;
7891 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
7892 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
7893 }
7894 return Ok((out, reason));
7895 }
7896 }
7897 }
7898 cache.pos += chunk;
7899 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) { kvl.len += chunk; }
7900 }
7901 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
7902 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
7903 }
7904 Ok((out, reason))
7905 }
7906
7907 pub(crate) fn gemma4_decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize,
7913 cache: &mut Cache)
7914 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7915 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
7916 }
7917
7918 pub(crate) fn gemma4_decode_step_t_am(&self, e: &Engine, tokens: &[u32], pos0: usize,
7922 cache: &mut Cache)
7923 -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7924 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
7925 let t = tokens.len();
7926 let n_vocab = self.output.out_features();
7927 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
7928 for i in 0..t {
7929 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
7930 }
7931 Ok((e.dtoh_u32(&toks)?, hn))
7932 }
7933
7934 pub(crate) fn gemma4_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
7937 pos0: usize, cache: &mut Cache)
7938 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7939 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
7940 let n_vocab = self.output.out_features();
7941 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
7942 for i in 0..t {
7943 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
7944 }
7945 Ok((vam, hn))
7946 }
7947
7948 pub(crate) fn gemma4_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
7951 cache: &mut Cache)
7952 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7953 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
7954 let t = tokens.len();
7955 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7956 e.softcap(&mut ld, cap, t * self.output.out_features())?;
7957 Ok((e.dtoh(&ld)?, hn))
7958 }
7959
7960 pub(crate) fn verify_stream_scratch(&self, e: &Engine, cap: usize)
7963 -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
7964 Ok(VerifyStreamScratch {
7965 pos_d: e.htod_i32(&vec![0i32; cap])?,
7966 row_ctrs: (0..cap).map(|_| e.htod_i32(&[0])).collect::<Result<_, _>>()?,
7967 })
7968 }
7969
7970 pub(crate) fn gemma4_verify_t_am_stream(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
7978 ctr: &CudaSlice<i32>, hint: usize,
7979 cache: &mut Cache,
7980 scr: &mut VerifyStreamScratch)
7981 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7982 let n_embd = self.cfg.n_embd as usize;
7983 let eps = self.cfg.rms_eps;
7984 assert!(t <= scr.row_ctrs.len() && t <= 64);
7985 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
7986 for i in 0..t {
7987 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
7988 }
7989 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
7990 let embd_gpu = self.embd_gpu.get_or_init(|| {
7991 e.upload_u8(&self.embd.raw).expect("embed table upload")
7992 });
7993 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
7994 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
7995 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7996 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7997 let n_layers = self.layers.len();
7998 for (il, layer) in self.layers.iter().enumerate() {
7999 let (hq, hdq) = match h_carry.take() {
8000 Some(p) => p,
8001 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
8002 };
8003 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8004 let o = self.gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache,
8005 hint, row_ctrs)?;
8006 let mut cur = e.uninit(t * n_embd)?;
8007 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
8008 let next_norm = if il + 1 < n_layers {
8009 Some(self.layers[il + 1].attn_norm.float_data())
8010 } else { None };
8011 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
8012 x = xn;
8013 h_carry = hn;
8014 self.dflash_tap(e, cache, il, &x, t)?;
8015 }
8016 let mut hn = e.uninit(t * n_embd)?;
8017 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8018 let ld = e.matmul(&self.output, &hn, t)?;
8019 let n_vocab = self.output.out_features();
8020 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
8021 for i in 0..t {
8022 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
8023 }
8024 Ok((vam, hn))
8025 }
8026
8027 fn dflash_tap(&self, e: &Engine, cache: &mut Cache, il: usize, x: &CudaSlice<f32>, t: usize)
8034 -> Result<(), Box<dyn std::error::Error>> {
8035 let Some(taps) = cache.dflash_taps.as_mut() else { return Ok(()) };
8036 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else { return Ok(()) };
8037 let h = taps.hidden;
8038 let n_taps = taps.layer_ids.len();
8039 debug_assert_eq!(taps.t, t);
8040 let xv = e.view(x, t * h);
8041 for r in 0..t {
8042 let row = xv.slice(r * h..(r + 1) * h);
8043 e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
8044 }
8045 Ok(())
8046 }
8047
8048 fn gemma4_verify_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
8049 tok_dev: Option<&CudaSlice<u32>>)
8050 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8051 let n_embd = self.cfg.n_embd as usize;
8052 let eps = self.cfg.rms_eps;
8053 let t = tokens.len();
8054 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8055 let pos_d = e.htod_i32(&pos)?;
8056 let mut x = match tok_dev {
8057 Some(td) => {
8058 let embd_gpu = self.embd_gpu.get_or_init(|| {
8059 e.upload_u8(&self.embd.raw).expect("embed table upload")
8060 });
8061 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8062 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
8063 }
8064 None => e.htod(&self.embd.gather(n_embd, tokens))?,
8065 };
8066 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
8067 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8068 let n_layers = self.layers.len();
8069 for (il, layer) in self.layers.iter().enumerate() {
8070 let (hq, hdq) = match h_carry.take() {
8071 Some(p) => p,
8072 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
8073 };
8074 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8075 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
8076 let mut cur = e.uninit(t * n_embd)?;
8077 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
8078 let next_norm = if il + 1 < n_layers {
8079 Some(self.layers[il + 1].attn_norm.float_data())
8080 } else { None };
8081 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
8082 x = xn;
8083 h_carry = hn;
8084 self.dflash_tap(e, cache, il, &x, t)?;
8085 }
8086 let mut hn = e.uninit(t * n_embd)?;
8087 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8088 let mut ld = e.matmul(&self.output, &hn, t)?;
8089 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
8091 Ok((ld, hn))
8092 }
8093
8094 #[allow(clippy::too_many_arguments)]
8102 fn gemma4_verify_attn_stream(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8103 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8104 pos_d: &CudaSlice<i32>, t: usize,
8105 cache: &mut Cache, hint: usize,
8106 row_ctrs: &[CudaSlice<i32>])
8107 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8108 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8109 let eps = self.cfg.rms_eps;
8110 let aux = self.gemma4_aux.as_ref().unwrap();
8111 let h0 = e.zeros(0)?;
8112 let h = &h0;
8113 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8116 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
8117 let fused_qkv = if f2b {
8118 if swa {
8119 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
8120 .map(|(a, b, c)| (a, b, Some(c)))
8121 } else {
8122 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
8123 .map(|(a, b)| (a, b, None))
8124 }
8125 } else { None };
8126 let (q0, k0, v0) = match fused_qkv {
8127 Some((a, b, cv)) => {
8128 let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
8129 (a, b, v)
8130 }
8131 None => {
8132 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
8133 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
8134 let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
8135 else { e.clone_dtod(&k0)? };
8136 (q0, k0, v0)
8137 }
8138 };
8139 let mut q = e.uninit(t * nh * hd)?;
8140 let mut k = e.uninit(t * nkv * hd)?;
8141 let mut v = e.uninit(t * nkv * hd)?;
8142 let ff = if swa { None } else {
8145 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
8146 };
8147 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8148 &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
8149 pos_d, nh, nkv, base, 1.0, ff, eps)?;
8150 let kvl = cache.kv[il].as_mut().unwrap();
8151 e.append_kv_quantized_rows_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d, t,
8153 kvl.kv_dim_k, kvl.kv_dim_v,
8154 kvl.k_tok_bytes, kvl.v_tok_bytes,
8155 (!swa && crate::Engine::gkv_on())
8156 || (swa && crate::Engine::wkv_on()))?;
8157 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8160 let mut attn = e.uninit(t * nh * hd)?;
8161 let k_view = e.view_u8(&kvl.k, kvl.k.len());
8162 let v_view = e.view_u8(&kvl.v, kvl.v.len());
8163 if swa && hint + 1 >= win {
8166 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8169 &kvl.len_d, 0, t, scale, win,
8170 kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8171 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
8172 let bucket = (hint + t + 2).next_power_of_two()
8185 .min(crate::fa512_min_tkv().saturating_sub(1));
8186 let qv = e.view(&q, t * nh * hd);
8187 for i in 0..t {
8188 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
8189 let mut q_one = e.uninit(nh * hd)?;
8190 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
8191 let mut a_one = e.uninit(nh * hd)?;
8192 e.fa_decode_dc(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv,
8193 &row_ctrs[i], bucket, scale,
8194 kvl.k_tok_bytes, kvl.v_tok_bytes, false)?;
8195 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
8196 }
8197 } else if hd == 512 {
8198 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, hint, t, scale,
8201 kvl.k_tok_bytes, kvl.v_tok_bytes,
8202 Some((&kvl.len_d, 0)), false, false, None)?;
8203 } else {
8204 e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8206 &kvl.len_d, hint + t, t, scale,
8207 kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
8208 swa && crate::Engine::wkv_on())?;
8209 }
8210 Ok(e.matmul(&fa.wo, &attn, t)?)
8211 }
8212
8213 fn gemma4_verify_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8214 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8215 pos_d: &CudaSlice<i32>, t: usize,
8216 cache: &mut Cache)
8217 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8218 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8219 let eps = self.cfg.rms_eps;
8220 let aux = self.gemma4_aux.as_ref().unwrap();
8221 let n_embd = self.cfg.n_embd as usize;
8222 let _ = n_embd;
8223
8224 let h0 = e.zeros(0)?;
8225 let h = &h0;
8226 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8229 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
8230 let fused_qkv = if f2b {
8231 if swa {
8232 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
8233 .map(|(a, b, c)| (a, b, Some(c)))
8234 } else {
8235 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
8236 .map(|(a, b)| (a, b, None))
8237 }
8238 } else { None };
8239 let (q0, k0, v0) = match fused_qkv {
8240 Some((a, b, cv)) => {
8241 let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
8242 (a, b, v)
8243 }
8244 None => {
8245 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
8246 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
8247 let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
8248 else { e.clone_dtod(&k0)? };
8249 (q0, k0, v0)
8250 }
8251 };
8252 let mut q = e.uninit(t * nh * hd)?;
8253 let mut k = e.uninit(t * nkv * hd)?;
8254 let mut v = e.uninit(t * nkv * hd)?;
8255 let ff = if swa { None } else {
8258 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
8259 };
8260 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8261 &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
8262 pos_d, nh, nkv, base, 1.0, ff, eps)?;
8263 let kvl = cache.kv[il].as_mut().unwrap();
8264 let base_len = kvl.len;
8265 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
8266 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()))?;
8267 kvl.len += t;
8268 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8269 let mut attn = e.uninit(t * nh * hd)?;
8270 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
8273 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
8276 if rows_ok && (!swa || base_len + t <= win) {
8277 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
8278 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
8279 if hd == 512 {
8280 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8282 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, base_len, t,
8283 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8284 Some((&kvl.len_d, 0)), false,
8285 swa && crate::Engine::wkv_on(), None)?;
8286 } else {
8287 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8291 e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8292 &kvl.len_d, base_len + t, t, scale,
8293 kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
8294 swa && crate::Engine::wkv_on())?;
8295 }
8296 return Ok(e.matmul(&fa.wo, &attn, t)?);
8297 }
8298 if hd == 256 && swa && base_len + 1 >= win
8306 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8307 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
8308 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
8309 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8310 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, 0,
8311 t, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8312 return Ok(e.matmul(&fa.wo, &attn, t)?);
8313 }
8314 for i in 0..t {
8315 let avail = base_len + i + 1;
8316 let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
8317 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
8318 (off_tok + t_kv) * kvl.k_tok_bytes);
8319 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
8320 (off_tok + t_kv) * kvl.v_tok_bytes);
8321 let qi = e.view(&q, t * nh * hd);
8322 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
8323 let mut q_one = e.uninit(nh * hd)?;
8324 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
8325 let mut a_one = e.uninit(nh * hd)?;
8326 if swa && avail > win && hd == 256
8330 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8331 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
8332 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
8333 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
8334 e.fa_decode_rows_w(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, &kvl.len_d, 0,
8335 1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8336 } else if !swa && hd == 512 && avail >= crate::fa512_min_tkv()
8337 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8338 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
8339 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
8340 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
8341 e.fa_decode_rows(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, avail - 1, 1,
8342 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8343 Some((&kvl.len_d, 0)), false, false, None)?;
8344 } else {
8345 e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
8346 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
8347 }
8348 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
8349 }
8350 Ok(e.matmul(&fa.wo, &attn, t)?)
8351 }
8352
8353 pub(crate) fn gemma4_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
8356 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8357 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
8362 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
8363 }
8364 if crate::pp::pp_cuts(self.layers.len()).is_some() {
8365 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
8366 }
8367 let n_embd = self.cfg.n_embd as usize;
8368 let eps = self.cfg.rms_eps;
8369 let pos_d = e.htod_i32(&[cache.pos as i32])?;
8370 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
8371 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8372 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8375 let n_layers = self.layers.len();
8376 for (il, layer) in self.layers.iter().enumerate() {
8377 let (hq, hdq) = match h_carry.take() {
8378 Some(p) => p,
8379 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
8380 };
8381 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8382 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
8383 let mut cur = e.uninit(n_embd)?;
8384 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
8385 let next_norm = if il + 1 < n_layers {
8386 Some(self.layers[il + 1].attn_norm.float_data())
8387 } else { None };
8388 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
8389 x = xn;
8390 h_carry = hn;
8391 }
8392 let mut hn = e.uninit(n_embd)?;
8393 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8394 let h_seed = e.clone_dtod(&x)?;
8395 let mut ld = e.matmul(&self.output, &hn, 1)?;
8396 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8397 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
8399 let logits = e.dtoh(&ld)?;
8400 cache.pos += 1;
8401 Ok((logits, h_seed))
8402 }
8403
8404 fn gemma4_decode_layers(&self, e: &Engine, mut x: CudaSlice<f32>, lo: usize, hi: usize,
8412 pos_d: &CudaSlice<i32>, cache: &mut Cache)
8413 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8414 let n_embd = self.cfg.n_embd as usize;
8415 let eps = self.cfg.rms_eps;
8416 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8417 for il in lo..hi {
8418 let layer = &self.layers[il];
8419 let (hq, hdq) = match h_carry.take() {
8420 Some(p) => p,
8421 None => e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?,
8423 };
8424 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8425 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
8426 let mut cur = e.uninit(n_embd)?;
8427 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
8428 let next_norm = if il + 1 < hi {
8429 Some(self.layers[il + 1].attn_norm.float_data())
8430 } else { None };
8431 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
8432 x = xn;
8433 h_carry = hn;
8434 }
8435 Ok(x)
8436 }
8437
8438 fn gemma4_decode_step_h_pp2(&self, e: &Engine, token: u32, cache: &mut Cache, split: usize)
8445 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8446 if crate::pp::pp2_streams_off() {
8447 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
8448 }
8449 let rt = crate::pp::Pp2Rt::get(e)?;
8450 let e0 = rt.engine(0, e);
8451 let e1 = rt.engine(1, e);
8452 let n_embd = self.cfg.n_embd as usize;
8453 let eps = self.cfg.rms_eps;
8454
8455 let (pos_d, slot) = {
8457 let _st0 = rt.enter(0);
8458 let pos_d = e0.htod_i32(&[cache.pos as i32])?;
8459 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
8460 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8461 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
8462 let slot = rt.tx(0, &x, n_embd)?;
8463 (pos_d, slot)
8464 };
8465
8466 let _st1 = rt.enter(1);
8468 let x = rt.rx(0, slot, n_embd)?;
8469 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
8470
8471 let mut hn = e1.uninit(n_embd)?;
8472 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8473 let h_seed = e1.clone_dtod(&x)?;
8474 let mut ld = e1.matmul(&self.output, &hn, 1)?;
8475 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8476 e1.softcap(&mut ld, cap, self.output.out_features())?;
8477 self.gemma4_suppress(e1, &mut ld, 1)?;
8478 let logits = e1.dtoh(&ld)?;
8479 cache.pos += 1;
8480 Ok((logits, h_seed))
8481 }
8482
8483 fn gemma4_decode_step_h_pp2_samestream(&self, e: &Engine, token: u32, cache: &mut Cache,
8486 split: usize)
8487 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8488 let n_embd = self.cfg.n_embd as usize;
8489 let eps = self.cfg.rms_eps;
8490 let pos_d = e.htod_i32(&[cache.pos as i32])?;
8491
8492 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
8494 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8495 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
8496
8497 let boundary_tx = e.clone_dtod(&x)?;
8499 let boundary_rx = e.clone_dtod(&boundary_tx)?;
8500
8501 let x = self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
8503
8504 let mut hn = e.uninit(n_embd)?;
8505 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8506 let h_seed = e.clone_dtod(&x)?;
8507 let mut ld = e.matmul(&self.output, &hn, 1)?;
8508 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8509 e.softcap(&mut ld, cap, self.output.out_features())?;
8510 self.gemma4_suppress(e, &mut ld, 1)?;
8511 let logits = e.dtoh(&ld)?;
8512 cache.pos += 1;
8513 Ok((logits, h_seed))
8514 }
8515}
8516
8517impl HybridModel {
8536 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
8539 let geometry = self.cfg.layer_geometry(il as u32)
8540 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
8541 debug_assert_eq!(
8542 geometry.attention_gate,
8543 memra_gguf::config::AttentionGateKind::SeparateHead
8544 );
8545 geometry
8546 }
8547
8548 #[allow(clippy::too_many_arguments)]
8608 fn step35_attn_pre_wo(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
8609 hg: Option<&CudaSlice<f32>>, gt_pre: Option<&CudaSlice<f32>>,
8610 pos_d: &CudaSlice<i32>, t: usize,
8611 cache: Option<&mut Cache>, il: usize, seq_end: usize)
8612 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8613 let geometry = self.step35_geom(il);
8614 let hd = geometry.head_dim_k as usize;
8615 let nkv = geometry.n_head_kv as usize;
8616 let nh = geometry.n_head as usize;
8617 let rbase = geometry.rope_base;
8618 let scale = geometry.attention_scale();
8619 let swa = geometry.window.is_some();
8620 let eps = self.cfg.rms_eps;
8621 let win = geometry.window.unwrap_or(0) as usize;
8622 let n_rot = geometry.n_rot as usize;
8623
8624 let v = g3.pop().unwrap();
8625 let k0 = g3.pop().unwrap();
8626 let q0 = g3.pop().unwrap();
8627
8628 let mut q = e.uninit(t * nh * hd)?;
8632 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
8633 let mut k = e.uninit(t * nkv * hd)?;
8634 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
8635 let ff = if geometry.rope_factors {
8636 self.step35_aux.as_ref().and_then(|a| a.rope_freqs.as_ref())
8637 } else {
8638 None
8639 };
8640 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
8641
8642 let mut attn = e.uninit(t * nh * hd)?;
8643 match cache {
8644 Some(cache) => {
8645 let base_len = {
8646 let kvl = cache.kv[il].as_mut().unwrap();
8647 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
8648 let base_len = kvl.len;
8649 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
8650 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
8651 kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
8652 kvl.len += t;
8653 let new_len = kvl.len as i32;
8654 e.set_i32_one(&mut kvl.len_d, new_len)?;
8655 base_len
8656 };
8657 let kvl = cache.kv[il].as_ref().unwrap();
8658 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
8668 let legacy_calllocal =
8669 std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
8670 let off = if swa {
8685 let raw = base_len.saturating_sub(win - 1);
8686 if legacy_tkv || legacy_calllocal { raw } else { raw & !31usize }
8687 } else {
8688 0
8689 };
8690 let t_kv = base_len + t - off;
8691 let k_view = e.view_u8_range(&kvl.k, off * kvl.k_tok_bytes,
8692 (off + t_kv) * kvl.k_tok_bytes);
8693 let v_view = e.view_u8_range(&kvl.v, off * kvl.v_tok_bytes,
8694 (off + t_kv) * kvl.v_tok_bytes);
8695 let swa_naive = if legacy_tkv { t_kv > win } else { seq_end > win };
8707 if swa && swa_naive {
8708 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
8721 e.sdpa_naive_w_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh,
8722 nkv, t, t_kv, scale, true, win,
8723 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
8724 } else {
8725 e.fa_prefill_view_ws_w_hd128(&q, &k_view, &v_view, &mut attn, hd, nh,
8726 nkv, t, t_kv, scale, true, win,
8727 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
8728 }
8729 } else if std::env::var("MEMRA_NOFA").is_ok() {
8730 e.sdpa_naive_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8731 t, t_kv, scale, true,
8732 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
8733 } else {
8734 e.fa_prefill_view_ws(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8739 t, t_kv, scale, true,
8740 kvl.k_tok_bytes, kvl.v_tok_bytes,
8741 crate::Engine::kv_fp8_on())?;
8742 }
8743 }
8744 None => {
8745 debug_assert_eq!(seq_end, t, "step35 cacheless prefill is monolithic (seq_end == t)");
8750 if swa && seq_end > win {
8751 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8752 } else if std::env::var("MEMRA_NOFA").is_ok() {
8753 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8754 } else {
8755 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8756 }
8757 }
8758 }
8759
8760 let gw = fa.attn_gate.as_ref()
8763 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
8764 let gt_owned = if gt_pre.is_none() {
8765 Some(e.matmul(
8766 gw,
8767 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
8768 t,
8769 )?)
8770 } else {
8771 None
8772 };
8773 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
8774 let mut ag = e.uninit(t * nh * hd)?;
8775 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
8776 Ok(ag)
8777 }
8778
8779 pub(crate) fn step35_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
8782 pos_d: &CudaSlice<i32>, t: usize, il: usize)
8783 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8784 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
8785 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
8787 Ok(e.matmul(&fa.wo, &ag, t)?)
8788 }
8789
8790 #[allow(clippy::too_many_arguments)]
8797 pub(crate) fn step35_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
8798 hx: Option<&CudaSlice<u8>>, pos_d: &CudaSlice<i32>, t: usize,
8799 cache: &mut Cache, il: usize, seq_end: usize)
8800 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8801 let g3 = match hx {
8802 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
8803 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
8804 };
8805 let ag = self.step35_attn_pre_wo(
8806 e,
8807 fa,
8808 g3,
8809 Some(h),
8810 None,
8811 pos_d,
8812 t,
8813 Some(cache),
8814 il,
8815 seq_end,
8816 )?;
8817 Ok(e.matmul(&fa.wo, &ag, t)?)
8818 }
8819
8820 #[allow(clippy::too_many_arguments)]
8830 pub(crate) fn step35_decode_attn(&self, e: &Engine, fa: &FullAttnLayer, il: usize,
8831 h: &CudaSlice<f32>,
8832 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8833 pos_d: &CudaSlice<i32>, cache: &mut Cache)
8834 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8835 let geometry = self.step35_geom(il);
8836 let hd = geometry.head_dim_k as usize;
8837 let nkv = geometry.n_head_kv as usize;
8838 let nh = geometry.n_head as usize;
8839 let rbase = geometry.rope_base;
8840 let scale = geometry.attention_scale();
8841 let swa = geometry.window.is_some();
8842 let eps = self.cfg.rms_eps;
8843 let win = geometry.window.unwrap_or(0) as usize;
8844 let n_rot = geometry.n_rot as usize;
8845 let n_embd = self.cfg.n_embd as usize;
8846 let gw = fa.attn_gate.as_ref()
8847 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
8848
8849 let (q0, k0, v0, gt) = match pre_q {
8850 Some((hq, hdq)) => {
8851 debug_assert!(e.uses_q8_1_fast(gw),
8852 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
8853 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast");
8854 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
8855 Some(t3) => t3,
8856 None => (e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
8857 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
8858 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?),
8859 };
8860 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
8861 (a, b, c, gt)
8862 }
8863 None => {
8864 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk)
8865 && e.uses_q8_1_fast(&fa.wv) && e.uses_q8_1_fast(gw) {
8866 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
8867 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
8868 Some(t3) => t3,
8869 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
8870 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
8871 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
8872 };
8873 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
8874 (a, b, c, gt)
8875 } else {
8876 (e.matmul(&fa.wq, h, 1)?, e.matmul(&fa.wk, h, 1)?,
8877 e.matmul(&fa.wv, h, 1)?, e.matmul(gw, h, 1)?)
8878 }
8879 }
8880 };
8881
8882 let mut q = e.uninit(nh * hd)?;
8883 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
8884 let mut k = e.uninit(nkv * hd)?;
8885 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
8886 let ff = if swa { None } else {
8887 self.step35_aux.as_ref().and_then(|a| a.rope_freqs.as_ref())
8888 };
8889 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
8890
8891 if std::env::var("MEMRA_NOFA").is_ok() {
8892 return Err("MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
8893 cache; unset MEMRA_NOFA to use fa_decode".into());
8894 }
8895 let kvl = cache.kv[il].as_mut().unwrap();
8896 e.append_kv_quantized(&k, &v0, &mut kvl.k, &mut kvl.v, kvl.len,
8897 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
8898 crate::Engine::kv_fp8_on())?;
8899 kvl.len += 1;
8900 let (off, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
8901 let k_view = e.view_u8_range(&kvl.k, off * kvl.k_tok_bytes,
8902 (off + t_kv) * kvl.k_tok_bytes);
8903 let v_view = e.view_u8_range(&kvl.v, off * kvl.v_tok_bytes,
8904 (off + t_kv) * kvl.v_tok_bytes);
8905 let mut attn = e.uninit(nh * hd)?;
8906 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
8907 kvl.k_tok_bytes, kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
8908
8909 let mut ag = e.uninit(nh * hd)?;
8910 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
8911 Ok(e.matmul(&fa.wo, &ag, 1)?)
8912 }
8913}
8914
8915impl HybridModel {
8924 pub fn is_gemma4_e4b(&self) -> bool {
8925 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
8926 }
8927
8928 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8932 let g = self.cfg.gemma4.as_ref().unwrap();
8933 let swa = g.swa_pattern[il];
8934 let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
8935 let Mixer::Full(fa) = &self.layers[il].mixer else { panic!("e4b layer {il} not full-attn") };
8936 let nh = fa.wq.out_features() / hd;
8937 let nkv = fa.wk.out_features() / hd;
8938 (hd, nkv, nh, if swa { g.rope_base_swa } else { g.rope_base_global }, 1.0, swa)
8939 }
8940
8941 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
8943 self.layers[il].gemma4.as_ref()
8944 .and_then(|b| b.e4b.as_ref())
8945 .and_then(|e4| e4.kv_share.map(|t| t as usize))
8946 }
8947
8948 fn gemma4_e4b_inp_pl(&self, e: &Engine, tokens: &[u32], x_scaled: &CudaSlice<f32>, t: usize)
8953 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8954 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
8955 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
8956 }
8957
8958 fn gemma4_e4b_inp_pl_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
8960 x_scaled: &CudaSlice<f32>, t: usize)
8961 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8962 let aux = self.gemma4_aux.as_ref().unwrap();
8963 let m = aux.e4b.as_ref().unwrap();
8964 let n_embd = self.cfg.n_embd as usize;
8965 let n_layer = self.layers.len();
8966 let width = m.n_epl * n_layer;
8967 let tbl = m.tok_tbl_gpu.get_or_init(|| {
8968 e.upload_u8(&m.tok_embd_bytes).expect("e4b per-layer token table upload")
8969 });
8970 let mut a = e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt,
8971 m.tok_embd_row_bytes)?;
8972 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
8973 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
8974 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
8975 let mut pn = e.uninit(t * width)?;
8976 e.rms_norm(&p, m.proj_norm.float_data(), &mut pn, m.n_epl, t * n_layer,
8977 self.cfg.rms_eps)?;
8978 let mut out = e.uninit(t * width)?;
8979 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
8980 Ok(out)
8981 }
8982
8983 #[allow(clippy::too_many_arguments)]
8988 fn gemma4_e4b_attn(&self, e: &Engine, il: usize,
8989 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8990 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
8991 dc_bucket: Option<usize>)
8992 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8993 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
8994 let eps = self.cfg.rms_eps;
8995 let aux = self.gemma4_aux.as_ref().unwrap();
8996 let Mixer::Full(fa) = &self.layers[il].mixer else { unreachable!() };
8997 let h0 = e.zeros(0)?;
9001 let h = &h0;
9002
9003 let ff = if swa { None } else {
9004 Some(aux.rope_freqs.as_ref().expect("e4b global rope needs rope_freqs.weight"))
9005 };
9006 let share = self.gemma4_e4b_kv_target(il);
9007 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
9009 let mut q;
9010 if let Some(_tgt) = share {
9011 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
9012 q = e.uninit(t * nh * hd)?;
9013 let mut kdummy = e.uninit(1)?;
9016 let mut vdummy = e.uninit(1)?;
9017 e.rms_norm_qkv_rope(&q0, &q0, &q0, fa.q_norm.float_data(),
9018 fa.q_norm.float_data(), &aux.ones,
9019 &mut q, &mut kdummy, &mut vdummy, hd, nh * t, 0,
9020 pos_d, nh, 1, base, 1.0, ff, eps)?;
9021 } else {
9022 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
9026 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
9027 q = e.uninit(t * nh * hd)?;
9028 let mut k = e.uninit(t * nkv * hd)?;
9029 let mut v = e.uninit(t * nkv * hd)?;
9030 if t == 1 && cat.is_some() {
9031 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
9032 e.rms_norm_qkv_rope_cat(&qkv0, fa.q_norm.float_data(), fa.k_norm.float_data(),
9033 &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
9034 pos_d, nh, nkv, base, 1.0, ff, eps)?;
9035 } else {
9036 let (q0, k0, v0) = match if t == 1 {
9037 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
9038 } else {
9039 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9042 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
9043 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
9044 } else { None }
9045 } {
9046 Some(triple) => triple,
9047 None => (e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
9048 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
9049 e.matmul_pre(&fa.wv, hq, hdq, h, t)?), };
9051 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(),
9054 fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v,
9055 hd, nh * t, nkv * t, pos_d, nh, nkv, base, 1.0, ff, eps)?;
9056 }
9057 let kvl = cache.kv[il].as_mut().unwrap();
9058 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9062 if dc_bucket.is_some() {
9063 debug_assert!(t == 1);
9068 e.append_kv_quantized_row_dc_inc(&k, &v, &mut kvl.k, &mut kvl.v,
9070 &mut kvl.len_d, kvl.kv_dim_k, kvl.kv_dim_v,
9071 kvl.k_tok_bytes, kvl.v_tok_bytes, cls)?;
9072 } else {
9073 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
9074 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
9075 kvl.v_tok_bytes, cls)?;
9076 kvl.len += t;
9077 }
9078 kv_f32 = Some((k, v));
9079 }
9080 let kvl_idx = share.unwrap_or(il);
9083 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
9084 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
9086 let mut attn = e.uninit(t * nh * hd)?;
9087 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
9099 if let Some((kf, vf)) = &kv_f32 {
9100 if hd == 256 && t <= win {
9101 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9102 return Ok(e.matmul(&fa.wo, &attn, t)?);
9103 }
9104 if hd == 256 && swa && t > win {
9105 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true,
9106 win)?;
9107 return Ok(e.matmul(&fa.wo, &attn, t)?);
9108 }
9109 if hd == 512 && !swa {
9110 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale,
9111 true)?;
9112 return Ok(e.matmul(&fa.wo, &attn, t)?);
9113 }
9114 } else if share.is_some() {
9115 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9116 let k_view = e.view_u8(&kvl.k, kvl.k.len());
9117 let v_view = e.view_u8(&kvl.v, kvl.v.len());
9118 if hd == 256 && (!swa || t <= win) {
9119 e.fa_prefill_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t, t,
9121 scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9122 return Ok(e.matmul(&fa.wo, &attn, t)?);
9123 }
9124 let kv_dim = nkv * hd;
9127 let mut kf = e.uninit(t * kv_dim)?;
9128 let mut vf = e.uninit(t * kv_dim)?;
9129 e.fa_dequant_kv_view_f32(&k_view, &v_view, &mut kf, &mut vf, kv_dim, kv_dim,
9130 t, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9131 if hd == 512 {
9132 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale,
9133 true)?;
9134 } else {
9135 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true,
9136 win)?;
9137 }
9138 return Ok(e.matmul(&fa.wo, &attn, t)?);
9139 }
9140 }
9141 if let Some(bucket) = dc_bucket {
9142 assert!(t == 1);
9147 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
9153 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
9154 } else { bucket };
9155 let k_view = e.view_u8(&kvl.k, kvl.k.len());
9156 let v_view = e.view_u8(&kvl.v, kvl.v.len());
9157 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9158 if crate::Engine::wpf_level() >= 1 {
9166 e.prefetch_weight_l2(&fa.wo)?;
9167 }
9168 if e.uses_q8_1_fast(&fa.wo) {
9171 let mut oq = e.alloc_i8_uninit(nh * hd)?;
9172 let mut od = e.zeros(nh * hd / 32)?;
9173 e.fa_decode_dc_q8(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9174 &kvl.len_d, bucket, scale,
9175 kvl.k_tok_bytes, kvl.v_tok_bytes, g,
9176 Some((&mut oq, &mut od)))?;
9177 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
9178 }
9179 e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9180 &kvl.len_d, bucket, scale,
9181 kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9182 return Ok(e.matmul(&fa.wo, &attn, t)?);
9183 }
9184 for i in 0..t {
9185 let avail = base_len + i + 1;
9186 let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
9187 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
9188 (off_tok + t_kv) * kvl.k_tok_bytes);
9189 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
9190 (off_tok + t_kv) * kvl.v_tok_bytes);
9191 let qv = e.view(&q, t * nh * hd);
9192 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
9193 let mut q_one = e.uninit(nh * hd)?;
9194 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
9195 let mut a_one = e.uninit(nh * hd)?;
9196 e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
9200 kvl.k_tok_bytes, kvl.v_tok_bytes,
9201 (!swa && crate::Engine::gkv_on())
9202 || (swa && crate::Engine::wkv_on()))?;
9203 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
9204 }
9205 Ok(e.matmul(&fa.wo, &attn, t)?)
9206 }
9207
9208 fn gemma4_e4b_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
9213 head_last: bool)
9214 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9215 let n_embd = self.cfg.n_embd as usize;
9216 let t = tokens.len();
9217 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9218 let pos_d = e.htod_i32(&pos)?;
9219 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
9220 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9221 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
9222 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
9223 }
9224
9225 fn gemma4_e4b_trunk_core(&self, e: &Engine, x_in: CudaSlice<f32>, inp_pl: CudaSlice<f32>,
9229 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
9230 dc_bucket: Option<usize>, cap_logits: bool, head_last: bool)
9231 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9232 let n_embd = self.cfg.n_embd as usize;
9233 let eps = self.cfg.rms_eps;
9234 let n_layer = self.layers.len();
9235 let mut x = x_in;
9236 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
9237 let n_epl = aux_e4b.n_epl;
9238
9239 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9245 for il in 0..n_layer {
9246 let layer = &self.layers[il];
9247 let (hq, hdq) = match h_carry.take() {
9248 Some(p) => p,
9249 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
9250 };
9251 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
9252 let bits = layer.gemma4.as_ref().unwrap();
9255 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
9256 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
9267 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
9268 e, layer, &o, &x, t, Some(layer.post_attn_norm.float_data()), fuse_exit)?;
9269 let mut resid = e.uninit(t * n_embd)?;
9270 let g = if fuse_exit {
9276 let (rq, rd) = e.rms_pre_add_q8_1(&sn, bits.post_ffw_norm.float_data(),
9278 &attn_out, &mut resid, n_embd, t,
9279 self.cfg.rms_eps)?;
9280 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
9281 } else {
9282 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
9283 e.matmul(&e4b.inp_gate, &resid, t)?
9284 };
9285 let mut act = e.uninit(t * n_epl)?;
9286 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
9287 let ipv = e.view(&inp_pl, n_epl * n_layer);
9288 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
9289 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
9290 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
9291 } else {
9292 let mut inp_this = e.uninit(t * n_epl)?;
9293 e.copy_rows_strided(&inp_pl, &mut inp_this, n_epl, t, n_epl * n_layer,
9294 il * n_epl)?;
9295 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
9296 e.matmul(&e4b.proj, &act, t)?
9297 };
9298 let next_norm = if il + 1 < n_layer {
9301 self.layers[il + 1].attn_norm.float_data()
9302 } else {
9303 self.output_norm.float_data()
9304 };
9305 let mut xn = e.uninit(t * n_embd)?;
9306 let pair = e.rms_pre_add_scale_rms_norm_q8_1(&y, e4b.post_norm.float_data(),
9307 &resid, bits.layer_scale, next_norm,
9308 &mut xn, n_embd, t, eps)?;
9309 h_carry = Some(pair);
9310 x = xn;
9311 }
9312 let (oq, odq) = h_carry.take().unwrap();
9316 let h0 = e.zeros(0)?;
9317 let hm = if head_last { 1 } else { t };
9318 let (hq, hd) = if head_last && t > 1 {
9319 let mut q1 = e.uninit_i8(n_embd)?;
9320 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
9321 let nb = n_embd / 32;
9322 let mut d1 = e.uninit(nb)?;
9323 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
9324 (q1, d1)
9325 } else {
9326 (oq, odq)
9327 };
9328 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
9329 if cap_logits {
9333 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9334 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
9335 }
9336 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
9338 }
9339
9340 pub fn gemma4_e4b_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
9347 t: usize, pos0: usize, cache: &mut Cache)
9348 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9349 let n_embd = self.cfg.n_embd as usize;
9350 let eps = self.cfg.rms_eps;
9351 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9352 let pos_d = e.htod_i32(&pos)?;
9353 let embd_gpu = self.embd_gpu.get_or_init(|| {
9354 e.upload_u8(&self.embd.raw).expect("embed table upload")
9355 });
9356 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
9357 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
9358 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9359 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
9360 let (ld, xp) = self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true,
9361 false)?;
9362 let n_vocab = self.output.out_features();
9365 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
9366 for i in 0..t {
9367 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
9368 }
9369 let mut hn = e.uninit(t * n_embd)?;
9370 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9371 cache.pos += t;
9372 Ok((vam, hn))
9373 }
9374
9375 pub(crate) fn gemma4_e4b_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
9378 cache: &mut Cache)
9379 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9380 let n_embd = self.cfg.n_embd as usize;
9381 let eps = self.cfg.rms_eps;
9382 let t = tokens.len();
9383 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
9384 let mut hn = e.uninit(t * n_embd)?;
9385 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9386 cache.pos += t;
9387 Ok((e.dtoh(&ld)?, hn))
9388 }
9389
9390 pub fn gemma4_e4b_decode_step_dcg(&self, e: &Engine, token_d: &mut CudaSlice<u32>,
9396 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
9397 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
9398 n_vocab: usize, bucket: usize)
9399 -> Result<(), Box<dyn std::error::Error>> {
9400 let n_embd = self.cfg.n_embd as usize;
9401 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
9402 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
9403 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
9404 let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket),
9405 false, false)?;
9406 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
9407 e.inc_seqlen(pos_d)?;
9408 Ok(())
9409 }
9410
9411 #[allow(clippy::too_many_arguments)]
9419 pub fn gemma4_e4b_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
9420 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
9421 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
9422 n_vocab: usize)
9423 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9424 let n_embd = self.cfg.n_embd as usize;
9425 let eps = self.cfg.rms_eps;
9426 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
9427 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
9428 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
9429 let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false,
9430 false)?;
9431 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
9432 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
9433 e.inc_seqlen(pos_d)?;
9434 cache.pos += 1;
9435 let _ = eps;
9436 Ok(tok_out)
9437 }
9438
9439 pub(crate) fn gemma4_e4b_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
9442 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9443 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
9444 let logits = e.dtoh(&ld)?;
9445 cache.pos += 1;
9446 Ok((logits, x))
9447 }
9448
9449 pub(crate) fn gemma4_e4b_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
9453 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9454 if cache.pos != 0 {
9457 return Err("e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
9458 call or decode tokenwise".into());
9459 }
9460 let n_embd = self.cfg.n_embd as usize;
9461 let t = tokens.len();
9462 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
9463 cache.pos += t;
9464 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
9466 let row = xv.slice((t - 1) * n_embd..t * n_embd);
9467 let mut h_seed = e.uninit(n_embd)?;
9468 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
9469 Ok((last, h_seed, x))
9470 }
9471
9472 pub(crate) fn gemma4_e4b_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
9474 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
9475 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
9476 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
9477 Ok(e.dtoh(&ld)?) }
9479}
9480
9481#[cfg(test)]
9482mod prime_chunk_schedule_tests {
9483 use super::{
9484 dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges, PRIME_MIN_T,
9485 PRIME_PIPE_MIN_CHUNK,
9486 };
9487
9488 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
9489 ranges.iter().map(|(start, end)| end - start).collect()
9490 }
9491
9492 fn auto_chunk(t: usize) -> usize {
9493 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
9494 }
9495
9496 #[test]
9497 fn fixed_schedule_retains_measured_geometry() {
9498 assert_eq!(
9499 sizes(&fixed_prime_chunk_ranges(461, 128)),
9500 vec![128, 128, 128, 77]
9501 );
9502 assert_eq!(
9503 sizes(&fixed_prime_chunk_ranges(1833, 230)),
9504 vec![230, 230, 230, 230, 230, 230, 230, 223]
9505 );
9506 assert_eq!(
9507 sizes(&fixed_prime_chunk_ranges(4096, 512)),
9508 vec![512; 8]
9509 );
9510 }
9511
9512 #[test]
9513 fn dynamic_schedule_matches_registered_shapes() {
9514 let cases = [
9515 (461, vec![64, 141, 132, 124]),
9516 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
9517 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
9518 ];
9519 for (t, expected) in cases {
9520 let chunk = auto_chunk(t);
9521 let fixed = fixed_prime_chunk_ranges(t, chunk);
9522 assert_eq!(
9523 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
9524 expected
9525 );
9526 }
9527 }
9528
9529 #[test]
9530 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
9531 for t in 256..=8192 {
9532 let chunk = auto_chunk(t);
9533 let fixed = fixed_prime_chunk_ranges(t, chunk);
9534 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
9535 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
9536 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
9537 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
9538 for pair in dynamic.windows(2) {
9539 assert_eq!(pair[0].1, pair[1].0, "T={t}");
9540 }
9541 assert!(
9542 dynamic
9543 .iter()
9544 .all(|(start, end)| end - start >= PRIME_MIN_T),
9545 "T={t} sizes={:?}",
9546 sizes(&dynamic)
9547 );
9548 if dynamic.len() >= 3 {
9549 let chunk_sizes = sizes(&dynamic);
9550 assert!(
9551 chunk_sizes[0] < chunk_sizes[1],
9552 "T={t} sizes={chunk_sizes:?}"
9553 );
9554 assert!(
9555 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
9556 "T={t} sizes={chunk_sizes:?}"
9557 );
9558 }
9559 }
9560 }
9561}
9562
9563#[cfg(test)]
9564mod page_prefetch_tests {
9565 use super::{
9566 grouped_worker_prefetch_position, page_prefetch_positions,
9567 page_prefetch_window_from_values, worker_prefetch_positions,
9568 };
9569
9570 #[test]
9571 fn page_prefetch_window_keeps_existing_opt_in_default() {
9572 assert_eq!(page_prefetch_window_from_values(false, None), 0);
9573 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
9574 assert_eq!(page_prefetch_window_from_values(true, None), 1);
9575 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
9576 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
9577 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
9578 }
9579
9580 #[test]
9581 fn rolling_page_prefetch_advises_each_future_expert_once() {
9582 let advised: Vec<_> = (0..7)
9583 .flat_map(|position| page_prefetch_positions(position, 7, 3))
9584 .collect();
9585 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
9586
9587 let one_ahead: Vec<_> = (0..4)
9588 .flat_map(|position| page_prefetch_positions(position, 4, 1))
9589 .collect();
9590 assert_eq!(one_ahead, vec![1, 2, 3]);
9591 assert!(page_prefetch_positions(0, 4, 0).is_empty());
9592 }
9593
9594 #[test]
9595 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
9596 assert_eq!(grouped_worker_prefetch_position(0, None), None);
9597 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
9598 .chain((0..4).filter_map(|position| {
9599 grouped_worker_prefetch_position(4, Some(position))
9600 }))
9601 .collect();
9602 assert_eq!(positions, vec![0, 1, 2, 3]);
9603 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
9604 }
9605
9606 #[test]
9607 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
9608 let queued: Vec<_> = (0..8)
9609 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
9610 .collect();
9611 assert_eq!(queued, (0..8).collect::<Vec<_>>());
9612
9613 let one_at_a_time: Vec<_> = (0..4)
9614 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
9615 .collect();
9616 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
9617 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
9618 }
9619}
9620
9621pub struct G4DcSlots {
9622 x: CudaSlice<f32>, xn: CudaSlice<f32>, cur: CudaSlice<f32>,
9623 hq: CudaSlice<i8>, hd_: CudaSlice<f32>,
9624 q0: CudaSlice<f32>, k0: CudaSlice<f32>, v0: CudaSlice<f32>,
9625 q: CudaSlice<f32>, k: CudaSlice<f32>, v: CudaSlice<f32>,
9626 attn: CudaSlice<f32>, o: CudaSlice<f32>,
9627 attn_out: CudaSlice<f32>, zsh: CudaSlice<f32>,
9628 zq: CudaSlice<i8>, zd: CudaSlice<f32>,
9629 gate: CudaSlice<f32>, up: CudaSlice<f32>,
9630 act: CudaSlice<f32>, actq: CudaSlice<i8>, actd: CudaSlice<f32>,
9631 f0: CudaSlice<f32>, sn: CudaSlice<f32>,
9632 hn: CudaSlice<f32>, logits: CudaSlice<f32>,
9633}