1use cudarc::driver::CudaSlice;
6use memra_gguf::config::ModelConfig;
7use crate::Engine;
8use crate::cache::Cache;
9
10pub struct PrimeSlabs {
13 pub t_cap: usize,
14 pub h: CudaSlice<f32>,
15 pub x1: CudaSlice<f32>,
16 pub z: CudaSlice<f32>,
17 pub act: CudaSlice<f32>,
18 pub xa: CudaSlice<f32>,
19 pub xb: CudaSlice<f32>,
20 pub h16: CudaSlice<u8>,
21 pub z16: CudaSlice<u8>,
22 pub gate: CudaSlice<f32>, pub up: CudaSlice<f32>, pub ffn_out: CudaSlice<f32>, pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
32 pub mixed: CudaSlice<f32>,
36 pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
37 pub seg_t: usize,
38}
39
40unsafe impl Send for PrimeSlabs {}
43
44fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
45 std::iter::repeat_with(|| None).take(n).collect()
46}
47
48struct PrimeCacheStages<'a> {
53 parent: &'a mut Cache,
54 cut: usize,
55 stage0: Cache,
56 stage1: Cache,
57}
58
59impl<'a> PrimeCacheStages<'a> {
60 fn new(parent: &'a mut Cache, cut: usize) -> Self {
61 let n = parent.kv.len();
62 assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
63 assert!(cut <= n, "PP-2 cache cut {cut} exceeds {n} layers");
64 let mut kv0 = empty_cache_layers(n);
65 let mut kv1 = empty_cache_layers(n);
66 let mut recur0 = empty_cache_layers(n);
67 let mut recur1 = empty_cache_layers(n);
68 for i in 0..cut {
69 kv0[i] = parent.kv[i].take();
70 recur0[i] = parent.recur[i].take();
71 }
72 for i in cut..n {
73 kv1[i] = parent.kv[i].take();
74 recur1[i] = parent.recur[i].take();
75 }
76 let pos = parent.pos;
77 let max_ctx = parent.max_ctx;
78 Self {
79 parent,
80 cut,
81 stage0: Cache {
82 kv: kv0,
83 recur: recur0,
84 pos,
85 max_ctx,
86 last_logits_dev: None,
87 dflash_taps: None,
88 },
89 stage1: Cache {
90 kv: kv1,
91 recur: recur1,
92 pos,
93 max_ctx,
94 last_logits_dev: None,
95 dflash_taps: None,
96 },
97 }
98 }
99
100 fn parts(&mut self) -> (&mut Cache, &mut Cache) {
101 (&mut self.stage0, &mut self.stage1)
102 }
103}
104
105impl Drop for PrimeCacheStages<'_> {
106 fn drop(&mut self) {
107 let n = self.parent.kv.len();
108 for i in 0..n {
109 let source = if i < self.cut {
110 &mut self.stage0
111 } else {
112 &mut self.stage1
113 };
114 debug_assert!(self.parent.kv[i].is_none());
115 debug_assert!(self.parent.recur[i].is_none());
116 self.parent.kv[i] = source.kv[i].take();
117 self.parent.recur[i] = source.recur[i].take();
118 }
119 self.parent.pos = self.stage0.pos.min(self.stage1.pos);
120 }
121}
122
123
124pub(crate) struct AttnPre {
126 pub q: cudarc::driver::CudaSlice<f32>,
127 pub k: cudarc::driver::CudaSlice<f32>,
128 pub v: cudarc::driver::CudaSlice<f32>,
129 pub gate: Option<cudarc::driver::CudaSlice<f32>>,
130}
131
132pub(crate) struct GdnPrep {
134 pub hk: usize,
135 pub q_l2: cudarc::driver::CudaSlice<f32>,
136 pub k_l2: cudarc::driver::CudaSlice<f32>,
137 pub v_g: cudarc::driver::CudaSlice<f32>,
138 pub beta: cudarc::driver::CudaSlice<f32>,
139 pub g_log: cudarc::driver::CudaSlice<f32>,
140 pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
141 pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
142}
143
144pub(crate) struct VerifyStreamScratch {
146 pub pos_d: CudaSlice<i32>,
147 pub row_ctrs: Vec<CudaSlice<i32>>,
148}
149use crate::hybrid::{HybridModel, Mixer, FullAttnLayer, LinearAttnLayer, MoeWeights};
150
151struct MoeInputTraceWriter {
152 dir: std::path::PathBuf,
153 index: std::fs::File,
154 payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
155}
156
157static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<
158 std::sync::Mutex<Option<MoeInputTraceWriter>>,
159> = std::sync::OnceLock::new();
160
161fn gdec_enabled() -> bool {
164 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
165 *E.get_or_init(|| std::env::var("MEMRA_MOE_GDEC").map(|v| v != "0").unwrap_or(true))
166}
167
168fn moe_slab_enabled() -> bool {
179 std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
180}
181
182fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
186 std::env::var("MEMRA_MOE_GROUPED")
187 .map(|value| value != "0")
188 .unwrap_or(false)
189}
190
191fn moe_prefetch_enabled() -> bool {
194 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
195 *E.get_or_init(|| std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
196 || crate::spill_pread::worker_enabled())
197}
198
199fn moe_page_prefetch_window() -> usize {
204 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
205 *W.get_or_init(|| page_prefetch_window_from_values(
206 std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
207 std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW").ok().as_deref(),
208 ))
209}
210
211fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
212 if !enabled {
213 return 0;
214 }
215 raw_window
216 .and_then(|value| value.parse().ok())
217 .unwrap_or(1)
218}
219
220fn page_prefetch_positions(
224 position: usize,
225 len: usize,
226 window: usize,
227) -> std::ops::Range<usize> {
228 if window == 0 || position >= len {
229 return len..len;
230 }
231 let (start, count) = if position == 0 {
232 (1, window)
233 } else {
234 (position.saturating_add(window), 1)
235 };
236 let start = start.min(len);
237 start..start.saturating_add(count).min(len)
238}
239
240fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
243 let position = current.map_or(0, |position| position.saturating_add(1));
244 (position < order_len).then_some(position)
245}
246
247fn worker_prefetch_window() -> usize {
252 static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
253 *WINDOW.get_or_init(|| {
254 let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
255 std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
256 .ok()
257 .and_then(|value| value.parse::<usize>().ok())
258 .unwrap_or(automatic.max(1))
259 })
260}
261
262fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
266 if window == 0 || position >= len {
267 return len..len;
268 }
269 let (start, count) = if position == 0 {
270 (0, window)
271 } else {
272 (position.saturating_add(window).saturating_sub(1), 1)
273 };
274 let start = start.min(len);
275 start..start.saturating_add(count).min(len)
276}
277
278fn moe_dev_enabled() -> bool {
283 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
284 *E.get_or_init(|| std::env::var("MEMRA_MOE_DEV").map(|v| v != "0").unwrap_or(true)
285 && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")))
286}
287
288fn sigmoid_router_enabled() -> bool {
291 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
292 *E.get_or_init(|| std::env::var("MEMRA_SIG_ROUTER").map(|v| v != "0").unwrap_or(true))
293}
294
295fn moe_q8_enabled() -> bool {
300 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
301 *E.get_or_init(|| std::env::var("MEMRA_MOE_Q8").map(|v| v != "0").unwrap_or(true))
302}
303
304fn expert_dp4a_supported(qt: i32) -> bool {
307 qt == crate::QT_Q4_0 || qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS
308 || qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K
309}
310
311fn q8_expert_supported(qt: i32) -> bool {
312 static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
318 let kq = *KQ.get_or_init(|| {
319 std::env::var("MEMRA_MOE_Q8_KQ").map(|v| v != "0").unwrap_or(true)
320 });
321 let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4").map(|v| v != "0").unwrap_or(true);
328 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || (nvfp4_q8 && qt == crate::QT_NVFP4)
329 || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
330}
331
332fn q8_expert_dec_supported(qt: i32) -> bool {
335 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
336}
337
338fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
344 match qt {
345 crate::QT_Q4_0 => in_f % 32 == 0,
346 crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K
347 | crate::QT_Q6_K => in_f % 256 == 0,
348 _ => false,
349 }
350}
351
352fn moe_prewarm_enabled() -> bool {
355 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
356 *E.get_or_init(|| std::env::var("MEMRA_MOE_PREWARM").map(|v| v != "0").unwrap_or(true))
357}
358
359fn cpu_expert_profile_admit_enabled() -> bool {
363 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
364 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
365}
366
367pub const PRIME_MIN_T: usize = 16;
371const PRIME_PIPE_MICROBATCHES: usize = 8;
372const PRIME_PIPE_MIN_CHUNK: usize = 128;
373const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
374const PRIME_PIPE_LINEAR_WORK: usize = 8;
375
376fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
377 crate::pp::prime_pp_on()
378 && !crate::pp::pp2_streams_off()
379 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
380}
381
382pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
386 if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
387 let parsed = value.parse::<usize>().unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
388 return if crate::cache::swa_ring_on() {
389 if parsed == 0 {
390 crate::cache::PRIME_CHUNK_MAX_TOKENS
391 } else {
392 parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
393 }
394 } else {
395 parsed
396 };
397 }
398 let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
399 if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
400 chunk.min(
401 t.div_ceil(PRIME_PIPE_MICROBATCHES)
402 .max(PRIME_PIPE_MIN_CHUNK),
403 )
404 } else {
405 chunk
406 }
407}
408
409fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
410 fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
411}
412
413fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
414 if chunk == 0 || t <= chunk {
415 return vec![(0, t)];
416 }
417 let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
418 let mut start = 0usize;
419 while start < t {
420 let mut end = (start + chunk).min(t);
421 if t - end > 0 && t - end < PRIME_MIN_T {
422 if ring_on {
423 let shifted = t - PRIME_MIN_T;
424 end = if shifted > start { shifted } else { t };
425 } else {
426 end = t;
427 }
428 }
429 ranges.push((start, end));
430 start = end;
431 }
432 ranges
433}
434
435fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
436 let prefix = prefix as u128;
437 prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
438}
439
440fn dynamic_prime_chunk_ranges(
441 t: usize,
442 fixed_chunk: usize,
443 fixed: &[(usize, usize)],
444) -> Vec<(usize, usize)> {
445 let n = fixed.len();
446 if n < 3 {
447 return fixed.to_vec();
448 }
449
450 let max_first = t - (n - 1) * PRIME_MIN_T;
451 let first = fixed_chunk
452 .div_ceil(2)
453 .max(PRIME_PIPE_EDGE_MIN_CHUNK)
454 .min(max_first);
455 let mut ranges = Vec::with_capacity(n);
456 ranges.push((0, first));
457
458 let first_work = prime_chunk_work(first, t);
459 let work_span = prime_chunk_work(t, t) - first_work;
460 let denominator = (n - 1) as u128;
461 let mut previous = first;
462 for boundary in 1..n - 1 {
463 let target = first_work * denominator + work_span * (boundary as u128);
464 let remaining = n - 1 - boundary;
465 let mut low = previous + PRIME_MIN_T;
466 let mut high = t - remaining * PRIME_MIN_T;
467 while low < high {
468 let mid = low + (high - low) / 2;
469 if prime_chunk_work(mid, t) * denominator >= target {
470 high = mid;
471 } else {
472 low = mid + 1;
473 }
474 }
475 ranges.push((previous, low));
476 previous = low;
477 }
478 ranges.push((previous, t));
479 ranges
480}
481
482pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
486 let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
487 let chunk = prime_chunk_tokens(t, n_layers);
488 let fixed = fixed_prime_chunk_ranges(t, chunk);
489 let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
490 Ok(value) => value == "dynamic",
491 Err(_) => true,
492 };
493 if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
494 fixed
495 } else {
496 dynamic_prime_chunk_ranges(t, chunk, &fixed)
497 }
498}
499
500impl HybridModel {
501 fn prime_trace_path() -> Option<&'static str> {
506 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
507 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
508 .as_deref()
509 }
510
511 pub fn forward(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
513 if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, false); }
514 if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, false); }
515 let cfg = &self.cfg;
516 let n_embd = cfg.n_embd as usize;
517 let t = tokens.len();
518 let eps = cfg.rms_eps;
519 let pos: Vec<i32> = (0..t as i32).collect();
520 let pos_d = e.htod_i32(&pos)?;
521
522 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
525 let mut h = e.uninit(t * n_embd)?;
527 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
528
529 let mixed = match &layer.mixer {
530 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
531 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
532 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
533 };
534
535 let mut x1 = e.uninit(t * n_embd)?;
537 e.add(&x, &mixed, &mut x1, t * n_embd)?;
538
539 let mut z = e.uninit(t * n_embd)?;
541 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
542 let ffn_out = match &layer.ffn {
543 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
544 let n_ff = ffn_gate.out_features();
545 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
546 let up = g2.pop().unwrap();
547 let gate = g2.pop().unwrap();
548 let mut act = e.uninit(t * n_ff)?;
549 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
554 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
555 e.matmul(ffn_down, &act, t)?
556 }
557 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
558 };
559 let mut x2 = e.uninit(t * n_embd)?;
560 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
561 x = x2;
562 }
563
564 let mut hn = e.uninit(t * n_embd)?;
565 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
566 let logits = e.matmul(&self.output, &hn, t)?;
567 Ok(e.dtoh(&logits)?)
568 }
569
570 pub fn forward_last(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
576 if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, true); }
577 let cfg = &self.cfg;
578 let n_embd = cfg.n_embd as usize;
579 let t = tokens.len();
580 let eps = cfg.rms_eps;
581 let pos: Vec<i32> = (0..t as i32).collect();
582 let pos_d = e.htod_i32(&pos)?;
583
584 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
588 for (il, layer) in self.layers.iter().enumerate() {
589 let mut h = e.uninit(t * n_embd)?;
590 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
591 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} norm ok"); }
592 let mixed = match &layer.mixer {
593 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
594 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
595 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
596 };
597 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} mixer ok"); }
598 let mut x1 = e.uninit(t * n_embd)?;
599 e.add(&x, &mixed, &mut x1, t * n_embd)?;
600 let mut z = e.uninit(t * n_embd)?;
601 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
602 let ffn_out = match &layer.ffn {
603 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
604 let n_ff = ffn_gate.out_features();
605 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
606 let up = g2.pop().unwrap();
607 let gate = g2.pop().unwrap();
608 let mut act = e.uninit(t * n_ff)?;
609 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
611 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
612 e.matmul(ffn_down, &act, t)?
613 }
614 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
615 };
616 if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} ffn ok"); }
617 let mut x2 = e.uninit(t * n_embd)?;
618 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
619 x = x2;
620 }
621 let mut hn = e.uninit(t * n_embd)?;
623 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
624 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)?;
627 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
628 let logits = e.matmul(&self.output, &hlast, 1)?; Ok(e.dtoh(&logits)?)
630 }
631
632 pub fn prime_cache(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize)
664 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
665 let n_embd = self.cfg.n_embd as usize;
666 let t = tokens.len();
667 assert!(t >= PRIME_MIN_T, "prime_cache needs T >= {PRIME_MIN_T} (caller gates)");
671 assert!(cache.pos + t <= cache.max_ctx, "prime_cache: prompt exceeds cache max_ctx");
672
673 if self.is_gemma4_e4b() {
685 return self.gemma4_e4b_prime(e, tokens, cache);
686 }
687 if self.cfg.gemma4.is_some() {
688 return self.gemma4_prime(e, tokens, cache);
690 }
691 let ranges = prime_chunk_ranges(t, self.layers.len());
692 let legacy_calllocal =
728 std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
729 let seq_end = if legacy_calllocal {
730 cache.pos + t
731 } else {
732 cache.pos + t + queued_after
733 };
734 if ranges.len() == 1 {
735 return self.prime_chunk(e, tokens, cache, seq_end);
736 }
737 if crate::pp::prime_pipe_on()
742 && crate::pp::prime_pp_on()
743 && !crate::pp::pp2_streams_off()
744 {
745 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
746 if crate::pp::pp_multi_stream_same_device() {
747 return Err(
748 "prime chunk pipeline refused with 2 stage streams on one device — \
749 that concurrent-stream placement remains quarantined by the deferred \
750 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
751 the serial split."
752 .into(),
753 );
754 }
755 return self.prime_cache_pp2_pipelined(
756 e, tokens, cache, seq_end, &ranges, &fence,
757 );
758 }
759 }
760 let mut hiddens = e.uninit(t * n_embd)?;
761 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
762 for &(start, end) in &ranges {
763 let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
764 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
765 last = Some((l, hs));
766 }
767 let (logits, h_seed) = last.unwrap();
768 Ok((logits, h_seed, hiddens))
769 }
770
771 fn prime_cache_pp2_pipelined(
776 &self,
777 e: &Engine,
778 tokens: &[u32],
779 cache: &mut Cache,
780 seq_end: usize,
781 ranges: &[(usize, usize)],
782 fence: &[usize],
783 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
784 debug_assert_eq!(fence.len(), 3);
785 debug_assert!(ranges.len() >= 2);
786 let rt = crate::pp::PpNRt::get(e)?;
787 assert_eq!(rt.n_stages(), 2, "prime pipeline requires exactly two PP stages");
788 let n_embd = self.cfg.n_embd as usize;
789 let t = tokens.len();
790 let initial_base = cache.pos;
791 let caller_stream = e.stream();
792
793 rt.fence_stages_behind(&caller_stream)?;
798 let max_payload = ranges
799 .iter()
800 .map(|(s, e)| (e - s) * n_embd)
801 .max()
802 .unwrap();
803 rt.prepare_overlap_slots(0, max_payload)?;
804
805 let mut hiddens = e.uninit(t * n_embd)?;
806 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
807 let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
808 let (cache0, cache1) = stage_caches.parts();
809 let (first_start, first_end) = ranges[0];
810 let mut slot = self.prime_pp2_stage0_enqueue(
811 e,
812 rt,
813 &tokens[first_start..first_end],
814 cache0,
815 seq_end,
816 fence,
817 initial_base + first_start,
818 true,
819 )?;
820 cache0.pos = initial_base + first_end;
821
822 for (i, &(start, end)) in ranges.iter().enumerate() {
823 let base = initial_base + start;
824 debug_assert_eq!(
825 cache1.pos, base,
826 "stage 1 must drain chunks in original position order"
827 );
828 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
829 let next_base = initial_base + next_start;
830 debug_assert_eq!(
831 cache0.pos, next_base,
832 "stage 0 must issue chunks in original position order"
833 );
834 let cache0_stage = &mut *cache0;
835 std::thread::scope(
840 |scope| -> Result<_, Box<dyn std::error::Error>> {
841 let stage0 = scope.spawn(move || -> Result<usize, String> {
842 let next = self
843 .prime_pp2_stage0_enqueue(
844 e,
845 rt,
846 &tokens[next_start..next_end],
847 cache0_stage,
848 seq_end,
849 fence,
850 next_base,
851 true,
852 )
853 .map_err(|err| err.to_string())?;
854 cache0_stage.pos = initial_base + next_end;
855 Ok(next)
856 });
857 let x = self.prime_pp2_stage1_enqueue(
858 e,
859 rt,
860 slot,
861 end - start,
862 cache1,
863 seq_end,
864 fence,
865 base,
866 true,
867 )?;
868 let out = {
869 rt.bind_stage(1)?;
870 let _st1 = rt.enter(1);
871 let e1 = rt.engine(1, e);
872 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
873 };
874 let next = stage0
875 .join()
876 .map_err(|_| "pipeprime stage-0 host walker panicked")?
877 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
878 Ok((out, Some(next)))
879 },
880 )?
881 } else {
882 let x = self.prime_pp2_stage1_enqueue(
883 e,
884 rt,
885 slot,
886 end - start,
887 cache1,
888 seq_end,
889 fence,
890 base,
891 true,
892 )?;
893 let out = {
894 rt.bind_stage(1)?;
895 let _st1 = rt.enter(1);
896 let e1 = rt.engine(1, e);
897 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
898 };
899 (out, None)
900 };
901
902 rt.publish_to(1, &caller_stream)?;
903 e.copy_into(
904 &mut hiddens,
905 start * n_embd,
906 &out.2,
907 (end - start) * n_embd,
908 )?;
909 last = Some((out.0, out.1));
910 crate::pp::PRIME_SPLIT_CHUNKS
911 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
912
913 if let Some(next) = next_slot {
914 rt.fence_stages_behind(&caller_stream)?;
919 slot = next;
920 }
921 }
922
923 debug_assert_eq!(cache0.pos, initial_base + t);
924 debug_assert_eq!(cache1.pos, initial_base + t);
925 let (logits, h_seed) = last.unwrap();
926 Ok((logits, h_seed, hiddens))
927 }
928
929 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
936 if Engine::gdn_db_on()
937 && Engine::gdn_chunked_enabled() && t >= 16
938 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
939 && num_k * 2 == num_v
940 {
941 num_k
942 } else {
943 num_v
944 }
945 }
946
947 fn f16out_on(e: &Engine, t: usize) -> bool {
952 crate::f16_ffi::pp_f16_enabled() && t >= 16 && !e.verify_exact_on()
953 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
954 }
955
956 pub fn prime_slabs_get(
964 &self,
965 e: &Engine,
966 t: usize,
967 n_embd: usize,
968 n_ff_max: usize,
969 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
970 let mut slabs = self.prime_slabs.lock().unwrap();
971 let dev = e.ctx().ordinal();
972 let need_new = match slabs.get(&dev) {
973 None => true,
974 Some(sl) => sl.lock().unwrap().t_cap < t,
975 };
976 if need_new {
977 slabs.insert(dev, std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
978 t_cap: t,
979 h: e.uninit(t * n_embd)?,
980 x1: e.uninit(t * n_embd)?,
981 z: e.uninit(t * n_embd)?,
982 act: e.uninit(t * n_ff_max)?,
983 xa: e.uninit(t * n_embd)?,
984 xb: e.uninit(t * n_embd)?,
985 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
986 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
987 gate: e.uninit(t * n_ff_max)?,
988 up: e.uninit(t * n_ff_max)?,
989 ffn_out: e.uninit(t * n_embd)?,
990 seg_glue: Vec::new(),
991 mixed: e.uninit(t * n_embd)?,
992 seg_mid: Vec::new(),
993 seg_t: 0,
994 })));
995 }
996 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
997 }
998
999 fn prime_chunk(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize)
1003 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1004 if crate::pp::pp_host_bounce_active()
1005 && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1006 {
1007 return Err(
1008 "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1009 has no active prime stage split and would peer-read remote weights; keep \
1010 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1011 .into(),
1012 );
1013 }
1014 if self.cfg.gemma4.is_none()
1023 && !crate::pp::pp2_streams_off()
1024 && crate::pp::prime_pp_on()
1025 {
1026 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1027 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1028 }
1029 }
1030 if crate::pp::pp_host_bounce_active() {
1031 return Err(
1032 "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1033 refusing an unsplit remote-weight walk"
1034 .into(),
1035 );
1036 }
1037 let t = tokens.len();
1038 let base = cache.pos;
1039 debug_assert!(seq_end >= base + t, "prime_chunk: seq_end must cover this chunk");
1040 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1041 let pos_d = e.htod_i32(&pos)?;
1042
1043 let x_embed = self.embed(e, tokens)?; let x = self.prime_layers(
1045 e, x_embed, 0, self.layers.len(), &pos_d, t, base, cache, seq_end,
1046 )?;
1047 self.prime_chunk_epilogue(e, x, t, cache)
1048 }
1049
1050 #[allow(clippy::too_many_arguments)]
1066 fn prime_layers(&self, e: &Engine, x_in: CudaSlice<f32>, lo: usize, hi: usize,
1067 pos_d: &CudaSlice<i32>, t: usize, base: usize, cache: &mut Cache,
1068 seq_end: usize)
1069 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1070 let cfg = &self.cfg;
1071 let n_embd = cfg.n_embd as usize;
1072 let eps = cfg.rms_eps;
1073 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1077 let n_ff_max = self.layers.iter().map(|l| match &l.ffn {
1083 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1084 _ => n_embd,
1085 }).max().unwrap_or(n_embd).max(n_embd);
1086 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1087 let slab = if use_slabs {
1088 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1089 } else {
1090 None
1091 };
1092 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1093 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>);
1095 let (mut x_cur, mut x_nxt, sl): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, Option<SlabRefs>);
1096 let mut seg: Option<(&mut Vec<Option<cudarc::driver::CudaGraph>>, &mut Vec<Option<cudarc::driver::CudaGraph>>, &mut CudaSlice<f32>, &mut usize)> = None;
1097 let mut x_own2;
1098 match slab_guard.as_mut() {
1099 Some(g) => {
1100 let slabs = &mut **g;
1101 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1102 let PrimeSlabs { xa, xb, h, x1, z, act, h16, z16, gate, up, ffn_out, seg_glue, mixed, seg_mid, seg_t, .. } = slabs;
1103 x_cur = xa;
1104 x_nxt = xb;
1105 seg = Some((seg_glue, seg_mid, mixed, seg_t));
1106 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1107 }
1108 None => {
1109 x_own = x_in;
1110 x_own2 = e.uninit(t * n_embd)?;
1111 x_cur = &mut x_own;
1112 x_nxt = &mut x_own2;
1113 sl = None;
1114 }
1115 }
1116 let mut alloc_h; let mut alloc_x1; let mut alloc_z; let mut alloc_act;
1117 let mut alloc_h16; let mut alloc_z16;
1118 let mut alloc_gate; let mut alloc_up; let mut alloc_fo;
1119 let (h, x1, z, act): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
1120 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1121 let (sl_gate, sl_up, sl_fo): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
1122 match sl {
1123 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1124 h = a; x1 = b; z = c; act = d; h16 = e16; z16 = f16b;
1125 sl_gate = g; sl_up = u; sl_fo = fo;
1126 }
1127 None => {
1128 alloc_h = e.uninit(t * n_embd)?;
1129 alloc_x1 = e.uninit(t * n_embd)?;
1130 alloc_z = e.uninit(t * n_embd)?;
1131 alloc_act = e.uninit(t * n_ff_max)?;
1132 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1133 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1134 alloc_gate = e.uninit(t * n_ff_max)?;
1135 alloc_up = e.uninit(t * n_ff_max)?;
1136 alloc_fo = e.uninit(t * n_embd)?;
1137 h = &mut alloc_h; x1 = &mut alloc_x1; z = &mut alloc_z; act = &mut alloc_act;
1138 h16 = &mut alloc_h16; z16 = &mut alloc_z16;
1139 sl_gate = &mut alloc_gate; sl_up = &mut alloc_up; sl_fo = &mut alloc_fo;
1140 }
1141 }
1142 let n_layers = self.layers.len();
1147 let use_seg = f16fuse && seg.is_some() && self.cfg.step35.is_none()
1157 && lo == 0 && hi == n_layers
1158 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1159 if let Some((sg, sm, _, st)) = seg.as_mut() {
1160 if **st != t {
1161 sg.clear();
1162 sg.extend((0..n_layers).map(|_| None));
1163 sm.clear();
1164 sm.extend((0..n_layers).map(|_| None));
1165 **st = t;
1166 }
1167 }
1168 {
1169 let layer_lo = &self.layers[lo];
1170 if f16fuse {
1171 e.rms_norm_f16out(x_cur, layer_lo.attn_norm.float_data(), h, h16, n_embd, t, eps)?;
1172 } else {
1173 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1174 }
1175 }
1176 for il in lo..hi {
1177 let layer = &self.layers[il];
1178 let hx16 = if f16fuse { Some(&*h16) } else { None };
1179 if use_seg {
1180 let (pre, pre16, w_out) = match &layer.mixer {
1183 Mixer::Full(fa) => {
1184 let g3 = match hx16 {
1185 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1186 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1187 };
1188 let (pre, pre16) = self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1189 (pre, pre16, &fa.wo)
1190 }
1191 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1192 Mixer::Linear(la) => {
1193 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1194 let g4 = match hx16 {
1195 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1196 None => e.matmul_group(&ws, h, t)?,
1197 };
1198 let (pre, pre16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1199 (pre, pre16, &la.ssm_out)
1200 }
1201 };
1202 {
1203 let (_, sm, mslab, _) = seg.as_mut().unwrap();
1204 let pre_n = pre.len() / t;
1205 let xh_pre = match pre16 {
1206 Some(x) => x,
1207 None => e.f16_act(&pre, t * pre_n, pre_n)?,
1208 };
1209 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1210 let y = e.matmul(w_out, &pre, t)?;
1211 e.copy_into(mslab, 0, &y, t * n_embd)?;
1212 }
1213 if sm[il].is_none() {
1214 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1215 let w_post = layer.post_attn_norm.float_data();
1216 e.stream().synchronize()?;
1217 e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1218 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1219 e.add(x_cur, mslab, x1, t * n_embd)?;
1220 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1221 Ok(())
1222 })();
1223 let g = e.stream().end_capture(
1224 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1225 r?;
1226 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1227 }
1228 sm[il].as_ref().unwrap().launch()?;
1229 }
1230 } else {
1231 let mixed = match &layer.mixer {
1232 Mixer::Full(fa) => self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il,
1233 seq_end)?,
1234 Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1235 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1236 };
1237 if f16fuse {
1238 e.add_rms_norm_f16out(x_cur, &mixed, layer.post_attn_norm.float_data(),
1241 x1, z, z16, n_embd, t, eps)?;
1242 } else {
1243 e.add(x_cur, &mixed, x1, t * n_embd)?;
1244 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1245 }
1246 }
1247 let zx16 = if f16fuse { Some(&*z16) } else { None };
1248 match &layer.ffn {
1249 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1250 let n_ff = ffn_gate.out_features();
1251 let mut into_ok = false;
1254 if let Some(xh) = zx16 {
1255 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1256 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1257 }
1258 if !into_ok {
1259 let mut g2 = match zx16 {
1260 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1261 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1262 };
1263 let up_y = g2.pop().unwrap();
1264 let gate_y = g2.pop().unwrap();
1265 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1266 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1267 }
1268 let d_lim = self.cfg.clamp_shexp_at(il as u32);
1273 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none()
1274 && d_lim.is_none() {
1275 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1276 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1277 Some(a16)
1278 } else {
1279 Self::ffn_act_lim(e, &self.cfg, sl_gate, sl_up, 1.0, 1.0, d_lim,
1280 act, t * n_ff)?;
1281 None
1282 };
1283 let xh_act = match act16 {
1285 Some(x) => x,
1286 None => e.f16_act(act, t * n_ff, n_ff)?,
1287 };
1288 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1289 let y = e.matmul(ffn_down, &*act, t)?;
1290 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1291 }
1292 }
1293 crate::hybrid::Ffn::Moe(m) => {
1294 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1295 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1296 }
1297 }
1298 if use_seg && il + 1 < hi {
1299 let w_next = self.layers[il + 1].attn_norm.float_data();
1301 let (sg, _, _, _) = seg.as_mut().unwrap();
1302 if sg[il].is_none() {
1303 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1304 e.stream().synchronize()?;
1305 e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1306 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1307 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1308 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1309 Ok(())
1310 })();
1311 let g = e.stream().end_capture(
1312 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1313 r?;
1314 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1315 }
1316 sg[il].as_ref().unwrap().launch()?;
1317 } else {
1318 if il + 1 < hi {
1319 let w_next = self.layers[il + 1].attn_norm.float_data();
1320 if f16fuse {
1321 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1322 } else {
1323 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1324 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1325 }
1326 } else {
1327 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1328 }
1329 }
1330 if let Some(path) = Self::prime_trace_path() {
1336 let row = (base + t - 1) as usize;
1337 let host = e.dtoh(x_nxt)?;
1338 let last = &host[(t - 1) * n_embd..t * n_embd];
1339 use std::io::Write as _;
1340 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
1341 let mut h64: u64 = 0xcbf29ce484222325;
1342 for v in last {
1343 h64 ^= v.to_bits() as u64;
1344 h64 = h64.wrapping_mul(0x100000001b3);
1345 }
1346 writeln!(f, "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1347 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1348 last[0], last[1], last[2])?;
1349 }
1350 std::mem::swap(&mut x_cur, &mut x_nxt);
1351 }
1352 let mut x = e.uninit(t * n_embd)?;
1354 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1355 drop(slab_guard);
1356 Ok(x)
1357 }
1358
1359 fn prime_chunk_epilogue(&self, e: &Engine, x: CudaSlice<f32>, t: usize, cache: &mut Cache)
1364 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1365 let n_embd = self.cfg.n_embd as usize;
1366 let eps = self.cfg.rms_eps;
1367 let mut h_seed = e.uninit(n_embd)?;
1371 if !crate::spec::spec_hpost() {
1372 e.copy_view_into(&mut h_seed, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
1373 }
1374 let mut hn = e.uninit(t * n_embd)?;
1376 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1377 if crate::spec::spec_hpost() {
1378 e.copy_view_into(&mut h_seed, 0, &hn.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
1379 }
1380 let last = e.view(&hn, t * n_embd);
1381 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1382 let mut hlast = e.uninit(n_embd)?;
1383 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1384 let logits = e.matmul(&self.output, &hlast, 1)?;
1385 cache.pos += t;
1386 Ok((e.dtoh(&logits)?, h_seed, if crate::spec::spec_hpost() { hn } else { x }))
1389 }
1390
1391 fn prime_chunk_ppn(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize,
1415 fence: &[usize])
1416 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1417 let rt = crate::pp::PpNRt::get(e)?;
1418 let n_st = fence.len() - 1;
1419 assert_eq!(
1420 rt.n_stages(), n_st,
1421 "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
1422 );
1423 let n_embd = self.cfg.n_embd as usize;
1424 let t = tokens.len();
1425 let base = cache.pos;
1426 debug_assert!(seq_end >= base + t, "prime_chunk_ppn: seq_end must cover this chunk");
1427 let payload = t * n_embd;
1428 let caller_stream = e.stream();
1432 rt.fence_stages_behind(&caller_stream)?;
1433
1434 if n_st == 2 {
1435 let slot = self.prime_pp2_stage0_enqueue(
1436 e, rt, tokens, cache, seq_end, fence, base, false,
1437 )?;
1438 let x = self.prime_pp2_stage1_enqueue(
1439 e, rt, slot, t, cache, seq_end, fence, base, false,
1440 )?;
1441 let out = {
1442 rt.bind_stage(1)?;
1443 let _st1 = rt.enter(1);
1444 let e1 = rt.engine(1, e);
1445 self.prime_chunk_epilogue(e1, x, t, cache)?
1446 };
1447 rt.publish_to(1, &caller_stream)?;
1448 crate::pp::PRIME_SPLIT_CHUNKS
1449 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1450 return Ok(out);
1451 }
1452
1453 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1454
1455 let mut slot = {
1457 let _st0 = rt.enter(0);
1458 let e0 = rt.engine(0, e);
1459 let pos_d = e0.htod_i32(&pos)?;
1460 let x = self.embed(e0, tokens)?;
1461 let x = self.prime_layers(
1462 e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end,
1463 )?;
1464 rt.tx(0, &x, payload)?
1465 };
1467
1468 for s in 1..n_st - 1 {
1470 let _st = rt.enter(s);
1471 let es = rt.engine(s, e);
1472 let pos_d = es.htod_i32(&pos)?;
1473 let x = rt.rx(s - 1, slot, payload)?;
1474 let x = self.prime_layers(
1475 es, x, fence[s], fence[s + 1], &pos_d, t, base, cache, seq_end,
1476 )?;
1477 slot = rt.tx(s, &x, payload)?;
1478 }
1479
1480 let _stl = rt.enter(n_st - 1);
1482 let el = rt.engine(n_st - 1, e);
1483 let pos_d = el.htod_i32(&pos)?;
1484 let x = rt.rx(n_st - 2, slot, payload)?;
1485 let x = self.prime_layers(
1486 el, x, fence[n_st - 1], fence[n_st], &pos_d, t, base, cache, seq_end,
1487 )?;
1488 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1489 rt.publish_to(n_st - 1, &caller_stream)?;
1495 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1496 Ok(out)
1497 }
1498
1499 fn prime_pp2_stage0_enqueue(
1500 &self,
1501 e: &Engine,
1502 rt: &crate::pp::PpNRt,
1503 tokens: &[u32],
1504 cache: &mut Cache,
1505 seq_end: usize,
1506 fence: &[usize],
1507 base: usize,
1508 pipelined: bool,
1509 ) -> Result<usize, Box<dyn std::error::Error>> {
1510 let t = tokens.len();
1511 let n_embd = self.cfg.n_embd as usize;
1512 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1513 rt.bind_stage(0)?;
1514 let _st0 = rt.enter(0);
1515 let e0 = rt.engine(0, e);
1516 let pos_d = e0.htod_i32(&pos)?;
1517 let x = self.embed(e0, tokens)?;
1518 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1519 let x = self.prime_layers(
1520 e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end,
1521 )?;
1522 if pipelined {
1523 rt.tx_pipelined(0, &x, t * n_embd)
1524 } else {
1525 rt.tx(0, &x, t * n_embd)
1526 }
1527 }
1528
1529 fn prime_pp2_stage1_enqueue(
1530 &self,
1531 e: &Engine,
1532 rt: &crate::pp::PpNRt,
1533 slot: usize,
1534 t: usize,
1535 cache: &mut Cache,
1536 seq_end: usize,
1537 fence: &[usize],
1538 base: usize,
1539 pipelined: bool,
1540 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1541 let n_embd = self.cfg.n_embd as usize;
1542 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1543 rt.bind_stage(1)?;
1544 let _st1 = rt.enter(1);
1545 let e1 = rt.engine(1, e);
1546 let pos_d = e1.htod_i32(&pos)?;
1547 let x = rt.rx(0, slot, t * n_embd)?;
1548 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1549 self.prime_layers(
1550 e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end,
1551 )
1552 }
1553
1554 pub fn prime_chunk_captured(&self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>,
1570 t: usize, cache: &mut Cache,
1571 len_d: &CudaSlice<i32>,
1572 logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>)
1573 -> Result<(), Box<dyn std::error::Error>> {
1574 let cfg = &self.cfg;
1575 let n_embd = cfg.n_embd as usize;
1576 let eps = cfg.rms_eps;
1577 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1578 let mut x = e.uninit(t * n_embd)?;
1579 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1580 for (il, layer) in self.layers.iter().enumerate() {
1581 let mut h = e.uninit(t * n_embd)?;
1582 let mut hx16: Option<CudaSlice<u8>> = None;
1583 if f16fuse {
1584 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1585 e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut b16, n_embd, t, eps)?;
1586 hx16 = Some(b16);
1587 } else {
1588 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1589 }
1590 let mixed = match &layer.mixer {
1591 Mixer::Full(fa) => self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache,
1595 il, t)?,
1596 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1597 Mixer::Linear(la) => {
1598 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1599 let g4 = match hx16.as_ref() {
1600 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1601 None => e.matmul_group(&ws, &h, t)?,
1602 };
1603 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1604 }
1605 };
1606 let mut x1 = e.uninit(t * n_embd)?;
1607 e.add(&x, &mixed, &mut x1, t * n_embd)?;
1608 let mut z = e.uninit(t * n_embd)?;
1609 let mut zx16: Option<CudaSlice<u8>> = None;
1610 if f16fuse {
1611 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1612 e.rms_norm_f16out(&x1, layer.post_attn_norm.float_data(), &mut z, &mut b16, n_embd, t, eps)?;
1613 zx16 = Some(b16);
1614 } else {
1615 e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
1616 }
1617 let ffn_out = match &layer.ffn {
1618 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1619 let n_ff = ffn_gate.out_features();
1620 let mut g2 = match &zx16 {
1621 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1622 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1623 };
1624 let up = g2.pop().unwrap();
1625 let gate = g2.pop().unwrap();
1626 let mut act = e.uninit(t * n_ff)?;
1627 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
1629 self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
1630 e.matmul(ffn_down, &act, t)?
1631 }
1632 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
1633 };
1634 let mut x2 = e.uninit(t * n_embd)?;
1635 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1636 x = x2;
1637 }
1638 if !crate::spec::spec_hpost() {
1640 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1641 }
1642 let mut hn = e.uninit(t * n_embd)?;
1643 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1644 if crate::spec::spec_hpost() {
1645 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1646 }
1647 let mut hlast = e.uninit(n_embd)?;
1648 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1649 let logits = e.matmul(&self.output, &hlast, 1)?;
1650 let nv = logits.len();
1651 e.copy_into(logits_out, 0, &logits, nv)?;
1652 Ok(())
1653 }
1654
1655 fn step35_prime_batch_on() -> bool {
1656 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
1657 }
1658
1659 #[allow(clippy::too_many_arguments)]
1662 fn step35_prime_batch_layers(
1663 &self,
1664 e: &Engine,
1665 mut x: CudaSlice<f32>,
1666 lo: usize,
1667 hi: usize,
1668 ts: &[usize],
1669 offs: &[usize],
1670 pos_ds: &[CudaSlice<i32>],
1671 caches: &mut [&mut Cache],
1672 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1673 let cfg = &self.cfg;
1674 let n_embd = cfg.n_embd as usize;
1675 let eps = cfg.rms_eps;
1676 let b = ts.len();
1677 let total: usize = ts.iter().sum();
1678 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
1679
1680 let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
1681 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1682 let mut out = Vec::with_capacity(b);
1683 for s in 0..b {
1684 let mut ys = e.uninit(ts[s] * dim)?;
1685 e.copy_view_into(
1686 &mut ys,
1687 0,
1688 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
1689 ts[s] * dim,
1690 )?;
1691 out.push(ys);
1692 }
1693 Ok(out)
1694 };
1695
1696 for il in lo..hi {
1697 let layer = &self.layers[il];
1698 let Mixer::Full(fa) = &layer.mixer else {
1699 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
1700 };
1701
1702 let mut h = e.uninit(total * n_embd)?;
1703 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1704 if f16fuse {
1705 e.rms_norm_f16out(
1706 &x,
1707 layer.attn_norm.float_data(),
1708 &mut h,
1709 &mut hx16,
1710 n_embd,
1711 total,
1712 eps,
1713 )?;
1714 } else {
1715 e.rms_norm(
1716 &x,
1717 layer.attn_norm.float_data(),
1718 &mut h,
1719 n_embd,
1720 total,
1721 eps,
1722 )?;
1723 }
1724
1725 let gate_w = fa
1729 .attn_gate
1730 .as_ref()
1731 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
1732 let mut g4 = if f16fuse {
1733 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
1734 } else {
1735 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
1736 };
1737 let gate = g4.pop().unwrap();
1738 let mut parts: Vec<Vec<CudaSlice<f32>>> =
1739 (0..b).map(|_| Vec::with_capacity(3)).collect();
1740 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
1741 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
1742 parts[s].push(ys);
1743 }
1744 }
1745 let gates = split(e, &gate, gate_w.out_features())?;
1746 let geometry = self.step35_geom(il);
1747 let hd = geometry.head_dim_k as usize;
1748 let nh = geometry.n_head as usize;
1749 let mut ag_cat = e.uninit(total * nh * hd)?;
1750 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
1751 let ag = self.step35_attn_pre_wo(
1752 e,
1753 fa,
1754 g3s,
1755 None,
1756 Some(&gate),
1757 &pos_ds[s],
1758 ts[s],
1759 Some(&mut *caches[s]),
1760 il,
1761 ts[s],
1762 )?;
1763 e.copy_into(
1764 &mut ag_cat,
1765 offs[s] * nh * hd,
1766 &ag,
1767 ts[s] * nh * hd,
1768 )?;
1769 }
1770 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
1771
1772 let mut x1 = e.uninit(total * n_embd)?;
1773 let mut z = e.uninit(total * n_embd)?;
1774 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1775 if f16fuse {
1776 e.add_rms_norm_f16out(
1777 &x,
1778 &mixed,
1779 layer.post_attn_norm.float_data(),
1780 &mut x1,
1781 &mut z,
1782 &mut zx16,
1783 n_embd,
1784 total,
1785 eps,
1786 )?;
1787 } else {
1788 e.add(&x, &mixed, &mut x1, total * n_embd)?;
1789 e.rms_norm(
1790 &x1,
1791 layer.post_attn_norm.float_data(),
1792 &mut z,
1793 n_embd,
1794 total,
1795 eps,
1796 )?;
1797 }
1798
1799 let ffn_out = match &layer.ffn {
1800 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1801 let n_ff = ffn_gate.out_features();
1802 let mut g2 = if f16fuse {
1803 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
1804 } else {
1805 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
1806 };
1807 let up = g2.pop().unwrap();
1808 let gate = g2.pop().unwrap();
1809 let mut act = e.uninit(total * n_ff)?;
1810 let d_lim = cfg.clamp_shexp_at(il as u32);
1811 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
1812 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
1813 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
1814 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
1815 Some(y) => y,
1816 None => e.matmul(ffn_down, &act, total)?,
1817 }
1818 } else {
1819 Self::ffn_act_lim(
1820 e,
1821 cfg,
1822 &gate,
1823 &up,
1824 1.0,
1825 1.0,
1826 d_lim,
1827 &mut act,
1828 total * n_ff,
1829 )?;
1830 e.matmul(ffn_down, &act, total)?
1831 }
1832 }
1833 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
1834 };
1835 let mut x2 = e.uninit(total * n_embd)?;
1836 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
1837 x = x2;
1838 }
1839 Ok(x)
1840 }
1841
1842 fn step35_prime_batch_epilogue(
1843 &self,
1844 e: &Engine,
1845 x: CudaSlice<f32>,
1846 ts: &[usize],
1847 offs: &[usize],
1848 caches: &mut [&mut Cache],
1849 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1850 let n_embd = self.cfg.n_embd as usize;
1851 let total: usize = ts.iter().sum();
1852 let mut hn = e.uninit(total * n_embd)?;
1853 e.rms_norm(
1854 &x,
1855 self.output_norm.float_data(),
1856 &mut hn,
1857 n_embd,
1858 total,
1859 self.cfg.rms_eps,
1860 )?;
1861
1862 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
1863 let mut out = Vec::with_capacity(ts.len());
1864 for s in 0..ts.len() {
1865 let mut hidden = e.uninit(ts[s] * n_embd)?;
1866 e.copy_view_into(
1867 &mut hidden,
1868 0,
1869 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
1870 ts[s] * n_embd,
1871 )?;
1872 let last0 = (offs[s] + ts[s] - 1) * n_embd;
1873 let mut h_seed = e.uninit(n_embd)?;
1874 e.copy_view_into(
1875 &mut h_seed,
1876 0,
1877 &hidden_src.slice(last0..last0 + n_embd),
1878 n_embd,
1879 )?;
1880 let mut hlast = e.uninit(n_embd)?;
1882 e.copy_view_into(
1883 &mut hlast,
1884 0,
1885 &hn.slice(last0..last0 + n_embd),
1886 n_embd,
1887 )?;
1888 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
1889 caches[s].pos += ts[s];
1890 out.push((logits, h_seed, hidden));
1891 }
1892 Ok(out)
1893 }
1894
1895 fn step35_prime_cache_batch(
1896 &self,
1897 e: &Engine,
1898 prompts: &[&[u32]],
1899 caches: &mut [&mut Cache],
1900 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1901 if crate::pp::pp_host_bounce_active()
1902 && (!crate::pp::prime_pp_on()
1903 || crate::pp::pp_cuts(self.layers.len()).is_none())
1904 {
1905 return Err(
1906 "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
1907 stage split; refusing an unsplit remote-weight walk"
1908 .into(),
1909 );
1910 }
1911 if !Self::step35_prime_batch_on() {
1912 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
1913 }
1914 if caches.iter().any(|c| c.pos != 0) {
1915 return Err(
1916 "step35 batched prime currently supports complete fresh prompts only; \
1917 continuation/tick chunks require per-request queued_after"
1918 .into(),
1919 );
1920 }
1921
1922 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
1923 for &t in &ts {
1924 assert!(t >= PRIME_MIN_T, "step35 batched prime needs T >= {PRIME_MIN_T}");
1925 }
1926 for (s, c) in caches.iter().enumerate() {
1927 assert!(ts[s] <= c.max_ctx, "step35 batched prime exceeds cache max_ctx");
1928 }
1929 let offs: Vec<usize> = ts
1930 .iter()
1931 .scan(0usize, |a, &t| {
1932 let o = *a;
1933 *a += t;
1934 Some(o)
1935 })
1936 .collect();
1937 let total: usize = ts.iter().sum();
1938 let payload = total * self.cfg.n_embd as usize;
1939 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
1940 let positions: Vec<Vec<i32>> = ts
1941 .iter()
1942 .map(|&t| (0..t as i32).collect())
1943 .collect();
1944 let upload_positions = |e: &Engine|
1945 -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
1946 positions
1947 .iter()
1948 .map(|p| e.htod_i32(p))
1949 .collect::<Result<_, _>>()
1950 };
1951
1952 static ONCE: std::sync::Once = std::sync::Once::new();
1953 ONCE.call_once(|| {
1954 eprintln!(
1955 "[step35-prime-batch] first concat prime: B={} tokens={total}",
1956 prompts.len()
1957 );
1958 });
1959
1960 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1961 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1962 let rt = crate::pp::PpNRt::get(e)?;
1963 let n_st = fence.len() - 1;
1964 assert_eq!(rt.n_stages(), n_st, "step35 prime batch stage count mismatch");
1965 let caller_stream = e.stream();
1966 rt.fence_stages_behind(&caller_stream)?;
1967
1968 let mut slot = {
1969 let _st0 = rt.enter(0);
1970 let e0 = rt.engine(0, e);
1971 let pos_ds = upload_positions(e0)?;
1972 let x = self.embed(e0, &cat_tokens)?;
1973 let x = self.step35_prime_batch_layers(
1974 e0,
1975 x,
1976 fence[0],
1977 fence[1],
1978 &ts,
1979 &offs,
1980 &pos_ds,
1981 caches,
1982 )?;
1983 rt.tx(0, &x, payload)?
1984 };
1985 for s in 1..n_st - 1 {
1986 let _st = rt.enter(s);
1987 let es = rt.engine(s, e);
1988 let pos_ds = upload_positions(es)?;
1989 let x = rt.rx(s - 1, slot, payload)?;
1990 let x = self.step35_prime_batch_layers(
1991 es,
1992 x,
1993 fence[s],
1994 fence[s + 1],
1995 &ts,
1996 &offs,
1997 &pos_ds,
1998 caches,
1999 )?;
2000 slot = rt.tx(s, &x, payload)?;
2001 }
2002
2003 let _stl = rt.enter(n_st - 1);
2004 let el = rt.engine(n_st - 1, e);
2005 let pos_ds = upload_positions(el)?;
2006 let x = rt.rx(n_st - 2, slot, payload)?;
2007 let x = self.step35_prime_batch_layers(
2008 el,
2009 x,
2010 fence[n_st - 1],
2011 fence[n_st],
2012 &ts,
2013 &offs,
2014 &pos_ds,
2015 caches,
2016 )?;
2017 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2018 rt.publish_to(n_st - 1, &caller_stream)?;
2019 crate::pp::STEP35_PRIME_BATCH_SPLITS.fetch_add(
2020 1,
2021 std::sync::atomic::Ordering::Relaxed,
2022 );
2023 out
2024 } else {
2025 let pos_ds = upload_positions(e)?;
2026 let x = self.embed(e, &cat_tokens)?;
2027 let x = self.step35_prime_batch_layers(
2028 e,
2029 x,
2030 0,
2031 self.layers.len(),
2032 &ts,
2033 &offs,
2034 &pos_ds,
2035 caches,
2036 )?;
2037 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2038 }
2039 } else {
2040 let pos_ds = upload_positions(e)?;
2041 let x = self.embed(e, &cat_tokens)?;
2042 let x = self.step35_prime_batch_layers(
2043 e,
2044 x,
2045 0,
2046 self.layers.len(),
2047 &ts,
2048 &offs,
2049 &pos_ds,
2050 caches,
2051 )?;
2052 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2053 };
2054 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2055 Ok(out)
2056 }
2057
2058 pub fn prime_cache_batch(&self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache])
2075 -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2076 let cfg = &self.cfg;
2077 let n_embd = cfg.n_embd as usize;
2078 let eps = cfg.rms_eps;
2079 let b = prompts.len();
2080 assert!(b >= 1 && b == caches.len());
2081 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2082 let carried = pos0s.iter().any(|&p| p > 0);
2083 if cfg.gemma4.is_some() {
2089 return Err("prime_cache_batch: gemma4 has no batched prime core (per-layer \
2090 swa/global geometry, softcapped head) — use gemma4_prime per sequence".into());
2091 }
2092 if cfg.step35.is_some() {
2095 return self.step35_prime_cache_batch(e, prompts, caches);
2096 }
2097 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2098 for &t in &ts { assert!(t >= PRIME_MIN_T, "prime_cache_batch needs T >= {PRIME_MIN_T}"); }
2099 for (s, c) in caches.iter().enumerate() {
2100 assert!(c.pos + ts[s] <= c.max_ctx, "prime_cache_batch: prompt exceeds cache max_ctx");
2101 }
2102 let total: usize = ts.iter().sum();
2103 let offs: Vec<usize> = ts.iter().scan(0usize, |a, &t| { let o = *a; *a += t; Some(o) }).collect();
2104 let pos_ds: Vec<CudaSlice<i32>> = ts.iter().zip(&pos0s)
2106 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2107 .collect::<Result<_, _>>()?;
2108 let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
2110 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2111 let mut out = Vec::with_capacity(b);
2112 for s in 0..b {
2113 let mut ys = e.uninit(ts[s] * dim)?;
2114 e.copy_view_into(&mut ys, 0, &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim), ts[s] * dim)?;
2115 out.push(ys);
2116 }
2117 Ok(out)
2118 };
2119
2120 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2121 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
2123 let mut h = e.uninit(total * n_embd)?;
2124 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2125 e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut hx16, n_embd, total, eps)?;
2126 let mut mixed = e.uninit(total * n_embd)?;
2128 match &layer.mixer {
2129 Mixer::Full(fa) => {
2130 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2131 let geometry = self.cfg.full_attention_geometry_at(il as u32);
2137 let (n_head, n_head_kv, head_dim) = (
2138 geometry.n_head as usize,
2139 geometry.n_head_kv as usize,
2140 geometry.head_dim_k as usize,
2141 );
2142 let fa_scale = geometry.attention_scale();
2143 let use_favl = !carried
2144 && (2..=8).contains(&b)
2145 && (head_dim == 256 || head_dim == 128)
2146 && geometry.attention_gate
2147 == memra_gguf::config::AttentionGateKind::FusedQ
2148 && std::env::var("MEMRA_NOFA").is_err()
2149 && std::env::var("MEMRA_FA_FLOOR").is_err()
2150 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2151 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2152 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2153 if use_favl {
2154 let (qf_w, kf_w, vf_w) =
2155 (fa.wq.out_features(), fa.wk.out_features(), fa.wv.out_features());
2156 struct APre {
2157 q: CudaSlice<f32>, gate: Option<CudaSlice<f32>>,
2158 qn: CudaSlice<f32>, kn: CudaSlice<f32>,
2159 }
2160 let mut aps = Vec::with_capacity(b);
2161 for &t in ts.iter().take(b) {
2162 aps.push(APre {
2163 q: e.uninit(t * n_head * head_dim)?,
2164 gate: Some(e.uninit(t * n_head * head_dim)?),
2165 qn: e.uninit(t * n_head * head_dim)?,
2166 kn: e.uninit(t * n_head_kv * head_dim)?,
2167 });
2168 }
2169 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2170 let kvl = caches[0].kv[il].as_ref().unwrap();
2171 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2172 };
2173 let pargs: Vec<crate::AttnPreVl> = (0..b).map(|s| {
2174 let (o, t) = (offs[s], ts[s]);
2175 let kvl = caches[s].kv[il].as_ref().unwrap();
2176 assert!(kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2177 "prime_cache_batch attn vl: fresh + capacity");
2178 crate::AttnPreVl {
2179 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2180 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2181 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2182 q: e.addr_f32(&aps[s].q),
2183 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2184 qn: e.addr_f32(&aps[s].qn), kn: e.addr_f32(&aps[s].kn),
2185 kc: e.addr_u8(&kvl.k), vc: e.addr_u8(&kvl.v),
2186 t: t as i32, pad: 0,
2187 }
2188 }).collect();
2189 e.attn_pre_vl8(&pargs, fa.q_norm.float_data(), fa.k_norm.float_data(),
2190 head_dim, geometry.n_rot as usize, n_head, n_head_kv,
2191 self.cfg.rms_eps, geometry.rope_base, 1.0,
2192 kv_dim_k, kv_dim_v, ktb, vtb)?;
2193 for s in 0..b {
2194 let kvl = caches[s].kv[il].as_mut().unwrap();
2195 kvl.len += ts[s];
2196 let new_len = kvl.len as i32;
2197 e.set_i32_one(&mut kvl.len_d, new_len)?;
2198 }
2199 let mut attns = Vec::with_capacity(b);
2200 let mut mirrors = Vec::with_capacity(b);
2201 for &t in ts.iter().take(b) {
2202 attns.push(e.uninit(t * n_head * head_dim)?);
2203 let n = t * n_head_kv * head_dim;
2204 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2205 }
2206 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2209 Ok("0") => false,
2210 Ok("1") => true,
2211 _ => cfg!(memra_hopper_mma),
2212 };
2213 if fa3_on {
2214 let mut q16s = Vec::with_capacity(b);
2215 let mut v16s = Vec::with_capacity(b);
2216 for s in 0..b {
2217 let t = ts[s];
2218 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2219 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2220 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2221 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2222 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2223 e.f32_to_bf16_v(&g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2224 &mut v16, t * n_head_kv * head_dim)?;
2225 q16s.push(q16);
2226 v16s.push((k16, v16));
2227 }
2228 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2229 let mut kp = qp;
2230 let mut vp = qp;
2231 let mut op = [core::ptr::null_mut::<f32>(); 8];
2232 let mut tsv = [0i32; 8];
2233 for s in 0..b {
2234 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2235 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2236 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2237 op[s] = e.addr_f32(&attns[s]) as *mut f32;
2238 tsv[s] = ts[s] as i32;
2239 }
2240 let rc = unsafe {
2241 crate::fa3_vl_raw(qp.as_ptr(), kp.as_ptr(), vp.as_ptr(), op.as_ptr(),
2242 tsv.as_ptr(), b as i32, n_head as i32,
2243 n_head_kv as i32, head_dim as i32, fa_scale,
2244 e.stream().cu_stream() as *mut core::ffi::c_void)
2245 };
2246 if rc != 0 {
2247 return Err(format!("memra_fa3_vl rc={rc}").into());
2248 }
2249 } else {
2250 let fargs: Vec<crate::FaSeqVl> = (0..b).map(|s| crate::FaSeqVl {
2251 q: e.addr_f32(&aps[s].qn), k16: e.addr_u8(&mirrors[s].0),
2252 v16: e.addr_u8(&mirrors[s].1), o: e.addr_f32(&attns[s]),
2253 kf: e.addr_f32(&aps[s].kn),
2254 vf: e.addr_f32v(&g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w)),
2255 t: ts[s] as i32, pad: 0,
2256 }).collect();
2257 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2258 }
2259 for (s, attn) in attns.into_iter().enumerate() {
2260 let (attn_g, ag16) = self.full_attn_prime_post_fa(
2261 e, attn, &aps[s].gate, ts[s], n_head, head_dim)?;
2262 let mut done = false;
2263 if let Some(xh) = &ag16 {
2264 done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
2265 }
2266 if !done {
2267 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2268 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2269 }
2270 }
2271 } else {
2272 let mut parts: Vec<Vec<CudaSlice<f32>>> = (0..b).map(|_| Vec::new()).collect();
2273 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2274 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2275 parts[s].push(ys);
2276 }
2277 }
2278 for (s, g3s) in parts.into_iter().enumerate() {
2279 let (attn_g, ag16) = self.full_attn_prime_core_inner(
2281 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il)?;
2282 let mut done = false;
2283 if let Some(xh) = &ag16 {
2284 done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
2285 }
2286 if !done {
2287 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2288 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2289 }
2290 }
2291 }
2292 }
2293 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2294 Mixer::Linear(la) => {
2295 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2300 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2301 let outs = self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2302 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2303 let (o, t) = (offs[s], ts[s]);
2304 let mut done = false;
2305 if let Some(xh) = &gn16 {
2306 done = e.try_f16_gemm_pre_into_off(&la.ssm_out, xh, t, &mut mixed, o * n_embd)?;
2307 }
2308 if !done {
2309 let m = e.matmul(&la.ssm_out, &gn, t)?;
2310 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2311 }
2312 }
2313 }
2314 }
2315 let mut x1 = e.uninit(total * n_embd)?;
2316 let mut z = e.uninit(total * n_embd)?;
2317 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2318 e.add_rms_norm_f16out(&x, &mixed, layer.post_attn_norm.float_data(),
2319 &mut x1, &mut z, &mut zx16, n_embd, total, eps)?;
2320 let ffn_out = match &layer.ffn {
2321 crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
2322 let n_ff = ffn_gate.out_features();
2323 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2324 let up = g2.pop().unwrap();
2325 let gate = g2.pop().unwrap();
2326 let mut act = e.uninit(total * n_ff)?;
2327 let d_lim = self.cfg.clamp_shexp_at(il as u32);
2331 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2332 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2333 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2334 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2335 Some(y) => y,
2336 None => e.matmul(ffn_down, &act, total)?,
2337 }
2338 } else {
2339 Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, d_lim,
2340 &mut act, total * n_ff)?;
2341 e.matmul(ffn_down, &act, total)?
2342 }
2343 }
2344 crate::hybrid::Ffn::Moe(m) => {
2345 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2346 }
2347 };
2348 let mut x2 = e.uninit(total * n_embd)?;
2349 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2350 x = x2;
2351 }
2352 let mut hn = e.uninit(total * n_embd)?;
2354 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, total, eps)?;
2355 let mut hcat = e.uninit(b * n_embd)?;
2361 for s in 0..b {
2362 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2363 e.copy_view_into(&mut hcat, s * n_embd, &hn.slice(last0..last0 + n_embd), n_embd)?;
2364 }
2365 let logits_cat = if b >= 2 { e.try_f16_gemm(&self.output, &hcat, b)? } else { None };
2366 let logits_host: Option<Vec<f32>> = match &logits_cat {
2367 Some(lc) => Some(e.dtoh(lc)?),
2368 None => None,
2369 };
2370 let n_vocab = self.output.out_features();
2371 let mut hidden_all = if crate::spec::spec_hpost() {
2372 split(e, &hn, n_embd)?
2373 } else {
2374 split(e, &x, n_embd)?
2375 };
2376 let mut out = Vec::with_capacity(b);
2377 for s in 0..b {
2378 let last0 = (offs[s] + ts[s] - 1) * n_embd;
2379 let mut h_seed = e.uninit(n_embd)?;
2380 if !crate::spec::spec_hpost() {
2381 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2382 } else {
2383 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2384 }
2385 let logits = match &logits_host {
2386 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2387 None => {
2388 let mut hlast = e.uninit(n_embd)?;
2389 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2390 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2391 }
2392 };
2393 caches[s].pos += ts[s];
2394 out.push((logits, h_seed, hidden_all.remove(0)));
2395 }
2396 Ok(out)
2397 }
2398
2399 #[allow(clippy::too_many_arguments)]
2410 fn full_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
2411 hx: Option<&CudaSlice<u8>>,
2412 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize,
2413 seq_end: usize)
2414 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2415 if self.cfg.step35.is_some() {
2416 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2417 }
2418 let g3 = match hx {
2423 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2424 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2425 };
2426 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2427 }
2428
2429 fn full_attn_prime_core(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
2433 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2434 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2435 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2436 if let Some(xh) = &ag16 {
2437 if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2438 return Ok(y);
2439 }
2440 }
2441 Ok(e.matmul(&fa.wo, &attn_g, t)?)
2442 }
2443
2444 fn full_attn_prime_core_inner(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
2445 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2446 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2447 let cfg = &self.cfg;
2448 let geometry = cfg.full_attention_geometry_at(il as u32);
2449 let n_head = geometry.n_head as usize;
2450 let n_head_kv = geometry.n_head_kv as usize;
2451 let head_dim = geometry.head_dim_k as usize;
2452 let scale = geometry.attention_scale();
2453 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2454 let AttnPre { q, k, v, gate } = pre;
2455 let mut attn = e.uninit(t * n_head * head_dim)?;
2456 self.full_attn_prime_fa_dispatch(e, &q, &k, &v, &mut attn, base_len, t, cache, il,
2457 head_dim, n_head, n_head_kv, scale)?;
2458 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2459 }
2460
2461 #[allow(clippy::type_complexity)]
2465 fn full_attn_prime_pre_fa(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
2466 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2467 -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
2468 let cfg = &self.cfg;
2469 let geometry = cfg.full_attention_geometry_at(il as u32);
2470 let n_head = geometry.n_head as usize;
2471 let n_head_kv = geometry.n_head_kv as usize;
2472 let head_dim = geometry.head_dim_k as usize;
2473 let eps = cfg.rms_eps;
2474
2475 let gated = geometry.attention_gate
2479 == memra_gguf::config::AttentionGateKind::FusedQ;
2480 let v = g3.pop().unwrap();
2481 let mut k = g3.pop().unwrap();
2482 let qf = g3.pop().unwrap();
2483 let (mut q, gate) = if gated {
2484 let mut q = e.uninit(t * n_head * head_dim)?;
2485 let mut gate = e.uninit(t * n_head * head_dim)?;
2486 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2487 (q, Some(gate))
2488 } else {
2489 (qf, None)
2490 };
2491
2492 let mut qn = e.uninit(t * n_head * head_dim)?;
2493 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2494 q = qn;
2495 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2496 e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2497 k = kn;
2498 let rope_dims = geometry.n_rot as usize;
2499 e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, geometry.rope_base, 1.0)?;
2500 e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, geometry.rope_base, 1.0)?;
2501
2502 {
2505 let kvl = cache.kv[il].as_mut().unwrap();
2506 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
2507 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
2508 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
2509 crate::Engine::kv_fp8_on())?;
2510 kvl.len += t;
2511 let new_len = kvl.len as i32;
2512 e.set_i32_one(&mut kvl.len_d, new_len)?;
2513 }
2514
2515 let base_len = {
2516 let kvl = cache.kv[il].as_ref().unwrap();
2517 kvl.len - t };
2519 Ok((AttnPre { q, k, v, gate }, base_len))
2520 }
2521
2522 #[allow(clippy::too_many_arguments)]
2529 fn full_attn_prime_fa_dispatch(&self, e: &Engine, q: &CudaSlice<f32>, k: &CudaSlice<f32>,
2530 v: &CudaSlice<f32>, attn: &mut CudaSlice<f32>, base_len: usize,
2531 t: usize, cache: &mut Cache, il: usize,
2532 head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32)
2533 -> Result<(), Box<dyn std::error::Error>> {
2534 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
2547 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2548 e.sdpa_naive(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2549 } else {
2550 e.fa_prefill(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2551 }
2552 return Ok(());
2553 }
2554 let kvl = cache.kv[il].as_ref().unwrap();
2555 let t_kv = base_len + t;
2556 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
2557 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
2558 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2562 e.sdpa_naive_quantized_view(q, &k_view, &v_view, attn, head_dim, n_head,
2563 n_head_kv, t, t_kv, scale, true,
2564 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
2565 return Ok(());
2566 }
2567 let deqw = std::env::var("MEMRA_PRIME_DEQW").map(|v| v != "0").unwrap_or(true);
2575 if deqw {
2576 e.fa_prefill_view_ws(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
2577 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
2578 crate::Engine::kv_fp8_on())?;
2579 } else {
2580 e.fa_prefill_view(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
2581 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
2582 crate::Engine::kv_fp8_on())?;
2583 }
2584 Ok(())
2585 }
2586
2587 fn full_attn_prime_post_fa(&self, e: &Engine, attn: CudaSlice<f32>,
2590 gate: &Option<CudaSlice<f32>>, t: usize,
2591 n_head: usize, head_dim: usize)
2592 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2593 let (attn_g, ag16) = match gate {
2594 Some(gate) => {
2595 let n = t * n_head * head_dim;
2596 let mut ag = e.uninit(n)?;
2597 if Self::f16out_on(e, t) {
2598 let mut a16 = e.alloc_u8_uninit(n * 2)?;
2599 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
2600 (ag, Some(a16))
2601 } else {
2602 let mut gsig = e.uninit(n)?;
2603 e.sigmoid(gate, &mut gsig, n)?;
2604 e.mul(&attn, &gsig, &mut ag, n)?;
2605 (ag, None)
2606 }
2607 }
2608 None => (attn, None),
2609 };
2610 Ok((attn_g, ag16))
2611 }
2612
2613 fn linear_attn_prime(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>,
2620 hx: Option<&CudaSlice<u8>>, t: usize,
2621 cache: &mut Cache, il: usize)
2622 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2623 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2625 let g4 = match hx {
2626 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
2627 None => e.matmul_group(&ws, h, t)?,
2628 };
2629 self.linear_attn_prime_core(e, la, g4, t, cache, il)
2630 }
2631
2632 fn linear_attn_prime_core(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
2634 t: usize, cache: &mut Cache, il: usize)
2635 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2636 self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
2637 }
2638
2639 #[allow(clippy::too_many_arguments)]
2643 fn linear_attn_prime_core_pad_inner(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
2644 t: usize, cache: &mut Cache, il: usize,
2645 pad_len: Option<&CudaSlice<i32>>)
2646 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2647 let ssm = self.cfg.ssm.as_ref().unwrap();
2649 let d_state = ssm.state_size as usize;
2650 let num_k = ssm.group_count as usize;
2651 let num_v = ssm.time_step_rank as usize;
2652 let key_dim = d_state * num_k;
2653 let value_dim = d_state * num_v;
2654 let conv_dim = key_dim * 2 + value_dim;
2655 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(
2660 e, la,
2661 &qkv_mixed.slice(0..t * conv_dim), &z.slice(0..t * value_dim),
2662 &beta_raw.slice(0..t * num_v), &alpha.slice(0..t * num_v),
2663 t, cache, il, pad_len)
2664 }
2665
2666 #[allow(clippy::too_many_arguments)]
2669 fn linear_attn_gdn_prep(&self, e: &Engine, la: &LinearAttnLayer,
2670 qkv_mixed: &cudarc::driver::CudaView<f32>,
2671 beta_raw: &cudarc::driver::CudaView<f32>,
2672 alpha: &cudarc::driver::CudaView<f32>,
2673 t: usize, cache: &mut Cache, il: usize,
2674 pad_len: Option<&CudaSlice<i32>>)
2675 -> Result<GdnPrep, Box<dyn std::error::Error>> {
2676 let cfg = &self.cfg;
2677 let ssm = cfg.ssm.as_ref().unwrap();
2678 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;
2686 debug_assert!(t >= d_conv - 1, "stateful conv needs T >= pad (PRIME_MIN_T gates)");
2687
2688 let rl = cache.recur[il].as_mut().unwrap();
2693 let hk = Self::gdn_hk(e, t, num_v, num_k);
2694 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
2695 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
2697 let mut k_g = e.uninit(d_state * hk * t)?;
2698 let mut v_g = e.uninit(d_state * num_v * t)?;
2699 if conv_fuse {
2700 e.ssm_conv1d_gdn_state_pad(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
2701 &mut q_g, &mut k_g, &mut v_g,
2702 conv_dim, t, d_conv, d_state, num_v, num_k, key_dim, hk, pad_len)?;
2703 } else {
2704 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(),
2706 &mut conv_out, conv_dim, t, d_conv, pad_len)?;
2707 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)?;
2708 }
2709 let mut q_l2 = e.uninit(d_state * hk * t)?;
2710 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
2714 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
2715 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
2716 Some(qb)
2717 } else {
2718 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
2719 None
2720 };
2721 let mut k_l2 = e.uninit(d_state * hk * t)?;
2722 let kb16 = if Engine::l2_v2_on(d_state) {
2724 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
2725 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
2726 Some(kb)
2727 } else {
2728 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
2729 None
2730 };
2731 let mut beta = e.uninit(t * num_v)?;
2732 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
2733 let mut g_log = e.uninit(t * num_v)?;
2734 e.gdn_glog_v(alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
2735 if let Some(len_d) = pad_len {
2736 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
2737 }
2738 Ok(GdnPrep { hk, q_l2, k_l2, v_g, beta, g_log, kb16, qb16 })
2739 }
2740
2741 #[allow(clippy::too_many_arguments)]
2746 fn linear_attn_prime_core_batch(&self, e: &Engine, la: &LinearAttnLayer,
2747 g4: &[CudaSlice<f32>], offs: &[usize], ts: &[usize],
2748 caches: &mut [&mut Cache], il: usize)
2749 -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
2750 let ssm = self.cfg.ssm.as_ref().unwrap();
2751 let d_state = ssm.state_size as usize;
2752 let num_k = ssm.group_count as usize;
2753 let num_v = ssm.time_step_rank as usize;
2754 let key_dim = d_state * num_k;
2755 let value_dim = d_state * num_v;
2756 let conv_dim = key_dim * 2 + value_dim;
2757 let eps = self.cfg.rms_eps;
2758 let scale = 1.0 / (d_state as f32).sqrt();
2759 let b = ts.len();
2760 let c = Engine::gdn_chunk_size();
2761 let carried = caches.iter().any(|c| c.pos > 0);
2764 let use_vl = !carried
2765 && (2..=8).contains(&b)
2766 && Engine::gdn_chunked_enabled() && ts.iter().all(|&t| t >= 16)
2767 && e.gdn_mma_enabled(c)
2768 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
2769 if !use_vl {
2770 return (0..b).map(|s| {
2771 let (o, t) = (offs[s], ts[s]);
2772 self.linear_attn_prime_core_pad_view(
2773 e, la,
2774 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
2775 &g4[1].slice(o * value_dim..(o + t) * value_dim),
2776 &g4[2].slice(o * num_v..(o + t) * num_v),
2777 &g4[3].slice(o * num_v..(o + t) * num_v),
2778 t, caches[s], il, None)
2779 }).collect();
2780 }
2781 struct SeqBufs {
2785 conv_out: CudaSlice<f32>, q_g: CudaSlice<f32>, k_g: CudaSlice<f32>, v_g: CudaSlice<f32>,
2786 q_l2: CudaSlice<f32>, k_l2: CudaSlice<f32>, beta: CudaSlice<f32>, g_log: CudaSlice<f32>,
2787 gn: CudaSlice<f32>, gn16: CudaSlice<u8>,
2788 }
2789 let d_conv = ssm.conv_kernel as usize;
2790 let f16o = Self::f16out_on(e, 16);
2791 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
2793 let mut pres = Vec::with_capacity(b);
2794 for &t in ts.iter().take(b) {
2795 sb.push(SeqBufs {
2796 conv_out: e.uninit(conv_dim * t)?,
2797 q_g: e.uninit(d_state * hk * t)?,
2798 k_g: e.uninit(d_state * hk * t)?,
2799 v_g: e.uninit(d_state * num_v * t)?,
2800 q_l2: e.uninit(d_state * hk * t)?,
2801 k_l2: e.uninit(d_state * hk * t)?,
2802 beta: e.uninit(t * num_v)?,
2803 g_log: e.uninit(t * num_v)?,
2804 gn: e.uninit(d_state * num_v * t)?,
2805 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
2806 });
2807 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
2808 }
2809 let prep_args: Vec<crate::GdnPrepVl> = (0..b).map(|s| {
2810 let (o, t) = (offs[s], ts[s]);
2811 let rl = caches[s].recur[il].as_ref().unwrap();
2812 crate::GdnPrepVl {
2813 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
2814 conv_state: e.addr_f32(&rl.conv_state),
2815 conv_out: e.addr_f32(&sb[s].conv_out),
2816 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),
2817 q_l2: e.addr_f32(&sb[s].q_l2), k_l2: e.addr_f32(&sb[s].k_l2),
2818 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
2819 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
2820 beta: e.addr_f32(&sb[s].beta), g_log: e.addr_f32(&sb[s].g_log),
2821 o: e.addr_f32(&pres[s].o),
2822 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
2823 gn: e.addr_f32(&sb[s].gn), gn16: e.addr_u8(&sb[s].gn16),
2824 kb16: if Engine::l2_v2_on(d_state) { e.addr_u8(&pres[s].kb16) } else { 0 },
2825 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) { e.addr_u8(&pres[s].qb16) } else { 0 },
2826 t: t as i32, pad: 0,
2827 }
2828 }).collect();
2829 let args: Vec<crate::GdnSeqVl> = (0..b).map(|s| {
2830 let rl = caches[s].recur[il].as_ref().unwrap();
2831 crate::GdnSeqVl {
2832 kb16: e.addr_u8(&pres[s].kb16), gcum: e.addr_f32(&pres[s].gcum),
2833 beta: e.addr_f32(&sb[s].beta), u: e.addr_f32(&pres[s].u),
2834 wb16: e.addr_u8(&pres[s].wb16), y: e.addr_u8(&pres[s].y16),
2835 ssnap: e.addr_u8(&pres[s].ssnap16),
2836 state_in: e.addr_f32(&rl.ssm_state), state_out: e.addr_f32(&rl.ssm_state_alt),
2837 q: e.addr_f32(&sb[s].q_l2), p: e.addr_f32(&pres[s].p),
2838 o: e.addr_f32(&pres[s].o),
2839 k: e.addr_f32(&sb[s].k_l2), v: e.addr_f32(&sb[s].v_g),
2840 g: e.addr_f32(&sb[s].g_log), a: e.addr_f32(&pres[s].a),
2841 w: e.addr_f32(&pres[s].w),
2842 t: ts[s] as i32, nc: pres[s].nc as i32,
2843 }
2844 }).collect();
2845 e.gdn_prep_vl8(&prep_args, la.ssm_conv1d.float_data(), la.ssm_dt.float_data(),
2846 la.ssm_a.float_data(), conv_dim, d_conv, d_state, num_v, num_k, key_dim, hk, eps)?;
2847 if !Engine::l2_v2_on(d_state) {
2850 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
2851 }
2852 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
2854 if !Engine::l2_v2_on(d_state) {
2856 for s in 0..b {
2857 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
2858 }
2859 }
2860 let mut wa = [crate::GdnWVl::default(); 8];
2861 for s in 0..b {
2862 wa[s] = crate::GdnWVl { qb16: e.addr_u8(&pres[s].qb16), pb16: e.addr_u8(&pres[s].pb16) };
2863 }
2864 Some(crate::GdnWVl8(wa))
2865 } else { None };
2866 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
2867 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
2868 if f16o {
2869 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
2870 }
2871 let mut out = Vec::with_capacity(b);
2873 for (s, bufs) in sb.into_iter().enumerate() {
2874 let rl = caches[s].recur[il].as_mut().unwrap();
2875 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2876 let (o, t) = (offs[s], ts[s]);
2877 let SeqBufs { mut gn, gn16, .. } = bufs;
2878 if f16o {
2879 out.push((gn, Some(gn16)));
2880 } else {
2881 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
2882 e.gated_rmsnorm_zv(&pres[s].o, la.ssm_norm.float_data(), &z_v, &mut gn,
2883 d_state, num_v * t, eps)?;
2884 out.push((gn, None));
2885 }
2886 }
2887 Ok(out)
2888 }
2889
2890 #[allow(clippy::too_many_arguments)]
2894 fn linear_attn_prime_core_pad_view(&self, e: &Engine, la: &LinearAttnLayer,
2895 qkv_mixed: &cudarc::driver::CudaView<f32>,
2896 z: &cudarc::driver::CudaView<f32>,
2897 beta_raw: &cudarc::driver::CudaView<f32>,
2898 alpha: &cudarc::driver::CudaView<f32>,
2899 t: usize, cache: &mut Cache, il: usize,
2900 pad_len: Option<&CudaSlice<i32>>)
2901 -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2902 let cfg = &self.cfg;
2903 let ssm = cfg.ssm.as_ref().unwrap();
2904 let d_state = ssm.state_size as usize; let num_v = ssm.time_step_rank as usize; let eps = cfg.rms_eps;
2907 let scale = 1.0 / (d_state as f32).sqrt();
2908
2909 let prep = self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
2910
2911 let mut o = e.uninit(d_state * num_v * t)?;
2917 let rl = cache.recur[il].as_mut().unwrap();
2918 {
2919 let crate::cache::RecurLayer { ssm_state, ssm_state_alt, .. } = rl;
2920 e.gdn_scan_prefill(&prep.q_l2, &prep.k_l2, &prep.v_g, &prep.g_log, &prep.beta,
2921 prep.kb16.as_ref(), prep.qb16.as_ref(), ssm_state, ssm_state_alt, &mut o, num_v, t, scale,
2922 prep.hk)?;
2923 }
2924 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2925
2926 let mut gn = e.uninit(d_state * num_v * t)?;
2929 let gn16 = if Self::f16out_on(e, t) {
2930 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
2931 e.gated_rmsnorm_f16out_zv(&o, la.ssm_norm.float_data(), z, &mut gn, &mut g16,
2932 d_state, num_v * t, eps)?;
2933 Some(g16)
2934 } else {
2935 e.gated_rmsnorm_zv(&o, la.ssm_norm.float_data(), z, &mut gn, d_state, num_v * t, eps)?;
2936 None
2937 };
2938 Ok((gn, gn16))
2939 }
2940
2941 #[allow(clippy::too_many_arguments)]
2943 fn linear_attn_prime_core_pad(&self, e: &Engine, la: &LinearAttnLayer, g4: Vec<CudaSlice<f32>>,
2944 t: usize, cache: &mut Cache, il: usize,
2945 pad_len: Option<&CudaSlice<i32>>)
2946 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2947 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
2948 if let Some(xh) = &gn16 {
2949 if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
2950 return Ok(y);
2951 }
2952 }
2953 Ok(e.matmul(&la.ssm_out, &gn, t)?)
2954 }
2955
2956 pub fn full_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, il: usize)
2961 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2962 if self.cfg.step35.is_some() {
2963 return self.step35_attn(e, fa, h, pos_d, t, il);
2964 }
2965 let cfg = &self.cfg;
2966 let _n_embd = cfg.n_embd as usize;
2967 let geometry = cfg.full_attention_geometry_at(il as u32);
2968 let n_head = geometry.n_head as usize;
2969 let n_head_kv = geometry.n_head_kv as usize;
2970 let head_dim = geometry.head_dim_k as usize;
2971 let eps = cfg.rms_eps;
2972 let scale = geometry.attention_scale();
2973
2974 let gated = geometry.attention_gate
2977 == memra_gguf::config::AttentionGateKind::FusedQ;
2978 let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
2980 let v = g3.pop().unwrap();
2981 let mut k = g3.pop().unwrap();
2982 let qf = g3.pop().unwrap();
2983 let (mut q, gate) = if gated {
2984 let mut q = e.uninit(t * n_head * head_dim)?;
2985 let mut gate = e.uninit(t * n_head * head_dim)?;
2986 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2987 (q, Some(gate))
2988 } else {
2989 (qf, None)
2990 };
2991
2992 let mut qn = e.uninit(t * n_head * head_dim)?;
2994 e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2995 q = qn;
2996 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2997 e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2998 k = kn;
2999 let rope_dims = geometry.n_rot as usize;
3000 e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, geometry.rope_base, 1.0)?;
3001 e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, geometry.rope_base, 1.0)?;
3002
3003 let mut attn = e.uninit(t * n_head * head_dim)?;
3005 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3008 e.sdpa_naive(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
3010 } else {
3011 e.fa_prefill(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
3012 }
3013
3014 let attn_g = match &gate {
3016 Some(gate) => {
3017 let mut gsig = e.uninit(t * n_head * head_dim)?;
3018 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3019 let mut ag = e.uninit(t * n_head * head_dim)?;
3020 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3021 ag
3022 }
3023 None => attn,
3024 };
3025
3026 let o = e.matmul(&fa.wo, &attn_g, t)?;
3028 Ok(o)
3029 }
3030
3031 pub fn linear_attn(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize)
3033 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3034 let cfg = &self.cfg;
3035 let _n_embd = cfg.n_embd as usize;
3036 let ssm = cfg.ssm.as_ref().unwrap();
3037 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;
3042 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;
3046 let scale = 1.0 / (d_state as f32).sqrt();
3047
3048 let mut g4 = e.matmul_group(&[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha], h, t)?;
3051 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);
3063 let mut q_g = e.uninit(d_state * num_v * t)?;
3064 let mut k_g = e.uninit(d_state * num_v * t)?;
3065 let mut v_g = e.uninit(d_state * num_v * t)?;
3066 e.ssm_conv1d_gdn(&qkv_mixed, la.ssm_conv1d.float_data(), &mut q_g, &mut k_g, &mut v_g,
3067 conv_dim, t, d_conv, d_state, num_v, num_k, key_dim)?;
3068 let mut q_l2 = e.uninit(d_state * num_v * t)?;
3070 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3071 let mut k_l2 = e.uninit(d_state * num_v * t)?;
3072 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3073 let v_gd = v_g;
3074
3075 let mut beta = e.uninit(t * num_v)?;
3078 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3079 let mut g_log = e.uninit(t * num_v)?;
3081 e.gdn_glog(&alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
3082
3083 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
3086 let mut o = e.uninit(d_state * num_v * t)?;
3087 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)?;
3088
3089 let mut gn = e.uninit(d_state * num_v * t)?;
3094 e.gated_rmsnorm(&o, la.ssm_norm.float_data(), &z, &mut gn, d_state, num_v * t, eps)?;
3095
3096 let out = e.matmul(&la.ssm_out, &gn, t)?;
3100 Ok(out)
3101 }
3102}
3103
3104impl HybridModel {
3105 pub fn moe_ffn_il(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16)
3116 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3117 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
3118 }
3119
3120 pub fn moe_ffn_il_prefill(
3123 &self,
3124 e: &Engine,
3125 m: &MoeWeights,
3126 z: &CudaSlice<f32>,
3127 t: usize,
3128 il: u16,
3129 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3130 Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
3131 }
3132
3133 pub fn moe_ffn_il_zq8(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
3137 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16)
3138 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3139 Self::moe_ffn_inner(
3140 e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false,
3141 )
3142 }
3143
3144 pub(crate) fn moe_ffn(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
3152 cfg: &ModelConfig, il: u16, max_block: usize)
3153 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3154 Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
3155 }
3156
3157 #[allow(clippy::too_many_arguments)]
3158 pub(crate) fn moe_ffn_inner(
3159 e: &Engine,
3160 m: &MoeWeights,
3161 z: &CudaSlice<f32>,
3162 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
3163 t: usize,
3164 cfg: &ModelConfig,
3165 il: u16,
3166 max_block: usize,
3167 prefill: bool,
3168 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3169 let worker_io = crate::spill_pread::worker_enabled();
3170 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
3171 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
3172 e.with_moe_cache(max_block, |cache, _| {
3173 cache.begin_forward_epoch(il, t);
3174 if worker_io {
3175 cache.begin_worker_scope();
3176 }
3177 Ok(())
3178 })?;
3179 }
3180 if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
3181 let moe = cfg.moe.as_ref().unwrap();
3182 let n_expert = moe.expert_count as usize;
3183 let n_used = moe.expert_used_count as usize;
3184 let sigmoid = cfg.sigmoid_router().unwrap();
3185 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
3186 Self::trace_sigmoid_router_logits(
3187 e, il, t, n_expert, n_used, &logits, m, sigmoid,
3188 )?;
3189 return Self::moe_ffn_sigmoid_dev(
3190 e, m, z, zq8, &logits, t, cfg, il, sigmoid,
3191 );
3192 }
3193 if t > 1 && moe_grouped_enabled(cfg, prefill) {
3196 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
3197 if std::env::var("MEMRA_MOE_GATE").is_ok() {
3202 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
3203 let g_host = e.dtoh(&grouped_out)?;
3204 let s_host = e.dtoh(&seq_out)?;
3205 let g_bytes: &[u8] = unsafe { std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4) };
3206 let s_bytes: &[u8] = unsafe { std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4) };
3207 if g_bytes == s_bytes {
3208 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
3209 } else {
3210 let diffs = g_host.iter().zip(s_host.iter()).enumerate()
3211 .filter(|(_, (a, b))| a != b).count();
3212 let maxdiff = g_host.iter().zip(s_host.iter())
3213 .map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
3214 panic!("moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}", g_host.len());
3215 }
3216 }
3217 return Ok(grouped_out);
3218 }
3219 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
3220 }
3221
3222 fn sigmoid_resident_dev_eligible(
3223 e: &Engine,
3224 m: &MoeWeights,
3225 cfg: &ModelConfig,
3226 ) -> bool {
3227 let Some(moe) = cfg.moe.as_ref() else { return false };
3228 static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3231 let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
3232 std::env::var("MEMRA_MOE_STATS").is_ok()
3233 || std::env::var("MEMRA_MOE_TRACE").is_ok()
3234 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
3235 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
3236 || std::env::var("MEMRA_MOE_GATE").is_ok()
3237 });
3238 cfg.step35.is_some()
3239 && sigmoid_router_enabled()
3240 && moe_dev_enabled()
3241 && moe_slab_enabled()
3242 && !observation_mode
3243 && moe.expert_used_count <= 8
3244 && m.has_uniform_expert_layout()
3245 && m.gate_exps.macros.is_none()
3246 && m.up_exps.macros.is_none()
3247 && m.down_exps.macros.is_none()
3248 && !m.has_macros
3249 && moe_q8_enabled()
3250 && q8_expert_supported(m.gate_exps.qtype)
3251 && q8_expert_supported(m.up_exps.qtype)
3252 && q8_expert_supported(m.down_exps.qtype)
3253 && m.dev_exps
3254 .as_ref()
3255 .is_some_and(|dev| dev.dev == e.ctx().ordinal())
3256 }
3257
3258 pub(crate) fn moe_ffn_sequential(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
3260 cfg: &ModelConfig, il: u16, max_block: usize)
3261 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3262 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
3263 }
3264
3265 fn moe_router_logits(
3269 e: &Engine,
3270 m: &MoeWeights,
3271 z: &CudaSlice<f32>,
3272 t: usize,
3273 cfg: &ModelConfig,
3274 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3275 if t < PRIME_MIN_T {
3276 if crate::router_kernel_on() {
3278 e.router_gemv(
3279 m.gate_inp.float_data(),
3280 z,
3281 cfg.n_embd as usize,
3282 m.gate_exps.n_expert,
3283 t,
3284 )
3285 } else {
3286 e.matmul_decode_exact(&m.gate_inp, z, t)
3287 }
3288 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
3289 e.router_gemv(
3290 m.gate_inp.float_data(),
3291 z,
3292 cfg.n_embd as usize,
3293 m.gate_exps.n_expert,
3294 t,
3295 )
3296 } else {
3297 e.matmul(&m.gate_inp, z, t)
3298 }
3299 }
3300
3301 fn trace_moe_routes(il: u16, t: usize, sel_all: &[u32], weights: &[f32])
3305 -> Result<(), Box<dyn std::error::Error>> {
3306 use std::io::Write as _;
3307 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
3308 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
3309 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
3310 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
3311 }
3312 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
3313 let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
3314 let pairs: Vec<String> = sel_all.iter().zip(weights)
3315 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
3316 .collect();
3317 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
3318 }
3319 Ok(())
3320 }
3321
3322 #[allow(clippy::too_many_arguments)]
3323 fn trace_sigmoid_router_logits(
3324 e: &Engine,
3325 il: u16,
3326 t: usize,
3327 n_expert: usize,
3328 n_used: usize,
3329 logits: &CudaSlice<f32>,
3330 m: &MoeWeights,
3331 (scaling_factor, route_norm): (f32, bool),
3332 ) -> Result<(), Box<dyn std::error::Error>> {
3333 if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
3334 return Ok(());
3335 }
3336 let logits = e.dtoh(logits)?;
3337 let active: Vec<u8> = m
3338 .active_experts
3339 .as_ref()
3340 .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
3341 .unwrap_or_else(|| vec![1; n_expert]);
3342 let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
3343 crate::sigrouter_contract::capture_served_logits(
3344 il as u32,
3345 t,
3346 n_expert,
3347 n_used,
3348 scaling_factor,
3349 route_norm,
3350 &active,
3351 &bias,
3352 &logits,
3353 )?;
3354 Ok(())
3355 }
3356
3357 fn trace_moe_input(e: &Engine, il: u16, t: usize, n_embd: usize, z: &CudaSlice<f32>)
3362 -> Result<(), Box<dyn std::error::Error>> {
3363 use std::io::Write as _;
3364 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else { return Ok(()) };
3365 let host = e.dtoh(z)?;
3366 if host.len() != t * n_embd {
3367 return Err(format!(
3368 "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
3369 host.len(), t, n_embd
3370 ).into());
3371 }
3372 let bytes = unsafe {
3373 std::slice::from_raw_parts(
3374 host.as_ptr().cast::<u8>(), host.len() * std::mem::size_of::<f32>()
3375 )
3376 };
3377 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
3378 let mut state = state.lock().map_err(|_| "MoE input trace writer lock is poisoned")?;
3379 if state.is_none() {
3380 let dir = std::path::PathBuf::from(&dir);
3381 std::fs::create_dir_all(&dir)?;
3382 let index = std::fs::OpenOptions::new().create(true).append(true)
3383 .open(dir.join("index.jsonl"))?;
3384 *state = Some(MoeInputTraceWriter {
3385 dir,
3386 index,
3387 payloads: std::collections::HashMap::new(),
3388 });
3389 }
3390 let writer = state.as_mut().unwrap();
3391 if writer.dir != std::path::Path::new(&dir) {
3392 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
3393 }
3394 let file_name = format!("layer-{il:03}.f32");
3395 if !writer.payloads.contains_key(&il) {
3396 let payload = std::fs::OpenOptions::new().create(true).append(true)
3397 .open(writer.dir.join(&file_name))?;
3398 let offset = payload.metadata()?.len();
3399 writer.payloads.insert(il, (payload, offset));
3400 }
3401 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
3402 let row_offset = *offset;
3403 payload.write_all(bytes)?;
3404 *offset += bytes.len() as u64;
3405 writeln!(
3406 writer.index,
3407 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
3408 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
3409 \"payload_bytes\":{}}}",
3410 bytes.len()
3411 )?;
3412 Ok(())
3413 }
3414
3415 #[allow(clippy::too_many_arguments)]
3416 pub(crate) fn moe_ffn_sequential_zq8(
3417 e: &Engine,
3418 m: &MoeWeights,
3419 z: &CudaSlice<f32>,
3420 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
3421 t: usize,
3422 cfg: &ModelConfig,
3423 il: u16,
3424 max_block: usize,
3425 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3426 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
3427 let moe = cfg.moe.as_ref().unwrap();
3428 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);
3435 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
3436 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);
3439
3440 let lim_exp = cfg.clamp_exp_at(il as u32);
3443 let lim_shexp = cfg.clamp_shexp_at(il as u32);
3444 let use_cache = Engine::moe_cache_enabled();
3445 let uniform_experts = m.has_uniform_expert_layout();
3446 let moe_q8 = uniform_experts && moe_q8_enabled()
3447 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3448 && q8_expert_supported(m.down_exps.qtype);
3449 let cpu_expert_requested = crate::cpu_experts::configured();
3456 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
3457 return Err(std::io::Error::other(
3458 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
3459 )
3460 .into());
3461 }
3462 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
3463 let freeze_cpu_residency = cpu_expert_requested
3469 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
3470 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
3471 .ok()
3472 .and_then(|value| value.parse::<usize>().ok())
3473 .is_some_and(|tokens| tokens > 0);
3474 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
3475 e.freeze_moe_cache();
3476 }
3477 let cache_frozen = use_cache && e.moe_cache_frozen();
3478 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
3479
3480 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
3483 if let Some(sig) = cfg.sigmoid_router() {
3484 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
3485 }
3486
3487 let no_exp_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
3526 && m.down_exps.macros.is_none();
3527 if cfg.sigmoid_router().is_none() && cfg.m3.is_none() && cfg.hy3.is_none()
3531 && !cfg.swiglu_clamped_at(il as u32)
3532 && no_exp_macros
3533 && t >= PRIME_MIN_T && m.dev_exps.is_some() && moe_q8_enabled()
3534 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3535 && q8_expert_supported(m.down_exps.qtype)
3536 && std::env::var("MEMRA_MOE_PAIRS").map(|v| v != "0").unwrap_or(true)
3537 && std::env::var("MEMRA_MOE_STATS").is_err() {
3538 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
3539 }
3540
3541 let dev_ok = uniform_experts && cfg.sigmoid_router().is_none()
3559 && cfg.m3.is_none() && cfg.hy3.is_none()
3560 && !cfg.swiglu_clamped_at(il as u32);
3561 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
3565 || std::env::var("MEMRA_MOE_TRACE").is_ok()
3566 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
3567 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
3568 if dev_ok && t < PRIME_MIN_T && m.dev_exps.is_some() && n_used <= 8 && moe_dev_enabled()
3569 && !observe_routes {
3570 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
3571 }
3572 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled()
3573 && !observe_routes {
3574 let row_ok = e.with_moe_cache(max_block, |c, eng| {
3575 if moe_prewarm_enabled() { c.prewarm_layer(il, m, eng)?; }
3576 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
3577 })?;
3578 if row_ok {
3579 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
3580 }
3581 }
3582
3583 let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
3585 if cpu_hybrid {
3586 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
3587 e,
3588 &logits,
3589 z,
3590 t,
3591 n_expert,
3592 n_used,
3593 m.exp_probs_b.as_deref(),
3594 sig,
3595 m.active_experts.as_deref(),
3596 )?;
3597 (sel, w, Some(input))
3598 } else {
3599 let (sel, w) = Self::moe_route_sigmoid_cfg(
3600 e, &logits, t, n_expert, n_used, m, sig,
3601 )?;
3602 (sel, w, None)
3603 }
3604 } else {
3605 let (sel, w) = Self::moe_route_cfg(
3606 e, &logits, t, n_expert, n_used, m.active_experts.as_deref(),
3607 )?;
3608 (sel, w, None)
3609 };
3610 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
3611
3612 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
3616 Self::trace_moe_input(e, il, t, n_embd, z)?;
3617
3618 let worker_disk_prefetch =
3630 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
3631 let promote_worker_h2d =
3632 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
3633 if promote_worker_h2d {
3634 let mut selected_blocks = Vec::with_capacity(n_used * 3);
3635 for &ex in sel_all.iter().take(n_used) {
3636 let ex = ex as u16;
3637 selected_blocks.extend([
3638 BlockId::new(il, PROJ_GATE, ex),
3639 BlockId::new(il, PROJ_UP, ex),
3640 BlockId::new(il, PROJ_DOWN, ex),
3641 ]);
3642 }
3643 for &ex in sel_all.iter().take(n_used) {
3644 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
3645 }
3646 e.with_moe_cache(max_block, |cache, eng| {
3647 cache.promote_worker_reads_at_safe_boundary(
3648 &selected_blocks,
3649 &selected_blocks,
3650 eng,
3651 )?;
3652 Ok(())
3653 })?;
3654 }
3655
3656 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
3659 let mut cnt = vec![0u32; n_expert];
3660 for &s in sel_all.iter() { cnt[s as usize] += 1; }
3661 let total = sel_all.len() as f64;
3662 let mut h = 0.0f64;
3663 let mut active = 0usize;
3664 for &c in &cnt { if c > 0 { active += 1; let p = c as f64 / total; h -= p * p.log2(); } }
3665 let maxc = cnt.iter().copied().max().unwrap_or(0);
3666 println!("moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
3667 il, t, sel_all.len(), active, n_expert, h, (n_expert as f64).log2(), total / active.max(1) as f64, maxc);
3668 }
3669
3670 let gdec_may_fire = uniform_experts && use_cache && n_used <= 8 && gdec_enabled()
3683 && !cfg.swiglu_clamped_at(il as u32);
3684 let slab_local = m.dev_exps.as_ref()
3700 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
3701 let slab_bases = slab_local.map(|d| {
3702 use cudarc::driver::DevicePtr;
3703 let s = e.stream();
3704 let (pg, _g0) = d.gate.device_ptr(&s);
3705 let (pu, _g1) = d.up.device_ptr(&s);
3706 let (pd, _g2) = d.down.device_ptr(&s);
3707 (pg as u64, pu as u64, pd as u64)
3708 });
3709 let slab_fused_may_fire = slab_bases.is_some() && n_used <= 8 && gdec_enabled()
3719 && !cfg.swiglu_clamped_at(il as u32) && cfg.m3.is_none()
3720 && no_exp_macros && moe_q8;
3721 let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
3724 e.uninit(t * n_embd)?
3725 } else {
3726 e.zeros(t * n_embd)?
3727 };
3728 let cpu_input = if cpu_hybrid {
3731 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
3732 } else {
3733 None
3734 };
3735
3736 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;
3744 let mut scratch_u: Option<CudaSlice<u8>> = None;
3745 let mut scratch_d: Option<CudaSlice<u8>> = None;
3746 let page_window = moe_page_prefetch_window();
3754
3755 for tok in 0..t {
3758 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
3759 let w = &w_all[tok * n_used..(tok + 1) * n_used];
3760 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
3762
3763 let no_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
3777 && m.down_exps.macros.is_none();
3778 if slab_fused_may_fire {
3788 let (pg, pu, pd) = slab_bases.unwrap();
3789 let mut gp = [0u64; 8];
3790 let mut up = [0u64; 8];
3791 let mut dp = [0u64; 8];
3792 for (j, &ex) in sel.iter().enumerate() {
3793 let ex = ex as usize;
3794 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
3795 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
3796 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
3797 }
3798 let mut wv = [0f32; 8];
3799 wv[..n_used].copy_from_slice(w);
3800 if tok_q8.is_none() {
3801 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3802 }
3803 let (zq, zd) = tok_q8.as_ref().unwrap();
3804 let act = e.moe_gate_up_silu8_q8(crate::WPtr8(gp), crate::WPtr8(up), zq, zd,
3805 n_embd, n_ff_exp, n_used,
3806 m.gate_exps.qtype, m.up_exps.qtype,
3807 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
3808 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
3809 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3810 e.moe_down8_fma_q8(crate::WPtr8(dp), crate::F32x8(wv), &aq2, &ad2, &mut dst,
3811 n_ff_exp, n_embd, n_used,
3812 m.down_exps.qtype, m.down_exps.row_bytes)?;
3813 continue;
3814 }
3815 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
3816 if tok_q8.is_none() {
3817 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3818 }
3819 let (zq, zd) = tok_q8.as_ref().unwrap();
3820 if Self::moe_gdec_token_q8(e, m, il, max_block, zq, zd, sel, w,
3821 &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
3822 continue;
3823 }
3824 } else if gdec_may_fire && cfg.m3.is_none() && no_macros
3825 && Self::moe_gdec_token(e, m, il, max_block, &zt, sel, w,
3826 &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
3827 continue;
3828 }
3829
3830 if gdec_may_fire || slab_fused_may_fire {
3836 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3837 e.memset_zeros_view(&mut row)?;
3838 }
3839
3840 let mut cpu_mask = vec![false; sel.len()];
3846 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
3847 let gpu_resident = if use_cache {
3848 e.with_moe_cache(max_block, |cache, _| {
3849 Ok(sel
3850 .iter()
3851 .map(|&expert| {
3852 let expert = expert as u16;
3853 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
3854 .into_iter()
3855 .filter(|&projection| {
3856 cache
3857 .resident(BlockId::new(il, projection, expert))
3858 .is_some()
3859 })
3860 .count()
3861 })
3862 .collect::<Vec<_>>())
3863 })?
3864 } else {
3865 vec![0; sel.len()]
3866 };
3867 let mut cpu_selected = Vec::new();
3868 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
3869 if gpu_resident[index] != 3 {
3870 cpu_mask[index] = true;
3871 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
3872 let expert = expert as usize;
3873 cpu_selected.push((expert, route_weight));
3874 }
3875 }
3876 if crate::cpu_experts::predictor_enabled() {
3877 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
3881 crate::cpu_experts::predictor_submit(il, row);
3882 }
3883 if cpu_selected.is_empty() {
3884 None
3885 } else {
3886 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
3887 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
3888 .map_err(std::io::Error::other)?;
3889 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
3890 }
3891 } else {
3892 None
3893 };
3894
3895 let worker_window = worker_disk_prefetch
3896 .then(worker_prefetch_window)
3897 .unwrap_or(0);
3898 for (j, &ex) in sel.iter().enumerate() {
3899 if cpu_mask[j] {
3900 continue;
3901 }
3902 let ex = ex as usize;
3903 if let Some(d) = slab_local {
3910 let gl = m.gate_exps.expert_layout(ex);
3911 let ul = m.up_exps.expert_layout(ex);
3912 let dl = m.down_exps.expert_layout(ex);
3913 let (g0, u0, d0) = (ex * m.gate_exps.expert_stride,
3914 ex * m.up_exps.expert_stride,
3915 ex * m.down_exps.expert_stride);
3916 let (gate, up) = if moe_q8 {
3917 if tok_q8.is_none() {
3918 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3919 }
3920 let (zq, zd) = tok_q8.as_ref().unwrap();
3921 (e.qmatvec_expert_q8(&d.gate, g0..g0 + gl.len, zq, zd, 1,
3922 m.gate_exps.in_f, m.gate_exps.out_f,
3923 gl.qtype, gl.row_bytes)?,
3924 e.qmatvec_expert_q8(&d.up, u0..u0 + ul.len, zq, zd, 1,
3925 m.up_exps.in_f, m.up_exps.out_f,
3926 ul.qtype, ul.row_bytes)?)
3927 } else {
3928 (e.qmatvec_view(&d.gate, g0..g0 + gl.len, &zt, 1,
3929 m.gate_exps.in_f, m.gate_exps.out_f,
3930 gl.qtype, gl.row_bytes)?,
3931 e.qmatvec_view(&d.up, u0..u0 + ul.len, &zt, 1,
3932 m.up_exps.in_f, m.up_exps.out_f,
3933 ul.qtype, ul.row_bytes)?)
3934 };
3935 let mut act = e.uninit(n_ff_exp)?;
3936 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3937 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3938 let y = if moe_q8 {
3939 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
3940 e.qmatvec_expert_q8(&d.down, d0..d0 + dl.len, &aq2, &ad2, 1,
3941 m.down_exps.in_f, m.down_exps.out_f,
3942 dl.qtype, dl.row_bytes)?
3943 } else {
3944 let actv = act.slice(0..n_ff_exp);
3945 e.qmatvec_view(&d.down, d0..d0 + dl.len, &actv, 1,
3946 m.down_exps.in_f, m.down_exps.out_f,
3947 dl.qtype, dl.row_bytes)?
3948 };
3949 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3950 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3951 continue;
3952 }
3953 for next in page_prefetch_positions(j, sel.len(), page_window) {
3954 Self::moe_prefetch_host_expert(sel[next] as usize, m);
3955 }
3956 let keep = [
3957 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
3958 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
3959 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
3960 ];
3961 if worker_disk_prefetch && worker_window > 0 {
3962 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
3963 Self::moe_prefetch_disk_expert(
3964 e,
3965 il,
3966 sel[next] as usize,
3967 m,
3968 max_block,
3969 &keep,
3970 )?;
3971 }
3972 } else if cache_dispatch
3973 && !cpu_hybrid
3974 && moe_prefetch_enabled()
3975 && j + 1 < sel.len()
3976 {
3977 let next = sel[j + 1] as usize;
3978 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
3979 }
3980 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
3981 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
3982 if (gate_q8 || up_q8) && tok_q8.is_none() {
3985 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3986 }
3987 let gate = if gate_q8 {
3988 let (zq, zd) = tok_q8.as_ref().unwrap();
3989 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
3990 } else {
3991 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
3992 };
3993 let up = if up_q8 {
3994 let (zq, zd) = tok_q8.as_ref().unwrap();
3995 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
3996 } else {
3997 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
3998 };
3999 let mut act = e.uninit(n_ff_exp)?;
4000 Self::ffn_act_lim(
4001 e,
4002 cfg,
4003 &gate,
4004 &up,
4005 m.gate_exps.macro_scale(ex),
4006 m.up_exps.macro_scale(ex),
4007 lim_exp,
4008 &mut act,
4009 n_ff_exp,
4010 )?;
4011 let y = if down_q8 {
4012 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
4013 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
4014 } else {
4015 let actv = act.slice(0..n_ff_exp);
4016 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
4017 };
4018 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4019 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4021 } else if cache_dispatch {
4022 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
4027 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
4028 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4030 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
4031 let actv = act.slice(0..n_ff_exp);
4032 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
4033 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4034 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4036 } else if cache_frozen {
4037 let gate = Self::moe_frozen_gemm(
4042 e,
4043 il,
4044 PROJ_GATE,
4045 ex,
4046 m,
4047 max_block,
4048 &zt,
4049 &mut scratch_g,
4050 g_len,
4051 )?;
4052 let up = Self::moe_frozen_gemm(
4053 e,
4054 il,
4055 PROJ_UP,
4056 ex,
4057 m,
4058 max_block,
4059 &zt,
4060 &mut scratch_u,
4061 u_len,
4062 )?;
4063 let mut act = e.uninit(n_ff_exp)?;
4064 Self::ffn_act_lim(
4065 e,
4066 cfg,
4067 &gate,
4068 &up,
4069 m.gate_exps.macro_scale(ex),
4070 m.up_exps.macro_scale(ex),
4071 lim_exp,
4072 &mut act,
4073 n_ff_exp,
4074 )?;
4075 let actv = act.slice(0..n_ff_exp);
4076 let y = Self::moe_frozen_gemm(
4077 e,
4078 il,
4079 PROJ_DOWN,
4080 ex,
4081 m,
4082 max_block,
4083 &actv,
4084 &mut scratch_d,
4085 d_len,
4086 )?;
4087 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4088 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4089 } else {
4090 if scratch_g.is_none() {
4094 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
4095 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
4096 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
4097 }
4098 let (sg, su, sd) = (scratch_g.as_mut().unwrap(), scratch_u.as_mut().unwrap(),
4099 scratch_d.as_mut().unwrap());
4100 let gl = m.gate_exps.expert_layout(ex);
4101 let ul = m.up_exps.expert_layout(ex);
4102 let dl = m.down_exps.expert_layout(ex);
4103 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
4104 let gate = e.qmatvec_view(sg, 0..gl.len, &zt, 1,
4105 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
4106
4107 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
4108 let up = e.qmatvec_view(su, 0..ul.len, &zt, 1,
4109 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
4110
4111 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4113 m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
4114
4115 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
4116 let actv = act.slice(0..n_ff_exp);
4117 let y = e.qmatvec_view(sd, 0..dl.len, &actv, 1,
4118 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?;
4119
4120 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4121 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4122 }
4123 }
4124 if let Some(worker) = cpu_worker {
4125 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
4126 let cpu_output = e.htod(&cpu_output)?;
4127 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4128 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
4129 }
4130 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
4131 for (j, &ex) in sel.iter().enumerate() {
4132 if cpu_mask[j] {
4133 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
4134 }
4135 }
4136 }
4137 }
4138
4139 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4144 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4145 {
4146 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
4155 let (sg_gate, sg_up) = if t == 1 {
4156 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
4157 Some(pair) => pair,
4158 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
4159 }
4160 } else if verify_t {
4161 (e.matmul_decode_exact(gate_shexp, z, t)?, e.matmul_decode_exact(up_shexp, z, t)?)
4162 } else {
4163 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
4165 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)?;
4167 let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
4168 else { e.matmul(down_shexp, &sa, t)? }; let g = match &m.gate_inp_shexp {
4182 Some(gate_inp_shexp) => {
4183 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4184 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4185 } else {
4186 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4187 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
4189 g
4190 }
4191 }
4192 None => e.htod(&vec![1.0f32; t])?,
4193 };
4194 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4196 }
4197
4198 Ok(moe_out)
4199 }
4200
4201 pub fn stage1_h2d_per_token(&self) -> u64 {
4204 use crate::hybrid::Ffn;
4205 let n_used = self.cfg.moe.as_ref().map(|m| m.expert_used_count as u64).unwrap_or(0);
4206 let mut bytes = 0u64;
4207 for l in self.layers.iter() {
4208 if let Ffn::Moe(m) = &l.ffn {
4209 bytes += n_used * (m.gate_exps.max_expert_bytes() + m.up_exps.max_expert_bytes()
4210 + m.down_exps.max_expert_bytes()) as u64;
4211 }
4212 }
4213 bytes
4214 }
4215
4216 pub(crate) fn max_moe_block(&self) -> usize {
4220 use crate::hybrid::Ffn;
4221 let mut mx = 0usize;
4222 let mut scan = |ffn: &Ffn| {
4223 if let Ffn::Moe(m) = ffn {
4224 mx = mx.max(m.gate_exps.max_expert_bytes())
4225 .max(m.up_exps.max_expert_bytes())
4226 .max(m.down_exps.max_expert_bytes());
4227 }
4228 };
4229 for l in self.layers.iter() { scan(&l.ffn); }
4230 if let Some(mtp) = self.mtp.as_ref() { scan(&mtp.ffn); }
4231 mx
4232 }
4233
4234 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
4237 use crate::hybrid::Ffn;
4238 let mut sizes = Vec::new();
4239 let mut scan = |ffn: &Ffn| {
4240 let Ffn::Moe(m) = ffn else { return };
4241 for ex in 0..m.gate_exps.n_expert {
4242 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
4243 continue;
4244 }
4245 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
4246 let len = exps.expert_layout(ex).len;
4247 if len > 0 {
4248 sizes.push(len);
4249 }
4250 }
4251 }
4252 };
4253 for layer in &self.layers {
4254 scan(&layer.ffn);
4255 }
4256 if let Some(mtp) = &self.mtp {
4257 scan(&mtp.ffn);
4258 }
4259 sizes
4260 }
4261
4262 pub fn save_cpu_expert_residency_profile(
4268 &self,
4269 e: &Engine,
4270 path: &std::path::Path,
4271 ) -> Result<(), Box<dyn std::error::Error>> {
4272 let Some(ids) = e.export_moe_residency() else {
4273 return Err("no MoE residency cache to persist".into());
4274 };
4275 let mut body = format!(
4276 "memra-freeze-profile v1 max_block={} blocks={}\n",
4277 self.max_moe_block(),
4278 ids.len()
4279 );
4280 for (layer, proj, ex) in &ids {
4281 body.push_str(&format!("{layer} {proj} {ex}\n"));
4282 }
4283 let tmp = path.with_extension("tmp");
4284 std::fs::write(&tmp, body)?;
4285 std::fs::rename(&tmp, path)?;
4286 println!(
4287 "[moe-cache] freeze profile saved: {} blocks -> {}",
4288 ids.len(),
4289 path.display()
4290 );
4291 Ok(())
4292 }
4293
4294 pub fn restore_cpu_expert_residency_profile(
4298 &self,
4299 e: &Engine,
4300 path: &std::path::Path,
4301 ) -> Result<bool, Box<dyn std::error::Error>> {
4302 use crate::hybrid::Ffn;
4303 use crate::moe_cache::BlockId;
4304 let Ok(content) = std::fs::read_to_string(path) else {
4305 return Ok(false);
4306 };
4307 let mut lines = content.lines();
4308 let Some(header) = lines.next() else { return Ok(false) };
4309 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
4310 if !header.starts_with(&expected) {
4311 println!(
4312 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
4313 path.display()
4314 );
4315 return Ok(false);
4316 }
4317 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
4318 std::collections::HashMap::new();
4319 for line in lines {
4320 let mut fields = line.split_whitespace();
4321 let (Some(layer), Some(proj), Some(ex)) =
4322 (fields.next(), fields.next(), fields.next())
4323 else {
4324 continue;
4325 };
4326 let (Ok(layer), Ok(proj), Ok(ex)) =
4327 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
4328 else {
4329 continue;
4330 };
4331 by_layer
4332 .entry(layer)
4333 .or_default()
4334 .push(BlockId::new(layer, proj, ex));
4335 }
4336 let requested: usize = by_layer.values().map(Vec::len).sum();
4337 if requested == 0 {
4338 return Ok(false);
4339 }
4340 let max_block = self.max_moe_block();
4341 let mut restaged = 0usize;
4342 let mut stage_layer = |layer_index: u16,
4343 ffn: &Ffn|
4344 -> Result<(), Box<dyn std::error::Error>> {
4345 let Ffn::Moe(m) = ffn else { return Ok(()) };
4346 let Some(ids) = by_layer.get(&layer_index) else {
4347 return Ok(());
4348 };
4349 e.with_moe_cache(max_block, |cache, eng| {
4350 for id in ids {
4351 if cache.restage_block(*id, m, eng)? {
4352 restaged += 1;
4353 }
4354 }
4355 Ok(())
4356 })
4357 };
4358 for (index, layer) in self.layers.iter().enumerate() {
4359 stage_layer(index as u16, &layer.ffn)?;
4360 }
4361 if let Some(mtp) = self.mtp.as_ref() {
4362 stage_layer(u16::MAX, &mtp.ffn)?;
4363 }
4364 e.freeze_moe_cache();
4365 println!(
4366 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
4367 path.display()
4368 );
4369 Ok(true)
4370 }
4371
4372 pub fn freeze_cpu_expert_residency(
4374 &self,
4375 e: &Engine,
4376 ) -> Result<(), Box<dyn std::error::Error>> {
4377 e.freeze_moe_cache();
4378 Ok(())
4379 }
4380
4381 pub fn ffn_act(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4389 act: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
4390 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
4391 }
4392
4393 #[allow(clippy::too_many_arguments)]
4397 pub(crate) fn ffn_act_scaled(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4398 gs: f32, us: f32, act: &mut CudaSlice<f32>, n: usize)
4399 -> Result<(), Box<dyn std::error::Error>> {
4400 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
4401 }
4402
4403 #[allow(clippy::too_many_arguments)]
4412 pub(crate) fn ffn_act_lim(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4413 gs: f32, us: f32, limit: Option<f32>, act: &mut CudaSlice<f32>, n: usize)
4414 -> Result<(), Box<dyn std::error::Error>> {
4415 if let Some(m3) = cfg.m3.as_ref() {
4416 debug_assert!(limit.is_none(), "m3 swigluoai and step35 clamp are different archs");
4417 return e.swigluoai_mul_scaled(gate, up, gs, us, m3.swiglu_alpha, m3.swiglu_limit, act, n);
4418 }
4419 if let Some(l) = limit {
4420 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
4421 }
4422 if gs == 1.0 && us == 1.0 { return e.silu_mul(gate, up, act, n); }
4423 e.silu_mul_scaled(gate, up, gs, us, act, n)
4424 }
4425
4426 fn moe_route(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
4432 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4433 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
4434 }
4435
4436 #[allow(clippy::too_many_arguments)]
4444 fn moe_route_sigmoid_cfg(
4445 e: &Engine,
4446 logits: &CudaSlice<f32>,
4447 t: usize,
4448 n_expert: usize,
4449 n_used: usize,
4450 m: &MoeWeights,
4451 (sf, route_norm): (f32, bool),
4452 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4453 if sigmoid_router_enabled() {
4454 return e.moe_router_sigmoid_topk_host(
4455 logits,
4456 t,
4457 n_expert,
4458 n_used,
4459 m.active_count(),
4460 &m.exp_probs_b_dev,
4461 &m.active_experts_dev,
4462 sf,
4463 route_norm,
4464 );
4465 }
4466 let lg = e.dtoh(logits)?;
4467 Self::moe_route_sigmoid_host(
4468 &lg,
4469 t,
4470 n_expert,
4471 n_used,
4472 m.exp_probs_b.as_deref(),
4473 sf,
4474 route_norm,
4475 m.active_experts.as_deref(),
4476 )
4477 }
4478
4479 fn moe_route_cfg(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize,
4482 active: Option<&[bool]>)
4483 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4484 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
4487 return e.moe_router_topk_host(logits, t, n_expert, n_used);
4488 }
4489 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
4492 let mut w_out = vec![0f32; t * n_used];
4493 for tok in 0..t {
4494 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
4495 let maxl = row.iter().enumerate()
4497 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
4498 .map(|(_, &x)| x).fold(f32::NEG_INFINITY, f32::max);
4499 let mut probs = vec![0f32; n_expert];
4500 let mut den = 0f32;
4501 for i in 0..n_expert {
4502 if active.is_some_and(|mask| !mask[i]) { continue; }
4503 let x = (row[i] - maxl).exp(); probs[i] = x; den += x;
4504 }
4505 for p in probs.iter_mut() { *p /= den; }
4506 let mut idx: Vec<usize> = (0..n_expert)
4508 .filter(|&i| active.is_none_or(|mask| mask[i])).collect();
4509 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
4510 let sl = &idx[..n_used];
4511 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
4512 let mut ws: f32 = wv.iter().sum();
4513 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() { *x /= ws; }
4515 for j in 0..n_used {
4516 sel[tok * n_used + j] = sl[j] as u32;
4517 w_out[tok * n_used + j] = wv[j];
4518 }
4519 }
4520 Ok((sel, w_out))
4521 }
4522
4523 #[allow(clippy::too_many_arguments)]
4524 fn moe_route_sigmoid_with_input(
4525 e: &Engine,
4526 logits: &CudaSlice<f32>,
4527 input: &CudaSlice<f32>,
4528 t: usize,
4529 n_expert: usize,
4530 n_used: usize,
4531 bias: Option<&[f32]>,
4532 (sf, route_norm): (f32, bool),
4533 active: Option<&[bool]>,
4534 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
4535 let (lg, input) = e.dtoh_pair(logits, input)?;
4536 let (sel, w) =
4537 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
4538 Ok((sel, w, input))
4539 }
4540
4541 pub fn start_moe_prefetch_predictor(
4546 &self,
4547 e: &Engine,
4548 cfg: &ModelConfig,
4549 ) -> Result<(), Box<dyn std::error::Error>> {
4550 use crate::hybrid::Ffn;
4551 let Some(sig) = cfg.sigmoid_router() else {
4552 return Err("prefetch predictor requires a sigmoid-router arch".into());
4553 };
4554 let resident: std::collections::HashSet<(u16, u8, u16)> = e
4555 .export_moe_residency()
4556 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
4557 .into_iter()
4558 .collect();
4559 let mut layers = Vec::new();
4560 for (index, layer) in self.layers.iter().enumerate() {
4561 let Ffn::Moe(m) = &layer.ffn else { continue };
4562 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else { continue };
4563 let router = e.dtoh(data)?;
4564 let n_expert = m.gate_exps.n_expert;
4565 let n_embd = m.gate_exps.in_f;
4566 if router.len() != n_embd * n_expert {
4567 continue;
4568 }
4569 let build = |exps: &crate::model::HostExps| {
4570 (0..n_expert)
4571 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
4572 .collect::<Vec<_>>()
4573 };
4574 layers.push((index as u16, crate::cpu_experts::PredictLayerInit {
4575 router,
4576 bias: m.exp_probs_b.clone(),
4577 active: m.active_experts.clone(),
4578 n_embd,
4579 n_used: cfg
4580 .moe
4581 .as_ref()
4582 .map(|moe| moe.expert_used_count as usize)
4583 .ok_or("prefetch predictor requires MoE config")?,
4584 sig,
4585 weights_n_expert: n_expert,
4586 gate: build(&m.gate_exps),
4587 up: build(&m.up_exps),
4588 down: build(&m.down_exps),
4589 }));
4590 }
4591 crate::cpu_experts::start_prefetch_predictor(layers, resident)
4592 .map_err(|error| error.into())
4593 }
4594
4595 #[allow(clippy::too_many_arguments)]
4598 pub fn moe_route_sigmoid_host_public(
4599 logits: &[f32],
4600 t: usize,
4601 n_expert: usize,
4602 n_used: usize,
4603 bias: Option<&[f32]>,
4604 sf: f32,
4605 route_norm: bool,
4606 active: Option<&[bool]>,
4607 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4608 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
4609 }
4610
4611 #[allow(clippy::too_many_arguments)]
4612 fn moe_route_sigmoid_host(
4613 lg: &[f32],
4614 t: usize,
4615 n_expert: usize,
4616 n_used: usize,
4617 bias: Option<&[f32]>,
4618 sf: f32,
4619 route_norm: bool,
4620 active: Option<&[bool]>,
4621 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4622 let active_count = active
4623 .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
4624 .unwrap_or(n_expert);
4625 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4626 if lg.len() != t * n_expert {
4627 return Err(format!(
4628 "sigmoid router logits length mismatch: got {}, expected {}",
4629 lg.len(),
4630 t * n_expert,
4631 )
4632 .into());
4633 }
4634 let mut sel = vec![0u32; t * n_used];
4635 let mut w_out = vec![0f32; t * n_used];
4636 for tok in 0..t {
4637 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
4638 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
4639 let selsc: Vec<f32> = match bias {
4641 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
4642 None => scores.clone(),
4643 };
4644 let mut idx: Vec<usize> = (0..n_expert)
4645 .filter(|&i| active.is_none_or(|mask| mask[i]))
4646 .collect();
4647 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
4648 let sl = &idx[..n_used];
4649 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
4650 if route_norm {
4651 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
4652 for x in wv.iter_mut() {
4653 *x = *x / ws * sf;
4654 }
4655 } else {
4656 for x in wv.iter_mut() {
4657 *x *= sf;
4658 }
4659 }
4660 for j in 0..n_used {
4661 sel[tok * n_used + j] = sl[j] as u32;
4662 w_out[tok * n_used + j] = wv[j];
4663 }
4664 }
4665 Ok((sel, w_out))
4666 }
4667
4668 #[allow(clippy::too_many_arguments)]
4672 fn moe_ffn_sigmoid_dev(
4673 e: &Engine,
4674 m: &MoeWeights,
4675 z: &CudaSlice<f32>,
4676 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4677 logits: &CudaSlice<f32>,
4678 t: usize,
4679 cfg: &ModelConfig,
4680 il: u16,
4681 (scaling_factor, route_norm): (f32, bool),
4682 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4683 let moe = cfg.moe.as_ref().unwrap();
4684 let n_embd = cfg.n_embd as usize;
4685 let n_expert = moe.expert_count as usize;
4686 let n_used = moe.expert_used_count as usize;
4687 let n_ff_exp = moe.expert_ff_length as usize;
4688 let dev = m.dev_exps.as_ref().unwrap();
4689 debug_assert!(cfg.step35.is_some());
4690 debug_assert_eq!(dev.dev, e.ctx().ordinal());
4691 debug_assert!(m.has_uniform_expert_layout());
4692 debug_assert!(!m.has_macros);
4693
4694 let (sel_d, w_d) = e.moe_router_sigmoid_topk(
4695 logits,
4696 t,
4697 n_expert,
4698 n_used,
4699 m.active_count(),
4700 &m.exp_probs_b_dev,
4701 &m.active_experts_dev,
4702 scaling_factor,
4703 route_norm,
4704 )?;
4705 crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
4706 let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
4707 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
4708 (combined, combined)
4709 } else {
4710 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
4711 };
4712 let (zq, zd) = match (t, zq8) {
4713 (1, Some((q, d))) => (q.clone(), d.clone()),
4714 _ => e.quantize_q8_1(z, t, n_embd)?,
4715 };
4716 let n_pairs = t * n_used;
4717 let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
4718 let pair_tok: Vec<i32> = (0..n_pairs)
4722 .map(|pair| (pair / n_used) as i32)
4723 .collect();
4724 let pair_tok_d = e.htod_i32(&pair_tok)?;
4725 let gate = e.moe_pairs_matvec_q8(
4726 &dev.ptr_row,
4727 0,
4728 &pair_tok_d,
4729 &sel_d,
4730 &zq,
4731 &zd,
4732 n_embd,
4733 n_ff_exp,
4734 n_expert,
4735 n_pairs,
4736 m.gate_exps.qtype,
4737 gate_row_bytes,
4738 )?;
4739 let up = e.moe_pairs_matvec_q8(
4740 &dev.ptr_row,
4741 1,
4742 &pair_tok_d,
4743 &sel_d,
4744 &zq,
4745 &zd,
4746 n_embd,
4747 n_ff_exp,
4748 n_expert,
4749 n_pairs,
4750 m.up_exps.qtype,
4751 up_row_bytes,
4752 )?;
4753 let mut act = e.uninit(n_pairs * n_ff_exp)?;
4754 Self::ffn_act_lim(
4755 e,
4756 cfg,
4757 &gate,
4758 &up,
4759 1.0,
4760 1.0,
4761 cfg.clamp_exp_at(il as u32),
4762 &mut act,
4763 n_pairs * n_ff_exp,
4764 )?;
4765 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4766 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
4767 let pair_self_d = e.htod_i32(&pair_self)?;
4768 let down = e.moe_pairs_matvec_q8(
4769 &dev.ptr_row,
4770 2,
4771 &pair_self_d,
4772 &sel_d,
4773 &aq2,
4774 &ad2,
4775 n_ff_exp,
4776 n_embd,
4777 n_expert,
4778 n_pairs,
4779 m.down_exps.qtype,
4780 m.down_exps.row_bytes,
4781 )?;
4782 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
4783 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
4784 let tok_off_d = e.htod_i32(&tok_off)?;
4785 let tok_ids_d = e.htod_i32(&tok_ids)?;
4786 let mut output = e.uninit(t * n_embd)?;
4787 e.moe_pairs_scatter(
4788 &down,
4789 &w_d,
4790 &tok_off_d,
4791 &tok_ids_d,
4792 &mut output,
4793 t,
4794 n_embd,
4795 )?;
4796 output
4797 } else {
4798 let act = e.moe_gate_up_silu8_dev_q8_rows(
4799 &dev.ptr_row,
4800 &sel_d,
4801 &zq,
4802 &zd,
4803 t,
4804 n_embd,
4805 n_ff_exp,
4806 n_used,
4807 n_expert,
4808 m.gate_exps.qtype,
4809 m.up_exps.qtype,
4810 gate_row_bytes,
4811 up_row_bytes,
4812 &m.dev_macros,
4813 )?;
4814 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4815 let mut output = e.uninit(t * n_embd)?;
4816 e.moe_down8_fma_dev_q8_rows_g(
4817 &dev.ptr_row,
4818 &sel_d,
4819 &w_d,
4820 &aq2,
4821 &ad2,
4822 &mut output,
4823 t,
4824 n_ff_exp,
4825 n_embd,
4826 n_used,
4827 n_expert,
4828 m.down_exps.qtype,
4829 m.down_exps.row_bytes,
4830 )?;
4831 output
4832 };
4833
4834 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
4835 eprintln!(
4836 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
4837 cfg.clamp_exp_at(il as u32).is_some(),
4838 dev.gu_il,
4839 );
4840 }
4841 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
4842 Ok(moe_out)
4843 }
4844
4845 fn moe_ffn_pairs(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, logits: &CudaSlice<f32>,
4854 t: usize, cfg: &ModelConfig)
4855 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4856 let moe = cfg.moe.as_ref().unwrap();
4857 let n_embd = cfg.n_embd as usize;
4858 let n_expert = moe.expert_count as usize;
4859 let n_used = moe.expert_used_count as usize;
4860 let n_ff_exp = moe.expert_ff_length as usize;
4861 debug_assert!(!cfg.swiglu_clamped_anywhere(),
4866 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU");
4867 let dev = m.dev_exps.as_ref().unwrap();
4868 let (rbg_d, rbu_d) = if dev.gu_il {
4870 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
4871 } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
4872
4873 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
4874 let n_pairs = t * n_used;
4875 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
4878 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
4879 let pair_w: Vec<f32> = w_all.clone();
4880 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
4881 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
4882 let pt = e.htod_i32(&pair_tok)?;
4883 let px = e.htod_i32(&pair_ex)?;
4884 let pw = e.htod(&pair_w)?;
4885 let toff = e.htod_i32(&tok_off)?;
4886 let tids = e.htod_i32(&tok_ids)?;
4887
4888 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
4892 for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
4893 let mut ex_ids: Vec<i32> = Vec::new();
4894 let mut ex_off: Vec<i32> = vec![0];
4895 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
4896 for (ex, list) in by_ex.iter().enumerate() {
4897 if list.is_empty() { continue; }
4898 ex_ids.push(ex as i32);
4899 ex_pairs.extend_from_slice(list);
4900 ex_off.push(ex_pairs.len() as i32);
4901 }
4902 let n_active = ex_ids.len();
4903 let exi = e.htod_i32(&ex_ids)?;
4904 let exo = e.htod_i32(&ex_off)?;
4905 let exp_d = e.htod_i32(&ex_pairs)?;
4906 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4927 let mma_t = *MMA_T.get_or_init(|| {
4928 std::env::var("MEMRA_MOE_MMA_T").ok().and_then(|v| v.parse().ok()).unwrap_or(16)
4929 });
4930 let use_mma = std::env::var("MEMRA_MOE_MMA").map(|v| v != "0").unwrap_or(true)
4931 && t >= mma_t
4932 && q8_expert_dec_supported(m.gate_exps.qtype) && q8_expert_dec_supported(m.up_exps.qtype)
4933 && q8_expert_dec_supported(m.down_exps.qtype)
4934 && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
4935 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
4951 && q8_expert_dec_supported(m.up_exps.qtype)
4952 && q8_expert_dec_supported(m.down_exps.qtype)
4953 && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
4954 let f16g_mode = crate::moe_f16g_mode();
4955 let f16g = f16g_mode != 0 && t >= mma_t
4956 && (f16g_mode != 3 || !mma_capable)
4957 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
4958 && f16g_proj_ok(m.up_exps.qtype, n_embd)
4959 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
4960 if use_mma || f16g {
4961 let y_down = if f16g {
4969 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
4973 let csr_tok_d = e.htod_i32(&csr_tok)?;
4974 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
4975 let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
4976 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4977 m.gate_exps.qtype, rbg_d)?;
4978 let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
4979 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4980 m.up_exps.qtype, rbu_d)?;
4981 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
4982 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
4983 let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
4984 &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
4985 m.down_exps.qtype, m.down_exps.row_bytes)?;
4986 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
4987 } else {
4988 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
4990 let gate = e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4991 n_embd, n_ff_exp, n_active, n_pairs, t,
4992 m.gate_exps.qtype, rbg_d)?;
4993 let up = e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4994 n_embd, n_ff_exp, n_active, n_pairs, t,
4995 m.up_exps.qtype, rbu_d)?;
4996 let a_scr = if crate::moe_fuse_actq_on() {
5002 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
5003 } else {
5004 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5005 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
5006 };
5007 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5008 let pself = e.htod_i32(&pair_self)?;
5009 e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
5010 n_ff_exp, n_embd, n_active, n_pairs, n_pairs,
5011 m.down_exps.qtype, m.down_exps.row_bytes)?
5012 };
5013 let mut moe_out = e.uninit(t * n_embd)?;
5014 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5015 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5016 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5017 {
5018 let n_ff_sh = gate_shexp.out_features();
5019 let sg_gate = e.matmul(gate_shexp, z, t)?;
5020 let sg_up = e.matmul(up_shexp, z, t)?;
5021 let mut sa = e.uninit(t * n_ff_sh)?;
5022 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5023 let sh = e.matmul(down_shexp, &sa, t)?;
5024 let g = match &m.gate_inp_shexp {
5030 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
5031 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5032 }
5033 Some(gate_inp_shexp) => {
5034 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5035 let mut g = e.uninit(t)?;
5036 e.sigmoid(&gs, &mut g, t)?;
5037 g
5038 }
5039 None => e.htod(&vec![1.0f32; t])?,
5040 };
5041 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5042 }
5043 return Ok(moe_out);
5044 }
5045
5046 let dec = std::env::var("MEMRA_MOE_DEC").map(|v| v != "0").unwrap_or(true);
5049 let matvec = |proj, exi: &_, exo: &_, exp_d: &_, pt: &_, aq: &_, ad: &_,
5050 inf, outf, qtype, rb| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5051 let dec = dec && q8_expert_dec_supported(qtype);
5053 if dec { e.moe_pairs_matvec_q8_dec(&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
5054 inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
5055 else { e.moe_pairs_matvec_q8_em (&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
5056 inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
5057 };
5058 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5059 let gate = matvec(0, &exi, &exo, &exp_d, &pt, &zq, &zd,
5060 n_embd, n_ff_exp, m.gate_exps.qtype, rbg_d)?;
5061 let up = matvec(1, &exi, &exo, &exp_d, &pt, &zq, &zd,
5062 n_embd, n_ff_exp, m.up_exps.qtype, rbu_d)?;
5063 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5064 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5065 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5067 let pself = e.htod_i32(&pair_self)?;
5068 let y_down = matvec(2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
5069 n_ff_exp, n_embd, m.down_exps.qtype, m.down_exps.row_bytes)?;
5070 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5072
5073 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5077 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5078 {
5079 let n_ff_sh = gate_shexp.out_features();
5080 let step_exact = cfg.step35.is_some();
5084 let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
5085 let (sg_gate, sg_up) = if step_exact && t == 1 {
5086 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5087 Some(pair) => pair,
5088 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5089 }
5090 } else if verify_t {
5091 let mut fused = None;
5092 if crate::spec::spec_fused_t() && (2..=4).contains(&t)
5093 && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp)
5094 {
5095 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5096 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
5097 }
5098 match fused {
5099 Some(pair) => pair,
5100 None => (
5101 e.matmul_decode_exact(gate_shexp, z, t)?,
5102 e.matmul_decode_exact(up_shexp, z, t)?,
5103 ),
5104 }
5105 } else {
5106 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
5107 };
5108 let mut sa = e.uninit(t * n_ff_sh)?;
5109 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5110 let sh = if verify_t {
5111 e.matmul_decode_exact(down_shexp, &sa, t)?
5112 } else {
5113 e.matmul(down_shexp, &sa, t)?
5114 };
5115 let g = match &m.gate_inp_shexp {
5120 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
5121 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5122 }
5123 Some(gate_inp_shexp) => {
5124 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5125 let mut g = e.uninit(t)?;
5126 e.sigmoid(&gs, &mut g, t)?;
5127 g
5128 }
5129 None => e.htod(&vec![1.0f32; t])?,
5130 };
5131 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5132 }
5133 Ok(moe_out)
5134 }
5135
5136 #[allow(clippy::too_many_arguments)]
5138 #[allow(clippy::too_many_arguments)]
5139 fn moe_ffn_dev(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
5140 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, logits: &CudaSlice<f32>,
5141 t: usize, cfg: &ModelConfig, il: u16, max_block: usize)
5142 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5143 let moe = cfg.moe.as_ref().unwrap();
5144 let n_embd = cfg.n_embd as usize;
5145 let n_expert = moe.expert_count as usize;
5146 let n_used = moe.expert_used_count as usize;
5147 let n_ff_exp = moe.expert_ff_length as usize;
5148 debug_assert!(cfg.sigmoid_router().is_none(),
5152 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts");
5153 debug_assert!(!cfg.swiglu_clamped_at(il as u32),
5154 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form");
5155
5156 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
5158 if m.has_macros {
5161 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
5162 }
5163
5164 let mut moe_out = e.uninit(t * n_embd)?;
5166
5167 if let Some(dev) = m.dev_exps.as_ref() {
5170 let (rbg_d, rbu_d) = if dev.gu_il {
5173 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
5174 } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
5175 let q8 = moe_q8_enabled()
5176 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
5177 && q8_expert_supported(m.down_exps.qtype);
5178 let rows_arm = q8 && t > 1 && crate::spec::spec_m2()
5187 && n_ff_exp == 512 && n_used <= 8
5188 && std::env::var("MEMRA_MOE_DEVQ8_GU").map(|v| v.is_empty() || v == "v").unwrap_or(true)
5189 && std::env::var("MEMRA_MOE_DEVQ8_DOWN").map(|v| v.is_empty() || v == "w8h2v").unwrap_or(true);
5190 let csr_mode = std::env::var("MEMRA_MOE_CSR").ok()
5199 .and_then(|v| v.parse::<i32>().ok()).unwrap_or(1);
5200 let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
5201 let csr_arm = rows_arm && csr_mode > 0 && t <= 10
5202 && csr_qt(m.gate_exps.qtype) && csr_qt(m.up_exps.qtype)
5203 && csr_qt(m.down_exps.qtype);
5204 if csr_arm {
5205 if csr_mode == 2 {
5206 static ENGAGED: std::sync::Once = std::sync::Once::new();
5207 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
5208 }
5209 let n_pairs = t * n_used;
5210 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5211 let act = e.moe_gate_up_silu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, n_pairs,
5212 n_embd, n_ff_exp, n_used, n_expert,
5213 m.gate_exps.qtype, m.up_exps.qtype,
5214 rbg_d, rbu_d)?;
5215 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5216 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
5220 t, n_ff_exp, n_embd, n_used, n_expert,
5221 m.down_exps.qtype, m.down_exps.row_bytes)?;
5222 if csr_mode == 2 {
5223 let act_r = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
5225 n_embd, n_ff_exp, n_used, n_expert,
5226 m.gate_exps.qtype, m.up_exps.qtype,
5227 rbg_d, rbu_d, &m.dev_macros)?;
5228 let mut out_r = e.uninit(t * n_embd)?;
5229 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
5230 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2r, &ad2r, &mut out_r,
5231 t, n_ff_exp, n_embd, n_used, n_expert,
5232 m.down_exps.qtype, m.down_exps.row_bytes)?;
5233 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
5234 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
5235 let ba = a1.iter().zip(&a2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
5236 let bo = o1.iter().zip(&o2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
5237 if ba + bo > 0 {
5238 eprintln!("[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
5239 a1.len(), o1.len());
5240 let sel_h = e.dtoh_i32(&sel_d)?;
5242 let mut shown = 0;
5243 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
5244 if x.to_bits() != y.to_bits() && shown < 4 {
5245 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
5246 let ex = sel_h[p];
5247 let npx = sel_h.iter().filter(|&&v| v == ex).count();
5248 eprintln!(" ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}");
5249 shown += 1;
5250 }
5251 }
5252 std::process::exit(3);
5253 }
5254 }
5255 } else if rows_arm {
5256 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
5259 use std::sync::atomic::{AtomicU64, Ordering};
5260 static PAIRS: AtomicU64 = AtomicU64::new(0);
5261 static UNIQ: AtomicU64 = AtomicU64::new(0);
5262 static CALLS: AtomicU64 = AtomicU64::new(0);
5263 let sel_h = e.dtoh_i32(&sel_d)?;
5264 let mut u: Vec<i32> = sel_h.clone(); u.sort_unstable(); u.dedup();
5265 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
5266 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
5267 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
5268 if c % 480 == 0 {
5269 let p = PAIRS.load(Ordering::Relaxed); let q = UNIQ.load(Ordering::Relaxed);
5270 eprintln!("[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
5271 q as f64 / p as f64);
5272 }
5273 }
5274 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5275 let act = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
5276 n_embd, n_ff_exp, n_used, n_expert,
5277 m.gate_exps.qtype, m.up_exps.qtype,
5278 rbg_d, rbu_d, &m.dev_macros)?;
5279 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
5280 e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
5281 t, n_ff_exp, n_embd, n_used, n_expert,
5282 m.down_exps.qtype, m.down_exps.row_bytes)?;
5283 } else {
5284 for tok in 0..t {
5285 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
5286 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
5287 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
5288 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5289 if q8 {
5290 let (zq, zd) = match (t, zq8) {
5291 (1, Some((q, d))) => (q.clone(), d.clone()),
5292 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
5293 };
5294 let act = e.moe_gate_up_silu8_dev_q8(&dev.ptr_row, &selt, &zq, &zd,
5295 n_embd, n_ff_exp, n_used, n_expert,
5296 m.gate_exps.qtype, m.up_exps.qtype,
5297 rbg_d, rbu_d, &m.dev_macros)?;
5298 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5299 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selt, &wt, &aq2, &ad2, &mut dst,
5300 n_ff_exp, n_embd, n_used, n_expert,
5301 m.down_exps.qtype, m.down_exps.row_bytes)?;
5302 } else {
5303 let act = e.moe_gate_up_silu8_dev(&dev.ptr_row, &selt, &zt, n_embd, n_ff_exp,
5304 n_used, n_expert,
5305 m.gate_exps.qtype, m.up_exps.qtype,
5306 rbg_d, rbu_d, &m.dev_macros)?;
5307 e.moe_down8_fma_dev(&dev.ptr_row, &selt, &wt, &act, &mut dst,
5308 n_ff_exp, n_embd, n_used, n_expert,
5309 m.down_exps.qtype, m.down_exps.row_bytes)?;
5310 }
5311 }
5312 }
5313 } else {
5314 let q8 = moe_q8_enabled()
5321 && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
5322 && q8_expert_supported(m.down_exps.qtype);
5323 e.with_moe_cache(max_block, |c, eng| {
5324 let row = c.layer_dev_row(il, n_expert, eng)?
5325 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
5326 for tok in 0..t {
5327 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
5328 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
5329 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
5330 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5331 if q8 {
5332 let (zq, zd) = match (t, zq8) {
5333 (1, Some((q, d))) => (q.clone(), d.clone()),
5334 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
5335 };
5336 let act = eng.moe_gate_up_silu8_dev_q8(row, &selt, &zq, &zd,
5337 n_embd, n_ff_exp, n_used, n_expert,
5338 m.gate_exps.qtype, m.up_exps.qtype,
5339 m.gate_exps.row_bytes, m.up_exps.row_bytes,
5340 &m.dev_macros)?;
5341 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
5342 eng.moe_down8_fma_dev_q8(row, &selt, &wt, &aq2, &ad2, &mut dst,
5343 n_ff_exp, n_embd, n_used, n_expert,
5344 m.down_exps.qtype, m.down_exps.row_bytes)?;
5345 } else {
5346 let act = eng.moe_gate_up_silu8_dev(row, &selt, &zt, n_embd, n_ff_exp,
5347 n_used, n_expert,
5348 m.gate_exps.qtype, m.up_exps.qtype,
5349 m.gate_exps.row_bytes, m.up_exps.row_bytes,
5350 &m.dev_macros)?;
5351 eng.moe_down8_fma_dev(row, &selt, &wt, &act, &mut dst,
5352 n_ff_exp, n_embd, n_used, n_expert,
5353 m.down_exps.qtype, m.down_exps.row_bytes)?;
5354 }
5355 }
5356 c.hits += (t * 3 * n_used) as u64;
5358 Ok(())
5359 })?;
5360 }
5361
5362 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5367 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5368 {
5369 let n_ff_sh = gate_shexp.out_features();
5370 let verify_t = t > 1 && t < PRIME_MIN_T;
5373 let (sg_gate, sg_up) = if t == 1 {
5374 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5375 Some(pair) => pair,
5376 None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5377 }
5378 } else if verify_t {
5379 let mut fused = None;
5383 if crate::spec::spec_fused_t() && (2..=4).contains(&t)
5384 && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp) {
5385 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5386 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
5387 }
5388 match fused {
5389 Some(pair) => pair,
5390 None => (e.matmul_decode_exact(gate_shexp, z, t)?,
5391 e.matmul_decode_exact(up_shexp, z, t)?),
5392 }
5393 } else {
5394 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
5395 };
5396 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5398 let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
5399 else { e.matmul(down_shexp, &sa, t)? };
5400 let g = match &m.gate_inp_shexp {
5404 Some(gate_inp_shexp) => {
5405 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5408 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5409 } else {
5410 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5411 let mut g = e.uninit(t)?;
5412 e.sigmoid(&gs, &mut g, t)?;
5413 g
5414 }
5415 }
5416 None => e.htod(&vec![1.0f32; t])?,
5417 };
5418 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5419 }
5420
5421 Ok(moe_out)
5422 }
5423
5424 #[allow(clippy::too_many_arguments)]
5434 #[allow(clippy::too_many_arguments)]
5437 fn moe_gdec_token_q8(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
5438 zq: &CudaSlice<i8>, zd: &CudaSlice<f32>, sel: &[u32], w: &[f32],
5439 moe_out: &mut CudaSlice<f32>, tok: usize,
5440 n_embd: usize, n_ff_exp: usize, n_used: usize)
5441 -> Result<bool, Box<dyn std::error::Error>> {
5442 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
5443 use cudarc::driver::DevicePtr;
5444 let ptrs = e.with_moe_cache(max_block, |c, eng| {
5445 let mut g = [0u64; 8];
5446 let mut u = [0u64; 8];
5447 let mut d = [0u64; 8];
5448 for (j, &ex) in sel.iter().enumerate() {
5449 let ex = ex as u16;
5450 let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
5451 c.resident(BlockId::new(il, PROJ_UP, ex)),
5452 c.resident(BlockId::new(il, PROJ_DOWN, ex)))
5453 else { return Ok(None); };
5454 let __s = eng.stream();
5455 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
5456 let (pu, _e1) = c.slot(su).device_ptr(&__s);
5457 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
5458 g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
5459 }
5460 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
5461 for &ex in sel {
5462 let ex = ex as u16;
5463 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
5464 c.note_profile_hit(BlockId::new(il, proj, ex));
5465 }
5466 }
5467 }
5468 c.hits += (3 * n_used) as u64;
5469 Ok(Some((g, u, d)))
5470 })?;
5471 let Some((g, u, d)) = ptrs else { return Ok(false) };
5472 let mut wv = [0f32; 8];
5473 wv[..n_used].copy_from_slice(w);
5474 let act = e.moe_gate_up_silu8_q8(crate::WPtr8(g), crate::WPtr8(u), zq, zd,
5475 n_embd, n_ff_exp, n_used,
5476 m.gate_exps.qtype, m.up_exps.qtype,
5477 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5478 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5480 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5481 e.moe_down8_fma_q8(crate::WPtr8(d), crate::F32x8(wv), &aq2, &ad2, &mut dst,
5482 n_ff_exp, n_embd, n_used,
5483 m.down_exps.qtype, m.down_exps.row_bytes)?;
5484 Ok(true)
5485 }
5486
5487 fn moe_gdec_token(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
5488 zt: &cudarc::driver::CudaView<f32>, sel: &[u32], w: &[f32],
5489 moe_out: &mut CudaSlice<f32>, tok: usize,
5490 n_embd: usize, n_ff_exp: usize, n_used: usize)
5491 -> Result<bool, Box<dyn std::error::Error>> {
5492 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
5493 use cudarc::driver::DevicePtr;
5494 let ptrs = e.with_moe_cache(max_block, |c, eng| {
5496 let mut g = [0u64; 8];
5497 let mut u = [0u64; 8];
5498 let mut d = [0u64; 8];
5499 for (j, &ex) in sel.iter().enumerate() {
5500 let ex = ex as u16;
5501 let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
5502 c.resident(BlockId::new(il, PROJ_UP, ex)),
5503 c.resident(BlockId::new(il, PROJ_DOWN, ex)))
5504 else { return Ok(None); };
5505 let __s = eng.stream();
5506 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
5507 let (pu, _e1) = c.slot(su).device_ptr(&__s);
5508 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
5509 g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
5510 }
5511 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
5512 for &ex in sel {
5513 let ex = ex as u16;
5514 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
5515 c.note_profile_hit(BlockId::new(il, proj, ex));
5516 }
5517 }
5518 }
5519 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
5521 })?;
5522 let Some((g, u, d)) = ptrs else { return Ok(false) };
5523 let mut wv = [0f32; 8];
5524 wv[..n_used].copy_from_slice(w);
5525 let act = e.moe_gate_up_silu8(crate::WPtr8(g), crate::WPtr8(u), zt,
5527 n_embd, n_ff_exp, n_used,
5528 m.gate_exps.qtype, m.up_exps.qtype,
5529 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5530 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5531 e.moe_down8_fma_into(crate::WPtr8(d), crate::F32x8(wv), &act, &mut dst,
5532 n_ff_exp, n_embd, n_used,
5533 m.down_exps.qtype, m.down_exps.row_bytes)?;
5534 Ok(true)
5535 }
5536
5537 fn moe_cached_gemm_q8(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
5542 max_block: usize, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
5543 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5544 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
5545 let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
5546 let layout = exps.expert_layout(ex);
5547 let id = BlockId::new(il, proj, ex as u16);
5548 let source = exps.expert_source(ex);
5549 e.with_moe_cache(max_block, |c, eng| {
5550 let slot = c.dispatch_source(id, source, eng)?;
5551 let DispatchSlot::Resident(sl) = slot;
5552 let buf = c.slot(sl);
5553 eng.qmatvec_expert_q8(buf, 0..layout.len, aq, ad, 1, exps.in_f, exps.out_f,
5554 layout.qtype, layout.row_bytes)
5555 })
5556 }
5557
5558 fn moe_cached_gemm(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
5559 max_block: usize, x: &cudarc::driver::CudaView<f32>)
5560 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5561 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
5562 let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
5563 let layout = exps.expert_layout(ex);
5564 let id = BlockId::new(il, proj, ex as u16);
5565 let source = exps.expert_source(ex);
5566 e.with_moe_cache(max_block, |c, eng| {
5568 let slot = c.dispatch_source(id, source, eng)?;
5569 let DispatchSlot::Resident(sl) = slot;
5572 let buf = c.slot(sl);
5573 eng.qmatvec_view(buf, 0..layout.len, x, 1, exps.in_f, exps.out_f,
5574 layout.qtype, layout.row_bytes)
5575 })
5576 }
5577
5578 fn moe_profile_admit_expert(
5582 e: &Engine,
5583 il: u16,
5584 ex: usize,
5585 m: &MoeWeights,
5586 max_block: usize,
5587 ) -> Result<(), Box<dyn std::error::Error>> {
5588 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5589 e.with_moe_cache(max_block, |cache, eng| {
5590 for (proj, exps) in [
5591 (PROJ_GATE, &m.gate_exps),
5592 (PROJ_UP, &m.up_exps),
5593 (PROJ_DOWN, &m.down_exps),
5594 ] {
5595 let id = BlockId::new(il, proj, ex as u16);
5596 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
5597 }
5598 Ok(())
5599 })
5600 }
5601
5602 #[allow(clippy::too_many_arguments)]
5605 fn moe_frozen_gemm(
5606 e: &Engine,
5607 il: u16,
5608 proj: u8,
5609 ex: usize,
5610 m: &MoeWeights,
5611 max_block: usize,
5612 x: &cudarc::driver::CudaView<f32>,
5613 scratch: &mut Option<CudaSlice<u8>>,
5614 scratch_len: usize,
5615 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5616 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
5617 let exps = match proj {
5618 PROJ_GATE => &m.gate_exps,
5619 PROJ_UP => &m.up_exps,
5620 _ => &m.down_exps,
5621 };
5622 let layout = exps.expert_layout(ex);
5623 let id = BlockId::new(il, proj, ex as u16);
5624 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
5625 let Some(slot) = cache.resident(id) else {
5626 return Ok(None);
5627 };
5628 let buf = cache.slot(slot);
5629 Ok(Some(eng.qmatvec_view(
5630 buf,
5631 0..layout.len,
5632 x,
5633 1,
5634 exps.in_f,
5635 exps.out_f,
5636 layout.qtype,
5637 layout.row_bytes,
5638 )?))
5639 })? {
5640 return Ok(output);
5641 }
5642 if scratch.is_none() {
5643 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
5644 }
5645 let scratch = scratch.as_mut().unwrap();
5646 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
5647 e.qmatvec_view(
5648 scratch,
5649 0..layout.len,
5650 x,
5651 1,
5652 exps.in_f,
5653 exps.out_f,
5654 layout.qtype,
5655 layout.row_bytes,
5656 )
5657 }
5658
5659 fn moe_prefetch_expert(
5660 e: &Engine,
5661 il: u16,
5662 ex: usize,
5663 m: &MoeWeights,
5664 max_block: usize,
5665 keep: &[crate::moe_cache::BlockId],
5666 ) -> Result<(), Box<dyn std::error::Error>> {
5667 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5668 e.with_moe_cache(max_block, |c, eng| {
5669 for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
5670 (PROJ_DOWN, &m.down_exps)] {
5671 let id = BlockId::new(il, proj, ex as u16);
5672 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
5673 }
5674 Ok(())
5675 })
5676 }
5677
5678 fn moe_prefetch_disk_expert(e: &Engine, il: u16, ex: usize, m: &MoeWeights,
5681 max_block: usize, keep: &[crate::moe_cache::BlockId])
5682 -> Result<(), Box<dyn std::error::Error>> {
5683 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5684 e.with_moe_cache(max_block, |c, eng| {
5685 for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
5686 (PROJ_DOWN, &m.down_exps)] {
5687 let source = exps.expert_source(ex);
5688 if let crate::model::ExpertSource::Disk { .. } = &source {
5689 let id = BlockId::new(il, proj, ex as u16);
5690 let _ = c.prefetch_source(id, source, keep, eng)?;
5691 }
5692 }
5693 Ok(())
5694 })
5695 }
5696
5697 #[inline]
5698 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
5699 let _ = m.gate_exps.prefetch_expert_pages(ex);
5700 let _ = m.up_exps.prefetch_expert_pages(ex);
5701 let _ = m.down_exps.prefetch_expert_pages(ex);
5702 }
5703}
5704
5705impl HybridModel {
5722 #[allow(clippy::too_many_arguments)]
5726 fn moe_ffn_grouped_resident_q8(
5727 e: &Engine,
5728 m: &MoeWeights,
5729 z: &CudaSlice<f32>,
5730 t: usize,
5731 cfg: &ModelConfig,
5732 il: u16,
5733 sel_all: &[u32],
5734 w_all: &[f32],
5735 table: &CudaSlice<u64>,
5736 gu_il: bool,
5737 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5738 let moe = cfg.moe.as_ref().unwrap();
5739 let n_embd = cfg.n_embd as usize;
5740 let n_expert = moe.expert_count as usize;
5741 let n_used = moe.expert_used_count as usize;
5742 let n_ff_exp = moe.expert_ff_length as usize;
5743 let n_pairs = t * n_used;
5744 debug_assert_eq!(sel_all.len(), n_pairs);
5745 debug_assert_eq!(w_all.len(), n_pairs);
5746 debug_assert!(
5747 m.gate_exps.macros.is_none()
5748 && m.up_exps.macros.is_none()
5749 && m.down_exps.macros.is_none(),
5750 "resident grouped q8 does not fold per-expert macro scales",
5751 );
5752
5753 if !cfg.swiglu_clamped_at(il as u32) {
5759 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
5760 let sel_d = e.htod_i32(&sel)?;
5761 let w_d = e.htod(w_all)?;
5762 let (gate_row_bytes, up_row_bytes) = if gu_il {
5763 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5764 (combined, combined)
5765 } else {
5766 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5767 };
5768 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5769 let act = e.moe_gate_up_silu8_dev_q8_rows(
5770 table,
5771 &sel_d,
5772 &zq,
5773 &zd,
5774 t,
5775 n_embd,
5776 n_ff_exp,
5777 n_used,
5778 n_expert,
5779 m.gate_exps.qtype,
5780 m.up_exps.qtype,
5781 gate_row_bytes,
5782 up_row_bytes,
5783 &m.dev_macros,
5784 )?;
5785 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5786 let mut moe_out = e.uninit(t * n_embd)?;
5787 e.moe_down8_fma_dev_q8_rows_g(
5788 table,
5789 &sel_d,
5790 &w_d,
5791 &aq2,
5792 &ad2,
5793 &mut moe_out,
5794 t,
5795 n_ff_exp,
5796 n_embd,
5797 n_used,
5798 n_expert,
5799 m.down_exps.qtype,
5800 m.down_exps.row_bytes,
5801 )?;
5802
5803 if std::env::var("MEMRA_MOE_STATS").is_ok() {
5804 let mut counts = vec![0usize; n_expert];
5805 for &expert in sel_all {
5806 counts[expert as usize] += 1;
5807 }
5808 let mut sizes: Vec<usize> =
5809 counts.into_iter().filter(|&count| count != 0).collect();
5810 sizes.sort_unstable();
5811 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
5812 println!(
5813 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
5814 m_e: min={} median={} mean={mean:.1} max={}",
5815 sizes.len(),
5816 n_expert,
5817 sizes.first().copied().unwrap_or(0),
5818 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
5819 sizes.last().copied().unwrap_or(0),
5820 );
5821 }
5822 return Ok(moe_out);
5823 }
5824
5825 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5829 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
5830 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5831 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5832
5833 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
5834 for (pair, &expert) in pair_ex.iter().enumerate() {
5835 by_expert[expert as usize].push(pair as i32);
5836 }
5837
5838 let pair_tok_d = e.htod_i32(&pair_tok)?;
5839 let pair_ex_d = e.htod_i32(&pair_ex)?;
5840 let pair_w_d = e.htod(w_all)?;
5841 let tok_off_d = e.htod_i32(&tok_off)?;
5842 let tok_ids_d = e.htod_i32(&tok_ids)?;
5843
5844 let matvec = |
5845 proj: i32,
5846 pair_rows: &CudaSlice<i32>,
5847 aq: &CudaSlice<i8>,
5848 ad: &CudaSlice<f32>,
5849 in_f: usize,
5850 out_f: usize,
5851 qtype: i32,
5852 row_bytes: usize,
5853 | -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5854 e.moe_pairs_matvec_q8(
5855 table,
5856 proj,
5857 pair_rows,
5858 &pair_ex_d,
5859 aq,
5860 ad,
5861 in_f,
5862 out_f,
5863 n_expert,
5864 n_pairs,
5865 qtype,
5866 row_bytes,
5867 )
5868 };
5869
5870 let (gate_row_bytes, up_row_bytes) = if gu_il {
5871 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5872 (combined, combined)
5873 } else {
5874 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5875 };
5876 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5877 let gate = matvec(
5878 0,
5879 &pair_tok_d,
5880 &zq,
5881 &zd,
5882 n_embd,
5883 n_ff_exp,
5884 m.gate_exps.qtype,
5885 gate_row_bytes,
5886 )?;
5887 let up = matvec(
5888 1,
5889 &pair_tok_d,
5890 &zq,
5891 &zd,
5892 n_embd,
5893 n_ff_exp,
5894 m.up_exps.qtype,
5895 up_row_bytes,
5896 )?;
5897 let mut act = e.uninit(n_pairs * n_ff_exp)?;
5898 Self::ffn_act_lim(
5899 e,
5900 cfg,
5901 &gate,
5902 &up,
5903 1.0,
5904 1.0,
5905 cfg.clamp_exp_at(il as u32),
5906 &mut act,
5907 n_pairs * n_ff_exp,
5908 )?;
5909 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5910 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5911 let pair_self_d = e.htod_i32(&pair_self)?;
5912 let down = matvec(
5913 2,
5914 &pair_self_d,
5915 &aq2,
5916 &ad2,
5917 n_ff_exp,
5918 n_embd,
5919 m.down_exps.qtype,
5920 m.down_exps.row_bytes,
5921 )?;
5922 let mut moe_out = e.uninit(t * n_embd)?;
5923 e.moe_pairs_scatter(
5924 &down,
5925 &pair_w_d,
5926 &tok_off_d,
5927 &tok_ids_d,
5928 &mut moe_out,
5929 t,
5930 n_embd,
5931 )?;
5932
5933 if std::env::var("MEMRA_MOE_STATS").is_ok() {
5934 let mut sizes: Vec<usize> = by_expert
5935 .iter()
5936 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
5937 .collect();
5938 sizes.sort_unstable();
5939 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
5940 println!(
5941 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
5942 m_e: min={} median={} mean={mean:.1} max={}",
5943 sizes.len(),
5944 n_expert,
5945 sizes.first().copied().unwrap_or(0),
5946 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
5947 sizes.last().copied().unwrap_or(0),
5948 );
5949 }
5950 Ok(moe_out)
5951 }
5952
5953 fn moe_ffn_grouped_add_shared(
5954 e: &Engine,
5955 m: &MoeWeights,
5956 z: &CudaSlice<f32>,
5957 t: usize,
5958 cfg: &ModelConfig,
5959 il: u16,
5960 moe_out: &mut CudaSlice<f32>,
5961 ) -> Result<(), Box<dyn std::error::Error>> {
5962 let n_embd = cfg.n_embd as usize;
5963 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5964 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5965 {
5966 let n_ff_sh = gate_shexp.out_features();
5967 let sg_gate = e.matmul(gate_shexp, z, t)?;
5968 let sg_up = e.matmul(up_shexp, z, t)?;
5969 let mut sa = e.uninit(t * n_ff_sh)?;
5970 Self::ffn_act_lim(
5971 e,
5972 cfg,
5973 &sg_gate,
5974 &sg_up,
5975 1.0,
5976 1.0,
5977 cfg.clamp_shexp_at(il as u32),
5978 &mut sa,
5979 t * n_ff_sh,
5980 )?;
5981 let sh = e.matmul(down_shexp, &sa, t)?;
5982 let gate = match &m.gate_inp_shexp {
5983 Some(gate_inp_shexp) => {
5984 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5985 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5986 } else {
5987 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5988 let mut gate = e.uninit(t)?;
5989 e.sigmoid(&raw, &mut gate, t)?;
5990 gate
5991 }
5992 }
5993 None => e.htod(&vec![1.0f32; t])?,
5994 };
5995 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
5996 }
5997 Ok(())
5998 }
5999
6000 pub(crate) fn moe_ffn_grouped(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
6003 cfg: &ModelConfig, il: u16, max_block: usize)
6004 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6005 let moe = cfg.moe.as_ref().unwrap();
6006 let n_embd = cfg.n_embd as usize;
6007 let n_expert = moe.expert_count as usize;
6008 let n_used = moe.expert_used_count as usize;
6009 let n_ff_exp = moe.expert_ff_length as usize;
6010 let lim_exp = cfg.clamp_exp_at(il as u32);
6012
6013 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
6017 if let Some(sig) = cfg.sigmoid_router() {
6018 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
6019 }
6020 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
6021 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
6022 } else {
6023 Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
6024 m.active_experts.as_deref())?
6025 };
6026 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
6027 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
6028 Self::trace_moe_input(e, il, t, n_embd, z)?;
6029
6030 let no_exp_macros = m.gate_exps.macros.is_none()
6035 && m.up_exps.macros.is_none()
6036 && m.down_exps.macros.is_none();
6037 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
6038 m.has_uniform_expert_layout()
6039 && no_exp_macros
6040 && moe_q8_enabled()
6041 && q8_expert_supported(m.gate_exps.qtype)
6042 && q8_expert_supported(m.up_exps.qtype)
6043 && q8_expert_supported(m.down_exps.qtype)
6044 && moe_slab_enabled()
6045 && dev.dev == e.ctx().ordinal()
6046 });
6047 if let Some(dev) = resident_q8 {
6048 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
6049 e,
6050 m,
6051 z,
6052 t,
6053 cfg,
6054 il,
6055 &sel_all,
6056 &w_all,
6057 &dev.ptr_row,
6058 dev.gu_il,
6059 )?;
6060 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6061 return Ok(moe_out);
6062 }
6063
6064 struct ExpertGroup {
6068 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
6072 let mut groups: Vec<ExpertGroup> = (0..n_expert).map(|_| ExpertGroup {
6073 tok_indices: Vec::new(), slot_indices: Vec::new(), weights: Vec::new(),
6074 }).collect();
6075
6076 for tok in 0..t {
6077 for j in 0..n_used {
6078 let ex = sel_all[tok * n_used + j] as usize;
6079 let w = w_all[tok * n_used + j];
6080 groups[ex].tok_indices.push(tok as i32);
6081 groups[ex].slot_indices.push(j as i32);
6082 groups[ex].weights.push(w);
6083 }
6084 }
6085
6086 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
6089 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
6093 let u_len = m.up_exps.max_expert_bytes();
6094 let d_len = m.down_exps.max_expert_bytes();
6095 let moe_q8 = m.has_uniform_expert_layout()
6096 && moe_q8_enabled()
6097 && q8_expert_supported(m.gate_exps.qtype)
6098 && q8_expert_supported(m.up_exps.qtype)
6099 && q8_expert_supported(m.down_exps.qtype);
6100 let slab_local = m.dev_exps.as_ref().filter(|dev| {
6103 !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal()
6104 });
6105 let use_cache =
6106 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
6107 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
6110
6111 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
6113 (Some(e.alloc_u8(g_len)?), Some(e.alloc_u8(u_len)?), Some(e.alloc_u8(d_len)?))
6114 } else {
6115 (None, None, None)
6116 };
6117
6118 let mut order: Vec<usize> =
6129 (0..n_expert).filter(|&ex| !groups[ex].tok_indices.is_empty()).collect();
6130 order.sort_by(|&a, &b| groups[b].tok_indices.len()
6131 .cmp(&groups[a].tok_indices.len()).then(a.cmp(&b)));
6132 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
6134 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
6135 if worker_disk_prefetch {
6136 if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
6137 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
6138 }
6139 }
6140 for (order_pos, &ex) in order.iter().enumerate() {
6141 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
6142 Self::moe_prefetch_host_expert(order[next], m);
6143 }
6144 if worker_disk_prefetch {
6145 if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
6146 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6147 let keep = [
6148 BlockId::new(il, PROJ_GATE, ex as u16),
6149 BlockId::new(il, PROJ_UP, ex as u16),
6150 BlockId::new(il, PROJ_DOWN, ex as u16),
6151 ];
6152 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
6153 }
6154 }
6155 let grp = &groups[ex];
6156 let m_e = grp.tok_indices.len();
6157 m_dist.push(m_e);
6158 let gl = m.gate_exps.expert_layout(ex);
6159 let ul = m.up_exps.expert_layout(ex);
6160 let dl = m.down_exps.expert_layout(ex);
6161
6162 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
6166 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
6167 let dmac = m.down_exps.macro_scale(ex);
6168 let weight_d = if dmac == 1.0 { e.htod(&grp.weights)? } else {
6169 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
6170 e.htod(&scaled)?
6171 };
6172
6173 let mut gathered = e.zeros(m_e * n_embd)?;
6175 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
6176 let gv = gathered.slice(0..m_e * n_embd);
6177
6178 let y = if let Some(dev) = slab_local {
6181 let gate_start = ex * m.gate_exps.expert_stride;
6182 let up_start = ex * m.up_exps.expert_stride;
6183 let down_start = ex * m.down_exps.expert_stride;
6184 if grouped_q8 {
6185 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6186 let gate = e.qmatvec_expert_q8(
6187 &dev.gate,
6188 gate_start..gate_start + gl.len,
6189 &zq,
6190 &zd,
6191 m_e,
6192 m.gate_exps.in_f,
6193 m.gate_exps.out_f,
6194 gl.qtype,
6195 gl.row_bytes,
6196 )?;
6197 let up = e.qmatvec_expert_q8(
6198 &dev.up,
6199 up_start..up_start + ul.len,
6200 &zq,
6201 &zd,
6202 m_e,
6203 m.up_exps.in_f,
6204 m.up_exps.out_f,
6205 ul.qtype,
6206 ul.row_bytes,
6207 )?;
6208 let mut act = e.uninit(m_e * n_ff_exp)?;
6209 Self::ffn_act_lim(
6210 e,
6211 cfg,
6212 &gate,
6213 &up,
6214 m.gate_exps.macro_scale(ex),
6215 m.up_exps.macro_scale(ex),
6216 lim_exp,
6217 &mut act,
6218 m_e * n_ff_exp,
6219 )?;
6220 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6221 e.qmatvec_expert_q8(
6222 &dev.down,
6223 down_start..down_start + dl.len,
6224 &aq2,
6225 &ad2,
6226 m_e,
6227 m.down_exps.in_f,
6228 m.down_exps.out_f,
6229 dl.qtype,
6230 dl.row_bytes,
6231 )?
6232 } else {
6233 let gate = e.qmatvec_view(
6234 &dev.gate,
6235 gate_start..gate_start + gl.len,
6236 &gv,
6237 m_e,
6238 m.gate_exps.in_f,
6239 m.gate_exps.out_f,
6240 gl.qtype,
6241 gl.row_bytes,
6242 )?;
6243 let up = e.qmatvec_view(
6244 &dev.up,
6245 up_start..up_start + ul.len,
6246 &gv,
6247 m_e,
6248 m.up_exps.in_f,
6249 m.up_exps.out_f,
6250 ul.qtype,
6251 ul.row_bytes,
6252 )?;
6253 let mut act = e.uninit(m_e * n_ff_exp)?;
6254 Self::ffn_act_lim(
6255 e,
6256 cfg,
6257 &gate,
6258 &up,
6259 m.gate_exps.macro_scale(ex),
6260 m.up_exps.macro_scale(ex),
6261 lim_exp,
6262 &mut act,
6263 m_e * n_ff_exp,
6264 )?;
6265 let actv = act.slice(0..m_e * n_ff_exp);
6266 e.qmatvec_view(
6267 &dev.down,
6268 down_start..down_start + dl.len,
6269 &actv,
6270 m_e,
6271 m.down_exps.in_f,
6272 m.down_exps.out_f,
6273 dl.qtype,
6274 dl.row_bytes,
6275 )?
6276 }
6277 } else if use_cache {
6278 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6279 if grouped_q8 {
6280 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6281 let gate = e.with_moe_cache(max_block, |cache, eng| {
6282 let id = BlockId::new(il, PROJ_GATE, ex as u16);
6283 let slot =
6284 cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
6285 eng.qmatvec_expert_q8(
6286 cache.buf(slot),
6287 0..gl.len,
6288 &zq,
6289 &zd,
6290 m_e,
6291 m.gate_exps.in_f,
6292 m.gate_exps.out_f,
6293 gl.qtype,
6294 gl.row_bytes,
6295 )
6296 })?;
6297 let up = e.with_moe_cache(max_block, |cache, eng| {
6298 let id = BlockId::new(il, PROJ_UP, ex as u16);
6299 let slot =
6300 cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
6301 eng.qmatvec_expert_q8(
6302 cache.buf(slot),
6303 0..ul.len,
6304 &zq,
6305 &zd,
6306 m_e,
6307 m.up_exps.in_f,
6308 m.up_exps.out_f,
6309 ul.qtype,
6310 ul.row_bytes,
6311 )
6312 })?;
6313 let mut act = e.uninit(m_e * n_ff_exp)?;
6314 Self::ffn_act_lim(
6315 e,
6316 cfg,
6317 &gate,
6318 &up,
6319 m.gate_exps.macro_scale(ex),
6320 m.up_exps.macro_scale(ex),
6321 lim_exp,
6322 &mut act,
6323 m_e * n_ff_exp,
6324 )?;
6325 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6326 e.with_moe_cache(max_block, |cache, eng| {
6327 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
6328 let slot =
6329 cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
6330 eng.qmatvec_expert_q8(
6331 cache.buf(slot),
6332 0..dl.len,
6333 &aq2,
6334 &ad2,
6335 m_e,
6336 m.down_exps.in_f,
6337 m.down_exps.out_f,
6338 dl.qtype,
6339 dl.row_bytes,
6340 )
6341 })?
6342 } else {
6343 let gate = e.with_moe_cache(max_block, |cache, eng| {
6344 let id = BlockId::new(il, PROJ_GATE, ex as u16);
6345 let slot =
6346 cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
6347 eng.qmatvec_view(
6348 cache.buf(slot),
6349 0..gl.len,
6350 &gv,
6351 m_e,
6352 m.gate_exps.in_f,
6353 m.gate_exps.out_f,
6354 gl.qtype,
6355 gl.row_bytes,
6356 )
6357 })?;
6358 let up = e.with_moe_cache(max_block, |cache, eng| {
6359 let id = BlockId::new(il, PROJ_UP, ex as u16);
6360 let slot =
6361 cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
6362 eng.qmatvec_view(
6363 cache.buf(slot),
6364 0..ul.len,
6365 &gv,
6366 m_e,
6367 m.up_exps.in_f,
6368 m.up_exps.out_f,
6369 ul.qtype,
6370 ul.row_bytes,
6371 )
6372 })?;
6373 let mut act = e.uninit(m_e * n_ff_exp)?;
6374 Self::ffn_act_lim(
6375 e,
6376 cfg,
6377 &gate,
6378 &up,
6379 m.gate_exps.macro_scale(ex),
6380 m.up_exps.macro_scale(ex),
6381 lim_exp,
6382 &mut act,
6383 m_e * n_ff_exp,
6384 )?;
6385 let actv = act.slice(0..m_e * n_ff_exp);
6386 e.with_moe_cache(max_block, |cache, eng| {
6387 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
6388 let slot =
6389 cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
6390 eng.qmatvec_view(
6391 cache.buf(slot),
6392 0..dl.len,
6393 &actv,
6394 m_e,
6395 m.down_exps.in_f,
6396 m.down_exps.out_f,
6397 dl.qtype,
6398 dl.row_bytes,
6399 )
6400 })?
6401 }
6402 } else {
6403 let sg = scratch_g.as_mut().unwrap();
6404 let su = scratch_u.as_mut().unwrap();
6405 let sd = scratch_d.as_mut().unwrap();
6406 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
6407 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
6408 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
6409 if grouped_q8 {
6410 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6411 let gate = e.qmatvec_expert_q8(
6412 sg,
6413 0..gl.len,
6414 &zq,
6415 &zd,
6416 m_e,
6417 m.gate_exps.in_f,
6418 m.gate_exps.out_f,
6419 gl.qtype,
6420 gl.row_bytes,
6421 )?;
6422 let up = e.qmatvec_expert_q8(
6423 su,
6424 0..ul.len,
6425 &zq,
6426 &zd,
6427 m_e,
6428 m.up_exps.in_f,
6429 m.up_exps.out_f,
6430 ul.qtype,
6431 ul.row_bytes,
6432 )?;
6433 let mut act = e.uninit(m_e * n_ff_exp)?;
6434 Self::ffn_act_lim(
6435 e,
6436 cfg,
6437 &gate,
6438 &up,
6439 m.gate_exps.macro_scale(ex),
6440 m.up_exps.macro_scale(ex),
6441 lim_exp,
6442 &mut act,
6443 m_e * n_ff_exp,
6444 )?;
6445 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6446 e.qmatvec_expert_q8(
6447 sd,
6448 0..dl.len,
6449 &aq2,
6450 &ad2,
6451 m_e,
6452 m.down_exps.in_f,
6453 m.down_exps.out_f,
6454 dl.qtype,
6455 dl.row_bytes,
6456 )?
6457 } else {
6458 let gate = e.qmatvec_view(
6459 sg,
6460 0..gl.len,
6461 &gv,
6462 m_e,
6463 m.gate_exps.in_f,
6464 m.gate_exps.out_f,
6465 gl.qtype,
6466 gl.row_bytes,
6467 )?;
6468 let up = e.qmatvec_view(
6469 su,
6470 0..ul.len,
6471 &gv,
6472 m_e,
6473 m.up_exps.in_f,
6474 m.up_exps.out_f,
6475 ul.qtype,
6476 ul.row_bytes,
6477 )?;
6478 let mut act = e.uninit(m_e * n_ff_exp)?;
6479 Self::ffn_act_lim(
6480 e,
6481 cfg,
6482 &gate,
6483 &up,
6484 m.gate_exps.macro_scale(ex),
6485 m.up_exps.macro_scale(ex),
6486 lim_exp,
6487 &mut act,
6488 m_e * n_ff_exp,
6489 )?;
6490 let actv = act.slice(0..m_e * n_ff_exp);
6491 e.qmatvec_view(
6492 sd,
6493 0..dl.len,
6494 &actv,
6495 m_e,
6496 m.down_exps.in_f,
6497 m.down_exps.out_f,
6498 dl.qtype,
6499 dl.row_bytes,
6500 )?
6501 }
6502 };
6503
6504 e.scatter_slot(&y, &tok_idx_d, &slot_idx_d, &weight_d,
6506 &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
6507 }
6508
6509 let mut moe_out = e.zeros(t * n_embd)?;
6511 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
6512
6513 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
6515 m_dist.sort_unstable();
6516 let active = m_dist.len();
6517 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
6518 let median = m_dist[active / 2];
6519 let max_m = *m_dist.last().unwrap();
6520 let min_m = m_dist[0];
6521 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
6522 println!("moe-grouped il={il} t={t} active={active}/{n_expert} \
6523 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
6524 above_gemm_threshold(>=16)={above16}/{active}");
6525 }
6526
6527 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6528 Ok(moe_out)
6529 }
6530
6531 pub(crate) fn moe_ffn_lockstep(
6538 &self,
6539 e: &Engine,
6540 m: &MoeWeights,
6541 zbatch: &CudaSlice<f32>,
6542 mrows: usize,
6543 il: u16,
6544 max_block: usize,
6545 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6546 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6547 let cfg = &self.cfg;
6548 let moe = cfg.moe.as_ref().unwrap();
6549 let n_embd = cfg.n_embd as usize;
6550 let n_expert = moe.expert_count as usize;
6551 let n_used = moe.expert_used_count as usize;
6552 let n_ff_exp = moe.expert_ff_length as usize;
6553 let lim_exp = cfg.clamp_exp_at(il as u32);
6555 let lim_shexp = cfg.clamp_shexp_at(il as u32);
6556
6557 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
6558 if let Some(sig) = cfg.sigmoid_router() {
6559 Self::trace_sigmoid_router_logits(
6560 e, il, mrows, n_expert, n_used, &logits, m, sig,
6561 )?;
6562 }
6563 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
6564 Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
6565 } else {
6566 Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
6567 m.active_experts.as_deref())?
6568 };
6569 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
6570
6571 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
6573 Ok((0..n_expert)
6574 .map(|ex| {
6575 [PROJ_GATE, PROJ_UP, PROJ_DOWN].into_iter().all(|p| {
6576 c.resident(BlockId::new(il, p, ex as u16)).is_some()
6577 })
6578 })
6579 .collect())
6580 })?;
6581
6582 struct Group {
6583 rows: Vec<i32>,
6584 slots: Vec<i32>,
6585 weights: Vec<f32>,
6586 }
6587 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
6588 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
6589 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
6590 Default::default();
6591 for row in 0..mrows {
6592 for j in 0..n_used {
6593 let ex = sel_all[row * n_used + j] as usize;
6594 let w = w_all[row * n_used + j];
6595 if resident_expert[ex] {
6596 let group = groups.entry(ex).or_insert_with(|| Group {
6597 rows: Vec::new(),
6598 slots: Vec::new(),
6599 weights: Vec::new(),
6600 });
6601 group.rows.push(row as i32);
6602 group.slots.push(j as i32);
6603 group.weights.push(w);
6604 } else {
6605 crate::cpu_experts::record_incomplete_gpu_residency(0);
6606 cpu_rows[row].push((ex, w));
6607 cpu_by_expert.entry(ex).or_default().push((row, w));
6608 }
6609 }
6610 }
6611
6612 let host_rows = e.dtoh(zbatch)?;
6618 let rows_ok = crate::cpu_experts::rows_supported();
6619 enum CpuPart {
6620 Single { row: usize },
6621 Rows { rows: Vec<usize> },
6622 }
6623 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
6624 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
6625 if rows_ok {
6626 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
6627 .into_iter()
6628 .filter(|(_, rows)| rows.len() >= 2)
6629 .collect();
6630 shared.sort_by_key(|(ex, _)| *ex);
6631 for (ex, mut row_weights) in shared {
6632 row_weights.sort_by_key(|(row, _)| *row);
6633 let inputs: Vec<(&[f32], f32)> = row_weights
6634 .iter()
6635 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
6636 .collect();
6637 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
6638 .map_err(std::io::Error::other)?;
6639 for &(row, _) in &row_weights {
6640 rows_served.insert((row, ex));
6641 }
6642 tickets.push((
6643 CpuPart::Rows {
6644 rows: row_weights.iter().map(|&(row, _)| row).collect(),
6645 },
6646 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
6647 ));
6648 }
6649 }
6650 for (row, selected) in cpu_rows.iter().enumerate() {
6651 let leftover: Vec<(usize, f32)> = selected
6652 .iter()
6653 .copied()
6654 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
6655 .collect();
6656 if leftover.is_empty() {
6657 continue;
6658 }
6659 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
6660 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
6661 .map_err(std::io::Error::other)?;
6662 tickets.push((
6663 CpuPart::Single { row },
6664 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
6665 ));
6666 }
6667
6668 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
6669 let mut wbuf = e.zeros(mrows * n_used)?;
6670 let mut order: Vec<usize> = groups.keys().copied().collect();
6671 order.sort_by(|&a, &b| {
6672 groups[&b].rows.len().cmp(&groups[&a].rows.len()).then(a.cmp(&b))
6673 });
6674 for &ex in &order {
6675 let group = &groups[&ex];
6676 let m_e = group.rows.len();
6677 let gl = m.gate_exps.expert_layout(ex);
6678 let ul = m.up_exps.expert_layout(ex);
6679 let dl = m.down_exps.expert_layout(ex);
6680 let row_idx_d = e.htod_i32(&group.rows)?;
6681 let slot_idx_d = e.htod_i32(&group.slots)?;
6682 let dmac = m.down_exps.macro_scale(ex);
6683 let weight_d = if dmac == 1.0 {
6684 e.htod(&group.weights)?
6685 } else {
6686 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
6687 e.htod(&scaled)?
6688 };
6689 let mut gathered = e.zeros(m_e * n_embd)?;
6690 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
6691 let gv = gathered.slice(0..m_e * n_embd);
6692 let gate = e.with_moe_cache(max_block, |c, eng| {
6693 let slot = c
6694 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
6695 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6696 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..gl.len, &gv, m_e,
6697 m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
6698 })?;
6699 let up = e.with_moe_cache(max_block, |c, eng| {
6700 let slot = c
6701 .resident(BlockId::new(il, PROJ_UP, ex as u16))
6702 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6703 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..ul.len, &gv, m_e,
6704 m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
6705 })?;
6706 let mut act = e.zeros(m_e * n_ff_exp)?;
6707 Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
6708 m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
6709 let actv = act.slice(0..m_e * n_ff_exp);
6710 let y = e.with_moe_cache(max_block, |c, eng| {
6711 let slot = c
6712 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
6713 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6714 eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..dl.len, &actv, m_e,
6715 m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
6716 })?;
6717 e.scatter_slot(&y, &row_idx_d, &slot_idx_d, &weight_d,
6718 &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
6719 }
6720 let mut moe_out = e.zeros(mrows * n_embd)?;
6721 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
6722
6723 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
6725 for (part, ticket) in tickets {
6726 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
6727 let mut add_row = |row: usize, chunk: &[f32]| {
6728 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
6729 for (accumulator, value) in sum.iter_mut().zip(chunk) {
6730 *accumulator += value;
6731 }
6732 };
6733 match part {
6734 CpuPart::Single { row } => add_row(row, &cpu_output),
6735 CpuPart::Rows { rows } => {
6736 for (slot, row) in rows.into_iter().enumerate() {
6737 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
6738 }
6739 }
6740 }
6741 }
6742 for (row, sum) in row_sums.into_iter().enumerate() {
6743 let Some(sum) = sum else { continue };
6744 let cpu_output = e.htod(&sum)?;
6745 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
6746 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
6747 }
6748
6749 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6750 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6751 {
6752 let n_ff_sh = gate_shexp.out_features();
6753 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
6754 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
6755 let mut sa = e.zeros(mrows * n_ff_sh)?;
6756 Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp,
6757 &mut sa, mrows * n_ff_sh)?;
6758 let sh = e.matmul(down_shexp, &sa, mrows)?;
6759 let g = match &m.gate_inp_shexp {
6762 Some(gate_inp_shexp) => {
6763 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
6764 }
6765 None => e.htod(&vec![1.0f32; mrows])?,
6766 };
6767 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
6768 }
6769
6770 Ok(moe_out)
6771 }
6772}
6773
6774impl HybridModel {
6780 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
6782 let g = self.cfg.gemma4.as_ref().unwrap();
6783 let swa = g.swa_pattern[il];
6784 let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
6785 (hd, g.head_count_kv[il] as usize, self.cfg.n_head as usize,
6789 if swa { g.rope_base_swa } else { g.rope_base_global },
6790 1.0, swa)
6791 }
6792
6793 fn gemma4_suppress(&self, e: &Engine, ld: &mut CudaSlice<f32>, t: usize)
6797 -> Result<(), Box<dyn std::error::Error>> {
6798 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
6799 #[cfg(debug_assertions)]
6804 crate::debug_assert_tensor_stream_device(ids, &e.stream(),
6805 "gemma4_suppress.suppress_d");
6806 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
6807 }
6808 Ok(())
6809 }
6810
6811 fn gemma4_attn_prime(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6816 h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
6817 cache: Option<&mut Cache>)
6818 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6819 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6820 let eps = self.cfg.rms_eps;
6821 let aux = self.gemma4_aux.as_ref().unwrap();
6822 let ones = aux.ones(e);
6823 #[cfg(debug_assertions)]
6824 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
6825 "gemma4_attn_prime.ones");
6826
6827 e.mmq_act_begin();
6830 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)? };
6835
6836 let mut q = e.uninit(t * nh * hd)?;
6837 let mut k = e.uninit(t * nkv * hd)?;
6838 let mut v = e.uninit(t * nkv * hd)?;
6840 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6844 let emit = t >= 16 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
6845 && *EMIT.get_or_init(|| std::env::var("MEMRA_FA_EMIT").map(|s| s != "0").unwrap_or(true));
6846 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
6847 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
6848 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
6849 let v_f16 = emit && crate::fa_f16pv_on() && match hd {
6852 512 => true,
6853 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
6854 _ => false,
6855 };
6856 if emit {
6857 e.rms_norm_qkv_w4b(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6858 ones, &mut q, &mut k, &mut v, &mut vb,
6859 hd, nh * t, nkv * t, eps, v_f16)?;
6860 } else {
6861 e.rms_norm_qkv(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6862 ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t, eps)?;
6863 }
6864
6865 let ff = if swa { None } else {
6866 Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
6867 };
6868 #[cfg(debug_assertions)]
6869 if let Some(ff) = ff {
6870 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
6871 "gemma4_attn_prime.rope_freqs");
6872 }
6873 if emit {
6874 e.rope_neox2_bf16e(&mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t,
6875 base, 1.0, ff)?;
6876 } else {
6877 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
6878 }
6879
6880 if let Some(cache) = cache {
6881 let kvl = cache.kv[il].as_mut().unwrap();
6882 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
6883 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
6884 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()))?;
6885 kvl.len += t;
6886 }
6887 let mut attn = e.zeros(t * nh * hd)?;
6888 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6892 if swa && t > win {
6893 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
6894 if emit { e.fa_prefill_w_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
6895 scale, true, win, v_f16)?; }
6896 else { e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true,
6897 win)?; }
6898 } else {
6899 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
6900 }
6901 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
6902 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6903 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
6904 if emit { e.fa_prefill_hd512_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
6905 scale, true, v_f16)?; }
6906 else { e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?; }
6907 } else {
6908 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6909 }
6910 Ok(e.matmul(&fa.wo, &attn, t)?)
6911 }
6912
6913 fn gemma4_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6915 h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
6916 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6917 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
6918 }
6919
6920 fn gemma4_moe_q8(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
6925 bits: &crate::hybrid::Gemma4MoeBits,
6926 mq: &(CudaSlice<i8>, CudaSlice<f32>),
6927 router_in: &CudaSlice<f32>, t: usize)
6928 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6929 let cfg = &self.cfg;
6930 let moe = cfg.moe.as_ref().unwrap();
6931 let n_embd = cfg.n_embd as usize;
6932 let n_expert = moe.expert_count as usize;
6933 let n_used = moe.expert_used_count as usize;
6934 let n_ff_exp = moe.expert_ff_length as usize;
6935 let logits = if crate::router_kernel_on() {
6939 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
6940 } else {
6941 e.matmul(&m.gate_inp, router_in, t)?
6942 };
6943 let dev = m.dev_exps.as_ref().unwrap();
6944 let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
6945 &bits.per_expert_scale_d)?;
6946 let (zq, zd) = mq;
6947 if t == 1 {
6948 let selv = sel_d.slice(0..n_used);
6949 let wv = w_d.slice(0..n_used);
6950 let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, zq, zd,
6951 n_embd, n_ff_exp, n_used, n_expert,
6952 m.gate_exps.qtype, m.up_exps.qtype,
6953 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
6954 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6955 let mut moe_out = e.uninit(n_embd)?;
6956 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
6957 &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
6958 n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
6959 return Ok(moe_out);
6960 }
6961 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
6962 let act = if csr {
6963 e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, zq, zd, t * n_used,
6964 n_embd, n_ff_exp, n_used, n_expert,
6965 m.gate_exps.qtype, m.up_exps.qtype,
6966 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6967 } else {
6968 e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, zq, zd, t,
6969 n_embd, n_ff_exp, n_used, n_expert,
6970 m.gate_exps.qtype, m.up_exps.qtype,
6971 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6972 };
6973 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6974 let mut moe_out = e.uninit(t * n_embd)?;
6975 e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
6978 n_ff_exp, n_embd, n_used, n_expert,
6979 m.down_exps.qtype, m.down_exps.row_bytes)?;
6980 Ok(moe_out)
6981 }
6982
6983 fn gemma4_moe(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
6987 bits: &crate::hybrid::Gemma4MoeBits, moe_in: &CudaSlice<f32>,
6988 router_in: &CudaSlice<f32>, t: usize)
6989 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6990 let cfg = &self.cfg;
6991 let moe = cfg.moe.as_ref().unwrap();
6992 let n_embd = cfg.n_embd as usize;
6993 let n_expert = moe.expert_count as usize;
6994 let n_used = moe.expert_used_count as usize;
6995 let n_ff_exp = moe.expert_ff_length as usize;
6996
6997 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
7001 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
7002 } else {
7003 e.matmul(&m.gate_inp, router_in, t)?
7004 };
7005
7006 if t < PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7011 && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
7012 && expert_dp4a_supported(m.down_exps.qtype)
7013 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0") {
7014 let dev = m.dev_exps.as_ref().unwrap();
7015 let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
7016 &bits.per_expert_scale_d)?;
7017 if t == 1 {
7018 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
7019 let selv = sel_d.slice(0..n_used);
7020 let wv = w_d.slice(0..n_used);
7021 let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, &zq, &zd,
7022 n_embd, n_ff_exp, n_used, n_expert,
7023 m.gate_exps.qtype, m.up_exps.qtype,
7024 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
7025 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7026 let mut moe_out = e.uninit(n_embd)?;
7027 e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
7028 &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
7029 n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
7030 return Ok(moe_out);
7031 }
7032 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
7037 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
7038 let act = if csr {
7039 e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, t * n_used,
7040 n_embd, n_ff_exp, n_used, n_expert,
7041 m.gate_exps.qtype, m.up_exps.qtype,
7042 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
7043 } else {
7044 e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
7045 n_embd, n_ff_exp, n_used, n_expert,
7046 m.gate_exps.qtype, m.up_exps.qtype,
7047 m.gate_exps.row_bytes, m.up_exps.row_bytes)?
7048 };
7049 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
7050 let mut moe_out = e.uninit(t * n_embd)?;
7051 e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
7052 n_ff_exp, n_embd, n_used, n_expert,
7053 m.down_exps.qtype, m.down_exps.row_bytes)?;
7054 return Ok(moe_out);
7055 }
7056
7057 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
7058 for (i, &sx) in sel_all.iter().enumerate() {
7059 w_all[i] *= bits.per_expert_scale[sx as usize];
7060 }
7061
7062 if t >= PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7066 && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
7067 && expert_dp4a_supported(m.down_exps.qtype)
7068 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0") {
7069 let dev = m.dev_exps.as_ref().unwrap();
7070 let n_pairs = t * n_used;
7071 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
7072 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
7073 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7074 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7075 let pt = e.htod_i32(&pair_tok)?;
7076 let pw = e.htod(&w_all)?;
7077 let toff = e.htod_i32(&tok_off)?;
7078 let tids = e.htod_i32(&tok_ids)?;
7079 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7080 for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
7081 let mut ex_ids: Vec<i32> = Vec::new();
7082 let mut ex_off: Vec<i32> = vec![0];
7083 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
7084 for (ex, list) in by_ex.iter().enumerate() {
7085 if list.is_empty() { continue; }
7086 ex_ids.push(ex as i32);
7087 ex_pairs.extend_from_slice(list);
7088 ex_off.push(ex_pairs.len() as i32);
7089 }
7090 let n_active = ex_ids.len();
7091 let exi = e.htod_i32(&ex_ids)?;
7092 let exo = e.htod_i32(&ex_off)?;
7093 let exp_d = e.htod_i32(&ex_pairs)?;
7094 if crate::moe_f16g_gemma_on()
7102 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
7103 && f16g_proj_ok(m.up_exps.qtype, n_embd)
7104 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp) {
7105 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
7106 let csr_tok_d = e.htod_i32(&csr_tok)?;
7107 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
7108 let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
7109 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
7110 m.gate_exps.qtype, m.gate_exps.row_bytes)?;
7111 let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
7112 &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
7113 m.up_exps.qtype, m.up_exps.row_bytes)?;
7114 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
7115 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
7116 let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
7117 &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
7118 m.down_exps.qtype, m.down_exps.row_bytes)?;
7119 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
7120 let mut moe_out = e.uninit(t * n_embd)?;
7121 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
7122 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
7123 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
7124 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
7125 eprintln!("[f16g-debug] post-permute bad={} post-scatter bad={}",
7126 scan(&yd), scan(&mo));
7127 }
7128 return Ok(moe_out);
7129 }
7130 let mma = n_embd % 256 == 0
7133 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
7134 let (gate, up) = if mma {
7135 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
7136 (e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
7137 n_embd, n_ff_exp, n_active, n_pairs, t,
7138 m.gate_exps.qtype, m.gate_exps.row_bytes)?,
7139 e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
7140 n_embd, n_ff_exp, n_active, n_pairs, t,
7141 m.up_exps.qtype, m.up_exps.row_bytes)?)
7142 } else {
7143 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
7144 (e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 0, &exi, &exo, &exp_d, &pt, &zq, &zd,
7145 n_embd, n_ff_exp, n_expert, n_active, n_pairs,
7146 m.gate_exps.qtype, m.gate_exps.row_bytes)?,
7147 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 1, &exi, &exo, &exp_d, &pt, &zq, &zd,
7148 n_embd, n_ff_exp, n_expert, n_active, n_pairs,
7149 m.up_exps.qtype, m.up_exps.row_bytes)?)
7150 };
7151 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7152 let pself = e.htod_i32(&pair_self)?;
7153 let y_down = if mma {
7165 let in_pad = n_ff_exp.div_ceil(256) * 256;
7166 let a_scr = if crate::moe_fuse_actq_on() {
7167 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
7168 } else {
7169 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
7170 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
7171 };
7172 e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
7173 in_pad, n_embd, n_active, n_pairs, n_pairs,
7174 m.down_exps.qtype, m.down_exps.row_bytes)?
7175 } else {
7176 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
7177 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7178 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
7179 n_ff_exp, n_embd, n_expert, n_active, n_pairs,
7180 m.down_exps.qtype, m.down_exps.row_bytes)?
7181 };
7182 let mut moe_out = e.uninit(t * n_embd)?;
7183 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
7184 return Ok(moe_out);
7185 }
7186
7187 let g_len = m.gate_exps.expert_stride;
7188 let u_len = m.up_exps.expert_stride;
7189 let d_len = m.down_exps.expert_stride;
7190 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
7194 let (mut sg, mut su, mut sd) = if dev.is_some() { (None, None, None) } else {
7195 (Some(e.alloc_u8_uninit(g_len)?), Some(e.alloc_u8_uninit(u_len)?), Some(e.alloc_u8_uninit(d_len)?))
7196 };
7197 let mut moe_out = e.zeros(t * n_embd)?;
7198 for tok in 0..t {
7199 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
7200 let w = &w_all[tok * n_used..(tok + 1) * n_used];
7201 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
7202 for (j, &ex) in sel.iter().enumerate() {
7203 let ex = ex as usize;
7204 let gate = match dev {
7205 Some(d) => e.qmatvec_view(&d.gate, ex * g_len..(ex + 1) * g_len, &zt, 1,
7206 m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?,
7207 None => {
7208 let sg = sg.as_mut().unwrap();
7209 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
7210 e.qmatvec_view(sg, 0..g_len, &zt, 1,
7211 m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?
7212 }
7213 };
7214 let up = match dev {
7215 Some(d) => e.qmatvec_view(&d.up, ex * u_len..(ex + 1) * u_len, &zt, 1,
7216 m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?,
7217 None => {
7218 let su = su.as_mut().unwrap();
7219 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
7220 e.qmatvec_view(su, 0..u_len, &zt, 1,
7221 m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?
7222 }
7223 };
7224 let mut act = e.uninit(n_ff_exp)?;
7225 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
7226 let actv = act.slice(0..n_ff_exp);
7227 let y = match dev {
7228 Some(d) => e.qmatvec_view(&d.down, ex * d_len..(ex + 1) * d_len, &actv, 1,
7229 m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?,
7230 None => {
7231 let sd = sd.as_mut().unwrap();
7232 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
7233 e.qmatvec_view(sd, 0..d_len, &actv, 1,
7234 m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?
7235 }
7236 };
7237 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7238 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
7239 }
7240 }
7241 Ok(moe_out)
7242 }
7243
7244 fn gemma4_layer(&self, e: &Engine, il: usize, layer: &crate::hybrid::HybridLayer,
7246 x: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
7247 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7248 let n_embd = self.cfg.n_embd as usize;
7249 let eps = self.cfg.rms_eps;
7250
7251 let mut h = e.zeros(t * n_embd)?;
7252 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7253 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7254 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
7255 let mut cur = e.zeros(t * n_embd)?;
7257 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
7258 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
7259 }
7260
7261 fn gemma4_layer_tail_add(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7265 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
7266 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7267 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
7268 }
7269
7270 fn gemma4_layer_tail_add_n(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7273 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7274 next_norm: Option<&CudaSlice<f32>>)
7275 -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
7276 let n_embd = self.cfg.n_embd as usize;
7277 let bits = layer.gemma4.as_ref().unwrap();
7278 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
7279 let mut xn = e.uninit(t * n_embd)?;
7280 match next_norm {
7281 Some(w) => {
7282 let mut hn = e.uninit(t * n_embd)?;
7283 e.add_scale_rms_norm(&sn, &attn_out, bits.layer_scale, w, &mut xn, &mut hn,
7284 n_embd, t, self.cfg.rms_eps)?;
7285 Ok((xn, Some(hn)))
7286 }
7287 None => {
7288 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
7289 Ok((xn, None))
7290 }
7291 }
7292 }
7293
7294 fn gemma4_layer_tail_core(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7297 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
7298 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7299 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
7300 }
7301
7302 fn gemma4_layer_tail_core_pn(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7309 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7310 pre_norm: Option<&CudaSlice<f32>>, defer_post_norm: bool)
7311 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7312 let n_embd = self.cfg.n_embd as usize;
7313 let eps = self.cfg.rms_eps;
7314 let bits = layer.gemma4.as_ref().unwrap();
7315
7316 let Some(mbits) = bits.moe_bits.as_ref() else {
7319 let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
7320 else { panic!("gemma4 dense layer without Dense ffn") };
7321 let mut attn_out = e.uninit(t * n_embd)?;
7322 let mut zsh = e.uninit(t * n_embd)?;
7323 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7326 match pre_norm {
7327 Some(wa) if t == 1 => {
7328 zpair = Some(e.rms_pre_add_rms_norm_q8z(cur, wa, x,
7329 bits.ffn_norm.float_data(),
7330 &mut attn_out, &mut zsh,
7331 n_embd, t, eps)?);
7332 }
7333 Some(wa) => e.rms_pre_add_rms_norm(cur, wa, x, bits.ffn_norm.float_data(),
7334 &mut attn_out, &mut zsh, n_embd, t, eps)?,
7335 None => e.add_rms_norm(cur, x, bits.ffn_norm.float_data(), &mut attn_out,
7336 &mut zsh, n_embd, t, eps)?,
7337 }
7338 let n_ff = ffn_gate.out_features();
7339 let (gate, up) = if t == 1 {
7345 let (zq, zd) = match zpair {
7346 Some(p) => p,
7347 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
7348 };
7349 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
7350 Some(p) => p,
7351 None => (e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
7352 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?),
7353 }
7354 } else {
7355 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7360 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
7361 let fused = if f2b {
7362 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
7363 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
7364 } else { None };
7365 match fused {
7366 Some(p) => p,
7367 None => {
7368 e.mmq_act_begin();
7370 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
7371 }
7372 }
7373 };
7374 let mut act = e.uninit(t * n_ff)?;
7375 let f0 = if e.uses_q8_1_fast(ffn_down) {
7378 let upv = e.view(&up, t * n_ff);
7379 let up_all = upv.slice(0..t * n_ff);
7380 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
7381 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
7382 } else {
7383 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
7384 e.matmul(ffn_down, &act, t)?
7385 };
7386 if defer_post_norm { return Ok((f0, attn_out)); }
7387 let mut sn = e.uninit(t * n_embd)?;
7388 e.rms_norm(&f0, bits.post_ffw_norm.float_data(), &mut sn, n_embd, t, eps)?;
7389 return Ok((sn, attn_out));
7390 };
7391
7392 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
7393 let mut attn_out = e.uninit(t * n_embd)?;
7398 let mut router_in = e.uninit(t * n_embd)?;
7399 let fast_moe = match &layer.ffn {
7400 crate::hybrid::Ffn::Moe(m) => m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7401 && expert_dp4a_supported(m.gate_exps.qtype)
7402 && expert_dp4a_supported(m.up_exps.qtype)
7403 && expert_dp4a_supported(m.down_exps.qtype)
7404 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0"),
7405 _ => false,
7406 };
7407 let q8z = t < PRIME_MIN_T && fast_moe;
7408 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
7409 let (z0, m2) = e.add_rms_norm3_q8z(cur, x, bits.ffn_norm.float_data(),
7410 &mbits.router_scale_pre,
7411 mbits.pre_ffw_norm_2.float_data(),
7412 &mut attn_out, &mut router_in, n_embd, t, eps)?;
7413 (None, Some(z0), Some(m2))
7414 } else {
7415 let mut zsh = e.uninit(t * n_embd)?;
7416 let mut moe_in = e.uninit(t * n_embd)?;
7417 e.add_rms_norm3(cur, x, bits.ffn_norm.float_data(), &mbits.router_scale_pre,
7418 mbits.pre_ffw_norm_2.float_data(), &mut attn_out, &mut zsh,
7419 &mut router_in, &mut moe_in, n_embd, t, eps)?;
7420 (Some((zsh, moe_in)), None, None)
7421 };
7422 let attn_out2 = attn_out;
7423 #[allow(unused_variables)]
7424 let attn_out = &attn_out2;
7425 let n_ff = mbits.shared_gate.out_features();
7426 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
7427 if t == 1 {
7428 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
7429 Some(p) => p,
7430 None => {
7431 let h0 = e.zeros(0)?;
7432 (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
7433 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?)
7434 }
7435 }
7436 } else {
7437 let h0 = e.zeros(0)?;
7439 (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
7440 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?)
7441 }
7442 } else {
7443 let (zsh, _) = zsh_f32.as_ref().unwrap();
7444 (e.matmul(&mbits.shared_gate, zsh, t)?, e.matmul(&mbits.shared_up, zsh, t)?)
7445 };
7446 let mut act = e.uninit(t * n_ff)?;
7447 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
7448 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
7449 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else { panic!("gemma4 layer not MoE") };
7450 let moe0 = match (&moe_q8, &zsh_f32) {
7451 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
7452 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
7453 _ => unreachable!(),
7454 };
7455 let mut mlp = e.uninit(t * n_embd)?;
7457 let mut moe = e.uninit(t * n_embd)?;
7458 e.rms_norm2x(&mlp0, &moe0, mbits.post_ffw_norm_1.float_data(),
7459 mbits.post_ffw_norm_2.float_data(), &mut mlp, &mut moe, n_embd, t, eps)?;
7460
7461 let mut sum = e.uninit(t * n_embd)?;
7464 let mut sn = e.uninit(t * n_embd)?;
7465 e.add_rms_norm(&mlp, &moe, bits.post_ffw_norm.float_data(), &mut sum, &mut sn,
7466 n_embd, t, eps)?;
7467 Ok((sn, attn_out2))
7468 }
7469
7470 fn gemma4_layer_tail_add_nq(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7472 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7473 next_norm: Option<&CudaSlice<f32>>)
7474 -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>> {
7475 let n_embd = self.cfg.n_embd as usize;
7476 let bits = layer.gemma4.as_ref().unwrap();
7477 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
7478 let mut xn = e.uninit(t * n_embd)?;
7479 match next_norm {
7480 Some(w) => {
7481 let pair = e.add_scale_rms_norm_q8_1(&sn, &attn_out, bits.layer_scale, w, &mut xn,
7482 n_embd, t, self.cfg.rms_eps)?;
7483 Ok((xn, Some(pair)))
7484 }
7485 None => {
7486 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
7487 Ok((xn, None))
7488 }
7489 }
7490 }
7491
7492 fn gemma4_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
7495 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7496 if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, last_only); }
7499 let n_embd = self.cfg.n_embd as usize;
7500 let t = tokens.len();
7501 let pos: Vec<i32> = (0..t as i32).collect();
7502 let pos_d = e.htod_i32(&pos)?;
7503
7504 let mut x = self.embed(e, tokens)?;
7505 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7506 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
7509 let stat = |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
7510 let h = e.dtoh(x)?;
7511 let bad = h.iter().filter(|v| !v.is_finite()).count();
7512 let mx = h.iter().filter(|v| v.is_finite()).fold(0.0f32, |m, v| m.max(v.abs()));
7513 eprintln!("[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}", &h[..3]);
7514 Ok(())
7515 };
7516 if probe { stat(e, &x, "embed")?; }
7517 for (il, layer) in self.layers.iter().enumerate() {
7518 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
7519 if probe { stat(e, &x, &format!("L{il}"))?; }
7520 }
7521 let mut hn = e.zeros(t * n_embd)?;
7522 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, self.cfg.rms_eps)?;
7523 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7524 let n_vocab = self.output.out_features();
7525 let logits = if last_only {
7526 let hv = e.view(&hn, t * n_embd);
7527 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
7528 let mut hlast = e.zeros(n_embd)?;
7529 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
7530 let mut ld = e.matmul(&self.output, &hlast, 1)?;
7531 e.softcap(&mut ld, cap, n_vocab)?;
7532 self.gemma4_suppress(e, &mut ld, 1)?;
7533 e.dtoh(&ld)?
7534 } else {
7535 let mut ld = e.matmul(&self.output, &hn, t)?;
7536 e.softcap(&mut ld, cap, t * n_vocab)?;
7537 self.gemma4_suppress(e, &mut ld, t)?;
7538 e.dtoh(&ld)?
7539 };
7540 Ok(logits)
7541 }
7542
7543 pub(crate) fn gemma4_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
7548 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7549 if cache.pos != 0 {
7554 return Err("gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
7555 — prime the full prompt in one call or decode tokenwise".into());
7556 }
7557 let n_embd = self.cfg.n_embd as usize;
7558 let eps = self.cfg.rms_eps;
7559 let t = tokens.len();
7560 let pos: Vec<i32> = (0..t as i32).collect();
7561 let pos_d = e.htod_i32(&pos)?;
7562 let mut x = self.embed(e, tokens)?;
7563 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7564 for (il, layer) in self.layers.iter().enumerate() {
7565 let mut h = e.zeros(t * n_embd)?;
7566 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7567 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer not full-attn") };
7568 let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
7569 let mut cur = e.zeros(t * n_embd)?;
7570 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
7571 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
7572 self.dflash_tap(e, cache, il, &x, t)?;
7573 }
7574 cache.pos += t;
7575 let hiddens = e.clone_dtod(&x)?;
7576 let xv = e.view(&x, t * n_embd);
7577 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
7578 let mut h_seed = e.zeros(n_embd)?;
7579 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
7580 let mut hn = e.uninit(n_embd)?;
7581 e.rms_norm(&h_seed, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7582 let mut ld = e.matmul(&self.output, &hn, 1)?;
7583 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7584 e.softcap(&mut ld, cap, self.output.out_features())?;
7585 self.gemma4_suppress(e, &mut ld, 1)?;
7586 let logits = e.dtoh(&ld)?;
7587 Ok((logits, h_seed, hiddens))
7588 }
7589
7590 fn gemma4_decode_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
7595 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
7596 pos_d: &CudaSlice<i32>, cache: &mut Cache)
7597 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7598 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7599 let eps = self.cfg.rms_eps;
7600 let aux = self.gemma4_aux.as_ref().unwrap();
7601 let ones = aux.ones(e);
7602 #[cfg(debug_assertions)]
7603 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
7604 "gemma4_decode_attn.ones");
7605 let (hq, hdq) = (hq, hdq);
7606 let h0 = e.zeros(0)?;
7607 let h = &h0;
7608 let (q0, k0, v0) = if swa {
7609 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
7610 Some(t3) => t3,
7611 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7612 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
7613 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
7614 }
7615 } else {
7616 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
7617 Some(p) => p,
7618 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7619 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?),
7620 };
7621 let v0 = e.clone_dtod(&k0)?;
7622 (q0, k0, v0)
7623 };
7624 let mut q = e.uninit(nh * hd)?;
7625 let mut k = e.uninit(nkv * hd)?;
7626 let mut v = e.uninit(nkv * hd)?;
7627 let ff = if swa { None } else {
7630 Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
7631 };
7632 #[cfg(debug_assertions)]
7633 if let Some(ff) = ff {
7634 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
7635 "gemma4_decode_attn.rope_freqs");
7636 }
7637 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7638 ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7639 pos_d, nh, nkv, base, 1.0, ff, eps)?;
7640 let kvl = cache.kv[il].as_mut().unwrap();
7641 e.append_kv_quantized(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len,
7642 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()))?;
7643 kvl.len += 1;
7644 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7648 let mut attn = e.uninit(nh * hd)?;
7649 if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
7651 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7652 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7653 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7654 let base = kvl.len as i32;
7656 e.i32_set_k(&mut kvl.len_d, base)?;
7657 e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1, scale,
7658 kvl.k_tok_bytes, kvl.v_tok_bytes, Some((&kvl.len_d, -1)), false,
7659 false, None)?;
7660 return Ok(e.matmul(&fa.wo, &attn, 1)?);
7661 }
7662 if swa && kvl.len > win && hd == 256
7664 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7665 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7666 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7667 let base = kvl.len as i32;
7668 e.i32_set_k(&mut kvl.len_d, base)?;
7669 e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1, 1, scale,
7670 win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
7671 return Ok(e.matmul(&fa.wo, &attn, 1)?);
7672 }
7673 let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
7674 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
7675 (off_tok + t_kv) * kvl.k_tok_bytes);
7676 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
7677 (off_tok + t_kv) * kvl.v_tok_bytes);
7678 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
7679 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
7680 Ok(e.matmul(&fa.wo, &attn, 1)?)
7681 }
7682
7683 #[allow(clippy::too_many_arguments)]
7690 pub fn gemma4_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
7691 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7692 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7693 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>)
7694 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7695 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
7696 self.gemma4_decode_step_dc_into(e, token_d, pos_d, embd_gpu, embd_qt, embd_rb, cache,
7697 n_vocab, cap_bucket_max, &mut tok_out)?;
7698 Ok(tok_out)
7699 }
7700
7701 #[allow(clippy::too_many_arguments)]
7704 pub fn gemma4_decode_step_dc_into(&self, e: &Engine, token_d: &CudaSlice<u32>,
7705 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7706 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7707 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
7708 tok_out: &mut CudaSlice<u32>)
7709 -> Result<(), Box<dyn std::error::Error>> {
7710 let n_embd = self.cfg.n_embd as usize;
7711 let eps = self.cfg.rms_eps;
7712 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7713 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7714 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7715 let n_layers = self.layers.len();
7716 for (il, layer) in self.layers.iter().enumerate() {
7717 let (hq, hdq) = match h_carry.take() {
7718 Some(p) => p,
7719 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
7720 };
7721 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7722 let o = self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
7723 let mut cur = e.uninit(n_embd)?;
7724 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
7725 let next_norm = if il + 1 < n_layers {
7726 Some(self.layers[il + 1].attn_norm.float_data())
7727 } else { None };
7728 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
7729 x = xn;
7730 h_carry = hn;
7731 }
7732 let mut hn = e.uninit(n_embd)?;
7733 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7734 let mut logits = e.matmul(&self.output, &hn, 1)?;
7735 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
7737 e.inc_seqlen(pos_d)?;
7738 if cap_bucket_max.is_none() { cache.pos += 1; }
7739 Ok(())
7740 }
7741
7742 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
7749 let n_embd = self.cfg.n_embd as usize;
7750 let n_vocab = self.output.out_features();
7751 let n_layers = self.layers.len();
7752 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
7753 for il in 0..n_layers {
7754 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
7755 qmax = qmax.max(nh * hd);
7756 kvmax = kvmax.max(nkv * hd);
7757 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
7758 ffmax = ffmax.max(ffn_gate.out_features());
7759 }
7760 }
7761 Ok(G4DcSlots {
7762 x: e.uninit(n_embd)?, xn: e.uninit(n_embd)?, cur: e.uninit(n_embd)?,
7763 hq: e.alloc_i8_uninit(n_embd)?, hd_: e.uninit(n_embd / 32)?,
7764 q0: e.uninit(qmax)?, k0: e.uninit(kvmax)?, v0: e.uninit(kvmax)?,
7765 q: e.uninit(qmax)?, k: e.uninit(kvmax)?, v: e.uninit(kvmax)?,
7766 attn: e.uninit(qmax)?, o: e.uninit(n_embd)?,
7767 attn_out: e.uninit(n_embd)?, zsh: e.uninit(n_embd)?,
7768 zq: e.alloc_i8_uninit(n_embd.max(qmax))?, zd: e.uninit(n_embd.max(qmax) / 32)?,
7771 gate: e.uninit(ffmax)?, up: e.uninit(ffmax)?,
7772 act: e.uninit(ffmax)?, actq: e.alloc_i8_uninit(ffmax)?, actd: e.uninit(ffmax / 32)?,
7773 f0: e.uninit(n_embd)?, sn: e.uninit(n_embd)?,
7774 hn: e.uninit(n_embd)?, logits: e.uninit(n_vocab)?,
7775 })
7776 }
7777
7778 fn g4_matvec_m1_into(&self, e: &Engine, w: &crate::model::GpuTensor,
7781 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y: &mut CudaSlice<f32>)
7782 -> Result<(), Box<dyn std::error::Error>> {
7783 use crate::model::GpuTensor;
7784 let (bytes, qtype, row_bytes, scale, rp) = match w {
7785 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
7786 (bytes, *qtype, *row_bytes, *scale, *rp),
7787 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
7788 };
7789 let (mbytes, mrp) = match w {
7790 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
7791 _ => (bytes, rp),
7792 };
7793 e.qmatvec_mmvq_into(mbytes, aq, ad, 1, w.in_features(), w.out_features(),
7794 qtype, row_bytes, scale, mrp, y)
7795 }
7796
7797 #[allow(clippy::too_many_arguments)]
7801 pub fn gemma4_decode_step_dc_slotted(&self, e: &Engine, token_d: &CudaSlice<u32>,
7802 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7803 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7804 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
7805 sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>,
7806 ring: Option<(&mut CudaSlice<u32>, usize)>)
7807 -> Result<(), Box<dyn std::error::Error>> {
7808 let n_embd = self.cfg.n_embd as usize;
7809 let eps = self.cfg.rms_eps;
7810 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
7811 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
7812 let n_layers = self.layers.len();
7813 let mut has_carry = false;
7814 for il in 0..n_layers {
7815 if !has_carry {
7816 e.rms_norm_q8_1_into(&sl.x, self.layers[il].attn_norm.float_data(), n_embd, 1,
7817 eps, &mut sl.hq, &mut sl.hd_)?;
7818 }
7819 has_carry = true;
7820 let layer = &self.layers[il];
7821 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7822 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
7823 e.rms_norm(&sl.o, layer.post_attn_norm.float_data(), &mut sl.cur, n_embd, 1, eps)?;
7824 let next_norm = if il + 1 < n_layers {
7825 Some(self.layers[il + 1].attn_norm.float_data())
7826 } else { None };
7827 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
7828 std::mem::swap(&mut sl.x, &mut sl.xn);
7829 }
7830 e.rms_norm(&sl.x, self.output_norm.float_data(), &mut sl.hn, n_embd, 1, eps)?;
7831 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
7832 {
7834 let (zq, zd) = (&sl.zq, &sl.zd);
7835 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
7836 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
7837 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
7838 }
7839 self.gemma4_suppress(e, &mut sl.logits, 1)?;
7840 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
7841 if let Some((ring, base)) = ring {
7842 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
7846 }
7847 e.inc_seqlen(pos_d)?;
7848 if cap_bucket_max.is_none() { cache.pos += 1; }
7849 Ok(())
7850 }
7851
7852 #[allow(clippy::too_many_arguments)]
7854 fn gemma4_decode_attn_dc_slotted(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer,
7855 il: usize, pos_d: &CudaSlice<i32>, cache: &mut Cache,
7856 cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots)
7857 -> Result<(), Box<dyn std::error::Error>> {
7858 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7859 let eps = self.cfg.rms_eps;
7860 let aux = self.gemma4_aux.as_ref().unwrap();
7861 let ones = aux.ones(e);
7862 #[cfg(debug_assertions)]
7863 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
7864 "gemma4_decode_attn_dc_slotted.ones");
7865 {
7866 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
7867 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
7868 if swa {
7869 if !e.matmul_q4_fused3_into(&fa.wq, &fa.wk, &fa.wv, hq, hdq,
7870 &mut sl.q0, &mut sl.k0, &mut sl.v0)? {
7871 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
7872 }
7873 } else {
7874 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
7875 return Err("slotted step: fused2 unavailable".into());
7876 }
7877 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
7878 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
7879 }
7880 }
7881 let ff = if swa { None } else {
7884 Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
7885 };
7886 #[cfg(debug_assertions)]
7887 if let Some(ff) = ff {
7888 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
7889 "gemma4_decode_attn_dc_slotted.rope_freqs");
7890 }
7891 let kvl = cache.kv[il].as_mut().unwrap();
7892 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7893 if crate::Engine::qkv_append_on() {
7894 e.rms_norm_qkv_rope_append_dc(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(),
7896 fa.k_norm.float_data(), ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
7897 pos_d, nh, nkv, base, 1.0, ff, eps,
7898 &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
7899 } else {
7900 e.rms_norm_qkv_rope(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7901 ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
7902 pos_d, nh, nkv, base, 1.0, ff, eps)?;
7903 e.append_kv_quantized_dc(&sl.k, &sl.v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
7904 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
7905 kv_fp8)?;
7906 }
7907 e.inc_seqlen(&mut kvl.len_d)?;
7908 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
7909 let k_view = e.view_u8(&kvl.k, kvl.k.len());
7910 let v_view = e.view_u8(&kvl.v, kvl.v.len());
7911 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
7912 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7913 let mut fa_q8 = false;
7917 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
7918 e.fa_decode_rows(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, b_glob - 1,
7919 1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7920 Some((&kvl.len_d, -1)), false, false,
7921 Some((&mut sl.zq, &mut sl.zd)))?;
7922 fa_q8 = true;
7923 } else if swa && b_swa > win && hd == 256 && rows_on {
7924 e.fa_decode_rows_w(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv,
7925 &kvl.len_d, -1, 1, scale, win,
7926 kvl.k_tok_bytes, kvl.v_tok_bytes,
7927 Some((&mut sl.zq, &mut sl.zd)))?;
7928 fa_q8 = true;
7929 } else {
7930 let b = if swa { b_swa } else { b_glob };
7931 e.fa_decode_dc(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, &kvl.len_d, b,
7932 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7933 swa && crate::Engine::wkv_on())?;
7934 }
7935 if !fa_q8 {
7936 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
7937 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
7938 }
7939 {
7940 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
7941 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
7942 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
7943 }
7944 Ok(())
7945 }
7946
7947 fn gemma4_layer_tail_slotted(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7950 next_norm: Option<&CudaSlice<f32>>, sl: &mut G4DcSlots)
7951 -> Result<(), Box<dyn std::error::Error>> {
7952 let n_embd = self.cfg.n_embd as usize;
7953 let eps = self.cfg.rms_eps;
7954 let bits = layer.gemma4.as_ref().unwrap();
7955 let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
7956 else { return Err("slotted tail: dense ffn only".into()) };
7957 e.add_rms_norm(&sl.cur, &sl.x, bits.ffn_norm.float_data(), &mut sl.attn_out,
7958 &mut sl.zsh, n_embd, 1, eps)?;
7959 let n_ff = ffn_gate.out_features();
7960 {
7961 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
7962 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
7963 }
7964 {
7965 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
7966 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
7967 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
7968 return Err("slotted tail: ffn fused2 unavailable".into());
7969 }
7970 }
7971 debug_assert!(e.uses_q8_1_fast(ffn_down));
7972 {
7973 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
7974 let upv = e.view(upr, n_ff);
7975 let up_all = upv.slice(0..n_ff);
7976 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
7977 e.gelu_tanh_mul_q8_1_into(gr, &up_all, &mut sl.act, n_ff, 1,
7978 &mut sl.actq, &mut sl.actd)?;
7979 }
7980 {
7981 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
7982 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
7983 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
7984 }
7985 e.rms_norm(&sl.f0, bits.post_ffw_norm.float_data(), &mut sl.sn, n_embd, 1, eps)?;
7986 match next_norm {
7987 Some(w) => {
7988 e.add_scale_rms_norm_q8_1_into(&sl.sn, &sl.attn_out, bits.layer_scale, w,
7989 &mut sl.xn, n_embd, 1, eps,
7990 &mut sl.hq, &mut sl.hd_)?;
7991 }
7992 None => {
7993 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
7994 }
7995 }
7996 Ok(())
7997 }
7998
7999 #[allow(clippy::too_many_arguments)]
8001 fn gemma4_decode_attn_dc(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8002 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8003 pos_d: &CudaSlice<i32>, cache: &mut Cache,
8004 cap_bucket_max: Option<(usize, usize)>)
8005 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8006 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8007 let eps = self.cfg.rms_eps;
8008 let aux = self.gemma4_aux.as_ref().unwrap();
8009 let ones = aux.ones(e);
8010 #[cfg(debug_assertions)]
8011 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
8012 "gemma4_decode_attn_dc.ones");
8013 let (q0, k0, v0) = if swa {
8014 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
8015 Some(t3) => t3,
8016 None => {
8017 let h0 = e.zeros(0)?;
8018 (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
8019 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
8020 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?)
8021 }
8022 }
8023 } else {
8024 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
8025 Some(p) => p,
8026 None => {
8027 let h0 = e.zeros(0)?;
8028 (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
8029 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?)
8030 }
8031 };
8032 let v0 = e.clone_dtod(&k0)?;
8033 (q0, k0, v0)
8034 };
8035 let mut q = e.uninit(nh * hd)?;
8036 let mut k = e.uninit(nkv * hd)?;
8037 let mut v = e.uninit(nkv * hd)?;
8038 let ff = if swa { None } else {
8040 Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
8041 };
8042 #[cfg(debug_assertions)]
8043 if let Some(ff) = ff {
8044 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
8045 "gemma4_decode_attn_dc.rope_freqs");
8046 }
8047 let kvl = cache.kv[il].as_mut().unwrap();
8048 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
8049 if crate::Engine::qkv_append_on() {
8050 e.rms_norm_qkv_rope_append_dc(&q0, &k0, &v0, fa.q_norm.float_data(),
8052 fa.k_norm.float_data(), ones, &mut q, &mut k, &mut v, hd, nh, nkv,
8053 pos_d, nh, nkv, base, 1.0, ff, eps,
8054 &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
8055 } else {
8056 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8057 ones, &mut q, &mut k, &mut v, hd, nh, nkv,
8058 pos_d, nh, nkv, base, 1.0, ff, eps)?;
8059 e.append_kv_quantized_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
8060 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
8061 }
8062 e.inc_seqlen(&mut kvl.len_d)?;
8063 let mut attn = e.uninit(nh * hd)?;
8064 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8067 match cap_bucket_max {
8072 None => {
8073 kvl.len += 1;
8077 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8078 if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
8079 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8080 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
8083 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
8084 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8085 e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1,
8086 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8087 Some((&kvl.len_d, -1)), false, false,
8088 Some((&mut aq8, &mut ad8)))?;
8089 fa_q8 = Some((aq8, ad8));
8090 } else if swa && kvl.len > win && hd == 256
8091 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8092 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
8094 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
8095 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8096 e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1,
8097 1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes,
8098 Some((&mut aq8, &mut ad8)))?;
8099 fa_q8 = Some((aq8, ad8));
8100 } else {
8101 let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) }
8102 else { (0, kvl.len) };
8103 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
8104 (off_tok + t_kv) * kvl.k_tok_bytes);
8105 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
8106 (off_tok + t_kv) * kvl.v_tok_bytes);
8107 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
8108 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
8109 }
8110 }
8111 Some((b_swa, b_glob)) => {
8112 let k_view = e.view_u8(&kvl.k, kvl.k.len());
8118 let v_view = e.view_u8(&kvl.v, kvl.v.len());
8119 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
8120 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8121 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
8122 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8123 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, b_glob - 1,
8124 1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8125 Some((&kvl.len_d, -1)), false, false,
8126 Some((&mut aq8, &mut ad8)))?;
8127 fa_q8 = Some((aq8, ad8));
8128 } else if swa && b_swa > win && hd == 256 && rows_on {
8129 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8130 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8131 &kvl.len_d, -1, 1, scale, win,
8132 kvl.k_tok_bytes, kvl.v_tok_bytes,
8133 Some((&mut aq8, &mut ad8)))?;
8134 fa_q8 = Some((aq8, ad8));
8135 } else {
8136 let b = if swa { b_swa } else { b_glob };
8137 e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, b,
8138 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8139 swa && crate::Engine::wkv_on())?;
8140 }
8141 }
8142 }
8143 if let Some((aq8, ad8)) = fa_q8 {
8146 let mut y = e.uninit(fa.wo.out_features())?;
8147 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
8148 return Ok(y);
8149 }
8150 Ok(e.matmul(&fa.wo, &attn, 1)?)
8151 }
8152
8153 pub fn gemma4_generate_graph(&self, e: &Engine, prompt_pos: usize, first_token: u32,
8158 cache: &mut Cache, max_new: usize, eos: &[u32],
8159 mut on_token: impl FnMut(u32) -> bool)
8160 -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
8161 if self.is_gemma4_e4b() {
8162 return Err("E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm".into());
8163 }
8164 use crate::decode::StopReason;
8165 let n_vocab = self.output.out_features();
8166 let n_embd = self.cfg.n_embd as usize;
8167 let embd_gpu = self.embd_gpu.get_or_init(|| {
8168 e.upload_u8(&self.embd.raw).expect("embed table upload")
8169 });
8170 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8171 for kvl in cache.kv.iter_mut().flatten() {
8172 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8173 }
8174 let mut token_d = e.stream().clone_htod(&[first_token])?;
8175 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
8176 let g4 = self.cfg.gemma4.as_ref().unwrap();
8177 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
8178 let nkv_s = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
8180 .find(|p| *p.1).map(|p| *p.0 as usize).unwrap_or(8);
8181 let nkv_g = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
8182 .find(|p| !*p.1).map(|p| *p.0 as usize).unwrap_or(2);
8183 let mut graphs: std::collections::HashMap<((bool, usize), (bool, usize), bool, bool),
8184 (cudarc::driver::CudaGraph,
8185 Vec<Box<dyn std::any::Any + Send>>)> = Default::default();
8186 let mut slots = self.g4_dc_slots(e)?;
8189 const RING: usize = 64;
8192 const DRAIN: usize = 1;
8198 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
8199 let ring_base = prompt_pos;
8200 let mut out = Vec::with_capacity(max_new);
8201 let mut reason = StopReason::MaxNew;
8202 let mut next = first_token;
8203 let mut captures = 0usize;
8204 for _ in 0..max_new {
8205 out.push(next);
8206 if eos.contains(&next) { reason = StopReason::Eos; break; }
8207 if !on_token(next) { reason = StopReason::Callback; break; }
8208 let t_kv = cache.pos + 1;
8209 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8217 let f512 = crate::fa512_min_tkv();
8218 let key_s = if t_kv > win { (true, usize::MAX) }
8219 else { e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on()) };
8220 let (key_g, rung_end) = if t_kv >= f512 {
8221 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
8224 ((true, end), end)
8225 } else { (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv) };
8226 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
8227 if !graphs.contains_key(&key) {
8228 let bucket_max = (t_kv, rung_end);
8229 let snap = cache.snapshot(e)?;
8231 let pos_save = e.dtoh_i32_one(&pos_d)?;
8232 let len_save: Vec<Option<i32>> = cache.kv.iter()
8233 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
8234 let tok_save = e.dtoh_u32_one(&token_d)?;
8235 let graph = {
8240 let tok_ref = &mut token_d;
8241 let pos_ref = &mut pos_d;
8242 let cache_ref = &mut *cache;
8243 let slots_ref = &mut slots;
8244 let ring_ref = &mut ring;
8245 e.capture_graph_retained_flags(
8246 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
8247 |e| {
8248 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
8250 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
8251 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
8252 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
8253 cache_ref, n_vocab, Some(bucket_max),
8254 sl, tok_ref, Some((rg, ring_base)))
8255 })?
8256 };
8257 cache.rollback(e, &snap, 0)?;
8258 e.set_i32_one(&mut pos_d, pos_save)?;
8259 for (il, ls) in len_save.iter().enumerate() {
8260 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
8261 e.set_i32_one(&mut kvl.len_d, *v)?;
8262 }
8263 }
8264 e.set_u32_one(&mut token_d, tok_save)?;
8265 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
8266 if let Ok(c) = crate::graph_update::node_census(&graph.0) {
8267 eprintln!("[graph-census] {c:?}");
8268 }
8269 }
8270 graphs.insert(key, graph);
8271 captures += 1;
8272 }
8273 let mut chunk = 1usize;
8278 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN").ok()
8279 .and_then(|v| v.parse().ok()).unwrap_or(DRAIN);
8280 while chunk < drain_cap && out.len() + chunk < max_new {
8281 let t_next = cache.pos + 1 + chunk;
8282 let key_s2 = if t_next > win { (true, usize::MAX) }
8283 else { e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on()) };
8284 let key_g2 = if t_next >= f512 {
8285 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
8286 } else { e.fa_bucket_key(t_next, hd_g, nkv_g, false) };
8287 if (key_s2, key_g2, t_next >= f512, t_next > win) != key { break; }
8288 chunk += 1;
8289 }
8290 let g = &graphs.get(&key).unwrap().0;
8291 for _ in 0..chunk { g.launch()?; }
8292 e.stream().synchronize()?;
8293 let ringh = e.dtoh_u32(&ring)?;
8294 for j in 0..chunk {
8295 let pos_j = cache.pos + j;
8296 let tok_j = ringh[(pos_j - ring_base) % RING];
8297 cache.pos += 0; if j + 1 == chunk { next = tok_j; }
8299 else {
8300 out.push(tok_j);
8301 if eos.contains(&tok_j) || !on_token(tok_j) {
8302 reason = if eos.contains(&tok_j) { StopReason::Eos }
8303 else { StopReason::Callback };
8304 let keep = cache.pos + j + 1;
8306 e.set_i32_one(&mut pos_d, keep as i32)?;
8307 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
8308 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
8309 kvl.len = keep;
8310 }
8311 cache.pos = keep;
8312 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
8313 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
8314 }
8315 return Ok((out, reason));
8316 }
8317 }
8318 }
8319 cache.pos += chunk;
8320 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) { kvl.len += chunk; }
8321 }
8322 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
8323 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
8324 }
8325 Ok((out, reason))
8326 }
8327
8328 pub(crate) fn gemma4_decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize,
8334 cache: &mut Cache)
8335 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8336 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
8337 }
8338
8339 pub(crate) fn gemma4_decode_step_t_am(&self, e: &Engine, tokens: &[u32], pos0: usize,
8343 cache: &mut Cache)
8344 -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8345 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
8346 let t = tokens.len();
8347 let n_vocab = self.output.out_features();
8348 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
8349 for i in 0..t {
8350 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
8351 }
8352 Ok((e.dtoh_u32(&toks)?, hn))
8353 }
8354
8355 pub(crate) fn gemma4_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
8358 pos0: usize, cache: &mut Cache)
8359 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8360 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
8361 let n_vocab = self.output.out_features();
8362 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
8363 for i in 0..t {
8364 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
8365 }
8366 Ok((vam, hn))
8367 }
8368
8369 pub(crate) fn gemma4_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
8372 cache: &mut Cache)
8373 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8374 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
8375 let t = tokens.len();
8376 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8377 e.softcap(&mut ld, cap, t * self.output.out_features())?;
8378 Ok((e.dtoh(&ld)?, hn))
8379 }
8380
8381 pub(crate) fn verify_stream_scratch(&self, e: &Engine, cap: usize)
8384 -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
8385 Ok(VerifyStreamScratch {
8386 pos_d: e.htod_i32(&vec![0i32; cap])?,
8387 row_ctrs: (0..cap).map(|_| e.htod_i32(&[0])).collect::<Result<_, _>>()?,
8388 })
8389 }
8390
8391 pub(crate) fn gemma4_verify_t_am_stream(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
8399 ctr: &CudaSlice<i32>, hint: usize,
8400 cache: &mut Cache,
8401 scr: &mut VerifyStreamScratch)
8402 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8403 let n_embd = self.cfg.n_embd as usize;
8404 let eps = self.cfg.rms_eps;
8405 assert!(t <= scr.row_ctrs.len() && t <= 64);
8406 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
8407 for i in 0..t {
8408 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
8409 }
8410 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
8411 let embd_gpu = self.embd_gpu.get_or_init(|| {
8412 e.upload_u8(&self.embd.raw).expect("embed table upload")
8413 });
8414 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8415 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
8416 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
8417 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8418 let n_layers = self.layers.len();
8419 for (il, layer) in self.layers.iter().enumerate() {
8420 let (hq, hdq) = match h_carry.take() {
8421 Some(p) => p,
8422 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
8423 };
8424 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8425 let o = self.gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache,
8426 hint, row_ctrs)?;
8427 let mut cur = e.uninit(t * n_embd)?;
8428 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
8429 let next_norm = if il + 1 < n_layers {
8430 Some(self.layers[il + 1].attn_norm.float_data())
8431 } else { None };
8432 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
8433 x = xn;
8434 h_carry = hn;
8435 self.dflash_tap(e, cache, il, &x, t)?;
8436 }
8437 let mut hn = e.uninit(t * n_embd)?;
8438 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8439 let ld = e.matmul(&self.output, &hn, t)?;
8440 let n_vocab = self.output.out_features();
8441 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
8442 for i in 0..t {
8443 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
8444 }
8445 Ok((vam, hn))
8446 }
8447
8448 fn dflash_tap(&self, e: &Engine, cache: &mut Cache, il: usize, x: &CudaSlice<f32>, t: usize)
8455 -> Result<(), Box<dyn std::error::Error>> {
8456 let Some(taps) = cache.dflash_taps.as_mut() else { return Ok(()) };
8457 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else { return Ok(()) };
8458 let h = taps.hidden;
8459 let n_taps = taps.layer_ids.len();
8460 debug_assert_eq!(taps.t, t);
8461 let xv = e.view(x, t * h);
8462 for r in 0..t {
8463 let row = xv.slice(r * h..(r + 1) * h);
8464 e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
8465 }
8466 Ok(())
8467 }
8468
8469 fn gemma4_verify_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
8470 tok_dev: Option<&CudaSlice<u32>>)
8471 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8472 let n_embd = self.cfg.n_embd as usize;
8473 let eps = self.cfg.rms_eps;
8474 let t = tokens.len();
8475 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8476 let pos_d = e.htod_i32(&pos)?;
8477 let mut x = match tok_dev {
8478 Some(td) => {
8479 let embd_gpu = self.embd_gpu.get_or_init(|| {
8480 e.upload_u8(&self.embd.raw).expect("embed table upload")
8481 });
8482 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8483 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
8484 }
8485 None => e.htod(&self.embd.gather(n_embd, tokens))?,
8486 };
8487 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
8488 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8489 let n_layers = self.layers.len();
8490 for (il, layer) in self.layers.iter().enumerate() {
8491 let (hq, hdq) = match h_carry.take() {
8492 Some(p) => p,
8493 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
8494 };
8495 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8496 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
8497 let mut cur = e.uninit(t * n_embd)?;
8498 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
8499 let next_norm = if il + 1 < n_layers {
8500 Some(self.layers[il + 1].attn_norm.float_data())
8501 } else { None };
8502 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
8503 x = xn;
8504 h_carry = hn;
8505 self.dflash_tap(e, cache, il, &x, t)?;
8506 }
8507 let mut hn = e.uninit(t * n_embd)?;
8508 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8509 let mut ld = e.matmul(&self.output, &hn, t)?;
8510 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
8512 Ok((ld, hn))
8513 }
8514
8515 #[allow(clippy::too_many_arguments)]
8523 fn gemma4_verify_attn_stream(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8524 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8525 pos_d: &CudaSlice<i32>, t: usize,
8526 cache: &mut Cache, hint: usize,
8527 row_ctrs: &[CudaSlice<i32>])
8528 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8529 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8530 let eps = self.cfg.rms_eps;
8531 let aux = self.gemma4_aux.as_ref().unwrap();
8532 let ones = aux.ones(e);
8533 #[cfg(debug_assertions)]
8534 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
8535 "gemma4_verify_attn_stream.ones");
8536 let h0 = e.zeros(0)?;
8537 let h = &h0;
8538 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8541 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
8542 let fused_qkv = if f2b {
8543 if swa {
8544 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
8545 .map(|(a, b, c)| (a, b, Some(c)))
8546 } else {
8547 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
8548 .map(|(a, b)| (a, b, None))
8549 }
8550 } else { None };
8551 let (q0, k0, v0) = match fused_qkv {
8552 Some((a, b, cv)) => {
8553 let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
8554 (a, b, v)
8555 }
8556 None => {
8557 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
8558 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
8559 let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
8560 else { e.clone_dtod(&k0)? };
8561 (q0, k0, v0)
8562 }
8563 };
8564 let mut q = e.uninit(t * nh * hd)?;
8565 let mut k = e.uninit(t * nkv * hd)?;
8566 let mut v = e.uninit(t * nkv * hd)?;
8567 let ff = if swa { None } else {
8570 Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
8571 };
8572 #[cfg(debug_assertions)]
8573 if let Some(ff) = ff {
8574 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
8575 "gemma4_verify_attn_stream.rope_freqs");
8576 }
8577 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8578 ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
8579 pos_d, nh, nkv, base, 1.0, ff, eps)?;
8580 let kvl = cache.kv[il].as_mut().unwrap();
8581 e.append_kv_quantized_rows_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d, t,
8583 kvl.kv_dim_k, kvl.kv_dim_v,
8584 kvl.k_tok_bytes, kvl.v_tok_bytes,
8585 (!swa && crate::Engine::gkv_on())
8586 || (swa && crate::Engine::wkv_on()))?;
8587 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8590 let mut attn = e.uninit(t * nh * hd)?;
8591 let k_view = e.view_u8(&kvl.k, kvl.k.len());
8592 let v_view = e.view_u8(&kvl.v, kvl.v.len());
8593 if swa && hint + 1 >= win {
8596 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8599 &kvl.len_d, 0, t, scale, win,
8600 kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8601 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
8602 let bucket = (hint + t + 2).next_power_of_two()
8615 .min(crate::fa512_min_tkv().saturating_sub(1));
8616 let qv = e.view(&q, t * nh * hd);
8617 for i in 0..t {
8618 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
8619 let mut q_one = e.uninit(nh * hd)?;
8620 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
8621 let mut a_one = e.uninit(nh * hd)?;
8622 e.fa_decode_dc(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv,
8623 &row_ctrs[i], bucket, scale,
8624 kvl.k_tok_bytes, kvl.v_tok_bytes, false)?;
8625 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
8626 }
8627 } else if hd == 512 {
8628 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, hint, t, scale,
8631 kvl.k_tok_bytes, kvl.v_tok_bytes,
8632 Some((&kvl.len_d, 0)), false, false, None)?;
8633 } else {
8634 e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8636 &kvl.len_d, hint + t, t, scale,
8637 kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
8638 swa && crate::Engine::wkv_on())?;
8639 }
8640 Ok(e.matmul(&fa.wo, &attn, t)?)
8641 }
8642
8643 fn gemma4_verify_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8644 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8645 pos_d: &CudaSlice<i32>, t: usize,
8646 cache: &mut Cache)
8647 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8648 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8649 let eps = self.cfg.rms_eps;
8650 let aux = self.gemma4_aux.as_ref().unwrap();
8651 let ones = aux.ones(e);
8652 #[cfg(debug_assertions)]
8653 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
8654 "gemma4_verify_attn.ones");
8655 let n_embd = self.cfg.n_embd as usize;
8656 let _ = n_embd;
8657
8658 let h0 = e.zeros(0)?;
8659 let h = &h0;
8660 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8663 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
8664 let fused_qkv = if f2b {
8665 if swa {
8666 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
8667 .map(|(a, b, c)| (a, b, Some(c)))
8668 } else {
8669 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
8670 .map(|(a, b)| (a, b, None))
8671 }
8672 } else { None };
8673 let (q0, k0, v0) = match fused_qkv {
8674 Some((a, b, cv)) => {
8675 let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
8676 (a, b, v)
8677 }
8678 None => {
8679 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
8680 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
8681 let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
8682 else { e.clone_dtod(&k0)? };
8683 (q0, k0, v0)
8684 }
8685 };
8686 let mut q = e.uninit(t * nh * hd)?;
8687 let mut k = e.uninit(t * nkv * hd)?;
8688 let mut v = e.uninit(t * nkv * hd)?;
8689 let ff = if swa { None } else {
8692 Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
8693 };
8694 #[cfg(debug_assertions)]
8695 if let Some(ff) = ff {
8696 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
8697 "gemma4_verify_attn.rope_freqs");
8698 }
8699 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8700 ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
8701 pos_d, nh, nkv, base, 1.0, ff, eps)?;
8702 let kvl = cache.kv[il].as_mut().unwrap();
8703 let base_len = kvl.len;
8704 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
8705 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()))?;
8706 kvl.len += t;
8707 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8708 let mut attn = e.uninit(t * nh * hd)?;
8709 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
8712 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
8715 if rows_ok && (!swa || base_len + t <= win) {
8716 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
8717 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
8718 if hd == 512 {
8719 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8721 e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, base_len, t,
8722 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8723 Some((&kvl.len_d, 0)), false,
8724 swa && crate::Engine::wkv_on(), None)?;
8725 } else {
8726 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8730 e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8731 &kvl.len_d, base_len + t, t, scale,
8732 kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
8733 swa && crate::Engine::wkv_on())?;
8734 }
8735 return Ok(e.matmul(&fa.wo, &attn, t)?);
8736 }
8737 if hd == 256 && swa && base_len + 1 >= win
8745 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8746 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
8747 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
8748 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8749 e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, 0,
8750 t, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8751 return Ok(e.matmul(&fa.wo, &attn, t)?);
8752 }
8753 for i in 0..t {
8754 let avail = base_len + i + 1;
8755 let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
8756 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
8757 (off_tok + t_kv) * kvl.k_tok_bytes);
8758 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
8759 (off_tok + t_kv) * kvl.v_tok_bytes);
8760 let qi = e.view(&q, t * nh * hd);
8761 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
8762 let mut q_one = e.uninit(nh * hd)?;
8763 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
8764 let mut a_one = e.uninit(nh * hd)?;
8765 if swa && avail > win && hd == 256
8769 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8770 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
8771 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
8772 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
8773 e.fa_decode_rows_w(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, &kvl.len_d, 0,
8774 1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8775 } else if !swa && hd == 512 && avail >= crate::fa512_min_tkv()
8776 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8777 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
8778 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
8779 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
8780 e.fa_decode_rows(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, avail - 1, 1,
8781 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8782 Some((&kvl.len_d, 0)), false, false, None)?;
8783 } else {
8784 e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
8785 kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
8786 }
8787 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
8788 }
8789 Ok(e.matmul(&fa.wo, &attn, t)?)
8790 }
8791
8792 pub(crate) fn gemma4_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
8795 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8796 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
8801 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
8802 }
8803 if crate::pp::pp_cuts(self.layers.len()).is_some() {
8804 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
8805 }
8806 let n_embd = self.cfg.n_embd as usize;
8807 let eps = self.cfg.rms_eps;
8808 let pos_d = e.htod_i32(&[cache.pos as i32])?;
8809 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
8810 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8811 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8814 let n_layers = self.layers.len();
8815 for (il, layer) in self.layers.iter().enumerate() {
8816 let (hq, hdq) = match h_carry.take() {
8817 Some(p) => p,
8818 None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
8819 };
8820 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8821 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
8822 let mut cur = e.uninit(n_embd)?;
8823 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
8824 let next_norm = if il + 1 < n_layers {
8825 Some(self.layers[il + 1].attn_norm.float_data())
8826 } else { None };
8827 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
8828 x = xn;
8829 h_carry = hn;
8830 }
8831 let mut hn = e.uninit(n_embd)?;
8832 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8833 let h_seed = e.clone_dtod(&x)?;
8834 let mut ld = e.matmul(&self.output, &hn, 1)?;
8835 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8836 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
8838 let logits = e.dtoh(&ld)?;
8839 cache.pos += 1;
8840 Ok((logits, h_seed))
8841 }
8842
8843 fn gemma4_decode_layers(&self, e: &Engine, mut x: CudaSlice<f32>, lo: usize, hi: usize,
8851 pos_d: &CudaSlice<i32>, cache: &mut Cache)
8852 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8853 let n_embd = self.cfg.n_embd as usize;
8854 let eps = self.cfg.rms_eps;
8855 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8856 for il in lo..hi {
8857 let layer = &self.layers[il];
8858 let (hq, hdq) = match h_carry.take() {
8859 Some(p) => p,
8860 None => e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?,
8862 };
8863 let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8864 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
8865 let mut cur = e.uninit(n_embd)?;
8866 e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
8867 let next_norm = if il + 1 < hi {
8868 Some(self.layers[il + 1].attn_norm.float_data())
8869 } else { None };
8870 let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
8871 x = xn;
8872 h_carry = hn;
8873 }
8874 Ok(x)
8875 }
8876
8877 fn gemma4_decode_step_h_pp2(&self, e: &Engine, token: u32, cache: &mut Cache, split: usize)
8885 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8886 if crate::pp::pp2_streams_off() {
8887 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
8888 }
8889 let rt = crate::pp::Pp2Rt::get(e)?;
8890 let e0 = rt.engine(0, e);
8891 let e1 = rt.engine(1, e);
8892 let n_embd = self.cfg.n_embd as usize;
8893 let eps = self.cfg.rms_eps;
8894 let pos = cache.pos as i32;
8895
8896 let slot = {
8898 let _st0 = rt.enter(0);
8899 let pos_d = e0.htod_i32(&[pos])?;
8900 #[cfg(debug_assertions)]
8901 crate::debug_assert_tensor_stream_device(&pos_d, &e0.stream(),
8902 "gemma4_decode_step_h_pp2.stage0.pos_d");
8903 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
8904 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8905 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
8906 rt.tx(0, &x, n_embd)?
8907 };
8908
8909 let _st1 = rt.enter(1);
8911 let pos_d = e1.htod_i32(&[pos])?;
8912 #[cfg(debug_assertions)]
8913 crate::debug_assert_tensor_stream_device(&pos_d, &e1.stream(),
8914 "gemma4_decode_step_h_pp2.stage1.pos_d");
8915 let x = rt.rx(0, slot, n_embd)?;
8916 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
8917
8918 let mut hn = e1.uninit(n_embd)?;
8919 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8920 let h_seed = e1.clone_dtod(&x)?;
8921 let mut ld = e1.matmul(&self.output, &hn, 1)?;
8922 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8923 e1.softcap(&mut ld, cap, self.output.out_features())?;
8924 self.gemma4_suppress(e1, &mut ld, 1)?;
8925 let logits = e1.dtoh(&ld)?;
8926 cache.pos += 1;
8927 Ok((logits, h_seed))
8928 }
8929
8930 fn gemma4_decode_step_h_pp2_samestream(&self, e: &Engine, token: u32, cache: &mut Cache,
8933 split: usize)
8934 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8935 let n_embd = self.cfg.n_embd as usize;
8936 let eps = self.cfg.rms_eps;
8937 let pos_d = e.htod_i32(&[cache.pos as i32])?;
8938
8939 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
8941 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8942 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
8943
8944 let boundary_tx = e.clone_dtod(&x)?;
8946 let boundary_rx = e.clone_dtod(&boundary_tx)?;
8947
8948 let x = self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
8950
8951 let mut hn = e.uninit(n_embd)?;
8952 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8953 let h_seed = e.clone_dtod(&x)?;
8954 let mut ld = e.matmul(&self.output, &hn, 1)?;
8955 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8956 e.softcap(&mut ld, cap, self.output.out_features())?;
8957 self.gemma4_suppress(e, &mut ld, 1)?;
8958 let logits = e.dtoh(&ld)?;
8959 cache.pos += 1;
8960 Ok((logits, h_seed))
8961 }
8962}
8963
8964impl HybridModel {
8983 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
8986 let geometry = self.cfg.layer_geometry(il as u32)
8987 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
8988 debug_assert_eq!(
8989 geometry.attention_gate,
8990 memra_gguf::config::AttentionGateKind::SeparateHead
8991 );
8992 geometry
8993 }
8994
8995 #[allow(clippy::too_many_arguments)]
9055 fn step35_attn_pre_wo(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
9056 hg: Option<&CudaSlice<f32>>, gt_pre: Option<&CudaSlice<f32>>,
9057 pos_d: &CudaSlice<i32>, t: usize,
9058 cache: Option<&mut Cache>, il: usize, seq_end: usize)
9059 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9060 let geometry = self.step35_geom(il);
9061 let hd = geometry.head_dim_k as usize;
9062 let nkv = geometry.n_head_kv as usize;
9063 let nh = geometry.n_head as usize;
9064 let rbase = geometry.rope_base;
9065 let scale = geometry.attention_scale();
9066 let swa = geometry.window.is_some();
9067 let eps = self.cfg.rms_eps;
9068 let win = geometry.window.unwrap_or(0) as usize;
9069 let n_rot = geometry.n_rot as usize;
9070
9071 let v = g3.pop().unwrap();
9072 let k0 = g3.pop().unwrap();
9073 let q0 = g3.pop().unwrap();
9074
9075 let mut q = e.uninit(t * nh * hd)?;
9079 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
9080 let mut k = e.uninit(t * nkv * hd)?;
9081 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
9082 let ff = if geometry.rope_factors {
9083 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
9084 } else {
9085 None
9086 };
9087 #[cfg(debug_assertions)]
9088 if let Some(ff) = ff {
9089 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
9090 "step35_attn_pre_wo.rope_freqs");
9091 }
9092 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
9093
9094 let mut attn = e.uninit(t * nh * hd)?;
9095 match cache {
9096 Some(cache) => {
9097 let base_len = cache.kv[il].as_ref().unwrap().len;
9098 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
9100 let legacy_calllocal =
9101 std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
9102 let off = if swa {
9103 let raw = base_len.saturating_sub(win - 1);
9104 if legacy_tkv || legacy_calllocal { raw } else { raw & !31usize }
9105 } else {
9106 0
9107 };
9108 {
9109 let kvl = cache.kv[il].as_mut().unwrap();
9110 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
9111 let write_row = e.prepare_kv_append(kvl, off, t)?;
9112 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, write_row, t,
9113 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
9114 kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
9115 kvl.len += t;
9116 let new_len = kvl.len as i32;
9117 e.set_i32_one(&mut kvl.len_d, new_len)?;
9118 }
9119 let kvl = cache.kv[il].as_ref().unwrap();
9120 let t_kv = base_len + t - off;
9143 let physical = kvl.physical_rows(off, off + t_kv)?;
9144 let k_view = e.view_u8_range(&kvl.k, physical.start * kvl.k_tok_bytes,
9145 physical.end * kvl.k_tok_bytes);
9146 let v_view = e.view_u8_range(&kvl.v, physical.start * kvl.v_tok_bytes,
9147 physical.end * kvl.v_tok_bytes);
9148 let swa_naive = if legacy_tkv { t_kv > win } else { seq_end > win };
9160 if swa && swa_naive {
9161 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
9174 e.sdpa_naive_w_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh,
9175 nkv, t, t_kv, scale, true, win,
9176 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
9177 } else {
9178 e.fa_prefill_view_ws_w_hd128(&q, &k_view, &v_view, &mut attn, hd, nh,
9179 nkv, t, t_kv, scale, true, win,
9180 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
9181 }
9182 } else if std::env::var("MEMRA_NOFA").is_ok() {
9183 e.sdpa_naive_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9184 t, t_kv, scale, true,
9185 kvl.k_tok_bytes, kvl.v_tok_bytes)?;
9186 } else {
9187 e.fa_prefill_view_ws(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9192 t, t_kv, scale, true,
9193 kvl.k_tok_bytes, kvl.v_tok_bytes,
9194 crate::Engine::kv_fp8_on())?;
9195 }
9196 }
9197 None => {
9198 debug_assert_eq!(seq_end, t, "step35 cacheless prefill is monolithic (seq_end == t)");
9203 if swa && seq_end > win {
9204 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
9205 } else if std::env::var("MEMRA_NOFA").is_ok() {
9206 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9207 } else {
9208 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9209 }
9210 }
9211 }
9212
9213 let gw = fa.attn_gate.as_ref()
9216 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
9217 let gt_owned = if gt_pre.is_none() {
9218 Some(e.matmul(
9219 gw,
9220 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
9221 t,
9222 )?)
9223 } else {
9224 None
9225 };
9226 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
9227 let mut ag = e.uninit(t * nh * hd)?;
9228 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
9229 Ok(ag)
9230 }
9231
9232 pub(crate) fn step35_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
9235 pos_d: &CudaSlice<i32>, t: usize, il: usize)
9236 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9237 let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
9238 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
9240 Ok(e.matmul(&fa.wo, &ag, t)?)
9241 }
9242
9243 #[allow(clippy::too_many_arguments)]
9250 pub(crate) fn step35_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
9251 hx: Option<&CudaSlice<u8>>, pos_d: &CudaSlice<i32>, t: usize,
9252 cache: &mut Cache, il: usize, seq_end: usize)
9253 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9254 let g3 = match hx {
9255 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
9256 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
9257 };
9258 let ag = self.step35_attn_pre_wo(
9259 e,
9260 fa,
9261 g3,
9262 Some(h),
9263 None,
9264 pos_d,
9265 t,
9266 Some(cache),
9267 il,
9268 seq_end,
9269 )?;
9270 Ok(e.matmul(&fa.wo, &ag, t)?)
9271 }
9272
9273 #[allow(clippy::too_many_arguments)]
9283 pub(crate) fn step35_decode_attn(&self, e: &Engine, fa: &FullAttnLayer, il: usize,
9284 h: &CudaSlice<f32>,
9285 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9286 pos_d: &CudaSlice<i32>, cache: &mut Cache)
9287 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9288 let geometry = self.step35_geom(il);
9289 let hd = geometry.head_dim_k as usize;
9290 let nkv = geometry.n_head_kv as usize;
9291 let nh = geometry.n_head as usize;
9292 let rbase = geometry.rope_base;
9293 let scale = geometry.attention_scale();
9294 let swa = geometry.window.is_some();
9295 let eps = self.cfg.rms_eps;
9296 let win = geometry.window.unwrap_or(0) as usize;
9297 let n_rot = geometry.n_rot as usize;
9298 let n_embd = self.cfg.n_embd as usize;
9299 let gw = fa.attn_gate.as_ref()
9300 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
9301
9302 let (q0, k0, v0, gt) = match pre_q {
9303 Some((hq, hdq)) => {
9304 debug_assert!(e.uses_q8_1_fast(gw),
9305 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
9306 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast");
9307 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
9308 Some(t3) => t3,
9309 None => (e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
9310 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
9311 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?),
9312 };
9313 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
9314 (a, b, c, gt)
9315 }
9316 None => {
9317 if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk)
9318 && e.uses_q8_1_fast(&fa.wv) && e.uses_q8_1_fast(gw) {
9319 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
9320 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
9321 Some(t3) => t3,
9322 None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
9323 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
9324 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
9325 };
9326 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
9327 (a, b, c, gt)
9328 } else {
9329 (e.matmul(&fa.wq, h, 1)?, e.matmul(&fa.wk, h, 1)?,
9330 e.matmul(&fa.wv, h, 1)?, e.matmul(gw, h, 1)?)
9331 }
9332 }
9333 };
9334
9335 let mut q = e.uninit(nh * hd)?;
9336 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
9337 let mut k = e.uninit(nkv * hd)?;
9338 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
9339 let ff = if swa { None } else {
9340 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
9341 };
9342 #[cfg(debug_assertions)]
9343 if let Some(ff) = ff {
9344 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
9345 "step35_decode_attn.rope_freqs");
9346 }
9347 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
9348
9349 if std::env::var("MEMRA_NOFA").is_ok() {
9350 return Err("MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
9351 cache; unset MEMRA_NOFA to use fa_decode".into());
9352 }
9353 let kvl = cache.kv[il].as_mut().unwrap();
9354 let next_len = kvl.len + 1;
9355 let (off, t_kv) = if swa && next_len > win { (next_len - win, win) } else { (0, next_len) };
9356 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
9357 e.append_kv_quantized(&k, &v0, &mut kvl.k, &mut kvl.v, write_row,
9358 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
9359 crate::Engine::kv_fp8_on())?;
9360 kvl.len = next_len;
9361 let physical = kvl.physical_rows(off, off + t_kv)?;
9362 let k_view = e.view_u8_range(&kvl.k, physical.start * kvl.k_tok_bytes,
9363 physical.end * kvl.k_tok_bytes);
9364 let v_view = e.view_u8_range(&kvl.v, physical.start * kvl.v_tok_bytes,
9365 physical.end * kvl.v_tok_bytes);
9366 let mut attn = e.uninit(nh * hd)?;
9367 e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
9368 kvl.k_tok_bytes, kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
9369
9370 let mut ag = e.uninit(nh * hd)?;
9371 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
9372 Ok(e.matmul(&fa.wo, &ag, 1)?)
9373 }
9374}
9375
9376impl HybridModel {
9385 pub fn is_gemma4_e4b(&self) -> bool {
9386 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
9387 }
9388
9389 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
9393 let g = self.cfg.gemma4.as_ref().unwrap();
9394 let swa = g.swa_pattern[il];
9395 let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
9396 let Mixer::Full(fa) = &self.layers[il].mixer else { panic!("e4b layer {il} not full-attn") };
9397 let nh = fa.wq.out_features() / hd;
9398 let nkv = fa.wk.out_features() / hd;
9399 (hd, nkv, nh, if swa { g.rope_base_swa } else { g.rope_base_global }, 1.0, swa)
9400 }
9401
9402 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
9404 self.layers[il].gemma4.as_ref()
9405 .and_then(|b| b.e4b.as_ref())
9406 .and_then(|e4| e4.kv_share.map(|t| t as usize))
9407 }
9408
9409 fn gemma4_e4b_inp_pl(&self, e: &Engine, tokens: &[u32], x_scaled: &CudaSlice<f32>, t: usize)
9414 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9415 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
9416 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
9417 }
9418
9419 fn gemma4_e4b_inp_pl_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
9421 x_scaled: &CudaSlice<f32>, t: usize)
9422 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9423 let aux = self.gemma4_aux.as_ref().unwrap();
9424 let m = aux.e4b.as_ref().unwrap();
9425 let n_embd = self.cfg.n_embd as usize;
9426 let n_layer = self.layers.len();
9427 let width = m.n_epl * n_layer;
9428 let tbl = m.tok_tbl_gpu.get_or_init(|| {
9429 e.upload_u8(&m.tok_embd_bytes).expect("e4b per-layer token table upload")
9430 });
9431 let mut a = e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt,
9432 m.tok_embd_row_bytes)?;
9433 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
9434 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
9435 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
9436 let mut pn = e.uninit(t * width)?;
9437 e.rms_norm(&p, m.proj_norm.float_data(), &mut pn, m.n_epl, t * n_layer,
9438 self.cfg.rms_eps)?;
9439 let mut out = e.uninit(t * width)?;
9440 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
9441 Ok(out)
9442 }
9443
9444 #[allow(clippy::too_many_arguments)]
9449 fn gemma4_e4b_attn(&self, e: &Engine, il: usize,
9450 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
9451 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
9452 dc_bucket: Option<usize>)
9453 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9454 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
9455 let eps = self.cfg.rms_eps;
9456 let aux = self.gemma4_aux.as_ref().unwrap();
9457 let ones = aux.ones(e);
9458 #[cfg(debug_assertions)]
9459 crate::debug_assert_tensor_stream_device(ones, &e.stream(),
9460 "gemma4_e4b_attn.ones");
9461 let Mixer::Full(fa) = &self.layers[il].mixer else { unreachable!() };
9462 let h0 = e.zeros(0)?;
9466 let h = &h0;
9467
9468 let ff = if swa { None } else {
9469 Some(aux.rope_freqs(e).expect("e4b global rope needs rope_freqs.weight"))
9470 };
9471 #[cfg(debug_assertions)]
9472 if let Some(ff) = ff {
9473 crate::debug_assert_tensor_stream_device(ff, &e.stream(),
9474 "gemma4_e4b_attn.rope_freqs");
9475 }
9476 let share = self.gemma4_e4b_kv_target(il);
9477 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
9479 let mut q;
9480 if let Some(_tgt) = share {
9481 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
9482 q = e.uninit(t * nh * hd)?;
9483 let mut kdummy = e.uninit(1)?;
9486 let mut vdummy = e.uninit(1)?;
9487 e.rms_norm_qkv_rope(&q0, &q0, &q0, fa.q_norm.float_data(),
9488 fa.q_norm.float_data(), ones,
9489 &mut q, &mut kdummy, &mut vdummy, hd, nh * t, 0,
9490 pos_d, nh, 1, base, 1.0, ff, eps)?;
9491 } else {
9492 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
9496 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
9497 q = e.uninit(t * nh * hd)?;
9498 let mut k = e.uninit(t * nkv * hd)?;
9499 let mut v = e.uninit(t * nkv * hd)?;
9500 if t == 1 && cat.is_some() {
9501 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
9502 e.rms_norm_qkv_rope_cat(&qkv0, fa.q_norm.float_data(), fa.k_norm.float_data(),
9503 ones, &mut q, &mut k, &mut v, hd, nh, nkv,
9504 pos_d, nh, nkv, base, 1.0, ff, eps)?;
9505 } else {
9506 let (q0, k0, v0) = match if t == 1 {
9507 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
9508 } else {
9509 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9512 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
9513 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
9514 } else { None }
9515 } {
9516 Some(triple) => triple,
9517 None => (e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
9518 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
9519 e.matmul_pre(&fa.wv, hq, hdq, h, t)?), };
9521 e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(),
9524 fa.k_norm.float_data(), ones, &mut q, &mut k, &mut v,
9525 hd, nh * t, nkv * t, pos_d, nh, nkv, base, 1.0, ff, eps)?;
9526 }
9527 let kvl = cache.kv[il].as_mut().unwrap();
9528 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9532 if dc_bucket.is_some() {
9533 debug_assert!(t == 1);
9538 e.append_kv_quantized_row_dc_inc(&k, &v, &mut kvl.k, &mut kvl.v,
9540 &mut kvl.len_d, kvl.kv_dim_k, kvl.kv_dim_v,
9541 kvl.k_tok_bytes, kvl.v_tok_bytes, cls)?;
9542 } else {
9543 e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
9544 kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
9545 kvl.v_tok_bytes, cls)?;
9546 kvl.len += t;
9547 }
9548 kv_f32 = Some((k, v));
9549 }
9550 let kvl_idx = share.unwrap_or(il);
9553 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
9554 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
9556 let mut attn = e.uninit(t * nh * hd)?;
9557 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
9569 if let Some((kf, vf)) = &kv_f32 {
9570 if hd == 256 && t <= win {
9571 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9572 return Ok(e.matmul(&fa.wo, &attn, t)?);
9573 }
9574 if hd == 256 && swa && t > win {
9575 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true,
9576 win)?;
9577 return Ok(e.matmul(&fa.wo, &attn, t)?);
9578 }
9579 if hd == 512 && !swa {
9580 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale,
9581 true)?;
9582 return Ok(e.matmul(&fa.wo, &attn, t)?);
9583 }
9584 } else if share.is_some() {
9585 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9586 let k_view = e.view_u8(&kvl.k, kvl.k.len());
9587 let v_view = e.view_u8(&kvl.v, kvl.v.len());
9588 if hd == 256 && (!swa || t <= win) {
9589 e.fa_prefill_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t, t,
9591 scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9592 return Ok(e.matmul(&fa.wo, &attn, t)?);
9593 }
9594 let kv_dim = nkv * hd;
9597 let mut kf = e.uninit(t * kv_dim)?;
9598 let mut vf = e.uninit(t * kv_dim)?;
9599 e.fa_dequant_kv_view_f32(&k_view, &v_view, &mut kf, &mut vf, kv_dim, kv_dim,
9600 t, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9601 if hd == 512 {
9602 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale,
9603 true)?;
9604 } else {
9605 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true,
9606 win)?;
9607 }
9608 return Ok(e.matmul(&fa.wo, &attn, t)?);
9609 }
9610 }
9611 if let Some(bucket) = dc_bucket {
9612 assert!(t == 1);
9617 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
9623 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
9624 } else { bucket };
9625 let k_view = e.view_u8(&kvl.k, kvl.k.len());
9626 let v_view = e.view_u8(&kvl.v, kvl.v.len());
9627 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9628 if crate::Engine::wpf_level() >= 1 {
9636 e.prefetch_weight_l2(&fa.wo)?;
9637 }
9638 if e.uses_q8_1_fast(&fa.wo) {
9641 let mut oq = e.alloc_i8_uninit(nh * hd)?;
9642 let mut od = e.zeros(nh * hd / 32)?;
9643 e.fa_decode_dc_q8(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9644 &kvl.len_d, bucket, scale,
9645 kvl.k_tok_bytes, kvl.v_tok_bytes, g,
9646 Some((&mut oq, &mut od)))?;
9647 return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
9648 }
9649 e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9650 &kvl.len_d, bucket, scale,
9651 kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9652 return Ok(e.matmul(&fa.wo, &attn, t)?);
9653 }
9654 for i in 0..t {
9655 let avail = base_len + i + 1;
9656 let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
9657 let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
9658 (off_tok + t_kv) * kvl.k_tok_bytes);
9659 let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
9660 (off_tok + t_kv) * kvl.v_tok_bytes);
9661 let qv = e.view(&q, t * nh * hd);
9662 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
9663 let mut q_one = e.uninit(nh * hd)?;
9664 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
9665 let mut a_one = e.uninit(nh * hd)?;
9666 e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
9670 kvl.k_tok_bytes, kvl.v_tok_bytes,
9671 (!swa && crate::Engine::gkv_on())
9672 || (swa && crate::Engine::wkv_on()))?;
9673 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
9674 }
9675 Ok(e.matmul(&fa.wo, &attn, t)?)
9676 }
9677
9678 fn gemma4_e4b_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
9683 head_last: bool)
9684 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9685 let n_embd = self.cfg.n_embd as usize;
9686 let t = tokens.len();
9687 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9688 let pos_d = e.htod_i32(&pos)?;
9689 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
9690 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9691 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
9692 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
9693 }
9694
9695 fn gemma4_e4b_trunk_core(&self, e: &Engine, x_in: CudaSlice<f32>, inp_pl: CudaSlice<f32>,
9699 pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
9700 dc_bucket: Option<usize>, cap_logits: bool, head_last: bool)
9701 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9702 let n_embd = self.cfg.n_embd as usize;
9703 let eps = self.cfg.rms_eps;
9704 let n_layer = self.layers.len();
9705 let mut x = x_in;
9706 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
9707 let n_epl = aux_e4b.n_epl;
9708
9709 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9715 for il in 0..n_layer {
9716 let layer = &self.layers[il];
9717 let (hq, hdq) = match h_carry.take() {
9718 Some(p) => p,
9719 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
9720 };
9721 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
9722 let bits = layer.gemma4.as_ref().unwrap();
9725 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
9726 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
9737 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
9738 e, layer, &o, &x, t, Some(layer.post_attn_norm.float_data()), fuse_exit)?;
9739 let mut resid = e.uninit(t * n_embd)?;
9740 let g = if fuse_exit {
9746 let (rq, rd) = e.rms_pre_add_q8_1(&sn, bits.post_ffw_norm.float_data(),
9748 &attn_out, &mut resid, n_embd, t,
9749 self.cfg.rms_eps)?;
9750 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
9751 } else {
9752 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
9753 e.matmul(&e4b.inp_gate, &resid, t)?
9754 };
9755 let mut act = e.uninit(t * n_epl)?;
9756 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
9757 let ipv = e.view(&inp_pl, n_epl * n_layer);
9758 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
9759 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
9760 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
9761 } else {
9762 let mut inp_this = e.uninit(t * n_epl)?;
9763 e.copy_rows_strided(&inp_pl, &mut inp_this, n_epl, t, n_epl * n_layer,
9764 il * n_epl)?;
9765 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
9766 e.matmul(&e4b.proj, &act, t)?
9767 };
9768 let next_norm = if il + 1 < n_layer {
9771 self.layers[il + 1].attn_norm.float_data()
9772 } else {
9773 self.output_norm.float_data()
9774 };
9775 let mut xn = e.uninit(t * n_embd)?;
9776 let pair = e.rms_pre_add_scale_rms_norm_q8_1(&y, e4b.post_norm.float_data(),
9777 &resid, bits.layer_scale, next_norm,
9778 &mut xn, n_embd, t, eps)?;
9779 h_carry = Some(pair);
9780 x = xn;
9781 }
9782 let (oq, odq) = h_carry.take().unwrap();
9786 let h0 = e.zeros(0)?;
9787 let hm = if head_last { 1 } else { t };
9788 let (hq, hd) = if head_last && t > 1 {
9789 let mut q1 = e.uninit_i8(n_embd)?;
9790 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
9791 let nb = n_embd / 32;
9792 let mut d1 = e.uninit(nb)?;
9793 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
9794 (q1, d1)
9795 } else {
9796 (oq, odq)
9797 };
9798 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
9799 if cap_logits {
9803 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9804 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
9805 }
9806 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
9808 }
9809
9810 pub fn gemma4_e4b_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
9817 t: usize, pos0: usize, cache: &mut Cache)
9818 -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9819 let n_embd = self.cfg.n_embd as usize;
9820 let eps = self.cfg.rms_eps;
9821 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9822 let pos_d = e.htod_i32(&pos)?;
9823 let embd_gpu = self.embd_gpu.get_or_init(|| {
9824 e.upload_u8(&self.embd.raw).expect("embed table upload")
9825 });
9826 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
9827 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
9828 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9829 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
9830 let (ld, xp) = self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true,
9831 false)?;
9832 let n_vocab = self.output.out_features();
9835 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
9836 for i in 0..t {
9837 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
9838 }
9839 let mut hn = e.uninit(t * n_embd)?;
9840 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9841 cache.pos += t;
9842 Ok((vam, hn))
9843 }
9844
9845 pub(crate) fn gemma4_e4b_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
9848 cache: &mut Cache)
9849 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9850 let n_embd = self.cfg.n_embd as usize;
9851 let eps = self.cfg.rms_eps;
9852 let t = tokens.len();
9853 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
9854 let mut hn = e.uninit(t * n_embd)?;
9855 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9856 cache.pos += t;
9857 Ok((e.dtoh(&ld)?, hn))
9858 }
9859
9860 pub fn gemma4_e4b_decode_step_dcg(&self, e: &Engine, token_d: &mut CudaSlice<u32>,
9866 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
9867 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
9868 n_vocab: usize, bucket: usize)
9869 -> Result<(), Box<dyn std::error::Error>> {
9870 let n_embd = self.cfg.n_embd as usize;
9871 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
9872 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
9873 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
9874 let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket),
9875 false, false)?;
9876 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
9877 e.inc_seqlen(pos_d)?;
9878 Ok(())
9879 }
9880
9881 #[allow(clippy::too_many_arguments)]
9889 pub fn gemma4_e4b_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
9890 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
9891 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
9892 n_vocab: usize)
9893 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9894 let n_embd = self.cfg.n_embd as usize;
9895 let eps = self.cfg.rms_eps;
9896 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
9897 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
9898 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
9899 let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false,
9900 false)?;
9901 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
9902 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
9903 e.inc_seqlen(pos_d)?;
9904 cache.pos += 1;
9905 let _ = eps;
9906 Ok(tok_out)
9907 }
9908
9909 pub(crate) fn gemma4_e4b_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
9912 -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9913 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
9914 let logits = e.dtoh(&ld)?;
9915 cache.pos += 1;
9916 Ok((logits, x))
9917 }
9918
9919 pub(crate) fn gemma4_e4b_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
9923 -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9924 if cache.pos != 0 {
9927 return Err("e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
9928 call or decode tokenwise".into());
9929 }
9930 let n_embd = self.cfg.n_embd as usize;
9931 let t = tokens.len();
9932 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
9933 cache.pos += t;
9934 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
9936 let row = xv.slice((t - 1) * n_embd..t * n_embd);
9937 let mut h_seed = e.uninit(n_embd)?;
9938 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
9939 Ok((last, h_seed, x))
9940 }
9941
9942 pub(crate) fn gemma4_e4b_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
9944 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
9945 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
9946 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
9947 Ok(e.dtoh(&ld)?) }
9949}
9950
9951#[cfg(test)]
9952mod prime_chunk_schedule_tests {
9953 use super::{
9954 dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges, fixed_prime_chunk_ranges_for_ring,
9955 PRIME_MIN_T,
9956 PRIME_PIPE_MIN_CHUNK,
9957 };
9958
9959 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
9960 ranges.iter().map(|(start, end)| end - start).collect()
9961 }
9962
9963 fn auto_chunk(t: usize) -> usize {
9964 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
9965 }
9966
9967 #[test]
9968 fn fixed_schedule_retains_measured_geometry() {
9969 assert_eq!(
9970 sizes(&fixed_prime_chunk_ranges(461, 128)),
9971 vec![128, 128, 128, 77]
9972 );
9973 assert_eq!(
9974 sizes(&fixed_prime_chunk_ranges(1833, 230)),
9975 vec![230, 230, 230, 230, 230, 230, 230, 223]
9976 );
9977 assert_eq!(
9978 sizes(&fixed_prime_chunk_ranges(4096, 512)),
9979 vec![512; 8]
9980 );
9981 let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
9982 assert_eq!(capped, vec![4096, 4088, 16]);
9983 assert!(capped.iter().all(|&rows| rows <= 4096));
9984 assert_eq!(
9985 sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
9986 vec![4100],
9987 "flag-off schedule remains byte-for-byte the legacy monolithic tail",
9988 );
9989 }
9990
9991 #[test]
9992 fn dynamic_schedule_matches_registered_shapes() {
9993 let cases = [
9994 (461, vec![64, 141, 132, 124]),
9995 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
9996 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
9997 ];
9998 for (t, expected) in cases {
9999 let chunk = auto_chunk(t);
10000 let fixed = fixed_prime_chunk_ranges(t, chunk);
10001 assert_eq!(
10002 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
10003 expected
10004 );
10005 }
10006 }
10007
10008 #[test]
10009 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
10010 for t in 256..=8192 {
10011 let chunk = auto_chunk(t);
10012 let fixed = fixed_prime_chunk_ranges(t, chunk);
10013 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
10014 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
10015 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
10016 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
10017 for pair in dynamic.windows(2) {
10018 assert_eq!(pair[0].1, pair[1].0, "T={t}");
10019 }
10020 assert!(
10021 dynamic
10022 .iter()
10023 .all(|(start, end)| end - start >= PRIME_MIN_T),
10024 "T={t} sizes={:?}",
10025 sizes(&dynamic)
10026 );
10027 if dynamic.len() >= 3 {
10028 let chunk_sizes = sizes(&dynamic);
10029 assert!(
10030 chunk_sizes[0] < chunk_sizes[1],
10031 "T={t} sizes={chunk_sizes:?}"
10032 );
10033 assert!(
10034 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
10035 "T={t} sizes={chunk_sizes:?}"
10036 );
10037 }
10038 }
10039 }
10040}
10041
10042#[cfg(test)]
10043mod page_prefetch_tests {
10044 use super::{
10045 grouped_worker_prefetch_position, page_prefetch_positions,
10046 page_prefetch_window_from_values, worker_prefetch_positions,
10047 };
10048
10049 #[test]
10050 fn page_prefetch_window_keeps_existing_opt_in_default() {
10051 assert_eq!(page_prefetch_window_from_values(false, None), 0);
10052 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
10053 assert_eq!(page_prefetch_window_from_values(true, None), 1);
10054 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
10055 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
10056 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
10057 }
10058
10059 #[test]
10060 fn rolling_page_prefetch_advises_each_future_expert_once() {
10061 let advised: Vec<_> = (0..7)
10062 .flat_map(|position| page_prefetch_positions(position, 7, 3))
10063 .collect();
10064 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
10065
10066 let one_ahead: Vec<_> = (0..4)
10067 .flat_map(|position| page_prefetch_positions(position, 4, 1))
10068 .collect();
10069 assert_eq!(one_ahead, vec![1, 2, 3]);
10070 assert!(page_prefetch_positions(0, 4, 0).is_empty());
10071 }
10072
10073 #[test]
10074 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
10075 assert_eq!(grouped_worker_prefetch_position(0, None), None);
10076 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
10077 .chain((0..4).filter_map(|position| {
10078 grouped_worker_prefetch_position(4, Some(position))
10079 }))
10080 .collect();
10081 assert_eq!(positions, vec![0, 1, 2, 3]);
10082 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
10083 }
10084
10085 #[test]
10086 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
10087 let queued: Vec<_> = (0..8)
10088 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
10089 .collect();
10090 assert_eq!(queued, (0..8).collect::<Vec<_>>());
10091
10092 let one_at_a_time: Vec<_> = (0..4)
10093 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
10094 .collect();
10095 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
10096 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
10097 }
10098}
10099
10100pub struct G4DcSlots {
10101 x: CudaSlice<f32>, xn: CudaSlice<f32>, cur: CudaSlice<f32>,
10102 hq: CudaSlice<i8>, hd_: CudaSlice<f32>,
10103 q0: CudaSlice<f32>, k0: CudaSlice<f32>, v0: CudaSlice<f32>,
10104 q: CudaSlice<f32>, k: CudaSlice<f32>, v: CudaSlice<f32>,
10105 attn: CudaSlice<f32>, o: CudaSlice<f32>,
10106 attn_out: CudaSlice<f32>, zsh: CudaSlice<f32>,
10107 zq: CudaSlice<i8>, zd: CudaSlice<f32>,
10108 gate: CudaSlice<f32>, up: CudaSlice<f32>,
10109 act: CudaSlice<f32>, actq: CudaSlice<i8>, actd: CudaSlice<f32>,
10110 f0: CudaSlice<f32>, sn: CudaSlice<f32>,
10111 hn: CudaSlice<f32>, logits: CudaSlice<f32>,
10112}