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
40
41pub(crate) struct AttnPre {
43 pub q: cudarc::driver::CudaSlice<f32>,
44 pub k: cudarc::driver::CudaSlice<f32>,
45 pub v: cudarc::driver::CudaSlice<f32>,
46 pub gate: Option<cudarc::driver::CudaSlice<f32>>,
47}
48
49pub(crate) struct GdnPrep {
51 pub hk: usize,
52 pub q_l2: cudarc::driver::CudaSlice<f32>,
53 pub k_l2: cudarc::driver::CudaSlice<f32>,
54 pub v_g: cudarc::driver::CudaSlice<f32>,
55 pub beta: cudarc::driver::CudaSlice<f32>,
56 pub g_log: cudarc::driver::CudaSlice<f32>,
57 pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
58 pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
59}
60
61pub(crate) struct VerifyStreamScratch {
63 pub pos_d: CudaSlice<i32>,
64 pub row_ctrs: Vec<CudaSlice<i32>>,
65}
66use crate::hybrid::{HybridModel, Mixer, FullAttnLayer, LinearAttnLayer, MoeWeights};
67
68struct MoeInputTraceWriter {
69 dir: std::path::PathBuf,
70 index: std::fs::File,
71 payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
72}
73
74static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<
75 std::sync::Mutex<Option<MoeInputTraceWriter>>,
76> = std::sync::OnceLock::new();
77
78fn gdec_enabled() -> bool {
81 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
82 *E.get_or_init(|| std::env::var("MEMRA_MOE_GDEC").map(|v| v != "0").unwrap_or(true))
83}
84
85fn moe_slab_enabled() -> bool {
96 std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
97}
98
99fn moe_prefetch_enabled() -> bool {
102 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
103 *E.get_or_init(|| std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
104 || crate::spill_pread::worker_enabled())
105}
106
107fn moe_page_prefetch_window() -> usize {
112 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
113 *W.get_or_init(|| page_prefetch_window_from_values(
114 std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
115 std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW").ok().as_deref(),
116 ))
117}
118
119fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
120 if !enabled {
121 return 0;
122 }
123 raw_window
124 .and_then(|value| value.parse().ok())
125 .unwrap_or(1)
126}
127
128fn page_prefetch_positions(
132 position: usize,
133 len: usize,
134 window: usize,
135) -> std::ops::Range<usize> {
136 if window == 0 || position >= len {
137 return len..len;
138 }
139 let (start, count) = if position == 0 {
140 (1, window)
141 } else {
142 (position.saturating_add(window), 1)
143 };
144 let start = start.min(len);
145 start..start.saturating_add(count).min(len)
146}
147
148fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
151 let position = current.map_or(0, |position| position.saturating_add(1));
152 (position < order_len).then_some(position)
153}
154
155fn worker_prefetch_window() -> usize {
160 static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
161 *WINDOW.get_or_init(|| {
162 let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
163 std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
164 .ok()
165 .and_then(|value| value.parse::<usize>().ok())
166 .unwrap_or(automatic.max(1))
167 })
168}
169
170fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
174 if window == 0 || position >= len {
175 return len..len;
176 }
177 let (start, count) = if position == 0 {
178 (0, window)
179 } else {
180 (position.saturating_add(window).saturating_sub(1), 1)
181 };
182 let start = start.min(len);
183 start..start.saturating_add(count).min(len)
184}
185
186fn moe_dev_enabled() -> bool {
191 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
192 *E.get_or_init(|| std::env::var("MEMRA_MOE_DEV").map(|v| v != "0").unwrap_or(true)
193 && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")))
194}
195
196fn moe_q8_enabled() -> bool {
201 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
202 *E.get_or_init(|| std::env::var("MEMRA_MOE_Q8").map(|v| v != "0").unwrap_or(true))
203}
204
205fn expert_dp4a_supported(qt: i32) -> bool {
208 qt == crate::QT_Q4_0 || qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS
209 || qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K
210}
211
212fn q8_expert_supported(qt: i32) -> bool {
213 static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
219 let kq = *KQ.get_or_init(|| {
220 std::env::var("MEMRA_MOE_Q8_KQ").map(|v| v != "0").unwrap_or(true)
221 });
222 let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4").map(|v| v != "0").unwrap_or(true);
229 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || (nvfp4_q8 && qt == crate::QT_NVFP4)
230 || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
231}
232
233fn q8_expert_dec_supported(qt: i32) -> bool {
236 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
237}
238
239fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
245 match qt {
246 crate::QT_Q4_0 => in_f % 32 == 0,
247 crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K
248 | crate::QT_Q6_K => in_f % 256 == 0,
249 _ => false,
250 }
251}
252
253fn moe_prewarm_enabled() -> bool {
256 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
257 *E.get_or_init(|| std::env::var("MEMRA_MOE_PREWARM").map(|v| v != "0").unwrap_or(true))
258}
259
260fn cpu_expert_profile_admit_enabled() -> bool {
264 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
265 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
266}
267
268pub const PRIME_MIN_T: usize = 16;
272
273impl HybridModel {
274 fn prime_trace_path() -> Option<&'static str> {
279 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
280 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
281 .as_deref()
282 }
283
284 pub fn forward(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
286 if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, false); }
287 if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, false); }
288 let cfg = &self.cfg;
289 let n_embd = cfg.n_embd as usize;
290 let t = tokens.len();
291 let eps = cfg.rms_eps;
292 let pos: Vec<i32> = (0..t as i32).collect();
293 let pos_d = e.htod_i32(&pos)?;
294
295 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
298 let mut h = e.uninit(t * n_embd)?;
300 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
301
302 let mixed = match &layer.mixer {
303 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
304 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
305 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
306 };
307
308 let mut x1 = e.uninit(t * n_embd)?;
310 e.add(&x, &mixed, &mut x1, t * n_embd)?;
311
312 let mut z = e.uninit(t * n_embd)?;
314 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
315 let ffn_out = match &layer.ffn {
316 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
317 let n_ff = ffn_gate.out_features();
318 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
319 let up = g2.pop().unwrap();
320 let gate = g2.pop().unwrap();
321 let mut act = e.uninit(t * n_ff)?;
322 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
327 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
328 e.matmul(ffn_down, &act, t)?
329 }
330 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
331 };
332 let mut x2 = e.uninit(t * n_embd)?;
333 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
334 x = x2;
335 }
336
337 let mut hn = e.uninit(t * n_embd)?;
338 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
339 let logits = e.matmul(&self.output, &hn, t)?;
340 Ok(e.dtoh(&logits)?)
341 }
342
343 pub fn forward_last(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
349 if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, true); }
350 let cfg = &self.cfg;
351 let n_embd = cfg.n_embd as usize;
352 let t = tokens.len();
353 let eps = cfg.rms_eps;
354 let pos: Vec<i32> = (0..t as i32).collect();
355 let pos_d = e.htod_i32(&pos)?;
356
357 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
361 for (il, layer) in self.layers.iter().enumerate() {
362 let mut h = e.uninit(t * n_embd)?;
363 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
364 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} norm ok"); }
365 let mixed = match &layer.mixer {
366 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
367 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
368 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
369 };
370 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} mixer ok"); }
371 let mut x1 = e.uninit(t * n_embd)?;
372 e.add(&x, &mixed, &mut x1, t * n_embd)?;
373 let mut z = e.uninit(t * n_embd)?;
374 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
375 let ffn_out = match &layer.ffn {
376 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
377 let n_ff = ffn_gate.out_features();
378 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
379 let up = g2.pop().unwrap();
380 let gate = g2.pop().unwrap();
381 let mut act = e.uninit(t * n_ff)?;
382 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
384 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
385 e.matmul(ffn_down, &act, t)?
386 }
387 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
388 };
389 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} ffn ok"); }
390 let mut x2 = e.uninit(t * n_embd)?;
391 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
392 x = x2;
393 }
394 let mut hn = e.uninit(t * n_embd)?;
396 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
397 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)?;
400 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
401 let logits = e.matmul(&self.output, &hlast, 1)?; Ok(e.dtoh(&logits)?)
403 }
404
405 pub fn prime_cache(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize)
437 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
438 let n_embd = self.cfg.n_embd as usize;
439 let t = tokens.len();
440 assert!(t >= PRIME_MIN_T, "prime_cache needs T >= {PRIME_MIN_T} (caller gates)");
444 assert!(cache.pos + t <= cache.max_ctx, "prime_cache: prompt exceeds cache max_ctx");
445
446 if self.is_gemma4_e4b() {
458 return self.gemma4_e4b_prime(e, tokens, cache);
459 }
460 if self.cfg.gemma4.is_some() {
461 return self.gemma4_prime(e, tokens, cache);
463 }
464 let chunk: usize = std::env::var("MEMRA_PRIME_CHUNK").ok()
465 .and_then(|v| v.parse().ok()).unwrap_or(4096);
466 let legacy_calllocal =
502 std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
503 let seq_end = if legacy_calllocal {
504 cache.pos + t
505 } else {
506 cache.pos + t + queued_after
507 };
508 if chunk == 0 || t <= chunk {
509 return self.prime_chunk(e, tokens, cache, seq_end);
510 }
511 let mut hiddens = e.uninit(t * n_embd)?;
512 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
513 let mut start = 0usize;
514 while start < t {
515 let mut end = (start + chunk).min(t);
517 if t - end > 0 && t - end < PRIME_MIN_T { end = t; }
518 let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
519 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
520 last = Some((l, hs));
521 start = end;
522 }
523 let (logits, h_seed) = last.unwrap();
524 Ok((logits, h_seed, hiddens))
525 }
526
527 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
534 if Engine::gdn_db_on()
535 && Engine::gdn_chunked_enabled() && t >= 16
536 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
537 && num_k * 2 == num_v
538 {
539 num_k
540 } else {
541 num_v
542 }
543 }
544
545 fn f16out_on(e: &Engine, t: usize) -> bool {
550 crate::f16_ffi::pp_f16_enabled() && t >= 16 && !e.verify_exact_on()
551 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
552 }
553
554 pub fn prime_slabs_get(&self, e: &Engine, t: usize, n_embd: usize, n_ff_max: usize)
562 -> Result<std::sync::MutexGuard<'_, std::collections::HashMap<usize, PrimeSlabs>>, Box<dyn std::error::Error>> {
563 let mut g = self.prime_slabs.lock().unwrap();
564 let dev = e.ctx().ordinal();
565 let need_new = match g.get(&dev) { None => true, Some(sl) => sl.t_cap < t };
566 if need_new {
567 g.insert(dev, PrimeSlabs {
568 t_cap: t,
569 h: e.uninit(t * n_embd)?,
570 x1: e.uninit(t * n_embd)?,
571 z: e.uninit(t * n_embd)?,
572 act: e.uninit(t * n_ff_max)?,
573 xa: e.uninit(t * n_embd)?,
574 xb: e.uninit(t * n_embd)?,
575 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
576 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
577 gate: e.uninit(t * n_ff_max)?,
578 up: e.uninit(t * n_ff_max)?,
579 ffn_out: e.uninit(t * n_embd)?,
580 seg_glue: Vec::new(),
581 mixed: e.uninit(t * n_embd)?,
582 seg_mid: Vec::new(),
583 seg_t: 0,
584 });
585 }
586 Ok(g)
587 }
588
589 fn prime_chunk(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize)
593 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
594 if self.cfg.gemma4.is_none()
603 && !crate::pp::pp2_streams_off()
604 && crate::pp::prime_pp_on()
605 {
606 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
607 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
608 }
609 }
610 let t = tokens.len();
611 let base = cache.pos;
612 debug_assert!(seq_end >= base + t, "prime_chunk: seq_end must cover this chunk");
613 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
614 let pos_d = e.htod_i32(&pos)?;
615
616 let x_embed = self.embed(e, tokens)?; let x = self.prime_layers(e, x_embed, 0, self.layers.len(), &pos_d, t, cache, seq_end)?;
618 self.prime_chunk_epilogue(e, x, t, cache)
619 }
620
621 #[allow(clippy::too_many_arguments)]
637 fn prime_layers(&self, e: &Engine, x_in: CudaSlice<f32>, lo: usize, hi: usize,
638 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, seq_end: usize)
639 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
640 let cfg = &self.cfg;
641 let n_embd = cfg.n_embd as usize;
642 let eps = cfg.rms_eps;
643 let base = cache.pos;
644 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
648 let n_ff_max = self.layers.iter().map(|l| match &l.ffn {
654 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
655 _ => n_embd,
656 }).max().unwrap_or(n_embd).max(n_embd);
657 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
658 let mut slab_guard = if use_slabs {
659 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
660 } else {
661 None
662 };
663 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>);
665 let (mut x_cur, mut x_nxt, sl): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, Option<SlabRefs>);
666 let mut seg: Option<(&mut Vec<Option<cudarc::driver::CudaGraph>>, &mut Vec<Option<cudarc::driver::CudaGraph>>, &mut CudaSlice<f32>, &mut usize)> = None;
667 let mut x_own2;
668 match slab_guard.as_mut() {
669 Some(g) => {
670 let slabs = g.get_mut(&e.ctx().ordinal()).expect("prime slabs for this device");
673 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
674 let PrimeSlabs { xa, xb, h, x1, z, act, h16, z16, gate, up, ffn_out, seg_glue, mixed, seg_mid, seg_t, .. } = slabs;
675 x_cur = xa;
676 x_nxt = xb;
677 seg = Some((seg_glue, seg_mid, mixed, seg_t));
678 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
679 }
680 None => {
681 x_own = x_in;
682 x_own2 = e.uninit(t * n_embd)?;
683 x_cur = &mut x_own;
684 x_nxt = &mut x_own2;
685 sl = None;
686 }
687 }
688 let mut alloc_h; let mut alloc_x1; let mut alloc_z; let mut alloc_act;
689 let mut alloc_h16; let mut alloc_z16;
690 let mut alloc_gate; let mut alloc_up; let mut alloc_fo;
691 let (h, x1, z, act): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
692 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
693 let (sl_gate, sl_up, sl_fo): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
694 match sl {
695 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
696 h = a; x1 = b; z = c; act = d; h16 = e16; z16 = f16b;
697 sl_gate = g; sl_up = u; sl_fo = fo;
698 }
699 None => {
700 alloc_h = e.uninit(t * n_embd)?;
701 alloc_x1 = e.uninit(t * n_embd)?;
702 alloc_z = e.uninit(t * n_embd)?;
703 alloc_act = e.uninit(t * n_ff_max)?;
704 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
705 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
706 alloc_gate = e.uninit(t * n_ff_max)?;
707 alloc_up = e.uninit(t * n_ff_max)?;
708 alloc_fo = e.uninit(t * n_embd)?;
709 h = &mut alloc_h; x1 = &mut alloc_x1; z = &mut alloc_z; act = &mut alloc_act;
710 h16 = &mut alloc_h16; z16 = &mut alloc_z16;
711 sl_gate = &mut alloc_gate; sl_up = &mut alloc_up; sl_fo = &mut alloc_fo;
712 }
713 }
714 let n_layers = self.layers.len();
719 let use_seg = f16fuse && seg.is_some() && self.cfg.step35.is_none()
729 && lo == 0 && hi == n_layers
730 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
731 if let Some((sg, sm, _, st)) = seg.as_mut() {
732 if **st != t {
733 sg.clear();
734 sg.extend((0..n_layers).map(|_| None));
735 sm.clear();
736 sm.extend((0..n_layers).map(|_| None));
737 **st = t;
738 }
739 }
740 {
741 let layer_lo = &self.layers[lo];
742 if f16fuse {
743 e.rms_norm_f16out(x_cur, layer_lo.attn_norm.float_data(), h, h16, n_embd, t, eps)?;
744 } else {
745 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
746 }
747 }
748 for il in lo..hi {
749 let layer = &self.layers[il];
750 let hx16 = if f16fuse { Some(&*h16) } else { None };
751 if use_seg {
752 let (pre, pre16, w_out) = match &layer.mixer {
755 Mixer::Full(fa) => {
756 let g3 = match hx16 {
757 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
758 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
759 };
760 let (pre, pre16) = self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
761 (pre, pre16, &fa.wo)
762 }
763 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
764 Mixer::Linear(la) => {
765 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
766 let g4 = match hx16 {
767 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
768 None => e.matmul_group(&ws, h, t)?,
769 };
770 let (pre, pre16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
771 (pre, pre16, &la.ssm_out)
772 }
773 };
774 {
775 let (_, sm, mslab, _) = seg.as_mut().unwrap();
776 let pre_n = pre.len() / t;
777 let xh_pre = match pre16 {
778 Some(x) => x,
779 None => e.f16_act(&pre, t * pre_n, pre_n)?,
780 };
781 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
782 let y = e.matmul(w_out, &pre, t)?;
783 e.copy_into(mslab, 0, &y, t * n_embd)?;
784 }
785 if sm[il].is_none() {
786 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
787 let w_post = layer.post_attn_norm.float_data();
788 e.stream().synchronize()?;
789 e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
790 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
791 e.add(x_cur, mslab, x1, t * n_embd)?;
792 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
793 Ok(())
794 })();
795 let g = e.stream().end_capture(
796 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
797 r?;
798 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
799 }
800 sm[il].as_ref().unwrap().launch()?;
801 }
802 } else {
803 let mixed = match &layer.mixer {
804 Mixer::Full(fa) => self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il,
805 seq_end)?,
806 Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
807 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
808 };
809 if f16fuse {
810 e.add_rms_norm_f16out(x_cur, &mixed, layer.post_attn_norm.float_data(),
813 x1, z, z16, n_embd, t, eps)?;
814 } else {
815 e.add(x_cur, &mixed, x1, t * n_embd)?;
816 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
817 }
818 }
819 let zx16 = if f16fuse { Some(&*z16) } else { None };
820 match &layer.ffn {
821 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
822 let n_ff = ffn_gate.out_features();
823 let mut into_ok = false;
826 if let Some(xh) = zx16 {
827 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
828 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
829 }
830 if !into_ok {
831 let mut g2 = match zx16 {
832 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
833 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
834 };
835 let up_y = g2.pop().unwrap();
836 let gate_y = g2.pop().unwrap();
837 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
838 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
839 }
840 let d_lim = self.cfg.clamp_shexp_at(il as u32);
845 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none()
846 && d_lim.is_none() {
847 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
848 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
849 Some(a16)
850 } else {
851 Self::ffn_act_lim(e, &self.cfg, sl_gate, sl_up, 1.0, 1.0, d_lim,
852 act, t * n_ff)?;
853 None
854 };
855 let xh_act = match act16 {
857 Some(x) => x,
858 None => e.f16_act(act, t * n_ff, n_ff)?,
859 };
860 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
861 let y = e.matmul(ffn_down, &*act, t)?;
862 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
863 }
864 }
865 crate::hybrid::Ffn::Moe(m) => {
866 let y = self.moe_ffn_il(e, m, z, t, il as u16)?;
867 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
868 }
869 }
870 if use_seg && il + 1 < hi {
871 let w_next = self.layers[il + 1].attn_norm.float_data();
873 let (sg, _, _, _) = seg.as_mut().unwrap();
874 if sg[il].is_none() {
875 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
876 e.stream().synchronize()?;
877 e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
878 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
879 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
880 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
881 Ok(())
882 })();
883 let g = e.stream().end_capture(
884 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
885 r?;
886 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
887 }
888 sg[il].as_ref().unwrap().launch()?;
889 } else {
890 if il + 1 < hi {
891 let w_next = self.layers[il + 1].attn_norm.float_data();
892 if f16fuse {
893 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
894 } else {
895 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
896 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
897 }
898 } else {
899 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
900 }
901 }
902 if let Some(path) = Self::prime_trace_path() {
908 let row = (base + t - 1) as usize;
909 let host = e.dtoh(x_nxt)?;
910 let last = &host[(t - 1) * n_embd..t * n_embd];
911 use std::io::Write as _;
912 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
913 let mut h64: u64 = 0xcbf29ce484222325;
914 for v in last {
915 h64 ^= v.to_bits() as u64;
916 h64 = h64.wrapping_mul(0x100000001b3);
917 }
918 writeln!(f, "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
919 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
920 last[0], last[1], last[2])?;
921 }
922 std::mem::swap(&mut x_cur, &mut x_nxt);
923 }
924 let mut x = e.uninit(t * n_embd)?;
926 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
927 drop(slab_guard);
928 Ok(x)
929 }
930
931 fn prime_chunk_epilogue(&self, e: &Engine, x: CudaSlice<f32>, t: usize, cache: &mut Cache)
936 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
937 let n_embd = self.cfg.n_embd as usize;
938 let eps = self.cfg.rms_eps;
939 let mut h_seed = e.uninit(n_embd)?;
943 if !crate::spec::spec_hpost() {
944 e.copy_view_into(&mut h_seed, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
945 }
946 let mut hn = e.uninit(t * n_embd)?;
948 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
949 if crate::spec::spec_hpost() {
950 e.copy_view_into(&mut h_seed, 0, &hn.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
951 }
952 let last = e.view(&hn, t * n_embd);
953 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
954 let mut hlast = e.uninit(n_embd)?;
955 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
956 let logits = e.matmul(&self.output, &hlast, 1)?;
957 cache.pos += t;
958 Ok((e.dtoh(&logits)?, h_seed, if crate::spec::spec_hpost() { hn } else { x }))
961 }
962
963 fn prime_chunk_ppn(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize,
987 fence: &[usize])
988 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
989 let rt = crate::pp::PpNRt::get(e)?;
990 let n_st = fence.len() - 1;
991 assert_eq!(
992 rt.n_stages(), n_st,
993 "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
994 );
995 let n_embd = self.cfg.n_embd as usize;
996 let t = tokens.len();
997 let base = cache.pos;
998 debug_assert!(seq_end >= base + t, "prime_chunk_ppn: seq_end must cover this chunk");
999 let payload = t * n_embd;
1000 let caller_stream = e.stream();
1004 rt.fence_stages_behind(&caller_stream)?;
1005 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1006
1007 let mut slot = {
1009 let _st0 = rt.enter(0);
1010 let e0 = rt.engine(0, e);
1011 let pos_d = e0.htod_i32(&pos)?;
1012 let x = self.embed(e0, tokens)?;
1013 let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, cache, seq_end)?;
1014 rt.tx(0, &x, payload)?
1015 };
1017
1018 for s in 1..n_st - 1 {
1020 let _st = rt.enter(s);
1021 let es = rt.engine(s, e);
1022 let pos_d = es.htod_i32(&pos)?;
1023 let x = rt.rx(s - 1, slot, payload)?;
1024 let x = self.prime_layers(es, x, fence[s], fence[s + 1], &pos_d, t, cache, seq_end)?;
1025 slot = rt.tx(s, &x, payload)?;
1026 }
1027
1028 let _stl = rt.enter(n_st - 1);
1030 let el = rt.engine(n_st - 1, e);
1031 let pos_d = el.htod_i32(&pos)?;
1032 let x = rt.rx(n_st - 2, slot, payload)?;
1033 let x = self.prime_layers(el, x, fence[n_st - 1], fence[n_st], &pos_d, t, cache, seq_end)?;
1034 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1035 rt.publish_to(n_st - 1, &caller_stream)?;
1041 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1042 Ok(out)
1043 }
1044
1045 pub fn prime_chunk_captured(&self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>,
1061 t: usize, cache: &mut Cache,
1062 len_d: &CudaSlice<i32>,
1063 logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>)
1064 -> Result<(), Box<dyn std::error::Error>> {
1065 let cfg = &self.cfg;
1066 let n_embd = cfg.n_embd as usize;
1067 let eps = cfg.rms_eps;
1068 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1069 let mut x = e.uninit(t * n_embd)?;
1070 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1071 for (il, layer) in self.layers.iter().enumerate() {
1072 let mut h = e.uninit(t * n_embd)?;
1073 let mut hx16: Option<CudaSlice<u8>> = None;
1074 if f16fuse {
1075 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1076 e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut b16, n_embd, t, eps)?;
1077 hx16 = Some(b16);
1078 } else {
1079 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1080 }
1081 let mixed = match &layer.mixer {
1082 Mixer::Full(fa) => self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache,
1086 il, t)?,
1087 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1088 Mixer::Linear(la) => {
1089 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1090 let g4 = match hx16.as_ref() {
1091 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1092 None => e.matmul_group(&ws, &h, t)?,
1093 };
1094 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1095 }
1096 };
1097 let mut x1 = e.uninit(t * n_embd)?;
1098 e.add(&x, &mixed, &mut x1, t * n_embd)?;
1099 let mut z = e.uninit(t * n_embd)?;
1100 let mut zx16: Option<CudaSlice<u8>> = None;
1101 if f16fuse {
1102 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1103 e.rms_norm_f16out(&x1, layer.post_attn_norm.float_data(), &mut z, &mut b16, n_embd, t, eps)?;
1104 zx16 = Some(b16);
1105 } else {
1106 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
1107 }
1108 let ffn_out = match &layer.ffn {
1109 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1110 let n_ff = ffn_gate.out_features();
1111 let mut g2 = match &zx16 {
1112 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1113 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1114 };
1115 let up = g2.pop().unwrap();
1116 let gate = g2.pop().unwrap();
1117 let mut act = e.uninit(t * n_ff)?;
1118 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
1120 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
1121 e.matmul(ffn_down, &act, t)?
1122 }
1123 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
1124 };
1125 let mut x2 = e.uninit(t * n_embd)?;
1126 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1127 x = x2;
1128 }
1129 if !crate::spec::spec_hpost() {
1131 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1132 }
1133 let mut hn = e.uninit(t * n_embd)?;
1134 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1135 if crate::spec::spec_hpost() {
1136 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1137 }
1138 let mut hlast = e.uninit(n_embd)?;
1139 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1140 let logits = e.matmul(&self.output, &hlast, 1)?;
1141 let nv = logits.len();
1142 e.copy_into(logits_out, 0, &logits, nv)?;
1143 Ok(())
1144 }
1145
1146 pub fn prime_cache_batch(&self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache])
1163 -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1164 let cfg = &self.cfg;
1165 let n_embd = cfg.n_embd as usize;
1166 let eps = cfg.rms_eps;
1167 let b = prompts.len();
1168 assert!(b >= 1 && b == caches.len());
1169 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
1170 let carried = pos0s.iter().any(|&p| p > 0);
1171 if cfg.gemma4.is_some() {
1177 return Err("prime_cache_batch: gemma4 has no batched prime core (per-layer \
1178 swa/global geometry, softcapped head) — use gemma4_prime per sequence".into());
1179 }
1180 if cfg.step35.is_some() {
1186 return Err("prime_cache_batch: step35 has no batched prime core (per-layer n_head / \
1187 partial rope / SWA / head-wise gate) — use prime_cache per sequence".into());
1188 }
1189 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
1190 for &t in &ts { assert!(t >= PRIME_MIN_T, "prime_cache_batch needs T >= {PRIME_MIN_T}"); }
1191 for (s, c) in caches.iter().enumerate() {
1192 assert!(c.pos + ts[s] <= c.max_ctx, "prime_cache_batch: prompt exceeds cache max_ctx");
1193 }
1194 let total: usize = ts.iter().sum();
1195 let offs: Vec<usize> = ts.iter().scan(0usize, |a, &t| { let o = *a; *a += t; Some(o) }).collect();
1196 let pos_ds: Vec<CudaSlice<i32>> = ts.iter().zip(&pos0s)
1198 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
1199 .collect::<Result<_, _>>()?;
1200 let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
1202 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1203 let mut out = Vec::with_capacity(b);
1204 for s in 0..b {
1205 let mut ys = e.uninit(ts[s] * dim)?;
1206 e.copy_view_into(&mut ys, 0, &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim), ts[s] * dim)?;
1207 out.push(ys);
1208 }
1209 Ok(out)
1210 };
1211
1212 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
1213 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
1215 let mut h = e.uninit(total * n_embd)?;
1216 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1217 e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut hx16, n_embd, total, eps)?;
1218 let mut mixed = e.uninit(total * n_embd)?;
1220 match &layer.mixer {
1221 Mixer::Full(fa) => {
1222 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
1223 let (n_head, n_head_kv, head_dim) =
1229 (self.cfg.n_head as usize, self.cfg.n_head_kv as usize, self.cfg.head_dim_k as usize);
1230 let fa_scale = 1.0 / (head_dim as f32).sqrt();
1231 let use_favl = !carried
1232 && (2..=8).contains(&b)
1233 && (head_dim == 256 || head_dim == 128)
1234 && self.cfg.attn_out_gate()
1235 && std::env::var("MEMRA_NOFA").is_err()
1236 && std::env::var("MEMRA_FA_FLOOR").is_err()
1237 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
1238 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
1239 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
1240 if use_favl {
1241 let (qf_w, kf_w, vf_w) =
1242 (fa.wq.out_features(), fa.wk.out_features(), fa.wv.out_features());
1243 struct APre {
1244 q: CudaSlice<f32>, gate: Option<CudaSlice<f32>>,
1245 qn: CudaSlice<f32>, kn: CudaSlice<f32>,
1246 }
1247 let mut aps = Vec::with_capacity(b);
1248 for &t in ts.iter().take(b) {
1249 aps.push(APre {
1250 q: e.uninit(t * n_head * head_dim)?,
1251 gate: Some(e.uninit(t * n_head * head_dim)?),
1252 qn: e.uninit(t * n_head * head_dim)?,
1253 kn: e.uninit(t * n_head_kv * head_dim)?,
1254 });
1255 }
1256 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
1257 let kvl = caches[0].kv[il].as_ref().unwrap();
1258 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
1259 };
1260 let pargs: Vec<crate::AttnPreVl> = (0..b).map(|s| {
1261 let (o, t) = (offs[s], ts[s]);
1262 let kvl = caches[s].kv[il].as_ref().unwrap();
1263 assert!(kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
1264 "prime_cache_batch attn vl: fresh + capacity");
1265 crate::AttnPreVl {
1266 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
1267 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
1268 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
1269 q: e.addr_f32(&aps[s].q),
1270 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
1271 qn: e.addr_f32(&aps[s].qn), kn: e.addr_f32(&aps[s].kn),
1272 kc: e.addr_u8(&kvl.k), vc: e.addr_u8(&kvl.v),
1273 t: t as i32, pad: 0,
1274 }
1275 }).collect();
1276 e.attn_pre_vl8(&pargs, fa.q_norm.float_data(), fa.k_norm.float_data(),
1277 head_dim, self.cfg.rope_dim_count as usize, n_head, n_head_kv,
1278 self.cfg.rms_eps, self.cfg.rope_freq_base, 1.0,
1279 kv_dim_k, kv_dim_v, ktb, vtb)?;
1280 for s in 0..b {
1281 let kvl = caches[s].kv[il].as_mut().unwrap();
1282 kvl.len += ts[s];
1283 let new_len = kvl.len as i32;
1284 e.set_i32_one(&mut kvl.len_d, new_len)?;
1285 }
1286 let mut attns = Vec::with_capacity(b);
1287 let mut mirrors = Vec::with_capacity(b);
1288 for &t in ts.iter().take(b) {
1289 attns.push(e.uninit(t * n_head * head_dim)?);
1290 let n = t * n_head_kv * head_dim;
1291 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
1292 }
1293 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
1296 Ok("0") => false,
1297 Ok("1") => true,
1298 _ => cfg!(memra_hopper_mma),
1299 };
1300 if fa3_on {
1301 let mut q16s = Vec::with_capacity(b);
1302 let mut v16s = Vec::with_capacity(b);
1303 for s in 0..b {
1304 let t = ts[s];
1305 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
1306 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
1307 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
1308 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
1309 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
1310 e.f32_to_bf16_v(&g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
1311 &mut v16, t * n_head_kv * head_dim)?;
1312 q16s.push(q16);
1313 v16s.push((k16, v16));
1314 }
1315 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
1316 let mut kp = qp;
1317 let mut vp = qp;
1318 let mut op = [core::ptr::null_mut::<f32>(); 8];
1319 let mut tsv = [0i32; 8];
1320 for s in 0..b {
1321 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
1322 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
1323 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
1324 op[s] = e.addr_f32(&attns[s]) as *mut f32;
1325 tsv[s] = ts[s] as i32;
1326 }
1327 let rc = unsafe {
1328 crate::fa3_vl_raw(qp.as_ptr(), kp.as_ptr(), vp.as_ptr(), op.as_ptr(),
1329 tsv.as_ptr(), b as i32, n_head as i32,
1330 n_head_kv as i32, head_dim as i32, fa_scale,
1331 e.stream().cu_stream() as *mut core::ffi::c_void)
1332 };
1333 if rc != 0 {
1334 return Err(format!("memra_fa3_vl rc={rc}").into());
1335 }
1336 } else {
1337 let fargs: Vec<crate::FaSeqVl> = (0..b).map(|s| crate::FaSeqVl {
1338 q: e.addr_f32(&aps[s].qn), k16: e.addr_u8(&mirrors[s].0),
1339 v16: e.addr_u8(&mirrors[s].1), o: e.addr_f32(&attns[s]),
1340 kf: e.addr_f32(&aps[s].kn),
1341 vf: e.addr_f32v(&g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w)),
1342 t: ts[s] as i32, pad: 0,
1343 }).collect();
1344 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
1345 }
1346 for (s, attn) in attns.into_iter().enumerate() {
1347 let (attn_g, ag16) = self.full_attn_prime_post_fa(
1348 e, attn, &aps[s].gate, ts[s], n_head, head_dim)?;
1349 let mut done = false;
1350 if let Some(xh) = &ag16 {
1351 done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
1352 }
1353 if !done {
1354 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
1355 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
1356 }
1357 }
1358 } else {
1359 let mut parts: Vec<Vec<CudaSlice<f32>>> = (0..b).map(|_| Vec::new()).collect();
1360 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
1361 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
1362 parts[s].push(ys);
1363 }
1364 }
1365 for (s, g3s) in parts.into_iter().enumerate() {
1366 let (attn_g, ag16) = self.full_attn_prime_core_inner(
1368 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il)?;
1369 let mut done = false;
1370 if let Some(xh) = &ag16 {
1371 done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
1372 }
1373 if !done {
1374 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
1375 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
1376 }
1377 }
1378 }
1379 }
1380 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1381 Mixer::Linear(la) => {
1382 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1387 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
1388 let outs = self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
1389 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
1390 let (o, t) = (offs[s], ts[s]);
1391 let mut done = false;
1392 if let Some(xh) = &gn16 {
1393 done = e.try_f16_gemm_pre_into_off(&la.ssm_out, xh, t, &mut mixed, o * n_embd)?;
1394 }
1395 if !done {
1396 let m = e.matmul(&la.ssm_out, &gn, t)?;
1397 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
1398 }
1399 }
1400 }
1401 }
1402 let mut x1 = e.uninit(total * n_embd)?;
1403 let mut z = e.uninit(total * n_embd)?;
1404 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1405 e.add_rms_norm_f16out(&x, &mixed, layer.post_attn_norm.float_data(),
1406 &mut x1, &mut z, &mut zx16, n_embd, total, eps)?;
1407 let ffn_out = match &layer.ffn {
1408 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1409 let n_ff = ffn_gate.out_features();
1410 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
1411 let up = g2.pop().unwrap();
1412 let gate = g2.pop().unwrap();
1413 let mut act = e.uninit(total * n_ff)?;
1414 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1418 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
1419 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
1420 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
1421 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
1422 Some(y) => y,
1423 None => e.matmul(ffn_down, &act, total)?,
1424 }
1425 } else {
1426 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, d_lim,
1427 &mut act, total * n_ff)?;
1428 e.matmul(ffn_down, &act, total)?
1429 }
1430 }
1431 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
1432 };
1433 let mut x2 = e.uninit(total * n_embd)?;
1434 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
1435 x = x2;
1436 }
1437 let mut hn = e.uninit(total * n_embd)?;
1439 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, total, eps)?;
1440 let mut hcat = e.uninit(b * n_embd)?;
1446 for s in 0..b {
1447 let last0 = (offs[s] + ts[s] - 1) * n_embd;
1448 e.copy_view_into(&mut hcat, s * n_embd, &hn.slice(last0..last0 + n_embd), n_embd)?;
1449 }
1450 let logits_cat = if b >= 2 { e.try_f16_gemm(&self.output, &hcat, b)? } else { None };
1451 let logits_host: Option<Vec<f32>> = match &logits_cat {
1452 Some(lc) => Some(e.dtoh(lc)?),
1453 None => None,
1454 };
1455 let n_vocab = self.output.out_features();
1456 let mut hidden_all = if crate::spec::spec_hpost() {
1457 split(e, &hn, n_embd)?
1458 } else {
1459 split(e, &x, n_embd)?
1460 };
1461 let mut out = Vec::with_capacity(b);
1462 for s in 0..b {
1463 let last0 = (offs[s] + ts[s] - 1) * n_embd;
1464 let mut h_seed = e.uninit(n_embd)?;
1465 if !crate::spec::spec_hpost() {
1466 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
1467 } else {
1468 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
1469 }
1470 let logits = match &logits_host {
1471 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
1472 None => {
1473 let mut hlast = e.uninit(n_embd)?;
1474 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
1475 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
1476 }
1477 };
1478 caches[s].pos += ts[s];
1479 out.push((logits, h_seed, hidden_all.remove(0)));
1480 }
1481 Ok(out)
1482 }
1483
1484 #[allow(clippy::too_many_arguments)]
1495 fn full_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
1496 hx: Option<&CudaSlice<u8>>,
1497 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize,
1498 seq_end: usize)
1499 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1500 if self.cfg.step35.is_some() {
1501 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
1502 }
1503 let g3 = match hx {
1508 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1509 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1510 };
1511 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
1512 }
1513
1514 fn full_attn_prime_core(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
1518 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1519 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1520 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
1521 if let Some(xh) = &ag16 {
1522 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
1523 return Ok(y);
1524 }
1525 }
1526 Ok(e.matmul(&fa.wo, &attn_g, t)?)
1527 }
1528
1529 fn full_attn_prime_core_inner(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
1530 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1531 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1532 let cfg = &self.cfg;
1533 let n_head = cfg.n_head as usize;
1534 let n_head_kv = cfg.n_head_kv as usize;
1535 let head_dim = cfg.head_dim_k as usize;
1536 let scale = 1.0 / (head_dim as f32).sqrt();
1537 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
1538 let AttnPre { q, k, v, gate } = pre;
1539 let mut attn = e.uninit(t * n_head * head_dim)?;
1540 self.full_attn_prime_fa_dispatch(e, &q, &k, &v, &mut attn, base_len, t, cache, il,
1541 head_dim, n_head, n_head_kv, scale)?;
1542 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
1543 }
1544
1545 #[allow(clippy::type_complexity)]
1549 fn full_attn_prime_pre_fa(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
1550 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1551 -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
1552 let cfg = &self.cfg;
1553 let n_head = cfg.n_head as usize;
1554 let n_head_kv = cfg.n_head_kv as usize;
1555 let head_dim = cfg.head_dim_k as usize;
1556 let eps = cfg.rms_eps;
1557
1558 let gated = cfg.attn_out_gate();
1562 let v = g3.pop().unwrap();
1563 let mut k = g3.pop().unwrap();
1564 let qf = g3.pop().unwrap();
1565 let (mut q, gate) = if gated {
1566 let mut q = e.uninit(t * n_head * head_dim)?;
1567 let mut gate = e.uninit(t * n_head * head_dim)?;
1568 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
1569 (q, Some(gate))
1570 } else {
1571 (qf, None)
1572 };
1573
1574 let mut qn = e.uninit(t * n_head * head_dim)?;
1575 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
1576 q = qn;
1577 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
1578 e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
1579 k = kn;
1580 let rope_dims = cfg.rope_dim_count as usize;
1581 e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, cfg.rope_freq_base, 1.0)?;
1582 e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, cfg.rope_freq_base, 1.0)?;
1583
1584 {
1587 let kvl = cache.kv[il].as_mut().unwrap();
1588 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
1589 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
1590 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
1591 crate::Engine::kv_fp8_on())?;
1592 kvl.len += t;
1593 let new_len = kvl.len as i32;
1594 e.set_i32_one(&mut kvl.len_d, new_len)?;
1595 }
1596
1597 let base_len = {
1598 let kvl = cache.kv[il].as_ref().unwrap();
1599 kvl.len - t };
1601 Ok((AttnPre { q, k, v, gate }, base_len))
1602 }
1603
1604 #[allow(clippy::too_many_arguments)]
1611 fn full_attn_prime_fa_dispatch(&self, e: &Engine, q: &CudaSlice<f32>, k: &CudaSlice<f32>,
1612 v: &CudaSlice<f32>, attn: &mut CudaSlice<f32>, base_len: usize,
1613 t: usize, cache: &mut Cache, il: usize,
1614 head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32)
1615 -> Result<(), Box<dyn std::error::Error>> {
1616 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
1629 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1630 e.sdpa_naive(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1631 } else {
1632 e.fa_prefill(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1633 }
1634 return Ok(());
1635 }
1636 let kvl = cache.kv[il].as_ref().unwrap();
1637 let t_kv = base_len + t;
1638 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
1639 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
1640 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1644 e.sdpa_naive_quantized_view(q, &k_view, &v_view, attn, head_dim, n_head,
1645 n_head_kv, t, t_kv, scale, true,
1646 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
1647 return Ok(());
1648 }
1649 let deqw = std::env::var("MEMRA_PRIME_DEQW").map(|v| v != "0").unwrap_or(true);
1657 if deqw {
1658 e.fa_prefill_view_ws(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
1659 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
1660 crate::Engine::kv_fp8_on())?;
1661 } else {
1662 e.fa_prefill_view(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
1663 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
1664 crate::Engine::kv_fp8_on())?;
1665 }
1666 Ok(())
1667 }
1668
1669 fn full_attn_prime_post_fa(&self, e: &Engine, attn: CudaSlice<f32>,
1672 gate: &Option<CudaSlice<f32>>, t: usize,
1673 n_head: usize, head_dim: usize)
1674 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1675 let (attn_g, ag16) = match gate {
1676 Some(gate) => {
1677 let n = t * n_head * head_dim;
1678 let mut ag = e.uninit(n)?;
1679 if Self::f16out_on(e, t) {
1680 let mut a16 = e.alloc_u8_uninit(n * 2)?;
1681 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
1682 (ag, Some(a16))
1683 } else {
1684 let mut gsig = e.uninit(n)?;
1685 e.sigmoid(gate, &mut gsig, n)?;
1686 e.mul(&attn, &gsig, &mut ag, n)?;
1687 (ag, None)
1688 }
1689 }
1690 None => (attn, None),
1691 };
1692 Ok((attn_g, ag16))
1693 }
1694
1695 fn linear_attn_prime(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>,
1702 hx: Option<&CudaSlice<u8>>, t: usize,
1703 cache: &mut Cache, il: usize)
1704 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1705 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1707 let g4 = match hx {
1708 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1709 None => e.matmul_group(&ws, h, t)?,
1710 };
1711 self.linear_attn_prime_core(e, la, g4, t, cache, il)
1712 }
1713
1714 fn linear_attn_prime_core(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
1716 t: usize, cache: &mut Cache, il: usize)
1717 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1718 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
1719 }
1720
1721 #[allow(clippy::too_many_arguments)]
1725 fn linear_attn_prime_core_pad_inner(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
1726 t: usize, cache: &mut Cache, il: usize,
1727 pad_len: Option<&CudaSlice<i32>>)
1728 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1729 let ssm = self.cfg.ssm.as_ref().unwrap();
1731 let d_state = ssm.state_size as usize;
1732 let num_k = ssm.group_count as usize;
1733 let num_v = ssm.time_step_rank as usize;
1734 let key_dim = d_state * num_k;
1735 let value_dim = d_state * num_v;
1736 let conv_dim = key_dim * 2 + value_dim;
1737 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(
1742 e, la,
1743 &qkv_mixed.slice(0..t * conv_dim), &z.slice(0..t * value_dim),
1744 &beta_raw.slice(0..t * num_v), &alpha.slice(0..t * num_v),
1745 t, cache, il, pad_len)
1746 }
1747
1748 #[allow(clippy::too_many_arguments)]
1751 fn linear_attn_gdn_prep(&self, e: &Engine, la: &LinearAttnLayer,
1752 qkv_mixed: &cudarc::driver::CudaView<f32>,
1753 beta_raw: &cudarc::driver::CudaView<f32>,
1754 alpha: &cudarc::driver::CudaView<f32>,
1755 t: usize, cache: &mut Cache, il: usize,
1756 pad_len: Option<&CudaSlice<i32>>)
1757 -> Result<GdnPrep, Box<dyn std::error::Error>> {
1758 let cfg = &self.cfg;
1759 let ssm = cfg.ssm.as_ref().unwrap();
1760 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;
1768 debug_assert!(t >= d_conv - 1, "stateful conv needs T >= pad (PRIME_MIN_T gates)");
1769
1770 let rl = cache.recur[il].as_mut().unwrap();
1775 let hk = Self::gdn_hk(e, t, num_v, num_k);
1776 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
1777 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
1779 let mut k_g = e.uninit(d_state * hk * t)?;
1780 let mut v_g = e.uninit(d_state * num_v * t)?;
1781 if conv_fuse {
1782 e.ssm_conv1d_gdn_state_pad(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
1783 &mut q_g, &mut k_g, &mut v_g,
1784 conv_dim, t, d_conv, d_state, num_v, num_k, key_dim, hk, pad_len)?;
1785 } else {
1786 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(),
1788 &mut conv_out, conv_dim, t, d_conv, pad_len)?;
1789 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)?;
1790 }
1791 let mut q_l2 = e.uninit(d_state * hk * t)?;
1792 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
1796 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
1797 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
1798 Some(qb)
1799 } else {
1800 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
1801 None
1802 };
1803 let mut k_l2 = e.uninit(d_state * hk * t)?;
1804 let kb16 = if Engine::l2_v2_on(d_state) {
1806 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
1807 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
1808 Some(kb)
1809 } else {
1810 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
1811 None
1812 };
1813 let mut beta = e.uninit(t * num_v)?;
1814 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
1815 let mut g_log = e.uninit(t * num_v)?;
1816 e.gdn_glog_v(alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
1817 if let Some(len_d) = pad_len {
1818 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
1819 }
1820 Ok(GdnPrep { hk, q_l2, k_l2, v_g, beta, g_log, kb16, qb16 })
1821 }
1822
1823 #[allow(clippy::too_many_arguments)]
1828 fn linear_attn_prime_core_batch(&self, e: &Engine, la: &LinearAttnLayer,
1829 g4: &[CudaSlice<f32>], offs: &[usize], ts: &[usize],
1830 caches: &mut [&mut Cache], il: usize)
1831 -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
1832 let ssm = self.cfg.ssm.as_ref().unwrap();
1833 let d_state = ssm.state_size as usize;
1834 let num_k = ssm.group_count as usize;
1835 let num_v = ssm.time_step_rank as usize;
1836 let key_dim = d_state * num_k;
1837 let value_dim = d_state * num_v;
1838 let conv_dim = key_dim * 2 + value_dim;
1839 let eps = self.cfg.rms_eps;
1840 let scale = 1.0 / (d_state as f32).sqrt();
1841 let b = ts.len();
1842 let c = Engine::gdn_chunk_size();
1843 let carried = caches.iter().any(|c| c.pos > 0);
1846 let use_vl = !carried
1847 && (2..=8).contains(&b)
1848 && Engine::gdn_chunked_enabled() && ts.iter().all(|&t| t >= 16)
1849 && e.gdn_mma_enabled(c)
1850 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
1851 if !use_vl {
1852 return (0..b).map(|s| {
1853 let (o, t) = (offs[s], ts[s]);
1854 self.linear_attn_prime_core_pad_view(
1855 e, la,
1856 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
1857 &g4[1].slice(o * value_dim..(o + t) * value_dim),
1858 &g4[2].slice(o * num_v..(o + t) * num_v),
1859 &g4[3].slice(o * num_v..(o + t) * num_v),
1860 t, caches[s], il, None)
1861 }).collect();
1862 }
1863 struct SeqBufs {
1867 conv_out: CudaSlice<f32>, q_g: CudaSlice<f32>, k_g: CudaSlice<f32>, v_g: CudaSlice<f32>,
1868 q_l2: CudaSlice<f32>, k_l2: CudaSlice<f32>, beta: CudaSlice<f32>, g_log: CudaSlice<f32>,
1869 gn: CudaSlice<f32>, gn16: CudaSlice<u8>,
1870 }
1871 let d_conv = ssm.conv_kernel as usize;
1872 let f16o = Self::f16out_on(e, 16);
1873 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
1875 let mut pres = Vec::with_capacity(b);
1876 for &t in ts.iter().take(b) {
1877 sb.push(SeqBufs {
1878 conv_out: e.uninit(conv_dim * t)?,
1879 q_g: e.uninit(d_state * hk * t)?,
1880 k_g: e.uninit(d_state * hk * t)?,
1881 v_g: e.uninit(d_state * num_v * t)?,
1882 q_l2: e.uninit(d_state * hk * t)?,
1883 k_l2: e.uninit(d_state * hk * t)?,
1884 beta: e.uninit(t * num_v)?,
1885 g_log: e.uninit(t * num_v)?,
1886 gn: e.uninit(d_state * num_v * t)?,
1887 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
1888 });
1889 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
1890 }
1891 let prep_args: Vec<crate::GdnPrepVl> = (0..b).map(|s| {
1892 let (o, t) = (offs[s], ts[s]);
1893 let rl = caches[s].recur[il].as_ref().unwrap();
1894 crate::GdnPrepVl {
1895 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
1896 conv_state: e.addr_f32(&rl.conv_state),
1897 conv_out: e.addr_f32(&sb[s].conv_out),
1898 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),
1899 q_l2: e.addr_f32(&sb[s].q_l2), k_l2: e.addr_f32(&sb[s].k_l2),
1900 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
1901 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
1902 beta: e.addr_f32(&sb[s].beta), g_log: e.addr_f32(&sb[s].g_log),
1903 o: e.addr_f32(&pres[s].o),
1904 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
1905 gn: e.addr_f32(&sb[s].gn), gn16: e.addr_u8(&sb[s].gn16),
1906 kb16: if Engine::l2_v2_on(d_state) { e.addr_u8(&pres[s].kb16) } else { 0 },
1907 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) { e.addr_u8(&pres[s].qb16) } else { 0 },
1908 t: t as i32, pad: 0,
1909 }
1910 }).collect();
1911 let args: Vec<crate::GdnSeqVl> = (0..b).map(|s| {
1912 let rl = caches[s].recur[il].as_ref().unwrap();
1913 crate::GdnSeqVl {
1914 kb16: e.addr_u8(&pres[s].kb16), gcum: e.addr_f32(&pres[s].gcum),
1915 beta: e.addr_f32(&sb[s].beta), u: e.addr_f32(&pres[s].u),
1916 wb16: e.addr_u8(&pres[s].wb16), y: e.addr_u8(&pres[s].y16),
1917 ssnap: e.addr_u8(&pres[s].ssnap16),
1918 state_in: e.addr_f32(&rl.ssm_state), state_out: e.addr_f32(&rl.ssm_state_alt),
1919 q: e.addr_f32(&sb[s].q_l2), p: e.addr_f32(&pres[s].p),
1920 o: e.addr_f32(&pres[s].o),
1921 k: e.addr_f32(&sb[s].k_l2), v: e.addr_f32(&sb[s].v_g),
1922 g: e.addr_f32(&sb[s].g_log), a: e.addr_f32(&pres[s].a),
1923 w: e.addr_f32(&pres[s].w),
1924 t: ts[s] as i32, nc: pres[s].nc as i32,
1925 }
1926 }).collect();
1927 e.gdn_prep_vl8(&prep_args, la.ssm_conv1d.float_data(), la.ssm_dt.float_data(),
1928 la.ssm_a.float_data(), conv_dim, d_conv, d_state, num_v, num_k, key_dim, hk, eps)?;
1929 if !Engine::l2_v2_on(d_state) {
1932 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
1933 }
1934 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
1936 if !Engine::l2_v2_on(d_state) {
1938 for s in 0..b {
1939 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
1940 }
1941 }
1942 let mut wa = [crate::GdnWVl::default(); 8];
1943 for s in 0..b {
1944 wa[s] = crate::GdnWVl { qb16: e.addr_u8(&pres[s].qb16), pb16: e.addr_u8(&pres[s].pb16) };
1945 }
1946 Some(crate::GdnWVl8(wa))
1947 } else { None };
1948 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
1949 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
1950 if f16o {
1951 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
1952 }
1953 let mut out = Vec::with_capacity(b);
1955 for (s, bufs) in sb.into_iter().enumerate() {
1956 let rl = caches[s].recur[il].as_mut().unwrap();
1957 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1958 let (o, t) = (offs[s], ts[s]);
1959 let SeqBufs { mut gn, gn16, .. } = bufs;
1960 if f16o {
1961 out.push((gn, Some(gn16)));
1962 } else {
1963 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
1964 e.gated_rmsnorm_zv(&pres[s].o, la.ssm_norm.float_data(), &z_v, &mut gn,
1965 d_state, num_v * t, eps)?;
1966 out.push((gn, None));
1967 }
1968 }
1969 Ok(out)
1970 }
1971
1972 #[allow(clippy::too_many_arguments)]
1976 fn linear_attn_prime_core_pad_view(&self, e: &Engine, la: &LinearAttnLayer,
1977 qkv_mixed: &cudarc::driver::CudaView<f32>,
1978 z: &cudarc::driver::CudaView<f32>,
1979 beta_raw: &cudarc::driver::CudaView<f32>,
1980 alpha: &cudarc::driver::CudaView<f32>,
1981 t: usize, cache: &mut Cache, il: usize,
1982 pad_len: Option<&CudaSlice<i32>>)
1983 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1984 let cfg = &self.cfg;
1985 let ssm = cfg.ssm.as_ref().unwrap();
1986 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
1989 let scale = 1.0 / (d_state as f32).sqrt();
1990
1991 let prep = self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
1992
1993 let mut o = e.uninit(d_state * num_v * t)?;
1999 let rl = cache.recur[il].as_mut().unwrap();
2000 {
2001 let crate::cache::RecurLayer { ssm_state, ssm_state_alt, .. } = rl;
2002 e.gdn_scan_prefill(&prep.q_l2, &prep.k_l2, &prep.v_g, &prep.g_log, &prep.beta,
2003 prep.kb16.as_ref(), prep.qb16.as_ref(), ssm_state, ssm_state_alt, &mut o, num_v, t, scale,
2004 prep.hk)?;
2005 }
2006 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2007
2008 let mut gn = e.uninit(d_state * num_v * t)?;
2011 let gn16 = if Self::f16out_on(e, t) {
2012 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
2013 e.gated_rmsnorm_f16out_zv(&o, la.ssm_norm.float_data(), z, &mut gn, &mut g16,
2014 d_state, num_v * t, eps)?;
2015 Some(g16)
2016 } else {
2017 e.gated_rmsnorm_zv(&o, la.ssm_norm.float_data(), z, &mut gn, d_state, num_v * t, eps)?;
2018 None
2019 };
2020 Ok((gn, gn16))
2021 }
2022
2023 #[allow(clippy::too_many_arguments)]
2025 fn linear_attn_prime_core_pad(&self, e: &Engine, la: &LinearAttnLayer, g4: Vec<CudaSlice<f32>>,
2026 t: usize, cache: &mut Cache, il: usize,
2027 pad_len: Option<&CudaSlice<i32>>)
2028 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2029 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
2030 if let Some(xh) = &gn16 {
2031 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
2032 return Ok(y);
2033 }
2034 }
2035 Ok(e.matmul(&la.ssm_out, &gn, t)?)
2036 }
2037
2038 pub fn full_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, il: usize)
2043 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2044 if self.cfg.step35.is_some() {
2045 return self.step35_attn(e, fa, h, pos_d, t, il);
2046 }
2047 let cfg = &self.cfg;
2048 let _n_embd = cfg.n_embd as usize;
2049 let n_head = cfg.n_head as usize;
2050 let n_head_kv = cfg.n_head_kv as usize;
2051 let head_dim = cfg.head_dim_k as usize;
2052 let eps = cfg.rms_eps;
2053 let scale = 1.0 / (head_dim as f32).sqrt();
2054
2055 let gated = cfg.attn_out_gate();
2058 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
2060 let v = g3.pop().unwrap();
2061 let mut k = g3.pop().unwrap();
2062 let qf = g3.pop().unwrap();
2063 let (mut q, gate) = if gated {
2064 let mut q = e.uninit(t * n_head * head_dim)?;
2065 let mut gate = e.uninit(t * n_head * head_dim)?;
2066 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2067 (q, Some(gate))
2068 } else {
2069 (qf, None)
2070 };
2071
2072 let mut qn = e.uninit(t * n_head * head_dim)?;
2074 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2075 q = qn;
2076 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2077 e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2078 k = kn;
2079 let rope_dims = cfg.rope_dim_count as usize;
2080 e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, cfg.rope_freq_base, 1.0)?;
2081 e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, cfg.rope_freq_base, 1.0)?;
2082
2083 let mut attn = e.uninit(t * n_head * head_dim)?;
2085 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2088 e.sdpa_naive(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2090 } else {
2091 e.fa_prefill(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2092 }
2093
2094 let attn_g = match &gate {
2096 Some(gate) => {
2097 let mut gsig = e.uninit(t * n_head * head_dim)?;
2098 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
2099 let mut ag = e.uninit(t * n_head * head_dim)?;
2100 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
2101 ag
2102 }
2103 None => attn,
2104 };
2105
2106 let o = e.matmul(&fa.wo, &attn_g, t)?;
2108 Ok(o)
2109 }
2110
2111 pub fn linear_attn(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize)
2113 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2114 let cfg = &self.cfg;
2115 let _n_embd = cfg.n_embd as usize;
2116 let ssm = cfg.ssm.as_ref().unwrap();
2117 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;
2122 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;
2126 let scale = 1.0 / (d_state as f32).sqrt();
2127
2128 let mut g4 = e.matmul_group(&[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha], h, t)?;
2131 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);
2143 let mut q_g = e.uninit(d_state * num_v * t)?;
2144 let mut k_g = e.uninit(d_state * num_v * t)?;
2145 let mut v_g = e.uninit(d_state * num_v * t)?;
2146 e.ssm_conv1d_gdn(&qkv_mixed, la.ssm_conv1d.float_data(), &mut q_g, &mut k_g, &mut v_g,
2147 conv_dim, t, d_conv, d_state, num_v, num_k, key_dim)?;
2148 let mut q_l2 = e.uninit(d_state * num_v * t)?;
2150 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
2151 let mut k_l2 = e.uninit(d_state * num_v * t)?;
2152 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
2153 let v_gd = v_g;
2154
2155 let mut beta = e.uninit(t * num_v)?;
2158 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
2159 let mut g_log = e.uninit(t * num_v)?;
2161 e.gdn_glog(&alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
2162
2163 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
2166 let mut o = e.uninit(d_state * num_v * t)?;
2167 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)?;
2168
2169 let mut gn = e.uninit(d_state * num_v * t)?;
2174 e.gated_rmsnorm(&o, la.ssm_norm.float_data(), &z, &mut gn, d_state, num_v * t, eps)?;
2175
2176 let out = e.matmul(&la.ssm_out, &gn, t)?;
2180 Ok(out)
2181 }
2182}
2183
2184impl HybridModel {
2185 pub fn moe_ffn_il(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16)
2196 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2197 Self::moe_ffn(e, m, z, t, &self.cfg, il, self.max_moe_block())
2198 }
2199
2200 pub fn moe_ffn_il_zq8(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
2204 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16)
2205 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2206 Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block())
2207 }
2208
2209 pub(crate) fn moe_ffn(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
2217 cfg: &ModelConfig, il: u16, max_block: usize)
2218 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2219 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block)
2220 }
2221
2222 #[allow(clippy::too_many_arguments)]
2223 pub(crate) fn moe_ffn_inner(
2224 e: &Engine,
2225 m: &MoeWeights,
2226 z: &CudaSlice<f32>,
2227 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
2228 t: usize,
2229 cfg: &ModelConfig,
2230 il: u16,
2231 max_block: usize,
2232 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2233 let worker_io = crate::spill_pread::worker_enabled();
2234 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
2235 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
2236 e.with_moe_cache(max_block, |cache, _| {
2237 cache.begin_forward_epoch(il, t);
2238 if worker_io {
2239 cache.begin_worker_scope();
2240 }
2241 Ok(())
2242 })?;
2243 }
2244 if t > 1 && std::env::var("MEMRA_MOE_GROUPED").is_ok() {
2246 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
2247 if std::env::var("MEMRA_MOE_GATE").is_ok() {
2254 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
2255 let g_host = e.dtoh(&grouped_out)?;
2256 let s_host = e.dtoh(&seq_out)?;
2257 let g_bytes: &[u8] = unsafe { std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4) };
2258 let s_bytes: &[u8] = unsafe { std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4) };
2259 if g_bytes == s_bytes {
2260 if il == 0 { println!("moe-gate il={il} t={t} BYTE-IDENTICAL (first layer only printed)"); }
2261 } else {
2262 let diffs = g_host.iter().zip(s_host.iter()).enumerate()
2263 .filter(|(_, (a, b))| a != b).count();
2264 let maxdiff = g_host.iter().zip(s_host.iter())
2265 .map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
2266 panic!("moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}", g_host.len());
2267 }
2268 }
2269 return Ok(grouped_out);
2270 }
2271 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
2272 }
2273
2274 pub(crate) fn moe_ffn_sequential(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
2276 cfg: &ModelConfig, il: u16, max_block: usize)
2277 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2278 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
2279 }
2280
2281 fn trace_moe_routes(il: u16, t: usize, sel_all: &[u32], weights: &[f32])
2285 -> Result<(), Box<dyn std::error::Error>> {
2286 use std::io::Write as _;
2287 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
2288 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
2289 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
2290 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
2291 }
2292 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
2293 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
2294 let pairs: Vec<String> = sel_all.iter().zip(weights)
2295 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
2296 .collect();
2297 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
2298 }
2299 Ok(())
2300 }
2301
2302 fn trace_moe_input(e: &Engine, il: u16, t: usize, n_embd: usize, z: &CudaSlice<f32>)
2307 -> Result<(), Box<dyn std::error::Error>> {
2308 use std::io::Write as _;
2309 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else { return Ok(()) };
2310 let host = e.dtoh(z)?;
2311 if host.len() != t * n_embd {
2312 return Err(format!(
2313 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
2314 host.len(), t, n_embd
2315 ).into());
2316 }
2317 let bytes = unsafe {
2318 std::slice::from_raw_parts(
2319 host.as_ptr().cast::<u8>(), host.len() * std::mem::size_of::<f32>()
2320 )
2321 };
2322 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
2323 let mut state = state.lock().map_err(|_| "MoE input trace writer lock is poisoned")?;
2324 if state.is_none() {
2325 let dir = std::path::PathBuf::from(&dir);
2326 std::fs::create_dir_all(&dir)?;
2327 let index = std::fs::OpenOptions::new().create(true).append(true)
2328 .open(dir.join("index.jsonl"))?;
2329 *state = Some(MoeInputTraceWriter {
2330 dir,
2331 index,
2332 payloads: std::collections::HashMap::new(),
2333 });
2334 }
2335 let writer = state.as_mut().unwrap();
2336 if writer.dir != std::path::Path::new(&dir) {
2337 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
2338 }
2339 let file_name = format!("layer-{il:03}.f32");
2340 if !writer.payloads.contains_key(&il) {
2341 let payload = std::fs::OpenOptions::new().create(true).append(true)
2342 .open(writer.dir.join(&file_name))?;
2343 let offset = payload.metadata()?.len();
2344 writer.payloads.insert(il, (payload, offset));
2345 }
2346 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
2347 let row_offset = *offset;
2348 payload.write_all(bytes)?;
2349 *offset += bytes.len() as u64;
2350 writeln!(
2351 writer.index,
2352 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
2353 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
2354 \"payload_bytes\":{}}}",
2355 bytes.len()
2356 )?;
2357 Ok(())
2358 }
2359
2360 #[allow(clippy::too_many_arguments)]
2361 pub(crate) fn moe_ffn_sequential_zq8(
2362 e: &Engine,
2363 m: &MoeWeights,
2364 z: &CudaSlice<f32>,
2365 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
2366 t: usize,
2367 cfg: &ModelConfig,
2368 il: u16,
2369 max_block: usize,
2370 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2371 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
2372 let moe = cfg.moe.as_ref().unwrap();
2373 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);
2380 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
2381 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);
2384
2385 let lim_exp = cfg.clamp_exp_at(il as u32);
2388 let lim_shexp = cfg.clamp_shexp_at(il as u32);
2389 let use_cache = Engine::moe_cache_enabled();
2390 let uniform_experts = m.has_uniform_expert_layout();
2391 let moe_q8 = uniform_experts && moe_q8_enabled()
2392 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
2393 && q8_expert_supported(m.down_exps.qtype);
2394 let cpu_expert_requested = crate::cpu_experts::configured();
2401 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
2402 return Err(std::io::Error::other(
2403 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
2404 )
2405 .into());
2406 }
2407 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
2408 let freeze_cpu_residency = cpu_expert_requested
2414 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
2415 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
2416 .ok()
2417 .and_then(|value| value.parse::<usize>().ok())
2418 .is_some_and(|tokens| tokens > 0);
2419 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
2420 e.freeze_moe_cache();
2421 }
2422 let cache_frozen = use_cache && e.moe_cache_frozen();
2423 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
2424
2425 let logits = if t < PRIME_MIN_T {
2432 if crate::router_kernel_on() {
2436 e.router_gemv(m.gate_inp.float_data(), z, cfg.n_embd as usize,
2439 m.gate_exps.n_expert, t)?
2440 } else {
2441 e.matmul_decode_exact(&m.gate_inp, z, t)?
2442 }
2443 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
2444 e.router_gemv(m.gate_inp.float_data(), z, cfg.n_embd as usize,
2459 m.gate_exps.n_expert, t)?
2460 } else {
2461 e.matmul(&m.gate_inp, z, t)?
2462 };
2463
2464 let no_exp_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
2502 && m.down_exps.macros.is_none();
2503 if cfg.sigmoid_router().is_none() && cfg.m3.is_none() && cfg.hy3.is_none()
2507 && !cfg.swiglu_clamped_at(il as u32)
2508 && no_exp_macros
2509 && t >= PRIME_MIN_T && m.dev_exps.is_some() && moe_q8_enabled()
2510 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
2511 && q8_expert_supported(m.down_exps.qtype)
2512 && std::env::var("MEMRA_MOE_PAIRS").map(|v| v != "0").unwrap_or(true)
2513 && std::env::var("MEMRA_MOE_STATS").is_err() {
2514 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
2515 }
2516
2517 let dev_ok = uniform_experts && cfg.sigmoid_router().is_none()
2534 && cfg.m3.is_none() && cfg.hy3.is_none()
2535 && !cfg.swiglu_clamped_at(il as u32);
2536 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
2540 || std::env::var("MEMRA_MOE_TRACE").is_ok()
2541 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
2542 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
2543 if dev_ok && t < PRIME_MIN_T && m.dev_exps.is_some() && n_used <= 8 && moe_dev_enabled()
2544 && !observe_routes {
2545 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
2546 }
2547 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled()
2548 && !observe_routes {
2549 let row_ok = e.with_moe_cache(max_block, |c, eng| {
2550 if moe_prewarm_enabled() { c.prewarm_layer(il, m, eng)?; }
2551 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
2552 })?;
2553 if row_ok {
2554 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
2555 }
2556 }
2557
2558 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
2560 if cpu_hybrid {
2561 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
2562 e,
2563 &logits,
2564 z,
2565 t,
2566 n_expert,
2567 n_used,
2568 m.exp_probs_b.as_deref(),
2569 sig,
2570 m.active_experts.as_deref(),
2571 )?;
2572 (sel, w, Some(input))
2573 } else {
2574 let (sel, w) = Self::moe_route_cfg(
2575 e,
2576 &logits,
2577 t,
2578 n_expert,
2579 n_used,
2580 m.exp_probs_b.as_deref(),
2581 Some(sig),
2582 m.active_experts.as_deref(),
2583 )?;
2584 (sel, w, None)
2585 }
2586 } else {
2587 let (sel, w) = Self::moe_route_cfg(
2588 e,
2589 &logits,
2590 t,
2591 n_expert,
2592 n_used,
2593 None,
2594 None,
2595 m.active_experts.as_deref(),
2596 )?;
2597 (sel, w, None)
2598 };
2599
2600 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
2604 Self::trace_moe_input(e, il, t, n_embd, z)?;
2605
2606 let worker_disk_prefetch =
2618 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
2619 let promote_worker_h2d =
2620 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
2621 if promote_worker_h2d {
2622 let mut selected_blocks = Vec::with_capacity(n_used * 3);
2623 for &ex in sel_all.iter().take(n_used) {
2624 let ex = ex as u16;
2625 selected_blocks.extend([
2626 BlockId::new(il, PROJ_GATE, ex),
2627 BlockId::new(il, PROJ_UP, ex),
2628 BlockId::new(il, PROJ_DOWN, ex),
2629 ]);
2630 }
2631 for &ex in sel_all.iter().take(n_used) {
2632 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
2633 }
2634 e.with_moe_cache(max_block, |cache, eng| {
2635 cache.promote_worker_reads_at_safe_boundary(
2636 &selected_blocks,
2637 &selected_blocks,
2638 eng,
2639 )?;
2640 Ok(())
2641 })?;
2642 }
2643
2644 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
2647 let mut cnt = vec![0u32; n_expert];
2648 for &s in sel_all.iter() { cnt[s as usize] += 1; }
2649 let total = sel_all.len() as f64;
2650 let mut h = 0.0f64;
2651 let mut active = 0usize;
2652 for &c in &cnt { if c > 0 { active += 1; let p = c as f64 / total; h -= p * p.log2(); } }
2653 let maxc = cnt.iter().copied().max().unwrap_or(0);
2654 println!("moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
2655 il, t, sel_all.len(), active, n_expert, h, (n_expert as f64).log2(), total / active.max(1) as f64, maxc);
2656 }
2657
2658 let gdec_may_fire = uniform_experts && use_cache && n_used <= 8 && gdec_enabled()
2671 && !cfg.swiglu_clamped_at(il as u32);
2672 let slab_local = m.dev_exps.as_ref()
2688 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
2689 let slab_bases = slab_local.map(|d| {
2690 use cudarc::driver::DevicePtr;
2691 let s = e.stream();
2692 let (pg, _g0) = d.gate.device_ptr(&s);
2693 let (pu, _g1) = d.up.device_ptr(&s);
2694 let (pd, _g2) = d.down.device_ptr(&s);
2695 (pg as u64, pu as u64, pd as u64)
2696 });
2697 let slab_fused_may_fire = slab_bases.is_some() && n_used <= 8 && gdec_enabled()
2707 && !cfg.swiglu_clamped_at(il as u32) && cfg.m3.is_none()
2708 && no_exp_macros && moe_q8;
2709 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
2712 e.uninit(t * n_embd)?
2713 } else {
2714 e.zeros(t * n_embd)?
2715 };
2716 let cpu_input = if cpu_hybrid {
2719 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
2720 } else {
2721 None
2722 };
2723
2724 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;
2732 let mut scratch_u: Option<CudaSlice<u8>> = None;
2733 let mut scratch_d: Option<CudaSlice<u8>> = None;
2734 let page_window = moe_page_prefetch_window();
2742
2743 for tok in 0..t {
2746 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
2747 let w = &w_all[tok * n_used..(tok + 1) * n_used];
2748 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
2750
2751 let no_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
2765 && m.down_exps.macros.is_none();
2766 if slab_fused_may_fire {
2776 let (pg, pu, pd) = slab_bases.unwrap();
2777 let mut gp = [0u64; 8];
2778 let mut up = [0u64; 8];
2779 let mut dp = [0u64; 8];
2780 for (j, &ex) in sel.iter().enumerate() {
2781 let ex = ex as usize;
2782 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
2783 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
2784 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
2785 }
2786 let mut wv = [0f32; 8];
2787 wv[..n_used].copy_from_slice(w);
2788 if tok_q8.is_none() {
2789 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2790 }
2791 let (zq, zd) = tok_q8.as_ref().unwrap();
2792 let act = e.moe_gate_up_silu8_q8(crate::WPtr8(gp), crate::WPtr8(up), zq, zd,
2793 n_embd, n_ff_exp, n_used,
2794 m.gate_exps.qtype, m.up_exps.qtype,
2795 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
2796 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
2797 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2798 e.moe_down8_fma_q8(crate::WPtr8(dp), crate::F32x8(wv), &aq2, &ad2, &mut dst,
2799 n_ff_exp, n_embd, n_used,
2800 m.down_exps.qtype, m.down_exps.row_bytes)?;
2801 continue;
2802 }
2803 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
2804 if tok_q8.is_none() {
2805 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2806 }
2807 let (zq, zd) = tok_q8.as_ref().unwrap();
2808 if Self::moe_gdec_token_q8(e, m, il, max_block, zq, zd, sel, w,
2809 &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
2810 continue;
2811 }
2812 } else if gdec_may_fire && cfg.m3.is_none() && no_macros
2813 && Self::moe_gdec_token(e, m, il, max_block, &zt, sel, w,
2814 &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
2815 continue;
2816 }
2817
2818 if gdec_may_fire || slab_fused_may_fire {
2824 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2825 e.memset_zeros_view(&mut row)?;
2826 }
2827
2828 let mut cpu_mask = vec![false; sel.len()];
2834 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
2835 let gpu_resident = if use_cache {
2836 e.with_moe_cache(max_block, |cache, _| {
2837 Ok(sel
2838 .iter()
2839 .map(|&expert| {
2840 let expert = expert as u16;
2841 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
2842 .into_iter()
2843 .filter(|&projection| {
2844 cache
2845 .resident(BlockId::new(il, projection, expert))
2846 .is_some()
2847 })
2848 .count()
2849 })
2850 .collect::<Vec<_>>())
2851 })?
2852 } else {
2853 vec![0; sel.len()]
2854 };
2855 let mut cpu_selected = Vec::new();
2856 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
2857 if gpu_resident[index] != 3 {
2858 cpu_mask[index] = true;
2859 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
2860 let expert = expert as usize;
2861 cpu_selected.push((expert, route_weight));
2862 }
2863 }
2864 if crate::cpu_experts::predictor_enabled() {
2865 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
2869 crate::cpu_experts::predictor_submit(il, row);
2870 }
2871 if cpu_selected.is_empty() {
2872 None
2873 } else {
2874 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
2875 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
2876 .map_err(std::io::Error::other)?;
2877 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
2878 }
2879 } else {
2880 None
2881 };
2882
2883 let worker_window = worker_disk_prefetch
2884 .then(worker_prefetch_window)
2885 .unwrap_or(0);
2886 for (j, &ex) in sel.iter().enumerate() {
2887 if cpu_mask[j] {
2888 continue;
2889 }
2890 let ex = ex as usize;
2891 if let Some(d) = slab_local {
2898 let gl = m.gate_exps.expert_layout(ex);
2899 let ul = m.up_exps.expert_layout(ex);
2900 let dl = m.down_exps.expert_layout(ex);
2901 let (g0, u0, d0) = (ex * m.gate_exps.expert_stride,
2902 ex * m.up_exps.expert_stride,
2903 ex * m.down_exps.expert_stride);
2904 let (gate, up) = if moe_q8 {
2905 if tok_q8.is_none() {
2906 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2907 }
2908 let (zq, zd) = tok_q8.as_ref().unwrap();
2909 (e.qmatvec_expert_q8(&d.gate, g0..g0 + gl.len, zq, zd, 1,
2910 m.gate_exps.in_f, m.gate_exps.out_f,
2911 gl.qtype, gl.row_bytes)?,
2912 e.qmatvec_expert_q8(&d.up, u0..u0 + ul.len, zq, zd, 1,
2913 m.up_exps.in_f, m.up_exps.out_f,
2914 ul.qtype, ul.row_bytes)?)
2915 } else {
2916 (e.qmatvec_view(&d.gate, g0..g0 + gl.len, &zt, 1,
2917 m.gate_exps.in_f, m.gate_exps.out_f,
2918 gl.qtype, gl.row_bytes)?,
2919 e.qmatvec_view(&d.up, u0..u0 + ul.len, &zt, 1,
2920 m.up_exps.in_f, m.up_exps.out_f,
2921 ul.qtype, ul.row_bytes)?)
2922 };
2923 let mut act = e.uninit(n_ff_exp)?;
2924 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
2925 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
2926 let y = if moe_q8 {
2927 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
2928 e.qmatvec_expert_q8(&d.down, d0..d0 + dl.len, &aq2, &ad2, 1,
2929 m.down_exps.in_f, m.down_exps.out_f,
2930 dl.qtype, dl.row_bytes)?
2931 } else {
2932 let actv = act.slice(0..n_ff_exp);
2933 e.qmatvec_view(&d.down, d0..d0 + dl.len, &actv, 1,
2934 m.down_exps.in_f, m.down_exps.out_f,
2935 dl.qtype, dl.row_bytes)?
2936 };
2937 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2938 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
2939 continue;
2940 }
2941 for next in page_prefetch_positions(j, sel.len(), page_window) {
2942 Self::moe_prefetch_host_expert(sel[next] as usize, m);
2943 }
2944 let keep = [
2945 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
2946 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
2947 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
2948 ];
2949 if worker_disk_prefetch && worker_window > 0 {
2950 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
2951 Self::moe_prefetch_disk_expert(
2952 e,
2953 il,
2954 sel[next] as usize,
2955 m,
2956 max_block,
2957 &keep,
2958 )?;
2959 }
2960 } else if cache_dispatch
2961 && !cpu_hybrid
2962 && moe_prefetch_enabled()
2963 && j + 1 < sel.len()
2964 {
2965 let next = sel[j + 1] as usize;
2966 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
2967 }
2968 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
2969 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
2970 if (gate_q8 || up_q8) && tok_q8.is_none() {
2973 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2974 }
2975 let gate = if gate_q8 {
2976 let (zq, zd) = tok_q8.as_ref().unwrap();
2977 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
2978 } else {
2979 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
2980 };
2981 let up = if up_q8 {
2982 let (zq, zd) = tok_q8.as_ref().unwrap();
2983 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
2984 } else {
2985 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
2986 };
2987 let mut act = e.uninit(n_ff_exp)?;
2988 Self::ffn_act_lim(
2989 e,
2990 cfg,
2991 &gate,
2992 &up,
2993 m.gate_exps.macro_scale(ex),
2994 m.up_exps.macro_scale(ex),
2995 lim_exp,
2996 &mut act,
2997 n_ff_exp,
2998 )?;
2999 let y = if down_q8 {
3000 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
3001 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
3002 } else {
3003 let actv = act.slice(0..n_ff_exp);
3004 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
3005 };
3006 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3007 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3009 } else if cache_dispatch {
3010 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
3015 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
3016 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3018 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3019 let actv = act.slice(0..n_ff_exp);
3020 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
3021 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3022 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3024 } else if cache_frozen {
3025 let gate = Self::moe_frozen_gemm(
3030 e,
3031 il,
3032 PROJ_GATE,
3033 ex,
3034 m,
3035 max_block,
3036 &zt,
3037 &mut scratch_g,
3038 g_len,
3039 )?;
3040 let up = Self::moe_frozen_gemm(
3041 e,
3042 il,
3043 PROJ_UP,
3044 ex,
3045 m,
3046 max_block,
3047 &zt,
3048 &mut scratch_u,
3049 u_len,
3050 )?;
3051 let mut act = e.uninit(n_ff_exp)?;
3052 Self::ffn_act_lim(
3053 e,
3054 cfg,
3055 &gate,
3056 &up,
3057 m.gate_exps.macro_scale(ex),
3058 m.up_exps.macro_scale(ex),
3059 lim_exp,
3060 &mut act,
3061 n_ff_exp,
3062 )?;
3063 let actv = act.slice(0..n_ff_exp);
3064 let y = Self::moe_frozen_gemm(
3065 e,
3066 il,
3067 PROJ_DOWN,
3068 ex,
3069 m,
3070 max_block,
3071 &actv,
3072 &mut scratch_d,
3073 d_len,
3074 )?;
3075 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3076 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3077 } else {
3078 if scratch_g.is_none() {
3082 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
3083 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
3084 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
3085 }
3086 let (sg, su, sd) = (scratch_g.as_mut().unwrap(), scratch_u.as_mut().unwrap(),
3087 scratch_d.as_mut().unwrap());
3088 let gl = m.gate_exps.expert_layout(ex);
3089 let ul = m.up_exps.expert_layout(ex);
3090 let dl = m.down_exps.expert_layout(ex);
3091 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
3092 let gate = e.qmatvec_view(sg, 0..gl.len, &zt, 1,
3093 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
3094
3095 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
3096 let up = e.qmatvec_view(su, 0..ul.len, &zt, 1,
3097 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
3098
3099 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3101 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3102
3103 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
3104 let actv = act.slice(0..n_ff_exp);
3105 let y = e.qmatvec_view(sd, 0..dl.len, &actv, 1,
3106 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?;
3107
3108 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3109 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3110 }
3111 }
3112 if let Some(worker) = cpu_worker {
3113 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
3114 let cpu_output = e.htod(&cpu_output)?;
3115 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3116 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
3117 }
3118 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
3119 for (j, &ex) in sel.iter().enumerate() {
3120 if cpu_mask[j] {
3121 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
3122 }
3123 }
3124 }
3125 }
3126
3127 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3132 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3133 {
3134 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
3143 let (sg_gate, sg_up) = if t == 1 {
3144 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
3145 Some(pair) => pair,
3146 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
3147 }
3148 } else if verify_t {
3149 (e.matmul_decode_exact(gate_shexp, z, t)?, e.matmul_decode_exact(up_shexp, z, t)?)
3150 } else {
3151 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
3153 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)?;
3155 let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
3156 else { e.matmul(down_shexp, &sa, t)? }; let g = match &m.gate_inp_shexp {
3170 Some(gate_inp_shexp) => {
3171 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
3172 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3173 } else {
3174 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3175 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
3177 g
3178 }
3179 }
3180 None => e.htod(&vec![1.0f32; t])?,
3181 };
3182 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3184 }
3185
3186 Ok(moe_out)
3187 }
3188
3189 pub fn stage1_h2d_per_token(&self) -> u64 {
3192 use crate::hybrid::Ffn;
3193 let n_used = self.cfg.moe.as_ref().map(|m| m.expert_used_count as u64).unwrap_or(0);
3194 let mut bytes = 0u64;
3195 for l in self.layers.iter() {
3196 if let Ffn::Moe(m) = &l.ffn {
3197 bytes += n_used * (m.gate_exps.max_expert_bytes() + m.up_exps.max_expert_bytes()
3198 + m.down_exps.max_expert_bytes()) as u64;
3199 }
3200 }
3201 bytes
3202 }
3203
3204 pub(crate) fn max_moe_block(&self) -> usize {
3208 use crate::hybrid::Ffn;
3209 let mut mx = 0usize;
3210 let mut scan = |ffn: &Ffn| {
3211 if let Ffn::Moe(m) = ffn {
3212 mx = mx.max(m.gate_exps.max_expert_bytes())
3213 .max(m.up_exps.max_expert_bytes())
3214 .max(m.down_exps.max_expert_bytes());
3215 }
3216 };
3217 for l in self.layers.iter() { scan(&l.ffn); }
3218 if let Some(mtp) = self.mtp.as_ref() { scan(&mtp.ffn); }
3219 mx
3220 }
3221
3222 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
3225 use crate::hybrid::Ffn;
3226 let mut sizes = Vec::new();
3227 let mut scan = |ffn: &Ffn| {
3228 let Ffn::Moe(m) = ffn else { return };
3229 for ex in 0..m.gate_exps.n_expert {
3230 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
3231 continue;
3232 }
3233 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
3234 let len = exps.expert_layout(ex).len;
3235 if len > 0 {
3236 sizes.push(len);
3237 }
3238 }
3239 }
3240 };
3241 for layer in &self.layers {
3242 scan(&layer.ffn);
3243 }
3244 if let Some(mtp) = &self.mtp {
3245 scan(&mtp.ffn);
3246 }
3247 sizes
3248 }
3249
3250 pub fn save_cpu_expert_residency_profile(
3256 &self,
3257 e: &Engine,
3258 path: &std::path::Path,
3259 ) -> Result<(), Box<dyn std::error::Error>> {
3260 let Some(ids) = e.export_moe_residency() else {
3261 return Err("no MoE residency cache to persist".into());
3262 };
3263 let mut body = format!(
3264 "memra-freeze-profile v1 max_block={} blocks={}\n",
3265 self.max_moe_block(),
3266 ids.len()
3267 );
3268 for (layer, proj, ex) in &ids {
3269 body.push_str(&format!("{layer} {proj} {ex}\n"));
3270 }
3271 let tmp = path.with_extension("tmp");
3272 std::fs::write(&tmp, body)?;
3273 std::fs::rename(&tmp, path)?;
3274 println!(
3275 "[moe-cache] freeze profile saved: {} blocks -> {}",
3276 ids.len(),
3277 path.display()
3278 );
3279 Ok(())
3280 }
3281
3282 pub fn restore_cpu_expert_residency_profile(
3286 &self,
3287 e: &Engine,
3288 path: &std::path::Path,
3289 ) -> Result<bool, Box<dyn std::error::Error>> {
3290 use crate::hybrid::Ffn;
3291 use crate::moe_cache::BlockId;
3292 let Ok(content) = std::fs::read_to_string(path) else {
3293 return Ok(false);
3294 };
3295 let mut lines = content.lines();
3296 let Some(header) = lines.next() else { return Ok(false) };
3297 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
3298 if !header.starts_with(&expected) {
3299 println!(
3300 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
3301 path.display()
3302 );
3303 return Ok(false);
3304 }
3305 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
3306 std::collections::HashMap::new();
3307 for line in lines {
3308 let mut fields = line.split_whitespace();
3309 let (Some(layer), Some(proj), Some(ex)) =
3310 (fields.next(), fields.next(), fields.next())
3311 else {
3312 continue;
3313 };
3314 let (Ok(layer), Ok(proj), Ok(ex)) =
3315 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
3316 else {
3317 continue;
3318 };
3319 by_layer
3320 .entry(layer)
3321 .or_default()
3322 .push(BlockId::new(layer, proj, ex));
3323 }
3324 let requested: usize = by_layer.values().map(Vec::len).sum();
3325 if requested == 0 {
3326 return Ok(false);
3327 }
3328 let max_block = self.max_moe_block();
3329 let mut restaged = 0usize;
3330 let mut stage_layer = |layer_index: u16,
3331 ffn: &Ffn|
3332 -> Result<(), Box<dyn std::error::Error>> {
3333 let Ffn::Moe(m) = ffn else { return Ok(()) };
3334 let Some(ids) = by_layer.get(&layer_index) else {
3335 return Ok(());
3336 };
3337 e.with_moe_cache(max_block, |cache, eng| {
3338 for id in ids {
3339 if cache.restage_block(*id, m, eng)? {
3340 restaged += 1;
3341 }
3342 }
3343 Ok(())
3344 })
3345 };
3346 for (index, layer) in self.layers.iter().enumerate() {
3347 stage_layer(index as u16, &layer.ffn)?;
3348 }
3349 if let Some(mtp) = self.mtp.as_ref() {
3350 stage_layer(u16::MAX, &mtp.ffn)?;
3351 }
3352 e.freeze_moe_cache();
3353 println!(
3354 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
3355 path.display()
3356 );
3357 Ok(true)
3358 }
3359
3360 pub fn freeze_cpu_expert_residency(
3362 &self,
3363 e: &Engine,
3364 ) -> Result<(), Box<dyn std::error::Error>> {
3365 e.freeze_moe_cache();
3366 Ok(())
3367 }
3368
3369 pub fn ffn_act(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
3377 act: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
3378 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
3379 }
3380
3381 #[allow(clippy::too_many_arguments)]
3385 pub(crate) fn ffn_act_scaled(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
3386 gs: f32, us: f32, act: &mut CudaSlice<f32>, n: usize)
3387 -> Result<(), Box<dyn std::error::Error>> {
3388 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
3389 }
3390
3391 #[allow(clippy::too_many_arguments)]
3400 pub(crate) fn ffn_act_lim(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
3401 gs: f32, us: f32, limit: Option<f32>, act: &mut CudaSlice<f32>, n: usize)
3402 -> Result<(), Box<dyn std::error::Error>> {
3403 if let Some(m3) = cfg.m3.as_ref() {
3404 debug_assert!(limit.is_none(), "m3 swigluoai and step35 clamp are different archs");
3405 return e.swigluoai_mul_scaled(gate, up, gs, us, m3.swiglu_alpha, m3.swiglu_limit, act, n);
3406 }
3407 if let Some(l) = limit {
3408 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
3409 }
3410 if gs == 1.0 && us == 1.0 { return e.silu_mul(gate, up, act, n); }
3411 e.silu_mul_scaled(gate, up, gs, us, act, n)
3412 }
3413
3414 fn moe_route(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
3420 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3421 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None, None, None)
3422 }
3423
3424 fn moe_route_cfg(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize,
3432 bias: Option<&[f32]>, sig: Option<(f32, bool)>, active: Option<&[bool]>)
3433 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3434 if let Some((sf, route_norm)) = sig {
3435 let lg = e.dtoh(logits)?;
3437 return Self::moe_route_sigmoid_host(
3438 &lg, t, n_expert, n_used, bias, sf, route_norm, active,
3439 );
3440 }
3441 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
3445 return e.moe_router_topk_host(logits, t, n_expert, n_used);
3446 }
3447 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
3450 let mut w_out = vec![0f32; t * n_used];
3451 for tok in 0..t {
3452 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
3453 let maxl = row.iter().enumerate()
3455 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
3456 .map(|(_, &x)| x).fold(f32::NEG_INFINITY, f32::max);
3457 let mut probs = vec![0f32; n_expert];
3458 let mut den = 0f32;
3459 for i in 0..n_expert {
3460 if active.is_some_and(|mask| !mask[i]) { continue; }
3461 let x = (row[i] - maxl).exp(); probs[i] = x; den += x;
3462 }
3463 for p in probs.iter_mut() { *p /= den; }
3464 let mut idx: Vec<usize> = (0..n_expert)
3466 .filter(|&i| active.is_none_or(|mask| mask[i])).collect();
3467 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
3468 let sl = &idx[..n_used];
3469 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
3470 let mut ws: f32 = wv.iter().sum();
3471 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() { *x /= ws; }
3473 for j in 0..n_used {
3474 sel[tok * n_used + j] = sl[j] as u32;
3475 w_out[tok * n_used + j] = wv[j];
3476 }
3477 }
3478 Ok((sel, w_out))
3479 }
3480
3481 #[allow(clippy::too_many_arguments)]
3482 fn moe_route_sigmoid_with_input(
3483 e: &Engine,
3484 logits: &CudaSlice<f32>,
3485 input: &CudaSlice<f32>,
3486 t: usize,
3487 n_expert: usize,
3488 n_used: usize,
3489 bias: Option<&[f32]>,
3490 (sf, route_norm): (f32, bool),
3491 active: Option<&[bool]>,
3492 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
3493 let (lg, input) = e.dtoh_pair(logits, input)?;
3494 let (sel, w) =
3495 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
3496 Ok((sel, w, input))
3497 }
3498
3499 pub fn start_moe_prefetch_predictor(
3504 &self,
3505 e: &Engine,
3506 cfg: &ModelConfig,
3507 ) -> Result<(), Box<dyn std::error::Error>> {
3508 use crate::hybrid::Ffn;
3509 let Some(sig) = cfg.sigmoid_router() else {
3510 return Err("prefetch predictor requires a sigmoid-router arch".into());
3511 };
3512 let resident: std::collections::HashSet<(u16, u8, u16)> = e
3513 .export_moe_residency()
3514 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
3515 .into_iter()
3516 .collect();
3517 let mut layers = Vec::new();
3518 for (index, layer) in self.layers.iter().enumerate() {
3519 let Ffn::Moe(m) = &layer.ffn else { continue };
3520 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else { continue };
3521 let router = e.dtoh(data)?;
3522 let n_expert = m.gate_exps.n_expert;
3523 let n_embd = m.gate_exps.in_f;
3524 if router.len() != n_embd * n_expert {
3525 continue;
3526 }
3527 let build = |exps: &crate::model::HostExps| {
3528 (0..n_expert)
3529 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
3530 .collect::<Vec<_>>()
3531 };
3532 layers.push((index as u16, crate::cpu_experts::PredictLayerInit {
3533 router,
3534 bias: m.exp_probs_b.clone(),
3535 active: m.active_experts.clone(),
3536 n_embd,
3537 n_used: cfg
3538 .moe
3539 .as_ref()
3540 .map(|moe| moe.expert_used_count as usize)
3541 .ok_or("prefetch predictor requires MoE config")?,
3542 sig,
3543 weights_n_expert: n_expert,
3544 gate: build(&m.gate_exps),
3545 up: build(&m.up_exps),
3546 down: build(&m.down_exps),
3547 }));
3548 }
3549 crate::cpu_experts::start_prefetch_predictor(layers, resident)
3550 .map_err(|error| error.into())
3551 }
3552
3553 #[allow(clippy::too_many_arguments)]
3556 pub(crate) fn moe_route_sigmoid_host_public(
3557 logits: &[f32],
3558 t: usize,
3559 n_expert: usize,
3560 n_used: usize,
3561 bias: Option<&[f32]>,
3562 sf: f32,
3563 route_norm: bool,
3564 active: Option<&[bool]>,
3565 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3566 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
3567 }
3568
3569 #[allow(clippy::too_many_arguments)]
3570 fn moe_route_sigmoid_host(
3571 lg: &[f32],
3572 t: usize,
3573 n_expert: usize,
3574 n_used: usize,
3575 bias: Option<&[f32]>,
3576 sf: f32,
3577 route_norm: bool,
3578 active: Option<&[bool]>,
3579 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3580 if lg.len() != t * n_expert {
3581 return Err(format!(
3582 "sigmoid router logits length mismatch: got {}, expected {}",
3583 lg.len(),
3584 t * n_expert,
3585 )
3586 .into());
3587 }
3588 let mut sel = vec![0u32; t * n_used];
3589 let mut w_out = vec![0f32; t * n_used];
3590 for tok in 0..t {
3591 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
3592 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
3593 let selsc: Vec<f32> = match bias {
3595 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
3596 None => scores.clone(),
3597 };
3598 let mut idx: Vec<usize> = (0..n_expert)
3599 .filter(|&i| active.is_none_or(|mask| mask[i]))
3600 .collect();
3601 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
3602 let sl = &idx[..n_used];
3603 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
3604 if route_norm {
3605 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
3606 for x in wv.iter_mut() {
3607 *x = *x / ws * sf;
3608 }
3609 } else {
3610 for x in wv.iter_mut() {
3611 *x *= sf;
3612 }
3613 }
3614 for j in 0..n_used {
3615 sel[tok * n_used + j] = sl[j] as u32;
3616 w_out[tok * n_used + j] = wv[j];
3617 }
3618 }
3619 Ok((sel, w_out))
3620 }
3621
3622 fn moe_ffn_pairs(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, logits: &CudaSlice<f32>,
3631 t: usize, cfg: &ModelConfig)
3632 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3633 let moe = cfg.moe.as_ref().unwrap();
3634 let n_embd = cfg.n_embd as usize;
3635 let n_expert = moe.expert_count as usize;
3636 let n_used = moe.expert_used_count as usize;
3637 let n_ff_exp = moe.expert_ff_length as usize;
3638 debug_assert!(!cfg.swiglu_clamped_anywhere(),
3643 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU");
3644 let dev = m.dev_exps.as_ref().unwrap();
3645 let (rbg_d, rbu_d) = if dev.gu_il {
3647 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
3648 } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
3649
3650 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
3651 let n_pairs = t * n_used;
3652 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
3655 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
3656 let pair_w: Vec<f32> = w_all.clone();
3657 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
3658 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
3659 let pt = e.htod_i32(&pair_tok)?;
3660 let px = e.htod_i32(&pair_ex)?;
3661 let pw = e.htod(&pair_w)?;
3662 let toff = e.htod_i32(&tok_off)?;
3663 let tids = e.htod_i32(&tok_ids)?;
3664
3665 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
3669 for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
3670 let mut ex_ids: Vec<i32> = Vec::new();
3671 let mut ex_off: Vec<i32> = vec![0];
3672 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
3673 for (ex, list) in by_ex.iter().enumerate() {
3674 if list.is_empty() { continue; }
3675 ex_ids.push(ex as i32);
3676 ex_pairs.extend_from_slice(list);
3677 ex_off.push(ex_pairs.len() as i32);
3678 }
3679 let n_active = ex_ids.len();
3680 let exi = e.htod_i32(&ex_ids)?;
3681 let exo = e.htod_i32(&ex_off)?;
3682 let exp_d = e.htod_i32(&ex_pairs)?;
3683 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
3704 let mma_t = *MMA_T.get_or_init(|| {
3705 std::env::var("MEMRA_MOE_MMA_T").ok().and_then(|v| v.parse().ok()).unwrap_or(16)
3706 });
3707 let use_mma = std::env::var("MEMRA_MOE_MMA").map(|v| v != "0").unwrap_or(true)
3708 && t >= mma_t
3709 && q8_expert_dec_supported(m.gate_exps.qtype) && q8_expert_dec_supported(m.up_exps.qtype)
3710 && q8_expert_dec_supported(m.down_exps.qtype)
3711 && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
3712 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
3728 && q8_expert_dec_supported(m.up_exps.qtype)
3729 && q8_expert_dec_supported(m.down_exps.qtype)
3730 && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
3731 let f16g_mode = crate::moe_f16g_mode();
3732 let f16g = f16g_mode != 0 && t >= mma_t
3733 && (f16g_mode != 3 || !mma_capable)
3734 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
3735 && f16g_proj_ok(m.up_exps.qtype, n_embd)
3736 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
3737 if use_mma || f16g {
3738 let y_down = if f16g {
3746 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
3750 let csr_tok_d = e.htod_i32(&csr_tok)?;
3751 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
3752 let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
3753 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
3754 m.gate_exps.qtype, rbg_d)?;
3755 let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
3756 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
3757 m.up_exps.qtype, rbu_d)?;
3758 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
3759 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
3760 let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
3761 &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
3762 m.down_exps.qtype, m.down_exps.row_bytes)?;
3763 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
3764 } else {
3765 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
3767 let gate = e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
3768 n_embd, n_ff_exp, n_active, n_pairs, t,
3769 m.gate_exps.qtype, rbg_d)?;
3770 let up = e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
3771 n_embd, n_ff_exp, n_active, n_pairs, t,
3772 m.up_exps.qtype, rbu_d)?;
3773 let a_scr = if crate::moe_fuse_actq_on() {
3779 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
3780 } else {
3781 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
3782 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
3783 };
3784 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
3785 let pself = e.htod_i32(&pair_self)?;
3786 e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
3787 n_ff_exp, n_embd, n_active, n_pairs, n_pairs,
3788 m.down_exps.qtype, m.down_exps.row_bytes)?
3789 };
3790 let mut moe_out = e.uninit(t * n_embd)?;
3791 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
3792 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3793 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3794 {
3795 let n_ff_sh = gate_shexp.out_features();
3796 let sg_gate = e.matmul(gate_shexp, z, t)?;
3797 let sg_up = e.matmul(up_shexp, z, t)?;
3798 let mut sa = e.uninit(t * n_ff_sh)?;
3799 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3800 let sh = e.matmul(down_shexp, &sa, t)?;
3801 let g = match &m.gate_inp_shexp {
3807 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
3808 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3809 }
3810 Some(gate_inp_shexp) => {
3811 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3812 let mut g = e.uninit(t)?;
3813 e.sigmoid(&gs, &mut g, t)?;
3814 g
3815 }
3816 None => e.htod(&vec![1.0f32; t])?,
3817 };
3818 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3819 }
3820 return Ok(moe_out);
3821 }
3822
3823 let dec = std::env::var("MEMRA_MOE_DEC").map(|v| v != "0").unwrap_or(true);
3826 let matvec = |proj, exi: &_, exo: &_, exp_d: &_, pt: &_, aq: &_, ad: &_,
3827 inf, outf, qtype, rb| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3828 let dec = dec && q8_expert_dec_supported(qtype);
3830 if dec { e.moe_pairs_matvec_q8_dec(&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
3831 inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
3832 else { e.moe_pairs_matvec_q8_em (&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
3833 inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
3834 };
3835 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3836 let gate = matvec(0, &exi, &exo, &exp_d, &pt, &zq, &zd,
3837 n_embd, n_ff_exp, m.gate_exps.qtype, rbg_d)?;
3838 let up = matvec(1, &exi, &exo, &exp_d, &pt, &zq, &zd,
3839 n_embd, n_ff_exp, m.up_exps.qtype, rbu_d)?;
3840 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
3841 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
3842 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
3844 let pself = e.htod_i32(&pair_self)?;
3845 let y_down = matvec(2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
3846 n_ff_exp, n_embd, m.down_exps.qtype, m.down_exps.row_bytes)?;
3847 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
3849
3850 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3854 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3855 {
3856 let n_ff_sh = gate_shexp.out_features();
3857 let sg_gate = e.matmul(gate_shexp, z, t)?;
3858 let sg_up = e.matmul(up_shexp, z, t)?;
3859 let mut sa = e.uninit(t * n_ff_sh)?;
3860 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3861 let sh = e.matmul(down_shexp, &sa, t)?;
3862 let g = match &m.gate_inp_shexp {
3867 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
3868 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3869 }
3870 Some(gate_inp_shexp) => {
3871 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3872 let mut g = e.uninit(t)?;
3873 e.sigmoid(&gs, &mut g, t)?;
3874 g
3875 }
3876 None => e.htod(&vec![1.0f32; t])?,
3877 };
3878 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3879 }
3880 Ok(moe_out)
3881 }
3882
3883 #[allow(clippy::too_many_arguments)]
3885 #[allow(clippy::too_many_arguments)]
3886 fn moe_ffn_dev(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
3887 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, logits: &CudaSlice<f32>,
3888 t: usize, cfg: &ModelConfig, il: u16, max_block: usize)
3889 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3890 let moe = cfg.moe.as_ref().unwrap();
3891 let n_embd = cfg.n_embd as usize;
3892 let n_expert = moe.expert_count as usize;
3893 let n_used = moe.expert_used_count as usize;
3894 let n_ff_exp = moe.expert_ff_length as usize;
3895 debug_assert!(cfg.sigmoid_router().is_none(),
3899 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts");
3900 debug_assert!(!cfg.swiglu_clamped_at(il as u32),
3901 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form");
3902
3903 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
3905 if m.has_macros {
3908 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
3909 }
3910
3911 let mut moe_out = e.uninit(t * n_embd)?;
3913
3914 if let Some(dev) = m.dev_exps.as_ref() {
3917 let (rbg_d, rbu_d) = if dev.gu_il {
3920 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
3921 } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
3922 let q8 = moe_q8_enabled()
3923 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3924 && q8_expert_supported(m.down_exps.qtype);
3925 let rows_arm = q8 && t > 1 && crate::spec::spec_m2()
3934 && n_ff_exp == 512 && n_used <= 8
3935 && std::env::var("MEMRA_MOE_DEVQ8_GU").map(|v| v.is_empty() || v == "v").unwrap_or(true)
3936 && std::env::var("MEMRA_MOE_DEVQ8_DOWN").map(|v| v.is_empty() || v == "w8h2v").unwrap_or(true);
3937 let csr_mode = std::env::var("MEMRA_MOE_CSR").ok()
3946 .and_then(|v| v.parse::<i32>().ok()).unwrap_or(1);
3947 let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
3948 let csr_arm = rows_arm && csr_mode > 0 && t <= 10
3949 && csr_qt(m.gate_exps.qtype) && csr_qt(m.up_exps.qtype)
3950 && csr_qt(m.down_exps.qtype);
3951 if csr_arm {
3952 if csr_mode == 2 {
3953 static ENGAGED: std::sync::Once = std::sync::Once::new();
3954 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
3955 }
3956 let n_pairs = t * n_used;
3957 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3958 let act = e.moe_gate_up_silu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, n_pairs,
3959 n_embd, n_ff_exp, n_used, n_expert,
3960 m.gate_exps.qtype, m.up_exps.qtype,
3961 rbg_d, rbu_d)?;
3962 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
3963 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
3967 t, n_ff_exp, n_embd, n_used, n_expert,
3968 m.down_exps.qtype, m.down_exps.row_bytes)?;
3969 if csr_mode == 2 {
3970 let act_r = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
3972 n_embd, n_ff_exp, n_used, n_expert,
3973 m.gate_exps.qtype, m.up_exps.qtype,
3974 rbg_d, rbu_d, &m.dev_macros)?;
3975 let mut out_r = e.uninit(t * n_embd)?;
3976 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
3977 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2r, &ad2r, &mut out_r,
3978 t, n_ff_exp, n_embd, n_used, n_expert,
3979 m.down_exps.qtype, m.down_exps.row_bytes)?;
3980 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
3981 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
3982 let ba = a1.iter().zip(&a2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
3983 let bo = o1.iter().zip(&o2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
3984 if ba + bo > 0 {
3985 eprintln!("[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
3986 a1.len(), o1.len());
3987 let sel_h = e.dtoh_i32(&sel_d)?;
3989 let mut shown = 0;
3990 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
3991 if x.to_bits() != y.to_bits() && shown < 4 {
3992 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
3993 let ex = sel_h[p];
3994 let npx = sel_h.iter().filter(|&&v| v == ex).count();
3995 eprintln!(" ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}");
3996 shown += 1;
3997 }
3998 }
3999 std::process::exit(3);
4000 }
4001 }
4002 } else if rows_arm {
4003 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
4006 use std::sync::atomic::{AtomicU64, Ordering};
4007 static PAIRS: AtomicU64 = AtomicU64::new(0);
4008 static UNIQ: AtomicU64 = AtomicU64::new(0);
4009 static CALLS: AtomicU64 = AtomicU64::new(0);
4010 let sel_h = e.dtoh_i32(&sel_d)?;
4011 let mut u: Vec<i32> = sel_h.clone(); u.sort_unstable(); u.dedup();
4012 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
4013 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
4014 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
4015 if c % 480 == 0 {
4016 let p = PAIRS.load(Ordering::Relaxed); let q = UNIQ.load(Ordering::Relaxed);
4017 eprintln!("[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
4018 q as f64 / p as f64);
4019 }
4020 }
4021 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4022 let act = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
4023 n_embd, n_ff_exp, n_used, n_expert,
4024 m.gate_exps.qtype, m.up_exps.qtype,
4025 rbg_d, rbu_d, &m.dev_macros)?;
4026 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
4027 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
4028 t, n_ff_exp, n_embd, n_used, n_expert,
4029 m.down_exps.qtype, m.down_exps.row_bytes)?;
4030 } else {
4031 for tok in 0..t {
4032 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
4033 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
4034 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
4035 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4036 if q8 {
4037 let (zq, zd) = match (t, zq8) {
4038 (1, Some((q, d))) => (q.clone(), d.clone()),
4039 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
4040 };
4041 let act = e.moe_gate_up_silu8_dev_q8(&dev.ptr_row, &selt, &zq, &zd,
4042 n_embd, n_ff_exp, n_used, n_expert,
4043 m.gate_exps.qtype, m.up_exps.qtype,
4044 rbg_d, rbu_d, &m.dev_macros)?;
4045 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4046 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selt, &wt, &aq2, &ad2, &mut dst,
4047 n_ff_exp, n_embd, n_used, n_expert,
4048 m.down_exps.qtype, m.down_exps.row_bytes)?;
4049 } else {
4050 let act = e.moe_gate_up_silu8_dev(&dev.ptr_row, &selt, &zt, n_embd, n_ff_exp,
4051 n_used, n_expert,
4052 m.gate_exps.qtype, m.up_exps.qtype,
4053 rbg_d, rbu_d, &m.dev_macros)?;
4054 e.moe_down8_fma_dev(&dev.ptr_row, &selt, &wt, &act, &mut dst,
4055 n_ff_exp, n_embd, n_used, n_expert,
4056 m.down_exps.qtype, m.down_exps.row_bytes)?;
4057 }
4058 }
4059 }
4060 } else {
4061 let q8 = moe_q8_enabled()
4068 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
4069 && q8_expert_supported(m.down_exps.qtype);
4070 e.with_moe_cache(max_block, |c, eng| {
4071 let row = c.layer_dev_row(il, n_expert, eng)?
4072 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
4073 for tok in 0..t {
4074 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
4075 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
4076 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
4077 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4078 if q8 {
4079 let (zq, zd) = match (t, zq8) {
4080 (1, Some((q, d))) => (q.clone(), d.clone()),
4081 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
4082 };
4083 let act = eng.moe_gate_up_silu8_dev_q8(row, &selt, &zq, &zd,
4084 n_embd, n_ff_exp, n_used, n_expert,
4085 m.gate_exps.qtype, m.up_exps.qtype,
4086 m.gate_exps.row_bytes, m.up_exps.row_bytes,
4087 &m.dev_macros)?;
4088 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
4089 eng.moe_down8_fma_dev_q8(row, &selt, &wt, &aq2, &ad2, &mut dst,
4090 n_ff_exp, n_embd, n_used, n_expert,
4091 m.down_exps.qtype, m.down_exps.row_bytes)?;
4092 } else {
4093 let act = eng.moe_gate_up_silu8_dev(row, &selt, &zt, n_embd, n_ff_exp,
4094 n_used, n_expert,
4095 m.gate_exps.qtype, m.up_exps.qtype,
4096 m.gate_exps.row_bytes, m.up_exps.row_bytes,
4097 &m.dev_macros)?;
4098 eng.moe_down8_fma_dev(row, &selt, &wt, &act, &mut dst,
4099 n_ff_exp, n_embd, n_used, n_expert,
4100 m.down_exps.qtype, m.down_exps.row_bytes)?;
4101 }
4102 }
4103 c.hits += (t * 3 * n_used) as u64;
4105 Ok(())
4106 })?;
4107 }
4108
4109 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4114 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4115 {
4116 let n_ff_sh = gate_shexp.out_features();
4117 let verify_t = t > 1 && t < PRIME_MIN_T;
4120 let (sg_gate, sg_up) = if t == 1 {
4121 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
4122 Some(pair) => pair,
4123 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
4124 }
4125 } else if verify_t {
4126 let mut fused = None;
4130 if crate::spec::spec_fused_t() && (2..=4).contains(&t)
4131 && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp) {
4132 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4133 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
4134 }
4135 match fused {
4136 Some(pair) => pair,
4137 None => (e.matmul_decode_exact(gate_shexp, z, t)?,
4138 e.matmul_decode_exact(up_shexp, z, t)?),
4139 }
4140 } else {
4141 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
4142 };
4143 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
4145 let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
4146 else { e.matmul(down_shexp, &sa, t)? };
4147 let g = match &m.gate_inp_shexp {
4151 Some(gate_inp_shexp) => {
4152 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4155 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4156 } else {
4157 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4158 let mut g = e.uninit(t)?;
4159 e.sigmoid(&gs, &mut g, t)?;
4160 g
4161 }
4162 }
4163 None => e.htod(&vec![1.0f32; t])?,
4164 };
4165 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4166 }
4167
4168 Ok(moe_out)
4169 }
4170
4171 #[allow(clippy::too_many_arguments)]
4181 #[allow(clippy::too_many_arguments)]
4184 fn moe_gdec_token_q8(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
4185 zq: &CudaSlice<i8>, zd: &CudaSlice<f32>, sel: &[u32], w: &[f32],
4186 moe_out: &mut CudaSlice<f32>, tok: usize,
4187 n_embd: usize, n_ff_exp: usize, n_used: usize)
4188 -> Result<bool, Box<dyn std::error::Error>> {
4189 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4190 use cudarc::driver::DevicePtr;
4191 let ptrs = e.with_moe_cache(max_block, |c, eng| {
4192 let mut g = [0u64; 8];
4193 let mut u = [0u64; 8];
4194 let mut d = [0u64; 8];
4195 for (j, &ex) in sel.iter().enumerate() {
4196 let ex = ex as u16;
4197 let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
4198 c.resident(BlockId::new(il, PROJ_UP, ex)),
4199 c.resident(BlockId::new(il, PROJ_DOWN, ex)))
4200 else { return Ok(None); };
4201 let __s = eng.stream();
4202 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
4203 let (pu, _e1) = c.slot(su).device_ptr(&__s);
4204 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
4205 g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
4206 }
4207 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
4208 for &ex in sel {
4209 let ex = ex as u16;
4210 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
4211 c.note_profile_hit(BlockId::new(il, proj, ex));
4212 }
4213 }
4214 }
4215 c.hits += (3 * n_used) as u64;
4216 Ok(Some((g, u, d)))
4217 })?;
4218 let Some((g, u, d)) = ptrs else { return Ok(false) };
4219 let mut wv = [0f32; 8];
4220 wv[..n_used].copy_from_slice(w);
4221 let act = e.moe_gate_up_silu8_q8(crate::WPtr8(g), crate::WPtr8(u), zq, zd,
4222 n_embd, n_ff_exp, n_used,
4223 m.gate_exps.qtype, m.up_exps.qtype,
4224 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
4225 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4227 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4228 e.moe_down8_fma_q8(crate::WPtr8(d), crate::F32x8(wv), &aq2, &ad2, &mut dst,
4229 n_ff_exp, n_embd, n_used,
4230 m.down_exps.qtype, m.down_exps.row_bytes)?;
4231 Ok(true)
4232 }
4233
4234 fn moe_gdec_token(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
4235 zt: &cudarc::driver::CudaView<f32>, sel: &[u32], w: &[f32],
4236 moe_out: &mut CudaSlice<f32>, tok: usize,
4237 n_embd: usize, n_ff_exp: usize, n_used: usize)
4238 -> Result<bool, Box<dyn std::error::Error>> {
4239 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4240 use cudarc::driver::DevicePtr;
4241 let ptrs = e.with_moe_cache(max_block, |c, eng| {
4243 let mut g = [0u64; 8];
4244 let mut u = [0u64; 8];
4245 let mut d = [0u64; 8];
4246 for (j, &ex) in sel.iter().enumerate() {
4247 let ex = ex as u16;
4248 let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
4249 c.resident(BlockId::new(il, PROJ_UP, ex)),
4250 c.resident(BlockId::new(il, PROJ_DOWN, ex)))
4251 else { return Ok(None); };
4252 let __s = eng.stream();
4253 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
4254 let (pu, _e1) = c.slot(su).device_ptr(&__s);
4255 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
4256 g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
4257 }
4258 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
4259 for &ex in sel {
4260 let ex = ex as u16;
4261 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
4262 c.note_profile_hit(BlockId::new(il, proj, ex));
4263 }
4264 }
4265 }
4266 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
4268 })?;
4269 let Some((g, u, d)) = ptrs else { return Ok(false) };
4270 let mut wv = [0f32; 8];
4271 wv[..n_used].copy_from_slice(w);
4272 let act = e.moe_gate_up_silu8(crate::WPtr8(g), crate::WPtr8(u), zt,
4274 n_embd, n_ff_exp, n_used,
4275 m.gate_exps.qtype, m.up_exps.qtype,
4276 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
4277 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4278 e.moe_down8_fma_into(crate::WPtr8(d), crate::F32x8(wv), &act, &mut dst,
4279 n_ff_exp, n_embd, n_used,
4280 m.down_exps.qtype, m.down_exps.row_bytes)?;
4281 Ok(true)
4282 }
4283
4284 fn moe_cached_gemm_q8(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
4289 max_block: usize, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
4290 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4291 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
4292 let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
4293 let layout = exps.expert_layout(ex);
4294 let id = BlockId::new(il, proj, ex as u16);
4295 let source = exps.expert_source(ex);
4296 e.with_moe_cache(max_block, |c, eng| {
4297 let slot = c.dispatch_source(id, source, eng)?;
4298 let DispatchSlot::Resident(sl) = slot;
4299 let buf = c.slot(sl);
4300 eng.qmatvec_expert_q8(buf, 0..layout.len, aq, ad, 1, exps.in_f, exps.out_f,
4301 layout.qtype, layout.row_bytes)
4302 })
4303 }
4304
4305 fn moe_cached_gemm(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
4306 max_block: usize, x: &cudarc::driver::CudaView<f32>)
4307 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4308 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
4309 let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
4310 let layout = exps.expert_layout(ex);
4311 let id = BlockId::new(il, proj, ex as u16);
4312 let source = exps.expert_source(ex);
4313 e.with_moe_cache(max_block, |c, eng| {
4315 let slot = c.dispatch_source(id, source, eng)?;
4316 let DispatchSlot::Resident(sl) = slot;
4319 let buf = c.slot(sl);
4320 eng.qmatvec_view(buf, 0..layout.len, x, 1, exps.in_f, exps.out_f,
4321 layout.qtype, layout.row_bytes)
4322 })
4323 }
4324
4325 fn moe_profile_admit_expert(
4329 e: &Engine,
4330 il: u16,
4331 ex: usize,
4332 m: &MoeWeights,
4333 max_block: usize,
4334 ) -> Result<(), Box<dyn std::error::Error>> {
4335 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4336 e.with_moe_cache(max_block, |cache, eng| {
4337 for (proj, exps) in [
4338 (PROJ_GATE, &m.gate_exps),
4339 (PROJ_UP, &m.up_exps),
4340 (PROJ_DOWN, &m.down_exps),
4341 ] {
4342 let id = BlockId::new(il, proj, ex as u16);
4343 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
4344 }
4345 Ok(())
4346 })
4347 }
4348
4349 #[allow(clippy::too_many_arguments)]
4352 fn moe_frozen_gemm(
4353 e: &Engine,
4354 il: u16,
4355 proj: u8,
4356 ex: usize,
4357 m: &MoeWeights,
4358 max_block: usize,
4359 x: &cudarc::driver::CudaView<f32>,
4360 scratch: &mut Option<CudaSlice<u8>>,
4361 scratch_len: usize,
4362 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4363 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
4364 let exps = match proj {
4365 PROJ_GATE => &m.gate_exps,
4366 PROJ_UP => &m.up_exps,
4367 _ => &m.down_exps,
4368 };
4369 let layout = exps.expert_layout(ex);
4370 let id = BlockId::new(il, proj, ex as u16);
4371 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
4372 let Some(slot) = cache.resident(id) else {
4373 return Ok(None);
4374 };
4375 let buf = cache.slot(slot);
4376 Ok(Some(eng.qmatvec_view(
4377 buf,
4378 0..layout.len,
4379 x,
4380 1,
4381 exps.in_f,
4382 exps.out_f,
4383 layout.qtype,
4384 layout.row_bytes,
4385 )?))
4386 })? {
4387 return Ok(output);
4388 }
4389 if scratch.is_none() {
4390 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
4391 }
4392 let scratch = scratch.as_mut().unwrap();
4393 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
4394 e.qmatvec_view(
4395 scratch,
4396 0..layout.len,
4397 x,
4398 1,
4399 exps.in_f,
4400 exps.out_f,
4401 layout.qtype,
4402 layout.row_bytes,
4403 )
4404 }
4405
4406 fn moe_prefetch_expert(
4407 e: &Engine,
4408 il: u16,
4409 ex: usize,
4410 m: &MoeWeights,
4411 max_block: usize,
4412 keep: &[crate::moe_cache::BlockId],
4413 ) -> Result<(), Box<dyn std::error::Error>> {
4414 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4415 e.with_moe_cache(max_block, |c, eng| {
4416 for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
4417 (PROJ_DOWN, &m.down_exps)] {
4418 let id = BlockId::new(il, proj, ex as u16);
4419 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
4420 }
4421 Ok(())
4422 })
4423 }
4424
4425 fn moe_prefetch_disk_expert(e: &Engine, il: u16, ex: usize, m: &MoeWeights,
4428 max_block: usize, keep: &[crate::moe_cache::BlockId])
4429 -> Result<(), Box<dyn std::error::Error>> {
4430 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4431 e.with_moe_cache(max_block, |c, eng| {
4432 for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
4433 (PROJ_DOWN, &m.down_exps)] {
4434 let source = exps.expert_source(ex);
4435 if let crate::model::ExpertSource::Disk { .. } = &source {
4436 let id = BlockId::new(il, proj, ex as u16);
4437 let _ = c.prefetch_source(id, source, keep, eng)?;
4438 }
4439 }
4440 Ok(())
4441 })
4442 }
4443
4444 #[inline]
4445 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
4446 let _ = m.gate_exps.prefetch_expert_pages(ex);
4447 let _ = m.up_exps.prefetch_expert_pages(ex);
4448 let _ = m.down_exps.prefetch_expert_pages(ex);
4449 }
4450}
4451
4452impl HybridModel {
4469 pub(crate) fn moe_ffn_grouped(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
4472 cfg: &ModelConfig, il: u16, _max_block: usize)
4473 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4474 let moe = cfg.moe.as_ref().unwrap();
4475 let n_embd = cfg.n_embd as usize;
4476 let n_expert = moe.expert_count as usize;
4477 let n_used = moe.expert_used_count as usize;
4478 let n_ff_exp = moe.expert_ff_length as usize;
4479 let lim_exp = cfg.clamp_exp_at(il as u32);
4481 let lim_shexp = cfg.clamp_shexp_at(il as u32);
4482
4483 let logits = e.matmul(&m.gate_inp, z, t)?;
4485 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
4486 Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
4487 m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
4488 } else {
4489 Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
4490 None, None, m.active_experts.as_deref())?
4491 };
4492 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4493
4494 struct ExpertGroup {
4498 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
4502 let mut groups: Vec<ExpertGroup> = (0..n_expert).map(|_| ExpertGroup {
4503 tok_indices: Vec::new(), slot_indices: Vec::new(), weights: Vec::new(),
4504 }).collect();
4505
4506 for tok in 0..t {
4507 for j in 0..n_used {
4508 let ex = sel_all[tok * n_used + j] as usize;
4509 let w = w_all[tok * n_used + j];
4510 groups[ex].tok_indices.push(tok as i32);
4511 groups[ex].slot_indices.push(j as i32);
4512 groups[ex].weights.push(w);
4513 }
4514 }
4515
4516 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
4519 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
4523 let u_len = m.up_exps.max_expert_bytes();
4524 let d_len = m.down_exps.max_expert_bytes();
4525 let use_cache = Engine::moe_cache_enabled();
4526 let max_block = _max_block;
4527
4528 let (mut scratch_g, mut scratch_u, mut scratch_d) = if !use_cache {
4530 (Some(e.alloc_u8(g_len)?), Some(e.alloc_u8(u_len)?), Some(e.alloc_u8(d_len)?))
4531 } else {
4532 (None, None, None)
4533 };
4534
4535 let mut order: Vec<usize> =
4546 (0..n_expert).filter(|&ex| !groups[ex].tok_indices.is_empty()).collect();
4547 order.sort_by(|&a, &b| groups[b].tok_indices.len()
4548 .cmp(&groups[a].tok_indices.len()).then(a.cmp(&b)));
4549 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
4551 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
4552 if worker_disk_prefetch {
4553 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
4554 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
4555 }
4556 }
4557 for (order_pos, &ex) in order.iter().enumerate() {
4558 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
4559 Self::moe_prefetch_host_expert(order[next], m);
4560 }
4561 if worker_disk_prefetch {
4562 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
4563 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4564 let keep = [
4565 BlockId::new(il, PROJ_GATE, ex as u16),
4566 BlockId::new(il, PROJ_UP, ex as u16),
4567 BlockId::new(il, PROJ_DOWN, ex as u16),
4568 ];
4569 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
4570 }
4571 }
4572 let grp = &groups[ex];
4573 let m_e = grp.tok_indices.len();
4574 m_dist.push(m_e);
4575 let gl = m.gate_exps.expert_layout(ex);
4576 let ul = m.up_exps.expert_layout(ex);
4577 let dl = m.down_exps.expert_layout(ex);
4578
4579 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
4583 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
4584 let dmac = m.down_exps.macro_scale(ex);
4585 let weight_d = if dmac == 1.0 { e.htod(&grp.weights)? } else {
4586 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
4587 e.htod(&scaled)?
4588 };
4589
4590 let mut gathered = e.zeros(m_e * n_embd)?;
4592 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
4593 let gv = gathered.slice(0..m_e * n_embd);
4594
4595 let y = if use_cache {
4597 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4598 let gate = e.with_moe_cache(max_block, |c, eng| {
4600 let id = BlockId::new(il, PROJ_GATE, ex as u16);
4601 let slot = c.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
4602 let buf = c.buf(slot);
4603 eng.qmatvec_view(buf, 0..gl.len, &gv, m_e,
4604 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
4605 })?;
4606 let up = e.with_moe_cache(max_block, |c, eng| {
4607 let id = BlockId::new(il, PROJ_UP, ex as u16);
4608 let slot = c.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
4609 let buf = c.buf(slot);
4610 eng.qmatvec_view(buf, 0..ul.len, &gv, m_e,
4611 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
4612 })?;
4613 let mut act = e.zeros(m_e * n_ff_exp)?;
4615 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4616 m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
4617 let actv = act.slice(0..m_e * n_ff_exp);
4618 e.with_moe_cache(max_block, |c, eng| {
4619 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
4620 let slot = c.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
4621 let buf = c.buf(slot);
4622 eng.qmatvec_view(buf, 0..dl.len, &actv, m_e,
4623 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
4624 })?
4625 } else {
4626 let sg = scratch_g.as_mut().unwrap();
4628 let su = scratch_u.as_mut().unwrap();
4629 let sd = scratch_d.as_mut().unwrap();
4630 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
4631 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
4632 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
4633 let gate = e.qmatvec_view(sg, 0..gl.len, &gv, m_e,
4634 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
4635 let up = e.qmatvec_view(su, 0..ul.len, &gv, m_e,
4636 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
4637 let mut act = e.zeros(m_e * n_ff_exp)?;
4639 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4640 m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
4641 let actv = act.slice(0..m_e * n_ff_exp);
4642 e.qmatvec_view(sd, 0..dl.len, &actv, m_e,
4643 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?
4644 };
4645
4646 e.scatter_slot(&y, &tok_idx_d, &slot_idx_d, &weight_d,
4648 &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
4649 }
4650
4651 let mut moe_out = e.zeros(t * n_embd)?;
4653 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
4654
4655 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
4657 m_dist.sort_unstable();
4658 let active = m_dist.len();
4659 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
4660 let median = m_dist[active / 2];
4661 let max_m = *m_dist.last().unwrap();
4662 let min_m = m_dist[0];
4663 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
4664 println!("moe-grouped il={il} t={t} active={active}/{n_expert} \
4665 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
4666 above_gemm_threshold(>=16)={above16}/{active}");
4667 }
4668
4669 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4673 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4674 {
4675 let n_ff_sh = gate_shexp.out_features();
4676 let sg_gate = e.matmul(gate_shexp, z, t)?;
4677 let sg_up = e.matmul(up_shexp, z, t)?;
4678 let mut sa = e.zeros(t * n_ff_sh)?;
4679 Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp, &mut sa, t * n_ff_sh)?;
4680 let sh = e.matmul(down_shexp, &sa, t)?;
4681 let g = match &m.gate_inp_shexp {
4685 Some(gate_inp_shexp) => {
4686 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4689 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4690 } else {
4691 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4692 let mut g = e.uninit(t)?;
4693 e.sigmoid(&gs, &mut g, t)?;
4694 g
4695 }
4696 }
4697 None => e.htod(&vec![1.0f32; t])?,
4698 };
4699 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4700 }
4701
4702 Ok(moe_out)
4703 }
4704
4705 pub(crate) fn moe_ffn_lockstep(
4712 &self,
4713 e: &Engine,
4714 m: &MoeWeights,
4715 zbatch: &CudaSlice<f32>,
4716 mrows: usize,
4717 il: u16,
4718 max_block: usize,
4719 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4720 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4721 let cfg = &self.cfg;
4722 let moe = cfg.moe.as_ref().unwrap();
4723 let n_embd = cfg.n_embd as usize;
4724 let n_expert = moe.expert_count as usize;
4725 let n_used = moe.expert_used_count as usize;
4726 let n_ff_exp = moe.expert_ff_length as usize;
4727 let lim_exp = cfg.clamp_exp_at(il as u32);
4729 let lim_shexp = cfg.clamp_shexp_at(il as u32);
4730
4731 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
4732 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
4733 Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
4734 m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
4735 } else {
4736 Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
4737 None, None, m.active_experts.as_deref())?
4738 };
4739 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
4740
4741 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
4743 Ok((0..n_expert)
4744 .map(|ex| {
4745 [PROJ_GATE, PROJ_UP, PROJ_DOWN].into_iter().all(|p| {
4746 c.resident(BlockId::new(il, p, ex as u16)).is_some()
4747 })
4748 })
4749 .collect())
4750 })?;
4751
4752 struct Group {
4753 rows: Vec<i32>,
4754 slots: Vec<i32>,
4755 weights: Vec<f32>,
4756 }
4757 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
4758 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
4759 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
4760 Default::default();
4761 for row in 0..mrows {
4762 for j in 0..n_used {
4763 let ex = sel_all[row * n_used + j] as usize;
4764 let w = w_all[row * n_used + j];
4765 if resident_expert[ex] {
4766 let group = groups.entry(ex).or_insert_with(|| Group {
4767 rows: Vec::new(),
4768 slots: Vec::new(),
4769 weights: Vec::new(),
4770 });
4771 group.rows.push(row as i32);
4772 group.slots.push(j as i32);
4773 group.weights.push(w);
4774 } else {
4775 crate::cpu_experts::record_incomplete_gpu_residency(0);
4776 cpu_rows[row].push((ex, w));
4777 cpu_by_expert.entry(ex).or_default().push((row, w));
4778 }
4779 }
4780 }
4781
4782 let host_rows = e.dtoh(zbatch)?;
4788 let rows_ok = crate::cpu_experts::rows_supported();
4789 enum CpuPart {
4790 Single { row: usize },
4791 Rows { rows: Vec<usize> },
4792 }
4793 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
4794 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
4795 if rows_ok {
4796 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
4797 .into_iter()
4798 .filter(|(_, rows)| rows.len() >= 2)
4799 .collect();
4800 shared.sort_by_key(|(ex, _)| *ex);
4801 for (ex, mut row_weights) in shared {
4802 row_weights.sort_by_key(|(row, _)| *row);
4803 let inputs: Vec<(&[f32], f32)> = row_weights
4804 .iter()
4805 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
4806 .collect();
4807 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
4808 .map_err(std::io::Error::other)?;
4809 for &(row, _) in &row_weights {
4810 rows_served.insert((row, ex));
4811 }
4812 tickets.push((
4813 CpuPart::Rows {
4814 rows: row_weights.iter().map(|&(row, _)| row).collect(),
4815 },
4816 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
4817 ));
4818 }
4819 }
4820 for (row, selected) in cpu_rows.iter().enumerate() {
4821 let leftover: Vec<(usize, f32)> = selected
4822 .iter()
4823 .copied()
4824 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
4825 .collect();
4826 if leftover.is_empty() {
4827 continue;
4828 }
4829 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
4830 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
4831 .map_err(std::io::Error::other)?;
4832 tickets.push((
4833 CpuPart::Single { row },
4834 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
4835 ));
4836 }
4837
4838 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
4839 let mut wbuf = e.zeros(mrows * n_used)?;
4840 let mut order: Vec<usize> = groups.keys().copied().collect();
4841 order.sort_by(|&a, &b| {
4842 groups[&b].rows.len().cmp(&groups[&a].rows.len()).then(a.cmp(&b))
4843 });
4844 for &ex in &order {
4845 let group = &groups[&ex];
4846 let m_e = group.rows.len();
4847 let gl = m.gate_exps.expert_layout(ex);
4848 let ul = m.up_exps.expert_layout(ex);
4849 let dl = m.down_exps.expert_layout(ex);
4850 let row_idx_d = e.htod_i32(&group.rows)?;
4851 let slot_idx_d = e.htod_i32(&group.slots)?;
4852 let dmac = m.down_exps.macro_scale(ex);
4853 let weight_d = if dmac == 1.0 {
4854 e.htod(&group.weights)?
4855 } else {
4856 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
4857 e.htod(&scaled)?
4858 };
4859 let mut gathered = e.zeros(m_e * n_embd)?;
4860 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
4861 let gv = gathered.slice(0..m_e * n_embd);
4862 let gate = e.with_moe_cache(max_block, |c, eng| {
4863 let slot = c
4864 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
4865 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4866 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..gl.len, &gv, m_e,
4867 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
4868 })?;
4869 let up = e.with_moe_cache(max_block, |c, eng| {
4870 let slot = c
4871 .resident(BlockId::new(il, PROJ_UP, ex as u16))
4872 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4873 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..ul.len, &gv, m_e,
4874 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
4875 })?;
4876 let mut act = e.zeros(m_e * n_ff_exp)?;
4877 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4878 m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
4879 let actv = act.slice(0..m_e * n_ff_exp);
4880 let y = e.with_moe_cache(max_block, |c, eng| {
4881 let slot = c
4882 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
4883 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4884 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..dl.len, &actv, m_e,
4885 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
4886 })?;
4887 e.scatter_slot(&y, &row_idx_d, &slot_idx_d, &weight_d,
4888 &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
4889 }
4890 let mut moe_out = e.zeros(mrows * n_embd)?;
4891 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
4892
4893 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
4895 for (part, ticket) in tickets {
4896 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
4897 let mut add_row = |row: usize, chunk: &[f32]| {
4898 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
4899 for (accumulator, value) in sum.iter_mut().zip(chunk) {
4900 *accumulator += value;
4901 }
4902 };
4903 match part {
4904 CpuPart::Single { row } => add_row(row, &cpu_output),
4905 CpuPart::Rows { rows } => {
4906 for (slot, row) in rows.into_iter().enumerate() {
4907 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
4908 }
4909 }
4910 }
4911 }
4912 for (row, sum) in row_sums.into_iter().enumerate() {
4913 let Some(sum) = sum else { continue };
4914 let cpu_output = e.htod(&sum)?;
4915 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
4916 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
4917 }
4918
4919 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4920 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4921 {
4922 let n_ff_sh = gate_shexp.out_features();
4923 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
4924 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
4925 let mut sa = e.zeros(mrows * n_ff_sh)?;
4926 Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp,
4927 &mut sa, mrows * n_ff_sh)?;
4928 let sh = e.matmul(down_shexp, &sa, mrows)?;
4929 let g = match &m.gate_inp_shexp {
4932 Some(gate_inp_shexp) => {
4933 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
4934 }
4935 None => e.htod(&vec![1.0f32; mrows])?,
4936 };
4937 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
4938 }
4939
4940 Ok(moe_out)
4941 }
4942}
4943
4944impl HybridModel {
4950 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
4952 let g = self.cfg.gemma4.as_ref().unwrap();
4953 let swa = g.swa_pattern[il];
4954 let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
4955 (hd, g.head_count_kv[il] as usize, self.cfg.n_head as usize,
4959 if swa { g.rope_base_swa } else { g.rope_base_global },
4960 1.0, swa)
4961 }
4962
4963 fn gemma4_suppress(&self, e: &Engine, ld: &mut CudaSlice<f32>, t: usize)
4967 -> Result<(), Box<dyn std::error::Error>> {
4968 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
4969 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
4970 }
4971 Ok(())
4972 }
4973
4974 fn gemma4_attn_prime(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
4979 h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
4980 cache: Option<&mut Cache>)
4981 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4982 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
4983 let eps = self.cfg.rms_eps;
4984 let aux = self.gemma4_aux.as_ref().unwrap();
4985
4986 e.mmq_act_begin();
4989 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)? };
4994
4995 let mut q = e.uninit(t * nh * hd)?;
4996 let mut k = e.uninit(t * nkv * hd)?;
4997 let mut v = e.uninit(t * nkv * hd)?;
4999 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5003 let emit = t >= 16 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
5004 && *EMIT.get_or_init(|| std::env::var("MEMRA_FA_EMIT").map(|s| s != "0").unwrap_or(true));
5005 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
5006 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
5007 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
5008 let v_f16 = emit && crate::fa_f16pv_on() && match hd {
5011 512 => true,
5012 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
5013 _ => false,
5014 };
5015 if emit {
5016 e.rms_norm_qkv_w4b(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5017 &aux.ones, &mut q, &mut k, &mut v, &mut vb,
5018 hd, nh * t, nkv * t, eps, v_f16)?;
5019 } else {
5020 e.rms_norm_qkv(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5021 &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t, eps)?;
5022 }
5023
5024 let ff = if swa { None } else {
5025 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5026 };
5027 if emit {
5028 e.rope_neox2_bf16e(&mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t,
5029 base, 1.0, ff)?;
5030 } else {
5031 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
5032 }
5033
5034 if let Some(cache) = cache {
5035 let kvl = cache.kv[il].as_mut().unwrap();
5036 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
5037 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
5038 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()))?;
5039 kvl.len += t;
5040 }
5041 let mut attn = e.zeros(t * nh * hd)?;
5042 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5046 if swa && t > win {
5047 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
5048 if emit { e.fa_prefill_w_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
5049 scale, true, win, v_f16)?; }
5050 else { e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true,
5051 win)?; }
5052 } else {
5053 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
5054 }
5055 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
5056 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
5057 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
5058 if emit { e.fa_prefill_hd512_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
5059 scale, true, v_f16)?; }
5060 else { e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?; }
5061 } else {
5062 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
5063 }
5064 Ok(e.matmul(&fa.wo, &attn, t)?)
5065 }
5066
5067 fn gemma4_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
5069 h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
5070 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5071 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
5072 }
5073
5074 fn gemma4_moe_q8(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
5079 bits: &crate::hybrid::Gemma4MoeBits,
5080 mq: &(CudaSlice<i8>, CudaSlice<f32>),
5081 router_in: &CudaSlice<f32>, t: usize)
5082 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5083 let cfg = &self.cfg;
5084 let moe = cfg.moe.as_ref().unwrap();
5085 let n_embd = cfg.n_embd as usize;
5086 let n_expert = moe.expert_count as usize;
5087 let n_used = moe.expert_used_count as usize;
5088 let n_ff_exp = moe.expert_ff_length as usize;
5089 let logits = if crate::router_kernel_on() {
5093 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
5094 } else {
5095 e.matmul(&m.gate_inp, router_in, t)?
5096 };
5097 let dev = m.dev_exps.as_ref().unwrap();
5098 let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
5099 &bits.per_expert_scale_d)?;
5100 let (zq, zd) = mq;
5101 if t == 1 {
5102 let selv = sel_d.slice(0..n_used);
5103 let wv = w_d.slice(0..n_used);
5104 let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, zq, zd,
5105 n_embd, n_ff_exp, n_used, n_expert,
5106 m.gate_exps.qtype, m.up_exps.qtype,
5107 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5108 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5109 let mut moe_out = e.uninit(n_embd)?;
5110 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
5111 &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
5112 n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
5113 return Ok(moe_out);
5114 }
5115 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
5116 let act = if csr {
5117 e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, zq, zd, t * n_used,
5118 n_embd, n_ff_exp, n_used, n_expert,
5119 m.gate_exps.qtype, m.up_exps.qtype,
5120 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5121 } else {
5122 e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, zq, zd, t,
5123 n_embd, n_ff_exp, n_used, n_expert,
5124 m.gate_exps.qtype, m.up_exps.qtype,
5125 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5126 };
5127 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
5128 let mut moe_out = e.uninit(t * n_embd)?;
5129 e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
5132 n_ff_exp, n_embd, n_used, n_expert,
5133 m.down_exps.qtype, m.down_exps.row_bytes)?;
5134 Ok(moe_out)
5135 }
5136
5137 fn gemma4_moe(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
5141 bits: &crate::hybrid::Gemma4MoeBits, moe_in: &CudaSlice<f32>,
5142 router_in: &CudaSlice<f32>, t: usize)
5143 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5144 let cfg = &self.cfg;
5145 let moe = cfg.moe.as_ref().unwrap();
5146 let n_embd = cfg.n_embd as usize;
5147 let n_expert = moe.expert_count as usize;
5148 let n_used = moe.expert_used_count as usize;
5149 let n_ff_exp = moe.expert_ff_length as usize;
5150
5151 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
5155 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
5156 } else {
5157 e.matmul(&m.gate_inp, router_in, t)?
5158 };
5159
5160 if t < PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5165 && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
5166 && expert_dp4a_supported(m.down_exps.qtype)
5167 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0") {
5168 let dev = m.dev_exps.as_ref().unwrap();
5169 let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
5170 &bits.per_expert_scale_d)?;
5171 if t == 1 {
5172 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
5173 let selv = sel_d.slice(0..n_used);
5174 let wv = w_d.slice(0..n_used);
5175 let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, &zq, &zd,
5176 n_embd, n_ff_exp, n_used, n_expert,
5177 m.gate_exps.qtype, m.up_exps.qtype,
5178 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5179 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5180 let mut moe_out = e.uninit(n_embd)?;
5181 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
5182 &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
5183 n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
5184 return Ok(moe_out);
5185 }
5186 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
5191 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
5192 let act = if csr {
5193 e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, t * n_used,
5194 n_embd, n_ff_exp, n_used, n_expert,
5195 m.gate_exps.qtype, m.up_exps.qtype,
5196 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5197 } else {
5198 e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
5199 n_embd, n_ff_exp, n_used, n_expert,
5200 m.gate_exps.qtype, m.up_exps.qtype,
5201 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5202 };
5203 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
5204 let mut moe_out = e.uninit(t * n_embd)?;
5205 e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
5206 n_ff_exp, n_embd, n_used, n_expert,
5207 m.down_exps.qtype, m.down_exps.row_bytes)?;
5208 return Ok(moe_out);
5209 }
5210
5211 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
5212 for (i, &sx) in sel_all.iter().enumerate() {
5213 w_all[i] *= bits.per_expert_scale[sx as usize];
5214 }
5215
5216 if t >= PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5220 && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
5221 && expert_dp4a_supported(m.down_exps.qtype)
5222 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0") {
5223 let dev = m.dev_exps.as_ref().unwrap();
5224 let n_pairs = t * n_used;
5225 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
5226 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
5227 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5228 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5229 let pt = e.htod_i32(&pair_tok)?;
5230 let pw = e.htod(&w_all)?;
5231 let toff = e.htod_i32(&tok_off)?;
5232 let tids = e.htod_i32(&tok_ids)?;
5233 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
5234 for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
5235 let mut ex_ids: Vec<i32> = Vec::new();
5236 let mut ex_off: Vec<i32> = vec![0];
5237 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
5238 for (ex, list) in by_ex.iter().enumerate() {
5239 if list.is_empty() { continue; }
5240 ex_ids.push(ex as i32);
5241 ex_pairs.extend_from_slice(list);
5242 ex_off.push(ex_pairs.len() as i32);
5243 }
5244 let n_active = ex_ids.len();
5245 let exi = e.htod_i32(&ex_ids)?;
5246 let exo = e.htod_i32(&ex_off)?;
5247 let exp_d = e.htod_i32(&ex_pairs)?;
5248 if crate::moe_f16g_gemma_on()
5256 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
5257 && f16g_proj_ok(m.up_exps.qtype, n_embd)
5258 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp) {
5259 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
5260 let csr_tok_d = e.htod_i32(&csr_tok)?;
5261 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
5262 let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
5263 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
5264 m.gate_exps.qtype, m.gate_exps.row_bytes)?;
5265 let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
5266 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
5267 m.up_exps.qtype, m.up_exps.row_bytes)?;
5268 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
5269 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
5270 let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
5271 &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
5272 m.down_exps.qtype, m.down_exps.row_bytes)?;
5273 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
5274 let mut moe_out = e.uninit(t * n_embd)?;
5275 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5276 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
5277 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
5278 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
5279 eprintln!("[f16g-debug] post-permute bad={} post-scatter bad={}",
5280 scan(&yd), scan(&mo));
5281 }
5282 return Ok(moe_out);
5283 }
5284 let mma = n_embd % 256 == 0
5287 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
5288 let (gate, up) = if mma {
5289 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
5290 (e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
5291 n_embd, n_ff_exp, n_active, n_pairs, t,
5292 m.gate_exps.qtype, m.gate_exps.row_bytes)?,
5293 e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
5294 n_embd, n_ff_exp, n_active, n_pairs, t,
5295 m.up_exps.qtype, m.up_exps.row_bytes)?)
5296 } else {
5297 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
5298 (e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 0, &exi, &exo, &exp_d, &pt, &zq, &zd,
5299 n_embd, n_ff_exp, n_expert, n_active, n_pairs,
5300 m.gate_exps.qtype, m.gate_exps.row_bytes)?,
5301 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 1, &exi, &exo, &exp_d, &pt, &zq, &zd,
5302 n_embd, n_ff_exp, n_expert, n_active, n_pairs,
5303 m.up_exps.qtype, m.up_exps.row_bytes)?)
5304 };
5305 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5306 let pself = e.htod_i32(&pair_self)?;
5307 let y_down = if mma {
5319 let in_pad = n_ff_exp.div_ceil(256) * 256;
5320 let a_scr = if crate::moe_fuse_actq_on() {
5321 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
5322 } else {
5323 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5324 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
5325 };
5326 e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
5327 in_pad, n_embd, n_active, n_pairs, n_pairs,
5328 m.down_exps.qtype, m.down_exps.row_bytes)?
5329 } else {
5330 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5331 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5332 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
5333 n_ff_exp, n_embd, n_expert, n_active, n_pairs,
5334 m.down_exps.qtype, m.down_exps.row_bytes)?
5335 };
5336 let mut moe_out = e.uninit(t * n_embd)?;
5337 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5338 return Ok(moe_out);
5339 }
5340
5341 let g_len = m.gate_exps.expert_stride;
5342 let u_len = m.up_exps.expert_stride;
5343 let d_len = m.down_exps.expert_stride;
5344 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
5348 let (mut sg, mut su, mut sd) = if dev.is_some() { (None, None, None) } else {
5349 (Some(e.alloc_u8_uninit(g_len)?), Some(e.alloc_u8_uninit(u_len)?), Some(e.alloc_u8_uninit(d_len)?))
5350 };
5351 let mut moe_out = e.zeros(t * n_embd)?;
5352 for tok in 0..t {
5353 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
5354 let w = &w_all[tok * n_used..(tok + 1) * n_used];
5355 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
5356 for (j, &ex) in sel.iter().enumerate() {
5357 let ex = ex as usize;
5358 let gate = match dev {
5359 Some(d) => e.qmatvec_view(&d.gate, ex * g_len..(ex + 1) * g_len, &zt, 1,
5360 m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?,
5361 None => {
5362 let sg = sg.as_mut().unwrap();
5363 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5364 e.qmatvec_view(sg, 0..g_len, &zt, 1,
5365 m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?
5366 }
5367 };
5368 let up = match dev {
5369 Some(d) => e.qmatvec_view(&d.up, ex * u_len..(ex + 1) * u_len, &zt, 1,
5370 m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?,
5371 None => {
5372 let su = su.as_mut().unwrap();
5373 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5374 e.qmatvec_view(su, 0..u_len, &zt, 1,
5375 m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?
5376 }
5377 };
5378 let mut act = e.uninit(n_ff_exp)?;
5379 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
5380 let actv = act.slice(0..n_ff_exp);
5381 let y = match dev {
5382 Some(d) => e.qmatvec_view(&d.down, ex * d_len..(ex + 1) * d_len, &actv, 1,
5383 m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?,
5384 None => {
5385 let sd = sd.as_mut().unwrap();
5386 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5387 e.qmatvec_view(sd, 0..d_len, &actv, 1,
5388 m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?
5389 }
5390 };
5391 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5392 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
5393 }
5394 }
5395 Ok(moe_out)
5396 }
5397
5398 fn gemma4_layer(&self, e: &Engine, il: usize, layer: &crate::hybrid::HybridLayer,
5400 x: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
5401 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5402 let n_embd = self.cfg.n_embd as usize;
5403 let eps = self.cfg.rms_eps;
5404
5405 let mut h = e.zeros(t * n_embd)?;
5406 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5407 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5408 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
5409 let mut cur = e.zeros(t * n_embd)?;
5411 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
5412 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
5413 }
5414
5415 fn gemma4_layer_tail_add(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5419 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
5420 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5421 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
5422 }
5423
5424 fn gemma4_layer_tail_add_n(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5427 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5428 next_norm: Option<&CudaSlice<f32>>)
5429 -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
5430 let n_embd = self.cfg.n_embd as usize;
5431 let bits = layer.gemma4.as_ref().unwrap();
5432 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
5433 let mut xn = e.uninit(t * n_embd)?;
5434 match next_norm {
5435 Some(w) => {
5436 let mut hn = e.uninit(t * n_embd)?;
5437 e.add_scale_rms_norm(&sn, &attn_out, bits.layer_scale, w, &mut xn, &mut hn,
5438 n_embd, t, self.cfg.rms_eps)?;
5439 Ok((xn, Some(hn)))
5440 }
5441 None => {
5442 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
5443 Ok((xn, None))
5444 }
5445 }
5446 }
5447
5448 fn gemma4_layer_tail_core(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5451 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
5452 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5453 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
5454 }
5455
5456 fn gemma4_layer_tail_core_pn(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5463 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5464 pre_norm: Option<&CudaSlice<f32>>, defer_post_norm: bool)
5465 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5466 let n_embd = self.cfg.n_embd as usize;
5467 let eps = self.cfg.rms_eps;
5468 let bits = layer.gemma4.as_ref().unwrap();
5469
5470 let Some(mbits) = bits.moe_bits.as_ref() else {
5473 let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
5474 else { panic!("gemma4 dense layer without Dense ffn") };
5475 let mut attn_out = e.uninit(t * n_embd)?;
5476 let mut zsh = e.uninit(t * n_embd)?;
5477 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5480 match pre_norm {
5481 Some(wa) if t == 1 => {
5482 zpair = Some(e.rms_pre_add_rms_norm_q8z(cur, wa, x,
5483 bits.ffn_norm.float_data(),
5484 &mut attn_out, &mut zsh,
5485 n_embd, t, eps)?);
5486 }
5487 Some(wa) => e.rms_pre_add_rms_norm(cur, wa, x, bits.ffn_norm.float_data(),
5488 &mut attn_out, &mut zsh, n_embd, t, eps)?,
5489 None => e.add_rms_norm(cur, x, bits.ffn_norm.float_data(), &mut attn_out,
5490 &mut zsh, n_embd, t, eps)?,
5491 }
5492 let n_ff = ffn_gate.out_features();
5493 let (gate, up) = if t == 1 {
5499 let (zq, zd) = match zpair {
5500 Some(p) => p,
5501 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
5502 };
5503 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
5504 Some(p) => p,
5505 None => (e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
5506 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?),
5507 }
5508 } else {
5509 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5514 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
5515 let fused = if f2b {
5516 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
5517 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
5518 } else { None };
5519 match fused {
5520 Some(p) => p,
5521 None => {
5522 e.mmq_act_begin();
5524 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
5525 }
5526 }
5527 };
5528 let mut act = e.uninit(t * n_ff)?;
5529 let f0 = if e.uses_q8_1_fast(ffn_down) {
5532 let upv = e.view(&up, t * n_ff);
5533 let up_all = upv.slice(0..t * n_ff);
5534 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
5535 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5536 } else {
5537 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
5538 e.matmul(ffn_down, &act, t)?
5539 };
5540 if defer_post_norm { return Ok((f0, attn_out)); }
5541 let mut sn = e.uninit(t * n_embd)?;
5542 e.rms_norm(&f0, bits.post_ffw_norm.float_data(), &mut sn, n_embd, t, eps)?;
5543 return Ok((sn, attn_out));
5544 };
5545
5546 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
5547 let mut attn_out = e.uninit(t * n_embd)?;
5552 let mut router_in = e.uninit(t * n_embd)?;
5553 let fast_moe = match &layer.ffn {
5554 crate::hybrid::Ffn::Moe(m) => m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5555 && expert_dp4a_supported(m.gate_exps.qtype)
5556 && expert_dp4a_supported(m.up_exps.qtype)
5557 && expert_dp4a_supported(m.down_exps.qtype)
5558 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0"),
5559 _ => false,
5560 };
5561 let q8z = t < PRIME_MIN_T && fast_moe;
5562 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
5563 let (z0, m2) = e.add_rms_norm3_q8z(cur, x, bits.ffn_norm.float_data(),
5564 &mbits.router_scale_pre,
5565 mbits.pre_ffw_norm_2.float_data(),
5566 &mut attn_out, &mut router_in, n_embd, t, eps)?;
5567 (None, Some(z0), Some(m2))
5568 } else {
5569 let mut zsh = e.uninit(t * n_embd)?;
5570 let mut moe_in = e.uninit(t * n_embd)?;
5571 e.add_rms_norm3(cur, x, bits.ffn_norm.float_data(), &mbits.router_scale_pre,
5572 mbits.pre_ffw_norm_2.float_data(), &mut attn_out, &mut zsh,
5573 &mut router_in, &mut moe_in, n_embd, t, eps)?;
5574 (Some((zsh, moe_in)), None, None)
5575 };
5576 let attn_out2 = attn_out;
5577 #[allow(unused_variables)]
5578 let attn_out = &attn_out2;
5579 let n_ff = mbits.shared_gate.out_features();
5580 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
5581 if t == 1 {
5582 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
5583 Some(p) => p,
5584 None => {
5585 let h0 = e.zeros(0)?;
5586 (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
5587 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?)
5588 }
5589 }
5590 } else {
5591 let h0 = e.zeros(0)?;
5593 (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
5594 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?)
5595 }
5596 } else {
5597 let (zsh, _) = zsh_f32.as_ref().unwrap();
5598 (e.matmul(&mbits.shared_gate, zsh, t)?, e.matmul(&mbits.shared_up, zsh, t)?)
5599 };
5600 let mut act = e.uninit(t * n_ff)?;
5601 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
5602 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
5603 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else { panic!("gemma4 layer not MoE") };
5604 let moe0 = match (&moe_q8, &zsh_f32) {
5605 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
5606 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
5607 _ => unreachable!(),
5608 };
5609 let mut mlp = e.uninit(t * n_embd)?;
5611 let mut moe = e.uninit(t * n_embd)?;
5612 e.rms_norm2x(&mlp0, &moe0, mbits.post_ffw_norm_1.float_data(),
5613 mbits.post_ffw_norm_2.float_data(), &mut mlp, &mut moe, n_embd, t, eps)?;
5614
5615 let mut sum = e.uninit(t * n_embd)?;
5618 let mut sn = e.uninit(t * n_embd)?;
5619 e.add_rms_norm(&mlp, &moe, bits.post_ffw_norm.float_data(), &mut sum, &mut sn,
5620 n_embd, t, eps)?;
5621 Ok((sn, attn_out2))
5622 }
5623
5624 fn gemma4_layer_tail_add_nq(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5626 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5627 next_norm: Option<&CudaSlice<f32>>)
5628 -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>> {
5629 let n_embd = self.cfg.n_embd as usize;
5630 let bits = layer.gemma4.as_ref().unwrap();
5631 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
5632 let mut xn = e.uninit(t * n_embd)?;
5633 match next_norm {
5634 Some(w) => {
5635 let pair = e.add_scale_rms_norm_q8_1(&sn, &attn_out, bits.layer_scale, w, &mut xn,
5636 n_embd, t, self.cfg.rms_eps)?;
5637 Ok((xn, Some(pair)))
5638 }
5639 None => {
5640 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
5641 Ok((xn, None))
5642 }
5643 }
5644 }
5645
5646 fn gemma4_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
5649 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5650 if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, last_only); }
5653 let n_embd = self.cfg.n_embd as usize;
5654 let t = tokens.len();
5655 let pos: Vec<i32> = (0..t as i32).collect();
5656 let pos_d = e.htod_i32(&pos)?;
5657
5658 let mut x = self.embed(e, tokens)?;
5659 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
5660 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
5663 let stat = |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
5664 let h = e.dtoh(x)?;
5665 let bad = h.iter().filter(|v| !v.is_finite()).count();
5666 let mx = h.iter().filter(|v| v.is_finite()).fold(0.0f32, |m, v| m.max(v.abs()));
5667 eprintln!("[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}", &h[..3]);
5668 Ok(())
5669 };
5670 if probe { stat(e, &x, "embed")?; }
5671 for (il, layer) in self.layers.iter().enumerate() {
5672 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
5673 if probe { stat(e, &x, &format!("L{il}"))?; }
5674 }
5675 let mut hn = e.zeros(t * n_embd)?;
5676 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, self.cfg.rms_eps)?;
5677 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
5678 let n_vocab = self.output.out_features();
5679 let logits = if last_only {
5680 let hv = e.view(&hn, t * n_embd);
5681 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
5682 let mut hlast = e.zeros(n_embd)?;
5683 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
5684 let mut ld = e.matmul(&self.output, &hlast, 1)?;
5685 e.softcap(&mut ld, cap, n_vocab)?;
5686 self.gemma4_suppress(e, &mut ld, 1)?;
5687 e.dtoh(&ld)?
5688 } else {
5689 let mut ld = e.matmul(&self.output, &hn, t)?;
5690 e.softcap(&mut ld, cap, t * n_vocab)?;
5691 self.gemma4_suppress(e, &mut ld, t)?;
5692 e.dtoh(&ld)?
5693 };
5694 Ok(logits)
5695 }
5696
5697 pub(crate) fn gemma4_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
5702 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5703 if cache.pos != 0 {
5708 return Err("gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
5709 — prime the full prompt in one call or decode tokenwise".into());
5710 }
5711 let n_embd = self.cfg.n_embd as usize;
5712 let eps = self.cfg.rms_eps;
5713 let t = tokens.len();
5714 let pos: Vec<i32> = (0..t as i32).collect();
5715 let pos_d = e.htod_i32(&pos)?;
5716 let mut x = self.embed(e, tokens)?;
5717 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
5718 for (il, layer) in self.layers.iter().enumerate() {
5719 let mut h = e.zeros(t * n_embd)?;
5720 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5721 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer not full-attn") };
5722 let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
5723 let mut cur = e.zeros(t * n_embd)?;
5724 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
5725 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
5726 self.dflash_tap(e, cache, il, &x, t)?;
5727 }
5728 cache.pos += t;
5729 let hiddens = e.clone_dtod(&x)?;
5730 let xv = e.view(&x, t * n_embd);
5731 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
5732 let mut h_seed = e.zeros(n_embd)?;
5733 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
5734 let mut hn = e.uninit(n_embd)?;
5735 e.rms_norm(&h_seed, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
5736 let mut ld = e.matmul(&self.output, &hn, 1)?;
5737 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
5738 e.softcap(&mut ld, cap, self.output.out_features())?;
5739 self.gemma4_suppress(e, &mut ld, 1)?;
5740 let logits = e.dtoh(&ld)?;
5741 Ok((logits, h_seed, hiddens))
5742 }
5743
5744 fn gemma4_decode_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
5749 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
5750 pos_d: &CudaSlice<i32>, cache: &mut Cache)
5751 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5752 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
5753 let eps = self.cfg.rms_eps;
5754 let aux = self.gemma4_aux.as_ref().unwrap();
5755 let (hq, hdq) = (hq, hdq);
5756 let h0 = e.zeros(0)?;
5757 let h = &h0;
5758 let (q0, k0, v0) = if swa {
5759 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5760 Some(t3) => t3,
5761 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5762 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5763 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
5764 }
5765 } else {
5766 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
5767 Some(p) => p,
5768 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5769 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?),
5770 };
5771 let v0 = e.clone_dtod(&k0)?;
5772 (q0, k0, v0)
5773 };
5774 let mut q = e.uninit(nh * hd)?;
5775 let mut k = e.uninit(nkv * hd)?;
5776 let mut v = e.uninit(nkv * hd)?;
5777 let ff = if swa { None } else {
5780 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5781 };
5782 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5783 &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
5784 pos_d, nh, nkv, base, 1.0, ff, eps)?;
5785 let kvl = cache.kv[il].as_mut().unwrap();
5786 e.append_kv_quantized(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len,
5787 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()))?;
5788 kvl.len += 1;
5789 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5793 let mut attn = e.uninit(nh * hd)?;
5794 if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
5796 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5797 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5798 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5799 let base = kvl.len as i32;
5801 e.i32_set_k(&mut kvl.len_d, base)?;
5802 e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1, scale,
5803 kvl.k_tok_bytes, kvl.v_tok_bytes, Some((&kvl.len_d, -1)), false,
5804 false, None)?;
5805 return Ok(e.matmul(&fa.wo, &attn, 1)?);
5806 }
5807 if swa && kvl.len > win && hd == 256
5809 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5810 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5811 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5812 let base = kvl.len as i32;
5813 e.i32_set_k(&mut kvl.len_d, base)?;
5814 e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1, 1, scale,
5815 win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
5816 return Ok(e.matmul(&fa.wo, &attn, 1)?);
5817 }
5818 let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
5819 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
5820 (off_tok + t_kv) * kvl.k_tok_bytes);
5821 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
5822 (off_tok + t_kv) * kvl.v_tok_bytes);
5823 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
5824 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
5825 Ok(e.matmul(&fa.wo, &attn, 1)?)
5826 }
5827
5828 #[allow(clippy::too_many_arguments)]
5835 pub fn gemma4_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
5836 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5837 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5838 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>)
5839 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
5840 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
5841 self.gemma4_decode_step_dc_into(e, token_d, pos_d, embd_gpu, embd_qt, embd_rb, cache,
5842 n_vocab, cap_bucket_max, &mut tok_out)?;
5843 Ok(tok_out)
5844 }
5845
5846 #[allow(clippy::too_many_arguments)]
5849 pub fn gemma4_decode_step_dc_into(&self, e: &Engine, token_d: &CudaSlice<u32>,
5850 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5851 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5852 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
5853 tok_out: &mut CudaSlice<u32>)
5854 -> Result<(), Box<dyn std::error::Error>> {
5855 let n_embd = self.cfg.n_embd as usize;
5856 let eps = self.cfg.rms_eps;
5857 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
5858 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
5859 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5860 let n_layers = self.layers.len();
5861 for (il, layer) in self.layers.iter().enumerate() {
5862 let (hq, hdq) = match h_carry.take() {
5863 Some(p) => p,
5864 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
5865 };
5866 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5867 let o = self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
5868 let mut cur = e.uninit(n_embd)?;
5869 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
5870 let next_norm = if il + 1 < n_layers {
5871 Some(self.layers[il + 1].attn_norm.float_data())
5872 } else { None };
5873 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
5874 x = xn;
5875 h_carry = hn;
5876 }
5877 let mut hn = e.uninit(n_embd)?;
5878 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
5879 let mut logits = e.matmul(&self.output, &hn, 1)?;
5880 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
5882 e.inc_seqlen(pos_d)?;
5883 if cap_bucket_max.is_none() { cache.pos += 1; }
5884 Ok(())
5885 }
5886
5887 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
5894 let n_embd = self.cfg.n_embd as usize;
5895 let n_vocab = self.output.out_features();
5896 let n_layers = self.layers.len();
5897 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
5898 for il in 0..n_layers {
5899 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
5900 qmax = qmax.max(nh * hd);
5901 kvmax = kvmax.max(nkv * hd);
5902 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
5903 ffmax = ffmax.max(ffn_gate.out_features());
5904 }
5905 }
5906 Ok(G4DcSlots {
5907 x: e.uninit(n_embd)?, xn: e.uninit(n_embd)?, cur: e.uninit(n_embd)?,
5908 hq: e.alloc_i8_uninit(n_embd)?, hd_: e.uninit(n_embd / 32)?,
5909 q0: e.uninit(qmax)?, k0: e.uninit(kvmax)?, v0: e.uninit(kvmax)?,
5910 q: e.uninit(qmax)?, k: e.uninit(kvmax)?, v: e.uninit(kvmax)?,
5911 attn: e.uninit(qmax)?, o: e.uninit(n_embd)?,
5912 attn_out: e.uninit(n_embd)?, zsh: e.uninit(n_embd)?,
5913 zq: e.alloc_i8_uninit(n_embd.max(qmax))?, zd: e.uninit(n_embd.max(qmax) / 32)?,
5916 gate: e.uninit(ffmax)?, up: e.uninit(ffmax)?,
5917 act: e.uninit(ffmax)?, actq: e.alloc_i8_uninit(ffmax)?, actd: e.uninit(ffmax / 32)?,
5918 f0: e.uninit(n_embd)?, sn: e.uninit(n_embd)?,
5919 hn: e.uninit(n_embd)?, logits: e.uninit(n_vocab)?,
5920 })
5921 }
5922
5923 fn g4_matvec_m1_into(&self, e: &Engine, w: &crate::model::GpuTensor,
5926 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y: &mut CudaSlice<f32>)
5927 -> Result<(), Box<dyn std::error::Error>> {
5928 use crate::model::GpuTensor;
5929 let (bytes, qtype, row_bytes, scale, rp) = match w {
5930 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
5931 (bytes, *qtype, *row_bytes, *scale, *rp),
5932 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
5933 };
5934 let (mbytes, mrp) = match w {
5935 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5936 _ => (bytes, rp),
5937 };
5938 e.qmatvec_mmvq_into(mbytes, aq, ad, 1, w.in_features(), w.out_features(),
5939 qtype, row_bytes, scale, mrp, y)
5940 }
5941
5942 #[allow(clippy::too_many_arguments)]
5946 pub fn gemma4_decode_step_dc_slotted(&self, e: &Engine, token_d: &CudaSlice<u32>,
5947 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5948 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5949 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
5950 sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>,
5951 ring: Option<(&mut CudaSlice<u32>, usize)>)
5952 -> Result<(), Box<dyn std::error::Error>> {
5953 let n_embd = self.cfg.n_embd as usize;
5954 let eps = self.cfg.rms_eps;
5955 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
5956 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
5957 let n_layers = self.layers.len();
5958 let mut has_carry = false;
5959 for il in 0..n_layers {
5960 if !has_carry {
5961 e.rms_norm_q8_1_into(&sl.x, self.layers[il].attn_norm.float_data(), n_embd, 1,
5962 eps, &mut sl.hq, &mut sl.hd_)?;
5963 }
5964 has_carry = true;
5965 let layer = &self.layers[il];
5966 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5967 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
5968 e.rms_norm(&sl.o, layer.post_attn_norm.float_data(), &mut sl.cur, n_embd, 1, eps)?;
5969 let next_norm = if il + 1 < n_layers {
5970 Some(self.layers[il + 1].attn_norm.float_data())
5971 } else { None };
5972 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
5973 std::mem::swap(&mut sl.x, &mut sl.xn);
5974 }
5975 e.rms_norm(&sl.x, self.output_norm.float_data(), &mut sl.hn, n_embd, 1, eps)?;
5976 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
5977 {
5979 let (zq, zd) = (&sl.zq, &sl.zd);
5980 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
5981 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
5982 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
5983 }
5984 self.gemma4_suppress(e, &mut sl.logits, 1)?;
5985 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
5986 if let Some((ring, base)) = ring {
5987 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
5991 }
5992 e.inc_seqlen(pos_d)?;
5993 if cap_bucket_max.is_none() { cache.pos += 1; }
5994 Ok(())
5995 }
5996
5997 #[allow(clippy::too_many_arguments)]
5999 fn gemma4_decode_attn_dc_slotted(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer,
6000 il: usize, pos_d: &CudaSlice<i32>, cache: &mut Cache,
6001 cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots)
6002 -> Result<(), Box<dyn std::error::Error>> {
6003 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6004 let eps = self.cfg.rms_eps;
6005 let aux = self.gemma4_aux.as_ref().unwrap();
6006 {
6007 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
6008 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
6009 if swa {
6010 if !e.matmul_q4_fused3_into(&fa.wq, &fa.wk, &fa.wv, hq, hdq,
6011 &mut sl.q0, &mut sl.k0, &mut sl.v0)? {
6012 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
6013 }
6014 } else {
6015 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
6016 return Err("slotted step: fused2 unavailable".into());
6017 }
6018 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
6019 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
6020 }
6021 }
6022 let ff = if swa { None } else {
6025 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6026 };
6027 let kvl = cache.kv[il].as_mut().unwrap();
6028 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6029 if crate::Engine::qkv_append_on() {
6030 e.rms_norm_qkv_rope_append_dc(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(),
6032 fa.k_norm.float_data(), &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
6033 pos_d, nh, nkv, base, 1.0, ff, eps,
6034 &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
6035 } else {
6036 e.rms_norm_qkv_rope(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6037 &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
6038 pos_d, nh, nkv, base, 1.0, ff, eps)?;
6039 e.append_kv_quantized_dc(&sl.k, &sl.v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
6040 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
6041 kv_fp8)?;
6042 }
6043 e.inc_seqlen(&mut kvl.len_d)?;
6044 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
6045 let k_view = e.view_u8(&kvl.k, kvl.k.len());
6046 let v_view = e.view_u8(&kvl.v, kvl.v.len());
6047 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
6048 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6049 let mut fa_q8 = false;
6053 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
6054 e.fa_decode_rows(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, b_glob - 1,
6055 1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6056 Some((&kvl.len_d, -1)), false, false,
6057 Some((&mut sl.zq, &mut sl.zd)))?;
6058 fa_q8 = true;
6059 } else if swa && b_swa > win && hd == 256 && rows_on {
6060 e.fa_decode_rows_w(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv,
6061 &kvl.len_d, -1, 1, scale, win,
6062 kvl.k_tok_bytes, kvl.v_tok_bytes,
6063 Some((&mut sl.zq, &mut sl.zd)))?;
6064 fa_q8 = true;
6065 } else {
6066 let b = if swa { b_swa } else { b_glob };
6067 e.fa_decode_dc(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, &kvl.len_d, b,
6068 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6069 swa && crate::Engine::wkv_on())?;
6070 }
6071 if !fa_q8 {
6072 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
6073 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
6074 }
6075 {
6076 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
6077 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
6078 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
6079 }
6080 Ok(())
6081 }
6082
6083 fn gemma4_layer_tail_slotted(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
6086 next_norm: Option<&CudaSlice<f32>>, sl: &mut G4DcSlots)
6087 -> Result<(), Box<dyn std::error::Error>> {
6088 let n_embd = self.cfg.n_embd as usize;
6089 let eps = self.cfg.rms_eps;
6090 let bits = layer.gemma4.as_ref().unwrap();
6091 let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
6092 else { return Err("slotted tail: dense ffn only".into()) };
6093 e.add_rms_norm(&sl.cur, &sl.x, bits.ffn_norm.float_data(), &mut sl.attn_out,
6094 &mut sl.zsh, n_embd, 1, eps)?;
6095 let n_ff = ffn_gate.out_features();
6096 {
6097 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
6098 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
6099 }
6100 {
6101 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
6102 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
6103 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
6104 return Err("slotted tail: ffn fused2 unavailable".into());
6105 }
6106 }
6107 debug_assert!(e.uses_q8_1_fast(ffn_down));
6108 {
6109 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
6110 let upv = e.view(upr, n_ff);
6111 let up_all = upv.slice(0..n_ff);
6112 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
6113 e.gelu_tanh_mul_q8_1_into(gr, &up_all, &mut sl.act, n_ff, 1,
6114 &mut sl.actq, &mut sl.actd)?;
6115 }
6116 {
6117 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
6118 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
6119 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
6120 }
6121 e.rms_norm(&sl.f0, bits.post_ffw_norm.float_data(), &mut sl.sn, n_embd, 1, eps)?;
6122 match next_norm {
6123 Some(w) => {
6124 e.add_scale_rms_norm_q8_1_into(&sl.sn, &sl.attn_out, bits.layer_scale, w,
6125 &mut sl.xn, n_embd, 1, eps,
6126 &mut sl.hq, &mut sl.hd_)?;
6127 }
6128 None => {
6129 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
6130 }
6131 }
6132 Ok(())
6133 }
6134
6135 #[allow(clippy::too_many_arguments)]
6137 fn gemma4_decode_attn_dc(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6138 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6139 pos_d: &CudaSlice<i32>, cache: &mut Cache,
6140 cap_bucket_max: Option<(usize, usize)>)
6141 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6142 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6143 let eps = self.cfg.rms_eps;
6144 let aux = self.gemma4_aux.as_ref().unwrap();
6145 let (q0, k0, v0) = if swa {
6146 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
6147 Some(t3) => t3,
6148 None => {
6149 let h0 = e.zeros(0)?;
6150 (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
6151 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
6152 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?)
6153 }
6154 }
6155 } else {
6156 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
6157 Some(p) => p,
6158 None => {
6159 let h0 = e.zeros(0)?;
6160 (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
6161 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?)
6162 }
6163 };
6164 let v0 = e.clone_dtod(&k0)?;
6165 (q0, k0, v0)
6166 };
6167 let mut q = e.uninit(nh * hd)?;
6168 let mut k = e.uninit(nkv * hd)?;
6169 let mut v = e.uninit(nkv * hd)?;
6170 let ff = if swa { None } else {
6172 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6173 };
6174 let kvl = cache.kv[il].as_mut().unwrap();
6175 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6176 if crate::Engine::qkv_append_on() {
6177 e.rms_norm_qkv_rope_append_dc(&q0, &k0, &v0, fa.q_norm.float_data(),
6179 fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
6180 pos_d, nh, nkv, base, 1.0, ff, eps,
6181 &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
6182 } else {
6183 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6184 &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
6185 pos_d, nh, nkv, base, 1.0, ff, eps)?;
6186 e.append_kv_quantized_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
6187 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
6188 }
6189 e.inc_seqlen(&mut kvl.len_d)?;
6190 let mut attn = e.uninit(nh * hd)?;
6191 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6194 match cap_bucket_max {
6199 None => {
6200 kvl.len += 1;
6204 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6205 if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
6206 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6207 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
6210 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
6211 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6212 e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1,
6213 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6214 Some((&kvl.len_d, -1)), false, false,
6215 Some((&mut aq8, &mut ad8)))?;
6216 fa_q8 = Some((aq8, ad8));
6217 } else if swa && kvl.len > win && hd == 256
6218 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6219 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
6221 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
6222 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6223 e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1,
6224 1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes,
6225 Some((&mut aq8, &mut ad8)))?;
6226 fa_q8 = Some((aq8, ad8));
6227 } else {
6228 let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) }
6229 else { (0, kvl.len) };
6230 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
6231 (off_tok + t_kv) * kvl.k_tok_bytes);
6232 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
6233 (off_tok + t_kv) * kvl.v_tok_bytes);
6234 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
6235 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
6236 }
6237 }
6238 Some((b_swa, b_glob)) => {
6239 let k_view = e.view_u8(&kvl.k, kvl.k.len());
6245 let v_view = e.view_u8(&kvl.v, kvl.v.len());
6246 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
6247 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6248 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
6249 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6250 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, b_glob - 1,
6251 1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6252 Some((&kvl.len_d, -1)), false, false,
6253 Some((&mut aq8, &mut ad8)))?;
6254 fa_q8 = Some((aq8, ad8));
6255 } else if swa && b_swa > win && hd == 256 && rows_on {
6256 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6257 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6258 &kvl.len_d, -1, 1, scale, win,
6259 kvl.k_tok_bytes, kvl.v_tok_bytes,
6260 Some((&mut aq8, &mut ad8)))?;
6261 fa_q8 = Some((aq8, ad8));
6262 } else {
6263 let b = if swa { b_swa } else { b_glob };
6264 e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, b,
6265 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6266 swa && crate::Engine::wkv_on())?;
6267 }
6268 }
6269 }
6270 if let Some((aq8, ad8)) = fa_q8 {
6273 let mut y = e.uninit(fa.wo.out_features())?;
6274 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
6275 return Ok(y);
6276 }
6277 Ok(e.matmul(&fa.wo, &attn, 1)?)
6278 }
6279
6280 pub fn gemma4_generate_graph(&self, e: &Engine, prompt_pos: usize, first_token: u32,
6285 cache: &mut Cache, max_new: usize, eos: &[u32],
6286 mut on_token: impl FnMut(u32) -> bool)
6287 -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
6288 if self.is_gemma4_e4b() {
6289 return Err("E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm".into());
6290 }
6291 use crate::decode::StopReason;
6292 let n_vocab = self.output.out_features();
6293 let n_embd = self.cfg.n_embd as usize;
6294 let embd_gpu = self.embd_gpu.get_or_init(|| {
6295 e.upload_u8(&self.embd.raw).expect("embed table upload")
6296 });
6297 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6298 for kvl in cache.kv.iter_mut().flatten() {
6299 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6300 }
6301 let mut token_d = e.stream().clone_htod(&[first_token])?;
6302 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
6303 let g4 = self.cfg.gemma4.as_ref().unwrap();
6304 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
6305 let nkv_s = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
6307 .find(|p| *p.1).map(|p| *p.0 as usize).unwrap_or(8);
6308 let nkv_g = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
6309 .find(|p| !*p.1).map(|p| *p.0 as usize).unwrap_or(2);
6310 let mut graphs: std::collections::HashMap<((bool, usize), (bool, usize), bool, bool),
6311 (cudarc::driver::CudaGraph,
6312 Vec<Box<dyn std::any::Any + Send>>)> = Default::default();
6313 let mut slots = self.g4_dc_slots(e)?;
6316 const RING: usize = 64;
6319 const DRAIN: usize = 1;
6325 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
6326 let ring_base = prompt_pos;
6327 let mut out = Vec::with_capacity(max_new);
6328 let mut reason = StopReason::MaxNew;
6329 let mut next = first_token;
6330 let mut captures = 0usize;
6331 for _ in 0..max_new {
6332 out.push(next);
6333 if eos.contains(&next) { reason = StopReason::Eos; break; }
6334 if !on_token(next) { reason = StopReason::Callback; break; }
6335 let t_kv = cache.pos + 1;
6336 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6344 let f512 = crate::fa512_min_tkv();
6345 let key_s = if t_kv > win { (true, usize::MAX) }
6346 else { e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on()) };
6347 let (key_g, rung_end) = if t_kv >= f512 {
6348 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
6351 ((true, end), end)
6352 } else { (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv) };
6353 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
6354 if !graphs.contains_key(&key) {
6355 let bucket_max = (t_kv, rung_end);
6356 let snap = cache.snapshot(e)?;
6358 let pos_save = e.dtoh_i32_one(&pos_d)?;
6359 let len_save: Vec<Option<i32>> = cache.kv.iter()
6360 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
6361 let tok_save = e.dtoh_u32_one(&token_d)?;
6362 let graph = {
6367 let tok_ref = &mut token_d;
6368 let pos_ref = &mut pos_d;
6369 let cache_ref = &mut *cache;
6370 let slots_ref = &mut slots;
6371 let ring_ref = &mut ring;
6372 e.capture_graph_retained_flags(
6373 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
6374 |e| {
6375 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
6377 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
6378 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
6379 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
6380 cache_ref, n_vocab, Some(bucket_max),
6381 sl, tok_ref, Some((rg, ring_base)))
6382 })?
6383 };
6384 cache.rollback(e, &snap, 0)?;
6385 e.set_i32_one(&mut pos_d, pos_save)?;
6386 for (il, ls) in len_save.iter().enumerate() {
6387 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
6388 e.set_i32_one(&mut kvl.len_d, *v)?;
6389 }
6390 }
6391 e.set_u32_one(&mut token_d, tok_save)?;
6392 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
6393 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
6394 eprintln!("[graph-census] {c:?}");
6395 }
6396 }
6397 graphs.insert(key, graph);
6398 captures += 1;
6399 }
6400 let mut chunk = 1usize;
6405 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN").ok()
6406 .and_then(|v| v.parse().ok()).unwrap_or(DRAIN);
6407 while chunk < drain_cap && out.len() + chunk < max_new {
6408 let t_next = cache.pos + 1 + chunk;
6409 let key_s2 = if t_next > win { (true, usize::MAX) }
6410 else { e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on()) };
6411 let key_g2 = if t_next >= f512 {
6412 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
6413 } else { e.fa_bucket_key(t_next, hd_g, nkv_g, false) };
6414 if (key_s2, key_g2, t_next >= f512, t_next > win) != key { break; }
6415 chunk += 1;
6416 }
6417 let g = &graphs.get(&key).unwrap().0;
6418 for _ in 0..chunk { g.launch()?; }
6419 e.stream().synchronize()?;
6420 let ringh = e.dtoh_u32(&ring)?;
6421 for j in 0..chunk {
6422 let pos_j = cache.pos + j;
6423 let tok_j = ringh[(pos_j - ring_base) % RING];
6424 cache.pos += 0; if j + 1 == chunk { next = tok_j; }
6426 else {
6427 out.push(tok_j);
6428 if eos.contains(&tok_j) || !on_token(tok_j) {
6429 reason = if eos.contains(&tok_j) { StopReason::Eos }
6430 else { StopReason::Callback };
6431 let keep = cache.pos + j + 1;
6433 e.set_i32_one(&mut pos_d, keep as i32)?;
6434 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
6435 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
6436 kvl.len = keep;
6437 }
6438 cache.pos = keep;
6439 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
6440 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
6441 }
6442 return Ok((out, reason));
6443 }
6444 }
6445 }
6446 cache.pos += chunk;
6447 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) { kvl.len += chunk; }
6448 }
6449 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
6450 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
6451 }
6452 Ok((out, reason))
6453 }
6454
6455 pub(crate) fn gemma4_decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize,
6461 cache: &mut Cache)
6462 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6463 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
6464 }
6465
6466 pub(crate) fn gemma4_decode_step_t_am(&self, e: &Engine, tokens: &[u32], pos0: usize,
6470 cache: &mut Cache)
6471 -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6472 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
6473 let t = tokens.len();
6474 let n_vocab = self.output.out_features();
6475 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
6476 for i in 0..t {
6477 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
6478 }
6479 Ok((e.dtoh_u32(&toks)?, hn))
6480 }
6481
6482 pub(crate) fn gemma4_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
6485 pos0: usize, cache: &mut Cache)
6486 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6487 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
6488 let n_vocab = self.output.out_features();
6489 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
6490 for i in 0..t {
6491 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
6492 }
6493 Ok((vam, hn))
6494 }
6495
6496 pub(crate) fn gemma4_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
6499 cache: &mut Cache)
6500 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6501 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
6502 let t = tokens.len();
6503 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6504 e.softcap(&mut ld, cap, t * self.output.out_features())?;
6505 Ok((e.dtoh(&ld)?, hn))
6506 }
6507
6508 pub(crate) fn verify_stream_scratch(&self, e: &Engine, cap: usize)
6511 -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
6512 Ok(VerifyStreamScratch {
6513 pos_d: e.htod_i32(&vec![0i32; cap])?,
6514 row_ctrs: (0..cap).map(|_| e.htod_i32(&[0])).collect::<Result<_, _>>()?,
6515 })
6516 }
6517
6518 pub(crate) fn gemma4_verify_t_am_stream(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
6526 ctr: &CudaSlice<i32>, hint: usize,
6527 cache: &mut Cache,
6528 scr: &mut VerifyStreamScratch)
6529 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6530 let n_embd = self.cfg.n_embd as usize;
6531 let eps = self.cfg.rms_eps;
6532 assert!(t <= scr.row_ctrs.len() && t <= 64);
6533 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
6534 for i in 0..t {
6535 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
6536 }
6537 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
6538 let embd_gpu = self.embd_gpu.get_or_init(|| {
6539 e.upload_u8(&self.embd.raw).expect("embed table upload")
6540 });
6541 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6542 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
6543 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6544 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6545 let n_layers = self.layers.len();
6546 for (il, layer) in self.layers.iter().enumerate() {
6547 let (hq, hdq) = match h_carry.take() {
6548 Some(p) => p,
6549 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
6550 };
6551 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6552 let o = self.gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache,
6553 hint, row_ctrs)?;
6554 let mut cur = e.uninit(t * n_embd)?;
6555 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6556 let next_norm = if il + 1 < n_layers {
6557 Some(self.layers[il + 1].attn_norm.float_data())
6558 } else { None };
6559 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
6560 x = xn;
6561 h_carry = hn;
6562 self.dflash_tap(e, cache, il, &x, t)?;
6563 }
6564 let mut hn = e.uninit(t * n_embd)?;
6565 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6566 let ld = e.matmul(&self.output, &hn, t)?;
6567 let n_vocab = self.output.out_features();
6568 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
6569 for i in 0..t {
6570 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
6571 }
6572 Ok((vam, hn))
6573 }
6574
6575 fn dflash_tap(&self, e: &Engine, cache: &mut Cache, il: usize, x: &CudaSlice<f32>, t: usize)
6582 -> Result<(), Box<dyn std::error::Error>> {
6583 let Some(taps) = cache.dflash_taps.as_mut() else { return Ok(()) };
6584 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else { return Ok(()) };
6585 let h = taps.hidden;
6586 let n_taps = taps.layer_ids.len();
6587 debug_assert_eq!(taps.t, t);
6588 let xv = e.view(x, t * h);
6589 for r in 0..t {
6590 let row = xv.slice(r * h..(r + 1) * h);
6591 e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
6592 }
6593 Ok(())
6594 }
6595
6596 fn gemma4_verify_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
6597 tok_dev: Option<&CudaSlice<u32>>)
6598 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6599 let n_embd = self.cfg.n_embd as usize;
6600 let eps = self.cfg.rms_eps;
6601 let t = tokens.len();
6602 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6603 let pos_d = e.htod_i32(&pos)?;
6604 let mut x = match tok_dev {
6605 Some(td) => {
6606 let embd_gpu = self.embd_gpu.get_or_init(|| {
6607 e.upload_u8(&self.embd.raw).expect("embed table upload")
6608 });
6609 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6610 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
6611 }
6612 None => e.htod(&self.embd.gather(n_embd, tokens))?,
6613 };
6614 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6615 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6616 let n_layers = self.layers.len();
6617 for (il, layer) in self.layers.iter().enumerate() {
6618 let (hq, hdq) = match h_carry.take() {
6619 Some(p) => p,
6620 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
6621 };
6622 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6623 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
6624 let mut cur = e.uninit(t * n_embd)?;
6625 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6626 let next_norm = if il + 1 < n_layers {
6627 Some(self.layers[il + 1].attn_norm.float_data())
6628 } else { None };
6629 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
6630 x = xn;
6631 h_carry = hn;
6632 self.dflash_tap(e, cache, il, &x, t)?;
6633 }
6634 let mut hn = e.uninit(t * n_embd)?;
6635 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6636 let mut ld = e.matmul(&self.output, &hn, t)?;
6637 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
6639 Ok((ld, hn))
6640 }
6641
6642 #[allow(clippy::too_many_arguments)]
6650 fn gemma4_verify_attn_stream(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6651 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6652 pos_d: &CudaSlice<i32>, t: usize,
6653 cache: &mut Cache, hint: usize,
6654 row_ctrs: &[CudaSlice<i32>])
6655 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6656 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6657 let eps = self.cfg.rms_eps;
6658 let aux = self.gemma4_aux.as_ref().unwrap();
6659 let h0 = e.zeros(0)?;
6660 let h = &h0;
6661 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6664 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6665 let fused_qkv = if f2b {
6666 if swa {
6667 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6668 .map(|(a, b, c)| (a, b, Some(c)))
6669 } else {
6670 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
6671 .map(|(a, b)| (a, b, None))
6672 }
6673 } else { None };
6674 let (q0, k0, v0) = match fused_qkv {
6675 Some((a, b, cv)) => {
6676 let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
6677 (a, b, v)
6678 }
6679 None => {
6680 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6681 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
6682 let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
6683 else { e.clone_dtod(&k0)? };
6684 (q0, k0, v0)
6685 }
6686 };
6687 let mut q = e.uninit(t * nh * hd)?;
6688 let mut k = e.uninit(t * nkv * hd)?;
6689 let mut v = e.uninit(t * nkv * hd)?;
6690 let ff = if swa { None } else {
6693 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6694 };
6695 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6696 &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
6697 pos_d, nh, nkv, base, 1.0, ff, eps)?;
6698 let kvl = cache.kv[il].as_mut().unwrap();
6699 e.append_kv_quantized_rows_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d, t,
6701 kvl.kv_dim_k, kvl.kv_dim_v,
6702 kvl.k_tok_bytes, kvl.v_tok_bytes,
6703 (!swa && crate::Engine::gkv_on())
6704 || (swa && crate::Engine::wkv_on()))?;
6705 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6708 let mut attn = e.uninit(t * nh * hd)?;
6709 let k_view = e.view_u8(&kvl.k, kvl.k.len());
6710 let v_view = e.view_u8(&kvl.v, kvl.v.len());
6711 if swa && hint + 1 >= win {
6714 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6717 &kvl.len_d, 0, t, scale, win,
6718 kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6719 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
6720 let bucket = (hint + t + 2).next_power_of_two()
6733 .min(crate::fa512_min_tkv().saturating_sub(1));
6734 let qv = e.view(&q, t * nh * hd);
6735 for i in 0..t {
6736 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
6737 let mut q_one = e.uninit(nh * hd)?;
6738 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6739 let mut a_one = e.uninit(nh * hd)?;
6740 e.fa_decode_dc(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv,
6741 &row_ctrs[i], bucket, scale,
6742 kvl.k_tok_bytes, kvl.v_tok_bytes, false)?;
6743 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6744 }
6745 } else if hd == 512 {
6746 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, hint, t, scale,
6749 kvl.k_tok_bytes, kvl.v_tok_bytes,
6750 Some((&kvl.len_d, 0)), false, false, None)?;
6751 } else {
6752 e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6754 &kvl.len_d, hint + t, t, scale,
6755 kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
6756 swa && crate::Engine::wkv_on())?;
6757 }
6758 Ok(e.matmul(&fa.wo, &attn, t)?)
6759 }
6760
6761 fn gemma4_verify_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6762 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6763 pos_d: &CudaSlice<i32>, t: usize,
6764 cache: &mut Cache)
6765 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6766 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6767 let eps = self.cfg.rms_eps;
6768 let aux = self.gemma4_aux.as_ref().unwrap();
6769 let n_embd = self.cfg.n_embd as usize;
6770 let _ = n_embd;
6771
6772 let h0 = e.zeros(0)?;
6773 let h = &h0;
6774 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6777 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6778 let fused_qkv = if f2b {
6779 if swa {
6780 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6781 .map(|(a, b, c)| (a, b, Some(c)))
6782 } else {
6783 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
6784 .map(|(a, b)| (a, b, None))
6785 }
6786 } else { None };
6787 let (q0, k0, v0) = match fused_qkv {
6788 Some((a, b, cv)) => {
6789 let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
6790 (a, b, v)
6791 }
6792 None => {
6793 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6794 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
6795 let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
6796 else { e.clone_dtod(&k0)? };
6797 (q0, k0, v0)
6798 }
6799 };
6800 let mut q = e.uninit(t * nh * hd)?;
6801 let mut k = e.uninit(t * nkv * hd)?;
6802 let mut v = e.uninit(t * nkv * hd)?;
6803 let ff = if swa { None } else {
6806 Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6807 };
6808 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6809 &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
6810 pos_d, nh, nkv, base, 1.0, ff, eps)?;
6811 let kvl = cache.kv[il].as_mut().unwrap();
6812 let base_len = kvl.len;
6813 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
6814 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()))?;
6815 kvl.len += t;
6816 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6817 let mut attn = e.uninit(t * nh * hd)?;
6818 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
6821 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
6824 if rows_ok && (!swa || base_len + t <= win) {
6825 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
6826 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
6827 if hd == 512 {
6828 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6830 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, base_len, t,
6831 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6832 Some((&kvl.len_d, 0)), false,
6833 swa && crate::Engine::wkv_on(), None)?;
6834 } else {
6835 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6839 e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6840 &kvl.len_d, base_len + t, t, scale,
6841 kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
6842 swa && crate::Engine::wkv_on())?;
6843 }
6844 return Ok(e.matmul(&fa.wo, &attn, t)?);
6845 }
6846 if hd == 256 && swa && base_len + 1 >= win
6854 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6855 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
6856 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
6857 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6858 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, 0,
6859 t, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6860 return Ok(e.matmul(&fa.wo, &attn, t)?);
6861 }
6862 for i in 0..t {
6863 let avail = base_len + i + 1;
6864 let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
6865 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
6866 (off_tok + t_kv) * kvl.k_tok_bytes);
6867 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
6868 (off_tok + t_kv) * kvl.v_tok_bytes);
6869 let qi = e.view(&q, t * nh * hd);
6870 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
6871 let mut q_one = e.uninit(nh * hd)?;
6872 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6873 let mut a_one = e.uninit(nh * hd)?;
6874 if swa && avail > win && hd == 256
6878 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6879 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
6880 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
6881 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
6882 e.fa_decode_rows_w(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, &kvl.len_d, 0,
6883 1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6884 } else if !swa && hd == 512 && avail >= crate::fa512_min_tkv()
6885 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6886 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
6887 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
6888 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
6889 e.fa_decode_rows(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, avail - 1, 1,
6890 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6891 Some((&kvl.len_d, 0)), false, false, None)?;
6892 } else {
6893 e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
6894 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
6895 }
6896 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6897 }
6898 Ok(e.matmul(&fa.wo, &attn, t)?)
6899 }
6900
6901 pub(crate) fn gemma4_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
6904 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6905 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
6910 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
6911 }
6912 if crate::pp::pp_cuts(self.layers.len()).is_some() {
6913 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
6914 }
6915 let n_embd = self.cfg.n_embd as usize;
6916 let eps = self.cfg.rms_eps;
6917 let pos_d = e.htod_i32(&[cache.pos as i32])?;
6918 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
6919 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
6920 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6923 let n_layers = self.layers.len();
6924 for (il, layer) in self.layers.iter().enumerate() {
6925 let (hq, hdq) = match h_carry.take() {
6926 Some(p) => p,
6927 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
6928 };
6929 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6930 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
6931 let mut cur = e.uninit(n_embd)?;
6932 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
6933 let next_norm = if il + 1 < n_layers {
6934 Some(self.layers[il + 1].attn_norm.float_data())
6935 } else { None };
6936 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
6937 x = xn;
6938 h_carry = hn;
6939 }
6940 let mut hn = e.uninit(n_embd)?;
6941 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
6942 let h_seed = e.clone_dtod(&x)?;
6943 let mut ld = e.matmul(&self.output, &hn, 1)?;
6944 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6945 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
6947 let logits = e.dtoh(&ld)?;
6948 cache.pos += 1;
6949 Ok((logits, h_seed))
6950 }
6951
6952 fn gemma4_decode_layers(&self, e: &Engine, mut x: CudaSlice<f32>, lo: usize, hi: usize,
6960 pos_d: &CudaSlice<i32>, cache: &mut Cache)
6961 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6962 let n_embd = self.cfg.n_embd as usize;
6963 let eps = self.cfg.rms_eps;
6964 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6965 for il in lo..hi {
6966 let layer = &self.layers[il];
6967 let (hq, hdq) = match h_carry.take() {
6968 Some(p) => p,
6969 None => e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?,
6971 };
6972 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6973 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
6974 let mut cur = e.uninit(n_embd)?;
6975 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
6976 let next_norm = if il + 1 < hi {
6977 Some(self.layers[il + 1].attn_norm.float_data())
6978 } else { None };
6979 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
6980 x = xn;
6981 h_carry = hn;
6982 }
6983 Ok(x)
6984 }
6985
6986 fn gemma4_decode_step_h_pp2(&self, e: &Engine, token: u32, cache: &mut Cache, split: usize)
6993 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6994 if crate::pp::pp2_streams_off() {
6995 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
6996 }
6997 let rt = crate::pp::Pp2Rt::get(e)?;
6998 let e0 = rt.engine(0, e);
6999 let e1 = rt.engine(1, e);
7000 let n_embd = self.cfg.n_embd as usize;
7001 let eps = self.cfg.rms_eps;
7002
7003 let (pos_d, slot) = {
7005 let _st0 = rt.enter(0);
7006 let pos_d = e0.htod_i32(&[cache.pos as i32])?;
7007 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
7008 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7009 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
7010 let slot = rt.tx(0, &x, n_embd)?;
7011 (pos_d, slot)
7012 };
7013
7014 let _st1 = rt.enter(1);
7016 let x = rt.rx(0, slot, n_embd)?;
7017 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
7018
7019 let mut hn = e1.uninit(n_embd)?;
7020 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7021 let h_seed = e1.clone_dtod(&x)?;
7022 let mut ld = e1.matmul(&self.output, &hn, 1)?;
7023 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7024 e1.softcap(&mut ld, cap, self.output.out_features())?;
7025 self.gemma4_suppress(e1, &mut ld, 1)?;
7026 let logits = e1.dtoh(&ld)?;
7027 cache.pos += 1;
7028 Ok((logits, h_seed))
7029 }
7030
7031 fn gemma4_decode_step_h_pp2_samestream(&self, e: &Engine, token: u32, cache: &mut Cache,
7034 split: usize)
7035 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7036 let n_embd = self.cfg.n_embd as usize;
7037 let eps = self.cfg.rms_eps;
7038 let pos_d = e.htod_i32(&[cache.pos as i32])?;
7039
7040 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
7042 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7043 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
7044
7045 let boundary_tx = e.clone_dtod(&x)?;
7047 let boundary_rx = e.clone_dtod(&boundary_tx)?;
7048
7049 let x = self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
7051
7052 let mut hn = e.uninit(n_embd)?;
7053 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7054 let h_seed = e.clone_dtod(&x)?;
7055 let mut ld = e.matmul(&self.output, &hn, 1)?;
7056 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7057 e.softcap(&mut ld, cap, self.output.out_features())?;
7058 self.gemma4_suppress(e, &mut ld, 1)?;
7059 let logits = e.dtoh(&ld)?;
7060 cache.pos += 1;
7061 Ok((logits, h_seed))
7062 }
7063}
7064
7065impl HybridModel {
7084 pub(crate) fn step35_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
7086 let s = self.cfg.step35.as_ref().unwrap();
7087 let hd = self.cfg.head_dim_k as usize;
7088 (hd,
7089 s.n_head_kv(il as u32) as usize, s.n_head(il as u32) as usize, s.rope_base(il as u32), 1.0 / (hd as f32).sqrt(), s.is_swa(il as u32))
7094 }
7095
7096 #[allow(clippy::too_many_arguments)]
7154 fn step35_attn_pre_wo(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
7155 hg: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
7156 cache: Option<&mut Cache>, il: usize, seq_end: usize)
7157 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7158 let (hd, nkv, nh, rbase, scale, swa) = self.step35_geom(il);
7159 let eps = self.cfg.rms_eps;
7160 let s = self.cfg.step35.as_ref().unwrap();
7161 let win = s.sliding_window as usize;
7162 let n_rot = s.n_rot(il as u32) as usize;
7163
7164 let v = g3.pop().unwrap();
7165 let k0 = g3.pop().unwrap();
7166 let q0 = g3.pop().unwrap();
7167
7168 let mut q = e.uninit(t * nh * hd)?;
7172 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
7173 let mut k = e.uninit(t * nkv * hd)?;
7174 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
7175 let ff = if swa { None } else {
7176 self.step35_aux.as_ref().and_then(|a| a.rope_freqs.as_ref())
7177 };
7178 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
7179
7180 let mut attn = e.uninit(t * nh * hd)?;
7181 match cache {
7182 Some(cache) => {
7183 let base_len = {
7184 let kvl = cache.kv[il].as_mut().unwrap();
7185 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
7186 let base_len = kvl.len;
7187 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
7188 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
7189 kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
7190 kvl.len += t;
7191 let new_len = kvl.len as i32;
7192 e.set_i32_one(&mut kvl.len_d, new_len)?;
7193 base_len
7194 };
7195 let kvl = cache.kv[il].as_ref().unwrap();
7196 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
7206 let legacy_calllocal =
7207 std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
7208 let off = if swa {
7223 let raw = base_len.saturating_sub(win - 1);
7224 if legacy_tkv || legacy_calllocal { raw } else { raw & !31usize }
7225 } else {
7226 0
7227 };
7228 let t_kv = base_len + t - off;
7229 let k_view = e.view_u8_range(&kvl.k, off * kvl.k_tok_bytes,
7230 (off + t_kv) * kvl.k_tok_bytes);
7231 let v_view = e.view_u8_range(&kvl.v, off * kvl.v_tok_bytes,
7232 (off + t_kv) * kvl.v_tok_bytes);
7233 let swa_naive = if legacy_tkv { t_kv > win } else { seq_end > win };
7245 if swa && swa_naive {
7246 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
7259 e.sdpa_naive_w_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh,
7260 nkv, t, t_kv, scale, true, win,
7261 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
7262 } else {
7263 e.fa_prefill_view_ws_w_hd128(&q, &k_view, &v_view, &mut attn, hd, nh,
7264 nkv, t, t_kv, scale, true, win,
7265 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
7266 }
7267 } else if std::env::var("MEMRA_NOFA").is_ok() {
7268 e.sdpa_naive_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7269 t, t_kv, scale, true,
7270 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
7271 } else {
7272 e.fa_prefill_view_ws(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7277 t, t_kv, scale, true,
7278 kvl.k_tok_bytes, kvl.v_tok_bytes,
7279 crate::Engine::kv_fp8_on())?;
7280 }
7281 }
7282 None => {
7283 debug_assert_eq!(seq_end, t, "step35 cacheless prefill is monolithic (seq_end == t)");
7288 if swa && seq_end > win {
7289 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
7290 } else if std::env::var("MEMRA_NOFA").is_ok() {
7291 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
7292 } else {
7293 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
7294 }
7295 }
7296 }
7297
7298 let gw = fa.attn_gate.as_ref()
7301 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
7302 let gt = e.matmul(gw, hg, t)?;
7303 let mut ag = e.uninit(t * nh * hd)?;
7304 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, t)?;
7305 Ok(ag)
7306 }
7307
7308 pub(crate) fn step35_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
7311 pos_d: &CudaSlice<i32>, t: usize, il: usize)
7312 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7313 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
7314 let ag = self.step35_attn_pre_wo(e, fa, g3, h, pos_d, t, None, il, t)?;
7316 Ok(e.matmul(&fa.wo, &ag, t)?)
7317 }
7318
7319 #[allow(clippy::too_many_arguments)]
7326 pub(crate) fn step35_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
7327 hx: Option<&CudaSlice<u8>>, pos_d: &CudaSlice<i32>, t: usize,
7328 cache: &mut Cache, il: usize, seq_end: usize)
7329 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7330 let g3 = match hx {
7331 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
7332 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
7333 };
7334 let ag = self.step35_attn_pre_wo(e, fa, g3, h, pos_d, t, Some(cache), il, seq_end)?;
7335 Ok(e.matmul(&fa.wo, &ag, t)?)
7336 }
7337
7338 #[allow(clippy::too_many_arguments)]
7348 pub(crate) fn step35_decode_attn(&self, e: &Engine, fa: &FullAttnLayer, il: usize,
7349 h: &CudaSlice<f32>,
7350 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7351 pos_d: &CudaSlice<i32>, cache: &mut Cache)
7352 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7353 let (hd, nkv, nh, rbase, scale, swa) = self.step35_geom(il);
7354 let eps = self.cfg.rms_eps;
7355 let s = self.cfg.step35.as_ref().unwrap();
7356 let win = s.sliding_window as usize;
7357 let n_rot = s.n_rot(il as u32) as usize;
7358 let n_embd = self.cfg.n_embd as usize;
7359 let gw = fa.attn_gate.as_ref()
7360 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
7361
7362 let (q0, k0, v0, gt) = match pre_q {
7363 Some((hq, hdq)) => {
7364 debug_assert!(e.uses_q8_1_fast(gw),
7365 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
7366 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast");
7367 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
7368 Some(t3) => t3,
7369 None => (e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
7370 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
7371 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?),
7372 };
7373 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
7374 (a, b, c, gt)
7375 }
7376 None => {
7377 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk)
7378 && e.uses_q8_1_fast(&fa.wv) && e.uses_q8_1_fast(gw) {
7379 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
7380 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
7381 Some(t3) => t3,
7382 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7383 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
7384 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
7385 };
7386 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
7387 (a, b, c, gt)
7388 } else {
7389 (e.matmul(&fa.wq, h, 1)?, e.matmul(&fa.wk, h, 1)?,
7390 e.matmul(&fa.wv, h, 1)?, e.matmul(gw, h, 1)?)
7391 }
7392 }
7393 };
7394
7395 let mut q = e.uninit(nh * hd)?;
7396 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
7397 let mut k = e.uninit(nkv * hd)?;
7398 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
7399 let ff = if swa { None } else {
7400 self.step35_aux.as_ref().and_then(|a| a.rope_freqs.as_ref())
7401 };
7402 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
7403
7404 if std::env::var("MEMRA_NOFA").is_ok() {
7405 return Err("MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
7406 cache; unset MEMRA_NOFA to use fa_decode".into());
7407 }
7408 let kvl = cache.kv[il].as_mut().unwrap();
7409 e.append_kv_quantized(&k, &v0, &mut kvl.k, &mut kvl.v, kvl.len,
7410 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
7411 crate::Engine::kv_fp8_on())?;
7412 kvl.len += 1;
7413 let (off, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
7414 let k_view = e.view_u8_range(&kvl.k, off * kvl.k_tok_bytes,
7415 (off + t_kv) * kvl.k_tok_bytes);
7416 let v_view = e.view_u8_range(&kvl.v, off * kvl.v_tok_bytes,
7417 (off + t_kv) * kvl.v_tok_bytes);
7418 let mut attn = e.uninit(nh * hd)?;
7419 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
7420 kvl.k_tok_bytes, kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
7421
7422 let mut ag = e.uninit(nh * hd)?;
7423 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
7424 Ok(e.matmul(&fa.wo, &ag, 1)?)
7425 }
7426}
7427
7428impl HybridModel {
7437 pub fn is_gemma4_e4b(&self) -> bool {
7438 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
7439 }
7440
7441 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
7445 let g = self.cfg.gemma4.as_ref().unwrap();
7446 let swa = g.swa_pattern[il];
7447 let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
7448 let Mixer::Full(fa) = &self.layers[il].mixer else { panic!("e4b layer {il} not full-attn") };
7449 let nh = fa.wq.out_features() / hd;
7450 let nkv = fa.wk.out_features() / hd;
7451 (hd, nkv, nh, if swa { g.rope_base_swa } else { g.rope_base_global }, 1.0, swa)
7452 }
7453
7454 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
7456 self.layers[il].gemma4.as_ref()
7457 .and_then(|b| b.e4b.as_ref())
7458 .and_then(|e4| e4.kv_share.map(|t| t as usize))
7459 }
7460
7461 fn gemma4_e4b_inp_pl(&self, e: &Engine, tokens: &[u32], x_scaled: &CudaSlice<f32>, t: usize)
7466 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7467 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
7468 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
7469 }
7470
7471 fn gemma4_e4b_inp_pl_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
7473 x_scaled: &CudaSlice<f32>, t: usize)
7474 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7475 let aux = self.gemma4_aux.as_ref().unwrap();
7476 let m = aux.e4b.as_ref().unwrap();
7477 let n_embd = self.cfg.n_embd as usize;
7478 let n_layer = self.layers.len();
7479 let width = m.n_epl * n_layer;
7480 let tbl = m.tok_tbl_gpu.get_or_init(|| {
7481 e.upload_u8(&m.tok_embd_bytes).expect("e4b per-layer token table upload")
7482 });
7483 let mut a = e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt,
7484 m.tok_embd_row_bytes)?;
7485 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
7486 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
7487 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
7488 let mut pn = e.uninit(t * width)?;
7489 e.rms_norm(&p, m.proj_norm.float_data(), &mut pn, m.n_epl, t * n_layer,
7490 self.cfg.rms_eps)?;
7491 let mut out = e.uninit(t * width)?;
7492 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
7493 Ok(out)
7494 }
7495
7496 #[allow(clippy::too_many_arguments)]
7501 fn gemma4_e4b_attn(&self, e: &Engine, il: usize,
7502 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
7503 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
7504 dc_bucket: Option<usize>)
7505 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7506 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
7507 let eps = self.cfg.rms_eps;
7508 let aux = self.gemma4_aux.as_ref().unwrap();
7509 let Mixer::Full(fa) = &self.layers[il].mixer else { unreachable!() };
7510 let h0 = e.zeros(0)?;
7514 let h = &h0;
7515
7516 let ff = if swa { None } else {
7517 Some(aux.rope_freqs.as_ref().expect("e4b global rope needs rope_freqs.weight"))
7518 };
7519 let share = self.gemma4_e4b_kv_target(il);
7520 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
7522 let mut q;
7523 if let Some(_tgt) = share {
7524 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
7525 q = e.uninit(t * nh * hd)?;
7526 let mut kdummy = e.uninit(1)?;
7529 let mut vdummy = e.uninit(1)?;
7530 e.rms_norm_qkv_rope(&q0, &q0, &q0, fa.q_norm.float_data(),
7531 fa.q_norm.float_data(), &aux.ones,
7532 &mut q, &mut kdummy, &mut vdummy, hd, nh * t, 0,
7533 pos_d, nh, 1, base, 1.0, ff, eps)?;
7534 } else {
7535 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
7539 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
7540 q = e.uninit(t * nh * hd)?;
7541 let mut k = e.uninit(t * nkv * hd)?;
7542 let mut v = e.uninit(t * nkv * hd)?;
7543 if t == 1 && cat.is_some() {
7544 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
7545 e.rms_norm_qkv_rope_cat(&qkv0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7546 &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7547 pos_d, nh, nkv, base, 1.0, ff, eps)?;
7548 } else {
7549 let (q0, k0, v0) = match if t == 1 {
7550 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
7551 } else {
7552 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7555 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
7556 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
7557 } else { None }
7558 } {
7559 Some(triple) => triple,
7560 None => (e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
7561 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
7562 e.matmul_pre(&fa.wv, hq, hdq, h, t)?), };
7564 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(),
7567 fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v,
7568 hd, nh * t, nkv * t, pos_d, nh, nkv, base, 1.0, ff, eps)?;
7569 }
7570 let kvl = cache.kv[il].as_mut().unwrap();
7571 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7575 if dc_bucket.is_some() {
7576 debug_assert!(t == 1);
7581 e.append_kv_quantized_row_dc_inc(&k, &v, &mut kvl.k, &mut kvl.v,
7583 &mut kvl.len_d, kvl.kv_dim_k, kvl.kv_dim_v,
7584 kvl.k_tok_bytes, kvl.v_tok_bytes, cls)?;
7585 } else {
7586 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
7587 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
7588 kvl.v_tok_bytes, cls)?;
7589 kvl.len += t;
7590 }
7591 kv_f32 = Some((k, v));
7592 }
7593 let kvl_idx = share.unwrap_or(il);
7596 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
7597 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7599 let mut attn = e.uninit(t * nh * hd)?;
7600 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
7612 if let Some((kf, vf)) = &kv_f32 {
7613 if hd == 256 && t <= win {
7614 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
7615 return Ok(e.matmul(&fa.wo, &attn, t)?);
7616 }
7617 if hd == 256 && swa && t > win {
7618 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true,
7619 win)?;
7620 return Ok(e.matmul(&fa.wo, &attn, t)?);
7621 }
7622 if hd == 512 && !swa {
7623 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale,
7624 true)?;
7625 return Ok(e.matmul(&fa.wo, &attn, t)?);
7626 }
7627 } else if share.is_some() {
7628 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7629 let k_view = e.view_u8(&kvl.k, kvl.k.len());
7630 let v_view = e.view_u8(&kvl.v, kvl.v.len());
7631 if hd == 256 && (!swa || t <= win) {
7632 e.fa_prefill_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t, t,
7634 scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
7635 return Ok(e.matmul(&fa.wo, &attn, t)?);
7636 }
7637 let kv_dim = nkv * hd;
7640 let mut kf = e.uninit(t * kv_dim)?;
7641 let mut vf = e.uninit(t * kv_dim)?;
7642 e.fa_dequant_kv_view_f32(&k_view, &v_view, &mut kf, &mut vf, kv_dim, kv_dim,
7643 t, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
7644 if hd == 512 {
7645 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale,
7646 true)?;
7647 } else {
7648 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true,
7649 win)?;
7650 }
7651 return Ok(e.matmul(&fa.wo, &attn, t)?);
7652 }
7653 }
7654 if let Some(bucket) = dc_bucket {
7655 assert!(t == 1);
7660 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
7666 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
7667 } else { bucket };
7668 let k_view = e.view_u8(&kvl.k, kvl.k.len());
7669 let v_view = e.view_u8(&kvl.v, kvl.v.len());
7670 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7671 if crate::Engine::wpf_level() >= 1 {
7679 e.prefetch_weight_l2(&fa.wo)?;
7680 }
7681 if e.uses_q8_1_fast(&fa.wo) {
7684 let mut oq = e.alloc_i8_uninit(nh * hd)?;
7685 let mut od = e.zeros(nh * hd / 32)?;
7686 e.fa_decode_dc_q8(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7687 &kvl.len_d, bucket, scale,
7688 kvl.k_tok_bytes, kvl.v_tok_bytes, g,
7689 Some((&mut oq, &mut od)))?;
7690 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
7691 }
7692 e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7693 &kvl.len_d, bucket, scale,
7694 kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
7695 return Ok(e.matmul(&fa.wo, &attn, t)?);
7696 }
7697 for i in 0..t {
7698 let avail = base_len + i + 1;
7699 let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
7700 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
7701 (off_tok + t_kv) * kvl.k_tok_bytes);
7702 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
7703 (off_tok + t_kv) * kvl.v_tok_bytes);
7704 let qv = e.view(&q, t * nh * hd);
7705 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
7706 let mut q_one = e.uninit(nh * hd)?;
7707 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
7708 let mut a_one = e.uninit(nh * hd)?;
7709 e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
7713 kvl.k_tok_bytes, kvl.v_tok_bytes,
7714 (!swa && crate::Engine::gkv_on())
7715 || (swa && crate::Engine::wkv_on()))?;
7716 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
7717 }
7718 Ok(e.matmul(&fa.wo, &attn, t)?)
7719 }
7720
7721 fn gemma4_e4b_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
7726 head_last: bool)
7727 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7728 let n_embd = self.cfg.n_embd as usize;
7729 let t = tokens.len();
7730 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7731 let pos_d = e.htod_i32(&pos)?;
7732 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7733 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7734 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
7735 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
7736 }
7737
7738 fn gemma4_e4b_trunk_core(&self, e: &Engine, x_in: CudaSlice<f32>, inp_pl: CudaSlice<f32>,
7742 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
7743 dc_bucket: Option<usize>, cap_logits: bool, head_last: bool)
7744 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7745 let n_embd = self.cfg.n_embd as usize;
7746 let eps = self.cfg.rms_eps;
7747 let n_layer = self.layers.len();
7748 let mut x = x_in;
7749 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
7750 let n_epl = aux_e4b.n_epl;
7751
7752 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7758 for il in 0..n_layer {
7759 let layer = &self.layers[il];
7760 let (hq, hdq) = match h_carry.take() {
7761 Some(p) => p,
7762 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
7763 };
7764 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
7765 let bits = layer.gemma4.as_ref().unwrap();
7768 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
7769 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
7780 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
7781 e, layer, &o, &x, t, Some(layer.post_attn_norm.float_data()), fuse_exit)?;
7782 let mut resid = e.uninit(t * n_embd)?;
7783 let g = if fuse_exit {
7789 let (rq, rd) = e.rms_pre_add_q8_1(&sn, bits.post_ffw_norm.float_data(),
7791 &attn_out, &mut resid, n_embd, t,
7792 self.cfg.rms_eps)?;
7793 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
7794 } else {
7795 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
7796 e.matmul(&e4b.inp_gate, &resid, t)?
7797 };
7798 let mut act = e.uninit(t * n_epl)?;
7799 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
7800 let ipv = e.view(&inp_pl, n_epl * n_layer);
7801 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
7802 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
7803 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
7804 } else {
7805 let mut inp_this = e.uninit(t * n_epl)?;
7806 e.copy_rows_strided(&inp_pl, &mut inp_this, n_epl, t, n_epl * n_layer,
7807 il * n_epl)?;
7808 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
7809 e.matmul(&e4b.proj, &act, t)?
7810 };
7811 let next_norm = if il + 1 < n_layer {
7814 self.layers[il + 1].attn_norm.float_data()
7815 } else {
7816 self.output_norm.float_data()
7817 };
7818 let mut xn = e.uninit(t * n_embd)?;
7819 let pair = e.rms_pre_add_scale_rms_norm_q8_1(&y, e4b.post_norm.float_data(),
7820 &resid, bits.layer_scale, next_norm,
7821 &mut xn, n_embd, t, eps)?;
7822 h_carry = Some(pair);
7823 x = xn;
7824 }
7825 let (oq, odq) = h_carry.take().unwrap();
7829 let h0 = e.zeros(0)?;
7830 let hm = if head_last { 1 } else { t };
7831 let (hq, hd) = if head_last && t > 1 {
7832 let mut q1 = e.uninit_i8(n_embd)?;
7833 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
7834 let nb = n_embd / 32;
7835 let mut d1 = e.uninit(nb)?;
7836 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
7837 (q1, d1)
7838 } else {
7839 (oq, odq)
7840 };
7841 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
7842 if cap_logits {
7846 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7847 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
7848 }
7849 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
7851 }
7852
7853 pub fn gemma4_e4b_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
7860 t: usize, pos0: usize, cache: &mut Cache)
7861 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7862 let n_embd = self.cfg.n_embd as usize;
7863 let eps = self.cfg.rms_eps;
7864 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7865 let pos_d = e.htod_i32(&pos)?;
7866 let embd_gpu = self.embd_gpu.get_or_init(|| {
7867 e.upload_u8(&self.embd.raw).expect("embed table upload")
7868 });
7869 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
7870 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
7871 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7872 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
7873 let (ld, xp) = self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true,
7874 false)?;
7875 let n_vocab = self.output.out_features();
7878 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
7879 for i in 0..t {
7880 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
7881 }
7882 let mut hn = e.uninit(t * n_embd)?;
7883 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7884 cache.pos += t;
7885 Ok((vam, hn))
7886 }
7887
7888 pub(crate) fn gemma4_e4b_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
7891 cache: &mut Cache)
7892 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7893 let n_embd = self.cfg.n_embd as usize;
7894 let eps = self.cfg.rms_eps;
7895 let t = tokens.len();
7896 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
7897 let mut hn = e.uninit(t * n_embd)?;
7898 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7899 cache.pos += t;
7900 Ok((e.dtoh(&ld)?, hn))
7901 }
7902
7903 pub fn gemma4_e4b_decode_step_dcg(&self, e: &Engine, token_d: &mut CudaSlice<u32>,
7909 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7910 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7911 n_vocab: usize, bucket: usize)
7912 -> Result<(), Box<dyn std::error::Error>> {
7913 let n_embd = self.cfg.n_embd as usize;
7914 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7915 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7916 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
7917 let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket),
7918 false, false)?;
7919 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
7920 e.inc_seqlen(pos_d)?;
7921 Ok(())
7922 }
7923
7924 #[allow(clippy::too_many_arguments)]
7932 pub fn gemma4_e4b_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
7933 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7934 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7935 n_vocab: usize)
7936 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7937 let n_embd = self.cfg.n_embd as usize;
7938 let eps = self.cfg.rms_eps;
7939 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7940 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7941 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
7942 let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false,
7943 false)?;
7944 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
7945 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
7946 e.inc_seqlen(pos_d)?;
7947 cache.pos += 1;
7948 let _ = eps;
7949 Ok(tok_out)
7950 }
7951
7952 pub(crate) fn gemma4_e4b_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
7955 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7956 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
7957 let logits = e.dtoh(&ld)?;
7958 cache.pos += 1;
7959 Ok((logits, x))
7960 }
7961
7962 pub(crate) fn gemma4_e4b_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
7966 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7967 if cache.pos != 0 {
7970 return Err("e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
7971 call or decode tokenwise".into());
7972 }
7973 let n_embd = self.cfg.n_embd as usize;
7974 let t = tokens.len();
7975 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
7976 cache.pos += t;
7977 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
7979 let row = xv.slice((t - 1) * n_embd..t * n_embd);
7980 let mut h_seed = e.uninit(n_embd)?;
7981 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
7982 Ok((last, h_seed, x))
7983 }
7984
7985 pub(crate) fn gemma4_e4b_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
7987 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7988 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
7989 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
7990 Ok(e.dtoh(&ld)?) }
7992}
7993
7994#[cfg(test)]
7995mod page_prefetch_tests {
7996 use super::{
7997 grouped_worker_prefetch_position, page_prefetch_positions,
7998 page_prefetch_window_from_values, worker_prefetch_positions,
7999 };
8000
8001 #[test]
8002 fn page_prefetch_window_keeps_existing_opt_in_default() {
8003 assert_eq!(page_prefetch_window_from_values(false, None), 0);
8004 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
8005 assert_eq!(page_prefetch_window_from_values(true, None), 1);
8006 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
8007 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
8008 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
8009 }
8010
8011 #[test]
8012 fn rolling_page_prefetch_advises_each_future_expert_once() {
8013 let advised: Vec<_> = (0..7)
8014 .flat_map(|position| page_prefetch_positions(position, 7, 3))
8015 .collect();
8016 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
8017
8018 let one_ahead: Vec<_> = (0..4)
8019 .flat_map(|position| page_prefetch_positions(position, 4, 1))
8020 .collect();
8021 assert_eq!(one_ahead, vec![1, 2, 3]);
8022 assert!(page_prefetch_positions(0, 4, 0).is_empty());
8023 }
8024
8025 #[test]
8026 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
8027 assert_eq!(grouped_worker_prefetch_position(0, None), None);
8028 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
8029 .chain((0..4).filter_map(|position| {
8030 grouped_worker_prefetch_position(4, Some(position))
8031 }))
8032 .collect();
8033 assert_eq!(positions, vec![0, 1, 2, 3]);
8034 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
8035 }
8036
8037 #[test]
8038 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
8039 let queued: Vec<_> = (0..8)
8040 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
8041 .collect();
8042 assert_eq!(queued, (0..8).collect::<Vec<_>>());
8043
8044 let one_at_a_time: Vec<_> = (0..4)
8045 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
8046 .collect();
8047 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
8048 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
8049 }
8050}
8051
8052pub struct G4DcSlots {
8053 x: CudaSlice<f32>, xn: CudaSlice<f32>, cur: CudaSlice<f32>,
8054 hq: CudaSlice<i8>, hd_: CudaSlice<f32>,
8055 q0: CudaSlice<f32>, k0: CudaSlice<f32>, v0: CudaSlice<f32>,
8056 q: CudaSlice<f32>, k: CudaSlice<f32>, v: CudaSlice<f32>,
8057 attn: CudaSlice<f32>, o: CudaSlice<f32>,
8058 attn_out: CudaSlice<f32>, zsh: CudaSlice<f32>,
8059 zq: CudaSlice<i8>, zd: CudaSlice<f32>,
8060 gate: CudaSlice<f32>, up: CudaSlice<f32>,
8061 act: CudaSlice<f32>, actq: CudaSlice<i8>, actd: CudaSlice<f32>,
8062 f0: CudaSlice<f32>, sn: CudaSlice<f32>,
8063 hn: CudaSlice<f32>, logits: CudaSlice<f32>,
8064}