1use crate::Engine;
6use crate::cache::Cache;
7use cudarc::driver::CudaSlice;
8use memra_gguf::config::{ModelConfig, SwigluClamp};
9
10pub struct PrimeSlabs {
14 pub t_cap: usize,
15 pub h: CudaSlice<f32>,
16 pub x1: CudaSlice<f32>,
17 pub z: CudaSlice<f32>,
18 pub act: CudaSlice<f32>,
19 pub xa: CudaSlice<f32>,
20 pub xb: CudaSlice<f32>,
21 pub h16: CudaSlice<u8>,
22 pub z16: CudaSlice<u8>,
23 pub gate: CudaSlice<f32>, pub up: CudaSlice<f32>, pub ffn_out: CudaSlice<f32>, pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
33 pub mixed: CudaSlice<f32>,
37 pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
38 pub seg_t: usize,
39}
40
41unsafe impl Send for PrimeSlabs {}
44
45fn shexp_gate_up_t1(
54 e: &Engine,
55 gate_shexp: &crate::model::GpuTensor,
56 up_shexp: &crate::model::GpuTensor,
57 z: &CudaSlice<f32>,
58 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
59) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
60 let is_nvfp4 = |w: &crate::model::GpuTensor| matches!(w, crate::model::GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_NVFP4);
61 if is_nvfp4(gate_shexp) && is_nvfp4(up_shexp) {
62 let pair = match zq8 {
66 Some((zq, zd)) => e.matmul_nvfp4_fused2(gate_shexp, up_shexp, zq, zd, 1)?,
67 None => {
68 let (zq, zd) = e.quantize_q8_1(z, 1, gate_shexp.in_features())?;
69 e.matmul_nvfp4_fused2(gate_shexp, up_shexp, &zq, &zd, 1)?
70 }
71 };
72 if let Some(pair) = pair {
73 return Ok(pair);
74 }
75 }
76 match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
77 Some(pair) => Ok(pair),
78 None => Ok((e.matmul(gate_shexp, z, 1)?, e.matmul(up_shexp, z, 1)?)),
79 }
80}
81
82fn active_matrix_values(
83 available: usize,
84 rows: usize,
85 columns: usize,
86 label: &str,
87) -> Result<usize, String> {
88 let required = rows
89 .checked_mul(columns)
90 .ok_or_else(|| format!("{label} shape overflows: {rows}x{columns}"))?;
91 if available < required {
92 return Err(format!(
93 "{label} has {available} values, fewer than the active {rows}x{columns} ({required})"
94 ));
95 }
96 Ok(required)
97}
98
99fn step_grouped_decode_shape(prefill: bool, tokens: usize) -> bool {
100 !prefill && tokens == 1
101}
102
103fn parse_step_ep_grouped_prefill(value: Option<&str>) -> Result<bool, String> {
104 match value {
105 None | Some("") | Some("0") => Ok(false),
106 Some("1") => Ok(true),
107 Some(value) => Err(format!(
108 "MEMRA_STEP_EP_GROUPED_PREFILL={value:?} is invalid; expected 0 or 1"
109 )),
110 }
111}
112
113fn step_ep_grouped_prefill_enabled() -> Result<bool, String> {
114 parse_step_ep_grouped_prefill(
115 std::env::var("MEMRA_STEP_EP_GROUPED_PREFILL")
116 .ok()
117 .as_deref(),
118 )
119}
120
121fn step_grouped_prefill_shape(enabled: bool, prefill: bool, tokens: usize) -> bool {
122 enabled && prefill && (PRIME_MIN_T..=crate::cache::PRIME_CHUNK_MAX_TOKENS).contains(&tokens)
123}
124
125fn parse_step_tp_prefill(value: Option<&str>) -> Result<bool, String> {
126 match value {
127 None | Some("") | Some("0") => Ok(false),
128 Some("1") => Ok(true),
129 Some(value) => Err(format!(
130 "MEMRA_STEP_TP_PREFILL={value:?} is invalid; expected 0 or 1"
131 )),
132 }
133}
134
135fn step_tp_prefill_enabled() -> Result<bool, String> {
136 parse_step_tp_prefill(std::env::var("MEMRA_STEP_TP_PREFILL").ok().as_deref())
137}
138
139fn validate_step_prime_batch_modes(tp_prefill: bool, grouped_prefill: bool) -> Result<(), String> {
140 if grouped_prefill && !tp_prefill {
141 return Err("MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into());
142 }
143 if tp_prefill {
144 return Err(
145 "Step TP4 cross-request prime batching did not clear the live-server performance \
146 gate; use per-session grouped prefill"
147 .into(),
148 );
149 }
150 Ok(())
151}
152
153fn step_tp_prefill_shape(
154 enabled: bool,
155 tokens: usize,
156 ranks: usize,
157 native_p2p: bool,
158 has_rank_local_attention: bool,
159 fp8_kv: bool,
160) -> bool {
161 enabled
169 && tokens >= PRIME_MIN_T
170 && matches!(ranks, 2 | 4)
171 && native_p2p
172 && has_rank_local_attention
173 && !fp8_kv
174}
175
176fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
177 std::iter::repeat_with(|| None).take(n).collect()
178}
179
180fn prime_cache_stage_for_layer(fence: &[usize], layer: usize) -> usize {
181 debug_assert!(fence.len() >= 3);
182 match fence[1..fence.len() - 1].binary_search(&layer) {
183 Ok(index) => index + 1,
184 Err(index) => index,
185 }
186}
187
188fn move_prime_cache_layers<T>(
189 parent: &mut [Option<T>],
190 stages: &mut [Vec<Option<T>>],
191 fence: &[usize],
192) {
193 assert_eq!(stages.len() + 1, fence.len());
194 assert!(stages.iter().all(|stage| stage.len() == parent.len()));
195 for (layer, value) in parent.iter_mut().enumerate() {
196 let stage = prime_cache_stage_for_layer(fence, layer);
197 debug_assert!(stages[stage][layer].is_none());
198 stages[stage][layer] = value.take();
199 }
200}
201
202#[cfg(test)]
203fn restore_prime_cache_layers<T>(
204 parent: &mut [Option<T>],
205 stages: &mut [Vec<Option<T>>],
206 fence: &[usize],
207) {
208 assert_eq!(stages.len() + 1, fence.len());
209 assert!(stages.iter().all(|stage| stage.len() == parent.len()));
210 for (layer, value) in parent.iter_mut().enumerate() {
211 let stage = prime_cache_stage_for_layer(fence, layer);
212 debug_assert!(value.is_none());
213 *value = stages[stage][layer].take();
214 }
215}
216
217struct PrimeCacheStages<'a> {
222 parent: &'a mut Cache,
223 fence: Vec<usize>,
224 stages: Vec<std::sync::Mutex<Cache>>,
225 committed: bool,
226}
227
228impl<'a> PrimeCacheStages<'a> {
229 fn new(parent: &'a mut Cache, fence: &[usize]) -> Self {
230 let n = parent.kv.len();
231 assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
232 assert_eq!(parent.tp_kv.len(), n, "cache layer vectors disagree");
233 assert_eq!(parent.latent.len(), n, "cache layer vectors disagree");
234 let n_stages = fence.len().checked_sub(1).expect("PP cache fence is empty");
235 assert!((2..=4).contains(&n_stages), "PP cache needs 2..=4 stages");
236 assert_eq!(fence[0], 0, "PP cache fence must start at layer zero");
237 assert!(
238 fence.windows(2).all(|pair| pair[0] < pair[1]),
239 "PP cache fence must be strictly increasing"
240 );
241 assert!(fence[n_stages] <= n, "PP cache fence exceeds {n} layers");
242
243 let mut latent: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
244 let mut g5_recur: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
245 let mut g5_latent: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
246 let mut kv: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
247 let mut tp_kv: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
248 let mut recur: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
249 move_prime_cache_layers(&mut parent.kv, &mut kv, fence);
250 move_prime_cache_layers(&mut parent.tp_kv, &mut tp_kv, fence);
251 move_prime_cache_layers(&mut parent.recur, &mut recur, fence);
252
253 move_prime_cache_layers(&mut parent.latent, &mut latent, fence);
254 move_prime_cache_layers(&mut parent.glm5_tp_recur, &mut g5_recur, fence);
255 move_prime_cache_layers(&mut parent.glm5_tp_latent_peer, &mut g5_latent, fence);
256 let pos = parent.pos;
257 let max_ctx = parent.max_ctx;
258 let stages = (0..n_stages)
262 .map(|stage| {
263 std::sync::Mutex::new(Cache {
264 kv: std::mem::take(&mut kv[stage]),
265 tp_kv: std::mem::take(&mut tp_kv[stage]),
266 recur: std::mem::take(&mut recur[stage]),
267 latent: std::mem::take(&mut latent[stage]),
268 glm5_tp_recur: std::mem::take(&mut g5_recur[stage]),
269 glm5_tp_latent_peer: std::mem::take(&mut g5_latent[stage]),
270 pos,
271 max_ctx,
272 tainted: false,
273 last_logits_dev: None,
274 dflash_taps: None,
275 hc_taps: None,
276 })
277 })
278 .collect();
279 Self {
280 parent,
281 fence: fence.to_vec(),
282 stages,
283 committed: false,
284 }
285 }
286
287 fn pp2_parts(&mut self) -> (&mut Cache, &mut Cache) {
288 assert_eq!(self.stages.len(), 2);
289 let (stage0, stage1) = self.stages.split_at_mut(1);
290 (
291 stage0[0]
292 .get_mut()
293 .unwrap_or_else(|poisoned| poisoned.into_inner()),
294 stage1[0]
295 .get_mut()
296 .unwrap_or_else(|poisoned| poisoned.into_inner()),
297 )
298 }
299
300 fn stages(&self) -> &[std::sync::Mutex<Cache>] {
301 &self.stages
302 }
303
304 fn commit(&mut self) {
305 self.committed = true;
306 }
307}
308
309impl Drop for PrimeCacheStages<'_> {
310 fn drop(&mut self) {
311 let n = self.parent.kv.len();
312 for i in 0..n {
313 let stage = prime_cache_stage_for_layer(&self.fence, i);
314 let source = self.stages[stage]
315 .get_mut()
316 .unwrap_or_else(|poisoned| poisoned.into_inner());
317 debug_assert!(self.parent.kv[i].is_none());
318 debug_assert!(self.parent.tp_kv[i].is_none());
319 debug_assert!(self.parent.recur[i].is_none());
320 debug_assert!(self.parent.latent[i].is_none());
321 self.parent.kv[i] = source.kv[i].take();
322 self.parent.tp_kv[i] = source.tp_kv[i].take();
323 self.parent.recur[i] = source.recur[i].take();
324 self.parent.latent[i] = source.latent[i].take();
325 self.parent.glm5_tp_recur[i] = source.glm5_tp_recur[i].take();
326 self.parent.glm5_tp_latent_peer[i] = source.glm5_tp_latent_peer[i].take();
327 }
328 self.parent.pos = self
329 .stages
330 .iter_mut()
331 .map(|stage| {
332 stage
333 .get_mut()
334 .unwrap_or_else(|poisoned| poisoned.into_inner())
335 .pos
336 })
337 .min()
338 .unwrap_or(self.parent.pos);
339 if !self.committed {
340 self.parent.mark_tainted();
341 }
342 }
343}
344
345struct CacheTaintGuard {
349 caches: Vec<*mut Cache>,
350 committed: bool,
351}
352
353impl CacheTaintGuard {
354 fn arm(caches: &mut [&mut Cache]) -> Self {
355 Self {
356 caches: caches
357 .iter_mut()
358 .map(|cache| *cache as *mut Cache)
359 .collect(),
360 committed: false,
361 }
362 }
363
364 fn commit(&mut self) {
365 self.committed = true;
366 }
367}
368
369impl Drop for CacheTaintGuard {
370 fn drop(&mut self) {
371 if self.committed {
372 return;
373 }
374 for cache in &self.caches {
375 unsafe { (&mut **cache).mark_tainted() };
378 }
379 }
380}
381
382#[derive(Debug, Clone, Copy, PartialEq, Eq)]
383struct PrimePpWaveSlot {
384 wave: usize,
385 slot: usize,
386}
387
388#[derive(Debug)]
389enum PrimePpSignal {
390 Slot(PrimePpWaveSlot),
391 Error(String),
392}
393
394#[derive(Default)]
395struct PrimePpWaveCredits {
396 next_wave: usize,
397 pending: std::collections::VecDeque<PrimePpWaveSlot>,
398}
399
400impl PrimePpWaveCredits {
401 fn release_required(&self) -> Option<PrimePpWaveSlot> {
402 (self.pending.len() == 2).then(|| self.pending[0])
403 }
404
405 fn record_release(&mut self, released: PrimePpWaveSlot) -> Result<(), String> {
406 let expected =
407 self.pending.front().copied().ok_or_else(|| {
408 "prime PP received a slot release with no pending wave".to_string()
409 })?;
410 if released != expected {
411 return Err(format!(
412 "prime PP slot release {:?} does not match oldest pending {:?}",
413 released, expected
414 ));
415 }
416 self.pending.pop_front();
417 Ok(())
418 }
419
420 fn record_send(&mut self, sent: PrimePpWaveSlot) -> Result<(), String> {
421 if sent.wave != self.next_wave {
422 return Err(format!(
423 "prime PP sent wave {} while wave {} was next",
424 sent.wave, self.next_wave
425 ));
426 }
427 if sent.slot >= 2 {
428 return Err(format!(
429 "prime PP boundary returned invalid slot {}",
430 sent.slot
431 ));
432 }
433 if self.pending.iter().any(|pending| pending.slot == sent.slot) {
434 return Err(format!(
435 "prime PP reused slot {} before its exact-wave release",
436 sent.slot
437 ));
438 }
439 self.pending.push_back(sent);
440 self.next_wave += 1;
441 Ok(())
442 }
443}
444
445fn recv_prime_pp_signal(
446 receiver: &std::sync::mpsc::Receiver<PrimePpSignal>,
447 expected: PrimePpWaveSlot,
448 exact_slot: bool,
449 label: &str,
450) -> Result<PrimePpWaveSlot, String> {
451 match receiver.recv() {
452 Ok(PrimePpSignal::Error(error)) => Err(error),
453 Ok(PrimePpSignal::Slot(received))
454 if received.wave == expected.wave
455 && (!exact_slot || received.slot == expected.slot) =>
456 {
457 if received.slot >= 2 {
458 Err(format!(
459 "{label}: wave {} carried invalid slot {}",
460 received.wave, received.slot
461 ))
462 } else {
463 Ok(received)
464 }
465 }
466 Ok(PrimePpSignal::Slot(received)) => Err(format!(
467 "{label}: expected wave/slot {:?}, received {:?}",
468 expected, received
469 )),
470 Err(_) => Err(format!(
471 "{label}: channel closed while waiting for wave {}",
472 expected.wave
473 )),
474 }
475}
476
477fn send_prime_pp_signal(
478 sender: &std::sync::mpsc::Sender<PrimePpSignal>,
479 signal: PrimePpSignal,
480 label: &str,
481) -> Result<(), String> {
482 sender
483 .send(signal)
484 .map_err(|_| format!("{label}: channel closed"))
485}
486
487struct PrimePpWave<'a> {
488 start: usize,
489 end: usize,
490 tokens: &'a [u32],
491}
492
493struct PrimePpStageChannels {
494 incoming: Option<std::sync::mpsc::Receiver<PrimePpSignal>>,
495 release_upstream: Option<std::sync::mpsc::Sender<PrimePpSignal>>,
496 outgoing: std::sync::mpsc::Sender<PrimePpSignal>,
497 released_downstream: std::sync::mpsc::Receiver<PrimePpSignal>,
498}
499
500impl PrimePpStageChannels {
501 fn notify_failure(&self, error: &str) {
502 if let Some(upstream) = &self.release_upstream {
503 let _ = upstream.send(PrimePpSignal::Error(error.to_string()));
504 }
505 let _ = self.outgoing.send(PrimePpSignal::Error(error.to_string()));
506 }
507}
508
509pub struct IndexerPlanes<'a> {
523 pub state: &'a mut CudaSlice<f32>,
524 pub pool_keys: &'a mut Option<CudaSlice<f32>>,
525 pub ready: &'a mut usize,
526 pub state_ring_rows: usize,
531 pub capacity_tokens: usize,
535}
536
537pub(crate) struct AttnPre {
539 pub q: cudarc::driver::CudaSlice<f32>,
540 pub k: cudarc::driver::CudaSlice<f32>,
541 pub v: cudarc::driver::CudaSlice<f32>,
542 pub gate: Option<cudarc::driver::CudaSlice<f32>>,
543}
544
545pub(crate) struct GdnPrep {
547 pub hk: usize,
548 pub q_l2: cudarc::driver::CudaSlice<f32>,
549 pub k_l2: cudarc::driver::CudaSlice<f32>,
550 pub v_g: cudarc::driver::CudaSlice<f32>,
551 pub beta: cudarc::driver::CudaSlice<f32>,
552 pub g_log: cudarc::driver::CudaSlice<f32>,
553 pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
554 pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
555}
556
557pub(crate) struct VerifyStreamScratch {
559 pub pos_d: CudaSlice<i32>,
560 pub row_ctrs: Vec<CudaSlice<i32>>,
561}
562use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
563
564struct MoeInputTraceWriter {
565 dir: std::path::PathBuf,
566 index: std::fs::File,
567 payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
568}
569
570static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
571 std::sync::OnceLock::new();
572
573fn gdec_enabled() -> bool {
576 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
577 *E.get_or_init(|| {
578 std::env::var("MEMRA_MOE_GDEC")
579 .map(|v| v != "0")
580 .unwrap_or(true)
581 })
582}
583
584fn moe_slab_enabled() -> bool {
595 std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
596}
597
598fn moe_fused_epi_enabled() -> bool {
611 std::env::var("MEMRA_MOE_FUSED_EPI")
612 .map(|v| v != "0")
613 .unwrap_or(false)
614}
615
616fn hyper_decode_ws_on() -> bool {
630 std::env::var("MEMRA_HC_DECODE_WS").as_deref() == Ok("1")
631}
632
633pub static HC_DECODE_WS_DISPATCHES: std::sync::atomic::AtomicU64 =
636 std::sync::atomic::AtomicU64::new(0);
637
638fn mla_tc_prefill_enabled() -> bool {
661 std::env::var("MEMRA_MLA_TC_PREFILL")
662 .map(|v| v != "0")
663 .unwrap_or(true)
664}
665
666fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
670 std::env::var("MEMRA_MOE_GROUPED")
671 .map(|value| value != "0")
672 .unwrap_or(false)
673}
674
675fn moe_grouped_prefill_enabled() -> bool {
696 std::env::var("MEMRA_MOE_GROUPED_PREFILL")
697 .map(|v| v != "0")
698 .unwrap_or(true)
699}
700
701fn moe_prefetch_enabled() -> bool {
704 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
705 *E.get_or_init(|| {
706 std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
707 || crate::spill_pread::worker_enabled()
708 })
709}
710
711fn moe_page_prefetch_window() -> usize {
716 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
717 *W.get_or_init(|| {
718 page_prefetch_window_from_values(
719 std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
720 std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
721 .ok()
722 .as_deref(),
723 )
724 })
725}
726
727fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
728 if !enabled {
729 return 0;
730 }
731 raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
732}
733
734fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
738 if window == 0 || position >= len {
739 return len..len;
740 }
741 let (start, count) = if position == 0 {
742 (1, window)
743 } else {
744 (position.saturating_add(window), 1)
745 };
746 let start = start.min(len);
747 start..start.saturating_add(count).min(len)
748}
749
750fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
753 let position = current.map_or(0, |position| position.saturating_add(1));
754 (position < order_len).then_some(position)
755}
756
757fn worker_prefetch_window() -> usize {
762 static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
763 *WINDOW.get_or_init(|| {
764 let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
765 std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
766 .ok()
767 .and_then(|value| value.parse::<usize>().ok())
768 .unwrap_or(automatic.max(1))
769 })
770}
771
772fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
776 if window == 0 || position >= len {
777 return len..len;
778 }
779 let (start, count) = if position == 0 {
780 (0, window)
781 } else {
782 (position.saturating_add(window).saturating_sub(1), 1)
783 };
784 let start = start.min(len);
785 start..start.saturating_add(count).min(len)
786}
787
788fn moe_dev_enabled() -> bool {
793 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
794 *E.get_or_init(|| {
795 std::env::var("MEMRA_MOE_DEV")
796 .map(|v| v != "0")
797 .unwrap_or(true)
798 && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
799 })
800}
801
802enum VrowsSel<'a> {
811 Host(&'a [u32], &'a [f32]),
812 Dev(&'a CudaSlice<i32>, &'a CudaSlice<f32>),
813}
814
815fn sigmoid_router_enabled() -> bool {
816 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
817 *E.get_or_init(|| {
818 std::env::var("MEMRA_SIG_ROUTER")
819 .map(|v| v != "0")
820 .unwrap_or(true)
821 })
822}
823
824fn moe_q8_enabled() -> bool {
829 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
830 *E.get_or_init(|| {
831 std::env::var("MEMRA_MOE_Q8")
832 .map(|v| v != "0")
833 .unwrap_or(true)
834 })
835}
836
837fn expert_dp4a_supported(qt: i32) -> bool {
840 qt == crate::QT_Q4_0
841 || qt == crate::QT_IQ3_S
842 || qt == crate::QT_IQ4_XS
843 || qt == crate::QT_Q3_K
844 || qt == crate::QT_Q4_K
845 || qt == crate::QT_Q6_K
846}
847
848fn q8_expert_supported(qt: i32) -> bool {
849 static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
855 let kq = *KQ.get_or_init(|| {
856 std::env::var("MEMRA_MOE_Q8_KQ")
857 .map(|v| v != "0")
858 .unwrap_or(true)
859 });
860 let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
867 .map(|v| v != "0")
868 .unwrap_or(true);
869 qt == crate::QT_IQ3_S
870 || qt == crate::QT_IQ4_XS
871 || (nvfp4_q8 && qt == crate::QT_NVFP4)
872 || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
873}
874
875fn q8_expert_supported_for_model(cfg: &ModelConfig, qt: i32) -> bool {
879 let weight_only_nvfp4 = cfg.hy3.as_ref().is_some_and(|hy3| hy3.weight_only_nvfp4);
880 q8_expert_supported(qt) && !(weight_only_nvfp4 && qt == crate::QT_NVFP4)
881}
882
883fn moe_q8_enabled_for_model(cfg: &ModelConfig, m: &MoeWeights) -> bool {
884 m.has_uniform_expert_layout()
885 && moe_q8_enabled()
886 && q8_expert_supported_for_model(cfg, m.gate_exps.qtype)
887 && q8_expert_supported_for_model(cfg, m.up_exps.qtype)
888 && q8_expert_supported_for_model(cfg, m.down_exps.qtype)
889}
890
891#[cfg(test)]
892mod w4a16_dispatch_tests {
893 use super::q8_expert_supported_for_model;
894 use memra_gguf::config::{HfConfig, ModelConfig};
895
896 #[test]
897 fn hy3_w4a16_never_admits_q8_activations() {
898 let hf = HfConfig::parse(
899 r#"{"model_type":"hy_v3","num_hidden_layers":2,"hidden_size":8,
900 "num_attention_heads":2,"intermediate_size":16,"vocab_size":32,
901 "max_position_embeddings":32,
902 "quantization_config":{"quant_method":"modelopt","quant_algo":"W4A16_NVFP4"}}"#,
903 );
904 let cfg = ModelConfig::from_hf(&hf);
905 assert!(!q8_expert_supported_for_model(&cfg, crate::QT_NVFP4));
906 assert!(q8_expert_supported_for_model(&cfg, crate::QT_IQ4_XS));
907 }
908}
909
910fn q8_expert_dec_supported(qt: i32) -> bool {
913 qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
914}
915
916fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
922 match qt {
923 crate::QT_Q4_0 => in_f.is_multiple_of(32),
924 crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
925 in_f.is_multiple_of(256)
926 }
927 crate::QT_NVFP4 => in_f.is_multiple_of(64),
932 _ => false,
933 }
934}
935
936fn moe_prewarm_enabled() -> bool {
939 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
940 *E.get_or_init(|| {
941 std::env::var("MEMRA_MOE_PREWARM")
942 .map(|v| v != "0")
943 .unwrap_or(true)
944 })
945}
946
947fn cpu_expert_profile_admit_enabled() -> bool {
951 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
952 *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
953}
954
955pub const PRIME_MIN_T: usize = 16;
959
960fn step_gemm_prime_suffix_on() -> bool {
990 std::env::var("MEMRA_STEP_GEMM_PRIME_SUFFIX").as_deref() != Ok("0")
991}
992
993const MOE_DEV_MAX_T: usize = 16;
1002const PRIME_PIPE_MICROBATCHES: usize = 8;
1003const PRIME_PIPE_MIN_CHUNK: usize = 128;
1004const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
1005const PRIME_PIPE_LINEAR_WORK: usize = 8;
1006
1007fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
1008 crate::pp::prime_pp_on()
1009 && !crate::pp::pp2_streams_off()
1010 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
1011}
1012
1013fn prime_ppn_wave_auto_geometry(n_layers: usize) -> bool {
1014 crate::pp::prime_pp_on()
1015 && !crate::pp::pp2_streams_off()
1016 && crate::pp::pp_wave_on() == Ok(true)
1017 && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| matches!(cuts.len(), 4 | 5))
1018}
1019
1020fn prime_pipeline_auto_geometry(n_layers: usize) -> bool {
1021 prime_pp2_auto_geometry(n_layers) || prime_ppn_wave_auto_geometry(n_layers)
1022}
1023
1024pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
1029 if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
1030 let parsed = value
1031 .parse::<usize>()
1032 .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
1033 return if crate::cache::swa_ring_on() {
1034 if parsed == 0 {
1035 crate::cache::PRIME_CHUNK_MAX_TOKENS
1036 } else {
1037 parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
1038 }
1039 } else {
1040 parsed
1041 };
1042 }
1043 let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
1044 if prime_pipeline_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
1045 chunk.min(
1046 t.div_ceil(PRIME_PIPE_MICROBATCHES)
1047 .max(PRIME_PIPE_MIN_CHUNK),
1048 )
1049 } else {
1050 chunk
1051 }
1052}
1053
1054fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
1055 fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
1056}
1057
1058fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
1059 if chunk == 0 || t <= chunk {
1060 return vec![(0, t)];
1061 }
1062 let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
1063 let mut start = 0usize;
1064 while start < t {
1065 let mut end = (start + chunk).min(t);
1066 if t - end > 0 && t - end < PRIME_MIN_T {
1067 if ring_on {
1068 let shifted = t - PRIME_MIN_T;
1069 end = if shifted > start { shifted } else { t };
1070 } else {
1071 end = t;
1072 }
1073 }
1074 ranges.push((start, end));
1075 start = end;
1076 }
1077 ranges
1078}
1079
1080fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
1081 let prefix = prefix as u128;
1082 prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
1083}
1084
1085fn dynamic_prime_chunk_ranges(
1086 t: usize,
1087 fixed_chunk: usize,
1088 fixed: &[(usize, usize)],
1089) -> Vec<(usize, usize)> {
1090 let n = fixed.len();
1091 if n < 3 {
1092 return fixed.to_vec();
1093 }
1094
1095 let max_first = t - (n - 1) * PRIME_MIN_T;
1096 let first = fixed_chunk
1097 .div_ceil(2)
1098 .max(PRIME_PIPE_EDGE_MIN_CHUNK)
1099 .min(max_first);
1100 let mut ranges = Vec::with_capacity(n);
1101 ranges.push((0, first));
1102
1103 let first_work = prime_chunk_work(first, t);
1104 let work_span = prime_chunk_work(t, t) - first_work;
1105 let denominator = (n - 1) as u128;
1106 let mut previous = first;
1107 for boundary in 1..n - 1 {
1108 let target = first_work * denominator + work_span * (boundary as u128);
1109 let remaining = n - 1 - boundary;
1110 let mut low = previous + PRIME_MIN_T;
1111 let mut high = t - remaining * PRIME_MIN_T;
1112 while low < high {
1113 let mid = low + (high - low) / 2;
1114 if prime_chunk_work(mid, t) * denominator >= target {
1115 high = mid;
1116 } else {
1117 low = mid + 1;
1118 }
1119 }
1120 ranges.push((previous, low));
1121 previous = low;
1122 }
1123 ranges.push((previous, t));
1124 ranges
1125}
1126
1127pub fn prime_chunk_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1137 let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
1138 let chunk = prime_chunk_tokens(t, n_layers);
1139 let fixed = fixed_prime_chunk_ranges(t, chunk);
1140 let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
1141 Ok(value) => value == "dynamic",
1142 Err(_) => true,
1143 };
1144 if explicit_chunk {
1145 return fixed;
1146 }
1147 let ranges = if !dynamic || !prime_pipeline_auto_geometry(n_layers) {
1148 fixed
1149 } else {
1150 dynamic_prime_chunk_ranges(t, chunk, &fixed)
1151 };
1152 if gdn_grid && std::env::var("MEMRA_PRIME_GRID_ALIGN").as_deref() != Ok("0") {
1156 align_prime_ranges_to_gdn(&ranges, t, Engine::gdn_chunk_size())
1157 } else {
1158 ranges
1159 }
1160}
1161
1162pub fn hyper_prime_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1216 prime_chunk_ranges(t, n_layers, gdn_grid)
1217}
1218
1219pub fn hyper_prime_call_rows(t: usize, n_layers: usize, gdn_grid: bool) -> usize {
1224 hyper_prime_ranges(t, n_layers, gdn_grid)
1225 .iter()
1226 .map(|&(start, end)| end - start)
1227 .max()
1228 .unwrap_or(0)
1229}
1230
1231#[derive(Debug, Clone, Copy)]
1243pub struct HyperPrimeWorkspaceShape {
1244 pub chunk_token_bytes: usize,
1252 pub prompt_bytes_per_token: usize,
1256 pub kpool_score_pool: usize,
1261 pub n_layers: usize,
1264 pub gdn_grid: bool,
1266}
1267
1268impl HyperPrimeWorkspaceShape {
1269 pub fn admission_bytes(&self, prompt_rows: usize) -> usize {
1280 let rows = hyper_prime_call_rows(prompt_rows, self.n_layers, self.gdn_grid);
1281 let chunk = self.chunk_token_bytes.saturating_mul(rows);
1282 let score = prompt_rows
1283 .checked_div(self.kpool_score_pool)
1284 .map(|pools| rows.saturating_mul(pools).saturating_mul(size_of::<f32>()))
1285 .unwrap_or(0);
1286 chunk
1287 .saturating_add(score)
1288 .saturating_add(self.prompt_bytes_per_token.saturating_mul(prompt_rows))
1289 }
1290}
1291
1292pub fn align_prime_ranges_to_gdn(
1312 ranges: &[(usize, usize)],
1313 t: usize,
1314 c: usize,
1315) -> Vec<(usize, usize)> {
1316 if c == 0 || ranges.len() < 2 {
1317 return ranges.to_vec();
1318 }
1319 let mut out: Vec<(usize, usize)> = Vec::with_capacity(ranges.len());
1320 let mut start = 0usize;
1321 for (i, &(_, end)) in ranges.iter().enumerate() {
1322 let e = if i + 1 == ranges.len() {
1323 t
1324 } else {
1325 end / c * c
1326 };
1327 if e > start {
1328 out.push((start, e));
1329 start = e;
1330 } }
1332 debug_assert_eq!(out.last().map(|&(_, e)| e), Some(t));
1333 out
1334}
1335
1336struct HeadSplit {
1337 pin: u64,
1338 w1: CudaSlice<u8>,
1339 hn1: CudaSlice<f32>,
1340 y1: CudaSlice<f32>,
1341 logits_e: CudaSlice<f32>,
1342 ev_hn: cudarc::driver::CudaEvent,
1343 ev_done: cudarc::driver::CudaEvent,
1344 raw_hn1: u64,
1345 raw_y1: u64,
1346 raw_logits_hi: u64,
1347 samp: Option<SampScratch>,
1352}
1353
1354struct SampScratch {
1355 pb: CudaSlice<f32>,
1356 th: CudaSlice<f32>,
1357 z: CudaSlice<f32>,
1358 mx: CudaSlice<f32>,
1359 rows: CudaSlice<i32>,
1360}
1361static HEAD_SPLIT_WS: std::sync::Mutex<Option<HeadSplit>> = std::sync::Mutex::new(None);
1363
1364#[allow(clippy::type_complexity)]
1368static DEV1_ROUTER_REPS: std::sync::Mutex<
1369 Option<(
1370 std::collections::HashMap<u16, (CudaSlice<f32>, CudaSlice<f32>, CudaSlice<u8>)>,
1371 Option<CudaSlice<f32>>,
1372 )>,
1373> = std::sync::Mutex::new(None);
1374
1375#[allow(clippy::type_complexity)]
1379static SHEXP_D1_REPS: std::sync::Mutex<
1380 Option<std::collections::HashMap<u16, (CudaSlice<u8>, CudaSlice<u8>, CudaSlice<u8>)>>,
1381> = std::sync::Mutex::new(None);
1382#[allow(clippy::type_complexity)]
1383static SHEXP_D1_WS: std::sync::Mutex<
1384 Option<(
1385 (usize, usize),
1386 CudaSlice<f32>,
1387 CudaSlice<f32>,
1388 CudaSlice<f32>,
1389 cudarc::driver::CudaEvent,
1390 cudarc::driver::CudaEvent,
1391 )>,
1392> = std::sync::Mutex::new(None);
1393
1394#[allow(clippy::type_complexity)] static SHEXP_OV_WS: std::sync::Mutex<
1397 Option<(usize, usize, usize, CudaSlice<f32>, CudaSlice<f32>)>,
1398> = std::sync::Mutex::new(None);
1399
1400impl HybridModel {
1401 pub fn gdn_prime_grid_on(&self) -> bool {
1407 Engine::gdn_chunked_enabled()
1408 && self
1409 .layers
1410 .iter()
1411 .any(|l| matches!(l.mixer, crate::hybrid::Mixer::Linear(_)))
1412 }
1413
1414 fn full_attn_tp_device_resident(e: &Engine, tp: &crate::hybrid::StepTpQkv) -> bool {
1419 tp.runtime.native_p2p() && tp.runtime.root_shares_ctx(e)
1420 }
1421
1422 pub(crate) fn full_attn_tp_qkv(
1423 &self,
1424 e: &Engine,
1425 fa: &FullAttnLayer,
1426 h: &CudaSlice<f32>,
1427 t: usize,
1428 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
1429 let Some(tp) = fa.step_tp_qkv.as_ref() else {
1430 return Ok(None);
1431 };
1432 let values = active_matrix_values(
1433 h.len(),
1434 t,
1435 self.cfg.n_embd as usize,
1436 "Step TP QKV activation",
1437 )?;
1438 if Self::full_attn_tp_device_resident(e, tp) {
1447 e.stream().synchronize()?;
1450 let q = tp
1451 .runtime
1452 .bf16_column_parallel_resident_native_device(&tp.q, h, t)?;
1453 let k = tp
1454 .runtime
1455 .bf16_column_parallel_resident_native_device(&tp.k, h, t)?;
1456 let v = tp
1457 .runtime
1458 .bf16_column_parallel_resident_native_device(&tp.v, h, t)?;
1459 Self::full_attn_tp_log_once(tp, "qkv", "device-resident");
1460 return Ok(Some(vec![q, k, v]));
1461 }
1462 let host = e.dtoh_view(&h.slice(0..values))?;
1463 let q = if tp.runtime.native_p2p() {
1464 tp.runtime
1465 .bf16_column_parallel_resident_native(&tp.q, &host, t)?
1466 } else {
1467 tp.runtime
1468 .bf16_column_parallel_resident(&tp.q, &host, t)?
1469 .gathered
1470 };
1471 let k = if tp.runtime.native_p2p() {
1472 tp.runtime
1473 .bf16_column_parallel_resident_native(&tp.k, &host, t)?
1474 } else {
1475 tp.runtime
1476 .bf16_column_parallel_resident(&tp.k, &host, t)?
1477 .gathered
1478 };
1479 let v = if tp.runtime.native_p2p() {
1480 tp.runtime
1481 .bf16_column_parallel_resident_native(&tp.v, &host, t)?
1482 } else {
1483 tp.runtime
1484 .bf16_column_parallel_resident(&tp.v, &host, t)?
1485 .gathered
1486 };
1487 Self::full_attn_tp_log_once(tp, "qkv", "host-canonical");
1488 Ok(Some(vec![e.htod(&q)?, e.htod(&k)?, e.htod(&v)?]))
1489 }
1490
1491 fn full_attn_tp_log_once(tp: &crate::hybrid::StepTpQkv, proj: &str, activation: &'static str) {
1495 use std::sync::atomic::{AtomicBool, Ordering};
1496 static LOGGED: [AtomicBool; 4] = [
1497 AtomicBool::new(false),
1498 AtomicBool::new(false),
1499 AtomicBool::new(false),
1500 AtomicBool::new(false),
1501 ];
1502 let idx = 2 * usize::from(proj == "o") + usize::from(activation == "device-resident");
1503 if LOGGED[idx].swap(true, Ordering::Relaxed) {
1504 return;
1505 }
1506 eprintln!(
1507 "[step-tp-{proj}] execute layer={} devices={:?} projections={proj} \
1508 tensor_parallel=true attention_local=true kv_local=true transport={} \
1509 native_p2p={} bulk_p2p={} activation={activation} \
1510 output={} performance_claim=false (logged once per transport)",
1511 tp.layer,
1512 tp.devices,
1513 tp.runtime.transport_label(),
1514 tp.runtime.native_p2p(),
1515 tp.runtime.bulk_p2p(),
1516 if activation == "device-resident" {
1517 "root-resident"
1518 } else {
1519 "root-readback"
1520 },
1521 );
1522 }
1523
1524 pub(crate) fn full_attn_tp_o(
1525 &self,
1526 e: &Engine,
1527 fa: &FullAttnLayer,
1528 activation: &CudaSlice<f32>,
1529 tokens: usize,
1530 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1531 let Some(tp) = fa.step_tp_qkv.as_ref() else {
1532 return Ok(None);
1533 };
1534 if Self::full_attn_tp_device_resident(e, tp) {
1538 e.stream().synchronize()?; let output = tp
1540 .runtime
1541 .step_bf16_row_parallel_resident_native_device(&tp.o, activation, tokens)?;
1542 Self::full_attn_tp_log_once(tp, "o", "device-resident");
1543 return Ok(Some(output));
1544 }
1545 let host = e.dtoh(activation)?;
1546 let output = if tp.runtime.native_p2p() {
1547 tp.runtime
1548 .step_bf16_row_parallel_resident_native(&tp.o, &host, tokens)?
1549 } else {
1550 tp.runtime
1551 .step_bf16_row_parallel_resident(&tp.o, &host, tokens)?
1552 };
1553 Self::full_attn_tp_log_once(tp, "o", "host-canonical");
1554 Ok(Some(e.htod(&output)?))
1555 }
1556
1557 fn full_attn_o(
1558 &self,
1559 e: &Engine,
1560 fa: &FullAttnLayer,
1561 activation: &CudaSlice<f32>,
1562 tokens: usize,
1563 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1564 match self.full_attn_tp_o(e, fa, activation, tokens)? {
1565 Some(output) => Ok(output),
1566 None => e.matmul(&fa.wo, activation, tokens),
1567 }
1568 }
1569
1570 fn prime_trace_path() -> Option<&'static str> {
1575 static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
1576 P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
1577 .as_deref()
1578 }
1579
1580 fn prime_anatomy_on() -> bool {
1586 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1587 *E.get_or_init(|| std::env::var("MEMRA_PRIME_ANATOMY").as_deref() == Ok("1"))
1588 }
1589
1590 fn prime_anatomy_slots() -> &'static [std::sync::atomic::AtomicU64; 5] {
1591 static S: [std::sync::atomic::AtomicU64; 5] = [
1592 std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), std::sync::atomic::AtomicU64::new(0), ];
1598 &S
1599 }
1600
1601 pub(crate) fn refuse_hyper(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
1608 if let Some(topology) = self.hyper.as_ref() {
1609 return Err(format!(
1610 "{path} runs a serial residual, but this model's ModelPlan declares \
1611 ResidualTopology::HyperConnections{{ streams: {}, collapse: {:?} }}. Refusing: \
1612 that path would compute a different model. Converted paths: forward, \
1613 forward_last, prime_cache, decode_step, and the batched serving chain \
1614 decode_step_batch / _sampled / _lean / _masked.",
1615 topology.streams, topology.collapse
1616 )
1617 .into());
1618 }
1619 Ok(())
1620 }
1621
1622 fn hyper_ffn_branch(
1628 &self,
1629 e: &Engine,
1630 layer: &crate::hybrid::HybridLayer,
1631 z: &CudaSlice<f32>,
1632 t: usize,
1633 il: usize,
1634 prefill: bool,
1635 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1636 match &layer.ffn {
1637 crate::hybrid::Ffn::Dense {
1638 ffn_gate,
1639 ffn_up,
1640 ffn_down,
1641 } => {
1642 let n_ff = ffn_gate.out_features();
1643 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], z, t)?;
1644 let up = g2.pop().unwrap();
1645 let gate = g2.pop().unwrap();
1646 let mut act = e.uninit(t * n_ff)?;
1647 Self::ffn_act_lim(
1649 e,
1650 &self.cfg,
1651 &gate,
1652 &up,
1653 1.0,
1654 1.0,
1655 self.cfg.clamp_shexp_at(il as u32),
1656 &mut act,
1657 t * n_ff,
1658 )?;
1659 e.matmul(ffn_down, &act, t)
1660 }
1661 crate::hybrid::Ffn::Moe(m) => {
1662 if prefill {
1663 self.moe_ffn_il_prefill(e, m, z, t, il as u16)
1664 } else {
1665 self.moe_ffn_il_zq8(e, m, z, None, t, il as u16)
1666 }
1667 }
1668 }
1669 }
1670
1671 fn forward_hyper(
1676 &self,
1677 e: &Engine,
1678 tokens: &[u32],
1679 last_only: bool,
1680 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
1681 let topology = *self
1682 .hyper
1683 .as_ref()
1684 .ok_or("forward_hyper on a model with no HyperConnections topology")?;
1685 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1691 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1692 return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1693 }
1694 return self.forward_hyper_ppn(e, tokens, last_only, &topology, &fence);
1695 }
1696 let n_embd = self.cfg.n_embd as usize;
1697 let t = tokens.len();
1698 let eps = self.cfg.rms_eps;
1699 let pos: Vec<i32> = (0..t as i32).collect();
1700 let pos_d = e.htod_i32(&pos)?;
1701
1702 let embedded = self.embed(e, tokens)?;
1703 let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
1704 let trace = memra_reference::hidden_trace::enabled();
1705 if trace {
1706 memra_reference::hidden_trace::emit_tokens(tokens);
1707 let streams = x.len() / (t * n_embd);
1708 memra_reference::hidden_trace::emit_last_row(
1709 "expand",
1710 -1,
1711 t,
1712 streams * n_embd,
1713 &e.dtoh(&x)?,
1714 );
1715 }
1716
1717 x = self.hyper_range_forward(e, &topology, x, 0, self.layers.len(), &pos_d, t, trace)?;
1718
1719 self.hyper_head_logits(e, &topology, &x, t, n_embd, eps, last_only)
1723 }
1724
1725 #[allow(clippy::type_complexity)] fn prime_cache_hyper(
1737 &self,
1738 e: &Engine,
1739 tokens: &[u32],
1740 cache: &mut Cache,
1741 queued_after: usize,
1742 overlay: Option<&crate::vision::EmbedOverlay>,
1743 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1744 let topology = *self
1745 .hyper
1746 .as_ref()
1747 .ok_or("prime_cache_hyper on a model with no HyperConnections topology")?;
1748 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1770 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1771 return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1772 }
1773 let n_embd = self.cfg.n_embd as usize;
1782 let t = tokens.len();
1783 if cache.pos + t > cache.max_ctx {
1784 return Err("prime_cache: prompt exceeds cache max_ctx".into());
1785 }
1786 let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
1787 if ranges.len() == 1 {
1788 return self.prime_cache_hyper_ppn(
1789 e,
1790 tokens,
1791 cache,
1792 queued_after,
1793 &topology,
1794 &fence,
1795 overlay,
1796 );
1797 }
1798 let mut hiddens = e.uninit(t * n_embd)?;
1799 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1800 for &(start, end) in &ranges {
1801 let ov = overlay.and_then(|o| o.window(start, end - start));
1802 let (l, hs, x) = self.prime_cache_hyper_ppn(
1803 e,
1804 &tokens[start..end],
1805 cache,
1806 queued_after + (t - end),
1807 &topology,
1808 &fence,
1809 ov.as_ref(),
1810 )?;
1811 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
1812 last = Some((l, hs));
1813 }
1814 let (logits, h_seed) =
1815 last.expect("hyper_prime_ranges never returns an empty schedule");
1816 return Ok((logits, h_seed, hiddens));
1817 }
1818 let n_embd = self.cfg.n_embd as usize;
1819 let t = tokens.len();
1820 if cache.pos + t > cache.max_ctx {
1821 return Err("prime_cache: prompt exceeds cache max_ctx".into());
1822 }
1823 let seq_end = cache.pos + t + queued_after;
1827 let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
1828 if ranges.len() == 1 {
1829 return self.prime_chunk_hyper(e, tokens, cache, seq_end, 0, overlay);
1830 }
1831 let mut hiddens = e.uninit(t * n_embd)?;
1832 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1833 for &(start, end) in &ranges {
1834 let (l, hs, x) =
1835 self.prime_chunk_hyper(e, &tokens[start..end], cache, seq_end, start, overlay)?;
1836 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
1837 last = Some((l, hs));
1838 }
1839 let (logits, h_seed) = last.expect("hyper_prime_ranges never returns an empty schedule");
1840 Ok((logits, h_seed, hiddens))
1841 }
1842
1843 pub fn hyper_prime_workspace_shape(&self) -> Option<HyperPrimeWorkspaceShape> {
1868 let topology = self.hyper.as_ref()?;
1869 let h = self.cfg.n_embd as usize;
1870 let s = topology.streams;
1871 let f32b = std::mem::size_of::<f32>();
1872 let mut chunk_token_bytes = 4 * s * h * f32b + 4 * h * f32b + 2 * h * f32b;
1874 if let Some(moe) = self.cfg.moe.as_ref() {
1875 let u = moe.expert_used_count as usize;
1876 let f = moe.expert_ff_length as usize;
1877 chunk_token_bytes += u * (10 * h + 14 * f) + 4 * h;
1878 }
1879 let mut kpool_score_pool = 0;
1880 if let Some(glm5) = self.cfg.glm5.as_ref() {
1881 let heads = self.cfg.n_head as usize;
1882 chunk_token_bytes +=
1883 heads * (glm5.qk_head_dim as usize + glm5.v_head_dim as usize) * f32b;
1884 if glm5.index_kpool > 0 {
1885 chunk_token_bytes += (glm5.index_topk as usize / glm5.index_kpool as usize + 1)
1886 * std::mem::size_of::<i32>();
1887 kpool_score_pool = glm5.index_kpool as usize;
1888 }
1889 }
1890 Some(HyperPrimeWorkspaceShape {
1891 chunk_token_bytes,
1892 prompt_bytes_per_token: h * f32b,
1893 kpool_score_pool,
1894 n_layers: self.layers.len(),
1895 gdn_grid: self.gdn_prime_grid_on(),
1896 })
1897 }
1898
1899 pub(crate) fn decode_step_hyper(
1904 &self,
1905 e: &Engine,
1906 token: u32,
1907 cache: &mut Cache,
1908 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1909 let topology = *self
1910 .hyper
1911 .as_ref()
1912 .ok_or("decode_step_hyper on a model with no HyperConnections topology")?;
1913 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1917 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1918 return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1919 }
1920 return self.decode_step_hyper_ppn(e, token, cache, &topology, &fence);
1921 }
1922 let n_embd = self.cfg.n_embd as usize;
1923 let eps = self.cfg.rms_eps;
1924 let pos = cache.pos;
1925 let pos_d = e.htod_i32(&[pos as i32])?;
1926
1927 let embedded = e.htod(&self.embd.gather(n_embd, &[token]))?;
1928 let mut x = crate::hyper::expand(e, &topology, &embedded, 1, n_embd)?;
1929
1930 x = self.hyper_range_decode(e, &topology, x, 0, self.layers.len(), &pos_d, pos, cache)?;
1931
1932 self.hyper_decode_tail(e, &topology, &x, n_embd, eps, cache)
1934 }
1935
1936 #[allow(clippy::too_many_arguments)]
1943 fn hyper_range_forward(
1944 &self,
1945 e: &Engine,
1946 topology: &crate::hyper::HyperTopology,
1947 mut x: CudaSlice<f32>,
1948 lo: usize,
1949 hi: usize,
1950 pos_d: &CudaSlice<i32>,
1951 t: usize,
1952 trace: bool,
1953 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1954 let n_embd = self.cfg.n_embd as usize;
1955 let eps = self.cfg.rms_eps;
1956 for il in lo..hi {
1957 let layer = &self.layers[il];
1958 let hyper = layer.hyper.as_ref().ok_or_else(|| {
1959 format!("layer {il} carries no hyper-connection weights under an hc plan")
1960 })?;
1961
1962 let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
1963 let mut h = e.uninit(t * n_embd)?;
1964 e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1965 let mixed = match &layer.mixer {
1966 Mixer::Full(fa) => self.full_attn(e, fa, &h, pos_d, t, il)?,
1967 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
1968 Mixer::Mla(mla) => self.mla_attn(e, mla, &h, pos_d, t, il)?,
1969 Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
1970 };
1971 x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
1972 if trace {
1973 let index = il as i64;
1974 memra_reference::hidden_trace::emit_last_row(
1975 "mixer",
1976 index,
1977 t,
1978 n_embd,
1979 &e.dtoh(&mixed)?,
1980 );
1981 let streams = x.len() / (t * n_embd);
1982 memra_reference::hidden_trace::emit_last_row(
1983 "attn",
1984 index,
1985 t,
1986 streams * n_embd,
1987 &e.dtoh(&x)?,
1988 );
1989 }
1990
1991 let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
1992 let mut z = e.uninit(t * n_embd)?;
1993 e.rms_norm(
1994 &y,
1995 layer.post_attn_norm.float_data(),
1996 &mut z,
1997 n_embd,
1998 t,
1999 eps,
2000 )?;
2001 let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2002 x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2003 if trace {
2004 let index = il as i64;
2005 memra_reference::hidden_trace::emit_last_row(
2006 "ffn",
2007 index,
2008 t,
2009 n_embd,
2010 &e.dtoh(&ffn_out)?,
2011 );
2012 let streams = x.len() / (t * n_embd);
2013 memra_reference::hidden_trace::emit_last_row(
2014 "layer",
2015 index,
2016 t,
2017 streams * n_embd,
2018 &e.dtoh(&x)?,
2019 );
2020 }
2021 }
2022 Ok(x)
2023 }
2024
2025 #[allow(clippy::too_many_arguments)]
2030 fn hyper_range_prime(
2031 &self,
2032 e: &Engine,
2033 topology: &crate::hyper::HyperTopology,
2034 mut x: CudaSlice<f32>,
2035 lo: usize,
2036 hi: usize,
2037 pos_d: &CudaSlice<i32>,
2038 t: usize,
2039 cache: &mut Cache,
2040 seq_end: usize,
2041 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2042 let n_embd = self.cfg.n_embd as usize;
2043 let eps = self.cfg.rms_eps;
2044 for il in lo..hi {
2045 let layer = &self.layers[il];
2046 let hyper = layer.hyper.as_ref().ok_or_else(|| {
2047 format!("layer {il} carries no hyper-connection weights under an hc plan")
2048 })?;
2049
2050 let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
2051 let mut h = e.uninit(t * n_embd)?;
2052 e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2053 let mixed = match &layer.mixer {
2054 Mixer::Full(fa) => {
2055 self.full_attn_prime(e, fa, &h, None, pos_d, t, cache, il, seq_end)?
2056 }
2057 Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2058 Mixer::Mla(mla) if mla.tp.is_some() => {
2059 self.mla_tp_attn_cached(e, mla, &h, pos_d, t, il, cache, false)?
2060 }
2061 Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, t, il, cache)?,
2062 Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2063 e,
2064 la,
2065 &h,
2066 t,
2067 eps,
2068 cache,
2069 il,
2070 crate::kda::ConvArm::Prefill,
2071 )?,
2072 Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2073 };
2074 x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
2075
2076 let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
2077 let mut z = e.uninit(t * n_embd)?;
2078 e.rms_norm(
2079 &y,
2080 layer.post_attn_norm.float_data(),
2081 &mut z,
2082 n_embd,
2083 t,
2084 eps,
2085 )?;
2086 let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2087 x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2088 self.glm5_hc_tap(e, cache, topology, il, &x, t)?;
2091 }
2092 Ok(x)
2093 }
2094
2095 #[allow(clippy::too_many_arguments)]
2097 fn hyper_range_decode(
2098 &self,
2099 e: &Engine,
2100 topology: &crate::hyper::HyperTopology,
2101 mut x: CudaSlice<f32>,
2102 lo: usize,
2103 hi: usize,
2104 pos_d: &CudaSlice<i32>,
2105 pos: usize,
2106 cache: &mut Cache,
2107 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2108 if hyper_decode_ws_on() {
2114 return self.hyper_range_decode_ws(e, topology, x, lo, hi, pos_d, pos, cache);
2115 }
2116 let n_embd = self.cfg.n_embd as usize;
2117 let eps = self.cfg.rms_eps;
2118 for il in lo..hi {
2119 let layer = &self.layers[il];
2120 let hyper = layer.hyper.as_ref().ok_or_else(|| {
2121 format!("layer {il} carries no hyper-connection weights under an hc plan")
2122 })?;
2123
2124 let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, 1, n_embd)?;
2125 let mut h = e.uninit(n_embd)?;
2126 e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, 1, eps)?;
2127 let mixed = match &layer.mixer {
2128 Mixer::Full(fa) => self.full_attn_decode(e, fa, &h, pos_d, pos, cache, il)?,
2129 Mixer::Linear(la) => self.linear_attn_decode(e, la, &h, cache, il)?,
2130 Mixer::Mla(mla) if mla.tp.is_some() => {
2131 self.mla_tp_attn_cached(e, mla, &h, pos_d, 1, il, cache, false)?
2132 }
2133 Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, 1, il, cache)?,
2134 Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2135 e,
2136 la,
2137 &h,
2138 1,
2139 eps,
2140 cache,
2141 il,
2142 crate::kda::ConvArm::Decode,
2143 )?,
2144 Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il)?,
2145 };
2146 x = crate::hyper::post(e, topology, &mixed, &x, &mix, 1, n_embd)?;
2147
2148 let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, 1, n_embd)?;
2149 let mut z = e.uninit(n_embd)?;
2150 e.rms_norm(
2151 &y,
2152 layer.post_attn_norm.float_data(),
2153 &mut z,
2154 n_embd,
2155 1,
2156 eps,
2157 )?;
2158 let ffn_out = self.hyper_ffn_branch(e, layer, &z, 1, il, false)?;
2159 x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, 1, n_embd)?;
2160 }
2161 Ok(x)
2162 }
2163
2164 #[allow(clippy::too_many_arguments)]
2171 fn hyper_range_decode_ws(
2172 &self,
2173 e: &Engine,
2174 topology: &crate::hyper::HyperTopology,
2175 x: CudaSlice<f32>,
2176 lo: usize,
2177 hi: usize,
2178 pos_d: &CudaSlice<i32>,
2179 pos: usize,
2180 cache: &mut Cache,
2181 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2182 let n_embd = self.cfg.n_embd as usize;
2183 let mut ws = match e.hyper_ws_take() {
2184 Some(ws) if ws.matches(topology, n_embd) => ws,
2185 _ => crate::hyper::HyperDecodeWs::new(e, topology, n_embd)?,
2186 };
2187 if HC_DECODE_WS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
2188 eprintln!(
2189 "[hc-decode-ws] engaged streams={} hidden={n_embd} (persistent hc-glue \
2190 workspace, per-engine pool; MEMRA_HC_DECODE_WS=1)",
2191 topology.streams
2192 );
2193 }
2194 let out =
2195 self.hyper_range_decode_ws_body(e, topology, x, lo, hi, pos_d, pos, cache, &mut ws);
2196 e.hyper_ws_put(ws);
2197 out
2198 }
2199
2200 #[allow(clippy::too_many_arguments)]
2205 fn hyper_range_decode_ws_body(
2206 &self,
2207 e: &Engine,
2208 topology: &crate::hyper::HyperTopology,
2209 mut x: CudaSlice<f32>,
2210 lo: usize,
2211 hi: usize,
2212 pos_d: &CudaSlice<i32>,
2213 pos: usize,
2214 cache: &mut Cache,
2215 ws: &mut crate::hyper::HyperDecodeWs,
2216 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2217 let n_embd = self.cfg.n_embd as usize;
2218 let eps = self.cfg.rms_eps;
2219 for il in lo..hi {
2220 let layer = &self.layers[il];
2221 let hyper = layer.hyper.as_ref().ok_or_else(|| {
2222 format!("layer {il} carries no hyper-connection weights under an hc plan")
2223 })?;
2224
2225 crate::hyper::pre_t1_ws(e, topology, &hyper.attn, &x, ws, n_embd)?;
2226 e.rms_norm(
2227 &ws.y,
2228 layer.attn_norm.float_data(),
2229 &mut ws.h,
2230 n_embd,
2231 1,
2232 eps,
2233 )?;
2234 let mixed = match &layer.mixer {
2235 Mixer::Full(fa) => self.full_attn_decode(e, fa, &ws.h, pos_d, pos, cache, il)?,
2236 Mixer::Linear(la) => self.linear_attn_decode(e, la, &ws.h, cache, il)?,
2237 Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &ws.h, pos_d, 1, il, cache)?,
2238 Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &ws.h, eps, cache, il)?,
2239 };
2240 crate::hyper::post_t1_ws(e, topology, &mixed, &x, ws, n_embd)?;
2241 std::mem::swap(&mut x, &mut ws.xb);
2242
2243 crate::hyper::pre_t1_ws(e, topology, &hyper.mlp, &x, ws, n_embd)?;
2244 e.rms_norm(
2245 &ws.y,
2246 layer.post_attn_norm.float_data(),
2247 &mut ws.z,
2248 n_embd,
2249 1,
2250 eps,
2251 )?;
2252 let ffn_out = self.hyper_ffn_branch(e, layer, &ws.z, 1, il, false)?;
2253 crate::hyper::post_t1_ws(e, topology, &ffn_out, &x, ws, n_embd)?;
2254 std::mem::swap(&mut x, &mut ws.xb);
2255 }
2256 Ok(x)
2257 }
2258
2259 #[allow(clippy::too_many_arguments)]
2294 pub(crate) fn hyper_batch_range_decode(
2295 &self,
2296 e: &Engine,
2297 topology: &crate::hyper::HyperTopology,
2298 mut x: CudaSlice<f32>,
2299 lo: usize,
2300 hi: usize,
2301 pos_rows: &[CudaSlice<i32>],
2302 caches: &mut [&mut Cache],
2303 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2304 let b_n = caches.len();
2305 assert_eq!(
2306 pos_rows.len(),
2307 b_n,
2308 "hyper_batch_range_decode: pos_rows built for a different batch width"
2309 );
2310 let n_embd = self.cfg.n_embd as usize;
2311 let eps = self.cfg.rms_eps;
2312 for il in lo..hi {
2313 let layer = &self.layers[il];
2314 let hyper = layer.hyper.as_ref().ok_or_else(|| {
2315 format!("layer {il} carries no hyper-connection weights under an hc plan")
2316 })?;
2317
2318 let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, b_n, n_embd)?;
2319 let mut h = e.uninit(b_n * n_embd)?;
2320 e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, b_n, eps)?;
2321 let mut mixed = e.uninit(b_n * n_embd)?;
2325 for bi in 0..b_n {
2326 let mut h_row = e.uninit(n_embd)?;
2327 e.dtod_copy_view(&h.slice(bi * n_embd..(bi + 1) * n_embd), &mut h_row)?;
2328 let cache: &mut Cache = &mut *caches[bi];
2329 let pos = cache.pos;
2330 let out_row = match &layer.mixer {
2331 Mixer::Full(fa) => {
2332 self.full_attn_decode(e, fa, &h_row, &pos_rows[bi], pos, cache, il)?
2333 }
2334 Mixer::Linear(la) => self.linear_attn_decode(e, la, &h_row, cache, il)?,
2335 Mixer::Mla(mla) => {
2336 self.mla_attn_cached(e, mla, &h_row, &pos_rows[bi], 1, il, cache)?
2337 }
2338 Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?,
2339 };
2340 e.copy_into(&mut mixed, bi * n_embd, &out_row, n_embd)?;
2341 }
2342 x = crate::hyper::post(e, topology, &mixed, &x, &mix, b_n, n_embd)?;
2343
2344 let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, b_n, n_embd)?;
2345 let mut z = e.uninit(b_n * n_embd)?;
2346 e.rms_norm(
2347 &y,
2348 layer.post_attn_norm.float_data(),
2349 &mut z,
2350 n_embd,
2351 b_n,
2352 eps,
2353 )?;
2354 let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, b_n, il, false)?;
2355 x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, b_n, n_embd)?;
2356 }
2357 Ok(x)
2358 }
2359
2360 pub(crate) fn hyper_ffn_branch_batch(
2380 &self,
2381 e: &Engine,
2382 layer: &crate::hybrid::HybridLayer,
2383 z: &CudaSlice<f32>,
2384 b_n: usize,
2385 il: usize,
2386 vrows: bool,
2387 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2388 let n_embd = self.cfg.n_embd as usize;
2389 match &layer.ffn {
2390 crate::hybrid::Ffn::Dense { .. } => {
2391 let mut out = e.uninit(b_n * n_embd)?;
2392 for bi in 0..b_n {
2393 let mut z_row = e.uninit(n_embd)?;
2394 e.dtod_copy_view(&z.slice(bi * n_embd..(bi + 1) * n_embd), &mut z_row)?;
2395 let row = self.hyper_ffn_branch(e, layer, &z_row, 1, il, false)?;
2396 e.copy_into(&mut out, bi * n_embd, &row, n_embd)?;
2397 }
2398 Ok(out)
2399 }
2400 crate::hybrid::Ffn::Moe(m) => {
2401 if vrows {
2402 self.moe_ffn_il_zq8_vrows(e, m, z, b_n, il as u16)
2403 } else {
2404 self.moe_ffn_il_zq8(e, m, z, None, b_n, il as u16)
2405 }
2406 }
2407 }
2408 }
2409
2410 fn forward_hyper_ppn(
2448 &self,
2449 e: &Engine,
2450 tokens: &[u32],
2451 last_only: bool,
2452 topology: &crate::hyper::HyperTopology,
2453 fence: &[usize],
2454 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2455 let n_embd = self.cfg.n_embd as usize;
2456 let eps = self.cfg.rms_eps;
2457 let t = tokens.len();
2458 let width = topology.streams * n_embd;
2459 let trace = memra_reference::hidden_trace::enabled();
2460 if trace {
2461 memra_reference::hidden_trace::emit_tokens(tokens);
2462 }
2463 let pos: Vec<i32> = (0..t as i32).collect();
2464
2465 if crate::pp::pp2_streams_off() {
2466 let pos_d = e.htod_i32(&pos)?;
2470 let embedded = self.embed(e, tokens)?;
2471 let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
2472 x = self.hyper_range_forward(e, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
2473 for s in 1..fence.len() - 1 {
2474 let boundary_tx = e.clone_dtod(&x)?;
2475 let boundary_rx = e.clone_dtod(&boundary_tx)?;
2476 x = self.hyper_range_forward(
2477 e,
2478 topology,
2479 boundary_rx,
2480 fence[s],
2481 fence[s + 1],
2482 &pos_d,
2483 t,
2484 trace,
2485 )?;
2486 }
2487 return self.hyper_head_logits(e, topology, &x, t, n_embd, eps, last_only);
2488 }
2489
2490 let rt = crate::pp::PpNRt::get(e)?;
2491 let n_st = fence.len() - 1;
2492 assert_eq!(
2493 rt.n_stages(),
2494 n_st,
2495 "PpNRt stage count {} != fence stages {n_st}",
2496 rt.n_stages()
2497 );
2498 rt.fence_stages_behind(&e.stream())?;
2501
2502 let mut slot = {
2503 let _st0 = rt.enter(0);
2504 let e0 = rt.engine(0, e);
2505 let pos_d = e0.htod_i32(&pos)?;
2508 let embedded = self.embed(e0, tokens)?;
2509 let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
2510 let x =
2511 self.hyper_range_forward(e0, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
2512 rt.tx(0, &x, t * width)?
2513 };
2514 for s in 1..n_st - 1 {
2515 let _st = rt.enter(s);
2516 let es = rt.engine(s, e);
2517 let pos_d = es.htod_i32(&pos)?;
2518 let x = rt.rx(s - 1, slot, t * width)?;
2519 let x = self.hyper_range_forward(
2520 es,
2521 topology,
2522 x,
2523 fence[s],
2524 fence[s + 1],
2525 &pos_d,
2526 t,
2527 trace,
2528 )?;
2529 slot = rt.tx(s, &x, t * width)?;
2530 }
2531 let _stl = rt.enter(n_st - 1);
2532 let el = rt.engine(n_st - 1, e);
2533 let pos_d = el.htod_i32(&pos)?;
2534 let x = rt.rx(n_st - 2, slot, t * width)?;
2535 let x = self.hyper_range_forward(
2536 el,
2537 topology,
2538 x,
2539 fence[n_st - 1],
2540 fence[n_st],
2541 &pos_d,
2542 t,
2543 trace,
2544 )?;
2545 self.hyper_head_logits(el, topology, &x, t, n_embd, eps, last_only)
2546 }
2547
2548 #[allow(clippy::too_many_arguments)]
2553 fn hyper_head_logits(
2554 &self,
2555 e: &Engine,
2556 topology: &crate::hyper::HyperTopology,
2557 x: &CudaSlice<f32>,
2558 t: usize,
2559 n_embd: usize,
2560 eps: f32,
2561 last_only: bool,
2562 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2563 let collapsed =
2564 crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
2565 if memra_reference::hidden_trace::enabled() {
2566 memra_reference::hidden_trace::emit_last_row(
2567 "collapse",
2568 -1,
2569 t,
2570 n_embd,
2571 &e.dtoh(&collapsed)?,
2572 );
2573 }
2574 let mut hn = e.uninit(t * n_embd)?;
2575 e.rms_norm(
2576 &collapsed,
2577 self.output_norm.float_data(),
2578 &mut hn,
2579 n_embd,
2580 t,
2581 eps,
2582 )?;
2583 let logits = if last_only {
2584 let last = e.view(&hn, t * n_embd);
2585 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2586 let mut hlast = e.uninit(n_embd)?;
2587 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2588 e.matmul(&self.output, &hlast, 1)?
2589 } else {
2590 e.matmul(&self.output, &hn, t)?
2591 };
2592 e.dtoh(&logits)
2593 }
2594
2595 #[allow(clippy::type_complexity)] #[allow(clippy::too_many_arguments)]
2600 fn prime_cache_hyper_ppn(
2601 &self,
2602 e: &Engine,
2603 tokens: &[u32],
2604 cache: &mut Cache,
2605 queued_after: usize,
2606 topology: &crate::hyper::HyperTopology,
2607 fence: &[usize],
2608 overlay: Option<&crate::vision::EmbedOverlay>,
2609 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2610 let n_embd = self.cfg.n_embd as usize;
2611 let eps = self.cfg.rms_eps;
2612 let t = tokens.len();
2613 let width = topology.streams * n_embd;
2614 if cache.pos + t > cache.max_ctx {
2615 return Err("prime_cache: prompt exceeds cache max_ctx".into());
2616 }
2617 let seq_end = cache.pos + t + queued_after;
2618 let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
2619 if let Some(sink) = cache.hc_taps.as_mut() {
2622 sink.base = cache.pos;
2623 }
2624
2625 if crate::pp::pp2_streams_off() {
2626 let pos_d = e.htod_i32(&pos)?;
2627 let mut embedded = self.embed(e, tokens)?;
2628 if let Some(ov) = overlay {
2629 ov.splice_into(e, &mut embedded, 0, t, n_embd)?;
2633 }
2634 let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
2635 x = self.hyper_range_prime(
2636 e, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end,
2637 )?;
2638 for s in 1..fence.len() - 1 {
2639 let boundary_tx = e.clone_dtod(&x)?;
2640 let boundary_rx = e.clone_dtod(&boundary_tx)?;
2641 x = self.hyper_range_prime(
2642 e,
2643 topology,
2644 boundary_rx,
2645 fence[s],
2646 fence[s + 1],
2647 &pos_d,
2648 t,
2649 cache,
2650 seq_end,
2651 )?;
2652 }
2653 return self.hyper_prime_tail(e, topology, &x, t, n_embd, eps, cache);
2654 }
2655
2656 {
2657 let rt = crate::pp::PpNRt::get(e)?;
2658 let n_st = fence.len() - 1;
2659 assert_eq!(
2660 rt.n_stages(),
2661 n_st,
2662 "PpNRt stage count {} != fence stages {n_st}",
2663 rt.n_stages()
2664 );
2665 rt.fence_stages_behind(&e.stream())?;
2666 if overlay.is_some() && !std::ptr::eq(rt.engine(0, e), e) {
2672 return Err(
2673 "vision embedding overlay requires stage 0 on the primary device \
2674 (overlay rows are primary-resident); place stage 0 on the primary \
2675 or run MEMRA_PP_STREAMS=0"
2676 .into(),
2677 );
2678 }
2679 let mut slot = {
2680 let _st0 = rt.enter(0);
2681 let e0 = rt.engine(0, e);
2682 let pos_d = e0.htod_i32(&pos)?;
2683 let mut embedded = self.embed(e0, tokens)?;
2684 if let Some(ov) = overlay {
2685 ov.splice_into(e0, &mut embedded, 0, t, n_embd)?;
2691 }
2692 let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
2693 let x = self.hyper_range_prime(
2694 e0, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end,
2695 )?;
2696 rt.tx(0, &x, t * width)?
2697 };
2698 for s in 1..n_st - 1 {
2699 let _st = rt.enter(s);
2700 let es = rt.engine(s, e);
2701 let pos_d = es.htod_i32(&pos)?;
2702 let x = rt.rx(s - 1, slot, t * width)?;
2703 let x = self.hyper_range_prime(
2704 es,
2705 topology,
2706 x,
2707 fence[s],
2708 fence[s + 1],
2709 &pos_d,
2710 t,
2711 cache,
2712 seq_end,
2713 )?;
2714 slot = rt.tx(s, &x, t * width)?;
2715 }
2716 let _stl = rt.enter(n_st - 1);
2717 let el = rt.engine(n_st - 1, e);
2718 let pos_d = el.htod_i32(&pos)?;
2719 let x = rt.rx(n_st - 2, slot, t * width)?;
2720 let x = self.hyper_range_prime(
2721 el,
2722 topology,
2723 x,
2724 fence[n_st - 1],
2725 fence[n_st],
2726 &pos_d,
2727 t,
2728 cache,
2729 seq_end,
2730 )?;
2731 self.hyper_prime_tail(el, topology, &x, t, n_embd, eps, cache)
2732 }
2733 }
2734
2735 #[allow(clippy::type_complexity)] fn prime_chunk_hyper(
2741 &self,
2742 e: &Engine,
2743 tokens: &[u32],
2744 cache: &mut Cache,
2745 seq_end: usize,
2746 chunk_off: usize,
2747 overlay: Option<&crate::vision::EmbedOverlay>,
2748 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2749 let topology = *self
2750 .hyper
2751 .as_ref()
2752 .ok_or("prime_chunk_hyper on a model with no HyperConnections topology")?;
2753 let n_embd = self.cfg.n_embd as usize;
2754 let t = tokens.len();
2755 let eps = self.cfg.rms_eps;
2756 let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
2757 let pos_d = e.htod_i32(&pos)?;
2758 if let Some(sink) = cache.hc_taps.as_mut() {
2761 sink.base = cache.pos;
2762 }
2763
2764 let mut embedded = self.embed(e, tokens)?;
2765 if let Some(ov) = overlay {
2766 ov.splice_into(e, &mut embedded, chunk_off, t, n_embd)?;
2772 }
2773 let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
2774
2775 for (il, layer) in self.layers.iter().enumerate() {
2776 let hyper = layer.hyper.as_ref().ok_or_else(|| {
2777 format!("layer {il} carries no hyper-connection weights under an hc plan")
2778 })?;
2779
2780 let (y, mix) = crate::hyper::pre(e, &topology, &hyper.attn, &x, t, n_embd)?;
2781 let mut h = e.uninit(t * n_embd)?;
2782 e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2783 let mixed = match &layer.mixer {
2784 Mixer::Full(fa) => {
2785 self.full_attn_prime(e, fa, &h, None, &pos_d, t, cache, il, seq_end)?
2786 }
2787 Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2788 Mixer::Mla(mla) if mla.tp.is_some() => {
2789 self.mla_tp_attn_cached(e, mla, &h, &pos_d, t, il, cache, false)?
2790 }
2791 Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, &pos_d, t, il, cache)?,
2792 Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2793 e,
2794 la,
2795 &h,
2796 t,
2797 eps,
2798 cache,
2799 il,
2800 crate::kda::ConvArm::Prefill,
2801 )?,
2802 Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2803 };
2804 x = crate::hyper::post(e, &topology, &mixed, &x, &mix, t, n_embd)?;
2805
2806 let (y, mix) = crate::hyper::pre(e, &topology, &hyper.mlp, &x, t, n_embd)?;
2807 let mut z = e.uninit(t * n_embd)?;
2808 e.rms_norm(
2809 &y,
2810 layer.post_attn_norm.float_data(),
2811 &mut z,
2812 n_embd,
2813 t,
2814 eps,
2815 )?;
2816 let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2817 x = crate::hyper::post(e, &topology, &ffn_out, &x, &mix, t, n_embd)?;
2818 self.glm5_hc_tap(e, cache, &topology, il, &x, t)?;
2821 }
2822
2823 let hiddens =
2824 crate::hyper::collapse(e, &topology, self.hyper_head.as_ref(), &x, t, n_embd)?;
2825 let mut hn = e.uninit(t * n_embd)?;
2826 e.rms_norm(
2827 &hiddens,
2828 self.output_norm.float_data(),
2829 &mut hn,
2830 n_embd,
2831 t,
2832 eps,
2833 )?;
2834 let last = e.view(&hn, t * n_embd);
2835 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2836 let mut hlast = e.uninit(n_embd)?;
2837 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2838 let logits = e.matmul(&self.output, &hlast, 1)?;
2839 let host = e.dtoh(&logits)?;
2840
2841 let stack = e.view(&hiddens, t * n_embd);
2844 let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
2845 let mut h_seed = e.uninit(n_embd)?;
2846 e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
2847 cache.pos += t;
2848 Ok((host, h_seed, hiddens))
2849 }
2850
2851 #[allow(clippy::too_many_arguments)]
2855 #[allow(clippy::type_complexity)] fn hyper_prime_tail(
2857 &self,
2858 e: &Engine,
2859 topology: &crate::hyper::HyperTopology,
2860 x: &CudaSlice<f32>,
2861 t: usize,
2862 n_embd: usize,
2863 eps: f32,
2864 cache: &mut Cache,
2865 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2866 let hiddens = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
2867 let mut hn = e.uninit(t * n_embd)?;
2868 e.rms_norm(
2869 &hiddens,
2870 self.output_norm.float_data(),
2871 &mut hn,
2872 n_embd,
2873 t,
2874 eps,
2875 )?;
2876 let last = e.view(&hn, t * n_embd);
2877 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2878 let mut hlast = e.uninit(n_embd)?;
2879 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2880 let logits = e.matmul(&self.output, &hlast, 1)?;
2881 let host = e.dtoh(&logits)?;
2882 let stack = e.view(&hiddens, t * n_embd);
2883 let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
2884 let mut h_seed = e.uninit(n_embd)?;
2885 e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
2886 cache.pos += t;
2887 Ok((host, h_seed, hiddens))
2888 }
2889
2890 fn decode_step_hyper_ppn(
2895 &self,
2896 e: &Engine,
2897 token: u32,
2898 cache: &mut Cache,
2899 topology: &crate::hyper::HyperTopology,
2900 fence: &[usize],
2901 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2902 let n_embd = self.cfg.n_embd as usize;
2903 let eps = self.cfg.rms_eps;
2904 let pos = cache.pos;
2905 let width = topology.streams * n_embd;
2906
2907 if crate::pp::pp2_streams_off() {
2908 let pos_d = e.htod_i32(&[pos as i32])?;
2909 let embedded = e.htod(&self.embd.gather(n_embd, &[token]))?;
2910 let mut x = crate::hyper::expand(e, topology, &embedded, 1, n_embd)?;
2911 x = self.hyper_range_decode(e, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
2912 for s in 1..fence.len() - 1 {
2913 let boundary_tx = e.clone_dtod(&x)?;
2914 let boundary_rx = e.clone_dtod(&boundary_tx)?;
2915 x = self.hyper_range_decode(
2916 e,
2917 topology,
2918 boundary_rx,
2919 fence[s],
2920 fence[s + 1],
2921 &pos_d,
2922 pos,
2923 cache,
2924 )?;
2925 }
2926 return self.hyper_decode_tail(e, topology, &x, n_embd, eps, cache);
2927 }
2928
2929 let rt = crate::pp::PpNRt::get(e)?;
2930 let n_st = fence.len() - 1;
2931 assert_eq!(
2932 rt.n_stages(),
2933 n_st,
2934 "PpNRt stage count {} != fence stages {n_st}",
2935 rt.n_stages()
2936 );
2937 rt.fence_stages_behind(&e.stream())?;
2938
2939 let mut slot = {
2940 let _st0 = rt.enter(0);
2941 let e0 = rt.engine(0, e);
2942 let pos_d = e0.htod_i32(&[pos as i32])?;
2943 let embedded = e0.htod(&self.embd.gather(n_embd, &[token]))?;
2944 let x = crate::hyper::expand(e0, topology, &embedded, 1, n_embd)?;
2945 let x =
2946 self.hyper_range_decode(e0, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
2947 rt.tx(0, &x, width)?
2948 };
2949 for s in 1..n_st - 1 {
2950 let _st = rt.enter(s);
2951 let es = rt.engine(s, e);
2952 let pos_d = es.htod_i32(&[pos as i32])?;
2953 let x = rt.rx(s - 1, slot, width)?;
2954 let x = self.hyper_range_decode(
2955 es,
2956 topology,
2957 x,
2958 fence[s],
2959 fence[s + 1],
2960 &pos_d,
2961 pos,
2962 cache,
2963 )?;
2964 slot = rt.tx(s, &x, width)?;
2965 }
2966 let _stl = rt.enter(n_st - 1);
2967 let el = rt.engine(n_st - 1, e);
2968 let pos_d = el.htod_i32(&[pos as i32])?;
2969 let x = rt.rx(n_st - 2, slot, width)?;
2970 let x = self.hyper_range_decode(
2971 el,
2972 topology,
2973 x,
2974 fence[n_st - 1],
2975 fence[n_st],
2976 &pos_d,
2977 pos,
2978 cache,
2979 )?;
2980 self.hyper_decode_tail(el, topology, &x, n_embd, eps, cache)
2981 }
2982
2983 fn hyper_decode_tail(
2987 &self,
2988 e: &Engine,
2989 topology: &crate::hyper::HyperTopology,
2990 x: &CudaSlice<f32>,
2991 n_embd: usize,
2992 eps: f32,
2993 cache: &mut Cache,
2994 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2995 let h_seed = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, 1, n_embd)?;
2996 let mut hn = e.uninit(n_embd)?;
2997 e.rms_norm(
2998 &h_seed,
2999 self.output_norm.float_data(),
3000 &mut hn,
3001 n_embd,
3002 1,
3003 eps,
3004 )?;
3005 let logits = e.matmul(&self.output, &hn, 1)?;
3006 let host = e.dtoh(&logits)?;
3007 cache.pos += 1;
3008 Ok((host, h_seed))
3009 }
3010
3011 pub fn forward(
3013 &self,
3014 e: &Engine,
3015 tokens: &[u32],
3016 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3017 if self.hyper.is_some() {
3018 return self.forward_hyper(e, tokens, false);
3019 }
3020 if self.is_gemma4_e4b() {
3021 return self.gemma4_e4b_forward(e, tokens, false);
3022 }
3023 if self.uses_gemma_program() {
3024 return self.gemma4_forward(e, tokens, false);
3025 }
3026 let cfg = &self.cfg;
3027 let n_embd = cfg.n_embd as usize;
3028 let t = tokens.len();
3029 let eps = cfg.rms_eps;
3030 let pos: Vec<i32> = (0..t as i32).collect();
3031 let pos_d = e.htod_i32(&pos)?;
3032
3033 let mut x = self.embed(e, tokens)?; for (il, layer) in self.layers.iter().enumerate() {
3036 let mut h = e.uninit(t * n_embd)?;
3038 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3039
3040 let mixed = match &layer.mixer {
3041 Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
3042 Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
3043 Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
3044 Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
3045 };
3046
3047 let mut x1 = e.uninit(t * n_embd)?;
3049 e.add(&x, &mixed, &mut x1, t * n_embd)?;
3050
3051 let mut z = e.uninit(t * n_embd)?;
3053 e.rms_norm(
3054 &x1,
3055 layer.post_attn_norm.float_data(),
3056 &mut z,
3057 n_embd,
3058 t,
3059 eps,
3060 )?;
3061 let ffn_out = match &layer.ffn {
3062 crate::hybrid::Ffn::Dense {
3063 ffn_gate,
3064 ffn_up,
3065 ffn_down,
3066 } => {
3067 let n_ff = ffn_gate.out_features();
3068 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
3069 let up = g2.pop().unwrap();
3070 let gate = g2.pop().unwrap();
3071 let mut act = e.uninit(t * n_ff)?;
3072 Self::ffn_act_lim(
3077 e,
3078 &self.cfg,
3079 &gate,
3080 &up,
3081 1.0,
3082 1.0,
3083 self.cfg.clamp_shexp_at(il as u32),
3084 &mut act,
3085 t * n_ff,
3086 )?;
3087 e.matmul(ffn_down, &act, t)?
3088 }
3089 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
3090 };
3091 let mut x2 = e.uninit(t * n_embd)?;
3092 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
3093 x = x2;
3094 }
3095
3096 let mut hn = e.uninit(t * n_embd)?;
3097 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3098 let logits = e.matmul(&self.output, &hn, t)?;
3099 e.dtoh(&logits)
3100 }
3101
3102 pub fn forward_last(
3108 &self,
3109 e: &Engine,
3110 tokens: &[u32],
3111 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3112 if self.hyper.is_some() {
3113 return self.forward_hyper(e, tokens, true);
3114 }
3115 if self.uses_gemma_program() {
3116 return self.gemma4_forward(e, tokens, true);
3117 }
3118 let cfg = &self.cfg;
3119 let n_embd = cfg.n_embd as usize;
3120 let t = tokens.len();
3121 let eps = cfg.rms_eps;
3122 let pos: Vec<i32> = (0..t as i32).collect();
3123 let pos_d = e.htod_i32(&pos)?;
3124
3125 let mut x = self.embed(e, tokens)?; let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
3129 let anat = Self::prime_anatomy_on();
3130 let mut anat_last = if anat {
3131 e.stream().synchronize()?;
3132 Some(std::time::Instant::now())
3133 } else {
3134 None
3135 };
3136 macro_rules! anat_mark {
3137 ($slot:expr) => {
3138 if let Some(ts) = anat_last.as_mut() {
3139 e.stream().synchronize()?;
3140 Self::prime_anatomy_slots()[$slot].fetch_add(
3141 ts.elapsed().as_nanos() as u64,
3142 std::sync::atomic::Ordering::Relaxed,
3143 );
3144 *ts = std::time::Instant::now();
3145 }
3146 };
3147 }
3148 for (il, layer) in self.layers.iter().enumerate() {
3149 let mut h = e.uninit(t * n_embd)?;
3150 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3151 if probe {
3152 e.stream().synchronize()?;
3153 eprintln!("[probe] L{il} norm ok");
3154 }
3155 anat_mark!(4);
3156 let mixed = match &layer.mixer {
3157 Mixer::Full(fa) => {
3158 let y = self.full_attn(e, fa, &h, &pos_d, t, il)?;
3159 anat_mark!(0);
3160 y
3161 }
3162 Mixer::Linear(la) => {
3163 let y = self.linear_attn(e, la, &h, t)?;
3164 anat_mark!(1);
3165 y
3166 }
3167 Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
3168 Mixer::Kda(la) => {
3169 let y = crate::kda::kda_attn(e, la, &h, t, eps)?;
3170 anat_mark!(1);
3172 y
3173 }
3174 };
3175 if probe {
3176 e.stream().synchronize()?;
3177 eprintln!("[probe] L{il} mixer ok");
3178 }
3179 let mut x1 = e.uninit(t * n_embd)?;
3180 e.add(&x, &mixed, &mut x1, t * n_embd)?;
3181 let mut z = e.uninit(t * n_embd)?;
3182 e.rms_norm(
3183 &x1,
3184 layer.post_attn_norm.float_data(),
3185 &mut z,
3186 n_embd,
3187 t,
3188 eps,
3189 )?;
3190 anat_mark!(4);
3191 let ffn_out = match &layer.ffn {
3192 crate::hybrid::Ffn::Dense {
3193 ffn_gate,
3194 ffn_up,
3195 ffn_down,
3196 } => {
3197 let n_ff = ffn_gate.out_features();
3198 let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
3199 let up = g2.pop().unwrap();
3200 let gate = g2.pop().unwrap();
3201 let mut act = e.uninit(t * n_ff)?;
3202 Self::ffn_act_lim(
3204 e,
3205 &self.cfg,
3206 &gate,
3207 &up,
3208 1.0,
3209 1.0,
3210 self.cfg.clamp_shexp_at(il as u32),
3211 &mut act,
3212 t * n_ff,
3213 )?;
3214 let y = e.matmul(ffn_down, &act, t)?;
3215 anat_mark!(3);
3216 y
3217 }
3218 crate::hybrid::Ffn::Moe(m) => {
3219 let y = self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?;
3220 anat_mark!(2);
3221 y
3222 }
3223 };
3224 if probe {
3225 e.stream().synchronize()?;
3226 eprintln!("[probe] L{il} ffn ok");
3227 }
3228 let mut x2 = e.uninit(t * n_embd)?;
3229 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
3230 x = x2;
3231 }
3232 if anat {
3233 let s = Self::prime_anatomy_slots();
3234 let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
3235 eprintln!(
3236 "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
3237 dense={:.1} norms_adds={:.1} (t={t}, forward_last)",
3238 ms(0),
3239 ms(1),
3240 ms(2),
3241 ms(3),
3242 ms(4)
3243 );
3244 }
3245 let mut hn = e.uninit(t * n_embd)?;
3247 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3248 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)?;
3251 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
3252 let logits = e.matmul(&self.output, &hlast, 1)?; e.dtoh(&logits)
3254 }
3255
3256 #[allow(clippy::type_complexity)] pub fn prime_cache(
3289 &self,
3290 e: &Engine,
3291 tokens: &[u32],
3292 cache: &mut Cache,
3293 queued_after: usize,
3294 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3295 self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
3296 }
3297
3298 #[allow(clippy::type_complexity)] pub fn prime_cache_overlaid(
3307 &self,
3308 e: &Engine,
3309 tokens: &[u32],
3310 cache: &mut Cache,
3311 queued_after: usize,
3312 overlay: Option<&crate::vision::EmbedOverlay>,
3313 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3314 cache.ensure_usable("prime_cache")?;
3315 if self.hyper.is_some() {
3316 return self.prime_cache_hyper(e, tokens, cache, queued_after, overlay);
3320 }
3321 let _pp_walk =
3322 if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
3323 let rt = crate::pp::PpNRt::get(e)?;
3324 Some(rt.acquire_walk("prime_cache")?)
3325 } else {
3326 None
3327 };
3328 let n_embd = self.cfg.n_embd as usize;
3329 let t = tokens.len();
3330 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
3359 let seq_end = if legacy_calllocal {
3360 cache.pos + t
3361 } else {
3362 cache.pos + t + queued_after
3363 };
3364 if overlay.is_none()
3369 && (cache.pos == 0 || step_gemm_prime_suffix_on())
3370 && t >= PRIME_MIN_T
3371 && crate::step_gemm_prime_on()
3372 && self.uses_sliding_gated_moe_program()
3373 {
3374 let n_embd = self.cfg.n_embd as usize;
3375 let base = cache.pos;
3376 let width = crate::cache::PRIME_CHUNK_MAX_TOKENS;
3377 let mut hiddens = e.uninit(t * n_embd)?;
3378 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3379 let mut start = 0usize;
3380 while start < t {
3381 let mut end = (start + width).min(t);
3384 if t - end > 0 && t - end < PRIME_MIN_T {
3385 end = t;
3386 }
3387 let mut out = self.step35_prime_cache_batch(
3388 e,
3389 &[&tokens[start..end]],
3390 &mut [cache],
3391 &[seq_end],
3392 )?;
3393 if out.len() != 1 {
3394 return Err("B=1 batched prime returned a non-singleton".into());
3395 }
3396 let (logits, h_seed, hidden) = out.remove(0);
3397 e.copy_into(
3398 &mut hiddens,
3399 start * n_embd,
3400 &hidden,
3401 (end - start) * n_embd,
3402 )?;
3403 last = Some((logits, h_seed));
3404 start = end;
3405 }
3406 let (logits, h_seed) = last.expect("prime produced no chunk");
3407 eprintln!(
3412 "[gemm-prime] ENGAGED t={t} base={base} seq_end={seq_end} chunks<={width} (GEMM trunk + grouped MoE)"
3413 );
3414 return Ok((logits, h_seed, hiddens));
3415 }
3416 if self.uses_sliding_gated_moe_program() {
3417 eprintln!(
3418 "[gemm-prime] WALK t={t} base={} seq_end={seq_end} (batched prime declined)",
3419 cache.pos
3420 );
3421 }
3422 if overlay.is_none()
3423 && let Some(out) = self.step35_prime_trows(e, tokens, cache)?
3424 {
3425 return Ok(out);
3426 }
3427 assert!(
3431 t >= PRIME_MIN_T,
3432 "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
3433 );
3434 assert!(
3435 cache.pos + t <= cache.max_ctx,
3436 "prime_cache: prompt exceeds cache max_ctx"
3437 );
3438
3439 if self.is_gemma4_e4b() || self.uses_gemma_program() {
3451 if self.is_gemma4_e4b() {
3452 if overlay.is_some() {
3453 return Err(
3454 "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
3455 );
3456 }
3457 return self.gemma4_e4b_prime(e, tokens, cache);
3458 }
3459 return self.gemma4_prime(e, tokens, cache, overlay);
3464 }
3465 if crate::pp::prime_pipe_on()
3466 && crate::pp::prime_pp_on()
3467 && !crate::pp::pp2_streams_off()
3468 && crate::pp::pp_cuts(self.layers.len())
3469 .is_some_and(|fence| matches!(fence.len(), 4 | 5))
3470 {
3471 crate::pp::pp_wave_on()
3472 .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3473 }
3474 let ranges = prime_chunk_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
3475 if ranges.len() == 1 {
3513 return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
3514 }
3515 if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
3519 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
3520 if overlay.is_some() {
3521 return Err(
3522 "vision embedding overlay + pipelined PP prime unsupported (v1); \
3523 run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
3524 .into(),
3525 );
3526 }
3527 if crate::pp::pp_multi_stream_same_device() {
3528 return Err(
3529 "prime chunk pipeline refused with 2 stage streams on one device — \
3530 that concurrent-stream placement remains quarantined by the deferred \
3531 pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
3532 the serial split."
3533 .into(),
3534 );
3535 }
3536 return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
3537 }
3538 if let Some(fence) =
3539 crate::pp::pp_cuts(self.layers.len()).filter(|f| matches!(f.len(), 4 | 5))
3540 {
3541 let wave_on = crate::pp::pp_wave_on()
3542 .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3543 let stages = fence.len() - 1;
3544 if crate::pp::pp_wave_route_enabled(
3545 wave_on,
3546 crate::pp::pp2_overlap(),
3547 stages,
3548 ranges.len(),
3549 ) {
3550 if overlay.is_some() {
3551 return Err(
3552 "vision embedding overlay + pipelined PP prime unsupported; \
3553 run the serial prime (MEMRA_PP_WAVE=0 or MEMRA_PRIME_PIPE=0)"
3554 .into(),
3555 );
3556 }
3557 let rt = crate::pp::PpNRt::get(e)?;
3558 let double_slot = crate::pp::pp2_overlap();
3559 crate::pp::pp_wave_eligibility(
3560 stages,
3561 double_slot,
3562 rt.host_bounce_active(),
3563 rt.repeated_stage_device(),
3564 )
3565 .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3566 return self
3567 .prime_cache_ppn_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
3568 }
3569 }
3570 }
3571 let mut hiddens = e.uninit(t * n_embd)?;
3572 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3573 for &(start, end) in &ranges {
3574 if let Some(taps) = cache.dflash_taps.as_mut() {
3576 taps.base = start;
3577 }
3578 let (l, hs, x) =
3579 self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
3580 e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
3581 last = Some((l, hs));
3582 }
3583 let (logits, h_seed) = last.unwrap();
3584 Ok((logits, h_seed, hiddens))
3585 }
3586
3587 #[allow(clippy::type_complexity)] fn prime_cache_pp2_pipelined(
3593 &self,
3594 e: &Engine,
3595 tokens: &[u32],
3596 cache: &mut Cache,
3597 seq_end: usize,
3598 ranges: &[(usize, usize)],
3599 fence: &[usize],
3600 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3601 debug_assert_eq!(fence.len(), 3);
3602 debug_assert!(ranges.len() >= 2);
3603 let rt = crate::pp::PpNRt::get(e)?;
3604 assert_eq!(
3605 rt.n_stages(),
3606 2,
3607 "prime pipeline requires exactly two PP stages"
3608 );
3609 let n_embd = self.cfg.n_embd as usize;
3610 let t = tokens.len();
3611 let initial_base = cache.pos;
3612 let caller_stream = e.stream();
3613
3614 rt.fence_stages_behind(&caller_stream)?;
3619 let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
3620 rt.prepare_overlap_slots(0, max_payload)?;
3621
3622 let mut hiddens = e.uninit(t * n_embd)?;
3623 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3624 let mut stage_caches = PrimeCacheStages::new(cache, fence);
3625 let (cache0, cache1) = stage_caches.pp2_parts();
3626 let (first_start, first_end) = ranges[0];
3627 let mut slot = self.prime_pp2_stage0_enqueue(
3628 e,
3629 rt,
3630 &tokens[first_start..first_end],
3631 cache0,
3632 seq_end,
3633 fence,
3634 initial_base + first_start,
3635 true,
3636 )?;
3637 cache0.pos = initial_base + first_end;
3638
3639 for (i, &(start, end)) in ranges.iter().enumerate() {
3640 let base = initial_base + start;
3641 debug_assert_eq!(
3642 cache1.pos, base,
3643 "stage 1 must drain chunks in original position order"
3644 );
3645 let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
3646 let next_base = initial_base + next_start;
3647 debug_assert_eq!(
3648 cache0.pos, next_base,
3649 "stage 0 must issue chunks in original position order"
3650 );
3651 let cache0_stage = &mut *cache0;
3652 std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
3657 let stage0 = scope.spawn(move || -> Result<usize, String> {
3658 let next = self
3659 .prime_pp2_stage0_enqueue(
3660 e,
3661 rt,
3662 &tokens[next_start..next_end],
3663 cache0_stage,
3664 seq_end,
3665 fence,
3666 next_base,
3667 true,
3668 )
3669 .map_err(|err| err.to_string())?;
3670 cache0_stage.pos = initial_base + next_end;
3671 Ok(next)
3672 });
3673 let x = self.prime_pp2_stage1_enqueue(
3674 e,
3675 rt,
3676 slot,
3677 end - start,
3678 cache1,
3679 seq_end,
3680 fence,
3681 base,
3682 true,
3683 )?;
3684 let out = {
3685 rt.bind_stage(1)?;
3686 let _st1 = rt.enter(1);
3687 let e1 = rt.engine(1, e);
3688 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
3689 };
3690 let next = match stage0.join() {
3691 Ok(result) => {
3692 result.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?
3693 }
3694 Err(payload) => std::panic::resume_unwind(payload),
3695 };
3696 Ok((out, Some(next)))
3697 })?
3698 } else {
3699 let x = self.prime_pp2_stage1_enqueue(
3700 e,
3701 rt,
3702 slot,
3703 end - start,
3704 cache1,
3705 seq_end,
3706 fence,
3707 base,
3708 true,
3709 )?;
3710 let out = {
3711 rt.bind_stage(1)?;
3712 let _st1 = rt.enter(1);
3713 let e1 = rt.engine(1, e);
3714 self.prime_chunk_epilogue(e1, x, end - start, cache1)?
3715 };
3716 (out, None)
3717 };
3718
3719 rt.publish_to(1, &caller_stream)?;
3720 e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
3721 last = Some((out.0, out.1));
3722 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3723
3724 if let Some(next) = next_slot {
3725 rt.fence_stages_behind(&caller_stream)?;
3730 slot = next;
3731 }
3732 }
3733
3734 debug_assert_eq!(cache0.pos, initial_base + t);
3735 debug_assert_eq!(cache1.pos, initial_base + t);
3736 let (logits, h_seed) = last.unwrap();
3737 stage_caches.commit();
3738 Ok((logits, h_seed, hiddens))
3739 }
3740
3741 #[allow(clippy::type_complexity)] fn prime_cache_ppn_pipelined(
3751 &self,
3752 e: &Engine,
3753 tokens: &[u32],
3754 cache: &mut Cache,
3755 seq_end: usize,
3756 ranges: &[(usize, usize)],
3757 fence: &[usize],
3758 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3759 let stages = fence.len().saturating_sub(1);
3760 debug_assert!((3..=4).contains(&stages));
3761 debug_assert!(ranges.len() >= 2);
3762 let rt = crate::pp::PpNRt::get(e)?;
3763 assert_eq!(
3764 rt.n_stages(),
3765 stages,
3766 "prime wavefront PpNRt/fence stage mismatch"
3767 );
3768 let n_embd = self.cfg.n_embd as usize;
3769 let initial_base = cache.pos;
3770 let caller_stream = e.stream();
3771 let primary_context = crate::pp::PrimaryContextRestore::new(e);
3772
3773 rt.fence_stages_behind(&caller_stream)?;
3774 let max_payload = ranges
3775 .iter()
3776 .map(|(start, end)| (end - start) * n_embd)
3777 .max()
3778 .unwrap_or(0);
3779 for boundary in 0..stages - 1 {
3780 rt.prepare_overlap_slots(boundary, max_payload)?;
3781 }
3782
3783 let mut stage_caches = PrimeCacheStages::new(cache, fence);
3784 let waves: Vec<_> = ranges
3785 .iter()
3786 .map(|&(start, end)| PrimePpWave {
3787 start,
3788 end,
3789 tokens: &tokens[start..end],
3790 })
3791 .collect();
3792 let mut forward_senders = Vec::with_capacity(stages - 1);
3793 let mut forward_receivers = Vec::with_capacity(stages - 1);
3794 let mut release_senders = Vec::with_capacity(stages - 1);
3795 let mut release_receivers = Vec::with_capacity(stages - 1);
3796 for _ in 0..stages - 1 {
3797 let (forward_sender, forward_receiver) = std::sync::mpsc::channel();
3798 let (release_sender, release_receiver) = std::sync::mpsc::channel();
3799 forward_senders.push(Some(forward_sender));
3800 forward_receivers.push(Some(forward_receiver));
3801 release_senders.push(Some(release_sender));
3802 release_receivers.push(Some(release_receiver));
3803 }
3804 let mut stage_channels = Vec::with_capacity(stages - 1);
3805 for stage in 0..stages - 1 {
3806 stage_channels.push(Some(PrimePpStageChannels {
3807 incoming: (stage > 0).then(|| forward_receivers[stage - 1].take().unwrap()),
3808 release_upstream: (stage > 0).then(|| release_senders[stage - 1].take().unwrap()),
3809 outgoing: forward_senders[stage].take().unwrap(),
3810 released_downstream: release_receivers[stage].take().unwrap(),
3811 }));
3812 }
3813 let head_incoming = forward_receivers[stages - 2].take().unwrap();
3814 let head_release = release_senders[stages - 2].take().unwrap();
3815 #[allow(clippy::type_complexity)]
3818 let mut results: Vec<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>> =
3819 std::iter::repeat_with(|| None).take(waves.len()).collect();
3820 let walk_result = std::thread::scope(|scope| -> Result<(), Box<dyn std::error::Error>> {
3821 let waves_ref = &waves;
3822 let mut handles = Vec::with_capacity(stages - 1);
3823 #[allow(clippy::needless_range_loop)]
3827 for stage in 0..stages - 1 {
3828 let channels = stage_channels[stage].take().unwrap();
3829 let cache_state = &stage_caches.stages()[stage];
3830 handles.push(scope.spawn(move || -> Result<(), String> {
3831 let result = self.prime_ppn_wave_worker(
3832 e,
3833 rt,
3834 waves_ref,
3835 cache_state,
3836 channels.incoming.as_ref(),
3837 channels.release_upstream.as_ref(),
3838 &channels.outgoing,
3839 &channels.released_downstream,
3840 stage,
3841 seq_end,
3842 fence,
3843 initial_base,
3844 );
3845 if let Err(error) = &result {
3846 channels.notify_failure(&error.to_string());
3847 }
3848 result.map_err(|error| error.to_string())
3849 }));
3850 }
3851
3852 let mut head_panic = None;
3853 let head_result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(
3854 || -> Result<(), Box<dyn std::error::Error>> {
3855 let mut head_cache = stage_caches.stages()[stages - 1]
3856 .lock()
3857 .map_err(|_| "prime PP head cache lock poisoned")?;
3858 for (wave_index, wave) in waves_ref.iter().enumerate() {
3859 let incoming = recv_prime_pp_signal(
3860 &head_incoming,
3861 PrimePpWaveSlot {
3862 wave: wave_index,
3863 slot: 0,
3864 },
3865 false,
3866 "prime PP head input",
3867 )?;
3868 results[wave_index] = Some(self.prime_ppn_wave_final(
3869 e,
3870 rt,
3871 wave,
3872 &mut head_cache,
3873 incoming,
3874 &head_release,
3875 seq_end,
3876 fence,
3877 initial_base,
3878 )?);
3879 }
3880 Ok(())
3881 },
3882 )) {
3883 Ok(result) => result,
3884 Err(payload) => {
3885 head_panic = Some(payload);
3886 Err("prime PP head-stage host walker panicked".into())
3887 }
3888 };
3889 if let Err(error) = &head_result {
3890 let _ = head_release.send(PrimePpSignal::Error(error.to_string()));
3891 }
3892 let mut first_error = head_result.err().map(|error| error.to_string());
3893 let mut worker_panic = None;
3894 for handle in handles {
3895 match handle.join() {
3896 Ok(Ok(())) => {}
3897 Ok(Err(error)) => {
3898 first_error.get_or_insert(error);
3899 }
3900 Err(payload) => {
3901 if worker_panic.is_none() {
3902 worker_panic = Some(payload);
3903 }
3904 }
3905 }
3906 }
3907 if let Some(payload) = head_panic {
3908 std::panic::resume_unwind(payload);
3909 }
3910 if let Some(payload) = worker_panic {
3911 std::panic::resume_unwind(payload);
3912 }
3913 if let Some(error) = first_error {
3914 return Err(error.into());
3915 }
3916 Ok(())
3917 });
3918 let publish_result = if walk_result.is_ok() {
3919 Some(rt.publish_to(stages - 1, &caller_stream))
3920 } else {
3921 None
3922 };
3923 let restore_result = primary_context.restore();
3924 walk_result?;
3925 if let Some(result) = publish_result {
3926 result?;
3927 }
3928 restore_result?;
3929 let mut hiddens = e.uninit(tokens.len() * n_embd)?;
3930 let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3931 for (wave_index, (wave, result)) in waves.iter().zip(results).enumerate() {
3932 debug_assert_eq!((wave.start, wave.end), ranges[wave_index]);
3933 let out = result.ok_or("prime PP wavefront completed without a head-stage result")?;
3934 e.copy_into(
3935 &mut hiddens,
3936 wave.start * n_embd,
3937 &out.2,
3938 (wave.end - wave.start) * n_embd,
3939 )?;
3940 last = Some((out.0, out.1));
3941 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3942 }
3943 stage_caches.commit();
3944 drop(stage_caches);
3945
3946 static LOGGED: std::sync::Once = std::sync::Once::new();
3947 LOGGED.call_once(|| {
3948 eprintln!(
3949 "[pp-wave] PP{stages} prime wavefront engaged: microchunks={} \
3950 (experimental, MEMRA_PP_WAVE=1)",
3951 ranges.len(),
3952 );
3953 });
3954 let (logits, h_seed) = last.expect("prime PP wavefront produced no microchunk");
3955 crate::pp::record_pp_wave_tick();
3956 Ok((logits, h_seed, hiddens))
3957 }
3958
3959 #[allow(clippy::too_many_arguments)]
3960 fn prime_ppn_wave_worker(
3961 &self,
3962 e: &Engine,
3963 rt: &crate::pp::PpNRt,
3964 waves: &[PrimePpWave<'_>],
3965 cache: &std::sync::Mutex<Cache>,
3966 incoming: Option<&std::sync::mpsc::Receiver<PrimePpSignal>>,
3967 release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
3968 outgoing: &std::sync::mpsc::Sender<PrimePpSignal>,
3969 released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
3970 stage: usize,
3971 seq_end: usize,
3972 fence: &[usize],
3973 initial_base: usize,
3974 ) -> Result<(), Box<dyn std::error::Error>> {
3975 debug_assert_eq!(incoming.is_some(), stage > 0);
3976 debug_assert_eq!(release_upstream.is_some(), stage > 0);
3977 let mut cache = cache
3978 .lock()
3979 .map_err(|_| "prime PP cache stage lock poisoned")?;
3980 let mut credits = PrimePpWaveCredits::default();
3981 for (wave_index, wave) in waves.iter().enumerate() {
3982 let incoming = match incoming {
3983 Some(receiver) => Some(recv_prime_pp_signal(
3984 receiver,
3985 PrimePpWaveSlot {
3986 wave: wave_index,
3987 slot: 0,
3988 },
3989 false,
3990 "prime PP stage input",
3991 )?),
3992 None => None,
3993 };
3994 let sent = self.prime_ppn_wave_stage(
3995 e,
3996 rt,
3997 wave,
3998 &mut cache,
3999 stage,
4000 incoming,
4001 release_upstream,
4002 &mut credits,
4003 released_downstream,
4004 seq_end,
4005 fence,
4006 initial_base,
4007 )?;
4008 send_prime_pp_signal(outgoing, PrimePpSignal::Slot(sent), "prime PP stage output")?;
4009 }
4010 while let Some(expected) = credits.pending.front().copied() {
4011 let released = recv_prime_pp_signal(
4012 released_downstream,
4013 expected,
4014 true,
4015 "prime PP final slot release",
4016 )?;
4017 credits.record_release(released)?;
4018 }
4019 Ok(())
4020 }
4021
4022 #[allow(clippy::too_many_arguments)]
4023 fn prime_ppn_wave_stage(
4024 &self,
4025 e: &Engine,
4026 rt: &crate::pp::PpNRt,
4027 wave: &PrimePpWave<'_>,
4028 cache: &mut Cache,
4029 stage: usize,
4030 incoming: Option<PrimePpWaveSlot>,
4031 release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
4032 credits: &mut PrimePpWaveCredits,
4033 released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
4034 seq_end: usize,
4035 fence: &[usize],
4036 initial_base: usize,
4037 ) -> Result<PrimePpWaveSlot, Box<dyn std::error::Error>> {
4038 debug_assert!(stage + 1 < fence.len() - 1);
4039 let t = wave.end - wave.start;
4040 let base = initial_base + wave.start;
4041 debug_assert_eq!(cache.pos, base, "prime PP stage advanced out of order");
4042 let n_embd = self.cfg.n_embd as usize;
4043 let payload = t * n_embd;
4044 let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
4045 rt.bind_stage(stage)?;
4046 let _stage = rt.enter(stage);
4047 let engine = rt.engine(stage, e);
4048 let positions_d = engine.htod_i32(&positions)?;
4049 let x = if stage == 0 {
4050 debug_assert!(incoming.is_none());
4051 self.embed(engine, wave.tokens)?
4052 } else {
4053 let incoming = incoming.ok_or("prime PP stage has no incoming boundary slot")?;
4054 let x = rt.rx(stage - 1, incoming.slot, payload)?;
4055 send_prime_pp_signal(
4056 release_upstream.ok_or("prime PP stage has no upstream release channel")?,
4057 PrimePpSignal::Slot(incoming),
4058 "prime PP upstream slot release",
4059 )?;
4060 x
4061 };
4062 let x = {
4063 let _wave_cell = crate::pp::enter_pp_wave_cell();
4064 let _overlap = crate::pp::enter_prime_pipe_stage();
4065 self.prime_layers(
4066 engine,
4067 x,
4068 fence[stage],
4069 fence[stage + 1],
4070 &positions_d,
4071 t,
4072 base,
4073 cache,
4074 seq_end,
4075 )?
4076 };
4077 if let Some(expected) = credits.release_required() {
4078 let released =
4079 recv_prime_pp_signal(released_downstream, expected, true, "prime PP slot credit")?;
4080 credits.record_release(released)?;
4081 }
4082 let sent = PrimePpWaveSlot {
4083 wave: credits.next_wave,
4084 slot: rt.tx_pipelined(stage, &x, payload)?,
4085 };
4086 credits.record_send(sent)?;
4087 cache.pos = base + t;
4088 Ok(sent)
4089 }
4090
4091 #[allow(clippy::too_many_arguments)]
4092 #[allow(clippy::type_complexity)] fn prime_ppn_wave_final(
4094 &self,
4095 e: &Engine,
4096 rt: &crate::pp::PpNRt,
4097 wave: &PrimePpWave<'_>,
4098 cache: &mut Cache,
4099 incoming: PrimePpWaveSlot,
4100 release_upstream: &std::sync::mpsc::Sender<PrimePpSignal>,
4101 seq_end: usize,
4102 fence: &[usize],
4103 initial_base: usize,
4104 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4105 let stage = fence.len() - 2;
4106 let t = wave.end - wave.start;
4107 let base = initial_base + wave.start;
4108 debug_assert_eq!(cache.pos, base, "prime PP head stage advanced out of order");
4109 let n_embd = self.cfg.n_embd as usize;
4110 let payload = t * n_embd;
4111 let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
4112 rt.bind_stage(stage)?;
4113 let _stage = rt.enter(stage);
4114 let engine = rt.engine(stage, e);
4115 let positions_d = engine.htod_i32(&positions)?;
4116 let x = rt.rx(stage - 1, incoming.slot, payload)?;
4117 send_prime_pp_signal(
4118 release_upstream,
4119 PrimePpSignal::Slot(incoming),
4120 "prime PP head slot release",
4121 )?;
4122 let _wave_cell = crate::pp::enter_pp_wave_cell();
4123 let _overlap = crate::pp::enter_prime_pipe_stage();
4124 let x = self.prime_layers(
4125 engine,
4126 x,
4127 fence[stage],
4128 fence[stage + 1],
4129 &positions_d,
4130 t,
4131 base,
4132 cache,
4133 seq_end,
4134 )?;
4135 self.prime_chunk_epilogue(engine, x, t, cache)
4136 }
4137
4138 fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
4145 if Engine::gdn_db_on()
4146 && Engine::gdn_chunked_enabled()
4147 && t >= 16
4148 && e.gdn_mma_enabled(Engine::gdn_chunk_size())
4149 && num_k * 2 == num_v
4150 {
4151 num_k
4152 } else {
4153 num_v
4154 }
4155 }
4156
4157 fn f16out_on(e: &Engine, t: usize) -> bool {
4162 crate::f16_ffi::pp_f16_enabled()
4163 && t >= 16
4164 && !e.verify_exact_on()
4165 && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
4166 }
4167
4168 pub fn prime_slabs_get(
4176 &self,
4177 e: &Engine,
4178 t: usize,
4179 n_embd: usize,
4180 n_ff_max: usize,
4181 ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
4182 let mut slabs = self.prime_slabs.lock().unwrap();
4183 let dev = e.ctx().ordinal();
4184 let need_new = match slabs.get(&dev) {
4185 None => true,
4186 Some(sl) => sl.lock().unwrap().t_cap < t,
4187 };
4188 if need_new {
4189 slabs.insert(
4190 dev,
4191 std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
4192 t_cap: t,
4193 h: e.uninit(t * n_embd)?,
4194 x1: e.uninit(t * n_embd)?,
4195 z: e.uninit(t * n_embd)?,
4196 act: e.uninit(t * n_ff_max)?,
4197 xa: e.uninit(t * n_embd)?,
4198 xb: e.uninit(t * n_embd)?,
4199 h16: e.alloc_u8_uninit(t * n_embd * 2)?,
4200 z16: e.alloc_u8_uninit(t * n_embd * 2)?,
4201 gate: e.uninit(t * n_ff_max)?,
4202 up: e.uninit(t * n_ff_max)?,
4203 ffn_out: e.uninit(t * n_embd)?,
4204 seg_glue: Vec::new(),
4205 mixed: e.uninit(t * n_embd)?,
4206 seg_mid: Vec::new(),
4207 seg_t: 0,
4208 })),
4209 );
4210 }
4211 Ok(slabs.get(&dev).expect("prime slab inserted").clone())
4212 }
4213
4214 #[allow(clippy::type_complexity)] fn prime_chunk(
4219 &self,
4220 e: &Engine,
4221 tokens: &[u32],
4222 cache: &mut Cache,
4223 seq_end: usize,
4224 chunk_off: usize,
4225 overlay: Option<&crate::vision::EmbedOverlay>,
4226 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4227 if crate::pp::pp_host_bounce_active()
4228 && (self.uses_gemma_program() || !crate::pp::prime_pp_on())
4229 {
4230 return Err(
4231 "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
4232 has no active prime stage split and would peer-read remote weights; keep \
4233 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
4234 .into(),
4235 );
4236 }
4237 if !self.uses_gemma_program()
4246 && !crate::pp::pp2_streams_off()
4247 && crate::pp::prime_pp_on()
4248 && let Some(fence) = crate::pp::pp_cuts(self.layers.len())
4249 {
4250 if overlay.is_some() {
4251 return Err("vision embedding overlay + PP prime unsupported (v1); \
4252 run single-device or MEMRA_PRIME_PP=0"
4253 .into());
4254 }
4255 return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
4256 }
4257 if crate::pp::pp_host_bounce_active() {
4258 return Err(
4259 "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
4260 refusing an unsplit remote-weight walk"
4261 .into(),
4262 );
4263 }
4264 let t = tokens.len();
4265 let base = cache.pos;
4266 debug_assert!(
4267 seq_end >= base + t,
4268 "prime_chunk: seq_end must cover this chunk"
4269 );
4270 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
4271 let pos_d = e.htod_i32(&pos)?;
4272
4273 let mut x_embed = self.embed(e, tokens)?; if let Some(ov) = overlay {
4275 let n_embd = self.cfg.n_embd as usize;
4279 for &(pos, row_off, n_rows) in &ov.spans {
4280 let lo = pos.max(chunk_off);
4281 let hi = (pos + n_rows).min(chunk_off + t);
4282 if lo < hi {
4283 let src_row = row_off + (lo - pos);
4284 let view = ov
4285 .rows
4286 .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
4287 e.copy_view_into(
4288 &mut x_embed,
4289 (lo - chunk_off) * n_embd,
4290 &view,
4291 (hi - lo) * n_embd,
4292 )?;
4293 }
4294 }
4295 }
4296 let x = self.prime_layers(
4297 e,
4298 x_embed,
4299 0,
4300 self.layers.len(),
4301 &pos_d,
4302 t,
4303 base,
4304 cache,
4305 seq_end,
4306 )?;
4307 self.prime_chunk_epilogue(e, x, t, cache)
4308 }
4309
4310 #[allow(clippy::too_many_arguments)]
4326 fn prime_layers(
4327 &self,
4328 e: &Engine,
4329 x_in: CudaSlice<f32>,
4330 lo: usize,
4331 hi: usize,
4332 pos_d: &CudaSlice<i32>,
4333 t: usize,
4334 base: usize,
4335 cache: &mut Cache,
4336 seq_end: usize,
4337 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4338 let cfg = &self.cfg;
4339 let n_embd = cfg.n_embd as usize;
4340 let eps = cfg.rms_eps;
4341 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
4345 let n_ff_max = self
4352 .layers
4353 .iter()
4354 .map(|l| match &l.ffn {
4355 crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
4356 _ => n_embd,
4357 })
4358 .max()
4359 .unwrap_or(n_embd)
4360 .max(n_embd);
4361 let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
4362 let slab = if use_slabs {
4363 Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
4364 } else {
4365 None
4366 };
4367 let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
4368 let mut x_own; type SlabRefs<'a> = (
4370 &'a mut CudaSlice<f32>,
4371 &'a mut CudaSlice<f32>,
4372 &'a mut CudaSlice<f32>,
4373 &'a mut CudaSlice<f32>,
4374 &'a mut CudaSlice<u8>,
4375 &'a mut CudaSlice<u8>,
4376 &'a mut CudaSlice<f32>,
4377 &'a mut CudaSlice<f32>,
4378 &'a mut CudaSlice<f32>,
4379 );
4380 let (mut x_cur, mut x_nxt, sl): (
4381 &mut CudaSlice<f32>,
4382 &mut CudaSlice<f32>,
4383 Option<SlabRefs>,
4384 );
4385 #[allow(clippy::type_complexity)]
4386 let mut seg: Option<(
4388 &mut Vec<Option<cudarc::driver::CudaGraph>>,
4389 &mut Vec<Option<cudarc::driver::CudaGraph>>,
4390 &mut CudaSlice<f32>,
4391 &mut usize,
4392 )> = None;
4393 let mut x_own2;
4394 match slab_guard.as_mut() {
4395 Some(g) => {
4396 let slabs = &mut **g;
4397 e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
4398 let PrimeSlabs {
4399 xa,
4400 xb,
4401 h,
4402 x1,
4403 z,
4404 act,
4405 h16,
4406 z16,
4407 gate,
4408 up,
4409 ffn_out,
4410 seg_glue,
4411 mixed,
4412 seg_mid,
4413 seg_t,
4414 ..
4415 } = slabs;
4416 x_cur = xa;
4417 x_nxt = xb;
4418 seg = Some((seg_glue, seg_mid, mixed, seg_t));
4419 sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
4420 }
4421 None => {
4422 x_own = x_in;
4423 x_own2 = e.uninit(t * n_embd)?;
4424 x_cur = &mut x_own;
4425 x_nxt = &mut x_own2;
4426 sl = None;
4427 }
4428 }
4429 let mut alloc_h;
4430 let mut alloc_x1;
4431 let mut alloc_z;
4432 let mut alloc_act;
4433 let mut alloc_h16;
4434 let mut alloc_z16;
4435 let mut alloc_gate;
4436 let mut alloc_up;
4437 let mut alloc_fo;
4438 let (h, x1, z, act): (
4439 &mut CudaSlice<f32>,
4440 &mut CudaSlice<f32>,
4441 &mut CudaSlice<f32>,
4442 &mut CudaSlice<f32>,
4443 );
4444 let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
4445 let (sl_gate, sl_up, sl_fo): (
4446 &mut CudaSlice<f32>,
4447 &mut CudaSlice<f32>,
4448 &mut CudaSlice<f32>,
4449 );
4450 match sl {
4451 Some((a, b, c, d, e16, f16b, g, u, fo)) => {
4452 h = a;
4453 x1 = b;
4454 z = c;
4455 act = d;
4456 h16 = e16;
4457 z16 = f16b;
4458 sl_gate = g;
4459 sl_up = u;
4460 sl_fo = fo;
4461 }
4462 None => {
4463 alloc_h = e.uninit(t * n_embd)?;
4464 alloc_x1 = e.uninit(t * n_embd)?;
4465 alloc_z = e.uninit(t * n_embd)?;
4466 alloc_act = e.uninit(t * n_ff_max)?;
4467 alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
4468 alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
4469 alloc_gate = e.uninit(t * n_ff_max)?;
4470 alloc_up = e.uninit(t * n_ff_max)?;
4471 alloc_fo = e.uninit(t * n_embd)?;
4472 h = &mut alloc_h;
4473 x1 = &mut alloc_x1;
4474 z = &mut alloc_z;
4475 act = &mut alloc_act;
4476 h16 = &mut alloc_h16;
4477 z16 = &mut alloc_z16;
4478 sl_gate = &mut alloc_gate;
4479 sl_up = &mut alloc_up;
4480 sl_fo = &mut alloc_fo;
4481 }
4482 }
4483 let n_layers = self.layers.len();
4488 let use_seg = f16fuse
4498 && seg.is_some()
4499 && !self.uses_sliding_gated_moe_program()
4500 && lo == 0
4501 && hi == n_layers
4502 && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1")
4503 && {
4509 let ok = crate::spec::graph_launch_headroom_ok(e);
4510 if !ok {
4511 static NOTED: std::sync::Once = std::sync::Once::new();
4512 NOTED.call_once(|| crate::spec::graph_replay_suspended_note("prime-seg"));
4513 }
4514 ok
4515 };
4516 if let Some((sg, sm, _, st)) = seg.as_mut()
4517 && **st != t
4518 {
4519 sg.clear();
4520 sg.extend((0..n_layers).map(|_| None));
4521 sm.clear();
4522 sm.extend((0..n_layers).map(|_| None));
4523 **st = t;
4524 }
4525 {
4526 let layer_lo = &self.layers[lo];
4527 if f16fuse {
4528 e.rms_norm_f16out(
4529 x_cur,
4530 layer_lo.attn_norm.float_data(),
4531 h,
4532 h16,
4533 n_embd,
4534 t,
4535 eps,
4536 )?;
4537 } else {
4538 e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
4539 }
4540 }
4541 let anat = Self::prime_anatomy_on();
4542 let mut anat_last = if anat {
4543 e.stream().synchronize()?;
4544 Some(std::time::Instant::now())
4545 } else {
4546 None
4547 };
4548 macro_rules! anat_mark {
4550 ($slot:expr) => {
4551 if let Some(ts) = anat_last.as_mut() {
4552 e.stream().synchronize()?;
4553 Self::prime_anatomy_slots()[$slot].fetch_add(
4554 ts.elapsed().as_nanos() as u64,
4555 std::sync::atomic::Ordering::Relaxed,
4556 );
4557 *ts = std::time::Instant::now();
4558 }
4559 };
4560 }
4561 for il in lo..hi {
4562 let layer = &self.layers[il];
4563 let hx16 = if f16fuse { Some(&*h16) } else { None };
4564 if use_seg {
4565 let (pre, pre16, w_out) = match &layer.mixer {
4568 Mixer::Full(fa) => {
4569 let g3 = match hx16 {
4570 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
4571 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
4572 };
4573 let (pre, pre16) =
4574 self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
4575 (pre, pre16, &fa.wo)
4576 }
4577 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("core-split prime"),
4578 Mixer::Kda(_) => {
4579 crate::hybrid::kda_path_unimplemented("core-split captured prime")
4580 }
4581 Mixer::Linear(la) => {
4582 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
4583 let g4 = match hx16 {
4584 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
4585 None => e.matmul_group(&ws, h, t)?,
4586 };
4587 let (pre, pre16) =
4588 self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
4589 (pre, pre16, &la.ssm_out)
4590 }
4591 };
4592 {
4593 let (_, sm, mslab, _) = seg.as_mut().unwrap();
4594 let pre_n = pre.len() / t;
4595 let xh_pre = match pre16 {
4596 Some(x) => x,
4597 None => e.f16_act(&pre, t * pre_n, pre_n)?,
4598 };
4599 if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
4600 let y = e.matmul(w_out, &pre, t)?;
4601 e.copy_into(mslab, 0, &y, t * n_embd)?;
4602 }
4603 if sm[il].is_none() {
4604 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
4605 let w_post = layer.post_attn_norm.float_data();
4606 e.stream().synchronize()?;
4607 e.stream()
4608 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
4609 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
4610 e.add(x_cur, mslab, x1, t * n_embd)?;
4611 e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
4612 Ok(())
4613 })();
4614 let g = e.stream().end_capture(
4615 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
4616 r?;
4617 sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
4618 }
4619 sm[il].as_ref().unwrap().launch()?;
4620 }
4621 } else {
4622 let mixed = match &layer.mixer {
4623 Mixer::Full(fa) => {
4624 let y =
4625 self.full_attn_prime(e, fa, h, hx16, pos_d, t, cache, il, seq_end)?;
4626 anat_mark!(0);
4627 y
4628 }
4629 Mixer::Linear(la) => {
4630 let y = self.linear_attn_prime(e, la, h, hx16, t, cache, il)?;
4631 anat_mark!(1);
4632 y
4633 }
4634 Mixer::Mla(mla) => self.mla_attn_cached(e, mla, h, pos_d, t, il, cache)?,
4635 Mixer::Kda(la) => {
4636 let y = crate::kda::kda_prime_cached(e, la, h, t, eps, cache, il)?;
4637 anat_mark!(1);
4638 y
4639 }
4640 };
4641 if f16fuse {
4642 e.add_rms_norm_f16out(
4645 x_cur,
4646 &mixed,
4647 layer.post_attn_norm.float_data(),
4648 x1,
4649 z,
4650 z16,
4651 n_embd,
4652 t,
4653 eps,
4654 )?;
4655 } else {
4656 e.add(x_cur, &mixed, x1, t * n_embd)?;
4657 e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
4658 }
4659 anat_mark!(4);
4660 }
4661 let zx16 = if f16fuse { Some(&*z16) } else { None };
4662 match &layer.ffn {
4663 crate::hybrid::Ffn::Dense {
4664 ffn_gate,
4665 ffn_up,
4666 ffn_down,
4667 } => {
4668 let n_ff = ffn_gate.out_features();
4669 let mut into_ok = false;
4672 if let Some(xh) = zx16 {
4673 into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
4674 && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
4675 }
4676 if !into_ok {
4677 let mut g2 = match zx16 {
4678 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
4679 None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
4680 };
4681 let up_y = g2.pop().unwrap();
4682 let gate_y = g2.pop().unwrap();
4683 e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
4684 e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
4685 }
4686 let d_lim = self.cfg.clamp_shexp_at(il as u32);
4691 let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
4692 {
4693 let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
4694 e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
4695 Some(a16)
4696 } else {
4697 Self::ffn_act_lim(
4698 e,
4699 &self.cfg,
4700 sl_gate,
4701 sl_up,
4702 1.0,
4703 1.0,
4704 d_lim,
4705 act,
4706 t * n_ff,
4707 )?;
4708 None
4709 };
4710 let xh_act = match act16 {
4712 Some(x) => x,
4713 None => e.f16_act(act, t * n_ff, n_ff)?,
4714 };
4715 if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
4716 let y = e.matmul(ffn_down, &*act, t)?;
4717 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
4718 }
4719 }
4720 crate::hybrid::Ffn::Moe(m) => {
4721 let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
4722 e.copy_into(sl_fo, 0, &y, t * n_embd)?;
4723 anat_mark!(2);
4724 }
4725 }
4726 if let (crate::hybrid::Ffn::Dense { .. }, true) = (&layer.ffn, anat) {
4727 anat_mark!(3);
4728 }
4729 if use_seg && il + 1 < hi {
4730 let w_next = self.layers[il + 1].attn_norm.float_data();
4732 let (sg, _, _, _) = seg.as_mut().unwrap();
4733 if sg[il].is_none() {
4734 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
4735 e.stream().synchronize()?;
4736 e.stream()
4737 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
4738 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
4739 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4740 e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
4741 Ok(())
4742 })();
4743 let g = e.stream().end_capture(
4744 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
4745 );
4746 r?;
4747 sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
4748 }
4749 sg[il].as_ref().unwrap().launch()?;
4750 } else {
4751 if il + 1 < hi {
4752 let w_next = self.layers[il + 1].attn_norm.float_data();
4753 if f16fuse {
4754 e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
4755 } else {
4756 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4757 e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
4758 }
4759 } else {
4760 e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4761 }
4762 }
4763 anat_mark!(4);
4764 if let Some(path) = Self::prime_trace_path() {
4770 let row = base + t - 1;
4771 let host = e.dtoh(x_nxt)?;
4772 let last = &host[(t - 1) * n_embd..t * n_embd];
4773 use std::io::Write as _;
4774 let mut f = std::fs::OpenOptions::new()
4775 .create(true)
4776 .append(true)
4777 .open(path)?;
4778 let mut h64: u64 = 0xcbf29ce484222325;
4779 for v in last {
4780 h64 ^= v.to_bits() as u64;
4781 h64 = h64.wrapping_mul(0x100000001b3);
4782 }
4783 writeln!(
4784 f,
4785 "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
4786 \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
4787 last[0], last[1], last[2]
4788 )?;
4789 }
4790 self.dflash_tap(e, cache, il, x_nxt, t)?;
4793 std::mem::swap(&mut x_cur, &mut x_nxt);
4794 }
4795 if anat {
4796 let s = Self::prime_anatomy_slots();
4797 let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
4798 eprintln!(
4799 "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
4800 dense={:.1} norms_adds={:.1} (t={t}, layers {lo}..{hi})",
4801 ms(0),
4802 ms(1),
4803 ms(2),
4804 ms(3),
4805 ms(4)
4806 );
4807 }
4808 let mut x = e.uninit(t * n_embd)?;
4810 e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
4811 drop(slab_guard);
4812 Ok(x)
4813 }
4814
4815 #[allow(clippy::type_complexity)] fn prime_chunk_epilogue(
4821 &self,
4822 e: &Engine,
4823 x: CudaSlice<f32>,
4824 t: usize,
4825 cache: &mut Cache,
4826 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4827 let n_embd = self.cfg.n_embd as usize;
4828 let eps = self.cfg.rms_eps;
4829 let mut h_seed = e.uninit(n_embd)?;
4833 if !crate::spec::spec_hpost() {
4834 e.copy_view_into(
4835 &mut h_seed,
4836 0,
4837 &x.slice((t - 1) * n_embd..t * n_embd),
4838 n_embd,
4839 )?;
4840 }
4841 let mut hn = e.uninit(t * n_embd)?;
4843 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4844 if crate::spec::spec_hpost() {
4845 e.copy_view_into(
4846 &mut h_seed,
4847 0,
4848 &hn.slice((t - 1) * n_embd..t * n_embd),
4849 n_embd,
4850 )?;
4851 }
4852 let last = e.view(&hn, t * n_embd);
4853 let last_row = last.slice((t - 1) * n_embd..t * n_embd);
4854 let mut hlast = e.uninit(n_embd)?;
4855 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
4856 let logits = e.matmul(&self.output, &hlast, 1)?;
4857 cache.pos += t;
4858 Ok((
4861 e.dtoh(&logits)?,
4862 h_seed,
4863 if crate::spec::spec_hpost() { hn } else { x },
4864 ))
4865 }
4866
4867 pub fn hidden_postnorm_row(
4873 &self,
4874 e: &Engine,
4875 hiddens: &CudaSlice<f32>,
4876 row: usize,
4877 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4878 let n_embd = self.cfg.n_embd as usize;
4879 let mut x1 = e.uninit(n_embd)?;
4880 e.copy_view_into(
4881 &mut x1,
4882 0,
4883 &hiddens.slice(row * n_embd..(row + 1) * n_embd),
4884 n_embd,
4885 )?;
4886 if crate::spec::spec_hpost() {
4887 return e.dtoh(&x1);
4888 }
4889 let mut hn = e.uninit(n_embd)?;
4890 e.rms_norm(
4891 &x1,
4892 self.output_norm.float_data(),
4893 &mut hn,
4894 n_embd,
4895 1,
4896 self.cfg.rms_eps,
4897 )?;
4898 e.dtoh(&hn)
4899 }
4900
4901 #[allow(clippy::type_complexity)] fn prime_chunk_ppn(
4926 &self,
4927 e: &Engine,
4928 tokens: &[u32],
4929 cache: &mut Cache,
4930 seq_end: usize,
4931 fence: &[usize],
4932 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4933 let rt = crate::pp::PpNRt::get(e)?;
4934 let n_st = fence.len() - 1;
4935 assert_eq!(
4936 rt.n_stages(),
4937 n_st,
4938 "PpNRt stage count {} != fence stages {n_st}",
4939 rt.n_stages()
4940 );
4941 let n_embd = self.cfg.n_embd as usize;
4942 let t = tokens.len();
4943 let base = cache.pos;
4944 debug_assert!(
4945 seq_end >= base + t,
4946 "prime_chunk_ppn: seq_end must cover this chunk"
4947 );
4948 let payload = t * n_embd;
4949 let caller_stream = e.stream();
4953 rt.fence_stages_behind(&caller_stream)?;
4954
4955 if n_st == 2 {
4956 let slot =
4957 self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
4958 let x =
4959 self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
4960 let out = {
4961 rt.bind_stage(1)?;
4962 let _st1 = rt.enter(1);
4963 let e1 = rt.engine(1, e);
4964 self.prime_chunk_epilogue(e1, x, t, cache)?
4965 };
4966 rt.publish_to(1, &caller_stream)?;
4967 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4968 return Ok(out);
4969 }
4970
4971 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
4972
4973 let mut slot = {
4975 let _st0 = rt.enter(0);
4976 let e0 = rt.engine(0, e);
4977 let pos_d = e0.htod_i32(&pos)?;
4978 let x = self.embed(e0, tokens)?;
4979 let x =
4980 self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
4981 rt.tx(0, &x, payload)?
4982 };
4984
4985 for s in 1..n_st - 1 {
4987 let _st = rt.enter(s);
4988 let es = rt.engine(s, e);
4989 let pos_d = es.htod_i32(&pos)?;
4990 let x = rt.rx(s - 1, slot, payload)?;
4991 let x = self.prime_layers(
4992 es,
4993 x,
4994 fence[s],
4995 fence[s + 1],
4996 &pos_d,
4997 t,
4998 base,
4999 cache,
5000 seq_end,
5001 )?;
5002 slot = rt.tx(s, &x, payload)?;
5003 }
5004
5005 let _stl = rt.enter(n_st - 1);
5007 let el = rt.engine(n_st - 1, e);
5008 let pos_d = el.htod_i32(&pos)?;
5009 let x = rt.rx(n_st - 2, slot, payload)?;
5010 let x = self.prime_layers(
5011 el,
5012 x,
5013 fence[n_st - 1],
5014 fence[n_st],
5015 &pos_d,
5016 t,
5017 base,
5018 cache,
5019 seq_end,
5020 )?;
5021 let out = self.prime_chunk_epilogue(el, x, t, cache)?;
5022 rt.publish_to(n_st - 1, &caller_stream)?;
5028 crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5029 Ok(out)
5030 }
5031
5032 #[allow(clippy::too_many_arguments)] fn prime_pp2_stage0_enqueue(
5034 &self,
5035 e: &Engine,
5036 rt: &crate::pp::PpNRt,
5037 tokens: &[u32],
5038 cache: &mut Cache,
5039 seq_end: usize,
5040 fence: &[usize],
5041 base: usize,
5042 pipelined: bool,
5043 ) -> Result<usize, Box<dyn std::error::Error>> {
5044 let t = tokens.len();
5045 let n_embd = self.cfg.n_embd as usize;
5046 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5047 rt.bind_stage(0)?;
5048 let _st0 = rt.enter(0);
5049 let e0 = rt.engine(0, e);
5050 let pos_d = e0.htod_i32(&pos)?;
5051 let x = self.embed(e0, tokens)?;
5052 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
5053 let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
5054 if pipelined {
5055 rt.tx_pipelined(0, &x, t * n_embd)
5056 } else {
5057 rt.tx(0, &x, t * n_embd)
5058 }
5059 }
5060
5061 #[allow(clippy::too_many_arguments)] fn prime_pp2_stage1_enqueue(
5063 &self,
5064 e: &Engine,
5065 rt: &crate::pp::PpNRt,
5066 slot: usize,
5067 t: usize,
5068 cache: &mut Cache,
5069 seq_end: usize,
5070 fence: &[usize],
5071 base: usize,
5072 pipelined: bool,
5073 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5074 let n_embd = self.cfg.n_embd as usize;
5075 let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5076 rt.bind_stage(1)?;
5077 let _st1 = rt.enter(1);
5078 let e1 = rt.engine(1, e);
5079 let pos_d = e1.htod_i32(&pos)?;
5080 let x = rt.rx(0, slot, t * n_embd)?;
5081 let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
5082 self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
5083 }
5084
5085 #[allow(clippy::too_many_arguments)] pub fn prime_chunk_captured(
5102 &self,
5103 e: &Engine,
5104 x_in: &CudaSlice<f32>,
5105 pos_d: &CudaSlice<i32>,
5106 t: usize,
5107 cache: &mut Cache,
5108 len_d: &CudaSlice<i32>,
5109 logits_out: &mut CudaSlice<f32>,
5110 h_seed_out: &mut CudaSlice<f32>,
5111 ) -> Result<(), Box<dyn std::error::Error>> {
5112 self.refuse_hyper("prime_chunk_captured")?;
5113 cache.ensure_usable("prime_chunk_captured")?;
5114 let cfg = &self.cfg;
5115 let n_embd = cfg.n_embd as usize;
5116 let eps = cfg.rms_eps;
5117 let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
5118 let mut x = e.uninit(t * n_embd)?;
5119 e.copy_into(&mut x, 0, x_in, t * n_embd)?;
5120 for (il, layer) in self.layers.iter().enumerate() {
5121 let mut h = e.uninit(t * n_embd)?;
5122 let mut hx16: Option<CudaSlice<u8>> = None;
5123 if f16fuse {
5124 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5125 e.rms_norm_f16out(
5126 &x,
5127 layer.attn_norm.float_data(),
5128 &mut h,
5129 &mut b16,
5130 n_embd,
5131 t,
5132 eps,
5133 )?;
5134 hx16 = Some(b16);
5135 } else {
5136 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5137 }
5138 let mixed = match &layer.mixer {
5139 Mixer::Full(fa) => {
5143 self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
5144 }
5145 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("captured-graph prime"),
5146 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("captured prime chunk"),
5147 Mixer::Linear(la) => {
5148 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
5149 let g4 = match hx16.as_ref() {
5150 Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
5151 None => e.matmul_group(&ws, &h, t)?,
5152 };
5153 self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
5154 }
5155 };
5156 let mut x1 = e.uninit(t * n_embd)?;
5157 e.add(&x, &mixed, &mut x1, t * n_embd)?;
5158 let mut z = e.uninit(t * n_embd)?;
5159 let mut zx16: Option<CudaSlice<u8>> = None;
5160 if f16fuse {
5161 let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5162 e.rms_norm_f16out(
5163 &x1,
5164 layer.post_attn_norm.float_data(),
5165 &mut z,
5166 &mut b16,
5167 n_embd,
5168 t,
5169 eps,
5170 )?;
5171 zx16 = Some(b16);
5172 } else {
5173 e.rms_norm(
5174 &x1,
5175 layer.post_attn_norm.float_data(),
5176 &mut z,
5177 n_embd,
5178 t,
5179 eps,
5180 )?;
5181 }
5182 let ffn_out = match &layer.ffn {
5183 crate::hybrid::Ffn::Dense {
5184 ffn_gate,
5185 ffn_up,
5186 ffn_down,
5187 } => {
5188 let n_ff = ffn_gate.out_features();
5189 let mut g2 = match &zx16 {
5190 Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
5191 None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
5192 };
5193 let up = g2.pop().unwrap();
5194 let gate = g2.pop().unwrap();
5195 let mut act = e.uninit(t * n_ff)?;
5196 Self::ffn_act_lim(
5198 e,
5199 &self.cfg,
5200 &gate,
5201 &up,
5202 1.0,
5203 1.0,
5204 self.cfg.clamp_shexp_at(il as u32),
5205 &mut act,
5206 t * n_ff,
5207 )?;
5208 e.matmul(ffn_down, &act, t)?
5209 }
5210 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
5211 };
5212 let mut x2 = e.uninit(t * n_embd)?;
5213 e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5214 x = x2;
5215 }
5216 if !crate::spec::spec_hpost() {
5218 e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
5219 }
5220 let mut hn = e.uninit(t * n_embd)?;
5221 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5222 if crate::spec::spec_hpost() {
5223 e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
5224 }
5225 let mut hlast = e.uninit(n_embd)?;
5226 e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
5227 let logits = e.matmul(&self.output, &hlast, 1)?;
5228 let nv = logits.len();
5229 e.copy_into(logits_out, 0, &logits, nv)?;
5230 Ok(())
5231 }
5232
5233 fn step35_prime_batch_on() -> bool {
5234 std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
5235 }
5236
5237 #[allow(clippy::too_many_arguments)]
5240 #[allow(clippy::too_many_arguments)]
5245 fn step35_prime_batch_layers(
5246 &self,
5247 e: &Engine,
5248 mut x: CudaSlice<f32>,
5249 lo: usize,
5250 hi: usize,
5251 ts: &[usize],
5252 offs: &[usize],
5253 seq_ends: &[usize],
5254 pos_ds: &[CudaSlice<i32>],
5255 caches: &mut [&mut Cache],
5256 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5257 let cfg = &self.cfg;
5258 let n_embd = cfg.n_embd as usize;
5259 let eps = cfg.rms_eps;
5260 let b = ts.len();
5261 let total: usize = ts.iter().sum();
5262 let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
5263
5264 let split = |e: &Engine,
5265 y: &CudaSlice<f32>,
5266 dim: usize|
5267 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
5268 let mut out = Vec::with_capacity(b);
5269 for s in 0..b {
5270 let mut ys = e.uninit(ts[s] * dim)?;
5271 e.copy_view_into(
5272 &mut ys,
5273 0,
5274 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
5275 ts[s] * dim,
5276 )?;
5277 out.push(ys);
5278 }
5279 Ok(out)
5280 };
5281
5282 let prof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
5287 let mut ph = [0f64; 4]; let mark = |e: &Engine, acc: usize, t0: &mut std::time::Instant, ph: &mut [f64; 4]| {
5289 if prof {
5290 let _ = e.stream().synchronize();
5291 ph[acc] += t0.elapsed().as_secs_f64() * 1e3;
5292 *t0 = std::time::Instant::now();
5293 }
5294 };
5295 let mut pt = std::time::Instant::now();
5296 for il in lo..hi {
5297 let layer = &self.layers[il];
5298 let Mixer::Full(fa) = &layer.mixer else {
5299 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
5300 };
5301
5302 let mut h = e.uninit(total * n_embd)?;
5303 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5304 if f16fuse {
5305 e.rms_norm_f16out(
5306 &x,
5307 layer.attn_norm.float_data(),
5308 &mut h,
5309 &mut hx16,
5310 n_embd,
5311 total,
5312 eps,
5313 )?;
5314 } else {
5315 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
5316 }
5317
5318 let gate_w = fa
5322 .attn_gate
5323 .as_ref()
5324 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
5325 let mut g4 = if f16fuse {
5326 e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
5327 } else {
5328 e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
5329 };
5330 let gate = g4.pop().unwrap();
5331 let mut parts: Vec<Vec<CudaSlice<f32>>> =
5332 (0..b).map(|_| Vec::with_capacity(3)).collect();
5333 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
5334 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
5335 parts[s].push(ys);
5336 }
5337 }
5338 let gates = split(e, &gate, gate_w.out_features())?;
5339 let geometry = self.step35_geom(il);
5340 let hd = geometry.head_dim_k as usize;
5341 let nh = geometry.n_head as usize;
5342 let mut ag_cat = e.uninit(total * nh * hd)?;
5343 for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
5344 mark(e, 0, &mut pt, &mut ph);
5345 let ag = self.step35_attn_pre_wo(
5346 e,
5347 fa,
5348 g3s,
5349 None,
5350 Some(&gate),
5351 &pos_ds[s],
5352 ts[s],
5353 Some(&mut *caches[s]),
5354 il,
5355 seq_ends[s],
5356 )?;
5357 e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
5358 }
5359 mark(e, 1, &mut pt, &mut ph);
5360 let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
5361
5362 let mut x1 = e.uninit(total * n_embd)?;
5363 let mut z = e.uninit(total * n_embd)?;
5364 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5365 if f16fuse {
5366 e.add_rms_norm_f16out(
5367 &x,
5368 &mixed,
5369 layer.post_attn_norm.float_data(),
5370 &mut x1,
5371 &mut z,
5372 &mut zx16,
5373 n_embd,
5374 total,
5375 eps,
5376 )?;
5377 } else {
5378 e.add(&x, &mixed, &mut x1, total * n_embd)?;
5379 e.rms_norm(
5380 &x1,
5381 layer.post_attn_norm.float_data(),
5382 &mut z,
5383 n_embd,
5384 total,
5385 eps,
5386 )?;
5387 }
5388
5389 mark(e, 2, &mut pt, &mut ph);
5390 let ffn_out = match &layer.ffn {
5391 crate::hybrid::Ffn::Dense {
5392 ffn_gate,
5393 ffn_up,
5394 ffn_down,
5395 } => {
5396 let n_ff = ffn_gate.out_features();
5397 let mut g2 = if f16fuse {
5398 e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
5399 } else {
5400 e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
5401 };
5402 let up = g2.pop().unwrap();
5403 let gate = g2.pop().unwrap();
5404 let mut act = e.uninit(total * n_ff)?;
5405 let d_lim = cfg.clamp_shexp_at(il as u32);
5406 if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
5407 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
5408 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
5409 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
5410 Some(y) => y,
5411 None => e.matmul(ffn_down, &act, total)?,
5412 }
5413 } else {
5414 Self::ffn_act_lim(
5415 e,
5416 cfg,
5417 &gate,
5418 &up,
5419 1.0,
5420 1.0,
5421 d_lim,
5422 &mut act,
5423 total * n_ff,
5424 )?;
5425 e.matmul(ffn_down, &act, total)?
5426 }
5427 }
5428 crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
5429 };
5430 let mut x2 = e.uninit(total * n_embd)?;
5431 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
5432 x = x2;
5433 mark(e, 3, &mut pt, &mut ph);
5434 }
5435 if prof {
5436 eprintln!(
5437 "[prime-prof] t={total} layers={} norm+qkv={:.0}ms attn={:.0}ms o_proj={:.0}ms moe={:.0}ms",
5438 hi - lo,
5439 ph[0],
5440 ph[1],
5441 ph[2],
5442 ph[3]
5443 );
5444 }
5445 Ok(x)
5446 }
5447
5448 #[allow(clippy::type_complexity)] fn step35_prime_batch_epilogue(
5450 &self,
5451 e: &Engine,
5452 x: CudaSlice<f32>,
5453 ts: &[usize],
5454 offs: &[usize],
5455 caches: &mut [&mut Cache],
5456 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5457 let n_embd = self.cfg.n_embd as usize;
5458 let total: usize = ts.iter().sum();
5459 let mut hn = e.uninit(total * n_embd)?;
5460 e.rms_norm(
5461 &x,
5462 self.output_norm.float_data(),
5463 &mut hn,
5464 n_embd,
5465 total,
5466 self.cfg.rms_eps,
5467 )?;
5468
5469 let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
5470 let mut out = Vec::with_capacity(ts.len());
5471 for s in 0..ts.len() {
5472 let mut hidden = e.uninit(ts[s] * n_embd)?;
5473 e.copy_view_into(
5474 &mut hidden,
5475 0,
5476 &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
5477 ts[s] * n_embd,
5478 )?;
5479 let last0 = (offs[s] + ts[s] - 1) * n_embd;
5480 let mut h_seed = e.uninit(n_embd)?;
5481 e.copy_view_into(
5482 &mut h_seed,
5483 0,
5484 &hidden_src.slice(last0..last0 + n_embd),
5485 n_embd,
5486 )?;
5487 let mut hlast = e.uninit(n_embd)?;
5489 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
5490 let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
5491 caches[s].pos += ts[s];
5492 out.push((logits, h_seed, hidden));
5493 }
5494 Ok(out)
5495 }
5496
5497 #[allow(clippy::type_complexity)] fn step35_prime_cache_batch(
5503 &self,
5504 e: &Engine,
5505 prompts: &[&[u32]],
5506 caches: &mut [&mut Cache],
5507 seq_ends: &[usize],
5508 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5509 assert_eq!(
5510 seq_ends.len(),
5511 prompts.len(),
5512 "step35 batched prime: one seq_end per sequence"
5513 );
5514 validate_step_prime_batch_modes(
5515 step_tp_prefill_enabled()?,
5516 step_ep_grouped_prefill_enabled()?,
5517 )?;
5518 if crate::pp::pp_host_bounce_active()
5519 && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
5520 {
5521 return Err(
5522 "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
5523 stage split; refusing an unsplit remote-weight walk"
5524 .into(),
5525 );
5526 }
5527 if !Self::step35_prime_batch_on() {
5528 return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
5529 }
5530 if prompts.len() > 1 && caches.iter().any(|c| c.pos != 0) {
5535 return Err(
5536 "step35 batched prime supports continuation only at B=1; a cross-request batch \
5537 at mixed positions requires per-request queued_after"
5538 .into(),
5539 );
5540 }
5541
5542 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
5543 for &t in &ts {
5544 assert!(
5545 t >= PRIME_MIN_T,
5546 "step35 batched prime needs T >= {PRIME_MIN_T}"
5547 );
5548 }
5549 for (s, c) in caches.iter().enumerate() {
5550 assert!(
5553 c.pos + ts[s] <= c.max_ctx,
5554 "step35 batched prime exceeds cache max_ctx"
5555 );
5556 assert!(
5557 seq_ends[s] >= c.pos + ts[s],
5558 "step35 batched prime: seq_end must cover this chunk"
5559 );
5560 }
5561 let mut transaction = CacheTaintGuard::arm(caches);
5562 let legacy_tsend = std::env::var("MEMRA_STEP35_PRIME_BATCH_TSEND").as_deref() == Ok("1");
5569 let seq_ends_eff: Vec<usize> = if legacy_tsend {
5570 ts.clone()
5571 } else {
5572 seq_ends.to_vec()
5573 };
5574 let offs: Vec<usize> = ts
5575 .iter()
5576 .scan(0usize, |a, &t| {
5577 let o = *a;
5578 *a += t;
5579 Some(o)
5580 })
5581 .collect();
5582 let total: usize = ts.iter().sum();
5583 let payload = total * self.cfg.n_embd as usize;
5584 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
5585 let positions: Vec<Vec<i32>> = ts
5591 .iter()
5592 .zip(caches.iter())
5593 .map(|(&t, c)| {
5594 let base = c.pos as i32;
5595 (0..t as i32).map(|i| base + i).collect()
5596 })
5597 .collect();
5598 let upload_positions =
5599 |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
5600 positions
5601 .iter()
5602 .map(|p| e.htod_i32(p))
5603 .collect::<Result<_, _>>()
5604 };
5605
5606 static ONCE: std::sync::Once = std::sync::Once::new();
5607 ONCE.call_once(|| {
5608 eprintln!(
5609 "[step35-prime-batch] first concat prime: B={} tokens={total}",
5610 prompts.len()
5611 );
5612 });
5613
5614 let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
5615 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5616 let rt = crate::pp::PpNRt::get(e)?;
5617 let n_st = fence.len() - 1;
5618 assert_eq!(
5619 rt.n_stages(),
5620 n_st,
5621 "step35 prime batch stage count mismatch"
5622 );
5623 let caller_stream = e.stream();
5624 rt.fence_stages_behind(&caller_stream)?;
5625
5626 let mut slot = {
5627 let _st0 = rt.enter(0);
5628 let e0 = rt.engine(0, e);
5629 let pos_ds = upload_positions(e0)?;
5630 let x = self.embed(e0, &cat_tokens)?;
5631 let x = self.step35_prime_batch_layers(
5632 e0,
5633 x,
5634 fence[0],
5635 fence[1],
5636 &ts,
5637 &offs,
5638 &seq_ends_eff,
5639 &pos_ds,
5640 caches,
5641 )?;
5642 rt.tx(0, &x, payload)?
5643 };
5644 for s in 1..n_st - 1 {
5645 let _st = rt.enter(s);
5646 let es = rt.engine(s, e);
5647 let pos_ds = upload_positions(es)?;
5648 let x = rt.rx(s - 1, slot, payload)?;
5649 let x = self.step35_prime_batch_layers(
5650 es,
5651 x,
5652 fence[s],
5653 fence[s + 1],
5654 &ts,
5655 &offs,
5656 &seq_ends_eff,
5657 &pos_ds,
5658 caches,
5659 )?;
5660 slot = rt.tx(s, &x, payload)?;
5661 }
5662
5663 let _stl = rt.enter(n_st - 1);
5664 let el = rt.engine(n_st - 1, e);
5665 let pos_ds = upload_positions(el)?;
5666 let x = rt.rx(n_st - 2, slot, payload)?;
5667 let x = self.step35_prime_batch_layers(
5668 el,
5669 x,
5670 fence[n_st - 1],
5671 fence[n_st],
5672 &ts,
5673 &offs,
5674 &seq_ends_eff,
5675 &pos_ds,
5676 caches,
5677 )?;
5678 let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
5679 rt.publish_to(n_st - 1, &caller_stream)?;
5680 crate::pp::STEP35_PRIME_BATCH_SPLITS
5681 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5682 out
5683 } else {
5684 let pos_ds = upload_positions(e)?;
5685 let x = self.embed(e, &cat_tokens)?;
5686 let x = self.step35_prime_batch_layers(
5687 e,
5688 x,
5689 0,
5690 self.layers.len(),
5691 &ts,
5692 &offs,
5693 &seq_ends_eff,
5694 &pos_ds,
5695 caches,
5696 )?;
5697 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
5698 }
5699 } else {
5700 let pos_ds = upload_positions(e)?;
5701 let x = self.embed(e, &cat_tokens)?;
5702 let x = self.step35_prime_batch_layers(
5703 e,
5704 x,
5705 0,
5706 self.layers.len(),
5707 &ts,
5708 &offs,
5709 &seq_ends_eff,
5710 &pos_ds,
5711 caches,
5712 )?;
5713 self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
5714 };
5715 crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5716 transaction.commit();
5717 Ok(out)
5718 }
5719
5720 #[allow(clippy::type_complexity)] pub fn prime_cache_batch(
5738 &self,
5739 e: &Engine,
5740 prompts: &[&[u32]],
5741 caches: &mut [&mut Cache],
5742 ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5743 self.refuse_hyper("prime_cache_batch")?;
5744 for cache in caches.iter() {
5745 cache.ensure_usable("prime_cache_batch")?;
5746 }
5747 if crate::pp::pp_cuts(self.layers.len()).is_some()
5748 && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
5749 {
5750 return Err("pipeline rewrite is not qualified for batched prime".into());
5751 }
5752 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::CarriedPrime) {
5753 if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeEager) {
5754 return Err("neither batched-prime nor eager rewrite is qualified".into());
5755 }
5756 if prompts.len() != caches.len() {
5757 return Err("prime fallback prompt/cache shape mismatch".into());
5758 }
5759 static ONCE: std::sync::Once = std::sync::Once::new();
5760 ONCE.call_once(|| {
5761 eprintln!(
5762 "[rewrite] carried-prime.v1 unqualified; using individual native eager primes"
5763 );
5764 });
5765 let mut transaction = CacheTaintGuard::arm(caches);
5766 let result: Result<Vec<_>, Box<dyn std::error::Error>> = prompts
5767 .iter()
5768 .copied()
5769 .zip(caches.iter_mut())
5770 .map(|(prompt, cache)| self.prime_cache(e, prompt, cache, 0))
5771 .collect();
5772 if result.is_ok() {
5773 transaction.commit();
5774 }
5775 return result;
5776 }
5777 let _pp_walk =
5778 if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
5779 let rt = crate::pp::PpNRt::get(e)?;
5780 Some(rt.acquire_walk("prime_cache_batch")?)
5781 } else {
5782 None
5783 };
5784 let cfg = &self.cfg;
5785 let n_embd = cfg.n_embd as usize;
5786 let eps = cfg.rms_eps;
5787 let b = prompts.len();
5788 assert!(b >= 1 && b == caches.len());
5789 let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
5790 let carried = pos0s.iter().any(|&p| p > 0);
5791 if self.uses_gemma_program() {
5797 return Err(
5798 "prime_cache_batch: gemma4 has no batched prime core (per-layer \
5799 swa/global geometry, softcapped head) — use gemma4_prime per sequence"
5800 .into(),
5801 );
5802 }
5803 if self.uses_sliding_gated_moe_program() {
5806 let seq_ends: Vec<usize> = caches
5811 .iter()
5812 .zip(prompts.iter())
5813 .map(|(c, p)| c.pos + p.len())
5814 .collect();
5815 return self.step35_prime_cache_batch(e, prompts, caches, &seq_ends);
5816 }
5817 if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
5818 let rt = crate::pp::PpNRt::get(e)?;
5819 if rt.cross_device() {
5820 return Err(
5821 "prime_cache_batch: generic dense concat prime has no cross-device PP split; use individual prime_cache calls"
5822 .into(),
5823 );
5824 }
5825 }
5826 let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
5827 for &t in &ts {
5828 assert!(
5829 t >= PRIME_MIN_T,
5830 "prime_cache_batch needs T >= {PRIME_MIN_T}"
5831 );
5832 }
5833 for (s, c) in caches.iter().enumerate() {
5834 assert!(
5835 c.pos + ts[s] <= c.max_ctx,
5836 "prime_cache_batch: prompt exceeds cache max_ctx"
5837 );
5838 }
5839 let mut transaction = CacheTaintGuard::arm(caches);
5840 let total: usize = ts.iter().sum();
5841 let offs: Vec<usize> = ts
5842 .iter()
5843 .scan(0usize, |a, &t| {
5844 let o = *a;
5845 *a += t;
5846 Some(o)
5847 })
5848 .collect();
5849 let pos_ds: Vec<CudaSlice<i32>> = ts
5851 .iter()
5852 .zip(&pos0s)
5853 .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
5854 .collect::<Result<_, _>>()?;
5855 let split = |e: &Engine,
5857 y: &CudaSlice<f32>,
5858 dim: usize|
5859 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
5860 let mut out = Vec::with_capacity(b);
5861 for s in 0..b {
5862 let mut ys = e.uninit(ts[s] * dim)?;
5863 e.copy_view_into(
5864 &mut ys,
5865 0,
5866 &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
5867 ts[s] * dim,
5868 )?;
5869 out.push(ys);
5870 }
5871 Ok(out)
5872 };
5873
5874 let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
5875 let mut x = self.embed(e, &cat_tokens)?; for (il, layer) in self.layers.iter().enumerate() {
5877 let mut h = e.uninit(total * n_embd)?;
5878 let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5879 e.rms_norm_f16out(
5880 &x,
5881 layer.attn_norm.float_data(),
5882 &mut h,
5883 &mut hx16,
5884 n_embd,
5885 total,
5886 eps,
5887 )?;
5888 let mut mixed = e.uninit(total * n_embd)?;
5890 match &layer.mixer {
5891 Mixer::Full(fa) => {
5892 let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
5893 let geometry = self.cfg.full_attention_geometry_at(il as u32);
5899 let (n_head, n_head_kv, head_dim) = (
5900 geometry.n_head as usize,
5901 geometry.n_head_kv as usize,
5902 geometry.head_dim_k as usize,
5903 );
5904 let fa_scale = geometry.attention_scale();
5905 let use_favl = !carried
5906 && (2..=8).contains(&b)
5907 && (head_dim == 256 || head_dim == 128)
5908 && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
5909 && std::env::var("MEMRA_NOFA").is_err()
5910 && std::env::var("MEMRA_FA_FLOOR").is_err()
5911 && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
5912 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
5913 && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
5914 if use_favl {
5915 let (qf_w, kf_w, vf_w) = (
5916 fa.wq.out_features(),
5917 fa.wk.out_features(),
5918 fa.wv.out_features(),
5919 );
5920 memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
5925 struct APre {
5926 q: CudaSlice<f32>,
5927 gate: Option<CudaSlice<f32>>,
5928 qn: CudaSlice<f32>,
5929 kn: CudaSlice<f32>,
5930 }
5931 let mut aps = Vec::with_capacity(b);
5932 for &t in ts.iter().take(b) {
5933 aps.push(APre {
5934 q: e.uninit(t * n_head * head_dim)?,
5935 gate: Some(e.uninit(t * n_head * head_dim)?),
5936 qn: e.uninit(t * n_head * head_dim)?,
5937 kn: e.uninit(t * n_head_kv * head_dim)?,
5938 });
5939 }
5940 let (kv_dim_k, kv_dim_v, ktb, vtb) = {
5941 let kvl = caches[0].kv[il].as_ref().unwrap();
5942 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
5943 };
5944 let pargs: Vec<crate::AttnPreVl> = (0..b)
5945 .map(|s| {
5946 let (o, t) = (offs[s], ts[s]);
5947 let kvl = caches[s].kv[il].as_ref().unwrap();
5948 assert!(
5949 kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
5950 "prime_cache_batch attn vl: fresh + capacity"
5951 );
5952 crate::AttnPreVl {
5953 qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
5954 kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
5955 vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
5956 q: e.addr_f32(&aps[s].q),
5957 gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
5958 qn: e.addr_f32(&aps[s].qn),
5959 kn: e.addr_f32(&aps[s].kn),
5960 kc: e.addr_u8(&kvl.k),
5961 vc: e.addr_u8(&kvl.v),
5962 t: t as i32,
5963 pad: 0,
5964 }
5965 })
5966 .collect();
5967 e.attn_pre_vl8(
5968 &pargs,
5969 fa.q_norm.float_data(),
5970 fa.k_norm.float_data(),
5971 head_dim,
5972 geometry.n_rot as usize,
5973 n_head,
5974 n_head_kv,
5975 self.cfg.rms_eps,
5976 geometry.rope_base,
5977 1.0,
5978 kv_dim_k,
5979 kv_dim_v,
5980 ktb,
5981 vtb,
5982 )?;
5983 for s in 0..b {
5984 let kvl = caches[s].kv[il].as_mut().unwrap();
5985 kvl.len += ts[s];
5986 let new_len = kvl.len as i32;
5987 e.set_i32_one(&mut kvl.len_d, new_len)?;
5988 }
5989 let mut attns = Vec::with_capacity(b);
5990 let mut mirrors = Vec::with_capacity(b);
5991 for &t in ts.iter().take(b) {
5992 attns.push(e.uninit(t * n_head * head_dim)?);
5993 let n = t * n_head_kv * head_dim;
5994 mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
5995 }
5996 let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
5999 Ok("0") => false,
6000 Ok("1") => {
6004 crate::refuse_portable_force(
6005 "MEMRA_FA3=1",
6006 "the sm_90a fa3/bf16 kernels",
6007 );
6008 true
6009 }
6010 _ => cfg!(memra_hopper_mma),
6011 };
6012 if fa3_on {
6013 let mut q16s = Vec::with_capacity(b);
6014 let mut v16s = Vec::with_capacity(b);
6015 for s in 0..b {
6016 let t = ts[s];
6017 let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
6018 e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
6019 let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
6020 e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
6021 let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
6022 e.f32_to_bf16_v(
6023 &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
6024 &mut v16,
6025 t * n_head_kv * head_dim,
6026 )?;
6027 q16s.push(q16);
6028 v16s.push((k16, v16));
6029 }
6030 let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
6031 let mut kp = qp;
6032 let mut vp = qp;
6033 let mut op = [core::ptr::null_mut::<f32>(); 8];
6034 let mut tsv = [0i32; 8];
6035 for s in 0..b {
6036 qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
6037 kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
6038 vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
6039 op[s] = e.addr_f32(&attns[s]) as *mut f32;
6040 tsv[s] = ts[s] as i32;
6041 }
6042 let rc = unsafe {
6043 crate::fa3_vl_raw(
6044 qp.as_ptr(),
6045 kp.as_ptr(),
6046 vp.as_ptr(),
6047 op.as_ptr(),
6048 tsv.as_ptr(),
6049 b as i32,
6050 n_head as i32,
6051 n_head_kv as i32,
6052 head_dim as i32,
6053 fa_scale,
6054 e.stream().cu_stream() as *mut core::ffi::c_void,
6055 )
6056 };
6057 if rc != 0 {
6058 return Err(format!("memra_fa3_vl rc={rc}").into());
6059 }
6060 } else {
6061 let fargs: Vec<crate::FaSeqVl> = (0..b)
6062 .map(|s| crate::FaSeqVl {
6063 q: e.addr_f32(&aps[s].qn),
6064 k16: e.addr_u8(&mirrors[s].0),
6065 v16: e.addr_u8(&mirrors[s].1),
6066 o: e.addr_f32(&attns[s]),
6067 kf: e.addr_f32(&aps[s].kn),
6068 vf: e.addr_f32v(
6069 &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
6070 ),
6071 t: ts[s] as i32,
6072 pad: 0,
6073 })
6074 .collect();
6075 e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
6076 }
6077 for (s, attn) in attns.into_iter().enumerate() {
6078 let (attn_g, ag16) = self.full_attn_prime_post_fa(
6079 e,
6080 attn,
6081 &aps[s].gate,
6082 ts[s],
6083 n_head,
6084 head_dim,
6085 )?;
6086 let mut done = false;
6087 if let Some(xh) = &ag16 {
6088 done = e.try_f16_gemm_pre_into_off(
6089 &fa.wo,
6090 xh,
6091 ts[s],
6092 &mut mixed,
6093 offs[s] * n_embd,
6094 )?;
6095 }
6096 if !done {
6097 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
6098 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
6099 }
6100 }
6101 } else {
6102 let mut parts: Vec<Vec<CudaSlice<f32>>> =
6103 (0..b).map(|_| Vec::new()).collect();
6104 for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
6105 for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
6106 parts[s].push(ys);
6107 }
6108 }
6109 for (s, g3s) in parts.into_iter().enumerate() {
6110 let (attn_g, ag16) = self.full_attn_prime_core_inner(
6112 e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
6113 )?;
6114 let mut done = false;
6115 if let Some(xh) = &ag16 {
6116 done = e.try_f16_gemm_pre_into_off(
6117 &fa.wo,
6118 xh,
6119 ts[s],
6120 &mut mixed,
6121 offs[s] * n_embd,
6122 )?;
6123 }
6124 if !done {
6125 let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
6126 e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
6127 }
6128 }
6129 }
6130 }
6131 Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("batched cache prime"),
6132 Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("batched prime"),
6133 Mixer::Linear(la) => {
6134 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6139 let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
6140 let outs =
6141 self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
6142 for (s, (gn, gn16)) in outs.into_iter().enumerate() {
6143 let (o, t) = (offs[s], ts[s]);
6144 let mut done = false;
6145 if let Some(xh) = &gn16 {
6146 done = e.try_f16_gemm_pre_into_off(
6147 &la.ssm_out,
6148 xh,
6149 t,
6150 &mut mixed,
6151 o * n_embd,
6152 )?;
6153 }
6154 if !done {
6155 let m = e.matmul(&la.ssm_out, &gn, t)?;
6156 e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
6157 }
6158 }
6159 }
6160 }
6161 let mut x1 = e.uninit(total * n_embd)?;
6162 let mut z = e.uninit(total * n_embd)?;
6163 let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
6164 e.add_rms_norm_f16out(
6165 &x,
6166 &mixed,
6167 layer.post_attn_norm.float_data(),
6168 &mut x1,
6169 &mut z,
6170 &mut zx16,
6171 n_embd,
6172 total,
6173 eps,
6174 )?;
6175 let ffn_out = match &layer.ffn {
6176 crate::hybrid::Ffn::Dense {
6177 ffn_gate,
6178 ffn_up,
6179 ffn_down,
6180 } => {
6181 let n_ff = ffn_gate.out_features();
6182 let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
6183 let up = g2.pop().unwrap();
6184 let gate = g2.pop().unwrap();
6185 let mut act = e.uninit(total * n_ff)?;
6186 let d_lim = self.cfg.clamp_shexp_at(il as u32);
6190 if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
6191 let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
6192 e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
6193 match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
6194 Some(y) => y,
6195 None => e.matmul(ffn_down, &act, total)?,
6196 }
6197 } else {
6198 Self::ffn_act_lim(
6199 e,
6200 &self.cfg,
6201 &gate,
6202 &up,
6203 1.0,
6204 1.0,
6205 d_lim,
6206 &mut act,
6207 total * n_ff,
6208 )?;
6209 e.matmul(ffn_down, &act, total)?
6210 }
6211 }
6212 crate::hybrid::Ffn::Moe(m) => {
6213 self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
6214 }
6215 };
6216 let mut x2 = e.uninit(total * n_embd)?;
6217 e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
6218 x = x2;
6219 }
6220 let mut hn = e.uninit(total * n_embd)?;
6222 e.rms_norm(
6223 &x,
6224 self.output_norm.float_data(),
6225 &mut hn,
6226 n_embd,
6227 total,
6228 eps,
6229 )?;
6230 let mut hcat = e.uninit(b * n_embd)?;
6236 for s in 0..b {
6237 let last0 = (offs[s] + ts[s] - 1) * n_embd;
6238 e.copy_view_into(
6239 &mut hcat,
6240 s * n_embd,
6241 &hn.slice(last0..last0 + n_embd),
6242 n_embd,
6243 )?;
6244 }
6245 let logits_cat = if b >= 2 {
6246 e.try_f16_gemm(&self.output, &hcat, b)?
6247 } else {
6248 None
6249 };
6250 let logits_host: Option<Vec<f32>> = match &logits_cat {
6251 Some(lc) => Some(e.dtoh(lc)?),
6252 None => None,
6253 };
6254 let n_vocab = self.output.out_features();
6255 let mut hidden_all = if crate::spec::spec_hpost() {
6256 split(e, &hn, n_embd)?
6257 } else {
6258 split(e, &x, n_embd)?
6259 };
6260 let mut out = Vec::with_capacity(b);
6261 for s in 0..b {
6262 let last0 = (offs[s] + ts[s] - 1) * n_embd;
6263 let mut h_seed = e.uninit(n_embd)?;
6264 if !crate::spec::spec_hpost() {
6265 e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
6266 } else {
6267 e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6268 }
6269 let logits = match &logits_host {
6270 Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
6271 None => {
6272 let mut hlast = e.uninit(n_embd)?;
6273 e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6274 e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
6275 }
6276 };
6277 caches[s].pos += ts[s];
6278 out.push((logits, h_seed, hidden_all.remove(0)));
6279 }
6280 transaction.commit();
6281 Ok(out)
6282 }
6283
6284 #[allow(clippy::too_many_arguments)]
6295 fn full_attn_prime(
6296 &self,
6297 e: &Engine,
6298 fa: &FullAttnLayer,
6299 h: &CudaSlice<f32>,
6300 hx: Option<&CudaSlice<u8>>,
6301 pos_d: &CudaSlice<i32>,
6302 t: usize,
6303 cache: &mut Cache,
6304 il: usize,
6305 seq_end: usize,
6306 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6307 if self.uses_sliding_gated_moe_program() {
6308 return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
6309 }
6310 let g3 = match hx {
6315 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
6316 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
6317 };
6318 self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
6319 }
6320
6321 #[allow(clippy::too_many_arguments)] fn full_attn_prime_core(
6326 &self,
6327 e: &Engine,
6328 fa: &FullAttnLayer,
6329 g3: Vec<CudaSlice<f32>>,
6330 pos_d: &CudaSlice<i32>,
6331 t: usize,
6332 cache: &mut Cache,
6333 il: usize,
6334 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6335 let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
6336 if let Some(xh) = &ag16
6337 && let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)?
6338 {
6339 return Ok(y);
6340 }
6341 e.matmul(&fa.wo, &attn_g, t)
6342 }
6343
6344 #[allow(clippy::type_complexity)] #[allow(clippy::too_many_arguments)] fn full_attn_prime_core_inner(
6347 &self,
6348 e: &Engine,
6349 fa: &FullAttnLayer,
6350 g3: Vec<CudaSlice<f32>>,
6351 pos_d: &CudaSlice<i32>,
6352 t: usize,
6353 cache: &mut Cache,
6354 il: usize,
6355 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6356 let cfg = &self.cfg;
6357 let geometry = cfg.full_attention_geometry_at(il as u32);
6358 let n_head = geometry.n_head as usize;
6359 let n_head_kv = geometry.n_head_kv as usize;
6360 let head_dim = geometry.head_dim_k as usize;
6361 let scale = geometry.attention_scale();
6362 let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
6363 let AttnPre { q, k, v, gate } = pre;
6364 let mut attn = e.uninit(t * n_head * head_dim)?;
6365 self.full_attn_prime_fa_dispatch(
6366 e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
6367 )?;
6368 self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
6369 }
6370
6371 #[allow(clippy::type_complexity)]
6375 #[allow(clippy::too_many_arguments)] fn full_attn_prime_pre_fa(
6377 &self,
6378 e: &Engine,
6379 fa: &FullAttnLayer,
6380 mut g3: Vec<CudaSlice<f32>>,
6381 pos_d: &CudaSlice<i32>,
6382 t: usize,
6383 cache: &mut Cache,
6384 il: usize,
6385 ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
6386 let cfg = &self.cfg;
6387 let geometry = cfg.full_attention_geometry_at(il as u32);
6388 let n_head = geometry.n_head as usize;
6389 let n_head_kv = geometry.n_head_kv as usize;
6390 let head_dim = geometry.head_dim_k as usize;
6391 let eps = cfg.rms_eps;
6392
6393 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6397 let v = g3.pop().unwrap();
6398 let mut k = g3.pop().unwrap();
6399 let qf = g3.pop().unwrap();
6400 let (mut q, gate) = if gated {
6401 let mut q = e.uninit(t * n_head * head_dim)?;
6402 let mut gate = e.uninit(t * n_head * head_dim)?;
6403 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
6404 (q, Some(gate))
6405 } else {
6406 (qf, None)
6407 };
6408
6409 let mut qn = e.uninit(t * n_head * head_dim)?;
6410 e.rms_norm(
6411 &q,
6412 fa.q_norm.float_data(),
6413 &mut qn,
6414 head_dim,
6415 n_head * t,
6416 eps,
6417 )?;
6418 q = qn;
6419 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6420 e.rms_norm(
6421 &k,
6422 fa.k_norm.float_data(),
6423 &mut kn,
6424 head_dim,
6425 n_head_kv * t,
6426 eps,
6427 )?;
6428 k = kn;
6429 let rope_dims = geometry.n_rot as usize;
6430 e.rope_neox(
6431 &mut q,
6432 pos_d,
6433 head_dim,
6434 rope_dims,
6435 n_head,
6436 t,
6437 geometry.rope_base,
6438 1.0,
6439 )?;
6440 e.rope_neox(
6441 &mut k,
6442 pos_d,
6443 head_dim,
6444 rope_dims,
6445 n_head_kv,
6446 t,
6447 geometry.rope_base,
6448 1.0,
6449 )?;
6450
6451 {
6454 let kvl = cache.kv[il].as_mut().unwrap();
6455 assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
6456 e.append_kv_quantized_rows(
6457 &k,
6458 &v,
6459 &mut kvl.k,
6460 &mut kvl.v,
6461 kvl.len,
6462 t,
6463 kvl.kv_dim_k,
6464 kvl.kv_dim_v,
6465 kvl.k_tok_bytes,
6466 kvl.v_tok_bytes,
6467 crate::Engine::kv_fp8_on(),
6468 )?;
6469 kvl.len += t;
6470 let new_len = kvl.len as i32;
6471 e.set_i32_one(&mut kvl.len_d, new_len)?;
6472 }
6473
6474 let base_len = {
6475 let kvl = cache.kv[il].as_ref().unwrap();
6476 kvl.len - t };
6478 Ok((AttnPre { q, k, v, gate }, base_len))
6479 }
6480
6481 #[allow(clippy::too_many_arguments)]
6488 fn full_attn_prime_fa_dispatch(
6489 &self,
6490 e: &Engine,
6491 q: &CudaSlice<f32>,
6492 k: &CudaSlice<f32>,
6493 v: &CudaSlice<f32>,
6494 attn: &mut CudaSlice<f32>,
6495 base_len: usize,
6496 t: usize,
6497 cache: &mut Cache,
6498 il: usize,
6499 head_dim: usize,
6500 n_head: usize,
6501 n_head_kv: usize,
6502 scale: f32,
6503 ) -> Result<(), Box<dyn std::error::Error>> {
6504 if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
6517 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
6518 e.sdpa_naive(
6519 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
6520 )?;
6521 } else {
6522 e.fa_prefill(
6523 q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
6524 )?;
6525 }
6526 return Ok(());
6527 }
6528 let kvl = cache.kv[il].as_ref().unwrap();
6529 let t_kv = base_len + t;
6530 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6531 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6532 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
6536 e.sdpa_naive_quantized_view(
6537 q,
6538 &k_view,
6539 &v_view,
6540 attn,
6541 head_dim,
6542 n_head,
6543 n_head_kv,
6544 t,
6545 t_kv,
6546 scale,
6547 true,
6548 kvl.k_tok_bytes,
6549 kvl.v_tok_bytes,
6550 )?;
6551 return Ok(());
6552 }
6553 let deqw = std::env::var("MEMRA_PRIME_DEQW")
6561 .map(|v| v != "0")
6562 .unwrap_or(true);
6563 if deqw {
6564 e.fa_prefill_view_ws(
6565 q,
6566 &k_view,
6567 &v_view,
6568 attn,
6569 head_dim,
6570 n_head,
6571 n_head_kv,
6572 t,
6573 t_kv,
6574 scale,
6575 true,
6576 kvl.k_tok_bytes,
6577 kvl.v_tok_bytes,
6578 crate::Engine::kv_fp8_on(),
6579 )?;
6580 } else {
6581 e.fa_prefill_view(
6582 q,
6583 &k_view,
6584 &v_view,
6585 attn,
6586 head_dim,
6587 n_head,
6588 n_head_kv,
6589 t,
6590 t_kv,
6591 scale,
6592 true,
6593 kvl.k_tok_bytes,
6594 kvl.v_tok_bytes,
6595 crate::Engine::kv_fp8_on(),
6596 )?;
6597 }
6598 Ok(())
6599 }
6600
6601 #[allow(clippy::type_complexity)] fn full_attn_prime_post_fa(
6605 &self,
6606 e: &Engine,
6607 attn: CudaSlice<f32>,
6608 gate: &Option<CudaSlice<f32>>,
6609 t: usize,
6610 n_head: usize,
6611 head_dim: usize,
6612 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6613 let (attn_g, ag16) = match gate {
6614 Some(gate) => {
6615 let n = t * n_head * head_dim;
6616 let mut ag = e.uninit(n)?;
6617 if Self::f16out_on(e, t) {
6618 let mut a16 = e.alloc_u8_uninit(n * 2)?;
6619 e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
6620 (ag, Some(a16))
6621 } else {
6622 let mut gsig = e.uninit(n)?;
6623 e.sigmoid(gate, &mut gsig, n)?;
6624 e.mul(&attn, &gsig, &mut ag, n)?;
6625 (ag, None)
6626 }
6627 }
6628 None => (attn, None),
6629 };
6630 Ok((attn_g, ag16))
6631 }
6632
6633 #[allow(clippy::too_many_arguments)] fn linear_attn_prime(
6641 &self,
6642 e: &Engine,
6643 la: &LinearAttnLayer,
6644 h: &CudaSlice<f32>,
6645 hx: Option<&CudaSlice<u8>>,
6646 t: usize,
6647 cache: &mut Cache,
6648 il: usize,
6649 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6650 let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6652 let g4 = match hx {
6653 Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
6654 None => e.matmul_group(&ws, h, t)?,
6655 };
6656 self.linear_attn_prime_core(e, la, g4, t, cache, il)
6657 }
6658
6659 fn linear_attn_prime_core(
6661 &self,
6662 e: &Engine,
6663 la: &LinearAttnLayer,
6664 mut g4: Vec<CudaSlice<f32>>,
6665 t: usize,
6666 cache: &mut Cache,
6667 il: usize,
6668 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6669 self.linear_attn_prime_core_pad(e, la, std::mem::take(&mut g4), t, cache, il, None)
6670 }
6671
6672 #[allow(clippy::too_many_arguments)]
6676 #[allow(clippy::type_complexity)] fn linear_attn_prime_core_pad_inner(
6678 &self,
6679 e: &Engine,
6680 la: &LinearAttnLayer,
6681 mut g4: Vec<CudaSlice<f32>>,
6682 t: usize,
6683 cache: &mut Cache,
6684 il: usize,
6685 pad_len: Option<&CudaSlice<i32>>,
6686 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6687 let geometry = la.geometry;
6689 let d_state = geometry.key_head_dim as usize;
6690 let num_k = geometry.key_heads as usize;
6691 let num_v = geometry.value_heads as usize;
6692 let key_dim = d_state * num_k;
6693 let value_dim = geometry.value_head_dim as usize * num_v;
6694 let conv_dim = key_dim * 2 + value_dim;
6695 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(
6700 e,
6701 la,
6702 &qkv_mixed.slice(0..t * conv_dim),
6703 &z.slice(0..t * value_dim),
6704 &beta_raw.slice(0..t * num_v),
6705 &alpha.slice(0..t * num_v),
6706 t,
6707 cache,
6708 il,
6709 pad_len,
6710 )
6711 }
6712
6713 #[allow(clippy::too_many_arguments)]
6716 fn linear_attn_gdn_prep(
6717 &self,
6718 e: &Engine,
6719 la: &LinearAttnLayer,
6720 qkv_mixed: &cudarc::driver::CudaView<f32>,
6721 beta_raw: &cudarc::driver::CudaView<f32>,
6722 alpha: &cudarc::driver::CudaView<f32>,
6723 t: usize,
6724 cache: &mut Cache,
6725 il: usize,
6726 pad_len: Option<&CudaSlice<i32>>,
6727 ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
6728 let cfg = &self.cfg;
6729 let geometry = la.geometry;
6730 let d_state = geometry.key_head_dim as usize;
6731 let num_k = geometry.key_heads as usize;
6732 let num_v = geometry.value_heads as usize;
6733 let d_conv = geometry.conv_kernel as usize;
6734 let key_dim = d_state * num_k; let value_dim = geometry.value_head_dim as usize * num_v;
6736 let conv_dim = key_dim * 2 + value_dim; let eps = cfg.rms_eps;
6738 debug_assert!(
6739 t >= d_conv - 1,
6740 "stateful conv needs T >= pad (PRIME_MIN_T gates)"
6741 );
6742
6743 let rl = cache.recur[il].as_mut().unwrap();
6748 let hk = Self::gdn_hk(e, t, num_v, num_k);
6749 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
6750 let hk = if conv_fuse { hk } else { num_v }; let mut q_g = e.uninit(d_state * hk * t)?;
6752 let mut k_g = e.uninit(d_state * hk * t)?;
6753 let mut v_g = e.uninit(d_state * num_v * t)?;
6754 if conv_fuse {
6755 e.ssm_conv1d_gdn_state_pad(
6756 qkv_mixed,
6757 &mut rl.conv_state,
6758 la.ssm_conv1d.float_data(),
6759 &mut q_g,
6760 &mut k_g,
6761 &mut v_g,
6762 conv_dim,
6763 t,
6764 d_conv,
6765 d_state,
6766 num_v,
6767 num_k,
6768 key_dim,
6769 hk,
6770 pad_len,
6771 )?;
6772 } else {
6773 let mut conv_out = e.uninit(conv_dim * t)?; e.ssm_conv1d_tm_state_pad_v(
6775 qkv_mixed,
6776 &mut rl.conv_state,
6777 la.ssm_conv1d.float_data(),
6778 &mut conv_out,
6779 conv_dim,
6780 t,
6781 d_conv,
6782 pad_len,
6783 )?;
6784 e.qkv_to_gdn_repack(
6785 &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
6786 )?;
6787 }
6788 let mut q_l2 = e.uninit(d_state * hk * t)?;
6789 let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
6793 let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
6794 e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
6795 Some(qb)
6796 } else {
6797 e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
6798 None
6799 };
6800 let mut k_l2 = e.uninit(d_state * hk * t)?;
6801 let kb16 = if Engine::l2_v2_on(d_state) {
6803 let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
6804 e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
6805 Some(kb)
6806 } else {
6807 e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
6808 None
6809 };
6810 let mut beta = e.uninit(t * num_v)?;
6811 e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
6812 let mut g_log = e.uninit(t * num_v)?;
6813 e.gdn_glog_v(
6814 alpha,
6815 la.ssm_dt.float_data(),
6816 la.ssm_a.float_data(),
6817 &mut g_log,
6818 num_v,
6819 t,
6820 )?;
6821 if let Some(len_d) = pad_len {
6822 e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
6823 }
6824 Ok(GdnPrep {
6825 hk,
6826 q_l2,
6827 k_l2,
6828 v_g,
6829 beta,
6830 g_log,
6831 kb16,
6832 qb16,
6833 })
6834 }
6835
6836 #[allow(clippy::too_many_arguments)]
6841 #[allow(clippy::type_complexity)] fn linear_attn_prime_core_batch(
6843 &self,
6844 e: &Engine,
6845 la: &LinearAttnLayer,
6846 g4: &[CudaSlice<f32>],
6847 offs: &[usize],
6848 ts: &[usize],
6849 caches: &mut [&mut Cache],
6850 il: usize,
6851 ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
6852 let geometry = la.geometry;
6853 let d_state = geometry.key_head_dim as usize;
6854 let num_k = geometry.key_heads as usize;
6855 let num_v = geometry.value_heads as usize;
6856 let d_conv = geometry.conv_kernel as usize;
6857 let key_dim = d_state * num_k;
6858 let value_dim = geometry.value_head_dim as usize * num_v;
6859 let conv_dim = key_dim * 2 + value_dim;
6860 let eps = self.cfg.rms_eps;
6861 let scale = 1.0 / (d_state as f32).sqrt();
6862 let b = ts.len();
6863 let c = Engine::gdn_chunk_size();
6864 let carried = caches.iter().any(|c| c.pos > 0);
6867 let use_vl = !carried
6868 && (2..=8).contains(&b)
6869 && Engine::gdn_chunked_enabled()
6870 && ts.iter().all(|&t| t >= 16)
6871 && e.gdn_mma_enabled(c)
6872 && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
6873 if !use_vl {
6874 return (0..b)
6875 .map(|s| {
6876 let (o, t) = (offs[s], ts[s]);
6877 self.linear_attn_prime_core_pad_view(
6878 e,
6879 la,
6880 &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
6881 &g4[1].slice(o * value_dim..(o + t) * value_dim),
6882 &g4[2].slice(o * num_v..(o + t) * num_v),
6883 &g4[3].slice(o * num_v..(o + t) * num_v),
6884 t,
6885 caches[s],
6886 il,
6887 None,
6888 )
6889 })
6890 .collect();
6891 }
6892 struct SeqBufs {
6896 conv_out: CudaSlice<f32>,
6897 q_g: CudaSlice<f32>,
6898 k_g: CudaSlice<f32>,
6899 v_g: CudaSlice<f32>,
6900 q_l2: CudaSlice<f32>,
6901 k_l2: CudaSlice<f32>,
6902 beta: CudaSlice<f32>,
6903 g_log: CudaSlice<f32>,
6904 gn: CudaSlice<f32>,
6905 gn16: CudaSlice<u8>,
6906 }
6907 let f16o = Self::f16out_on(e, 16);
6908 let hk = Self::gdn_hk(e, 16, num_v, num_k); let mut sb = Vec::with_capacity(b);
6910 let mut pres = Vec::with_capacity(b);
6911 for &t in ts.iter().take(b) {
6912 sb.push(SeqBufs {
6913 conv_out: e.uninit(conv_dim * t)?,
6914 q_g: e.uninit(d_state * hk * t)?,
6915 k_g: e.uninit(d_state * hk * t)?,
6916 v_g: e.uninit(d_state * num_v * t)?,
6917 q_l2: e.uninit(d_state * hk * t)?,
6918 k_l2: e.uninit(d_state * hk * t)?,
6919 beta: e.uninit(t * num_v)?,
6920 g_log: e.uninit(t * num_v)?,
6921 gn: e.uninit(d_state * num_v * t)?,
6922 gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
6923 });
6924 pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
6925 }
6926 let prep_args: Vec<crate::GdnPrepVl> = (0..b)
6927 .map(|s| {
6928 let (o, t) = (offs[s], ts[s]);
6929 let rl = caches[s].recur[il].as_ref().unwrap();
6930 crate::GdnPrepVl {
6931 qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
6932 conv_state: e.addr_f32(&rl.conv_state),
6933 conv_out: e.addr_f32(&sb[s].conv_out),
6934 q_g: e.addr_f32(&sb[s].q_g),
6935 k_g: e.addr_f32(&sb[s].k_g),
6936 v_g: e.addr_f32(&sb[s].v_g),
6937 q_l2: e.addr_f32(&sb[s].q_l2),
6938 k_l2: e.addr_f32(&sb[s].k_l2),
6939 beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
6940 alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
6941 beta: e.addr_f32(&sb[s].beta),
6942 g_log: e.addr_f32(&sb[s].g_log),
6943 o: e.addr_f32(&pres[s].o),
6944 z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
6945 gn: e.addr_f32(&sb[s].gn),
6946 gn16: e.addr_u8(&sb[s].gn16),
6947 kb16: if Engine::l2_v2_on(d_state) {
6948 e.addr_u8(&pres[s].kb16)
6949 } else {
6950 0
6951 },
6952 qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
6953 e.addr_u8(&pres[s].qb16)
6954 } else {
6955 0
6956 },
6957 t: t as i32,
6958 pad: 0,
6959 }
6960 })
6961 .collect();
6962 let args: Vec<crate::GdnSeqVl> = (0..b)
6963 .map(|s| {
6964 let rl = caches[s].recur[il].as_ref().unwrap();
6965 crate::GdnSeqVl {
6966 kb16: e.addr_u8(&pres[s].kb16),
6967 gcum: e.addr_f32(&pres[s].gcum),
6968 beta: e.addr_f32(&sb[s].beta),
6969 u: e.addr_f32(&pres[s].u),
6970 wb16: e.addr_u8(&pres[s].wb16),
6971 y: e.addr_u8(&pres[s].y16),
6972 ssnap: e.addr_u8(&pres[s].ssnap16),
6973 state_in: e.addr_f32(&rl.ssm_state),
6974 state_out: e.addr_f32(&rl.ssm_state_alt),
6975 q: e.addr_f32(&sb[s].q_l2),
6976 p: e.addr_f32(&pres[s].p),
6977 o: e.addr_f32(&pres[s].o),
6978 k: e.addr_f32(&sb[s].k_l2),
6979 v: e.addr_f32(&sb[s].v_g),
6980 g: e.addr_f32(&sb[s].g_log),
6981 a: e.addr_f32(&pres[s].a),
6982 w: e.addr_f32(&pres[s].w),
6983 t: ts[s] as i32,
6984 nc: pres[s].nc as i32,
6985 }
6986 })
6987 .collect();
6988 e.gdn_prep_vl8(
6989 &prep_args,
6990 la.ssm_conv1d.float_data(),
6991 la.ssm_dt.float_data(),
6992 la.ssm_a.float_data(),
6993 conv_dim,
6994 d_conv,
6995 d_state,
6996 num_v,
6997 num_k,
6998 key_dim,
6999 hk,
7000 eps,
7001 )?;
7002 if !Engine::l2_v2_on(d_state) {
7005 e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
7006 }
7007 let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
7009 if !Engine::l2_v2_on(d_state) {
7011 for s in 0..b {
7012 e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
7013 }
7014 }
7015 let mut wa = [crate::GdnWVl::default(); 8];
7016 for s in 0..b {
7017 wa[s] = crate::GdnWVl {
7018 qb16: e.addr_u8(&pres[s].qb16),
7019 pb16: e.addr_u8(&pres[s].pb16),
7020 };
7021 }
7022 Some(crate::GdnWVl8(wa))
7023 } else {
7024 None
7025 };
7026 e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
7027 e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
7028 if f16o {
7029 e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
7030 }
7031 let mut out = Vec::with_capacity(b);
7033 for (s, bufs) in sb.into_iter().enumerate() {
7034 let rl = caches[s].recur[il].as_mut().unwrap();
7035 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7036 let (o, t) = (offs[s], ts[s]);
7037 let SeqBufs { mut gn, gn16, .. } = bufs;
7038 if f16o {
7039 out.push((gn, Some(gn16)));
7040 } else {
7041 let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
7042 e.gated_rmsnorm_zv(
7043 &pres[s].o,
7044 la.ssm_norm.float_data(),
7045 &z_v,
7046 &mut gn,
7047 d_state,
7048 num_v * t,
7049 eps,
7050 )?;
7051 out.push((gn, None));
7052 }
7053 }
7054 Ok(out)
7055 }
7056
7057 #[allow(clippy::too_many_arguments)]
7061 #[allow(clippy::type_complexity)] fn linear_attn_prime_core_pad_view(
7063 &self,
7064 e: &Engine,
7065 la: &LinearAttnLayer,
7066 qkv_mixed: &cudarc::driver::CudaView<f32>,
7067 z: &cudarc::driver::CudaView<f32>,
7068 beta_raw: &cudarc::driver::CudaView<f32>,
7069 alpha: &cudarc::driver::CudaView<f32>,
7070 t: usize,
7071 cache: &mut Cache,
7072 il: usize,
7073 pad_len: Option<&CudaSlice<i32>>,
7074 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
7075 let cfg = &self.cfg;
7076 let geometry = la.geometry;
7077 let d_state = geometry.key_head_dim as usize;
7078 let num_v = geometry.value_heads as usize;
7079 let eps = cfg.rms_eps;
7080 let scale = 1.0 / (d_state as f32).sqrt();
7081
7082 let prep =
7083 self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
7084
7085 let mut o = e.uninit(d_state * num_v * t)?;
7091 let rl = cache.recur[il].as_mut().unwrap();
7092 {
7093 let crate::cache::RecurLayer {
7094 ssm_state,
7095 ssm_state_alt,
7096 ..
7097 } = rl;
7098 e.gdn_scan_prefill(
7099 &prep.q_l2,
7100 &prep.k_l2,
7101 &prep.v_g,
7102 &prep.g_log,
7103 &prep.beta,
7104 prep.kb16.as_ref(),
7105 prep.qb16.as_ref(),
7106 ssm_state,
7107 ssm_state_alt,
7108 &mut o,
7109 num_v,
7110 t,
7111 scale,
7112 prep.hk,
7113 )?;
7114 }
7115 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7116
7117 let mut gn = e.uninit(d_state * num_v * t)?;
7120 let gn16 = if Self::f16out_on(e, t) {
7121 let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
7122 e.gated_rmsnorm_f16out_zv(
7123 &o,
7124 la.ssm_norm.float_data(),
7125 z,
7126 &mut gn,
7127 &mut g16,
7128 d_state,
7129 num_v * t,
7130 eps,
7131 )?;
7132 Some(g16)
7133 } else {
7134 e.gated_rmsnorm_zv(
7135 &o,
7136 la.ssm_norm.float_data(),
7137 z,
7138 &mut gn,
7139 d_state,
7140 num_v * t,
7141 eps,
7142 )?;
7143 None
7144 };
7145 Ok((gn, gn16))
7146 }
7147
7148 #[allow(clippy::too_many_arguments)]
7150 fn linear_attn_prime_core_pad(
7151 &self,
7152 e: &Engine,
7153 la: &LinearAttnLayer,
7154 g4: Vec<CudaSlice<f32>>,
7155 t: usize,
7156 cache: &mut Cache,
7157 il: usize,
7158 pad_len: Option<&CudaSlice<i32>>,
7159 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7160 let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
7161 if let Some(xh) = &gn16
7162 && let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)?
7163 {
7164 return Ok(y);
7165 }
7166 e.matmul(&la.ssm_out, &gn, t)
7167 }
7168
7169 pub fn full_attn(
7174 &self,
7175 e: &Engine,
7176 fa: &FullAttnLayer,
7177 h: &CudaSlice<f32>,
7178 pos_d: &CudaSlice<i32>,
7179 t: usize,
7180 il: usize,
7181 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7182 if self.uses_sliding_gated_moe_program() {
7183 return self.step35_attn(e, fa, h, pos_d, t, il);
7184 }
7185 let cfg = &self.cfg;
7186 let _n_embd = cfg.n_embd as usize;
7187 let geometry = cfg.full_attention_geometry_at(il as u32);
7188 let n_head = geometry.n_head as usize;
7189 let n_head_kv = geometry.n_head_kv as usize;
7190 let head_dim = geometry.head_dim_k as usize;
7191 let eps = cfg.rms_eps;
7192 let scale = geometry.attention_scale();
7193
7194 let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7197 let mut g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
7201 Some(g3) => g3,
7202 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
7203 };
7204 let v = g3.pop().unwrap();
7205 let mut k = g3.pop().unwrap();
7206 let qf = g3.pop().unwrap();
7207 let (mut q, gate) = if gated {
7208 let mut q = e.uninit(t * n_head * head_dim)?;
7209 let mut gate = e.uninit(t * n_head * head_dim)?;
7210 e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7211 (q, Some(gate))
7212 } else {
7213 (qf, None)
7214 };
7215
7216 let mut qn = e.uninit(t * n_head * head_dim)?;
7218 e.rms_norm(
7219 &q,
7220 fa.q_norm.float_data(),
7221 &mut qn,
7222 head_dim,
7223 n_head * t,
7224 eps,
7225 )?;
7226 q = qn;
7227 let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7228 e.rms_norm(
7229 &k,
7230 fa.k_norm.float_data(),
7231 &mut kn,
7232 head_dim,
7233 n_head_kv * t,
7234 eps,
7235 )?;
7236 k = kn;
7237 let rope_dims = geometry.n_rot as usize;
7238 e.rope_neox(
7239 &mut q,
7240 pos_d,
7241 head_dim,
7242 rope_dims,
7243 n_head,
7244 t,
7245 geometry.rope_base,
7246 1.0,
7247 )?;
7248 e.rope_neox(
7249 &mut k,
7250 pos_d,
7251 head_dim,
7252 rope_dims,
7253 n_head_kv,
7254 t,
7255 geometry.rope_base,
7256 1.0,
7257 )?;
7258
7259 let mut attn = e.uninit(t * n_head * head_dim)?;
7261 if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
7264 e.sdpa_naive(
7266 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7267 )?;
7268 } else {
7269 e.fa_prefill(
7270 &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7271 )?;
7272 }
7273
7274 let attn_g = match &gate {
7276 Some(gate) => {
7277 let mut gsig = e.uninit(t * n_head * head_dim)?;
7278 e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
7279 let mut ag = e.uninit(t * n_head * head_dim)?;
7280 e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
7281 ag
7282 }
7283 None => attn,
7284 };
7285
7286 self.full_attn_o(e, fa, &attn_g, t)
7288 }
7289
7290 fn mla_split_operand<'w>(
7298 w: &'w crate::model::GpuTensor,
7299 name: &str,
7300 il: usize,
7301 ) -> &'w CudaSlice<f32> {
7302 match w {
7303 crate::model::GpuTensor::Float { data, .. } => data,
7304 _ => panic!(
7305 "layer {il}: MLA conversion-split operand {name} is not f32-resident. The 3D \
7306 (d_nope|kv_rank, kv_rank|d_v, n_head) splits have no quantized resident layout: \
7307 a quantized 3D tensor mis-derives row_bytes in the generic 2D Quant arm, so the \
7308 source must dequantize the fused kv_b_proj (TensorTransform::SplitMlaKv). \
7309 Reaching this means both the loader rank guard and MlaAttnLayer::load's \
7310 residency audit were bypassed"
7311 ),
7312 }
7313 }
7314
7315 #[allow(clippy::too_many_arguments)]
7323 fn mla_attn_core(
7327 &self,
7328 e: &Engine,
7329 mla: &crate::hybrid::MlaAttnLayer,
7330 h: &CudaSlice<f32>,
7331 pos_d: &CudaSlice<i32>,
7332 t: usize,
7333 il: usize,
7334 latent: &mut CudaSlice<f32>,
7335 index_plane: Option<IndexerPlanes<'_>>,
7336 slot: usize,
7337 rows_exact: bool,
7338 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7339 let attn = self.mla_attn_core_pre_wo(
7340 e,
7341 mla,
7342 h,
7343 pos_d,
7344 t,
7345 il,
7346 latent,
7347 index_plane,
7348 slot,
7349 rows_exact,
7350 )?;
7351 if rows_exact {
7355 e.matmul_rows_exact(&mla.wo, &attn, t)
7356 } else {
7357 e.matmul(&mla.wo, &attn, t)
7358 }
7359 }
7360
7361 #[allow(clippy::too_many_arguments)]
7367 fn mla_attn_core_pre_wo(
7368 &self,
7369 e: &Engine,
7370 mla: &crate::hybrid::MlaAttnLayer,
7371 h: &CudaSlice<f32>,
7372 pos_d: &CudaSlice<i32>,
7373 t: usize,
7374 il: usize,
7375 latent: &mut CudaSlice<f32>,
7376 index_plane: Option<IndexerPlanes<'_>>,
7377 slot: usize,
7378 rows_exact: bool,
7379 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7380 let g = mla.geom;
7381 let cfg = &self.cfg;
7382 let eps = cfg.rms_eps;
7383 let base = cfg.rope_freq_base;
7384 let (nh, dn, dr, dv, r) = (g.n_head, g.d_nope, g.d_rope, g.d_v, g.kv_rank);
7385 assert_eq!(
7386 g.latent_dim,
7387 r + dr,
7388 "layer {il}: MlaGeom latent_dim disagrees with kv_rank + d_rope"
7389 );
7390 let t_kv = slot + t;
7391 let mm = |w: &crate::model::GpuTensor,
7396 x: &CudaSlice<f32>|
7397 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7398 if rows_exact {
7399 e.matmul_rows_exact(w, x, t)
7400 } else {
7401 e.matmul(w, x, t)
7402 }
7403 };
7404
7405 let q_a = mm(&mla.wq_a, h)?;
7407 let q_lora = mla.wq_b.in_features();
7408 let mut q_an = e.uninit(t * q_lora)?;
7409 e.rms_norm(&q_a, mla.q_a_norm.float_data(), &mut q_an, q_lora, t, eps)?;
7410 let q = mm(&mla.wq_b, &q_an)?;
7411 let mut q_nope = e.uninit(t * nh * dn)?;
7414 let mut q_pe = e.uninit((t * nh * dr).max(1))?;
7415 e.mla_split_latent(&q, &mut q_nope, &mut q_pe, t * nh, dn, dr)?;
7416 e.mla_rope_interleaved(&mut q_pe, pos_d, t, nh, dr, base)?;
7419
7420 let kv = mm(&mla.wkv_a, h)?;
7422 let mut c_kv = e.uninit(t * r)?;
7423 let mut k_pe = e.uninit((t * dr).max(1))?;
7424 e.mla_split_latent(&kv, &mut c_kv, &mut k_pe, t, r, dr)?;
7425 let mut c_kv_n = e.uninit(t * r)?;
7426 e.rms_norm(&c_kv, mla.kv_a_norm.float_data(), &mut c_kv_n, r, t, eps)?;
7427 e.mla_rope_interleaved(&mut k_pe, pos_d, t, 1, dr, base)?;
7428 e.mla_append_latent(latent, &c_kv_n, &k_pe, slot, t, r, dr)?;
7429
7430 let gathered = match (&mla.index, index_plane) {
7434 (Some(indexer), Some(plane)) => {
7435 Some(self.mla_kpool_select(e, indexer, h, &q_an, plane, t, slot, il, rows_exact)?)
7436 }
7437 (Some(_), None) => {
7438 return Err(format!(
7439 "layer {il} declares a DSA k-pool indexer but no indexer state plane was \
7440 supplied — the ModelPlan must declare StatePlan::LatentKvCache with a \
7441 non-zero index_width for it"
7442 )
7443 .into());
7444 }
7445 (None, _) => None,
7446 };
7447
7448 let wk_b = Self::mla_split_operand(&mla.wk_b, "attn_k_b", il);
7450 let wv_b = Self::mla_split_operand(&mla.wv_b, "attn_v_b", il);
7451
7452 if let Some((idx, slots)) = &gathered
7466 && dr == 0
7467 && r == 512
7468 && t >= 16
7469 && !rows_exact && !crate::portable_mma_gated()
7471 && mla_tc_prefill_enabled()
7472 && let Some(attn) = self.mla_tc_prefill_chain(
7473 e, wk_b, wv_b, &q_nope, latent, idx, *slots, t, t_kv, nh, dn, dv, r, g.scale,
7474 )?
7475 {
7476 return Ok(attn);
7477 }
7478
7479 let mut q_lat = e.uninit(t * nh * r)?;
7480 e.mla_absorb_q(&q_nope, wk_b, &mut q_lat, t, nh, dn, r)?;
7481 let mut o_lat = e.uninit(t * nh * r)?;
7482 match &gathered {
7483 Some((idx, slots)) => e.mla_attn_gathered(
7484 &q_lat, &q_pe, latent, idx, &mut o_lat, nh, r, dr, t, *slots, g.scale,
7485 )?,
7486 None => e.mla_attn_absorbed(
7487 &q_lat, &q_pe, latent, &mut o_lat, nh, r, dr, t, t_kv, g.scale,
7488 )?,
7489 }
7490 let mut attn = e.uninit(t * nh * dv)?;
7491 e.mla_decompress_v(&o_lat, wv_b, &mut attn, t, nh, dv, r)?;
7492
7493 Ok(attn)
7494 }
7495
7496 #[allow(clippy::too_many_arguments)]
7520 fn mla_tc_prefill_chain(
7521 &self,
7522 e: &Engine,
7523 wk_b: &CudaSlice<f32>,
7524 wv_b: &CudaSlice<f32>,
7525 q_nope: &CudaSlice<f32>,
7526 latent: &CudaSlice<f32>,
7527 idx: &CudaSlice<i32>,
7528 width: usize,
7529 t: usize,
7530 t_kv: usize,
7531 nh: usize,
7532 dn: usize,
7533 dv: usize,
7534 r: usize,
7535 scale: f32,
7536 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7537 fn declined(stage: &str, m: usize, n: usize, k: usize, batch: usize) {
7540 type ShapeSet = std::collections::HashSet<(usize, usize, usize, usize)>;
7541 static SAID: std::sync::Mutex<Option<ShapeSet>> = std::sync::Mutex::new(None);
7542 let mut g = SAID.lock().unwrap();
7543 if g.get_or_insert_with(std::collections::HashSet::new)
7544 .insert((m, n, k, batch))
7545 {
7546 eprintln!(
7547 "[mla-tc-prefill] DECLINED at {stage} m={m} n={n} k={k} batch={batch} \
7548 (no cuBLASLt heuristic) — this call falls back to the f32 MLA kernels"
7549 );
7550 }
7551 }
7552 for (name, n) in [
7556 ("wk_b", nh * r * dn),
7557 ("wv_b", nh * dv * r),
7558 ("q_nope", t * nh * dn),
7559 ("latent", t_kv * r),
7560 ] {
7561 debug_assert!(
7562 n.is_multiple_of(4),
7563 "mla-tc-prefill: {name} elems {n} % 4 != 0"
7564 );
7565 let _ = (name, n);
7566 }
7567 let wk_bf = e.f32_to_bf16(wk_b, nh * r * dn)?;
7568 let wv_bf = e.f32_to_bf16(wv_b, nh * dv * r)?;
7569 let qn_bf = e.f32_to_bf16(q_nope, t * nh * dn)?;
7570 let mut q_lat_bf = e.alloc_u8_uninit(t * nh * r * 2)?;
7574 if !e.mla_bf16_gemm_sb_bf16out(
7575 &wk_bf,
7576 &qn_bf,
7577 &mut q_lat_bf,
7578 t,
7579 r,
7580 dn,
7581 nh * dn,
7582 dn,
7583 nh * r,
7584 r,
7585 nh,
7586 )? {
7587 declined("absorb", t, r, dn, nh);
7588 return Ok(None);
7589 }
7590 let cache_bf = e.f32_to_bf16(latent, t_kv * r)?;
7592 let mut o_lat = e.uninit(t * nh * r)?;
7593 e.mla_attn_gathered_tc(
7594 &q_lat_bf, &cache_bf, idx, &mut o_lat, nh, r, t, width, scale,
7595 )?;
7596 let o_bf = e.f32_to_bf16(&o_lat, t * nh * r)?;
7599 let mut attn = e.uninit(t * nh * dv)?;
7600 if !e.mla_bf16_gemm_sb_f32out(
7601 &wv_bf,
7602 &o_bf,
7603 &mut attn,
7604 t,
7605 dv,
7606 r,
7607 nh * r,
7608 r,
7609 nh * dv,
7610 dv,
7611 nh,
7612 )? {
7613 declined("decompress", t, dv, r, nh);
7614 return Ok(None);
7615 }
7616 crate::MLA_TC_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7617 {
7618 static ANNOUNCED: std::sync::Once = std::sync::Once::new();
7619 ANNOUNCED.call_once(|| {
7620 eprintln!(
7621 "[mla-tc-prefill] engaged: absorb/decompress = strided-batched bf16 TC \
7622 GEMMs, attention = fa_mla_gathered_bf16 (t={t}, t_kv={t_kv}, nh={nh}, \
7623 width={width}); dispatches counted in MLA_TC_PREFILL_DISPATCHES"
7624 );
7625 });
7626 }
7627 Ok(Some(attn))
7628 }
7629
7630 #[allow(clippy::too_many_arguments)]
7632 fn mla_kpool_select(
7633 &self,
7634 e: &Engine,
7635 indexer: &crate::hybrid::MlaIndexer,
7636 h: &CudaSlice<f32>,
7637 q_resid: &CudaSlice<f32>,
7638 plane: IndexerPlanes<'_>,
7639 t: usize,
7640 slot: usize,
7641 il: usize,
7642 rows_exact: bool,
7643 ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7644 Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, rows_exact).map_err(
7645 |source| -> Box<dyn std::error::Error> {
7646 format!("layer {il}: DSA k-pool selection failed: {source}").into()
7647 },
7648 )
7649 }
7650
7651 #[allow(clippy::too_many_arguments)]
7668 pub fn mla_kpool_indices(
7669 e: &Engine,
7670 indexer: &crate::hybrid::MlaIndexer,
7671 h: &CudaSlice<f32>,
7672 q_resid: &CudaSlice<f32>,
7673 plane: IndexerPlanes<'_>,
7674 t: usize,
7675 slot: usize,
7676 ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7677 Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, false)
7678 }
7679
7680 #[allow(clippy::too_many_arguments)]
7685 pub fn mla_kpool_indices_ex(
7686 e: &Engine,
7687 indexer: &crate::hybrid::MlaIndexer,
7688 h: &CudaSlice<f32>,
7689 q_resid: &CudaSlice<f32>,
7690 plane: IndexerPlanes<'_>,
7691 t: usize,
7692 slot: usize,
7693 rows_exact: bool,
7694 ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7695 let mm = |w: &crate::model::GpuTensor,
7696 x: &CudaSlice<f32>|
7697 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7698 if rows_exact {
7699 e.matmul_rows_exact(w, x, t)
7700 } else {
7701 e.matmul(w, x, t)
7702 }
7703 };
7704 const INDEX_NORM_EPS: f32 = 1e-5;
7708
7709 let ig = indexer.geom;
7710 let d = ig.head_dim;
7711 let t_kv = slot + t;
7712 let IndexerPlanes {
7713 state: plane,
7714 pool_keys: pool_key_plane,
7715 ready: pools_ready,
7716 state_ring_rows,
7717 capacity_tokens,
7718 } = plane;
7719
7720 let ring = if state_ring_rows == 0 {
7728 0
7729 } else {
7730 state_ring_rows / ig.pool * ig.pool
7731 };
7732 if state_ring_rows > 0 && ring == 0 {
7733 return Err(format!(
7734 "indexer tail ring of {state_ring_rows} rows cannot hold one pool of {}; \
7735 raise MEMRA_DSA_INDEX_RING or set it to 0 for the flat plane",
7736 ig.pool
7737 )
7738 .into());
7739 }
7740 if *pools_ready > slot / ig.pool {
7745 return Err(format!(
7746 "resident k-pool key plane claims {} finished pools but the cache holds only {} \
7747 complete pools before this call ({slot} rows / pool {}) — a rewind reduced the \
7748 latent length without clamping index_pools_ready",
7749 *pools_ready,
7750 slot / ig.pool,
7751 ig.pool
7752 )
7753 .into());
7754 }
7755
7756 let k_raw = mm(&indexer.wk, h)?;
7759 let mut k_norm = e.uninit(t * d)?;
7760 e.layer_norm_bias(
7761 &k_raw,
7762 indexer.k_norm_w.float_data(),
7763 indexer.k_norm_b.float_data(),
7764 &mut k_norm,
7765 d,
7766 t,
7767 INDEX_NORM_EPS,
7768 )?;
7769 let gate = mm(&indexer.kpool_gate, h)?;
7770
7771 let n_pools = t_kv / ig.pool;
7777 let select_k = ig.select_k(n_pools);
7778 let width = ig.index_width(n_pools);
7779 let capacity_pools = capacity_tokens / ig.pool;
7787 let need = (capacity_pools * d).max(n_pools * d).max(1);
7788 if pool_key_plane.as_ref().is_none_or(|k| k.len() < need) {
7789 *pool_key_plane = Some(e.uninit(need)?);
7790 *pools_ready = 0;
7791 }
7792 let pool_keys = pool_key_plane
7793 .as_mut()
7794 .expect("resident pool-key plane just allocated");
7795
7796 let ape = indexer.kpool_ape.float_data();
7809 let mut cur = slot;
7810 let mut appended = 0usize;
7811 while appended < t {
7812 let take =
7813 crate::cache::index_ring_take(ring, ig.pool, *pools_ready, cur, t - appended)
7814 .ok_or_else(|| -> Box<dyn std::error::Error> {
7815 format!(
7816 "indexer tail ring lapped: {ring} rows cannot hold the {} rows still \
7817 owed to unbuilt pools at row {cur} (pools_ready {}, pool {}, slot \
7818 {slot}, t {t}). The pool-key plane was reset or the cache rewound \
7819 without clamping index_pools_ready, so rows this call must read were \
7820 already overwritten. Raise MEMRA_DSA_INDEX_RING, or set \
7821 MEMRA_DSA_INDEX_RING=0 for the flat plane",
7822 cur.saturating_sub((*pools_ready).saturating_mul(ig.pool)),
7823 *pools_ready,
7824 ig.pool
7825 )
7826 .into()
7827 })?;
7828 debug_assert!(take > 0 && appended + take <= t);
7829 e.mla_index_append(plane, &k_norm, &gate, appended, cur, take, d, d, ring)?;
7830 cur += take;
7831 appended += take;
7832 let ready_now = cur / ig.pool;
7833 e.mla_kpool_pool_keys(
7834 plane,
7835 ape,
7836 pool_keys,
7837 (*pools_ready).min(ready_now),
7838 ready_now,
7839 ig.pool,
7840 d,
7841 ring,
7842 )?;
7843 *pools_ready = ready_now;
7844 }
7845 debug_assert!(t == 0 || *pools_ready == n_pools);
7846 let pool_keys = &*pool_keys;
7847
7848 let q_index = mm(&indexer.wq_b, q_resid)?;
7850 let head_weights = mm(&indexer.weights_proj, h)?;
7851 let mut score = e.uninit((t * n_pools).max(1))?;
7852 e.mla_kpool_score(
7853 &q_index,
7854 pool_keys,
7855 &head_weights,
7856 &mut score,
7857 t,
7858 ig.heads,
7859 d,
7860 n_pools,
7861 ig.pool,
7862 slot,
7863 (d as f32).powf(-0.5),
7864 (ig.heads as f32).powf(-0.5),
7865 )?;
7866 let mut idx = e.uninit_i32(t * width)?;
7867 e.mla_kpool_select(
7868 &score,
7869 &mut idx,
7870 t,
7871 n_pools,
7872 ig.pool,
7873 select_k,
7874 width,
7875 slot,
7876 ig.always_select_tail,
7877 )?;
7878 Ok((idx, width))
7879 }
7880
7881 pub fn mla_attn(
7884 &self,
7885 e: &Engine,
7886 mla: &crate::hybrid::MlaAttnLayer,
7887 h: &CudaSlice<f32>,
7888 pos_d: &CudaSlice<i32>,
7889 t: usize,
7890 il: usize,
7891 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7892 if mla.tp.is_some() {
7893 return Err(format!(
7894 "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the stateless \
7895 mixer path is unwired for a head shard"
7896 )
7897 .into());
7898 }
7899 let mut latent = e.uninit(t * mla.geom.latent_dim)?;
7900 let mut index_plane = match mla.index.as_ref() {
7901 Some(indexer) => Some(e.uninit(t * indexer.geom.state_width())?),
7902 None => None,
7903 };
7904 let mut pool_keys = None;
7907 let mut pools_ready = 0usize;
7908 let planes = index_plane.as_mut().map(|state| IndexerPlanes {
7909 state,
7910 pool_keys: &mut pool_keys,
7911 ready: &mut pools_ready,
7912 state_ring_rows: 0,
7914 capacity_tokens: t,
7915 });
7916 self.mla_attn_core(e, mla, h, pos_d, t, il, &mut latent, planes, 0, false)
7917 }
7918
7919 #[allow(clippy::too_many_arguments)] pub fn mla_attn_cached(
7924 &self,
7925 e: &Engine,
7926 mla: &crate::hybrid::MlaAttnLayer,
7927 h: &CudaSlice<f32>,
7928 pos_d: &CudaSlice<i32>,
7929 t: usize,
7930 il: usize,
7931 cache: &mut Cache,
7932 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7933 self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, false)
7934 }
7935
7936 #[allow(clippy::too_many_arguments)] pub fn mla_attn_cached_rows_exact(
7947 &self,
7948 e: &Engine,
7949 mla: &crate::hybrid::MlaAttnLayer,
7950 h: &CudaSlice<f32>,
7951 pos_d: &CudaSlice<i32>,
7952 t: usize,
7953 il: usize,
7954 cache: &mut Cache,
7955 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7956 self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, true)
7957 }
7958
7959 #[allow(clippy::too_many_arguments)] fn mla_attn_cached_inner(
7961 &self,
7962 e: &Engine,
7963 mla: &crate::hybrid::MlaAttnLayer,
7964 h: &CudaSlice<f32>,
7965 pos_d: &CudaSlice<i32>,
7966 t: usize,
7967 il: usize,
7968 cache: &mut Cache,
7969 rows_exact: bool,
7970 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7971 if mla.tp.is_some() {
7976 return Err(format!(
7977 "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the plain mixer \
7978 path is unwired for a head shard — only the TP decode/prime walk may \
7979 execute it (rows_exact={rows_exact})"
7980 )
7981 .into());
7982 }
7983 let max_ctx = cache.max_ctx;
7986 let layer = cache.latent[il].as_mut().ok_or_else(|| {
7987 format!(
7988 "layer {il} is Mixer::Mla but the cache has no latent plane — the ModelPlan \
7989 must declare StatePlan::LatentKvCache for it"
7990 )
7991 })?;
7992 let attn =
7993 self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?;
7994 if rows_exact {
7997 e.matmul_rows_exact(&mla.wo, &attn, t)
7998 } else {
7999 e.matmul(&mla.wo, &attn, t)
8000 }
8001 }
8002
8003 #[allow(clippy::too_many_arguments)]
8008 pub(crate) fn mla_attn_cached_pre_wo(
8009 &self,
8010 e: &Engine,
8011 mla: &crate::hybrid::MlaAttnLayer,
8012 h: &CudaSlice<f32>,
8013 pos_d: &CudaSlice<i32>,
8014 t: usize,
8015 il: usize,
8016 layer: &mut memra_kv::LatentKvLayer,
8017 max_ctx: usize,
8018 rows_exact: bool,
8019 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8020 let slot = layer.len;
8021 let width = layer.width;
8022 assert_eq!(
8023 width, mla.geom.latent_dim,
8024 "layer {il}: cache latent width {width} != MlaGeom latent_dim {}",
8025 mla.geom.latent_dim
8026 );
8027 let capacity = layer.rows.len() / width;
8028 if slot + t > capacity {
8029 return Err(format!(
8030 "layer {il}: latent cache overflow — {slot} + {t} rows exceeds capacity {capacity}"
8031 )
8032 .into());
8033 }
8034 if mla.index.is_some() && layer.index_rows.is_none() {
8035 return Err(format!(
8036 "layer {il} loaded a DSA k-pool indexer but its latent cache carries no indexer \
8037 state plane — StatePlan::LatentKvCache declared index_width 0 for a layer whose \
8038 SparseIndexPlan is Own {{ kpool: Some(..) }}"
8039 )
8040 .into());
8041 }
8042 let mut rows = std::mem::replace(&mut layer.rows, e.uninit(0)?);
8047 let mut index_rows = layer.index_rows.take();
8048 let mut pool_keys = layer.index_pool_keys.take();
8049 let mut pools_ready = layer.index_pools_ready;
8050 let index_ring_rows = layer.index_ring_rows.unwrap_or(0);
8051 let planes = index_rows.as_mut().map(|state| IndexerPlanes {
8052 state,
8053 pool_keys: &mut pool_keys,
8054 ready: &mut pools_ready,
8055 state_ring_rows: index_ring_rows,
8056 capacity_tokens: max_ctx,
8057 });
8058 let out =
8059 self.mla_attn_core_pre_wo(e, mla, h, pos_d, t, il, &mut rows, planes, slot, rows_exact);
8060 layer.rows = rows;
8061 layer.index_rows = index_rows;
8062 layer.index_pool_keys = pool_keys;
8063 layer.index_pools_ready = if out.is_ok() {
8068 pools_ready
8069 } else if let Some(indexer) = mla.index.as_ref() {
8070 pools_ready.min(layer.len / indexer.geom.pool)
8071 } else {
8072 pools_ready
8073 };
8074 let out = out?;
8075 if let Some(indexer) = mla.index.as_ref() {
8080 let pool = indexer.geom.pool;
8081 if layer.index_pool != 0 && layer.index_pool != pool {
8082 return Err(format!(
8083 "layer {il}: resident indexer pool {} != loaded geometry pool {pool}",
8084 layer.index_pool,
8085 )
8086 .into());
8087 }
8088 layer.index_pool = pool;
8089 }
8090 layer.len = slot + t;
8091 let len_i32 = i32::try_from(layer.len).map_err(|_| "latent length exceeds i32 mirror")?;
8092 e.i32_mirror_store(&mut layer.len_d, len_i32)?;
8095 Ok(out)
8096 }
8097
8098 #[allow(clippy::too_many_arguments)]
8106 pub(crate) fn mla_tp_attn_cached(
8107 &self,
8108 e: &Engine,
8109 mla: &crate::hybrid::MlaAttnLayer,
8110 h: &CudaSlice<f32>,
8111 pos_d: &CudaSlice<i32>,
8112 t: usize,
8113 il: usize,
8114 cache: &mut Cache,
8115 rows_exact: bool,
8116 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8117 let tp = mla
8118 .tp
8119 .as_ref()
8120 .ok_or("mla_tp_attn_cached called on an unsharded layer")?;
8121 let rt = &tp.rt;
8122 let ranks = tp.ranks();
8123 let g = mla.geom; let hl = g.n_head;
8125 let dv = g.d_v;
8126 let full_heads = tp.full_heads;
8127 let n_embd = tp.n_embd;
8128 let hh = n_embd / ranks;
8129 let max_ctx = cache.max_ctx;
8130
8131 let hop = rt.hop(e);
8135 let h_peers = crate::glm5_tp_transport::fanout_f32(&hop, h, h.len())?;
8136 let pos_peers = crate::glm5_tp_transport::fanout_i32(&hop, pos_d, pos_d.len())?;
8137
8138 {
8140 let canonical = cache.latent[il].as_ref().ok_or_else(|| {
8141 format!("layer {il}: glm5 TP MLA walk found no canonical latent plane")
8142 })?;
8143 crate::glm5_tp::ensure_mla_peer_latent(
8144 rt,
8145 canonical,
8146 &mut cache.glm5_tp_latent_peer[il],
8147 )?;
8148 }
8149
8150 let mut attn: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
8157 for r in 1..ranks {
8158 let layer = &mut cache.glm5_tp_latent_peer[il].as_mut().unwrap()[r - 1];
8159 attn[r] = Some(self.mla_attn_cached_pre_wo(
8160 &rt.peers[r - 1],
8161 &tp.peers[r - 1],
8162 &h_peers[r - 1],
8163 &pos_peers[r - 1],
8164 t,
8165 il,
8166 layer,
8167 max_ctx,
8168 rows_exact,
8169 )?);
8170 }
8171 attn[0] = {
8172 let layer = cache.latent[il].as_mut().unwrap();
8173 Some(self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?)
8174 };
8175
8176 let part = hl * dv;
8179 debug_assert_eq!(full_heads * dv, ranks * part);
8180 let attn_refs: Vec<&CudaSlice<f32>> = attn
8181 .iter()
8182 .map(|a| a.as_ref().expect("filled above"))
8183 .collect();
8184 let fulls = crate::glm5_tp_transport::gather_parts(&hop, &attn_refs, t, part)?;
8185
8186 let mut ys = Vec::with_capacity(ranks);
8189 if rows_exact {
8190 ys.push(e.matmul_rows_exact(&mla.wo, &fulls[0], t)?);
8191 for r in 1..ranks {
8192 ys.push(rt.peers[r - 1].matmul_rows_exact(&tp.peers[r - 1].wo, &fulls[r], t)?);
8193 }
8194 } else {
8195 ys.push(e.matmul(&mla.wo, &fulls[0], t)?);
8196 for r in 1..ranks {
8197 ys.push(rt.peers[r - 1].matmul(&tp.peers[r - 1].wo, &fulls[r], t)?);
8198 }
8199 }
8200 debug_assert_eq!(n_embd, ranks * hh);
8202 let y_refs: Vec<&CudaSlice<f32>> = ys.iter().collect();
8203 crate::glm5_tp_transport::concat_parts_on_root(&hop, &y_refs, t, hh)
8204 }
8205
8206 pub fn linear_attn(
8208 &self,
8209 e: &Engine,
8210 la: &LinearAttnLayer,
8211 h: &CudaSlice<f32>,
8212 t: usize,
8213 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8214 let cfg = &self.cfg;
8215 let _n_embd = cfg.n_embd as usize;
8216 let geometry = la.geometry;
8217 let d_state = geometry.key_head_dim as usize;
8218 let num_k = geometry.key_heads as usize;
8219 let num_v = geometry.value_heads as usize;
8220 let d_conv = geometry.conv_kernel as usize;
8221 let head_k = d_state;
8222 let head_v = geometry.value_head_dim as usize;
8223 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;
8227 let scale = 1.0 / (d_state as f32).sqrt();
8228
8229 let mut g4 = e.matmul_group(
8232 &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
8233 h,
8234 t,
8235 )?;
8236 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);
8248 let mut q_g = e.uninit(d_state * num_v * t)?;
8249 let mut k_g = e.uninit(d_state * num_v * t)?;
8250 let mut v_g = e.uninit(d_state * num_v * t)?;
8251 e.ssm_conv1d_gdn(
8252 &qkv_mixed,
8253 la.ssm_conv1d.float_data(),
8254 &mut q_g,
8255 &mut k_g,
8256 &mut v_g,
8257 conv_dim,
8258 t,
8259 d_conv,
8260 d_state,
8261 num_v,
8262 num_k,
8263 key_dim,
8264 )?;
8265 let mut q_l2 = e.uninit(d_state * num_v * t)?;
8267 e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8268 let mut k_l2 = e.uninit(d_state * num_v * t)?;
8269 e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8270 let v_gd = v_g;
8271
8272 let mut beta = e.uninit(t * num_v)?;
8275 e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8276 let mut g_log = e.uninit(t * num_v)?;
8278 e.gdn_glog(
8279 &alpha,
8280 la.ssm_dt.float_data(),
8281 la.ssm_a.float_data(),
8282 &mut g_log,
8283 num_v,
8284 t,
8285 )?;
8286
8287 let state_in = e.zeros(d_state * d_state * num_v)?; let mut state_out = e.zeros(d_state * d_state * num_v)?;
8290 let mut o = e.uninit(d_state * num_v * t)?;
8291 e.gdn_scan_prefill(
8292 &q_l2,
8293 &k_l2,
8294 &v_gd,
8295 &g_log,
8296 &beta,
8297 None,
8298 None,
8299 &state_in,
8300 &mut state_out,
8301 &mut o,
8302 num_v,
8303 t,
8304 scale,
8305 num_v,
8306 )?;
8307
8308 let mut gn = e.uninit(d_state * num_v * t)?;
8313 e.gated_rmsnorm(
8314 &o,
8315 la.ssm_norm.float_data(),
8316 &z,
8317 &mut gn,
8318 d_state,
8319 num_v * t,
8320 eps,
8321 )?;
8322
8323 let out = e.matmul(&la.ssm_out, &gn, t)?;
8327 Ok(out)
8328 }
8329}
8330
8331impl HybridModel {
8332 pub fn moe_ffn_il(
8343 &self,
8344 e: &Engine,
8345 m: &MoeWeights,
8346 z: &CudaSlice<f32>,
8347 t: usize,
8348 il: u16,
8349 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8350 Self::moe_ffn_inner(
8351 e,
8352 m,
8353 z,
8354 None,
8355 t,
8356 &self.cfg,
8357 il,
8358 self.max_moe_block(),
8359 false,
8360 None,
8361 self.uses_sliding_gated_moe_program(),
8362 false,
8363 )
8364 }
8365
8366 pub fn moe_ffn_il_prefill(
8369 &self,
8370 e: &Engine,
8371 m: &MoeWeights,
8372 z: &CudaSlice<f32>,
8373 t: usize,
8374 il: u16,
8375 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8376 Self::moe_ffn_inner(
8377 e,
8378 m,
8379 z,
8380 None,
8381 t,
8382 &self.cfg,
8383 il,
8384 self.max_moe_block(),
8385 true,
8386 Some(&self.step_grouped_prefill),
8387 self.uses_sliding_gated_moe_program(),
8388 false,
8389 )
8390 }
8391
8392 pub fn moe_ffn_il_zq8(
8396 &self,
8397 e: &Engine,
8398 m: &MoeWeights,
8399 z: &CudaSlice<f32>,
8400 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
8401 t: usize,
8402 il: u16,
8403 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8404 Self::moe_ffn_inner(
8405 e,
8406 m,
8407 z,
8408 zq8,
8409 t,
8410 &self.cfg,
8411 il,
8412 self.max_moe_block(),
8413 false,
8414 None,
8415 self.uses_sliding_gated_moe_program(),
8416 false,
8417 )
8418 }
8419
8420 pub(crate) fn moe_ffn_il_zq8_vrows(
8425 &self,
8426 e: &Engine,
8427 m: &MoeWeights,
8428 z: &CudaSlice<f32>,
8429 t: usize,
8430 il: u16,
8431 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8432 Self::moe_ffn_inner(
8433 e,
8434 m,
8435 z,
8436 None,
8437 t,
8438 &self.cfg,
8439 il,
8440 self.max_moe_block(),
8441 false,
8442 None,
8443 self.uses_sliding_gated_moe_program(),
8444 true,
8445 )
8446 }
8447
8448 pub(crate) fn moe_ffn(
8456 e: &Engine,
8457 m: &MoeWeights,
8458 z: &CudaSlice<f32>,
8459 t: usize,
8460 cfg: &ModelConfig,
8461 il: u16,
8462 max_block: usize,
8463 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8464 Self::moe_ffn_inner(
8465 e, m, z, None, t, cfg, il, max_block, false, None, false, false,
8466 )
8467 }
8468
8469 #[allow(clippy::too_many_arguments)]
8470 #[allow(clippy::map_entry)] pub(crate) fn moe_ffn_inner(
8472 e: &Engine,
8473 m: &MoeWeights,
8474 z: &CudaSlice<f32>,
8475 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
8476 t: usize,
8477 cfg: &ModelConfig,
8478 il: u16,
8479 max_block: usize,
8480 prefill: bool,
8481 grouped_prefill: Option<&std::sync::Mutex<crate::hybrid::StepEpGroupedPrefill>>,
8482 sliding_gated_moe: bool,
8483 vrows: bool,
8484 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8485 let worker_io = crate::spill_pread::worker_enabled();
8486 let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
8487 if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
8488 e.with_moe_cache(max_block, |cache, _| {
8489 cache.begin_forward_epoch(il, t);
8490 if worker_io {
8491 cache.begin_worker_scope();
8492 }
8493 Ok(())
8494 })?;
8495 }
8496 if let Some(ep) = &m.glm5_ep {
8497 return Self::moe_ffn_glm5_ep(e, m, ep, z, zq8, t, cfg, il, prefill);
8502 }
8503 if m.step_ep.is_some() || m.step_tp.is_some() {
8504 let moe = cfg
8505 .moe
8506 .as_ref()
8507 .ok_or("Step distributed execution requires MoE model metadata")?;
8508 let n_embd = cfg.n_embd as usize;
8509 let n_expert = moe.expert_count as usize;
8510 let n_used = moe.expert_used_count as usize;
8511 let sigmoid = cfg
8512 .sigmoid_router()
8513 .ok_or("Step distributed execution requires the Step sigmoid router")?;
8514 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
8515 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
8516 let grouped_prefill_requested = prefill && step_ep_grouped_prefill_enabled()?;
8517 if grouped_prefill_requested && !step_tp_prefill_enabled()? {
8518 return Err(
8519 "MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into(),
8520 );
8521 }
8522 if grouped_prefill_requested && !step_grouped_prefill_shape(true, prefill, t) {
8523 return Err(format!(
8524 "Step grouped prefill tokens {t} are outside the qualified {}..={} range",
8525 PRIME_MIN_T,
8526 crate::cache::PRIME_CHUNK_MAX_TOKENS,
8527 )
8528 .into());
8529 }
8530 let grouped_decode_shape = step_grouped_decode_shape(prefill, t);
8531 let grouped_prefill_shape =
8532 step_grouped_prefill_shape(grouped_prefill_requested, prefill, t);
8533 if let Some(ep) = m.step_ep.as_ref().filter(|ep| {
8534 ep.grouped_decode.is_some() && (grouped_decode_shape || grouped_prefill_shape)
8535 }) {
8536 let (selected, route_weights) =
8537 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
8538 crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
8539 Self::trace_moe_routes(il, t, &selected, &route_weights)?;
8540 Self::trace_moe_input(e, il, t, n_embd, z)?;
8541 let selected = selected
8542 .iter()
8543 .map(|&expert| expert as usize)
8544 .collect::<Vec<_>>();
8545
8546 e.stream().synchronize()?;
8549 let execute = |state: &mut crate::hybrid::StepEpGroupedDecode| {
8550 state.projection.set_activation_limit(ep.activation_limit)?;
8551 ep.runtime
8552 .refresh_step_grouped_expert_parallel_gate_from_root_device(
8553 ep.experts.e4m3()?,
8554 &mut state.projection,
8555 z,
8556 t,
8557 &selected,
8558 )?;
8559 ep.runtime.refresh_step_grouped_expert_parallel_combine(
8560 &state.projection,
8561 &mut state.combine,
8562 &route_weights,
8563 )?;
8564 ep.runtime.execute_step_grouped_expert_parallel_gate(
8565 ep.experts.e4m3()?,
8566 &mut state.projection,
8567 )?;
8568 ep.runtime.execute_step_grouped_expert_parallel_combine(
8569 &state.projection,
8570 &mut state.combine,
8571 )?;
8572 let mut output = ep.runtime.copy_step_grouped_expert_parallel_combine_root(
8573 &state.projection,
8574 &state.combine,
8575 e,
8576 )?;
8577 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
8578 if prefill {
8579 e.stream().synchronize()?;
8582 }
8583 eprintln!(
8584 "[step-tp-ep-grouped] execute layer={il} tokens={t} devices={:?} \
8585 attention_layout=tensor-parallel expert_layout=expert-parallel \
8586 expert_transport={} native_p2p=true route_control=host-narrow \
8587 input=root-device projection_workspaces=persistent \
8588 combine=root-device output=owning-stage-device \
8589 prefill={prefill} batched_decode=false capacity={} \
8590 performance_claim=false",
8591 ep.devices,
8592 ep.runtime.transport_label(),
8593 state.projection.max_tokens(),
8594 );
8595 Ok::<_, Box<dyn std::error::Error>>(output)
8596 };
8597
8598 if grouped_prefill_shape {
8599 let grouped_prefill = grouped_prefill
8600 .ok_or("Step grouped prefill has no model-scoped executor")?;
8601 let mut shared = grouped_prefill
8602 .lock()
8603 .map_err(|_| "Step grouped prefill state lock is poisoned")?;
8604 let needs_prepare = shared.state.as_ref().is_none_or(|state| {
8605 state.devices != ep.devices
8606 || state.grouped.projection.max_tokens() < t
8607 || state.grouped.projection.input_width() != n_embd
8608 || state.grouped.projection.expert_width()
8609 != moe.expert_ff_length as usize
8610 });
8611 if needs_prepare {
8612 let seed_input = vec![0.0f32; n_embd];
8613 let seed_selected = &selected[..n_used];
8614 let seed_weights = &route_weights[..n_used];
8615 let projection = ep
8616 .runtime
8617 .prepare_step_grouped_expert_parallel_gate_with_capacity(
8618 ep.experts.e4m3()?,
8619 &seed_input,
8620 1,
8621 seed_selected,
8622 ep.activation_limit,
8623 t,
8624 )?;
8625 let combine = ep.runtime.prepare_step_grouped_expert_parallel_combine(
8626 &projection,
8627 seed_weights,
8628 )?;
8629 shared.state = Some(crate::hybrid::StepEpGroupedPrefillState {
8630 devices: ep.devices.clone(),
8631 grouped: crate::hybrid::StepEpGroupedDecode {
8632 projection,
8633 combine,
8634 },
8635 });
8636 eprintln!(
8637 "[step-tp-ep-grouped-prefill] prepare capacity={t} devices={:?} \
8638 shared_across_layers=true performance_claim=false",
8639 ep.devices,
8640 );
8641 }
8642 return execute(
8643 &mut shared
8644 .state
8645 .as_mut()
8646 .expect("Step grouped prefill state prepared above")
8647 .grouped,
8648 );
8649 }
8650
8651 let mut grouped = ep
8652 .grouped_decode
8653 .as_ref()
8654 .expect("grouped decode presence checked above")
8655 .lock()
8656 .map_err(|_| "Step grouped decode state lock is poisoned")?;
8657 return execute(&mut grouped);
8658 }
8659 if grouped_prefill_shape {
8660 return Err(
8661 "Step grouped prefill requires native-P2P expert-owner device arithmetic"
8662 .into(),
8663 );
8664 }
8665 if t >= 16
8680 && crate::step_gemm_prime_on()
8681 && let Some(tp) = &m.step_tp
8682 && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
8683 {
8684 let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1") && t >= 16;
8692 let mut mt = std::time::Instant::now();
8693 let (selected, route_weights) =
8694 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
8695 let sel_i32: Vec<i32> = selected.iter().map(|&x| x as i32).collect();
8696 let d_router = if mprof {
8697 let _ = e.stream().synchronize();
8698 let v = mt.elapsed().as_secs_f64() * 1e3;
8699 mt = std::time::Instant::now();
8700 v
8701 } else {
8702 0.0
8703 };
8704 let mdet =
8718 std::env::var("MEMRA_MOE_DETERM").as_deref() == Ok("1") && t >= 16 && il < 4;
8719 if mdet {
8720 let a = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8721 bank,
8722 e,
8723 z,
8724 t,
8725 &sel_i32,
8726 &route_weights,
8727 n_used,
8728 tp.activation_limit,
8729 )?;
8730 let b = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8731 bank,
8732 e,
8733 z,
8734 t,
8735 &sel_i32,
8736 &route_weights,
8737 n_used,
8738 tp.activation_limit,
8739 )?;
8740 let (ha, hb) = (e.dtoh(&a)?, e.dtoh(&b)?);
8741 let mut md = 0.0f32;
8742 let mut ndiff = 0usize;
8743 for (x, y) in ha.iter().zip(hb.iter()) {
8744 let d = (x - y).abs();
8745 if d > 0.0 {
8746 ndiff += 1;
8747 }
8748 if d > md {
8749 md = d;
8750 }
8751 }
8752 eprintln!(
8753 "[moe-determ] il={il} t={t} maxdiff={md:.3e} \
8754 differing={ndiff}/{} -> {}",
8755 ha.len(),
8756 if ndiff == 0 {
8757 "IDENTICAL"
8758 } else {
8759 "NONDETERMINISTIC"
8760 }
8761 );
8762 }
8763 let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8764 bank,
8765 e,
8766 z,
8767 t,
8768 &sel_i32,
8769 &route_weights,
8770 n_used,
8771 tp.activation_limit,
8772 )?;
8773 let d_gemm = if mprof {
8774 let _ = e.stream().synchronize();
8775 let v = mt.elapsed().as_secs_f64() * 1e3;
8776 mt = std::time::Instant::now();
8777 v
8778 } else {
8779 0.0
8780 };
8781 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
8782 if mprof {
8783 let _ = e.stream().synchronize();
8784 let d_shared = mt.elapsed().as_secs_f64() * 1e3;
8785 eprintln!(
8789 "[moe-prof] il={il} t={t} router={d_router:.1}ms \
8790 gemm={d_gemm:.1}ms shared={d_shared:.1}ms"
8791 );
8792 }
8793 return Ok(output);
8794 }
8795 if t == 1
8796 && crate::tp::step_nvfp4_dev_routes_enabled()?
8797 && crate::tp::step_tp_dev_router_enabled()?
8798 && let Some(tp) = &m.step_tp
8799 && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
8800 {
8801 let (sf, route_norm) = sigmoid;
8802 static D1_ROUTER: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8809 let d1_router = *D1_ROUTER
8810 .get_or_init(|| std::env::var("MEMRA_DEV1_ROUTER").as_deref() == Ok("1"));
8811 if d1_router {
8812 let (sf_h, rn_h) = sigmoid;
8813 let n_ex = m.gate_exps.n_expert;
8814 let act_ct = m.active_count();
8815 let _ = tp.runtime.nvfp4_routes_prestage_with(
8816 bank,
8817 e,
8818 z,
8819 |rank1, in1, sel1, w1| {
8820 let mut guard = DEV1_ROUTER_REPS
8821 .lock()
8822 .map_err(|_| "dev1 router replica lock")?;
8823 let (reps, scratch) =
8824 guard.get_or_insert_with(|| (Default::default(), None));
8825 if !reps.contains_key(&il) {
8826 use cudarc::driver::DevicePtr;
8827 let (g1, p1, a1) = (
8828 rank1.htod(&vec![0.0f32; n_ex * n_embd])?,
8829 rank1.htod(&vec![0.0f32; n_ex])?,
8830 rank1.alloc_u8_uninit(n_ex)?,
8831 );
8832 for (src, dst_len, dst) in [
8833 (
8834 {
8835 let s = e.stream();
8836 let (p, _g) = m.gate_inp.float_data().device_ptr(&s);
8837 p
8838 },
8839 n_ex * n_embd * 4,
8840 {
8841 let s = rank1.stream();
8842 let (p, _g) = g1.device_ptr(&s);
8843 p
8844 },
8845 ),
8846 (
8847 {
8848 let s = e.stream();
8849 let (p, _g) = m.exp_probs_b_dev.device_ptr(&s);
8850 p
8851 },
8852 n_ex * 4,
8853 {
8854 let s = rank1.stream();
8855 let (p, _g) = p1.device_ptr(&s);
8856 p
8857 },
8858 ),
8859 (
8860 {
8861 let s = e.stream();
8862 let (p, _g) = m.active_experts_dev.device_ptr(&s);
8863 p
8864 },
8865 n_ex,
8866 {
8867 let s = rank1.stream();
8868 let (p, _g) = a1.device_ptr(&s);
8869 p
8870 },
8871 ),
8872 ] {
8873 crate::tp::raw_copy_bytes(dst, src, dst_len, rank1)?;
8874 }
8875 rank1.stream().synchronize()?;
8876 reps.insert(il, (g1, p1, a1));
8877 }
8878 if scratch.is_none() {
8879 *scratch = Some(rank1.htod(&vec![0.0f32; n_ex])?);
8880 }
8881 let (g1, p1, a1) = reps.get(&il).expect("armed above");
8882 let logits1 = scratch.as_mut().expect("armed above");
8883 rank1.router_gemv_into(g1, in1, logits1, n_embd, n_ex, 1)?;
8884 rank1.moe_router_sigmoid_topk_into(
8885 logits1, 1, n_ex, n_used, act_ct, p1, a1, sf_h, rn_h, sel1, w1,
8886 )?;
8887 Ok(true)
8888 },
8889 )?;
8890 } else {
8891 let _ = tp.runtime.nvfp4_routes_prestage(bank, e, z)?;
8892 }
8893 #[allow(clippy::type_complexity)] static SELW: std::sync::Mutex<Option<(usize, CudaSlice<i32>, CudaSlice<f32>)>> =
8898 std::sync::Mutex::new(None);
8899 let mut selw = SELW.lock().map_err(|_| "selw lock poisoned")?;
8900 if selw.as_ref().is_none_or(|(d, ..)| *d != e.ctx().ordinal()) {
8901 *selw = Some((
8902 e.ctx().ordinal(),
8903 e.htod_i32(&vec![0i32; n_used])?,
8904 e.htod(&vec![0.0f32; n_used])?,
8905 ));
8906 }
8907 let (_, sel_d, w_d) = selw.as_mut().expect("armed above");
8908 e.moe_router_sigmoid_topk_into(
8909 &logits,
8910 t,
8911 n_expert,
8912 n_used,
8913 m.active_count(),
8914 &m.exp_probs_b_dev,
8915 &m.active_experts_dev,
8916 sf,
8917 route_norm,
8918 sel_d,
8919 w_d,
8920 )?;
8921 crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
8922 if std::env::var("MEMRA_MOE_TRACE").is_ok()
8931 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
8932 {
8933 return Err("MEMRA_MOE_TRACE/MEMRA_MOE_WEIGHT_TRACE cannot trace the \
8934 device-routed step TP walk (selection never returns to host; \
8935 tracing would add a new sync). Route through the host-router \
8936 arm — refused rather than silently dropping rows"
8937 .into());
8938 }
8939 static SHEXP_OV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8943 let shexp_ov = *SHEXP_OV
8944 .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
8945 static SHEXP_D1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8949 let shexp_d1 = *SHEXP_D1
8950 .get_or_init(|| std::env::var("MEMRA_SHEXP_DEV1").as_deref() == Ok("1"))
8951 && tp.runtime.rank_engine(1).is_some();
8952 static TAIL3: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8956 let tail3 =
8957 *TAIL3.get_or_init(|| std::env::var("MEMRA_TAIL_ADD3").as_deref() != Ok("0"));
8958 let mut ov_issued = false;
8959 let mut d1_issued = false;
8960 let mut tail_folded = false;
8961 let mut output = if shexp_d1 {
8962 let rank1 = tp.runtime.rank_engine(1).expect("checked above");
8963 tp.runtime
8964 .run_tensor_parallel_routes_nvfp4_device_routed_prejoin(
8965 bank,
8966 e,
8967 z,
8968 sel_d,
8969 w_d,
8970 n_used,
8971 tp.activation_limit,
8972 || {
8973 d1_issued =
8974 Self::shexp_dev1_issue(e, rank1, m, z, cfg, il, n_embd)?;
8975 Ok(())
8976 },
8977 )?
8978 } else if shexp_ov {
8979 let post_add = if tail3 {
8984 Self::shexp_overlap_tail_ptrs(e, m, cfg, n_embd)?
8985 } else {
8986 None
8987 };
8988 let used_post = post_add.is_some();
8989 let out = tp
8990 .runtime
8991 .run_tensor_parallel_routes_nvfp4_device_routed_prejoin_add3(
8992 bank,
8993 e,
8994 z,
8995 sel_d,
8996 w_d,
8997 n_used,
8998 tp.activation_limit,
8999 || {
9000 ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
9001 Ok(())
9002 },
9003 post_add,
9004 )?;
9005 if used_post && ov_issued {
9010 tail_folded = true; }
9012 out
9013 } else {
9014 tp.runtime.run_tensor_parallel_routes_nvfp4_device_routed(
9015 bank,
9016 e,
9017 z,
9018 sel_d,
9019 w_d,
9020 n_used,
9021 tp.activation_limit,
9022 )?
9023 };
9024 if output.len() != t * n_embd {
9025 return Err(format!(
9026 "Step tp routed output has {} values, expected {t}x{n_embd}",
9027 output.len()
9028 )
9029 .into());
9030 }
9031 if tail_folded {
9032 } else if d1_issued {
9034 Self::shexp_dev1_apply(e, &mut output, n_embd)?;
9035 } else if ov_issued {
9036 Self::shexp_overlap_apply(e, &mut output, n_embd)?;
9037 } else {
9038 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9039 }
9040 static DR_LOGGED: std::sync::atomic::AtomicU64 =
9041 std::sync::atomic::AtomicU64::new(0);
9042 let layer_bit = 1u64 << (il as u64 % 64);
9043 if DR_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9044 == 0
9045 {
9046 eprintln!(
9047 "[step-tp] execute layer={il} tokens={t} devices={:?} \
9048 expert_transport={} native_p2p={} router=device \
9049 activation=host-canonical accumulation=host-canonical \
9050 output=e-device io=device performance_claim=false \
9051 (logged once per layer)",
9052 tp.devices,
9053 tp.runtime.transport_label(),
9054 tp.runtime.native_p2p(),
9055 );
9056 }
9057 return Ok(output);
9058 }
9059 let automatic_ep_device_router = crate::tp::parallel_ep_device_router_enabled()?;
9060 let automatic_ep_q8_act = crate::tp::parallel_ep_q8_act_enabled()?;
9061 let automatic_ep_q8_scope = crate::tp::parallel_ep_q8_scope()?;
9062 let automatic_ep_q8_active =
9063 automatic_ep_q8_act && t <= crate::tp::NVFP4_EP_Q8_BATCH_CAP;
9064 if automatic_ep_q8_scope.is_some() && !automatic_ep_q8_act {
9065 return Err(
9066 "MEMRA_PARALLEL_EP_Q8_SCOPE requires MEMRA_PARALLEL_EP_Q8_ACT=1".into(),
9067 );
9068 }
9069 if automatic_ep_q8_act && !automatic_ep_device_router {
9070 return Err(
9071 "MEMRA_PARALLEL_EP_Q8_ACT=1 requires MEMRA_PARALLEL_EP_DEVICE_ROUTER=1".into(),
9072 );
9073 }
9074 if automatic_ep_q8_act && m.step_ep.as_ref().is_none_or(|ep| !ep.nvfp4_device_routes) {
9075 return Err(
9076 "MEMRA_PARALLEL_EP_Q8_ACT=1 requires automatic W4A16 whole-expert EP".into(),
9077 );
9078 }
9079 if t <= crate::tp::NVFP4_EP_DEVICE_ROUTER_BATCH_CAP
9080 && automatic_ep_device_router
9081 && let Some(ep) = &m.step_ep
9082 && ep.nvfp4_device_routes
9083 {
9084 let bank = match &ep.experts {
9085 crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
9086 crate::hybrid::StepEpExpertBank::E4m3(_) => {
9087 return Err("W4A16 device-routed EP reached an E4M3 expert bank".into());
9088 }
9089 };
9090 let pairs = t
9091 .checked_mul(n_used)
9092 .ok_or("W4A16 device-routed EP pair count overflow")?;
9093 let capacity = crate::tp::NVFP4_EP_DEVICE_BATCH_CAP * n_used;
9094 type EpSelwByDevice =
9098 std::collections::HashMap<usize, (usize, CudaSlice<i32>, CudaSlice<f32>)>;
9099 static EP_SELW: std::sync::Mutex<Option<EpSelwByDevice>> =
9100 std::sync::Mutex::new(None);
9101 let mut selw = EP_SELW
9102 .lock()
9103 .map_err(|_| "automatic EP device-router workspace lock poisoned")?;
9104 let device = e.ctx().ordinal();
9105 let workspaces = selw.get_or_insert_with(Default::default);
9106 if workspaces
9107 .get(&device)
9108 .is_none_or(|(cap, ..)| *cap < capacity)
9109 {
9110 workspaces.insert(
9111 device,
9112 (
9113 capacity,
9114 e.htod_i32(&vec![0i32; capacity])?,
9115 e.htod(&vec![0.0f32; capacity])?,
9116 ),
9117 );
9118 }
9119 let (_, sel_d, w_d) = workspaces.get_mut(&device).expect("armed above");
9120 let (sf, route_norm) = sigmoid;
9121 e.moe_router_sigmoid_topk_into(
9122 &logits,
9123 t,
9124 n_expert,
9125 n_used,
9126 m.active_count(),
9127 &m.exp_probs_b_dev,
9128 &m.active_experts_dev,
9129 sf,
9130 route_norm,
9131 sel_d,
9132 w_d,
9133 )?;
9134 crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
9135 static SHEXP_OV_AUTO: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9136 let shexp_ov = t == 1
9137 && *SHEXP_OV_AUTO
9138 .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
9139 let mut ov_issued = false;
9140 let mut output = if shexp_ov {
9141 ep.runtime
9142 .run_routed_experts_nvfp4_w4a16_device_routed_prejoin(
9143 bank,
9144 e,
9145 z,
9146 sel_d,
9147 w_d,
9148 t,
9149 n_used,
9150 ep.activation_limit,
9151 || {
9152 ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
9153 Ok(())
9154 },
9155 )?
9156 } else {
9157 ep.runtime.run_routed_experts_nvfp4_w4a16_device_routed(
9158 bank,
9159 e,
9160 z,
9161 sel_d,
9162 w_d,
9163 t,
9164 n_used,
9165 ep.activation_limit,
9166 )?
9167 };
9168 if output.len() != t * n_embd {
9169 return Err(format!(
9170 "W4A16 device-routed EP output has {} values, expected \
9171 {t}x{n_embd}={}",
9172 output.len(),
9173 t * n_embd,
9174 )
9175 .into());
9176 }
9177 if ov_issued {
9178 Self::shexp_overlap_apply(e, &mut output, n_embd)?;
9179 } else {
9180 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9181 }
9182 static DEVICE_ROUTER_LOGGED: std::sync::atomic::AtomicU64 =
9183 std::sync::atomic::AtomicU64::new(0);
9184 let layer_bit = 1u64 << (il as u64 % 64);
9185 if DEVICE_ROUTER_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
9186 & layer_bit
9187 == 0
9188 {
9189 eprintln!(
9190 "[parallel-ep] execute layer={il} tokens={t} devices={:?} \
9191 router=device expert_transport={} native_p2p={} \
9192 activation=bf16-rounded accumulation={} output=e-device \
9193 performance_claim=false (logged once per layer)",
9194 ep.devices,
9195 ep.runtime.transport_label(),
9196 ep.runtime.native_p2p(),
9197 if automatic_ep_q8_active {
9198 "token-slot-order-q8"
9199 } else {
9200 "token-slot-order"
9201 },
9202 );
9203 }
9204 debug_assert!(pairs <= capacity);
9205 return Ok(output);
9206 }
9207 static ROUTE_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9211 static ROUTE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9212 let route_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
9213 let route_started = route_timing.then(std::time::Instant::now);
9214 let (selected, route_weights, input) = Self::moe_route_sigmoid_with_input(
9215 e,
9216 &logits,
9217 z,
9218 t,
9219 n_embd,
9220 n_expert,
9221 n_used,
9222 m.exp_probs_b.as_deref(),
9223 sigmoid,
9224 m.active_experts.as_deref(),
9225 )?;
9226 if let Some(started) = route_started {
9227 use std::sync::atomic::Ordering;
9228 let ns = ROUTE_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
9229 + started.elapsed().as_nanos() as u64;
9230 let calls = ROUTE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
9231 if calls.is_multiple_of(430) {
9232 eprintln!(
9233 "[moe-route-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
9234 ns as f64 / 1.0e6,
9235 ns as f64 / calls as f64 / 1.0e3,
9236 );
9237 }
9238 }
9239 crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
9240 Self::trace_moe_routes(il, t, &selected, &route_weights)?;
9241 Self::trace_moe_input(e, il, t, n_embd, z)?;
9242 let selected = selected
9243 .iter()
9244 .map(|&expert| expert as usize)
9245 .collect::<Vec<_>>();
9246 if t <= crate::tp::NVFP4_EP_DEVICE_BATCH_CAP
9247 && let Some(ep) = &m.step_ep
9248 && ep.nvfp4_device_routes
9249 {
9250 let bank = match &ep.experts {
9251 crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
9252 crate::hybrid::StepEpExpertBank::E4m3(_) => {
9253 return Err("W4A16 NVFP4 device EP reached an E4M3 expert bank".into());
9254 }
9255 };
9256 let mut output = ep.runtime.run_routed_experts_nvfp4_w4a16_device_io(
9257 bank,
9258 e,
9259 z,
9260 t,
9261 &selected,
9262 &route_weights,
9263 n_used,
9264 ep.activation_limit,
9265 )?;
9266 if output.len() != t * n_embd {
9267 return Err(format!(
9268 "W4A16 NVFP4 EP routed output has {} values, expected \
9269 {t}x{n_embd}={}",
9270 output.len(),
9271 t * n_embd,
9272 )
9273 .into());
9274 }
9275 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9276 static W4A16_EP_LOGGED: std::sync::atomic::AtomicU64 =
9277 std::sync::atomic::AtomicU64::new(0);
9278 let layer_bit = 1u64 << (il as u64 % 64);
9279 if W4A16_EP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
9280 & layer_bit
9281 == 0
9282 {
9283 eprintln!(
9284 "[step-ep] execute layer={il} tokens={t} devices={:?} \
9285 expert_transport={} native_p2p={} activation=bf16-rounded \
9286 accumulation={} output=e-device \
9287 performance_claim=false (logged once per layer)",
9288 ep.devices,
9289 ep.runtime.transport_label(),
9290 ep.runtime.native_p2p(),
9291 if t == 1 {
9292 "owner-grouped-rank-order"
9293 } else {
9294 "token-slot-order"
9295 },
9296 );
9297 }
9298 return Ok(output);
9299 }
9300 if t == 1
9305 && crate::tp::step_nvfp4_dev_routes_enabled()?
9306 && let Some(tp) = &m.step_tp
9307 && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
9308 {
9309 let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_device_io(
9310 bank,
9311 e,
9312 z,
9313 &selected,
9314 &route_weights,
9315 n_used,
9316 tp.activation_limit,
9317 )?;
9318 if output.len() != t * n_embd {
9319 return Err(format!(
9320 "Step tp routed output has {} values, expected {t}x{n_embd}",
9321 output.len()
9322 )
9323 .into());
9324 }
9325 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9326 static IO_LOGGED: std::sync::atomic::AtomicU64 =
9327 std::sync::atomic::AtomicU64::new(0);
9328 let layer_bit = 1u64 << (il as u64 % 64);
9329 if IO_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9330 == 0
9331 {
9332 eprintln!(
9333 "[step-tp] execute layer={il} tokens={t} devices={:?} \
9334 expert_transport={} native_p2p={} activation=host-canonical \
9335 accumulation=host-canonical output=e-device io=device \
9336 performance_claim=false (logged once per layer)",
9337 tp.devices,
9338 tp.runtime.transport_label(),
9339 tp.runtime.native_p2p(),
9340 );
9341 }
9342 return Ok(output);
9343 }
9344 let (routed, mode, devices, transport, native_p2p) = if let Some(tp) = &m.step_tp {
9345 (
9346 match &tp.experts {
9347 crate::hybrid::StepTpExpertBank::E4m3(bank) => {
9348 tp.runtime.run_tensor_parallel_routes(
9349 bank,
9350 &input,
9351 t,
9352 &selected,
9353 &route_weights,
9354 n_used,
9355 )?
9356 }
9357 crate::hybrid::StepTpExpertBank::Nvfp4(bank) => {
9358 if t == 1 && crate::tp::step_nvfp4_dev_routes_enabled()? {
9359 tp.runtime.run_tensor_parallel_routes_nvfp4_device(
9360 bank,
9361 &input,
9362 &selected,
9363 &route_weights,
9364 n_used,
9365 tp.activation_limit,
9366 )?
9367 } else {
9368 tp.runtime.run_tensor_parallel_routes_nvfp4(
9369 bank,
9370 &input,
9371 t,
9372 &selected,
9373 &route_weights,
9374 n_used,
9375 tp.activation_limit,
9376 )?
9377 }
9378 }
9379 },
9380 "tp",
9381 &tp.devices,
9382 tp.runtime.transport_label(),
9383 tp.runtime.native_p2p(),
9384 )
9385 } else {
9386 let ep = m
9387 .step_ep
9388 .as_ref()
9389 .ok_or("Step distributed runtime has no EP or TP state")?;
9390 (
9391 match &ep.experts {
9392 crate::hybrid::StepEpExpertBank::E4m3(bank) => {
9393 ep.runtime.run_routed_experts(
9394 bank,
9395 &input,
9396 t,
9397 &selected,
9398 &route_weights,
9399 n_used,
9400 ep.activation_limit,
9401 )?
9402 }
9403 crate::hybrid::StepEpExpertBank::Nvfp4(bank) => {
9404 ep.runtime.run_routed_experts_nvfp4(
9405 bank,
9406 &input,
9407 t,
9408 &selected,
9409 &route_weights,
9410 n_used,
9411 ep.activation_limit,
9412 )?
9413 }
9414 },
9415 if ep.configured_by_tp { "tp-ep" } else { "ep" },
9416 &ep.devices,
9417 ep.runtime.transport_label(),
9418 ep.runtime.native_p2p(),
9419 )
9420 };
9421 if routed.len() != t * n_embd {
9422 return Err(format!(
9423 "Step {mode} routed output has {} values, expected {t}x{n_embd}",
9424 routed.len()
9425 )
9426 .into());
9427 }
9428 let mut output = e.htod(&routed)?;
9429 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9430 static STEP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9433 let layer_bit = 1u64 << (il as u64 % 64);
9434 if STEP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9435 == 0
9436 {
9437 eprintln!(
9438 "[step-{mode}] execute layer={il} tokens={t} devices={devices:?} \
9439 expert_transport={transport} native_p2p={native_p2p} \
9440 activation={} accumulation={} output={} \
9441 performance_claim=false (logged once per layer)",
9442 if let Some(ep) = &m.step_ep {
9443 ep.runtime.expert_activation_label()
9444 } else {
9445 "host-canonical"
9446 },
9447 if let Some(ep) = &m.step_ep {
9448 ep.runtime.expert_accumulation_label()
9449 } else {
9450 "host-canonical"
9451 },
9452 if let Some(ep) = &m.step_ep {
9453 ep.runtime.expert_output_label()
9454 } else {
9455 "host-accumulated"
9456 },
9457 );
9458 if let Some(ep) = &m.step_ep
9459 && let Some(limit) = ep.activation_limit
9460 {
9461 eprintln!(
9462 "[step-ep-clamp] execute layer={il} tokens={t} routed_clamp={limit} \
9463 formula=min-silu-times-clamped-up performance_claim=false"
9464 );
9465 }
9466 }
9467 return Ok(output);
9468 }
9469 if Self::sigmoid_resident_dev_eligible(e, m, cfg, sliding_gated_moe) {
9470 let moe = cfg.moe.as_ref().unwrap();
9471 let n_expert = moe.expert_count as usize;
9472 let n_used = moe.expert_used_count as usize;
9473 let sigmoid = cfg.sigmoid_router().unwrap();
9474 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9475 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
9476 return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
9477 }
9478 if prefill && t > MOE_DEV_MAX_T && cfg.sigmoid_router().is_some() && cfg.glm5.is_some() {
9500 static GPF_ANNOUNCED: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
9503 let enabled = moe_grouped_prefill_enabled();
9504 let bit = 1u8 << u8::from(enabled);
9505 if GPF_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
9506 eprintln!(
9507 "[moe-grouped-prefill] flag={} t={t} il={il} (announce printed in both \
9508 arms; engagement is the per-layer execute line + the dispatch counter)",
9509 if enabled { "on" } else { "off" },
9510 );
9511 }
9512 if enabled
9513 && let Some(out) = Self::moe_ffn_grouped_prefill_sigmoid(e, m, z, t, cfg, il)?
9514 {
9515 return Ok(out);
9516 }
9517 }
9518 if t > 1 && moe_grouped_enabled(cfg, prefill) {
9521 let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
9522 if std::env::var("MEMRA_MOE_GATE").is_ok() {
9527 let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
9528 let g_host = e.dtoh(&grouped_out)?;
9529 let s_host = e.dtoh(&seq_out)?;
9530 let g_bytes: &[u8] = unsafe {
9531 std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
9532 };
9533 let s_bytes: &[u8] = unsafe {
9534 std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
9535 };
9536 if g_bytes == s_bytes {
9537 println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
9538 } else {
9539 let diffs = g_host
9540 .iter()
9541 .zip(s_host.iter())
9542 .enumerate()
9543 .filter(|(_, (a, b))| a != b)
9544 .count();
9545 let maxdiff = g_host
9546 .iter()
9547 .zip(s_host.iter())
9548 .map(|(a, b)| (a - b).abs())
9549 .fold(0.0f32, f32::max);
9550 panic!(
9551 "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
9552 g_host.len()
9553 );
9554 }
9555 }
9556 return Ok(grouped_out);
9557 }
9558 Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block, vrows)
9559 }
9560
9561 fn sigmoid_resident_dev_eligible(
9562 e: &Engine,
9563 m: &MoeWeights,
9564 cfg: &ModelConfig,
9565 sliding_gated_moe: bool,
9566 ) -> bool {
9567 let Some(moe) = cfg.moe.as_ref() else {
9568 return false;
9569 };
9570 static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9573 let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
9574 std::env::var("MEMRA_MOE_STATS").is_ok()
9575 || std::env::var("MEMRA_MOE_TRACE").is_ok()
9576 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
9577 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
9578 || std::env::var("MEMRA_MOE_GATE").is_ok()
9579 });
9580 let resident_layout_supported = m.dev_exps.as_ref().is_some_and(|dev| {
9581 if dev.dev != e.ctx().ordinal() {
9582 return false;
9583 }
9584 let q8 = moe_q8_enabled_for_model(cfg, m);
9585 let fp8 = dev.fp8_blk.is_some()
9586 && m.gate_exps.qtype == crate::QT_F8_E4M3_BLK
9587 && m.up_exps.qtype == crate::QT_F8_E4M3_BLK
9588 && m.down_exps.qtype == crate::QT_F8_E4M3_BLK;
9589 q8 || fp8
9590 });
9591 sliding_gated_moe
9592 && sigmoid_router_enabled()
9593 && moe_dev_enabled()
9594 && moe_slab_enabled()
9595 && !observation_mode
9596 && moe.expert_used_count <= 8
9597 && m.has_uniform_expert_layout()
9598 && m.gate_exps.macros.is_none()
9599 && m.up_exps.macros.is_none()
9600 && m.down_exps.macros.is_none()
9601 && !m.has_macros
9602 && resident_layout_supported
9603 }
9604
9605 pub(crate) fn moe_ffn_sequential(
9607 e: &Engine,
9608 m: &MoeWeights,
9609 z: &CudaSlice<f32>,
9610 t: usize,
9611 cfg: &ModelConfig,
9612 il: u16,
9613 max_block: usize,
9614 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9615 Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block, false)
9616 }
9617
9618 fn moe_router_logits(
9622 e: &Engine,
9623 m: &MoeWeights,
9624 z: &CudaSlice<f32>,
9625 t: usize,
9626 cfg: &ModelConfig,
9627 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9628 if t < PRIME_MIN_T {
9629 if crate::router_kernel_on() {
9631 e.router_gemv(
9632 m.gate_inp.float_data(),
9633 z,
9634 cfg.n_embd as usize,
9635 m.gate_exps.n_expert,
9636 t,
9637 )
9638 } else {
9639 e.matmul_decode_exact(&m.gate_inp, z, t)
9640 }
9641 } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
9642 e.router_gemv(
9643 m.gate_inp.float_data(),
9644 z,
9645 cfg.n_embd as usize,
9646 m.gate_exps.n_expert,
9647 t,
9648 )
9649 } else {
9650 e.matmul(&m.gate_inp, z, t)
9651 }
9652 }
9653
9654 fn trace_moe_routes(
9662 il: u16,
9663 t: usize,
9664 sel_all: &[u32],
9665 weights: &[f32],
9666 ) -> Result<(), Box<dyn std::error::Error>> {
9667 use std::io::Write as _;
9668 if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
9669 let mut f = std::fs::OpenOptions::new()
9670 .create(true)
9671 .append(true)
9672 .open(path)?;
9673 let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
9674 writeln!(f, "{} {} {}", il, t, ids.join(","))?;
9675 }
9676 if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
9677 let mut f = std::fs::OpenOptions::new()
9678 .create(true)
9679 .append(true)
9680 .open(path)?;
9681 let pairs: Vec<String> = sel_all
9682 .iter()
9683 .zip(weights)
9684 .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
9685 .collect();
9686 writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
9687 }
9688 Ok(())
9689 }
9690
9691 #[allow(clippy::too_many_arguments)]
9692 fn trace_sigmoid_router_logits(
9693 e: &Engine,
9694 il: u16,
9695 t: usize,
9696 n_expert: usize,
9697 n_used: usize,
9698 logits: &CudaSlice<f32>,
9699 m: &MoeWeights,
9700 (scaling_factor, route_norm): (f32, bool),
9701 ) -> Result<(), Box<dyn std::error::Error>> {
9702 if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
9703 return Ok(());
9704 }
9705 let logits = e.dtoh(logits)?;
9706 let active: Vec<u8> = m
9707 .active_experts
9708 .as_ref()
9709 .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
9710 .unwrap_or_else(|| vec![1; n_expert]);
9711 let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
9712 crate::sigrouter_contract::capture_served_logits(
9713 il as u32,
9714 t,
9715 n_expert,
9716 n_used,
9717 scaling_factor,
9718 route_norm,
9719 &active,
9720 &bias,
9721 &logits,
9722 )?;
9723 Ok(())
9724 }
9725
9726 fn trace_moe_input(
9731 e: &Engine,
9732 il: u16,
9733 t: usize,
9734 n_embd: usize,
9735 z: &CudaSlice<f32>,
9736 ) -> Result<(), Box<dyn std::error::Error>> {
9737 use std::io::Write as _;
9738 let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
9739 return Ok(());
9740 };
9741 let values = active_matrix_values(z.len(), t, n_embd, "MoE input trace activation")?;
9742 let host = e.dtoh_view(&z.slice(0..values))?;
9743 let bytes = unsafe {
9744 std::slice::from_raw_parts(
9745 host.as_ptr().cast::<u8>(),
9746 host.len() * std::mem::size_of::<f32>(),
9747 )
9748 };
9749 let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
9750 let mut state = state
9751 .lock()
9752 .map_err(|_| "MoE input trace writer lock is poisoned")?;
9753 if state.is_none() {
9754 let dir = std::path::PathBuf::from(&dir);
9755 std::fs::create_dir_all(&dir)?;
9756 let index = std::fs::OpenOptions::new()
9757 .create(true)
9758 .append(true)
9759 .open(dir.join("index.jsonl"))?;
9760 *state = Some(MoeInputTraceWriter {
9761 dir,
9762 index,
9763 payloads: std::collections::HashMap::new(),
9764 });
9765 }
9766 let writer = state.as_mut().unwrap();
9767 if writer.dir != std::path::Path::new(&dir) {
9768 return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
9769 }
9770 let file_name = format!("layer-{il:03}.f32");
9771 if !writer.payloads.contains_key(&il) {
9772 let payload = std::fs::OpenOptions::new()
9773 .create(true)
9774 .append(true)
9775 .open(writer.dir.join(&file_name))?;
9776 let offset = payload.metadata()?.len();
9777 writer.payloads.insert(il, (payload, offset));
9778 }
9779 let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
9780 let row_offset = *offset;
9781 payload.write_all(bytes)?;
9782 *offset += bytes.len() as u64;
9783 writeln!(
9784 writer.index,
9785 "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
9786 \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
9787 \"payload_bytes\":{}}}",
9788 bytes.len()
9789 )?;
9790 Ok(())
9791 }
9792
9793 #[allow(clippy::too_many_arguments)]
9794 #[allow(clippy::too_many_arguments)]
9795 pub(crate) fn moe_ffn_sequential_zq8(
9797 e: &Engine,
9798 m: &MoeWeights,
9799 z: &CudaSlice<f32>,
9800 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
9801 t: usize,
9802 cfg: &ModelConfig,
9803 il: u16,
9804 max_block: usize,
9805 vrows: bool,
9806 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9807 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
9808 let moe = cfg.moe.as_ref().unwrap();
9809 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);
9816 debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
9817 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);
9820
9821 let lim_exp = cfg.clamp_exp_at(il as u32);
9824 let lim_shexp = cfg.clamp_shexp_at(il as u32);
9825 let use_cache = Engine::moe_cache_enabled();
9826 let uniform_experts = m.has_uniform_expert_layout();
9827 let moe_q8 = uniform_experts && moe_q8_enabled_for_model(cfg, m);
9828 let cpu_expert_requested = crate::cpu_experts::configured();
9835 if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
9836 return Err(std::io::Error::other(
9837 "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
9838 )
9839 .into());
9840 }
9841 let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
9842 let freeze_cpu_residency = cpu_expert_requested
9848 && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
9849 let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
9850 .ok()
9851 .and_then(|value| value.parse::<usize>().ok())
9852 .is_some_and(|tokens| tokens > 0);
9853 if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
9854 e.freeze_moe_cache();
9855 }
9856 let cache_frozen = use_cache && e.moe_cache_frozen();
9857 let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
9858
9859 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9862 if let Some(sig) = cfg.sigmoid_router() {
9863 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
9864 }
9865
9866 let no_exp_macros = m.gate_exps.macros.is_none()
9905 && m.up_exps.macros.is_none()
9906 && m.down_exps.macros.is_none();
9907 if cfg.sigmoid_router().is_none()
9911 && cfg.m3.is_none()
9912 && cfg.hy3.is_none()
9913 && !cfg.swiglu_clamped_at(il as u32)
9914 && no_exp_macros
9915 && t > MOE_DEV_MAX_T
9919 && m.dev_exps.is_some()
9920 && moe_q8_enabled_for_model(cfg, m)
9921 && std::env::var("MEMRA_MOE_PAIRS")
9922 .map(|v| v != "0")
9923 .unwrap_or(true)
9924 && std::env::var("MEMRA_MOE_STATS").is_err()
9925 {
9926 return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
9927 }
9928
9929 let dev_ok = uniform_experts
9947 && cfg.sigmoid_router().is_none()
9948 && cfg.m3.is_none()
9949 && cfg.hy3.is_none()
9950 && !cfg.swiglu_clamped_at(il as u32);
9951 let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
9955 || std::env::var("MEMRA_MOE_TRACE").is_ok()
9956 || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
9957 || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
9958 if dev_ok
9959 && t <= MOE_DEV_MAX_T
9960 && m.dev_exps.is_some()
9961 && n_used <= 8
9962 && moe_dev_enabled()
9963 && !observe_routes
9964 {
9965 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
9966 }
9967 if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
9968 let row_ok = e.with_moe_cache(max_block, |c, eng| {
9969 if moe_prewarm_enabled() {
9970 c.prewarm_layer(il, m, eng)?;
9971 }
9972 Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
9973 })?;
9974 if row_ok {
9975 return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
9976 }
9977 }
9978
9979 let slab_local = m
9985 .dev_exps
9986 .as_ref()
9987 .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
9988 let slab_bases = slab_local.map(|d| {
9989 use cudarc::driver::DevicePtr;
9990 let s = e.stream();
9991 let (pg, _g0) = d.gate.device_ptr(&s);
9992 let (pu, _g1) = d.up.device_ptr(&s);
9993 let (pd, _g2) = d.down.device_ptr(&s);
9994 (pg, pu, pd)
9995 });
9996 let vrows_dev = vrows
10012 && t >= 2
10013 && crate::moe_vrows_dev_tables_on()
10014 && slab_bases.is_some()
10015 && moe_q8
10016 && uniform_experts
10017 && n_used <= 8
10018 && cfg.sigmoid_router().is_some()
10019 && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10020 && !cpu_hybrid
10021 && sigmoid_router_enabled()
10022 && !observe_routes
10023 && !memra_reference::hidden_trace::enabled()
10024 && !crate::moesd::capture_active();
10025 let mut sel_dev: Option<(CudaSlice<i32>, CudaSlice<f32>)> = None;
10028 let (sel_all, w_all, routed_cpu_input) = if vrows_dev {
10029 let (sf, route_norm) = cfg
10030 .sigmoid_router()
10031 .expect("vrows_dev carries cfg.sigmoid_router().is_some()");
10032 sel_dev = Some(e.moe_router_sigmoid_topk(
10033 &logits,
10034 t,
10035 n_expert,
10036 n_used,
10037 m.active_count(),
10038 &m.exp_probs_b_dev,
10039 &m.active_experts_dev,
10040 sf,
10041 route_norm,
10042 )?);
10043 crate::MOE_VROWS_ROUTER_SYNCS_AVOIDED
10044 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10045 (Vec::new(), Vec::new(), None)
10046 } else if let Some(sig) = cfg.sigmoid_router() {
10047 if cpu_hybrid {
10048 let (sel, w, input) = Self::moe_route_sigmoid_with_input(
10049 e,
10050 &logits,
10051 z,
10052 t,
10053 n_embd,
10054 n_expert,
10055 n_used,
10056 m.exp_probs_b.as_deref(),
10057 sig,
10058 m.active_experts.as_deref(),
10059 )?;
10060 (sel, w, Some(input))
10061 } else {
10062 let (sel, w) =
10063 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
10064 (sel, w, None)
10065 }
10066 } else {
10067 let (sel, w) =
10068 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
10069 (sel, w, None)
10070 };
10071 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
10072 if memra_reference::hidden_trace::enabled() {
10073 memra_reference::hidden_trace::emit_last_row(
10074 "router",
10075 il as i64,
10076 t,
10077 n_expert,
10078 &e.dtoh(&logits)?,
10079 );
10080 let last = (t - 1) * n_used;
10081 let mut route = Vec::with_capacity(n_used * 2);
10082 for slot in 0..n_used {
10083 route.push(sel_all[last + slot] as f32);
10084 route.push(w_all[last + slot]);
10085 }
10086 memra_reference::hidden_trace::emit_last_row("route", il as i64, 1, n_used * 2, &route);
10087 }
10088
10089 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
10093 Self::trace_moe_input(e, il, t, n_embd, z)?;
10094
10095 let worker_disk_prefetch =
10107 cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
10108 let promote_worker_h2d =
10109 t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
10110 if promote_worker_h2d {
10111 let mut selected_blocks = Vec::with_capacity(n_used * 3);
10112 for &ex in sel_all.iter().take(n_used) {
10113 let ex = ex as u16;
10114 selected_blocks.extend([
10115 BlockId::new(il, PROJ_GATE, ex),
10116 BlockId::new(il, PROJ_UP, ex),
10117 BlockId::new(il, PROJ_DOWN, ex),
10118 ]);
10119 }
10120 for &ex in sel_all.iter().take(n_used) {
10121 Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
10122 }
10123 e.with_moe_cache(max_block, |cache, eng| {
10124 cache.promote_worker_reads_at_safe_boundary(
10125 &selected_blocks,
10126 &selected_blocks,
10127 eng,
10128 )?;
10129 Ok(())
10130 })?;
10131 }
10132
10133 if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
10136 let mut cnt = vec![0u32; n_expert];
10137 for &s in sel_all.iter() {
10138 cnt[s as usize] += 1;
10139 }
10140 let total = sel_all.len() as f64;
10141 let mut h = 0.0f64;
10142 let mut active = 0usize;
10143 for &c in &cnt {
10144 if c > 0 {
10145 active += 1;
10146 let p = c as f64 / total;
10147 h -= p * p.log2();
10148 }
10149 }
10150 let maxc = cnt.iter().copied().max().unwrap_or(0);
10151 println!(
10152 "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
10153 il,
10154 t,
10155 sel_all.len(),
10156 active,
10157 n_expert,
10158 h,
10159 (n_expert as f64).log2(),
10160 total / active.max(1) as f64,
10161 maxc
10162 );
10163 }
10164
10165 let gdec_may_fire = uniform_experts
10178 && use_cache
10179 && n_used <= 8
10180 && gdec_enabled()
10181 && !cfg.swiglu_clamped_at(il as u32);
10182 let slab_fused_may_fire = slab_bases.is_some()
10210 && n_used <= 8
10211 && gdec_enabled()
10212 && !cfg.swiglu_clamped_at(il as u32)
10213 && cfg.m3.is_none()
10214 && no_exp_macros
10215 && moe_q8;
10216 let fused_epi_common = n_used <= 8
10251 && moe_q8
10252 && cfg.m3.is_none()
10253 && cfg.sigmoid_router().is_some()
10254 && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10255 && moe_fused_epi_enabled();
10256 let fused_epi_may_fire = fused_epi_common
10257 && uniform_experts
10258 && use_cache
10259 && cache_dispatch
10260 && slab_local.is_none();
10261 let fused_epi_slab_may_fire = fused_epi_common && slab_bases.is_some();
10262 let vrows_fires = vrows
10278 && t >= 2
10279 && slab_bases.is_some()
10280 && moe_q8
10281 && uniform_experts
10282 && n_used <= 8
10283 && cfg.sigmoid_router().is_some()
10284 && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10285 && !cpu_hybrid;
10286 let mut moe_out = if gdec_may_fire
10291 || slab_fused_may_fire
10292 || fused_epi_may_fire
10293 || fused_epi_slab_may_fire
10294 || vrows_fires
10295 {
10296 e.uninit(t * n_embd)?
10297 } else {
10298 e.zeros(t * n_embd)?
10299 };
10300 let cpu_input = if cpu_hybrid {
10303 Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
10304 } else {
10305 None
10306 };
10307
10308 if vrows_dev && !vrows_fires {
10313 return Err(
10314 "MEMRA_MOE_VROWS_DEV_TABLES routed device-only but the verify-rows arm did not \
10315 fire: the door-D and vrows_fires predicates disagree"
10316 .into(),
10317 );
10318 }
10319 if vrows_fires {
10320 let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10321 return Err(
10322 "verify-rows MoE arm fired without a live PRE clamp: the predicate and \
10323 the dispatch disagree"
10324 .into(),
10325 );
10326 };
10327 let bases = slab_bases.expect("vrows_fires carries slab_bases.is_some()");
10328 let sel = match sel_dev.as_ref() {
10329 Some((si, sw)) => VrowsSel::Dev(si, sw),
10330 None => VrowsSel::Host(&sel_all, &w_all),
10331 };
10332 Self::moe_vrows_pairs_q8(
10333 e,
10334 m,
10335 z,
10336 sel,
10337 il,
10338 bases,
10339 t,
10340 n_embd,
10341 n_ff_exp,
10342 n_used,
10343 limit,
10344 &mut moe_out,
10345 )?;
10346 if memra_reference::hidden_trace::enabled() {
10347 memra_reference::hidden_trace::emit_last_row(
10348 "routed",
10349 il as i64,
10350 t,
10351 n_embd,
10352 &e.dtoh(&moe_out)?,
10353 );
10354 }
10355 Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
10356 return Ok(moe_out);
10357 }
10358
10359 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;
10367 let mut scratch_u: Option<CudaSlice<u8>> = None;
10368 let mut scratch_d: Option<CudaSlice<u8>> = None;
10369 let page_window = moe_page_prefetch_window();
10377
10378 for tok in 0..t {
10381 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
10382 let w = &w_all[tok * n_used..(tok + 1) * n_used];
10383 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10385
10386 let no_macros = m.gate_exps.macros.is_none()
10400 && m.up_exps.macros.is_none()
10401 && m.down_exps.macros.is_none();
10402 if slab_fused_may_fire {
10412 let (pg, pu, pd) = slab_bases.unwrap();
10413 let mut gp = [0u64; 8];
10414 let mut up = [0u64; 8];
10415 let mut dp = [0u64; 8];
10416 for (j, &ex) in sel.iter().enumerate() {
10417 let ex = ex as usize;
10418 gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
10419 up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
10420 dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
10421 }
10422 let mut wv = [0f32; 8];
10423 wv[..n_used].copy_from_slice(w);
10424 if tok_q8.is_none() {
10425 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10426 }
10427 let (zq, zd) = tok_q8.as_ref().unwrap();
10428 let act = e.moe_gate_up_silu8_q8(
10429 crate::WPtr8(gp),
10430 crate::WPtr8(up),
10431 zq,
10432 zd,
10433 n_embd,
10434 n_ff_exp,
10435 n_used,
10436 m.gate_exps.qtype,
10437 m.up_exps.qtype,
10438 m.gate_exps.row_bytes,
10439 m.up_exps.row_bytes,
10440 )?;
10441 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
10442 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10443 e.moe_down8_fma_q8(
10444 crate::WPtr8(dp),
10445 crate::F32x8(wv),
10446 &aq2,
10447 &ad2,
10448 &mut dst,
10449 n_ff_exp,
10450 n_embd,
10451 n_used,
10452 m.down_exps.qtype,
10453 m.down_exps.row_bytes,
10454 )?;
10455 continue;
10456 }
10457 if fused_epi_slab_may_fire {
10462 let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10463 return Err(
10464 "fused MoE epilogue (slab) fired without a live PRE clamp: the \
10465 predicate and the dispatch disagree"
10466 .into(),
10467 );
10468 };
10469 let (pg, pu, pd) = slab_bases.unwrap();
10470 let mut g = [0u64; 8];
10471 let mut u = [0u64; 8];
10472 let mut d = [0u64; 8];
10473 for (j, &ex) in sel.iter().enumerate() {
10474 let ex = ex as usize;
10475 g[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
10476 u[j] = pu + (ex * m.up_exps.expert_stride) as u64;
10477 d[j] = pd + (ex * m.down_exps.expert_stride) as u64;
10478 }
10479 if tok_q8.is_none() {
10480 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10481 }
10482 let (zq, zd) = tok_q8.as_ref().unwrap();
10483 Self::moe_fused_epi_launch(
10484 e,
10485 m,
10486 zq,
10487 zd,
10488 sel,
10489 w,
10490 g,
10491 u,
10492 d,
10493 &mut moe_out,
10494 tok,
10495 n_embd,
10496 n_ff_exp,
10497 n_used,
10498 limit,
10499 )?;
10500 continue;
10501 }
10502 if fused_epi_may_fire {
10507 let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10508 return Err(
10509 "fused MoE epilogue fired without a live PRE clamp: the predicate and \
10510 the dispatch disagree"
10511 .into(),
10512 );
10513 };
10514 if tok_q8.is_none() {
10515 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10516 }
10517 let (zq, zd) = tok_q8.as_ref().unwrap();
10518 if Self::moe_fused_epi_token_q8(
10519 e,
10520 m,
10521 il,
10522 max_block,
10523 zq,
10524 zd,
10525 sel,
10526 w,
10527 &mut moe_out,
10528 tok,
10529 n_embd,
10530 n_ff_exp,
10531 n_used,
10532 limit,
10533 )? {
10534 continue;
10535 }
10536 }
10537 if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
10538 if tok_q8.is_none() {
10539 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10540 }
10541 let (zq, zd) = tok_q8.as_ref().unwrap();
10542 if Self::moe_gdec_token_q8(
10543 e,
10544 m,
10545 il,
10546 max_block,
10547 zq,
10548 zd,
10549 sel,
10550 w,
10551 &mut moe_out,
10552 tok,
10553 n_embd,
10554 n_ff_exp,
10555 n_used,
10556 )? {
10557 continue;
10558 }
10559 } else if gdec_may_fire
10560 && cfg.m3.is_none()
10561 && no_macros
10562 && Self::moe_gdec_token(
10563 e,
10564 m,
10565 il,
10566 max_block,
10567 &zt,
10568 sel,
10569 w,
10570 &mut moe_out,
10571 tok,
10572 n_embd,
10573 n_ff_exp,
10574 n_used,
10575 )?
10576 {
10577 continue;
10578 }
10579
10580 if gdec_may_fire || slab_fused_may_fire || fused_epi_may_fire || fused_epi_slab_may_fire
10586 {
10587 let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10588 e.memset_zeros_view(&mut row)?;
10589 }
10590
10591 let mut cpu_mask = vec![false; sel.len()];
10597 let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
10598 let gpu_resident = if use_cache {
10599 e.with_moe_cache(max_block, |cache, _| {
10600 Ok(sel
10601 .iter()
10602 .map(|&expert| {
10603 let expert = expert as u16;
10604 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
10605 .into_iter()
10606 .filter(|&projection| {
10607 cache
10608 .resident(BlockId::new(il, projection, expert))
10609 .is_some()
10610 })
10611 .count()
10612 })
10613 .collect::<Vec<_>>())
10614 })?
10615 } else {
10616 vec![0; sel.len()]
10617 };
10618 let mut cpu_selected = Vec::new();
10619 for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
10620 if gpu_resident[index] != 3 {
10621 cpu_mask[index] = true;
10622 crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
10623 let expert = expert as usize;
10624 cpu_selected.push((expert, route_weight));
10625 }
10626 }
10627 if crate::cpu_experts::predictor_enabled() {
10628 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
10632 crate::cpu_experts::predictor_submit(il, row);
10633 }
10634 if cpu_selected.is_empty() {
10635 None
10636 } else {
10637 let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
10638 let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
10639 .map_err(std::io::Error::other)?;
10640 Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
10641 }
10642 } else {
10643 None
10644 };
10645
10646 let worker_window = worker_disk_prefetch
10647 .then(worker_prefetch_window)
10648 .unwrap_or(0);
10649 for (j, &ex) in sel.iter().enumerate() {
10650 if cpu_mask[j] {
10651 continue;
10652 }
10653 let ex = ex as usize;
10654 if let Some(d) = slab_local {
10661 let gl = m.gate_exps.expert_layout(ex);
10662 let ul = m.up_exps.expert_layout(ex);
10663 let dl = m.down_exps.expert_layout(ex);
10664 let (g0, u0, d0) = (
10665 ex * m.gate_exps.expert_stride,
10666 ex * m.up_exps.expert_stride,
10667 ex * m.down_exps.expert_stride,
10668 );
10669 let (gate, up) = if moe_q8 {
10670 if tok_q8.is_none() {
10671 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10672 }
10673 let (zq, zd) = tok_q8.as_ref().unwrap();
10674 (
10675 e.qmatvec_expert_q8(
10676 &d.gate,
10677 g0..g0 + gl.len,
10678 zq,
10679 zd,
10680 1,
10681 m.gate_exps.in_f,
10682 m.gate_exps.out_f,
10683 gl.qtype,
10684 gl.row_bytes,
10685 )?,
10686 e.qmatvec_expert_q8(
10687 &d.up,
10688 u0..u0 + ul.len,
10689 zq,
10690 zd,
10691 1,
10692 m.up_exps.in_f,
10693 m.up_exps.out_f,
10694 ul.qtype,
10695 ul.row_bytes,
10696 )?,
10697 )
10698 } else {
10699 (
10700 m.qmatvec_view(
10701 e,
10702 &d.gate,
10703 g0..g0 + gl.len,
10704 &zt,
10705 1,
10706 m.gate_exps.in_f,
10707 m.gate_exps.out_f,
10708 gl.qtype,
10709 gl.row_bytes,
10710 )?,
10711 m.qmatvec_view(
10712 e,
10713 &d.up,
10714 u0..u0 + ul.len,
10715 &zt,
10716 1,
10717 m.up_exps.in_f,
10718 m.up_exps.out_f,
10719 ul.qtype,
10720 ul.row_bytes,
10721 )?,
10722 )
10723 };
10724 let mut act = e.uninit(n_ff_exp)?;
10725 Self::ffn_act_lim(
10726 e,
10727 cfg,
10728 &gate,
10729 &up,
10730 m.gate_exps.macro_scale(ex),
10731 m.up_exps.macro_scale(ex),
10732 lim_exp,
10733 &mut act,
10734 n_ff_exp,
10735 )?;
10736 let y = if moe_q8 {
10737 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
10738 e.qmatvec_expert_q8(
10739 &d.down,
10740 d0..d0 + dl.len,
10741 &aq2,
10742 &ad2,
10743 1,
10744 m.down_exps.in_f,
10745 m.down_exps.out_f,
10746 dl.qtype,
10747 dl.row_bytes,
10748 )?
10749 } else {
10750 let actv = act.slice(0..n_ff_exp);
10751 m.qmatvec_view(
10752 e,
10753 &d.down,
10754 d0..d0 + dl.len,
10755 &actv,
10756 1,
10757 m.down_exps.in_f,
10758 m.down_exps.out_f,
10759 dl.qtype,
10760 dl.row_bytes,
10761 )?
10762 };
10763 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10764 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10765 continue;
10766 }
10767 for next in page_prefetch_positions(j, sel.len(), page_window) {
10768 Self::moe_prefetch_host_expert(sel[next] as usize, m);
10769 }
10770 let keep = [
10771 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
10772 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
10773 crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
10774 ];
10775 if worker_disk_prefetch && worker_window > 0 {
10776 for next in worker_prefetch_positions(j, sel.len(), worker_window) {
10777 Self::moe_prefetch_disk_expert(
10778 e,
10779 il,
10780 sel[next] as usize,
10781 m,
10782 max_block,
10783 &keep,
10784 )?;
10785 }
10786 } else if cache_dispatch
10787 && !cpu_hybrid
10788 && moe_prefetch_enabled()
10789 && j + 1 < sel.len()
10790 {
10791 let next = sel[j + 1] as usize;
10792 Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
10793 }
10794 let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
10795 if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
10796 if (gate_q8 || up_q8) && tok_q8.is_none() {
10799 tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10800 }
10801 let gate = if gate_q8 {
10802 let (zq, zd) = tok_q8.as_ref().unwrap();
10803 Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
10804 } else {
10805 Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
10806 };
10807 let up = if up_q8 {
10808 let (zq, zd) = tok_q8.as_ref().unwrap();
10809 Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
10810 } else {
10811 Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
10812 };
10813 let mut act = e.uninit(n_ff_exp)?;
10814 Self::ffn_act_lim(
10815 e,
10816 cfg,
10817 &gate,
10818 &up,
10819 m.gate_exps.macro_scale(ex),
10820 m.up_exps.macro_scale(ex),
10821 lim_exp,
10822 &mut act,
10823 n_ff_exp,
10824 )?;
10825 let y = if down_q8 {
10826 let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
10827 Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
10828 } else {
10829 let actv = act.slice(0..n_ff_exp);
10830 Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
10831 };
10832 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10833 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10835 } else if cache_dispatch {
10836 let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
10841 let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
10842 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
10844 e,
10845 cfg,
10846 &gate,
10847 &up,
10848 m.gate_exps.macro_scale(ex),
10849 m.up_exps.macro_scale(ex),
10850 lim_exp,
10851 &mut act,
10852 n_ff_exp,
10853 )?;
10854 let actv = act.slice(0..n_ff_exp);
10855 let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
10856 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10857 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10859 } else if cache_frozen {
10860 let gate = Self::moe_frozen_gemm(
10865 e,
10866 il,
10867 PROJ_GATE,
10868 ex,
10869 m,
10870 max_block,
10871 &zt,
10872 &mut scratch_g,
10873 g_len,
10874 )?;
10875 let up = Self::moe_frozen_gemm(
10876 e,
10877 il,
10878 PROJ_UP,
10879 ex,
10880 m,
10881 max_block,
10882 &zt,
10883 &mut scratch_u,
10884 u_len,
10885 )?;
10886 let mut act = e.uninit(n_ff_exp)?;
10887 Self::ffn_act_lim(
10888 e,
10889 cfg,
10890 &gate,
10891 &up,
10892 m.gate_exps.macro_scale(ex),
10893 m.up_exps.macro_scale(ex),
10894 lim_exp,
10895 &mut act,
10896 n_ff_exp,
10897 )?;
10898 let actv = act.slice(0..n_ff_exp);
10899 let y = Self::moe_frozen_gemm(
10900 e,
10901 il,
10902 PROJ_DOWN,
10903 ex,
10904 m,
10905 max_block,
10906 &actv,
10907 &mut scratch_d,
10908 d_len,
10909 )?;
10910 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10911 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10912 } else {
10913 if scratch_g.is_none() {
10917 scratch_g = Some(e.alloc_u8_uninit(g_len)?);
10918 scratch_u = Some(e.alloc_u8_uninit(u_len)?);
10919 scratch_d = Some(e.alloc_u8_uninit(d_len)?);
10920 }
10921 let (sg, su, sd) = (
10922 scratch_g.as_mut().unwrap(),
10923 scratch_u.as_mut().unwrap(),
10924 scratch_d.as_mut().unwrap(),
10925 );
10926 let gl = m.gate_exps.expert_layout(ex);
10927 let ul = m.up_exps.expert_layout(ex);
10928 let dl = m.down_exps.expert_layout(ex);
10929 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
10930 let gate = m.qmatvec_view(
10931 e,
10932 sg,
10933 0..gl.len,
10934 &zt,
10935 1,
10936 m.gate_exps.in_f,
10937 m.gate_exps.out_f,
10938 gl.qtype,
10939 gl.row_bytes,
10940 )?;
10941
10942 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
10943 let up = m.qmatvec_view(
10944 e,
10945 su,
10946 0..ul.len,
10947 &zt,
10948 1,
10949 m.up_exps.in_f,
10950 m.up_exps.out_f,
10951 ul.qtype,
10952 ul.row_bytes,
10953 )?;
10954
10955 let mut act = e.uninit(n_ff_exp)?; Self::ffn_act_lim(
10957 e,
10958 cfg,
10959 &gate,
10960 &up,
10961 m.gate_exps.macro_scale(ex),
10962 m.up_exps.macro_scale(ex),
10963 lim_exp,
10964 &mut act,
10965 n_ff_exp,
10966 )?;
10967
10968 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
10969 let actv = act.slice(0..n_ff_exp);
10970 let y = m.qmatvec_view(
10971 e,
10972 sd,
10973 0..dl.len,
10974 &actv,
10975 1,
10976 m.down_exps.in_f,
10977 m.down_exps.out_f,
10978 dl.qtype,
10979 dl.row_bytes,
10980 )?;
10981
10982 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10983 e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10984 }
10985 }
10986 if let Some(worker) = cpu_worker {
10987 let cpu_output = worker.wait().map_err(std::io::Error::other)?;
10988 let cpu_output = e.htod(&cpu_output)?;
10989 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10990 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
10991 }
10992 if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
10993 for (j, &ex) in sel.iter().enumerate() {
10994 if cpu_mask[j] {
10995 Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
10996 }
10997 }
10998 }
10999 }
11000
11001 if memra_reference::hidden_trace::enabled() {
11002 memra_reference::hidden_trace::emit_last_row(
11003 "routed",
11004 il as i64,
11005 t,
11006 n_embd,
11007 &e.dtoh(&moe_out)?,
11008 );
11009 }
11010
11011 Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11012
11013 Ok(moe_out)
11014 }
11015
11016 #[allow(clippy::too_many_arguments)]
11036 fn moe_ffn_glm5_ep(
11037 e: &Engine,
11038 m: &MoeWeights,
11039 ep: &crate::glm5_tp::Glm5EpExps,
11040 z: &CudaSlice<f32>,
11041 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11042 t: usize,
11043 cfg: &ModelConfig,
11044 il: u16,
11045 prefill: bool,
11046 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11047 let moe = cfg
11048 .moe
11049 .as_ref()
11050 .ok_or("glm5 EP execution requires MoE model metadata")?;
11051 let n_embd = cfg.n_embd as usize;
11052 let n_expert = moe.expert_count as usize;
11053 let n_used = moe.expert_used_count as usize;
11054 let n_ff_exp = moe.expert_ff_length as usize;
11055 let sig = cfg
11056 .sigmoid_router()
11057 .ok_or("glm5 EP execution requires the sigmoid router")?;
11058 let lim_exp = cfg.clamp_exp_at(il as u32);
11059 let lim_shexp = cfg.clamp_shexp_at(il as u32);
11060 if ep.slabs.iter().map(|s| s.n_experts).sum::<usize>() != n_expert {
11061 return Err(format!(
11062 "glm5 EP slabs cover {:?} experts, model declares {n_expert}",
11063 ep.slabs.iter().map(|s| s.n_experts).collect::<Vec<_>>()
11064 )
11065 .into());
11066 }
11067 let rt = &ep.rt;
11068 let ranks = ep.ranks();
11069
11070 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
11072 let (sel_all, w_all) =
11073 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
11074 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
11075 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
11078
11079 if prefill && t > MOE_DEV_MAX_T {
11085 static EPGP_ANNOUNCED: std::sync::atomic::AtomicU8 =
11086 std::sync::atomic::AtomicU8::new(0);
11087 let enabled = crate::glm5_ep_grouped_prime_on() && moe_grouped_prefill_enabled();
11088 let bit = 1u8 << u8::from(enabled);
11089 if EPGP_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
11090 eprintln!(
11091 "[glm5-ep-grouped-prime] flag={} t={t} il={il} (announce printed in both \
11092 arms; engagement is the dispatch counter + per-layer execute line)",
11093 if enabled { "on" } else { "off" },
11094 );
11095 }
11096 if enabled
11097 && let Some(mut out) =
11098 Self::moe_ffn_glm5_ep_grouped_prime(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?
11099 {
11100 Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
11101 return Ok(out);
11102 }
11103 }
11104
11105 if crate::glm5_ep_diet_on() {
11108 let mut out = Self::moe_ffn_glm5_ep_diet(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?;
11109 Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
11110 return Ok(out);
11111 }
11112
11113 let mut moe_out = e.zeros(t * n_embd)?;
11114 use crate::glm5_tp_transport::Glm5TpTransport as TpXport;
11115 let hop = ep.rt.hop(e);
11116 let z_host = match hop.transport {
11125 TpXport::HostCanonical => Some(crate::glm5_tp_transport::host_stage_block(
11126 &hop,
11127 0,
11128 z,
11129 t * n_embd,
11130 )?),
11131 TpXport::PeerPull => None,
11132 };
11133 let z_peer_bulks = match hop.transport {
11134 TpXport::PeerPull => Some(crate::glm5_tp_transport::fanout_f32(&hop, z, t * n_embd)?),
11135 TpXport::HostCanonical => None,
11136 };
11137 for tok in 0..t {
11138 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11139 let w = &w_all[tok * n_used..(tok + 1) * n_used];
11140 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
11141 let z_peer_row_holders: Option<Vec<CudaSlice<f32>>> = match &z_host {
11145 Some(h) => {
11146 let mut rows = Vec::with_capacity(ranks - 1);
11147 for r in 1..ranks {
11148 rows.push(crate::glm5_tp_transport::host_row_to(
11149 &hop,
11150 r,
11151 &h[tok * n_embd..(tok + 1) * n_embd],
11152 )?);
11153 }
11154 Some(rows)
11155 }
11156 None => None,
11157 };
11158 for (j, &ex) in sel.iter().enumerate() {
11161 let ex = ex as usize;
11162 let owner = ep.owner(ex);
11163 if owner != 0 {
11164 crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES
11166 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11167 if matches!(
11170 crate::glm5_tp::gate_red(),
11171 Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11172 ) {
11173 continue;
11174 }
11175 }
11176 let zin_holder;
11177 let (dev, slab, zin) = if owner == 0 {
11178 (e, &ep.slabs[0], &zt)
11179 } else {
11180 zin_holder = match (&z_peer_row_holders, &z_peer_bulks) {
11181 (Some(rows), _) => rows[owner - 1].slice(0..n_embd),
11182 (None, Some(bulks)) => {
11183 bulks[owner - 1].slice(tok * n_embd..(tok + 1) * n_embd)
11184 }
11185 (None, None) => {
11186 return Err(
11187 "glm5 EP: neither transport arm staged the peer activation".into(),
11188 );
11189 }
11190 };
11191 (
11192 crate::glm5_tp::rank_engine(e, rt, owner),
11193 &ep.slabs[owner],
11194 &zin_holder,
11195 )
11196 };
11197 let local = ep.local_of[ex] as usize;
11201 let gl = m.gate_exps.expert_stride;
11202 let ul = m.up_exps.expert_stride;
11203 let dl = m.down_exps.expert_stride;
11204 let gate = dev.qmatvec_view(
11205 &slab.gate,
11206 local * gl..(local + 1) * gl,
11207 zin,
11208 1,
11209 m.gate_exps.in_f,
11210 m.gate_exps.out_f,
11211 m.gate_exps.qtype,
11212 m.gate_exps.row_bytes,
11213 )?;
11214 let up = dev.qmatvec_view(
11215 &slab.up,
11216 local * ul..(local + 1) * ul,
11217 zin,
11218 1,
11219 m.up_exps.in_f,
11220 m.up_exps.out_f,
11221 m.up_exps.qtype,
11222 m.up_exps.row_bytes,
11223 )?;
11224 let mut act = dev.uninit(n_ff_exp)?; Self::ffn_act_lim(
11226 dev,
11227 cfg,
11228 &gate,
11229 &up,
11230 m.gate_exps.macro_scale(ex),
11231 m.up_exps.macro_scale(ex),
11232 lim_exp,
11233 &mut act,
11234 n_ff_exp,
11235 )?;
11236 let actv = act.slice(0..n_ff_exp);
11237 let y = dev.qmatvec_view(
11238 &slab.down,
11239 local * dl..(local + 1) * dl,
11240 &actv,
11241 1,
11242 m.down_exps.in_f,
11243 m.down_exps.out_f,
11244 m.down_exps.qtype,
11245 m.down_exps.row_bytes,
11246 )?;
11247 let y_root = if owner == 0 {
11252 y
11253 } else {
11254 crate::glm5_tp_transport::return_row_to_root(&hop, owner, &y, n_embd)?
11255 };
11256 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
11257 e.axpy_into(
11258 &y_root,
11259 w[j] * m.down_exps.macro_scale(ex),
11260 &mut dst,
11261 n_embd,
11262 )?;
11263 }
11264 }
11265
11266 Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11267 Ok(moe_out)
11268 }
11269
11270 #[allow(clippy::too_many_arguments)]
11298 fn moe_ffn_glm5_ep_diet(
11299 e: &Engine,
11300 m: &MoeWeights,
11301 ep: &crate::glm5_tp::Glm5EpExps,
11302 z: &CudaSlice<f32>,
11303 sel_all: &[u32],
11304 w_all: &[f32],
11305 t: usize,
11306 cfg: &ModelConfig,
11307 il: u16,
11308 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11309 use std::sync::atomic::Ordering;
11310 let moe = cfg
11311 .moe
11312 .as_ref()
11313 .ok_or("glm5 EP execution requires MoE model metadata")?;
11314 let n_embd = cfg.n_embd as usize;
11315 let n_used = moe.expert_used_count as usize;
11316 let n_ff_exp = moe.expert_ff_length as usize;
11317 let lim_exp = cfg.clamp_exp_at(il as u32);
11318 let rt = &ep.rt;
11319 let n_pairs = t * n_used;
11320 if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
11321 return Err("glm5 EP diet geometry".into());
11322 }
11323 let red_skip_peer = matches!(
11324 crate::glm5_tp::gate_red(),
11325 Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11326 );
11327
11328 let ranks = ep.ranks();
11334 let mut per_rank = vec![0usize; ranks];
11335 for &s in sel_all.iter().take(n_pairs) {
11336 let ex = s as usize;
11337 if ex >= ep.owner_of.len() {
11338 return Err(format!("glm5 EP diet: selection {ex} outside the bank").into());
11339 }
11340 per_rank[ep.owner(ex)] += 1;
11341 }
11342 let mut base = vec![0usize; ranks];
11343 for r in 1..ranks {
11344 base[r] = base[r - 1] + per_rank[r - 1];
11345 }
11346 let mut ids = vec![0i32; n_pairs];
11347 {
11348 let mut k = vec![0usize; ranks];
11349 for (p, id) in ids.iter_mut().enumerate() {
11350 let r = ep.owner(sel_all[p] as usize);
11351 *id = (base[r] + k[r]) as i32;
11352 k[r] += 1;
11353 }
11354 }
11355
11356 crate::glm5_tp::GLM5_EP_DIET_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11357 for r in 1..ranks {
11358 crate::glm5_tp::GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED.fetch_add(
11359 if per_rank[r] > 0 {
11360 (t - 1) as u64
11361 } else {
11362 t as u64
11363 },
11364 Ordering::Relaxed,
11365 );
11366 }
11367 static EP_DIET_MARKED: std::sync::atomic::AtomicBool =
11368 std::sync::atomic::AtomicBool::new(false);
11369 if !EP_DIET_MARKED.swap(true, Ordering::Relaxed) {
11370 eprintln!(
11371 "[glm5-ep-diet] engaged: bulk fan-out + compact peer staging + single \
11372 slot-ordered scatter combine; per-slot host round-trips removed \
11373 transport={} performance_claim=false",
11374 ep.rt.transport.name(),
11375 );
11376 }
11377
11378 let expert_row = |dev: &Engine,
11381 slab: &crate::glm5_tp::EpRankSlab,
11382 zin: &cudarc::driver::CudaView<f32>,
11383 ex: usize|
11384 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11385 let local = ep.local_of[ex] as usize;
11386 let gl = m.gate_exps.expert_stride;
11387 let ul = m.up_exps.expert_stride;
11388 let dl = m.down_exps.expert_stride;
11389 let gate = dev.qmatvec_view(
11390 &slab.gate,
11391 local * gl..(local + 1) * gl,
11392 zin,
11393 1,
11394 m.gate_exps.in_f,
11395 m.gate_exps.out_f,
11396 m.gate_exps.qtype,
11397 m.gate_exps.row_bytes,
11398 )?;
11399 let up = dev.qmatvec_view(
11400 &slab.up,
11401 local * ul..(local + 1) * ul,
11402 zin,
11403 1,
11404 m.up_exps.in_f,
11405 m.up_exps.out_f,
11406 m.up_exps.qtype,
11407 m.up_exps.row_bytes,
11408 )?;
11409 let mut act = dev.uninit(n_ff_exp)?; Self::ffn_act_lim(
11411 dev,
11412 cfg,
11413 &gate,
11414 &up,
11415 m.gate_exps.macro_scale(ex),
11416 m.up_exps.macro_scale(ex),
11417 lim_exp,
11418 &mut act,
11419 n_ff_exp,
11420 )?;
11421 let actv = act.slice(0..n_ff_exp);
11422 dev.qmatvec_view(
11423 &slab.down,
11424 local * dl..(local + 1) * dl,
11425 &actv,
11426 1,
11427 m.down_exps.in_f,
11428 m.down_exps.out_f,
11429 m.down_exps.qtype,
11430 m.down_exps.row_bytes,
11431 )
11432 };
11433
11434 let hop = ep.rt.hop(e);
11438 let mut y_peer_blks: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
11439 for r in 1..ranks {
11440 if per_rank[r] == 0 {
11441 continue;
11442 }
11443 let dev = crate::glm5_tp::rank_engine(e, rt, r);
11444 let z_r = crate::glm5_tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
11447 let mut blk = if red_skip_peer {
11450 dev.zeros(per_rank[r] * n_embd)?
11451 } else {
11452 dev.uninit(per_rank[r] * n_embd)?
11453 };
11454 let mut k = 0usize;
11455 for tok in 0..t {
11456 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11457 for &ex in sel.iter() {
11458 let ex = ex as usize;
11459 if ep.owner(ex) != r {
11460 continue;
11461 }
11462 crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11464 crate::glm5_tp::GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED
11465 .fetch_add(1, Ordering::Relaxed);
11466 if red_skip_peer {
11467 k += 1;
11468 continue;
11469 }
11470 let zt_r = z_r.slice(tok * n_embd..(tok + 1) * n_embd);
11471 let y = expert_row(dev, &ep.slabs[r], &zt_r, ex)?;
11472 dev.copy_into(&mut blk, k * n_embd, &y, n_embd)?;
11473 k += 1;
11474 }
11475 }
11476 y_peer_blks[r] = Some(blk);
11477 }
11478
11479 let mut y_all = e.uninit(n_pairs * n_embd)?;
11482 {
11483 let mut k = 0usize;
11484 for tok in 0..t {
11485 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11486 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
11487 for &ex in sel.iter() {
11488 let ex = ex as usize;
11489 if ep.owner(ex) != 0 {
11490 continue;
11491 }
11492 let y = expert_row(e, &ep.slabs[0], &zt, ex)?;
11493 e.copy_into(&mut y_all, k * n_embd, &y, n_embd)?;
11494 k += 1;
11495 }
11496 }
11497 }
11498
11499 for r in 1..ranks {
11504 if let Some(blk) = &y_peer_blks[r] {
11505 crate::glm5_tp_transport::return_block_to_root(
11506 &hop,
11507 r,
11508 blk,
11509 &mut y_all,
11510 base[r] * n_embd,
11511 per_rank[r] * n_embd,
11512 )?;
11513 crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
11514 }
11515 }
11516
11517 let mut wd = vec![0f32; n_pairs];
11521 for ((&id, &w), &s) in ids.iter().zip(w_all.iter()).zip(sel_all.iter()) {
11522 wd[id as usize] = w * m.down_exps.macro_scale(s as usize);
11523 }
11524 let pw = e.htod(&wd)?;
11525 let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
11526 let toff_d = e.htod_i32(&toff)?;
11527 let ids_d = e.htod_i32(&ids)?;
11528 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_all, &pw, &toff_d, &ids_d, &mut moe_out, t, n_embd)?;
11530 Ok(moe_out)
11531 }
11532
11533 #[allow(clippy::too_many_arguments)]
11552 fn moe_ffn_glm5_ep_grouped_prime(
11553 e: &Engine,
11554 m: &MoeWeights,
11555 ep: &crate::glm5_tp::Glm5EpExps,
11556 z: &CudaSlice<f32>,
11557 sel_all: &[u32],
11558 w_all: &[f32],
11559 t: usize,
11560 cfg: &ModelConfig,
11561 il: u16,
11562 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
11563 use std::sync::atomic::Ordering;
11564 let moe = cfg
11565 .moe
11566 .as_ref()
11567 .ok_or("glm5 EP grouped prime requires MoE model metadata")?;
11568 let n_embd = cfg.n_embd as usize;
11569 let n_expert = moe.expert_count as usize;
11570 let n_used = moe.expert_used_count as usize;
11571 let n_ff_exp = moe.expert_ff_length as usize;
11572 if crate::moe_f16g_mode() == 0 || std::env::var("MEMRA_MOE_GATE").is_ok() {
11576 return Ok(None);
11577 }
11578 if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
11579 && f16g_proj_ok(m.up_exps.qtype, n_embd)
11580 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
11581 {
11582 return Ok(None);
11583 }
11584 if n_expert > 512 || n_used == 0 || n_used > 8 {
11585 return Ok(None);
11586 }
11587 let lim_exp = cfg.clamp_exp_at(il as u32);
11588 if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
11589 return Err(
11590 "EP grouped prime is qualified for the PRE-clamped SwiGLU form only; \
11591 a POST-clamp layer must ride the sequential arm"
11592 .into(),
11593 );
11594 }
11595 let n_pairs = t * n_used;
11596 if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
11597 return Err("EP grouped prime geometry".into());
11598 }
11599 let rt = &ep.rt;
11600 let red_skip_peer = matches!(
11601 crate::glm5_tp::gate_red(),
11602 Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11603 );
11604
11605 let rank_pass = |dev: &Engine,
11610 rank: u8,
11611 ptr_row: &CudaSlice<u64>,
11612 z_dev: &CudaSlice<f32>|
11613 -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
11614 let mut buckets_l: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
11617 let mut local_tok = Vec::new(); let mut local_ex = Vec::new(); let mut local_wd = Vec::new(); let mut local_count_per_tok = vec![0i32; t];
11621 for p in 0..n_pairs {
11622 let ex = sel_all[p] as usize;
11623 if ex >= n_expert {
11624 return Err(format!("EP grouped prime selection {ex} >= {n_expert}").into());
11625 }
11626 if ep.owner(ex) != rank as usize {
11627 continue;
11628 }
11629 let l = local_tok.len() as i32;
11630 buckets_l[ex].push(l);
11631 let tok = p / n_used;
11632 local_tok.push(tok as i32);
11633 local_ex.push(ex);
11634 local_wd.push(w_all[p] * m.down_exps.macro_scale(ex));
11635 local_count_per_tok[tok] += 1;
11636 }
11637 let n_owned = local_tok.len();
11638 if n_owned == 0 {
11639 return Ok(None);
11640 }
11641 let mut ex_ids: Vec<i32> = Vec::new();
11642 let mut ex_off: Vec<i32> = vec![0];
11643 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_owned); let mut csr_tok: Vec<i32> = Vec::with_capacity(n_owned);
11645 for (e_id, b) in buckets_l.iter().enumerate() {
11646 if !b.is_empty() {
11647 ex_ids.push(e_id as i32);
11648 for &l in b {
11649 ex_pairs.push(l);
11650 csr_tok.push(local_tok[l as usize]);
11651 }
11652 ex_off.push(ex_pairs.len() as i32);
11653 }
11654 }
11655 let n_active = ex_ids.len();
11656 if n_active == 0 || n_active > 512 {
11657 return Err(format!("EP grouped prime n_active {n_active} outside 1..=512").into());
11658 }
11659
11660 let exi = dev.htod_i32(&ex_ids)?;
11661 let exo = dev.htod_i32(&ex_off)?;
11662 let exp_d = dev.htod_i32(&ex_pairs)?;
11663 let csr_tok_d = dev.htod_i32(&csr_tok)?;
11664
11665 let (z16, zs) = dev.moe_f16g_act(z_dev, Some(&csr_tok_d), n_embd, n_owned)?;
11667 let mut g = dev.moe_f16_grouped(
11668 ptr_row,
11669 0,
11670 n_expert,
11671 &exi,
11672 &ex_off,
11673 &exo,
11674 &z16,
11675 &zs,
11676 n_embd,
11677 n_ff_exp,
11678 n_active,
11679 n_owned,
11680 m.gate_exps.qtype,
11681 m.gate_exps.row_bytes,
11682 )?;
11683 if m.gate_exps.macros.is_some() {
11684 let mg: Vec<f32> = ex_pairs
11685 .iter()
11686 .map(|&l| m.gate_exps.macro_scale(local_ex[l as usize]))
11687 .collect();
11688 let mg_d = dev.htod(&mg)?;
11689 dev.scale_rows(&mut g, &mg_d, n_ff_exp, n_owned)?;
11690 }
11691 let mut u = dev.moe_f16_grouped(
11692 ptr_row,
11693 1,
11694 n_expert,
11695 &exi,
11696 &ex_off,
11697 &exo,
11698 &z16,
11699 &zs,
11700 n_embd,
11701 n_ff_exp,
11702 n_active,
11703 n_owned,
11704 m.up_exps.qtype,
11705 m.up_exps.row_bytes,
11706 )?;
11707 if m.up_exps.macros.is_some() {
11708 let mu: Vec<f32> = ex_pairs
11709 .iter()
11710 .map(|&l| m.up_exps.macro_scale(local_ex[l as usize]))
11711 .collect();
11712 let mu_d = dev.htod(&mu)?;
11713 dev.scale_rows(&mut u, &mu_d, n_ff_exp, n_owned)?;
11714 }
11715
11716 let act = match lim_exp {
11718 Some(SwigluClamp::Pre(limit)) => {
11719 let mut a = dev.uninit(n_owned * n_ff_exp)?;
11720 dev.swiglu_preclamped_mul_scaled(
11721 &g,
11722 &u,
11723 1.0,
11724 1.0,
11725 limit,
11726 &mut a,
11727 n_owned * n_ff_exp,
11728 )?;
11729 a
11730 }
11731 None => dev.moe_pairs_silu_mul(&g, &u, n_owned * n_ff_exp)?,
11732 Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
11733 };
11734
11735 let (a16, a_s) = dev.moe_f16g_act(&act, None, n_ff_exp, n_owned)?;
11737 let d_csr = dev.moe_f16_grouped(
11738 ptr_row,
11739 2,
11740 n_expert,
11741 &exi,
11742 &ex_off,
11743 &exo,
11744 &a16,
11745 &a_s,
11746 n_ff_exp,
11747 n_embd,
11748 n_active,
11749 n_owned,
11750 m.down_exps.qtype,
11751 m.down_exps.row_bytes,
11752 )?;
11753 let y_local = dev.rows_permute(&d_csr, &exp_d, n_owned, n_embd)?;
11754 let mut toff: Vec<i32> = Vec::with_capacity(t + 1);
11755 let mut acc = 0i32;
11756 toff.push(0);
11757 for &c in &local_count_per_tok {
11758 acc += c;
11759 toff.push(acc);
11760 }
11761 let tids: Vec<i32> = (0..n_owned as i32).collect();
11762 let pw = dev.htod(&local_wd)?;
11763 let toff_d = dev.htod_i32(&toff)?;
11764 let tids_d = dev.htod_i32(&tids)?;
11765 let mut partial = dev.uninit(t * n_embd)?; dev.moe_pairs_scatter(&y_local, &pw, &toff_d, &tids_d, &mut partial, t, n_embd)?;
11767 Ok(Some(partial))
11768 };
11769
11770 let ranks = ep.ranks();
11775 let n_peer_pairs = sel_all
11776 .iter()
11777 .take(n_pairs)
11778 .filter(|&&ex| ep.owner(ex as usize) != 0)
11779 .count() as u64;
11780 crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(n_peer_pairs, Ordering::Relaxed);
11781 let hop = ep.rt.hop(e);
11782 let mut peer_partials: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
11783 for r in 1..ranks {
11784 let rank_owns_pairs = sel_all
11785 .iter()
11786 .take(n_pairs)
11787 .any(|&ex| ep.owner(ex as usize) == r);
11788 if !rank_owns_pairs {
11789 continue;
11790 }
11791 let dev = crate::glm5_tp::rank_engine(e, rt, r);
11792 let z_r = crate::glm5_tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
11794 dev.bind_runtime_device(dev.ctx().ordinal() as i32)?;
11795 let res = rank_pass(dev, r as u8, &ep.ptr_rows[r], &z_r);
11796 e.bind_runtime_device(e.ctx().ordinal() as i32)?;
11797 peer_partials[r] = res?;
11798 }
11799 let root_partial = rank_pass(e, 0, &ep.ptr_rows[0], z)?;
11800
11801 let mut out = match root_partial {
11807 Some(p) => p,
11808 None => e.zeros(t * n_embd)?,
11809 };
11810 for r in 1..ranks {
11811 if let Some(pp) = &peer_partials[r]
11812 && !red_skip_peer
11813 {
11814 let pp_root =
11815 crate::glm5_tp_transport::return_row_to_root(&hop, r, pp, t * n_embd)?;
11816 let mut dst = out.slice_mut(0..t * n_embd);
11817 e.axpy_into(&pp_root, 1.0, &mut dst, t * n_embd)?;
11818 crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
11819 }
11820 }
11821 crate::glm5_tp::GLM5_EP_GROUPED_PRIME_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11822 static EPGP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
11823 let layer_bit = 1u64 << (il as u64 % 64);
11824 if EPGP_LOGGED.fetch_or(layer_bit, Ordering::Relaxed) & layer_bit == 0 {
11825 eprintln!(
11826 "[glm5-ep-grouped-prime] execute layer={il} tokens={t} \
11827 provenance=ep-rank-slabs router=sigmoid-host-oracle epilogue=pre-clamped \
11828 combine=rank-partial-add transport={} performance_claim=false \
11829 (logged once per layer)",
11830 hop.transport.name(),
11831 );
11832 }
11833 Ok(Some(out))
11834 }
11835
11836 #[allow(clippy::too_many_arguments)]
11844 fn moe_shexp_add(
11845 e: &Engine,
11846 m: &MoeWeights,
11847 z: &CudaSlice<f32>,
11848 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11849 t: usize,
11850 cfg: &ModelConfig,
11851 lim_shexp: Option<memra_gguf::config::SwigluClamp>,
11852 moe_out: &mut CudaSlice<f32>,
11853 ) -> Result<(), Box<dyn std::error::Error>> {
11854 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
11855 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
11856 {
11857 let n_embd = cfg.n_embd as usize;
11858 let n_ff_sh = gate_shexp.out_features(); let verify_t = t > 1 && t < PRIME_MIN_T;
11867 let (sg_gate, sg_up) = if t == 1 {
11868 shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
11869 } else if verify_t {
11870 (
11871 e.matmul_decode_exact(gate_shexp, z, t)?,
11872 e.matmul_decode_exact(up_shexp, z, t)?,
11873 )
11874 } else {
11875 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) };
11877 let mut sa = e.uninit(t * n_ff_sh)?; Self::ffn_act_lim(
11879 e,
11880 cfg,
11881 &sg_gate,
11882 &sg_up,
11883 1.0,
11884 1.0,
11885 lim_shexp,
11886 &mut sa,
11887 t * n_ff_sh,
11888 )?;
11889 let sh = if verify_t {
11890 e.matmul_decode_exact(down_shexp, &sa, t)?
11891 } else {
11892 e.matmul(down_shexp, &sa, t)?
11893 }; if m.gate_inp_shexp.is_none() && crate::glm5_htod_diet_on() {
11912 crate::GLM5_HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11913 e.add_scaled_rows_ones(&sh, moe_out, n_embd, t)?;
11914 return Ok(());
11915 }
11916 let g = match &m.gate_inp_shexp {
11917 Some(gate_inp_shexp) => {
11918 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
11919 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
11920 } else {
11921 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
11922 let mut g = e.uninit(t)?; e.sigmoid(&gs, &mut g, t)?;
11924 g
11925 }
11926 }
11927 None => e.htod(&vec![1.0f32; t])?,
11928 };
11929 e.add_scaled_rows(&sh, &g, moe_out, n_embd, t)?;
11931 }
11932 Ok(())
11933 }
11934
11935 pub fn stage1_h2d_per_token(&self) -> u64 {
11938 use crate::hybrid::Ffn;
11939 let n_used = self
11940 .cfg
11941 .moe
11942 .as_ref()
11943 .map(|m| m.expert_used_count as u64)
11944 .unwrap_or(0);
11945 let mut bytes = 0u64;
11946 for l in self.layers.iter() {
11947 if let Ffn::Moe(m) = &l.ffn {
11948 bytes += n_used
11949 * (m.gate_exps.max_expert_bytes()
11950 + m.up_exps.max_expert_bytes()
11951 + m.down_exps.max_expert_bytes()) as u64;
11952 }
11953 }
11954 bytes
11955 }
11956
11957 pub(crate) fn max_moe_block(&self) -> usize {
11961 use crate::hybrid::Ffn;
11962 let mut mx = 0usize;
11963 let mut scan = |ffn: &Ffn| {
11964 if let Ffn::Moe(m) = ffn {
11965 mx = mx
11966 .max(m.gate_exps.max_expert_bytes())
11967 .max(m.up_exps.max_expert_bytes())
11968 .max(m.down_exps.max_expert_bytes());
11969 }
11970 };
11971 for l in self.layers.iter() {
11972 scan(&l.ffn);
11973 }
11974 if let Some(mtp) = self.mtp.as_ref() {
11975 scan(&mtp.ffn);
11976 }
11977 mx
11978 }
11979
11980 pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
11983 use crate::hybrid::Ffn;
11984 let mut sizes = Vec::new();
11985 let mut scan = |ffn: &Ffn| {
11986 let Ffn::Moe(m) = ffn else { return };
11987 for ex in 0..m.gate_exps.n_expert {
11988 if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
11989 continue;
11990 }
11991 for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
11992 let len = exps.expert_layout(ex).len;
11993 if len > 0 {
11994 sizes.push(len);
11995 }
11996 }
11997 }
11998 };
11999 for layer in &self.layers {
12000 scan(&layer.ffn);
12001 }
12002 if let Some(mtp) = &self.mtp {
12003 scan(&mtp.ffn);
12004 }
12005 sizes
12006 }
12007
12008 pub fn save_cpu_expert_residency_profile(
12014 &self,
12015 e: &Engine,
12016 path: &std::path::Path,
12017 ) -> Result<(), Box<dyn std::error::Error>> {
12018 let Some(ids) = e.export_moe_residency() else {
12019 return Err("no MoE residency cache to persist".into());
12020 };
12021 let mut body = format!(
12022 "memra-freeze-profile v1 max_block={} blocks={}\n",
12023 self.max_moe_block(),
12024 ids.len()
12025 );
12026 for (layer, proj, ex) in &ids {
12027 body.push_str(&format!("{layer} {proj} {ex}\n"));
12028 }
12029 let tmp = path.with_extension("tmp");
12030 std::fs::write(&tmp, body)?;
12031 std::fs::rename(&tmp, path)?;
12032 println!(
12033 "[moe-cache] freeze profile saved: {} blocks -> {}",
12034 ids.len(),
12035 path.display()
12036 );
12037 Ok(())
12038 }
12039
12040 pub fn restore_cpu_expert_residency_profile(
12044 &self,
12045 e: &Engine,
12046 path: &std::path::Path,
12047 ) -> Result<bool, Box<dyn std::error::Error>> {
12048 use crate::hybrid::Ffn;
12049 use crate::moe_cache::BlockId;
12050 let Ok(content) = std::fs::read_to_string(path) else {
12051 return Ok(false);
12052 };
12053 let mut lines = content.lines();
12054 let Some(header) = lines.next() else {
12055 return Ok(false);
12056 };
12057 let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
12058 if !header.starts_with(&expected) {
12059 println!(
12060 "[moe-cache] freeze profile ignored (geometry mismatch): {}",
12061 path.display()
12062 );
12063 return Ok(false);
12064 }
12065 let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
12066 std::collections::HashMap::new();
12067 for line in lines {
12068 let mut fields = line.split_whitespace();
12069 let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
12070 else {
12071 continue;
12072 };
12073 let (Ok(layer), Ok(proj), Ok(ex)) =
12074 (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
12075 else {
12076 continue;
12077 };
12078 by_layer
12079 .entry(layer)
12080 .or_default()
12081 .push(BlockId::new(layer, proj, ex));
12082 }
12083 let requested: usize = by_layer.values().map(Vec::len).sum();
12084 if requested == 0 {
12085 return Ok(false);
12086 }
12087 let max_block = self.max_moe_block();
12088 let mut restaged = 0usize;
12089 let mut stage_layer =
12090 |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
12091 let Ffn::Moe(m) = ffn else { return Ok(()) };
12092 let Some(ids) = by_layer.get(&layer_index) else {
12093 return Ok(());
12094 };
12095 e.with_moe_cache(max_block, |cache, eng| {
12096 for id in ids {
12097 if cache.restage_block(*id, m, eng)? {
12098 restaged += 1;
12099 }
12100 }
12101 Ok(())
12102 })
12103 };
12104 for (index, layer) in self.layers.iter().enumerate() {
12105 stage_layer(index as u16, &layer.ffn)?;
12106 }
12107 if let Some(mtp) = self.mtp.as_ref() {
12108 stage_layer(u16::MAX, &mtp.ffn)?;
12109 }
12110 e.freeze_moe_cache();
12111 println!(
12112 "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
12113 path.display()
12114 );
12115 Ok(true)
12116 }
12117
12118 pub fn freeze_cpu_expert_residency(
12120 &self,
12121 e: &Engine,
12122 ) -> Result<(), Box<dyn std::error::Error>> {
12123 e.freeze_moe_cache();
12124 Ok(())
12125 }
12126
12127 pub fn ffn_act(
12135 e: &Engine,
12136 cfg: &ModelConfig,
12137 gate: &CudaSlice<f32>,
12138 up: &CudaSlice<f32>,
12139 act: &mut CudaSlice<f32>,
12140 n: usize,
12141 ) -> Result<(), Box<dyn std::error::Error>> {
12142 Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
12143 }
12144
12145 #[allow(clippy::too_many_arguments)]
12149 pub(crate) fn ffn_act_scaled(
12150 e: &Engine,
12151 cfg: &ModelConfig,
12152 gate: &CudaSlice<f32>,
12153 up: &CudaSlice<f32>,
12154 gs: f32,
12155 us: f32,
12156 act: &mut CudaSlice<f32>,
12157 n: usize,
12158 ) -> Result<(), Box<dyn std::error::Error>> {
12159 Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
12160 }
12161
12162 #[allow(clippy::too_many_arguments)]
12173 pub(crate) fn ffn_act_lim(
12174 e: &Engine,
12175 cfg: &ModelConfig,
12176 gate: &CudaSlice<f32>,
12177 up: &CudaSlice<f32>,
12178 gs: f32,
12179 us: f32,
12180 limit: Option<SwigluClamp>,
12181 act: &mut CudaSlice<f32>,
12182 n: usize,
12183 ) -> Result<(), Box<dyn std::error::Error>> {
12184 if let Some(m3) = cfg.m3.as_ref() {
12185 debug_assert!(
12186 limit.is_none(),
12187 "m3 swigluoai and the step35/glm5_next clamps are different archs"
12188 );
12189 return e.swigluoai_mul_scaled(
12190 gate,
12191 up,
12192 gs,
12193 us,
12194 m3.swiglu_alpha,
12195 m3.swiglu_limit,
12196 act,
12197 n,
12198 );
12199 }
12200 match limit {
12201 Some(SwigluClamp::Post(l)) => {
12202 return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
12203 }
12204 Some(SwigluClamp::Pre(l)) => {
12205 return e.swiglu_preclamped_mul_scaled(gate, up, gs, us, l, act, n);
12206 }
12207 None => {}
12208 }
12209 if gs == 1.0 && us == 1.0 {
12210 return e.silu_mul(gate, up, act, n);
12211 }
12212 e.silu_mul_scaled(gate, up, gs, us, act, n)
12213 }
12214
12215 fn fused_post_limit(lim: Option<SwigluClamp>) -> Result<Option<f32>, ()> {
12221 match lim {
12222 None => Ok(None),
12223 Some(SwigluClamp::Post(l)) => Ok(Some(l)),
12224 Some(SwigluClamp::Pre(_)) => Err(()),
12225 }
12226 }
12227
12228 fn moe_route(
12234 e: &Engine,
12235 logits: &CudaSlice<f32>,
12236 t: usize,
12237 n_expert: usize,
12238 n_used: usize,
12239 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12240 Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
12241 }
12242
12243 #[allow(clippy::too_many_arguments)]
12251 fn moe_route_sigmoid_cfg(
12252 e: &Engine,
12253 logits: &CudaSlice<f32>,
12254 t: usize,
12255 n_expert: usize,
12256 n_used: usize,
12257 m: &MoeWeights,
12258 (sf, route_norm): (f32, bool),
12259 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12260 if sigmoid_router_enabled() {
12261 return e.moe_router_sigmoid_topk_host(
12262 logits,
12263 t,
12264 n_expert,
12265 n_used,
12266 m.active_count(),
12267 &m.exp_probs_b_dev,
12268 &m.active_experts_dev,
12269 sf,
12270 route_norm,
12271 );
12272 }
12273 let lg = e.dtoh(logits)?;
12274 Self::moe_route_sigmoid_host(
12275 &lg,
12276 t,
12277 n_expert,
12278 n_used,
12279 m.exp_probs_b.as_deref(),
12280 sf,
12281 route_norm,
12282 m.active_experts.as_deref(),
12283 )
12284 }
12285
12286 #[allow(clippy::excessive_precision)] fn moe_route_cfg(
12290 e: &Engine,
12291 logits: &CudaSlice<f32>,
12292 t: usize,
12293 n_expert: usize,
12294 n_used: usize,
12295 active: Option<&[bool]>,
12296 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12297 if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
12300 return e.moe_router_topk_host(logits, t, n_expert, n_used);
12301 }
12302 let lg = e.dtoh(logits)?; let mut sel = vec![0u32; t * n_used];
12305 let mut w_out = vec![0f32; t * n_used];
12306 for tok in 0..t {
12307 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
12308 let maxl = row
12310 .iter()
12311 .enumerate()
12312 .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
12313 .map(|(_, &x)| x)
12314 .fold(f32::NEG_INFINITY, f32::max);
12315 let mut probs = vec![0f32; n_expert];
12316 let mut den = 0f32;
12317 for i in 0..n_expert {
12318 if active.is_some_and(|mask| !mask[i]) {
12319 continue;
12320 }
12321 let x = (row[i] - maxl).exp();
12322 probs[i] = x;
12323 den += x;
12324 }
12325 for p in probs.iter_mut() {
12326 *p /= den;
12327 }
12328 let mut idx: Vec<usize> = (0..n_expert)
12330 .filter(|&i| active.is_none_or(|mask| mask[i]))
12331 .collect();
12332 idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
12333 let sl = &idx[..n_used];
12334 let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
12335 let mut ws: f32 = wv.iter().sum();
12336 ws = ws.max(6.103515625e-5_f32); for x in wv.iter_mut() {
12338 *x /= ws;
12339 }
12340 for j in 0..n_used {
12341 sel[tok * n_used + j] = sl[j] as u32;
12342 w_out[tok * n_used + j] = wv[j];
12343 }
12344 }
12345 Ok((sel, w_out))
12346 }
12347
12348 #[allow(clippy::too_many_arguments)]
12349 #[allow(clippy::type_complexity)] fn moe_route_sigmoid_with_input(
12351 e: &Engine,
12352 logits: &CudaSlice<f32>,
12353 input: &CudaSlice<f32>,
12354 t: usize,
12355 in_features: usize,
12356 n_expert: usize,
12357 n_used: usize,
12358 bias: Option<&[f32]>,
12359 (sf, route_norm): (f32, bool),
12360 active: Option<&[bool]>,
12361 ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12362 let logit_values =
12363 active_matrix_values(logits.len(), t, n_expert, "sigmoid router logits")?;
12364 let input_values =
12365 active_matrix_values(input.len(), t, in_features, "sigmoid router input")?;
12366 let (lg, input) = e.dtoh_pair_views(
12367 &logits.slice(0..logit_values),
12368 &input.slice(0..input_values),
12369 )?;
12370 let (sel, w) =
12371 Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
12372 Ok((sel, w, input))
12373 }
12374
12375 pub fn start_moe_prefetch_predictor(
12380 &self,
12381 e: &Engine,
12382 cfg: &ModelConfig,
12383 ) -> Result<(), Box<dyn std::error::Error>> {
12384 use crate::hybrid::Ffn;
12385 let Some(sig) = cfg.sigmoid_router() else {
12386 return Err("prefetch predictor requires a sigmoid-router arch".into());
12387 };
12388 let resident: std::collections::HashSet<(u16, u8, u16)> = e
12389 .export_moe_residency()
12390 .ok_or("prefetch predictor needs the frozen MoE residency cache")?
12391 .into_iter()
12392 .collect();
12393 let mut layers = Vec::new();
12394 for (index, layer) in self.layers.iter().enumerate() {
12395 let Ffn::Moe(m) = &layer.ffn else { continue };
12396 let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
12397 continue;
12398 };
12399 let router = e.dtoh(data)?;
12400 let n_expert = m.gate_exps.n_expert;
12401 let n_embd = m.gate_exps.in_f;
12402 if router.len() != n_embd * n_expert {
12403 continue;
12404 }
12405 let build = |exps: &crate::model::HostExps| {
12406 (0..n_expert)
12407 .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
12408 .collect::<Vec<_>>()
12409 };
12410 layers.push((
12411 index as u16,
12412 crate::cpu_experts::PredictLayerInit {
12413 router,
12414 bias: m.exp_probs_b.clone(),
12415 active: m.active_experts.clone(),
12416 n_embd,
12417 n_used: cfg
12418 .moe
12419 .as_ref()
12420 .map(|moe| moe.expert_used_count as usize)
12421 .ok_or("prefetch predictor requires MoE config")?,
12422 sig,
12423 weights_n_expert: n_expert,
12424 gate: build(&m.gate_exps),
12425 up: build(&m.up_exps),
12426 down: build(&m.down_exps),
12427 },
12428 ));
12429 }
12430 crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
12431 }
12432
12433 #[allow(clippy::too_many_arguments)]
12436 pub fn moe_route_sigmoid_host_public(
12437 logits: &[f32],
12438 t: usize,
12439 n_expert: usize,
12440 n_used: usize,
12441 bias: Option<&[f32]>,
12442 sf: f32,
12443 route_norm: bool,
12444 active: Option<&[bool]>,
12445 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12446 Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
12447 }
12448
12449 #[allow(clippy::too_many_arguments)]
12450 fn moe_route_sigmoid_host(
12451 lg: &[f32],
12452 t: usize,
12453 n_expert: usize,
12454 n_used: usize,
12455 bias: Option<&[f32]>,
12456 sf: f32,
12457 route_norm: bool,
12458 active: Option<&[bool]>,
12459 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12460 let active_count = active
12461 .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
12462 .unwrap_or(n_expert);
12463 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
12464 if lg.len() != t * n_expert {
12465 return Err(format!(
12466 "sigmoid router logits length mismatch: got {}, expected {}",
12467 lg.len(),
12468 t * n_expert,
12469 )
12470 .into());
12471 }
12472 let mut sel = vec![0u32; t * n_used];
12473 let mut w_out = vec![0f32; t * n_used];
12474 for tok in 0..t {
12475 let row = &lg[tok * n_expert..(tok + 1) * n_expert];
12476 let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
12477 let selsc: Vec<f32> = match bias {
12479 Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
12480 None => scores.clone(),
12481 };
12482 let mut idx: Vec<usize> = (0..n_expert)
12483 .filter(|&i| active.is_none_or(|mask| mask[i]))
12484 .collect();
12485 idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
12486 let sl = &idx[..n_used];
12487 let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
12488 if route_norm {
12489 let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
12490 for x in wv.iter_mut() {
12491 *x = *x / ws * sf;
12492 }
12493 } else {
12494 for x in wv.iter_mut() {
12495 *x *= sf;
12496 }
12497 }
12498 for j in 0..n_used {
12499 sel[tok * n_used + j] = sl[j] as u32;
12500 w_out[tok * n_used + j] = wv[j];
12501 }
12502 }
12503 Ok((sel, w_out))
12504 }
12505
12506 #[allow(clippy::too_many_arguments)]
12510 fn moe_ffn_sigmoid_dev(
12511 e: &Engine,
12512 m: &MoeWeights,
12513 z: &CudaSlice<f32>,
12514 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
12515 logits: &CudaSlice<f32>,
12516 t: usize,
12517 cfg: &ModelConfig,
12518 il: u16,
12519 (scaling_factor, route_norm): (f32, bool),
12520 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12521 let moe = cfg.moe.as_ref().unwrap();
12522 let n_embd = cfg.n_embd as usize;
12523 let n_expert = moe.expert_count as usize;
12524 let n_used = moe.expert_used_count as usize;
12525 let n_ff_exp = moe.expert_ff_length as usize;
12526 let dev = m.dev_exps.as_ref().unwrap();
12527 debug_assert_eq!(dev.dev, e.ctx().ordinal());
12528 debug_assert!(m.has_uniform_expert_layout());
12529 debug_assert!(!m.has_macros);
12530
12531 let (sel_d, w_d) = e.moe_router_sigmoid_topk(
12532 logits,
12533 t,
12534 n_expert,
12535 n_used,
12536 m.active_count(),
12537 &m.exp_probs_b_dev,
12538 &m.active_experts_dev,
12539 scaling_factor,
12540 route_norm,
12541 )?;
12542 crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
12543 if let Some(fp8) = dev.fp8_blk.as_ref() {
12544 debug_assert_eq!(m.gate_exps.qtype, crate::QT_F8_E4M3_BLK);
12545 debug_assert_eq!(m.up_exps.qtype, crate::QT_F8_E4M3_BLK);
12546 debug_assert_eq!(m.down_exps.qtype, crate::QT_F8_E4M3_BLK);
12547 debug_assert_eq!(fp8.gate.rows, m.gate_exps.out_f.div_ceil(128));
12548 debug_assert_eq!(fp8.up.rows, m.up_exps.out_f.div_ceil(128));
12549 debug_assert_eq!(fp8.down.rows, m.down_exps.out_f.div_ceil(128));
12550
12551 let selected = e.dtoh_i32(&sel_d)?;
12558 let route_weights = e.dtoh(&w_d)?;
12559 let mut moe_out = e.zeros(t * n_embd)?;
12560 for tok in 0..t {
12561 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
12562 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12563 for j in 0..n_used {
12564 let pair = tok * n_used + j;
12565 let expert = selected[pair] as usize;
12566 let gate = Self::moe_resident_fp8_e4m3(
12567 e,
12568 &m.gate_exps,
12569 &dev.gate,
12570 &fp8.gate,
12571 expert,
12572 &zt,
12573 1,
12574 )?;
12575 let up = Self::moe_resident_fp8_e4m3(
12576 e, &m.up_exps, &dev.up, &fp8.up, expert, &zt, 1,
12577 )?;
12578 let mut act = e.uninit(n_ff_exp)?;
12579 Self::ffn_act_lim(
12580 e,
12581 cfg,
12582 &gate,
12583 &up,
12584 1.0,
12585 1.0,
12586 cfg.clamp_exp_at(il as u32),
12587 &mut act,
12588 n_ff_exp,
12589 )?;
12590 let act = act.slice(0..n_ff_exp);
12591 let down = Self::moe_resident_fp8_e4m3(
12592 e,
12593 &m.down_exps,
12594 &dev.down,
12595 &fp8.down,
12596 expert,
12597 &act,
12598 1,
12599 )?;
12600 e.axpy_into(&down, route_weights[pair], &mut dst, n_embd)?;
12601 }
12602 }
12603 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
12604 eprintln!(
12605 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} \
12606 native=fp8blk-w8a8-e4m3-reference clamp={}",
12607 cfg.clamp_exp_at(il as u32).is_some(),
12608 );
12609 }
12610 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
12611 return Ok(moe_out);
12612 }
12613 let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
12614 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
12615 (combined, combined)
12616 } else {
12617 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
12618 };
12619 let (zq, zd) = match (t, zq8) {
12620 (1, Some((q, d))) => (q.clone(), d.clone()),
12621 _ => e.quantize_q8_1(z, t, n_embd)?,
12622 };
12623 let n_pairs = t * n_used;
12624 let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
12625 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
12629 let pair_tok_d = e.htod_i32(&pair_tok)?;
12630 let gate = e.moe_pairs_matvec_q8(
12631 &dev.ptr_row,
12632 0,
12633 &pair_tok_d,
12634 &sel_d,
12635 &zq,
12636 &zd,
12637 n_embd,
12638 n_ff_exp,
12639 n_expert,
12640 n_pairs,
12641 m.gate_exps.qtype,
12642 gate_row_bytes,
12643 )?;
12644 let up = e.moe_pairs_matvec_q8(
12645 &dev.ptr_row,
12646 1,
12647 &pair_tok_d,
12648 &sel_d,
12649 &zq,
12650 &zd,
12651 n_embd,
12652 n_ff_exp,
12653 n_expert,
12654 n_pairs,
12655 m.up_exps.qtype,
12656 up_row_bytes,
12657 )?;
12658 let mut act = e.uninit(n_pairs * n_ff_exp)?;
12659 Self::ffn_act_lim(
12660 e,
12661 cfg,
12662 &gate,
12663 &up,
12664 1.0,
12665 1.0,
12666 cfg.clamp_exp_at(il as u32),
12667 &mut act,
12668 n_pairs * n_ff_exp,
12669 )?;
12670 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
12671 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
12672 let pair_self_d = e.htod_i32(&pair_self)?;
12673 let down = e.moe_pairs_matvec_q8(
12674 &dev.ptr_row,
12675 2,
12676 &pair_self_d,
12677 &sel_d,
12678 &aq2,
12679 &ad2,
12680 n_ff_exp,
12681 n_embd,
12682 n_expert,
12683 n_pairs,
12684 m.down_exps.qtype,
12685 m.down_exps.row_bytes,
12686 )?;
12687 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
12688 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
12689 let tok_off_d = e.htod_i32(&tok_off)?;
12690 let tok_ids_d = e.htod_i32(&tok_ids)?;
12691 let mut output = e.uninit(t * n_embd)?;
12692 e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
12693 output
12694 } else {
12695 let act = e.moe_gate_up_silu8_dev_q8_rows(
12696 &dev.ptr_row,
12697 &sel_d,
12698 &zq,
12699 &zd,
12700 t,
12701 n_embd,
12702 n_ff_exp,
12703 n_used,
12704 n_expert,
12705 m.gate_exps.qtype,
12706 m.up_exps.qtype,
12707 gate_row_bytes,
12708 up_row_bytes,
12709 &m.dev_macros,
12710 )?;
12711 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
12712 let mut output = e.uninit(t * n_embd)?;
12713 e.moe_down8_fma_dev_q8_rows_g(
12714 &dev.ptr_row,
12715 &sel_d,
12716 &w_d,
12717 &aq2,
12718 &ad2,
12719 &mut output,
12720 t,
12721 n_ff_exp,
12722 n_embd,
12723 n_used,
12724 n_expert,
12725 m.down_exps.qtype,
12726 m.down_exps.row_bytes,
12727 )?;
12728 output
12729 };
12730
12731 if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
12732 eprintln!(
12733 "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
12734 cfg.clamp_exp_at(il as u32).is_some(),
12735 dev.gu_il,
12736 );
12737 }
12738 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
12739 Ok(moe_out)
12740 }
12741
12742 #[allow(clippy::too_many_arguments)]
12743 fn moe_resident_fp8_e4m3(
12744 e: &Engine,
12745 exps: &crate::model::HostExps,
12746 bytes: &CudaSlice<u8>,
12747 scales: &crate::hybrid::DevExpertFp8ProjectionScales,
12748 expert: usize,
12749 x: &cudarc::driver::CudaView<f32>,
12750 m: usize,
12751 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12752 let layout = exps.expert_layout(expert);
12753 debug_assert_eq!(layout.qtype, crate::QT_F8_E4M3_BLK);
12754 debug_assert_eq!(scales.rows * scales.cols, scales.expert_stride);
12755 let byte_start = expert * exps.expert_stride;
12756 let scale_start = expert * scales.expert_stride;
12757 let weight = bytes.slice(byte_start..byte_start + layout.len);
12758 let scale = scales
12759 .scales
12760 .slice(scale_start..scale_start + scales.expert_stride);
12761 e.qmatvec_mmq_fp8_blk_view(&weight, &scale, x, m, exps.in_f, exps.out_f)
12762 }
12763
12764 fn moe_ffn_pairs(
12773 e: &Engine,
12774 m: &MoeWeights,
12775 z: &CudaSlice<f32>,
12776 logits: &CudaSlice<f32>,
12777 t: usize,
12778 cfg: &ModelConfig,
12779 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12780 let moe = cfg.moe.as_ref().unwrap();
12781 let n_embd = cfg.n_embd as usize;
12782 let n_expert = moe.expert_count as usize;
12783 let n_used = moe.expert_used_count as usize;
12784 let n_ff_exp = moe.expert_ff_length as usize;
12785 debug_assert!(
12790 !cfg.swiglu_clamped_anywhere(),
12791 "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
12792 );
12793 let dev = m.dev_exps.as_ref().unwrap();
12794 let (rbg_d, rbu_d) = if dev.gu_il {
12796 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
12797 (sxx, sxx)
12798 } else {
12799 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
12800 };
12801
12802 let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
12803 let n_pairs = t * n_used;
12804 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
12807 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
12808 let pair_w: Vec<f32> = w_all.clone();
12809 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
12810 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
12811 let pt = e.htod_i32(&pair_tok)?;
12812 let px = e.htod_i32(&pair_ex)?;
12813 let pw = e.htod(&pair_w)?;
12814 let toff = e.htod_i32(&tok_off)?;
12815 let tids = e.htod_i32(&tok_ids)?;
12816
12817 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
12821 for p in 0..n_pairs {
12822 by_ex[pair_ex[p] as usize].push(p as i32);
12823 }
12824 let mut ex_ids: Vec<i32> = Vec::new();
12825 let mut ex_off: Vec<i32> = vec![0];
12826 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
12827 for (ex, list) in by_ex.iter().enumerate() {
12828 if list.is_empty() {
12829 continue;
12830 }
12831 ex_ids.push(ex as i32);
12832 ex_pairs.extend_from_slice(list);
12833 ex_off.push(ex_pairs.len() as i32);
12834 }
12835 let n_active = ex_ids.len();
12836 let exi = e.htod_i32(&ex_ids)?;
12837 let exo = e.htod_i32(&ex_off)?;
12838 let exp_d = e.htod_i32(&ex_pairs)?;
12839 let _ = &px; static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
12860 let mma_t = *MMA_T.get_or_init(|| {
12861 std::env::var("MEMRA_MOE_MMA_T")
12862 .ok()
12863 .and_then(|v| v.parse().ok())
12864 .unwrap_or(16)
12865 });
12866 let use_mma = std::env::var("MEMRA_MOE_MMA")
12867 .map(|v| v != "0")
12868 .unwrap_or(true)
12869 && t >= mma_t
12870 && q8_expert_dec_supported(m.gate_exps.qtype)
12871 && q8_expert_dec_supported(m.up_exps.qtype)
12872 && q8_expert_dec_supported(m.down_exps.qtype)
12873 && n_embd.is_multiple_of(256)
12874 && n_ff_exp.is_multiple_of(256);
12875 let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
12891 && q8_expert_dec_supported(m.up_exps.qtype)
12892 && q8_expert_dec_supported(m.down_exps.qtype)
12893 && n_embd.is_multiple_of(256)
12894 && n_ff_exp.is_multiple_of(256);
12895 let f16g_mode = crate::moe_f16g_mode();
12896 let f16g = f16g_mode != 0
12897 && t >= mma_t
12898 && (f16g_mode != 3 || !mma_capable)
12899 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
12900 && f16g_proj_ok(m.up_exps.qtype, n_embd)
12901 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
12902 if use_mma || f16g {
12903 let y_down = if f16g {
12911 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
12915 let csr_tok_d = e.htod_i32(&csr_tok)?;
12916 let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
12917 let g_csr = e.moe_f16_grouped(
12918 &dev.ptr_row,
12919 0,
12920 n_expert,
12921 &exi,
12922 &ex_off,
12923 &exo,
12924 &z_f16,
12925 &z_s,
12926 n_embd,
12927 n_ff_exp,
12928 n_active,
12929 n_pairs,
12930 m.gate_exps.qtype,
12931 rbg_d,
12932 )?;
12933 let u_csr = e.moe_f16_grouped(
12934 &dev.ptr_row,
12935 1,
12936 n_expert,
12937 &exi,
12938 &ex_off,
12939 &exo,
12940 &z_f16,
12941 &z_s,
12942 n_embd,
12943 n_ff_exp,
12944 n_active,
12945 n_pairs,
12946 m.up_exps.qtype,
12947 rbu_d,
12948 )?;
12949 let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
12950 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
12951 let d_csr = e.moe_f16_grouped(
12952 &dev.ptr_row,
12953 2,
12954 n_expert,
12955 &exi,
12956 &ex_off,
12957 &exo,
12958 &a_f16,
12959 &a_s,
12960 n_ff_exp,
12961 n_embd,
12962 n_active,
12963 n_pairs,
12964 m.down_exps.qtype,
12965 m.down_exps.row_bytes,
12966 )?;
12967 e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
12968 } else {
12969 let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
12971 let gate = e.mmq_iq_experts(
12972 &dev.ptr_row,
12973 0,
12974 n_expert,
12975 &exi,
12976 &exo,
12977 &exp_d,
12978 &pt,
12979 &z_scr,
12980 n_embd,
12981 n_ff_exp,
12982 n_active,
12983 n_pairs,
12984 t,
12985 m.gate_exps.qtype,
12986 rbg_d,
12987 )?;
12988 let up = e.mmq_iq_experts(
12989 &dev.ptr_row,
12990 1,
12991 n_expert,
12992 &exi,
12993 &exo,
12994 &exp_d,
12995 &pt,
12996 &z_scr,
12997 n_embd,
12998 n_ff_exp,
12999 n_active,
13000 n_pairs,
13001 t,
13002 m.up_exps.qtype,
13003 rbu_d,
13004 )?;
13005 let a_scr = if crate::moe_fuse_actq_on() {
13011 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
13012 } else {
13013 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
13014 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
13015 };
13016 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
13017 let pself = e.htod_i32(&pair_self)?;
13018 e.mmq_iq_experts(
13019 &dev.ptr_row,
13020 2,
13021 n_expert,
13022 &exi,
13023 &exo,
13024 &exp_d,
13025 &pself,
13026 &a_scr,
13027 n_ff_exp,
13028 n_embd,
13029 n_active,
13030 n_pairs,
13031 n_pairs,
13032 m.down_exps.qtype,
13033 m.down_exps.row_bytes,
13034 )?
13035 };
13036 let mut moe_out = e.uninit(t * n_embd)?;
13037 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
13038 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13039 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13040 {
13041 let n_ff_sh = gate_shexp.out_features();
13042 let sg_gate = e.matmul(gate_shexp, z, t)?;
13043 let sg_up = e.matmul(up_shexp, z, t)?;
13044 let mut sa = e.uninit(t * n_ff_sh)?;
13045 Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13046 let sh = e.matmul(down_shexp, &sa, t)?;
13047 let g = match &m.gate_inp_shexp {
13053 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
13054 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13055 }
13056 Some(gate_inp_shexp) => {
13057 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13058 let mut g = e.uninit(t)?;
13059 e.sigmoid(&gs, &mut g, t)?;
13060 g
13061 }
13062 None => e.htod(&vec![1.0f32; t])?,
13063 };
13064 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13065 }
13066 return Ok(moe_out);
13067 }
13068
13069 let dec = std::env::var("MEMRA_MOE_DEC")
13072 .map(|v| v != "0")
13073 .unwrap_or(true);
13074 let matvec = |proj,
13075 exi: &_,
13076 exo: &_,
13077 exp_d: &_,
13078 pt: &_,
13079 aq: &_,
13080 ad: &_,
13081 inf,
13082 outf,
13083 qtype,
13084 rb|
13085 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13086 let dec = dec && q8_expert_dec_supported(qtype);
13088 if dec {
13089 e.moe_pairs_matvec_q8_dec(
13090 &dev.ptr_row,
13091 proj,
13092 exi,
13093 exo,
13094 exp_d,
13095 pt,
13096 aq,
13097 ad,
13098 inf,
13099 outf,
13100 n_expert,
13101 n_active,
13102 n_pairs,
13103 qtype,
13104 rb,
13105 )
13106 } else {
13107 e.moe_pairs_matvec_q8_em(
13108 &dev.ptr_row,
13109 proj,
13110 exi,
13111 exo,
13112 exp_d,
13113 pt,
13114 aq,
13115 ad,
13116 inf,
13117 outf,
13118 n_expert,
13119 n_active,
13120 n_pairs,
13121 qtype,
13122 rb,
13123 )
13124 }
13125 };
13126 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13127 let gate = matvec(
13128 0,
13129 &exi,
13130 &exo,
13131 &exp_d,
13132 &pt,
13133 &zq,
13134 &zd,
13135 n_embd,
13136 n_ff_exp,
13137 m.gate_exps.qtype,
13138 rbg_d,
13139 )?;
13140 let up = matvec(
13141 1,
13142 &exi,
13143 &exo,
13144 &exp_d,
13145 &pt,
13146 &zq,
13147 &zd,
13148 n_embd,
13149 n_ff_exp,
13150 m.up_exps.qtype,
13151 rbu_d,
13152 )?;
13153 let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
13154 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
13155 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
13157 let pself = e.htod_i32(&pair_self)?;
13158 let y_down = matvec(
13159 2,
13160 &exi,
13161 &exo,
13162 &exp_d,
13163 &pself,
13164 &aq2,
13165 &ad2,
13166 n_ff_exp,
13167 n_embd,
13168 m.down_exps.qtype,
13169 m.down_exps.row_bytes,
13170 )?;
13171 let mut moe_out = e.uninit(t * n_embd)?; e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
13173
13174 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13178 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13179 {
13180 let n_ff_sh = gate_shexp.out_features();
13181 let step_exact = true;
13185 let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
13186 let (sg_gate, sg_up) = if step_exact && t == 1 {
13187 shexp_gate_up_t1(e, gate_shexp, up_shexp, z, None)?
13188 } else if verify_t {
13189 let mut fused = None;
13190 if crate::spec::spec_fused_t()
13191 && (2..=4).contains(&t)
13192 && e.uses_q8_1_fast(gate_shexp)
13193 && e.uses_q8_1_fast(up_shexp)
13194 {
13195 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13196 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
13197 }
13198 match fused {
13199 Some(pair) => pair,
13200 None => (
13201 e.matmul_decode_exact(gate_shexp, z, t)?,
13202 e.matmul_decode_exact(up_shexp, z, t)?,
13203 ),
13204 }
13205 } else {
13206 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
13207 };
13208 let mut sa = e.uninit(t * n_ff_sh)?;
13209 e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13210 let sh = if verify_t {
13211 e.matmul_decode_exact(down_shexp, &sa, t)?
13212 } else {
13213 e.matmul(down_shexp, &sa, t)?
13214 };
13215 let g = match &m.gate_inp_shexp {
13220 Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
13221 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13222 }
13223 Some(gate_inp_shexp) => {
13224 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13225 let mut g = e.uninit(t)?;
13226 e.sigmoid(&gs, &mut g, t)?;
13227 g
13228 }
13229 None => e.htod(&vec![1.0f32; t])?,
13230 };
13231 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13232 }
13233 Ok(moe_out)
13234 }
13235
13236 #[allow(clippy::too_many_arguments)]
13238 #[allow(clippy::too_many_arguments)]
13239 fn moe_ffn_dev(
13240 e: &Engine,
13241 m: &MoeWeights,
13242 z: &CudaSlice<f32>,
13243 zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
13244 logits: &CudaSlice<f32>,
13245 t: usize,
13246 cfg: &ModelConfig,
13247 il: u16,
13248 max_block: usize,
13249 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13250 let moe = cfg.moe.as_ref().unwrap();
13251 let n_embd = cfg.n_embd as usize;
13252 let n_expert = moe.expert_count as usize;
13253 let n_used = moe.expert_used_count as usize;
13254 let n_ff_exp = moe.expert_ff_length as usize;
13255 debug_assert!(
13259 cfg.sigmoid_router().is_none(),
13260 "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
13261 );
13262 debug_assert!(
13263 !cfg.swiglu_clamped_at(il as u32),
13264 "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
13265 );
13266
13267 let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
13269 if m.has_macros {
13272 e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
13273 }
13274
13275 let mut moe_out = e.uninit(t * n_embd)?;
13277
13278 if let Some(dev) = m.dev_exps.as_ref() {
13281 let (rbg_d, rbu_d) = if dev.gu_il {
13284 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
13285 (sxx, sxx)
13286 } else {
13287 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
13288 };
13289 let q8 = moe_q8_enabled_for_model(cfg, m);
13290 let rows_arm = q8
13299 && t > 1
13300 && crate::spec::spec_m2()
13301 && n_ff_exp == 512
13302 && n_used <= 8
13303 && std::env::var("MEMRA_MOE_DEVQ8_GU")
13304 .map(|v| v.is_empty() || v == "v")
13305 .unwrap_or(true)
13306 && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
13307 .map(|v| v.is_empty() || v == "w8h2v")
13308 .unwrap_or(true);
13309 let csr_mode = std::env::var("MEMRA_MOE_CSR")
13318 .ok()
13319 .and_then(|v| v.parse::<i32>().ok())
13320 .unwrap_or(1);
13321 let csr_nvfp4_probe = std::env::var("MEMRA_MOE_CSR_NVFP4").as_deref() == Ok("1");
13338 let csr_qt = |qt: i32| {
13339 qt == crate::QT_IQ4_XS
13340 || qt == crate::QT_IQ3_S
13341 || (csr_nvfp4_probe && qt == crate::QT_NVFP4)
13342 };
13343 let csr_t_max = if csr_nvfp4_probe { MOE_DEV_MAX_T } else { 10 };
13344 let csr_uniform = m.gate_exps.qtype == m.up_exps.qtype;
13345 let csr_arm = rows_arm
13346 && csr_mode > 0
13347 && t <= csr_t_max
13348 && csr_uniform
13349 && csr_qt(m.gate_exps.qtype)
13350 && csr_qt(m.up_exps.qtype)
13351 && csr_qt(m.down_exps.qtype);
13352 if csr_arm {
13353 if csr_mode == 2 {
13354 static ENGAGED: std::sync::Once = std::sync::Once::new();
13355 ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
13356 }
13357 let n_pairs = t * n_used;
13358 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13359 let act = e.moe_gate_up_silu8_dev_q8_csr(
13360 &dev.ptr_row,
13361 &sel_d,
13362 &zq,
13363 &zd,
13364 n_pairs,
13365 n_embd,
13366 n_ff_exp,
13367 n_used,
13368 n_expert,
13369 m.gate_exps.qtype,
13370 m.up_exps.qtype,
13371 rbg_d,
13372 rbu_d,
13373 )?;
13374 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
13375 e.moe_down8_fma_dev_q8_rows(
13379 &dev.ptr_row,
13380 &sel_d,
13381 &w_d,
13382 &aq2,
13383 &ad2,
13384 &mut moe_out,
13385 t,
13386 n_ff_exp,
13387 n_embd,
13388 n_used,
13389 n_expert,
13390 m.down_exps.qtype,
13391 m.down_exps.row_bytes,
13392 )?;
13393 if csr_mode == 2 {
13394 let act_r = e.moe_gate_up_silu8_dev_q8_rows(
13396 &dev.ptr_row,
13397 &sel_d,
13398 &zq,
13399 &zd,
13400 t,
13401 n_embd,
13402 n_ff_exp,
13403 n_used,
13404 n_expert,
13405 m.gate_exps.qtype,
13406 m.up_exps.qtype,
13407 rbg_d,
13408 rbu_d,
13409 &m.dev_macros,
13410 )?;
13411 let mut out_r = e.uninit(t * n_embd)?;
13412 let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
13413 e.moe_down8_fma_dev_q8_rows(
13414 &dev.ptr_row,
13415 &sel_d,
13416 &w_d,
13417 &aq2r,
13418 &ad2r,
13419 &mut out_r,
13420 t,
13421 n_ff_exp,
13422 n_embd,
13423 n_used,
13424 n_expert,
13425 m.down_exps.qtype,
13426 m.down_exps.row_bytes,
13427 )?;
13428 let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
13429 let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
13430 let ba = a1
13431 .iter()
13432 .zip(&a2)
13433 .filter(|(x, y)| x.to_bits() != y.to_bits())
13434 .count();
13435 let bo = o1
13436 .iter()
13437 .zip(&o2)
13438 .filter(|(x, y)| x.to_bits() != y.to_bits())
13439 .count();
13440 if ba + bo > 0 {
13441 eprintln!(
13442 "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
13443 a1.len(),
13444 o1.len()
13445 );
13446 let sel_h = e.dtoh_i32(&sel_d)?;
13448 let mut shown = 0;
13449 for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
13450 if x.to_bits() != y.to_bits() && shown < 4 {
13451 let (p, o) = (i / n_ff_exp, i % n_ff_exp);
13452 let ex = sel_h[p];
13453 let npx = sel_h.iter().filter(|&&v| v == ex).count();
13454 eprintln!(
13455 " ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
13456 );
13457 shown += 1;
13458 }
13459 }
13460 std::process::exit(3);
13461 }
13462 }
13463 } else if rows_arm {
13464 if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
13467 use std::sync::atomic::{AtomicU64, Ordering};
13468 static PAIRS: AtomicU64 = AtomicU64::new(0);
13469 static UNIQ: AtomicU64 = AtomicU64::new(0);
13470 static CALLS: AtomicU64 = AtomicU64::new(0);
13471 let sel_h = e.dtoh_i32(&sel_d)?;
13472 let mut u: Vec<i32> = sel_h.clone();
13473 u.sort_unstable();
13474 u.dedup();
13475 PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
13476 UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
13477 let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
13478 if c.is_multiple_of(480) {
13479 let p = PAIRS.load(Ordering::Relaxed);
13480 let q = UNIQ.load(Ordering::Relaxed);
13481 eprintln!(
13482 "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
13483 q as f64 / p as f64
13484 );
13485 }
13486 }
13487 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13488 let act = e.moe_gate_up_silu8_dev_q8_rows(
13489 &dev.ptr_row,
13490 &sel_d,
13491 &zq,
13492 &zd,
13493 t,
13494 n_embd,
13495 n_ff_exp,
13496 n_used,
13497 n_expert,
13498 m.gate_exps.qtype,
13499 m.up_exps.qtype,
13500 rbg_d,
13501 rbu_d,
13502 &m.dev_macros,
13503 )?;
13504 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
13505 e.moe_down8_fma_dev_q8_rows(
13506 &dev.ptr_row,
13507 &sel_d,
13508 &w_d,
13509 &aq2,
13510 &ad2,
13511 &mut moe_out,
13512 t,
13513 n_ff_exp,
13514 n_embd,
13515 n_used,
13516 n_expert,
13517 m.down_exps.qtype,
13518 m.down_exps.row_bytes,
13519 )?;
13520 } else {
13521 for tok in 0..t {
13522 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13523 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
13524 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
13525 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13526 if q8 {
13527 let (zq, zd) = match (t, zq8) {
13528 (1, Some((q, d))) => (q.clone(), d.clone()),
13529 _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
13530 };
13531 let act = e.moe_gate_up_silu8_dev_q8(
13532 &dev.ptr_row,
13533 &selt,
13534 &zq,
13535 &zd,
13536 n_embd,
13537 n_ff_exp,
13538 n_used,
13539 n_expert,
13540 m.gate_exps.qtype,
13541 m.up_exps.qtype,
13542 rbg_d,
13543 rbu_d,
13544 &m.dev_macros,
13545 )?;
13546 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
13547 e.moe_down8_fma_dev_q8(
13548 &dev.ptr_row,
13549 &selt,
13550 &wt,
13551 &aq2,
13552 &ad2,
13553 &mut dst,
13554 n_ff_exp,
13555 n_embd,
13556 n_used,
13557 n_expert,
13558 m.down_exps.qtype,
13559 m.down_exps.row_bytes,
13560 )?;
13561 } else {
13562 let act = e.moe_gate_up_silu8_dev(
13563 &dev.ptr_row,
13564 &selt,
13565 &zt,
13566 n_embd,
13567 n_ff_exp,
13568 n_used,
13569 n_expert,
13570 m.gate_exps.qtype,
13571 m.up_exps.qtype,
13572 rbg_d,
13573 rbu_d,
13574 &m.dev_macros,
13575 )?;
13576 e.moe_down8_fma_dev(
13577 &dev.ptr_row,
13578 &selt,
13579 &wt,
13580 &act,
13581 &mut dst,
13582 n_ff_exp,
13583 n_embd,
13584 n_used,
13585 n_expert,
13586 m.down_exps.qtype,
13587 m.down_exps.row_bytes,
13588 )?;
13589 }
13590 }
13591 }
13592 } else {
13593 let q8 = moe_q8_enabled_for_model(cfg, m);
13600 e.with_moe_cache(max_block, |c, eng| {
13601 let row = c
13602 .layer_dev_row(il, n_expert, eng)?
13603 .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
13604 for tok in 0..t {
13605 let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13606 let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
13607 let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
13608 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13609 if q8 {
13610 let (zq, zd) = match (t, zq8) {
13611 (1, Some((q, d))) => (q.clone(), d.clone()),
13612 _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
13613 };
13614 let act = eng.moe_gate_up_silu8_dev_q8(
13615 row,
13616 &selt,
13617 &zq,
13618 &zd,
13619 n_embd,
13620 n_ff_exp,
13621 n_used,
13622 n_expert,
13623 m.gate_exps.qtype,
13624 m.up_exps.qtype,
13625 m.gate_exps.row_bytes,
13626 m.up_exps.row_bytes,
13627 &m.dev_macros,
13628 )?;
13629 let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
13630 eng.moe_down8_fma_dev_q8(
13631 row,
13632 &selt,
13633 &wt,
13634 &aq2,
13635 &ad2,
13636 &mut dst,
13637 n_ff_exp,
13638 n_embd,
13639 n_used,
13640 n_expert,
13641 m.down_exps.qtype,
13642 m.down_exps.row_bytes,
13643 )?;
13644 } else {
13645 let act = eng.moe_gate_up_silu8_dev(
13646 row,
13647 &selt,
13648 &zt,
13649 n_embd,
13650 n_ff_exp,
13651 n_used,
13652 n_expert,
13653 m.gate_exps.qtype,
13654 m.up_exps.qtype,
13655 m.gate_exps.row_bytes,
13656 m.up_exps.row_bytes,
13657 &m.dev_macros,
13658 )?;
13659 eng.moe_down8_fma_dev(
13660 row,
13661 &selt,
13662 &wt,
13663 &act,
13664 &mut dst,
13665 n_ff_exp,
13666 n_embd,
13667 n_used,
13668 n_expert,
13669 m.down_exps.qtype,
13670 m.down_exps.row_bytes,
13671 )?;
13672 }
13673 }
13674 c.hits += (t * 3 * n_used) as u64;
13676 Ok(())
13677 })?;
13678 }
13679
13680 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13685 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13686 {
13687 let n_ff_sh = gate_shexp.out_features();
13688 let verify_t = t > 1 && t < PRIME_MIN_T;
13691 let (sg_gate, sg_up) = if t == 1 {
13692 shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
13693 } else if verify_t {
13694 let mut fused = None;
13698 if crate::spec::spec_fused_t()
13699 && (2..=4).contains(&t)
13700 && e.uses_q8_1_fast(gate_shexp)
13701 && e.uses_q8_1_fast(up_shexp)
13702 {
13703 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13704 fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
13705 }
13706 match fused {
13707 Some(pair) => pair,
13708 None => (
13709 e.matmul_decode_exact(gate_shexp, z, t)?,
13710 e.matmul_decode_exact(up_shexp, z, t)?,
13711 ),
13712 }
13713 } else {
13714 (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
13715 };
13716 let mut sa = e.uninit(t * n_ff_sh)?; e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13718 let sh = if verify_t {
13719 e.matmul_decode_exact(down_shexp, &sa, t)?
13720 } else {
13721 e.matmul(down_shexp, &sa, t)?
13722 };
13723 let g = match &m.gate_inp_shexp {
13727 Some(gate_inp_shexp) => {
13728 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
13731 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13732 } else {
13733 let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13734 let mut g = e.uninit(t)?;
13735 e.sigmoid(&gs, &mut g, t)?;
13736 g
13737 }
13738 }
13739 None => e.htod(&vec![1.0f32; t])?,
13740 };
13741 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13742 }
13743
13744 Ok(moe_out)
13745 }
13746
13747 #[allow(clippy::too_many_arguments)]
13757 #[allow(clippy::too_many_arguments)]
13760 fn moe_gdec_token_q8(
13761 e: &Engine,
13762 m: &MoeWeights,
13763 il: u16,
13764 max_block: usize,
13765 zq: &CudaSlice<i8>,
13766 zd: &CudaSlice<f32>,
13767 sel: &[u32],
13768 w: &[f32],
13769 moe_out: &mut CudaSlice<f32>,
13770 tok: usize,
13771 n_embd: usize,
13772 n_ff_exp: usize,
13773 n_used: usize,
13774 ) -> Result<bool, Box<dyn std::error::Error>> {
13775 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
13776 use cudarc::driver::DevicePtr;
13777 let ptrs = e.with_moe_cache(max_block, |c, eng| {
13778 let mut g = [0u64; 8];
13779 let mut u = [0u64; 8];
13780 let mut d = [0u64; 8];
13781 for (j, &ex) in sel.iter().enumerate() {
13782 let ex = ex as u16;
13783 let (Some(sg), Some(su), Some(sd)) = (
13784 c.resident(BlockId::new(il, PROJ_GATE, ex)),
13785 c.resident(BlockId::new(il, PROJ_UP, ex)),
13786 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
13787 ) else {
13788 return Ok(None);
13789 };
13790 let __s = eng.stream();
13791 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
13792 let (pu, _e1) = c.slot(su).device_ptr(&__s);
13793 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
13794 g[j] = pg;
13795 u[j] = pu;
13796 d[j] = pd;
13797 }
13798 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
13799 for &ex in sel {
13800 let ex = ex as u16;
13801 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
13802 c.note_profile_hit(BlockId::new(il, proj, ex));
13803 }
13804 }
13805 }
13806 c.hits += (3 * n_used) as u64;
13807 Ok(Some((g, u, d)))
13808 })?;
13809 let Some((g, u, d)) = ptrs else {
13810 return Ok(false);
13811 };
13812 let mut wv = [0f32; 8];
13813 wv[..n_used].copy_from_slice(w);
13814 let act = e.moe_gate_up_silu8_q8(
13815 crate::WPtr8(g),
13816 crate::WPtr8(u),
13817 zq,
13818 zd,
13819 n_embd,
13820 n_ff_exp,
13821 n_used,
13822 m.gate_exps.qtype,
13823 m.up_exps.qtype,
13824 m.gate_exps.row_bytes,
13825 m.up_exps.row_bytes,
13826 )?;
13827 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
13829 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13830 e.moe_down8_fma_q8(
13831 crate::WPtr8(d),
13832 crate::F32x8(wv),
13833 &aq2,
13834 &ad2,
13835 &mut dst,
13836 n_ff_exp,
13837 n_embd,
13838 n_used,
13839 m.down_exps.qtype,
13840 m.down_exps.row_bytes,
13841 )?;
13842 Ok(true)
13843 }
13844
13845 #[allow(clippy::too_many_arguments)]
13875 fn moe_fused_epi_token_q8(
13876 e: &Engine,
13877 m: &MoeWeights,
13878 il: u16,
13879 max_block: usize,
13880 zq: &CudaSlice<i8>,
13881 zd: &CudaSlice<f32>,
13882 sel: &[u32],
13883 w: &[f32],
13884 moe_out: &mut CudaSlice<f32>,
13885 tok: usize,
13886 n_embd: usize,
13887 n_ff_exp: usize,
13888 n_used: usize,
13889 limit: f32,
13890 ) -> Result<bool, Box<dyn std::error::Error>> {
13891 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_DOWN, PROJ_GATE, PROJ_UP};
13892 use cudarc::driver::DevicePtr;
13893 debug_assert!(
13894 limit > 1e-6,
13895 "the fused epilogue's kernel collapses every gate to silu(0) at limit 0"
13896 );
13897 debug_assert_eq!(sel.len(), n_used);
13898 debug_assert_eq!(w.len(), n_used);
13899
13900 let ptrs = e.with_moe_cache(max_block, |c, eng| {
13901 if c.n_slots() < 3 * n_used {
13904 return Ok(None);
13905 }
13906 for &ex in sel.iter() {
13909 let ex_usize = ex as usize;
13910 for (proj, exps) in [
13911 (PROJ_GATE, &m.gate_exps),
13912 (PROJ_UP, &m.up_exps),
13913 (PROJ_DOWN, &m.down_exps),
13914 ] {
13915 let id = BlockId::new(il, proj, ex as u16);
13916 let DispatchSlot::Resident(_) =
13917 c.dispatch_source(id, exps.expert_source(ex_usize), eng)?;
13918 }
13919 }
13920 let mut g = [0u64; 8];
13922 let mut u = [0u64; 8];
13923 let mut d = [0u64; 8];
13924 for (j, &ex) in sel.iter().enumerate() {
13925 let ex = ex as u16;
13926 let (Some(sg), Some(su), Some(sd)) = (
13927 c.resident(BlockId::new(il, PROJ_GATE, ex)),
13928 c.resident(BlockId::new(il, PROJ_UP, ex)),
13929 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
13930 ) else {
13931 return Ok(None);
13932 };
13933 let __s = eng.stream();
13934 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
13935 let (pu, _e1) = c.slot(su).device_ptr(&__s);
13936 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
13937 g[j] = pg;
13938 u[j] = pu;
13939 d[j] = pd;
13940 }
13941 Ok(Some((g, u, d)))
13942 })?;
13943 let Some((g, u, d)) = ptrs else {
13944 return Ok(false);
13945 };
13946
13947 Self::moe_fused_epi_launch(
13948 e, m, zq, zd, sel, w, g, u, d, moe_out, tok, n_embd, n_ff_exp, n_used, limit,
13949 )?;
13950 Ok(true)
13951 }
13952
13953 #[allow(clippy::too_many_arguments)]
13964 fn moe_fused_epi_launch(
13965 e: &Engine,
13966 m: &MoeWeights,
13967 zq: &CudaSlice<i8>,
13968 zd: &CudaSlice<f32>,
13969 sel: &[u32],
13970 w: &[f32],
13971 g: [u64; 8],
13972 u: [u64; 8],
13973 d: [u64; 8],
13974 moe_out: &mut CudaSlice<f32>,
13975 tok: usize,
13976 n_embd: usize,
13977 n_ff_exp: usize,
13978 n_used: usize,
13979 limit: f32,
13980 ) -> Result<(), Box<dyn std::error::Error>> {
13981 let mut gs = [0f32; 8];
13982 let mut us = [0f32; 8];
13983 let mut wv = [0f32; 8];
13984 for (j, &ex) in sel.iter().enumerate() {
13985 let ex = ex as usize;
13986 gs[j] = m.gate_exps.macro_scale(ex);
13987 us[j] = m.up_exps.macro_scale(ex);
13988 wv[j] = w[j] * m.down_exps.macro_scale(ex);
13989 }
13990 let act = e.moe_gate_up_preclamp8_q8(
13991 crate::WPtr8(g),
13992 crate::WPtr8(u),
13993 zq,
13994 zd,
13995 crate::F32x8(gs),
13996 crate::F32x8(us),
13997 limit,
13998 n_embd,
13999 n_ff_exp,
14000 n_used,
14001 m.gate_exps.qtype,
14002 m.up_exps.qtype,
14003 m.gate_exps.row_bytes,
14004 m.up_exps.row_bytes,
14005 )?;
14006 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
14008 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14009 e.moe_down8_fma_q8(
14010 crate::WPtr8(d),
14011 crate::F32x8(wv),
14012 &aq2,
14013 &ad2,
14014 &mut dst,
14015 n_ff_exp,
14016 n_embd,
14017 n_used,
14018 m.down_exps.qtype,
14019 m.down_exps.row_bytes,
14020 )?;
14021 crate::MOE_FUSED_EPI_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14022 Ok(())
14023 }
14024
14025 #[allow(clippy::too_many_arguments)]
14032 fn moe_vrows_pairs_q8(
14034 e: &Engine,
14035 m: &MoeWeights,
14036 z: &CudaSlice<f32>,
14037 sel: VrowsSel<'_>,
14038 il: u16,
14039 (pg, pu, pd): (u64, u64, u64),
14040 t: usize,
14041 n_embd: usize,
14042 n_ff_exp: usize,
14043 n_used: usize,
14044 limit: f32,
14045 moe_out: &mut CudaSlice<f32>,
14046 ) -> Result<(), Box<dyn std::error::Error>> {
14047 let n_pairs = t * n_used;
14048 let order_on = crate::moe_vrows_dedup_order_on() && !crate::moe_vrows_pack_on();
14056 let n_planes = if order_on { 4 } else { 3 };
14057 let mut ptrs_d = e.vws_uninit_u64(n_planes * n_pairs)?;
14063 let mut scl_d = e.vws_uninit(3 * n_pairs)?;
14064 match sel {
14070 VrowsSel::Host(sel_all, w_all) => {
14071 debug_assert_eq!(sel_all.len(), n_pairs);
14072 debug_assert_eq!(w_all.len(), n_pairs);
14073 let mut ptrs = vec![0u64; n_planes * n_pairs];
14074 let mut scl = vec![0f32; 3 * n_pairs];
14075 for (p, (&ex, &w)) in sel_all.iter().zip(w_all).enumerate() {
14076 let ex = ex as usize;
14077 ptrs[p] = pg + (ex * m.gate_exps.expert_stride) as u64;
14078 ptrs[n_pairs + p] = pu + (ex * m.up_exps.expert_stride) as u64;
14079 ptrs[2 * n_pairs + p] = pd + (ex * m.down_exps.expert_stride) as u64;
14080 scl[p] = m.gate_exps.macro_scale(ex);
14081 scl[n_pairs + p] = m.up_exps.macro_scale(ex);
14082 scl[2 * n_pairs + p] = w * m.down_exps.macro_scale(ex);
14085 }
14086 if order_on {
14087 ptrs[3 * n_pairs..].copy_from_slice(&crate::vrows_expert_major_order(sel_all));
14089 let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
14092 crate::MOE_VROWS_SLAB_READS_AVOIDED
14093 .fetch_add(visits - distinct, std::sync::atomic::Ordering::Relaxed);
14094 }
14095 e.htod_u64_into(&ptrs, &mut ptrs_d)?;
14096 e.htod_f32_into(&scl, &mut scl_d)?;
14097 if crate::moe_vrows_dedup_stat_on() {
14102 let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
14103 debug_assert_eq!(visits, n_pairs as u64);
14104 crate::MOE_VROWS_PAIR_VISITS
14105 .fetch_add(visits, std::sync::atomic::Ordering::Relaxed);
14106 crate::MOE_VROWS_PAIR_DISTINCT
14107 .fetch_add(distinct, std::sync::atomic::Ordering::Relaxed);
14108 crate::moe_vrows_dedup_report();
14109 }
14110 }
14111 VrowsSel::Dev(sel_d, selw_d) => {
14112 let macros = match (
14113 m.gate_exps.macros.as_deref(),
14114 m.up_exps.macros.as_deref(),
14115 m.down_exps.macros.as_deref(),
14116 ) {
14117 (Some(g), Some(u), Some(d)) => Some((g, u, d)),
14118 (None, None, None) => None,
14122 _ => {
14123 return Err("vrows device tables: expert macro planes are not uniform \
14124 across gate/up/down"
14125 .into());
14126 }
14127 };
14128 e.moe_vrows_tables_from_sel(
14129 sel_d,
14130 selw_d,
14131 il,
14132 macros,
14133 (pg, pu, pd),
14134 (
14135 m.gate_exps.expert_stride,
14136 m.up_exps.expert_stride,
14137 m.down_exps.expert_stride,
14138 ),
14139 n_pairs,
14140 &mut ptrs_d,
14141 &mut scl_d,
14142 )?;
14143 if order_on {
14144 e.moe_vrows_order_from_sel(sel_d, n_pairs, &mut ptrs_d)?;
14148 }
14149 if crate::MOE_VROWS_DEV_TABLES_DISPATCHES
14150 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
14151 == 0
14152 {
14153 eprintln!(
14154 "[moe-vrows-dev-tables] engaged: pointer/scale tables built on device \
14155 from the router's own sel/w; the per-layer pinned readback and its \
14156 cuStreamSynchronize are skipped (MEMRA_MOE_VROWS_DEV_TABLES=1)"
14157 );
14158 }
14159 }
14160 }
14161 let (mut zq, mut zd) = (
14164 e.vws_uninit_i8(t * n_embd)?,
14165 e.vws_uninit(t * (n_embd / 32))?,
14166 );
14167 e.quantize_q8_1_into(z, t, n_embd, &mut zq, &mut zd)?;
14168 let act = e.moe_gate_up_preclamp8_q8_rows(
14169 &ptrs_d,
14170 &scl_d,
14171 &zq,
14172 &zd,
14173 limit,
14174 n_embd,
14175 n_ff_exp,
14176 n_used,
14177 n_pairs,
14178 m.gate_exps.qtype,
14179 m.up_exps.qtype,
14180 m.gate_exps.row_bytes,
14181 m.up_exps.row_bytes,
14182 )?;
14183 let (mut aq2, mut ad2) = (
14185 e.vws_uninit_i8(n_pairs * n_ff_exp)?,
14186 e.vws_uninit(n_pairs * (n_ff_exp / 32))?,
14187 );
14188 e.quantize_q8_1_into(&act, n_pairs, n_ff_exp, &mut aq2, &mut ad2)?;
14189 e.moe_down8_fma_q8_rows(
14190 &ptrs_d,
14191 &scl_d,
14192 &aq2,
14193 &ad2,
14194 moe_out,
14195 n_ff_exp,
14196 n_embd,
14197 n_used,
14198 n_pairs,
14199 m.down_exps.qtype,
14200 m.down_exps.row_bytes,
14201 )?;
14202 e.vws_recycle_u64(ptrs_d);
14205 e.vws_recycle(scl_d);
14206 e.vws_recycle_i8(zq);
14207 e.vws_recycle(zd);
14208 e.vws_recycle(act);
14209 e.vws_recycle_i8(aq2);
14210 e.vws_recycle(ad2);
14211 if crate::MOE_VROWS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
14212 eprintln!(
14213 "[glm5-vrows] verify MoE batched across rows: pairs={n_pairs} (t={t} x \
14214 {n_used}), one gate/up+preclamp launch + one down/FMA launch per layer-call \
14215 (rides MEMRA_GLM5_VERIFY_BATCH)"
14216 );
14217 }
14218 Ok(())
14219 }
14220
14221 fn moe_ffn_grouped_prefill_sigmoid(
14259 e: &Engine,
14260 m: &MoeWeights,
14261 z: &CudaSlice<f32>,
14262 t: usize,
14263 cfg: &ModelConfig,
14264 il: u16,
14265 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14266 let Some(dev) = m
14271 .dev_exps
14272 .as_ref()
14273 .filter(|d| moe_slab_enabled() && d.dev == e.ctx().ordinal())
14274 else {
14275 return Ok(None);
14276 };
14277 if crate::moe_f16g_mode() == 0 {
14278 return Ok(None);
14279 }
14280 if std::env::var("MEMRA_MOE_GATE").is_ok() {
14284 return Ok(None);
14285 }
14286 let moe = cfg
14287 .moe
14288 .as_ref()
14289 .ok_or("grouped sigmoid prefill requires MoE model metadata")?;
14290 let n_embd = cfg.n_embd as usize;
14291 let n_expert = moe.expert_count as usize;
14292 let n_used = moe.expert_used_count as usize;
14293 let n_ff_exp = moe.expert_ff_length as usize;
14294 if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
14295 && f16g_proj_ok(m.up_exps.qtype, n_embd)
14296 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
14297 {
14298 return Ok(None);
14299 }
14300 if n_expert > 512 || n_used == 0 || n_used > 8 {
14303 return Ok(None);
14304 }
14305 let sigmoid = cfg
14306 .sigmoid_router()
14307 .ok_or("grouped sigmoid prefill requires the sigmoid router")?;
14308 let lim_exp = cfg.clamp_exp_at(il as u32);
14312 if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
14313 return Err(
14314 "grouped sigmoid prefill is qualified for the PRE-clamped SwiGLU form only; \
14315 a POST-clamp layer must ride the sequential arm"
14316 .into(),
14317 );
14318 }
14319
14320 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
14325 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
14326 let (sel_all, w_all) =
14327 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
14328 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
14329 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
14330 Self::trace_moe_input(e, il, t, n_embd, z)?;
14331
14332 let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
14333 let mut mt = std::time::Instant::now();
14334 let mut phase = |on: bool| -> f64 {
14335 if on {
14336 let _ = e.stream().synchronize();
14337 let v = mt.elapsed().as_secs_f64() * 1e3;
14338 mt = std::time::Instant::now();
14339 v
14340 } else {
14341 0.0
14342 }
14343 };
14344 let d_router = phase(mprof);
14345
14346 let n_pairs = t * n_used;
14348 if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
14349 return Err("grouped sigmoid prefill geometry".into());
14350 }
14351 let mut buckets: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
14352 for (p, &s_id) in sel_all.iter().take(n_pairs).enumerate() {
14353 let s_id = s_id as usize;
14354 if s_id >= n_expert {
14355 return Err(format!("grouped prefill selection {s_id} >= {n_expert}").into());
14356 }
14357 buckets[s_id].push(p as i32);
14358 }
14359 let mut ex_ids: Vec<i32> = Vec::new();
14360 let mut ex_off: Vec<i32> = vec![0];
14361 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
14362 for (e_id, b) in buckets.iter().enumerate() {
14363 if !b.is_empty() {
14364 ex_ids.push(e_id as i32);
14365 ex_pairs.extend_from_slice(b);
14366 ex_off.push(ex_pairs.len() as i32);
14367 }
14368 }
14369 let n_active = ex_ids.len();
14370 if n_active == 0 || n_active > 512 {
14371 return Err(format!("grouped prefill n_active {n_active} outside 1..=512").into());
14372 }
14373 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
14374
14375 let wd: Vec<f32> = (0..n_pairs)
14379 .map(|p| w_all[p] * m.down_exps.macro_scale(sel_all[p] as usize))
14380 .collect();
14381
14382 let (rbg_d, rbu_d) = if dev.gu_il {
14384 let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14385 (sxx, sxx)
14386 } else {
14387 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14388 };
14389
14390 let exi = e.htod_i32(&ex_ids)?;
14391 let exo = e.htod_i32(&ex_off)?;
14392 let exp_d = e.htod_i32(&ex_pairs)?;
14393 let csr_tok_d = e.htod_i32(&csr_tok)?;
14394 let pw = e.htod(&wd)?;
14395
14396 let (z16, zs) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
14398 let mut g = e.moe_f16_grouped(
14399 &dev.ptr_row,
14400 0,
14401 n_expert,
14402 &exi,
14403 &ex_off,
14404 &exo,
14405 &z16,
14406 &zs,
14407 n_embd,
14408 n_ff_exp,
14409 n_active,
14410 n_pairs,
14411 m.gate_exps.qtype,
14412 rbg_d,
14413 )?;
14414 if m.gate_exps.macros.is_some() {
14415 let mg: Vec<f32> = ex_pairs
14416 .iter()
14417 .map(|&p| m.gate_exps.macro_scale(sel_all[p as usize] as usize))
14418 .collect();
14419 let mg_d = e.htod(&mg)?;
14420 e.scale_rows(&mut g, &mg_d, n_ff_exp, n_pairs)?;
14421 }
14422 let mut u = e.moe_f16_grouped(
14423 &dev.ptr_row,
14424 1,
14425 n_expert,
14426 &exi,
14427 &ex_off,
14428 &exo,
14429 &z16,
14430 &zs,
14431 n_embd,
14432 n_ff_exp,
14433 n_active,
14434 n_pairs,
14435 m.up_exps.qtype,
14436 rbu_d,
14437 )?;
14438 if m.up_exps.macros.is_some() {
14439 let mu: Vec<f32> = ex_pairs
14440 .iter()
14441 .map(|&p| m.up_exps.macro_scale(sel_all[p as usize] as usize))
14442 .collect();
14443 let mu_d = e.htod(&mu)?;
14444 e.scale_rows(&mut u, &mu_d, n_ff_exp, n_pairs)?;
14445 }
14446
14447 let act = match lim_exp {
14450 Some(SwigluClamp::Pre(limit)) => {
14451 let mut a = e.uninit(n_pairs * n_ff_exp)?;
14452 e.swiglu_preclamped_mul_scaled(
14455 &g,
14456 &u,
14457 1.0,
14458 1.0,
14459 limit,
14460 &mut a,
14461 n_pairs * n_ff_exp,
14462 )?;
14463 a
14464 }
14465 None => e.moe_pairs_silu_mul(&g, &u, n_pairs * n_ff_exp)?,
14466 Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
14467 };
14468 let d_gemm_gu = phase(mprof);
14469
14470 let (a16, a_s) = e.moe_f16g_act(&act, None, n_ff_exp, n_pairs)?;
14472 let d_csr = e.moe_f16_grouped(
14473 &dev.ptr_row,
14474 2,
14475 n_expert,
14476 &exi,
14477 &ex_off,
14478 &exo,
14479 &a16,
14480 &a_s,
14481 n_ff_exp,
14482 n_embd,
14483 n_active,
14484 n_pairs,
14485 m.down_exps.qtype,
14486 m.down_exps.row_bytes,
14487 )?;
14488 let y_pair = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
14489 let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14490 let tids: Vec<i32> = (0..n_pairs as i32).collect();
14491 let toff_d = e.htod_i32(&toff)?;
14492 let tids_d = e.htod_i32(&tids)?;
14493 let mut moe_out = e.uninit(t * n_embd)?;
14496 e.moe_pairs_scatter(&y_pair, &pw, &toff_d, &tids_d, &mut moe_out, t, n_embd)?;
14497 let d_down = phase(mprof);
14498
14499 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
14502 if mprof {
14503 let d_shared = phase(true);
14504 eprintln!(
14505 "[moe-grouped-prefill-prof] il={il} t={t} router={d_router:.1}ms \
14506 gemm_gu={d_gemm_gu:.1}ms down_scatter={d_down:.1}ms shared={d_shared:.1}ms"
14507 );
14508 }
14509
14510 crate::MOE_GROUPED_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14511 static GPF_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
14514 let layer_bit = 1u64 << (il as u64 % 64);
14515 if GPF_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit == 0 {
14516 eprintln!(
14517 "[moe-grouped-prefill] execute layer={il} tokens={t} n_active={n_active} \
14518 provenance=resident-slab router=sigmoid-host-oracle epilogue=pre-clamped \
14519 macro_fold=gate-up-rows+down-weight performance_claim=false \
14520 (logged once per layer)"
14521 );
14522 }
14523 Ok(Some(moe_out))
14524 }
14525
14526 #[allow(clippy::too_many_arguments)] fn moe_gdec_token(
14528 e: &Engine,
14529 m: &MoeWeights,
14530 il: u16,
14531 max_block: usize,
14532 zt: &cudarc::driver::CudaView<f32>,
14533 sel: &[u32],
14534 w: &[f32],
14535 moe_out: &mut CudaSlice<f32>,
14536 tok: usize,
14537 n_embd: usize,
14538 n_ff_exp: usize,
14539 n_used: usize,
14540 ) -> Result<bool, Box<dyn std::error::Error>> {
14541 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14542 use cudarc::driver::DevicePtr;
14543 let ptrs = e.with_moe_cache(max_block, |c, eng| {
14545 let mut g = [0u64; 8];
14546 let mut u = [0u64; 8];
14547 let mut d = [0u64; 8];
14548 for (j, &ex) in sel.iter().enumerate() {
14549 let ex = ex as u16;
14550 let (Some(sg), Some(su), Some(sd)) = (
14551 c.resident(BlockId::new(il, PROJ_GATE, ex)),
14552 c.resident(BlockId::new(il, PROJ_UP, ex)),
14553 c.resident(BlockId::new(il, PROJ_DOWN, ex)),
14554 ) else {
14555 return Ok(None);
14556 };
14557 let __s = eng.stream();
14558 let (pg, _e0) = c.slot(sg).device_ptr(&__s);
14559 let (pu, _e1) = c.slot(su).device_ptr(&__s);
14560 let (pd, _e2) = c.slot(sd).device_ptr(&__s);
14561 g[j] = pg;
14562 u[j] = pu;
14563 d[j] = pd;
14564 }
14565 if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
14566 for &ex in sel {
14567 let ex = ex as u16;
14568 for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
14569 c.note_profile_hit(BlockId::new(il, proj, ex));
14570 }
14571 }
14572 }
14573 c.hits += (3 * n_used) as u64; Ok(Some((g, u, d)))
14575 })?;
14576 let Some((g, u, d)) = ptrs else {
14577 return Ok(false);
14578 };
14579 let mut wv = [0f32; 8];
14580 wv[..n_used].copy_from_slice(w);
14581 let act = e.moe_gate_up_silu8(
14583 crate::WPtr8(g),
14584 crate::WPtr8(u),
14585 zt,
14586 n_embd,
14587 n_ff_exp,
14588 n_used,
14589 m.gate_exps.qtype,
14590 m.up_exps.qtype,
14591 m.gate_exps.row_bytes,
14592 m.up_exps.row_bytes,
14593 )?;
14594 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14595 e.moe_down8_fma_into(
14596 crate::WPtr8(d),
14597 crate::F32x8(wv),
14598 &act,
14599 &mut dst,
14600 n_ff_exp,
14601 n_embd,
14602 n_used,
14603 m.down_exps.qtype,
14604 m.down_exps.row_bytes,
14605 )?;
14606 Ok(true)
14607 }
14608
14609 #[allow(clippy::too_many_arguments)] fn moe_cached_gemm_q8(
14615 e: &Engine,
14616 il: u16,
14617 proj: u8,
14618 ex: usize,
14619 m: &MoeWeights,
14620 max_block: usize,
14621 aq: &CudaSlice<i8>,
14622 ad: &CudaSlice<f32>,
14623 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14624 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
14625 let exps = match proj {
14626 PROJ_GATE => &m.gate_exps,
14627 PROJ_UP => &m.up_exps,
14628 _ => &m.down_exps,
14629 };
14630 let layout = exps.expert_layout(ex);
14631 let id = BlockId::new(il, proj, ex as u16);
14632 let source = exps.expert_source(ex);
14633 e.with_moe_cache(max_block, |c, eng| {
14634 let slot = c.dispatch_source(id, source, eng)?;
14635 let DispatchSlot::Resident(sl) = slot;
14636 let buf = c.slot(sl);
14637 eng.qmatvec_expert_q8(
14638 buf,
14639 0..layout.len,
14640 aq,
14641 ad,
14642 1,
14643 exps.in_f,
14644 exps.out_f,
14645 layout.qtype,
14646 layout.row_bytes,
14647 )
14648 })
14649 }
14650
14651 fn moe_cached_gemm(
14652 e: &Engine,
14653 il: u16,
14654 proj: u8,
14655 ex: usize,
14656 m: &MoeWeights,
14657 max_block: usize,
14658 x: &cudarc::driver::CudaView<f32>,
14659 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14660 use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
14661 let exps = match proj {
14662 PROJ_GATE => &m.gate_exps,
14663 PROJ_UP => &m.up_exps,
14664 _ => &m.down_exps,
14665 };
14666 let layout = exps.expert_layout(ex);
14667 let id = BlockId::new(il, proj, ex as u16);
14668 let source = exps.expert_source(ex);
14669 e.with_moe_cache(max_block, |c, eng| {
14671 let slot = c.dispatch_source(id, source, eng)?;
14672 let DispatchSlot::Resident(sl) = slot;
14675 let buf = c.slot(sl);
14676 m.qmatvec_view(
14677 eng,
14678 buf,
14679 0..layout.len,
14680 x,
14681 1,
14682 exps.in_f,
14683 exps.out_f,
14684 layout.qtype,
14685 layout.row_bytes,
14686 )
14687 })
14688 }
14689
14690 fn moe_profile_admit_expert(
14694 e: &Engine,
14695 il: u16,
14696 ex: usize,
14697 m: &MoeWeights,
14698 max_block: usize,
14699 ) -> Result<(), Box<dyn std::error::Error>> {
14700 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14701 e.with_moe_cache(max_block, |cache, eng| {
14702 for (proj, exps) in [
14703 (PROJ_GATE, &m.gate_exps),
14704 (PROJ_UP, &m.up_exps),
14705 (PROJ_DOWN, &m.down_exps),
14706 ] {
14707 let id = BlockId::new(il, proj, ex as u16);
14708 let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
14709 }
14710 Ok(())
14711 })
14712 }
14713
14714 #[allow(clippy::too_many_arguments)]
14717 fn moe_frozen_gemm(
14718 e: &Engine,
14719 il: u16,
14720 proj: u8,
14721 ex: usize,
14722 m: &MoeWeights,
14723 max_block: usize,
14724 x: &cudarc::driver::CudaView<f32>,
14725 scratch: &mut Option<CudaSlice<u8>>,
14726 scratch_len: usize,
14727 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14728 use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
14729 let exps = match proj {
14730 PROJ_GATE => &m.gate_exps,
14731 PROJ_UP => &m.up_exps,
14732 _ => &m.down_exps,
14733 };
14734 let layout = exps.expert_layout(ex);
14735 let id = BlockId::new(il, proj, ex as u16);
14736 if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
14737 let Some(slot) = cache.resident(id) else {
14738 return Ok(None);
14739 };
14740 let buf = cache.slot(slot);
14741 Ok(Some(m.qmatvec_view(
14742 eng,
14743 buf,
14744 0..layout.len,
14745 x,
14746 1,
14747 exps.in_f,
14748 exps.out_f,
14749 layout.qtype,
14750 layout.row_bytes,
14751 )?))
14752 })? {
14753 return Ok(output);
14754 }
14755 if scratch.is_none() {
14756 *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
14757 }
14758 let scratch = scratch.as_mut().unwrap();
14759 e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
14760 m.qmatvec_view(
14761 e,
14762 scratch,
14763 0..layout.len,
14764 x,
14765 1,
14766 exps.in_f,
14767 exps.out_f,
14768 layout.qtype,
14769 layout.row_bytes,
14770 )
14771 }
14772
14773 fn moe_prefetch_expert(
14774 e: &Engine,
14775 il: u16,
14776 ex: usize,
14777 m: &MoeWeights,
14778 max_block: usize,
14779 keep: &[crate::moe_cache::BlockId],
14780 ) -> Result<(), Box<dyn std::error::Error>> {
14781 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14782 e.with_moe_cache(max_block, |c, eng| {
14783 for (proj, exps) in [
14784 (PROJ_GATE, &m.gate_exps),
14785 (PROJ_UP, &m.up_exps),
14786 (PROJ_DOWN, &m.down_exps),
14787 ] {
14788 let id = BlockId::new(il, proj, ex as u16);
14789 let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
14790 }
14791 Ok(())
14792 })
14793 }
14794
14795 fn moe_prefetch_disk_expert(
14798 e: &Engine,
14799 il: u16,
14800 ex: usize,
14801 m: &MoeWeights,
14802 max_block: usize,
14803 keep: &[crate::moe_cache::BlockId],
14804 ) -> Result<(), Box<dyn std::error::Error>> {
14805 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14806 e.with_moe_cache(max_block, |c, eng| {
14807 for (proj, exps) in [
14808 (PROJ_GATE, &m.gate_exps),
14809 (PROJ_UP, &m.up_exps),
14810 (PROJ_DOWN, &m.down_exps),
14811 ] {
14812 let source = exps.expert_source(ex);
14813 if let crate::model::ExpertSource::Disk { .. } = &source {
14814 let id = BlockId::new(il, proj, ex as u16);
14815 let _ = c.prefetch_source(id, source, keep, eng)?;
14816 }
14817 }
14818 Ok(())
14819 })
14820 }
14821
14822 #[inline]
14823 fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
14824 let _ = m.gate_exps.prefetch_expert_pages(ex);
14825 let _ = m.up_exps.prefetch_expert_pages(ex);
14826 let _ = m.down_exps.prefetch_expert_pages(ex);
14827 }
14828}
14829
14830impl HybridModel {
14847 #[allow(clippy::too_many_arguments)]
14851 fn moe_ffn_grouped_resident_q8(
14852 e: &Engine,
14853 m: &MoeWeights,
14854 z: &CudaSlice<f32>,
14855 t: usize,
14856 cfg: &ModelConfig,
14857 il: u16,
14858 sel_all: &[u32],
14859 w_all: &[f32],
14860 table: &CudaSlice<u64>,
14861 gu_il: bool,
14862 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14863 let moe = cfg.moe.as_ref().unwrap();
14864 let n_embd = cfg.n_embd as usize;
14865 let n_expert = moe.expert_count as usize;
14866 let n_used = moe.expert_used_count as usize;
14867 let n_ff_exp = moe.expert_ff_length as usize;
14868 let n_pairs = t * n_used;
14869 debug_assert_eq!(sel_all.len(), n_pairs);
14870 debug_assert_eq!(w_all.len(), n_pairs);
14871 debug_assert!(
14872 m.gate_exps.macros.is_none()
14873 && m.up_exps.macros.is_none()
14874 && m.down_exps.macros.is_none(),
14875 "resident grouped q8 does not fold per-expert macro scales",
14876 );
14877
14878 if !cfg.swiglu_clamped_at(il as u32) {
14884 let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
14885 let sel_d = e.htod_i32(&sel)?;
14886 let w_d = e.htod(w_all)?;
14887 let (gate_row_bytes, up_row_bytes) = if gu_il {
14888 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14889 (combined, combined)
14890 } else {
14891 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14892 };
14893 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
14894 let act = e.moe_gate_up_silu8_dev_q8_rows(
14895 table,
14896 &sel_d,
14897 &zq,
14898 &zd,
14899 t,
14900 n_embd,
14901 n_ff_exp,
14902 n_used,
14903 n_expert,
14904 m.gate_exps.qtype,
14905 m.up_exps.qtype,
14906 gate_row_bytes,
14907 up_row_bytes,
14908 &m.dev_macros,
14909 )?;
14910 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
14911 let mut moe_out = e.uninit(t * n_embd)?;
14912 e.moe_down8_fma_dev_q8_rows_g(
14913 table,
14914 &sel_d,
14915 &w_d,
14916 &aq2,
14917 &ad2,
14918 &mut moe_out,
14919 t,
14920 n_ff_exp,
14921 n_embd,
14922 n_used,
14923 n_expert,
14924 m.down_exps.qtype,
14925 m.down_exps.row_bytes,
14926 )?;
14927
14928 if std::env::var("MEMRA_MOE_STATS").is_ok() {
14929 let mut counts = vec![0usize; n_expert];
14930 for &expert in sel_all {
14931 counts[expert as usize] += 1;
14932 }
14933 let mut sizes: Vec<usize> =
14934 counts.into_iter().filter(|&count| count != 0).collect();
14935 sizes.sort_unstable();
14936 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
14937 println!(
14938 "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
14939 m_e: min={} median={} mean={mean:.1} max={}",
14940 sizes.len(),
14941 n_expert,
14942 sizes.first().copied().unwrap_or(0),
14943 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
14944 sizes.last().copied().unwrap_or(0),
14945 );
14946 }
14947 return Ok(moe_out);
14948 }
14949
14950 let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
14954 let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
14955 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14956 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
14957
14958 let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
14959 for (pair, &expert) in pair_ex.iter().enumerate() {
14960 by_expert[expert as usize].push(pair as i32);
14961 }
14962
14963 let pair_tok_d = e.htod_i32(&pair_tok)?;
14964 let pair_ex_d = e.htod_i32(&pair_ex)?;
14965 let pair_w_d = e.htod(w_all)?;
14966 let tok_off_d = e.htod_i32(&tok_off)?;
14967 let tok_ids_d = e.htod_i32(&tok_ids)?;
14968
14969 let matvec = |proj: i32,
14970 pair_rows: &CudaSlice<i32>,
14971 aq: &CudaSlice<i8>,
14972 ad: &CudaSlice<f32>,
14973 in_f: usize,
14974 out_f: usize,
14975 qtype: i32,
14976 row_bytes: usize|
14977 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14978 e.moe_pairs_matvec_q8(
14979 table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
14980 row_bytes,
14981 )
14982 };
14983
14984 let (gate_row_bytes, up_row_bytes) = if gu_il {
14985 let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14986 (combined, combined)
14987 } else {
14988 (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14989 };
14990 let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
14991 let gate = matvec(
14992 0,
14993 &pair_tok_d,
14994 &zq,
14995 &zd,
14996 n_embd,
14997 n_ff_exp,
14998 m.gate_exps.qtype,
14999 gate_row_bytes,
15000 )?;
15001 let up = matvec(
15002 1,
15003 &pair_tok_d,
15004 &zq,
15005 &zd,
15006 n_embd,
15007 n_ff_exp,
15008 m.up_exps.qtype,
15009 up_row_bytes,
15010 )?;
15011 let mut act = e.uninit(n_pairs * n_ff_exp)?;
15012 Self::ffn_act_lim(
15013 e,
15014 cfg,
15015 &gate,
15016 &up,
15017 1.0,
15018 1.0,
15019 cfg.clamp_exp_at(il as u32),
15020 &mut act,
15021 n_pairs * n_ff_exp,
15022 )?;
15023 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
15024 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
15025 let pair_self_d = e.htod_i32(&pair_self)?;
15026 let down = matvec(
15027 2,
15028 &pair_self_d,
15029 &aq2,
15030 &ad2,
15031 n_ff_exp,
15032 n_embd,
15033 m.down_exps.qtype,
15034 m.down_exps.row_bytes,
15035 )?;
15036 let mut moe_out = e.uninit(t * n_embd)?;
15037 e.moe_pairs_scatter(
15038 &down,
15039 &pair_w_d,
15040 &tok_off_d,
15041 &tok_ids_d,
15042 &mut moe_out,
15043 t,
15044 n_embd,
15045 )?;
15046
15047 if std::env::var("MEMRA_MOE_STATS").is_ok() {
15048 let mut sizes: Vec<usize> = by_expert
15049 .iter()
15050 .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
15051 .collect();
15052 sizes.sort_unstable();
15053 let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
15054 println!(
15055 "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
15056 m_e: min={} median={} mean={mean:.1} max={}",
15057 sizes.len(),
15058 n_expert,
15059 sizes.first().copied().unwrap_or(0),
15060 sizes.get(sizes.len() / 2).copied().unwrap_or(0),
15061 sizes.last().copied().unwrap_or(0),
15062 );
15063 }
15064 Ok(moe_out)
15065 }
15066
15067 #[allow(clippy::too_many_arguments)]
15073 #[allow(clippy::map_entry)] fn shexp_split_matvec(
15075 e: &Engine,
15076 rank1: &Engine,
15077 wg: &CudaSlice<u8>,
15078 wu: &CudaSlice<u8>,
15079 wd: &CudaSlice<u8>,
15080 z: &CudaSlice<f32>,
15081 lim: Option<SwigluClamp>,
15082 cfg: &ModelConfig,
15083 il: u16,
15084 n_embd: usize,
15085 n_ff_sh: usize,
15086 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15087 use cudarc::driver::DevicePtr;
15088 if !n_ff_sh.is_multiple_of(2) || !n_embd.is_multiple_of(2) {
15089 return Ok(None);
15090 }
15091 let hf = n_ff_sh / 2;
15092 let nd = n_embd / 2;
15093 struct Rep {
15094 wg1: CudaSlice<u8>,
15095 wu1: CudaSlice<u8>,
15096 wd1: CudaSlice<u8>,
15097 }
15098 struct SplitWs {
15099 pin_dev: usize,
15100 gate0: CudaSlice<f32>,
15102 up0: CudaSlice<f32>,
15103 act: CudaSlice<f32>,
15104 sh_buf: CudaSlice<f32>,
15105 ev_z: cudarc::driver::CudaEvent,
15106 ev_act0: cudarc::driver::CudaEvent,
15107 z1: CudaSlice<f32>,
15109 g1: CudaSlice<f32>,
15110 u1: CudaSlice<f32>,
15111 a1h: CudaSlice<f32>,
15112 act1: CudaSlice<f32>,
15113 y1: CudaSlice<f32>,
15114 ev_act1: cudarc::driver::CudaEvent,
15115 ev_y1: cudarc::driver::CudaEvent,
15116 raw_act_e: u64,
15117 raw_sh_e: u64,
15118 raw_z1: u64,
15119 raw_a1h: u64,
15120 raw_act1: u64,
15121 raw_y1: u64,
15122 }
15123 static WS: std::sync::Mutex<Option<SplitWs>> = std::sync::Mutex::new(None);
15124 static REPS: std::sync::Mutex<Option<std::collections::HashMap<u64, Rep>>> =
15125 std::sync::Mutex::new(None);
15126 let mut guard = WS.lock().map_err(|_| "shexp split lock is poisoned")?;
15127 let mut reps_guard = REPS.lock().map_err(|_| "shexp reps lock is poisoned")?;
15128 let reps = reps_guard.get_or_insert_with(std::collections::HashMap::new);
15129 let pins = e.ctx().ordinal();
15130 if guard.as_ref().is_none_or(|w| w.pin_dev != pins) {
15131 let (gate0, up0, act, sh_buf, ev_z, ev_act0) = {
15132 let _m = e.gpu.enter_main()?;
15133 (
15134 e.htod(&vec![0.0f32; hf])?,
15135 e.htod(&vec![0.0f32; hf])?,
15136 e.htod(&vec![0.0f32; n_ff_sh])?,
15137 e.htod(&vec![0.0f32; n_embd])?,
15138 e.ctx().new_event(None)?,
15139 e.ctx().new_event(None)?,
15140 )
15141 };
15142 let (z1, g1, u1, a1h, act1, y1, ev_act1, ev_y1) = {
15143 let _r = rank1.gpu.enter_main()?;
15144 (
15145 rank1.htod(&vec![0.0f32; n_embd])?,
15146 rank1.htod(&vec![0.0f32; hf])?,
15147 rank1.htod(&vec![0.0f32; hf])?,
15148 rank1.htod(&vec![0.0f32; hf])?,
15149 rank1.htod(&vec![0.0f32; n_ff_sh])?,
15150 rank1.htod(&vec![0.0f32; nd])?,
15151 rank1.ctx().new_event(None)?,
15152 rank1.ctx().new_event(None)?,
15153 )
15154 };
15155 let (raw_act_e, raw_sh_e) = {
15156 let _m = e.gpu.enter_main()?;
15157 let stream = e.stream();
15158 let (a, _g0) = act.device_ptr(&stream);
15159 let (b, _g1) = sh_buf.device_ptr(&stream);
15160 (a, b)
15161 };
15162 let (raw_z1, raw_a1h, raw_act1, raw_y1) = {
15163 let _r = rank1.gpu.enter_main()?;
15164 let rs = rank1.stream();
15165 let (a, _g0) = z1.device_ptr(&rs);
15166 let (b, _g1) = a1h.device_ptr(&rs);
15167 let (c, _g2) = act1.device_ptr(&rs);
15168 let (d, _g3) = y1.device_ptr(&rs);
15169 (a, b, c, d)
15170 };
15171 *guard = Some(SplitWs {
15172 pin_dev: pins,
15173 gate0,
15174 up0,
15175 act,
15176 sh_buf,
15177 ev_z,
15178 ev_act0,
15179 z1,
15180 g1,
15181 u1,
15182 a1h,
15183 act1,
15184 y1,
15185 ev_act1,
15186 ev_y1,
15187 raw_act_e,
15188 raw_sh_e,
15189 raw_z1,
15190 raw_a1h,
15191 raw_act1,
15192 raw_y1,
15193 });
15194 }
15195 let ws = guard.as_mut().expect("armed above");
15196 let wg_pin = {
15197 let _m = e.gpu.enter_main()?;
15198 let stream = e.stream();
15199 let (p, _g) = wg.device_ptr(&stream);
15200 p
15201 };
15202 if !reps.contains_key(&wg_pin) {
15203 let up = |src: &CudaSlice<u8>,
15205 off_bytes: usize,
15206 len: usize|
15207 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
15208 use cudarc::driver::sys;
15209 let sptr = {
15210 let _m = e.gpu.enter_main()?;
15211 let stream = e.stream();
15212 let (p, _g) = src.device_ptr(&stream);
15213 p + off_bytes as u64
15214 };
15215 let dst = {
15216 let _r = rank1.gpu.enter_main()?;
15217 rank1.alloc_u8_uninit(len)?
15218 };
15219 let dptr = {
15220 let _r = rank1.gpu.enter_main()?;
15221 let rs = rank1.stream();
15222 let (p, _g) = dst.device_ptr(&rs);
15223 p
15224 };
15225 let _r = rank1.gpu.enter_main()?;
15226 let r = unsafe {
15227 sys::cuMemcpyAsync(
15228 dptr as sys::CUdeviceptr,
15229 sptr as sys::CUdeviceptr,
15230 len,
15231 rank1.stream().cu_stream() as sys::CUstream,
15232 )
15233 };
15234 if r != sys::CUresult::CUDA_SUCCESS {
15235 return Err(format!("shexp split replica upload: {r:?}").into());
15236 }
15237 rank1.stream().synchronize()?;
15238 Ok(dst)
15239 };
15240 let wg1 = up(wg, hf * n_embd * 2, hf * n_embd * 2)?;
15241 let wu1 = up(wu, hf * n_embd * 2, hf * n_embd * 2)?;
15242 let wd1 = up(wd, nd * n_ff_sh * 2, nd * n_ff_sh * 2)?;
15243 reps.insert(wg_pin, Rep { wg1, wu1, wd1 });
15244 }
15245 let _ = il;
15246 let raw_z = {
15248 let _m = e.gpu.enter_main()?;
15249 let stream = e.stream();
15250 let (p, _g) = z.device_ptr(&stream);
15251 ws.ev_z.record(&stream)?;
15252 p
15253 };
15254 {
15256 let rep = reps.get(&wg_pin).expect("uploaded above");
15257 let _r = rank1.gpu.enter_main()?;
15258 rank1.stream().wait(&ws.ev_z)?;
15259 crate::tp::raw_copy_bytes(ws.raw_z1, raw_z, n_embd * 4, rank1)?;
15260 let SplitWs {
15261 z1, g1, u1, a1h, ..
15262 } = &mut *ws;
15263 rank1.matvec_bf16_dual_into(&rep.wg1, &rep.wu1, z1, g1, u1, n_embd, hf)?;
15264 Self::ffn_act_lim(rank1, cfg, g1, u1, 1.0, 1.0, lim, a1h, hf)?;
15265 crate::tp::raw_copy_bytes(ws.raw_act1 + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
15267 crate::tp::raw_copy_bytes(ws.raw_act_e + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
15268 ws.ev_act1.record(&rank1.stream())?;
15269 }
15270 {
15272 let _m = e.gpu.enter_main()?;
15273 let SplitWs {
15274 gate0, up0, act, ..
15275 } = &mut *ws;
15276 let wg_lo = wg.slice(0..hf * n_embd * 2);
15277 let wu_lo = wu.slice(0..hf * n_embd * 2);
15278 e.matvec_bf16_dual_view_into(&wg_lo, &wu_lo, z, gate0, up0, n_embd, hf)?;
15279 Self::ffn_act_lim(e, cfg, gate0, up0, 1.0, 1.0, lim, act, hf)?;
15280 ws.ev_act0.record(&e.stream())?;
15281 }
15282 {
15284 let rep = reps.get(&wg_pin).expect("uploaded above");
15285 let _r = rank1.gpu.enter_main()?;
15286 rank1.stream().wait(&ws.ev_act0)?;
15287 crate::tp::raw_copy_bytes(ws.raw_act1, ws.raw_act_e, hf * 4, rank1)?;
15288 let SplitWs { act1, y1, .. } = &mut *ws;
15289 rank1.matvec_bf16_into(&rep.wd1, act1, y1, n_ff_sh, nd)?;
15290 crate::tp::raw_copy_bytes(ws.raw_sh_e + (nd * 4) as u64, ws.raw_y1, nd * 4, rank1)?;
15291 ws.ev_y1.record(&rank1.stream())?;
15292 }
15293 {
15295 let _m = e.gpu.enter_main()?;
15296 e.stream().wait(&ws.ev_act1)?;
15297 let SplitWs { act, sh_buf, .. } = &mut *ws;
15298 let wd_lo = wd.slice(0..nd * n_ff_sh * 2);
15299 e.matvec_bf16_view_into(&wd_lo, act, sh_buf, n_ff_sh, nd)?;
15300 e.stream().wait(&ws.ev_y1)?;
15301 let mut sh = e.uninit(n_embd)?;
15302 {
15303 let mut dst = sh.slice_mut(0..n_embd);
15304 e.stream()
15305 .memcpy_dtod(&ws.sh_buf.slice(0..n_embd), &mut dst)?;
15306 }
15307 Ok(Some(sh))
15308 }
15309 }
15310
15311 fn shexp_overlap_issue(
15318 e: &Engine,
15319 m: &MoeWeights,
15320 z: &CudaSlice<f32>,
15321 cfg: &ModelConfig,
15322 il: u16,
15323 n_embd: usize,
15324 ) -> Result<bool, Box<dyn std::error::Error>> {
15325 if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15326 return Ok(false);
15327 }
15328 let (
15329 Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
15330 Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
15331 Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
15332 ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15333 else {
15334 return Ok(false);
15335 };
15336 let n_ff_sh = m
15337 .gate_shexp
15338 .as_ref()
15339 .expect("matched Some above")
15340 .out_features();
15341 let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
15344 return Ok(false);
15345 };
15346 let mut guard = SHEXP_OV_WS
15347 .lock()
15348 .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15349 let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
15350 if guard
15351 .as_ref()
15352 .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
15353 {
15354 *guard = Some((
15355 pins.0,
15356 pins.1,
15357 pins.2,
15358 e.uninit(n_ff_sh)?,
15359 e.uninit(n_embd)?,
15360 ));
15361 }
15362 let (_, _, _, act, sh) = guard.as_mut().expect("armed above");
15363 e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim)?;
15364 e.matvec_bf16_into(wd, act, sh, n_ff_sh, n_embd)?;
15365 drop(guard);
15366 Ok(true)
15367 }
15368
15369 #[allow(clippy::too_many_arguments)]
15375 #[allow(clippy::map_entry)] fn shexp_dev1_issue(
15377 e: &Engine,
15378 rank1: &Engine,
15379 m: &MoeWeights,
15380 z: &CudaSlice<f32>,
15381 cfg: &ModelConfig,
15382 il: u16,
15383 n_embd: usize,
15384 ) -> Result<bool, Box<dyn std::error::Error>> {
15385 use cudarc::driver::DevicePtr;
15386 if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15387 return Ok(false);
15388 }
15389 let (
15390 Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
15391 Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
15392 Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
15393 ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15394 else {
15395 return Ok(false);
15396 };
15397 let n_ff_sh = m
15398 .gate_shexp
15399 .as_ref()
15400 .expect("matched Some above")
15401 .out_features();
15402 let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
15405 return Ok(false);
15406 };
15407 let mut ws_guard = SHEXP_D1_WS
15409 .lock()
15410 .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
15411 if ws_guard
15412 .as_ref()
15413 .is_none_or(|(k, ..)| *k != (n_embd, n_ff_sh))
15414 {
15415 let (act1, z1, ev_done) = {
15416 let _r1 = rank1.gpu.enter_main()?;
15417 (
15418 rank1.htod(&vec![0.0f32; n_ff_sh])?,
15419 rank1.htod(&vec![0.0f32; n_embd])?,
15420 rank1.ctx().new_event(None)?,
15421 )
15422 };
15423 let (sh_root, ev_z) = {
15424 let _main = e.gpu.enter_main()?;
15425 (e.htod(&vec![0.0f32; n_embd])?, e.ctx().new_event(None)?)
15426 };
15427 *ws_guard = Some(((n_embd, n_ff_sh), act1, z1, sh_root, ev_z, ev_done));
15428 }
15429 let mut reps_guard = SHEXP_D1_REPS
15431 .lock()
15432 .map_err(|_| "shexp dev1 replica lock is poisoned")?;
15433 let reps = reps_guard.get_or_insert_with(Default::default);
15434 if !reps.contains_key(&il) {
15435 let (wg1, wu1, wd1) = {
15436 let _r1 = rank1.gpu.enter_main()?;
15437 (
15438 rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
15439 rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
15440 rank1.alloc_u8_uninit(n_embd * n_ff_sh * 2)?,
15441 )
15442 };
15443 for (src, dst) in [(wg, &wg1), (wu, &wu1), (wd, &wd1)] {
15444 let s_ptr = {
15445 let _main = e.gpu.enter_main()?;
15446 let stream = e.stream();
15447 let (p, _g) = src.device_ptr(&stream);
15448 p
15449 };
15450 let d_ptr = {
15451 let _r1 = rank1.gpu.enter_main()?;
15452 let stream = rank1.stream();
15453 let (p, _g) = dst.device_ptr(&stream);
15454 p
15455 };
15456 let _r1 = rank1.gpu.enter_main()?;
15457 crate::tp::raw_copy_bytes(d_ptr, s_ptr, src.len(), rank1)?;
15458 }
15459 {
15460 let _r1 = rank1.gpu.enter_main()?;
15461 rank1.stream().synchronize()?;
15462 }
15463 reps.insert(il, (wg1, wu1, wd1));
15464 }
15465 let (wg1, wu1, wd1) = reps.get(&il).expect("armed above");
15466 let (_, act1, z1, sh_root, ev_z, ev_done) = ws_guard.as_mut().expect("armed above");
15467 let (raw_z, raw_sh) = {
15470 let _main = e.gpu.enter_main()?;
15471 let stream = e.stream();
15472 let (a, _g0) = z.device_ptr(&stream);
15473 let (b, _g1) = sh_root.device_ptr(&stream);
15474 ev_z.record(&stream)?;
15475 (a, b)
15476 };
15477 {
15478 let _r1 = rank1.gpu.enter_main()?;
15479 rank1.stream().wait(ev_z)?;
15480 let raw_z1 = {
15481 let stream = rank1.stream();
15482 let (p, _g) = z1.device_ptr(&stream);
15483 p
15484 };
15485 crate::tp::raw_copy_bytes(raw_z1, raw_z, n_embd * 4, rank1)?;
15486 rank1.matvec_bf16_dual_silu_into(wg1, wu1, z1, act1, n_embd, n_ff_sh, lim)?;
15487 rank1.matvec_bf16_raw_out(wd1, act1, raw_sh, n_ff_sh, n_embd)?;
15491 ev_done.record(&rank1.stream())?;
15492 }
15493 Ok(true)
15494 }
15495
15496 fn shexp_dev1_apply(
15498 e: &Engine,
15499 output: &mut CudaSlice<f32>,
15500 n_embd: usize,
15501 ) -> Result<(), Box<dyn std::error::Error>> {
15502 let guard = SHEXP_D1_WS
15503 .lock()
15504 .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
15505 let (pin, _, _, sh_root, _, ev_done) =
15506 guard.as_ref().ok_or("shexp dev1 apply without issue")?;
15507 if pin.0 != n_embd {
15508 return Err("shexp dev1 width drifted".into());
15509 }
15510 let _main = e.gpu.enter_main()?;
15511 e.stream().wait(ev_done)?;
15512 static ONES_D1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15513 std::sync::Mutex::new(None);
15514 let mut og = ONES_D1.lock().map_err(|_| "ones lock is poisoned")?;
15515 if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15516 *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15517 }
15518 let ones = &og.as_ref().expect("armed above").1;
15519 e.add_scaled_rows(sh_root, ones, output, n_embd, 1)?;
15520 Ok(())
15521 }
15522
15523 fn shexp_overlap_tail_ptrs(
15527 e: &Engine,
15528 m: &MoeWeights,
15529 cfg: &ModelConfig,
15530 n_embd: usize,
15531 ) -> Result<Option<(u64, u64)>, Box<dyn std::error::Error>> {
15532 use cudarc::driver::DevicePtr;
15533 if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15534 return Ok(None);
15535 }
15536 let (
15537 Some(crate::model::GpuTensor::FloatBf16 { .. }),
15538 Some(crate::model::GpuTensor::FloatBf16 { .. }),
15539 Some(crate::model::GpuTensor::FloatBf16 { .. }),
15540 ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15541 else {
15542 return Ok(None);
15543 };
15544 let n_ff_sh = m
15545 .gate_shexp
15546 .as_ref()
15547 .expect("matched Some above")
15548 .out_features();
15549 let mut guard = SHEXP_OV_WS
15550 .lock()
15551 .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15552 let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
15553 if guard
15554 .as_ref()
15555 .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
15556 {
15557 *guard = Some((
15558 pins.0,
15559 pins.1,
15560 pins.2,
15561 e.uninit(n_ff_sh)?,
15562 e.uninit(n_embd)?,
15563 ));
15564 }
15565 let sh_raw = {
15566 let (_, _, _, _, sh) = guard.as_ref().expect("armed above");
15567 let stream = e.stream();
15568 let (p, _g) = sh.device_ptr(&stream);
15569 p
15570 };
15571 static ONES_T3: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15572 std::sync::Mutex::new(None);
15573 let mut og = ONES_T3.lock().map_err(|_| "ones lock is poisoned")?;
15574 if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15575 *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15576 }
15577 let ones_raw = {
15578 let stream = e.stream();
15579 let (p, _g) = og.as_ref().expect("armed above").1.device_ptr(&stream);
15580 p
15581 };
15582 Ok(Some((sh_raw, ones_raw)))
15583 }
15584
15585 fn shexp_overlap_apply(
15588 e: &Engine,
15589 output: &mut CudaSlice<f32>,
15590 n_embd: usize,
15591 ) -> Result<(), Box<dyn std::error::Error>> {
15592 let guard = SHEXP_OV_WS
15593 .lock()
15594 .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15595 let (_, ne, _, _, sh) = guard.as_ref().ok_or("shexp overlap apply without issue")?;
15596 if *ne != n_embd {
15597 return Err("shexp overlap width drifted".into());
15598 }
15599 static ONES_OV: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15600 std::sync::Mutex::new(None);
15601 let mut og = ONES_OV.lock().map_err(|_| "ones lock is poisoned")?;
15602 if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15603 *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15604 }
15605 let ones = &og.as_ref().expect("armed above").1;
15606 e.add_scaled_rows(sh, ones, output, n_embd, 1)?;
15607 Ok(())
15608 }
15609
15610 fn moe_ffn_grouped_add_shared(
15611 e: &Engine,
15612 m: &MoeWeights,
15613 z: &CudaSlice<f32>,
15614 t: usize,
15615 cfg: &ModelConfig,
15616 il: u16,
15617 moe_out: &mut CudaSlice<f32>,
15618 ) -> Result<(), Box<dyn std::error::Error>> {
15619 static SHEXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
15622 static SHEXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
15623 let shexp_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
15624 let shexp_started = shexp_timing.then(std::time::Instant::now);
15625 let result = Self::moe_ffn_grouped_add_shared_inner(e, m, z, t, cfg, il, moe_out);
15626 if let Some(started) = shexp_started {
15627 use std::sync::atomic::Ordering;
15628 e.stream().synchronize()?;
15629 let ns = SHEXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
15630 + started.elapsed().as_nanos() as u64;
15631 let calls = SHEXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
15632 if calls.is_multiple_of(430) {
15633 eprintln!(
15634 "[moe-shexp-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
15635 ns as f64 / 1.0e6,
15636 ns as f64 / calls as f64 / 1.0e3,
15637 );
15638 }
15639 }
15640 result
15641 }
15642
15643 #[allow(clippy::too_many_arguments)]
15644 fn moe_ffn_grouped_add_shared_inner(
15645 e: &Engine,
15646 m: &MoeWeights,
15647 z: &CudaSlice<f32>,
15648 t: usize,
15649 cfg: &ModelConfig,
15650 il: u16,
15651 moe_out: &mut CudaSlice<f32>,
15652 ) -> Result<(), Box<dyn std::error::Error>> {
15653 let n_embd = cfg.n_embd as usize;
15654 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
15655 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15656 {
15657 let n_ff_sh = gate_shexp.out_features();
15658 let lim = cfg.clamp_shexp_at(il as u32);
15659 let fused = t == 1
15666 && lim.is_none()
15667 && cfg.m3.is_none()
15668 && e.uses_q8_1_fast(gate_shexp)
15669 && e.uses_q8_1_fast(up_shexp);
15670 let canonical_w4a16_rows =
15671 t <= 32 && m.step_ep.as_ref().is_some_and(|ep| ep.nvfp4_device_routes);
15672 let bf16_dual = if (t == 1 || canonical_w4a16_rows)
15677 && crate::Engine::bf16_mmv_on()
15678 && n_embd.is_multiple_of(8)
15679 {
15680 match (gate_shexp, up_shexp) {
15681 (
15682 crate::model::GpuTensor::FloatBf16 { data: wg, .. },
15683 crate::model::GpuTensor::FloatBf16 { data: wu, .. },
15684 ) => Some((wg, wu)),
15685 _ => None,
15686 }
15687 } else {
15688 None
15689 };
15690 let sh = if let Some((wg, wu)) = bf16_dual {
15691 type SharedExpertWorkspace = (
15695 usize,
15696 usize,
15697 usize,
15698 CudaSlice<f32>,
15699 CudaSlice<f32>,
15700 CudaSlice<f32>,
15701 CudaSlice<f32>,
15702 );
15703 static SHEXP_WS: std::sync::Mutex<
15704 Option<std::collections::HashMap<usize, SharedExpertWorkspace>>,
15705 > = std::sync::Mutex::new(None);
15706 let down_bf16 = match down_shexp {
15707 crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
15708 _ => None,
15709 };
15710 let mut guard = SHEXP_WS
15711 .lock()
15712 .map_err(|_| "shexp workspace lock is poisoned")?;
15713 let capacity = if canonical_w4a16_rows { 32 } else { 1 };
15714 let device = e.ctx().ordinal();
15715 let workspaces = guard.get_or_insert_with(Default::default);
15716 if workspaces
15717 .get(&device)
15718 .is_none_or(|(ne, nf, cap, ..)| (*ne, *nf, *cap) != (n_embd, n_ff_sh, capacity))
15719 {
15720 workspaces.insert(
15721 device,
15722 (
15723 n_embd,
15724 n_ff_sh,
15725 capacity,
15726 e.uninit(capacity * n_ff_sh)?,
15727 e.uninit(capacity * n_ff_sh)?,
15728 e.uninit(capacity * n_ff_sh)?,
15729 e.uninit(capacity * n_embd)?,
15730 ),
15731 );
15732 }
15733 {
15736 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15737 let split_on = *ON
15738 .get_or_init(|| std::env::var("MEMRA_SHEXP_SPLIT").as_deref() == Ok("1"));
15739 if split_on
15740 && t == 1
15741 && let (Some(wd), Some(rank1)) = (
15742 match down_shexp {
15743 crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
15744 _ => None,
15745 },
15746 m.step_tp.as_ref().and_then(|st| st.runtime.rank_engine(1)),
15747 )
15748 && let Some(sh) = Self::shexp_split_matvec(
15749 e, rank1, wg, wu, wd, z, lim, cfg, il, n_embd, n_ff_sh,
15750 )?
15751 {
15752 drop(guard);
15753 let gate = match &m.gate_inp_shexp {
15754 Some(gate_inp_shexp) => {
15755 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
15756 }
15757 None => e.htod(&vec![1.0f32; t])?,
15758 };
15759 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
15760 return Ok(());
15761 }
15762 }
15763 let (_, _, _, gate, up, act, sh_buf) = workspaces
15764 .get_mut(&device)
15765 .expect("shexp workspace initialized above");
15766 if let (true, Ok(lim_post)) = (cfg.m3.is_none(), Self::fused_post_limit(lim)) {
15769 if canonical_w4a16_rows {
15772 e.matvec_bf16_dual_silu_rows_into(
15773 wg, wu, z, act, n_embd, n_ff_sh, lim_post, t,
15774 )?;
15775 } else {
15776 e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim_post)?;
15777 }
15778 let _ = (&gate, &up);
15779 } else {
15780 e.matvec_bf16_dual_into(wg, wu, z, gate, up, n_embd, n_ff_sh)?;
15781 Self::ffn_act_lim(e, cfg, gate, up, 1.0, 1.0, lim, act, n_ff_sh)?;
15782 }
15783 if let Some(down) = down_bf16 {
15784 if canonical_w4a16_rows {
15785 e.matvec_bf16_rows_into(down, act, sh_buf, n_ff_sh, n_embd, t)?;
15786 let mut sh = e.uninit(t * n_embd)?;
15787 {
15788 let mut dst = sh.slice_mut(0..t * n_embd);
15789 e.stream()
15790 .memcpy_dtod(&sh_buf.slice(0..t * n_embd), &mut dst)?;
15791 }
15792 sh
15793 } else {
15794 static FUSE_DA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15801 let fuse_da = *FUSE_DA.get_or_init(|| {
15802 std::env::var("MEMRA_FUSE_DOWN_ADDSCALE").as_deref() != Ok("0")
15803 });
15804 if fuse_da && m.gate_inp_shexp.is_none() {
15805 static ONES1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15806 std::sync::Mutex::new(None);
15807 let mut og = ONES1.lock().map_err(|_| "shexp ones lock is poisoned")?;
15808 if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15809 *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15810 }
15811 let ones = &og.as_ref().expect("armed above").1;
15812 e.matvec_bf16_down_addscale_into(
15813 down, act, ones, moe_out, n_ff_sh, n_embd,
15814 )?;
15815 return Ok(());
15816 }
15817 e.matvec_bf16_into(down, act, sh_buf, n_ff_sh, n_embd)?;
15818 let sh = e.uninit(n_embd)?;
15819 let mut sh = sh;
15821 {
15822 let mut dst = sh.slice_mut(0..n_embd);
15823 e.stream().memcpy_dtod(&sh_buf.slice(0..n_embd), &mut dst)?;
15824 }
15825 sh
15826 }
15827 } else {
15828 e.matmul(down_shexp, act, t)?
15829 }
15830 } else if fused {
15831 let (zq, zd) = e.quantize_q8_1(z, 1, n_embd)?;
15832 let pair = match e.matmul_pre_dual_noscale(gate_shexp, up_shexp, &zq, &zd, 1)? {
15833 Some((gate, up)) => Some((gate, up)),
15834 None => {
15835 match (
15836 e.matmul_pre_noscale(gate_shexp, &zq, &zd, 1)?,
15837 e.matmul_pre_noscale(up_shexp, &zq, &zd, 1)?,
15838 ) {
15839 (Some(gate), Some(up)) => Some((gate, up)),
15840 _ => None,
15841 }
15842 }
15843 };
15844 match pair {
15845 Some(((gate, gs), (up, us))) => {
15846 if e.uses_q8_1_fast(down_shexp) {
15847 let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, n_ff_sh)?;
15848 e.matmul_pre(down_shexp, &aq, &ad, &gate, 1)?
15849 } else {
15850 let mut act = e.uninit(n_ff_sh)?;
15851 e.silu_mul_scaled(&gate, &up, gs, us, &mut act, n_ff_sh)?;
15852 e.matmul(down_shexp, &act, 1)?
15853 }
15854 }
15855 None => {
15856 let gate = e.matmul_pre(gate_shexp, &zq, &zd, z, 1)?;
15857 let up = e.matmul_pre(up_shexp, &zq, &zd, z, 1)?;
15858 let mut act = e.uninit(n_ff_sh)?;
15859 Self::ffn_act(e, cfg, &gate, &up, &mut act, n_ff_sh)?;
15860 e.matmul(down_shexp, &act, 1)?
15861 }
15862 }
15863 } else {
15864 let sg_gate = e.matmul(gate_shexp, z, t)?;
15865 let sg_up = e.matmul(up_shexp, z, t)?;
15866 let mut sa = e.uninit(t * n_ff_sh)?;
15867 Self::ffn_act_lim(
15868 e,
15869 cfg,
15870 &sg_gate,
15871 &sg_up,
15872 1.0,
15873 1.0,
15874 lim,
15875 &mut sa,
15876 t * n_ff_sh,
15877 )?;
15878 e.matmul(down_shexp, &sa, t)?
15879 };
15880 let gate = match &m.gate_inp_shexp {
15881 Some(gate_inp_shexp) => {
15882 if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
15883 e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
15884 } else {
15885 let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
15886 let mut gate = e.uninit(t)?;
15887 e.sigmoid(&raw, &mut gate, t)?;
15888 gate
15889 }
15890 }
15891 None if t == 1 => {
15896 static ONES: std::sync::Mutex<
15897 Option<std::collections::HashMap<usize, CudaSlice<f32>>>,
15898 > = std::sync::Mutex::new(None);
15899 let mut guard = ONES.lock().map_err(|_| "shexp ones lock is poisoned")?;
15900 let device = e.ctx().ordinal();
15901 let rows = guard.get_or_insert_with(Default::default);
15902 use std::collections::hash_map::Entry;
15905 let ones = match rows.entry(device) {
15906 Entry::Occupied(occupied) => occupied.into_mut(),
15907 Entry::Vacant(vacant) => vacant.insert(e.htod(&[1.0f32])?),
15908 };
15909 e.add_scaled_rows(&sh, ones, moe_out, n_embd, t)?;
15910 return Ok(());
15911 }
15912 None => e.htod(&vec![1.0f32; t])?,
15913 };
15914 e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
15915 }
15916 Ok(())
15917 }
15918
15919 pub(crate) fn moe_ffn_grouped(
15922 e: &Engine,
15923 m: &MoeWeights,
15924 z: &CudaSlice<f32>,
15925 t: usize,
15926 cfg: &ModelConfig,
15927 il: u16,
15928 max_block: usize,
15929 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15930 let moe = cfg.moe.as_ref().unwrap();
15931 let n_embd = cfg.n_embd as usize;
15932 let n_expert = moe.expert_count as usize;
15933 let n_used = moe.expert_used_count as usize;
15934 let n_ff_exp = moe.expert_ff_length as usize;
15935 let lim_exp = cfg.clamp_exp_at(il as u32);
15937
15938 let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
15942 if let Some(sig) = cfg.sigmoid_router() {
15943 Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
15944 }
15945 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
15946 Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
15947 } else {
15948 Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
15949 };
15950 crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
15951 Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
15952 Self::trace_moe_input(e, il, t, n_embd, z)?;
15953
15954 let no_exp_macros = m.gate_exps.macros.is_none()
15959 && m.up_exps.macros.is_none()
15960 && m.down_exps.macros.is_none();
15961 let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
15962 m.has_uniform_expert_layout()
15963 && no_exp_macros
15964 && moe_q8_enabled_for_model(cfg, m)
15965 && moe_slab_enabled()
15966 && dev.dev == e.ctx().ordinal()
15967 });
15968 if let Some(dev) = resident_q8 {
15969 let mut moe_out = Self::moe_ffn_grouped_resident_q8(
15970 e,
15971 m,
15972 z,
15973 t,
15974 cfg,
15975 il,
15976 &sel_all,
15977 &w_all,
15978 &dev.ptr_row,
15979 dev.gu_il,
15980 )?;
15981 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
15982 return Ok(moe_out);
15983 }
15984
15985 struct ExpertGroup {
15989 tok_indices: Vec<i32>, slot_indices: Vec<i32>, weights: Vec<f32>, }
15993 let mut groups: Vec<ExpertGroup> = (0..n_expert)
15994 .map(|_| ExpertGroup {
15995 tok_indices: Vec::new(),
15996 slot_indices: Vec::new(),
15997 weights: Vec::new(),
15998 })
15999 .collect();
16000
16001 for tok in 0..t {
16002 for j in 0..n_used {
16003 let ex = sel_all[tok * n_used + j] as usize;
16004 let w = w_all[tok * n_used + j];
16005 groups[ex].tok_indices.push(tok as i32);
16006 groups[ex].slot_indices.push(j as i32);
16007 groups[ex].weights.push(w);
16008 }
16009 }
16010
16011 let mut slot_buf = e.zeros(t * n_used * n_embd)?;
16014 let mut wbuf = e.zeros(t * n_used)?; let g_len = m.gate_exps.max_expert_bytes();
16018 let u_len = m.up_exps.max_expert_bytes();
16019 let d_len = m.down_exps.max_expert_bytes();
16020 let moe_q8 = moe_q8_enabled_for_model(cfg, m);
16021 let slab_local = m
16024 .dev_exps
16025 .as_ref()
16026 .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
16027 let use_cache =
16028 slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
16029 let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
16032
16033 let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
16035 (
16036 Some(e.alloc_u8(g_len)?),
16037 Some(e.alloc_u8(u_len)?),
16038 Some(e.alloc_u8(d_len)?),
16039 )
16040 } else {
16041 (None, None, None)
16042 };
16043
16044 let mut order: Vec<usize> = (0..n_expert)
16055 .filter(|&ex| !groups[ex].tok_indices.is_empty())
16056 .collect();
16057 order.sort_by(|&a, &b| {
16058 groups[b]
16059 .tok_indices
16060 .len()
16061 .cmp(&groups[a].tok_indices.len())
16062 .then(a.cmp(&b))
16063 });
16064 let mut m_dist: Vec<usize> = Vec::new(); let page_window = moe_page_prefetch_window();
16066 let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
16067 if worker_disk_prefetch
16068 && let Some(first) = grouped_worker_prefetch_position(order.len(), None)
16069 {
16070 Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
16071 }
16072 for (order_pos, &ex) in order.iter().enumerate() {
16073 for next in page_prefetch_positions(order_pos, order.len(), page_window) {
16074 Self::moe_prefetch_host_expert(order[next], m);
16075 }
16076 if worker_disk_prefetch
16077 && let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos))
16078 {
16079 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16080 let keep = [
16081 BlockId::new(il, PROJ_GATE, ex as u16),
16082 BlockId::new(il, PROJ_UP, ex as u16),
16083 BlockId::new(il, PROJ_DOWN, ex as u16),
16084 ];
16085 Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
16086 }
16087 let grp = &groups[ex];
16088 let m_e = grp.tok_indices.len();
16089 m_dist.push(m_e);
16090 let gl = m.gate_exps.expert_layout(ex);
16091 let ul = m.up_exps.expert_layout(ex);
16092 let dl = m.down_exps.expert_layout(ex);
16093
16094 let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
16098 let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
16099 let dmac = m.down_exps.macro_scale(ex);
16100 let weight_d = if dmac == 1.0 {
16101 e.htod(&grp.weights)?
16102 } else {
16103 let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
16104 e.htod(&scaled)?
16105 };
16106
16107 let mut gathered = e.zeros(m_e * n_embd)?;
16109 e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
16110 let gv = gathered.slice(0..m_e * n_embd);
16111
16112 let y = if let Some(dev) = slab_local {
16115 let gate_start = ex * m.gate_exps.expert_stride;
16116 let up_start = ex * m.up_exps.expert_stride;
16117 let down_start = ex * m.down_exps.expert_stride;
16118 if grouped_q8 {
16119 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16120 let gate = e.qmatvec_expert_q8(
16121 &dev.gate,
16122 gate_start..gate_start + gl.len,
16123 &zq,
16124 &zd,
16125 m_e,
16126 m.gate_exps.in_f,
16127 m.gate_exps.out_f,
16128 gl.qtype,
16129 gl.row_bytes,
16130 )?;
16131 let up = e.qmatvec_expert_q8(
16132 &dev.up,
16133 up_start..up_start + ul.len,
16134 &zq,
16135 &zd,
16136 m_e,
16137 m.up_exps.in_f,
16138 m.up_exps.out_f,
16139 ul.qtype,
16140 ul.row_bytes,
16141 )?;
16142 let mut act = e.uninit(m_e * n_ff_exp)?;
16143 Self::ffn_act_lim(
16144 e,
16145 cfg,
16146 &gate,
16147 &up,
16148 m.gate_exps.macro_scale(ex),
16149 m.up_exps.macro_scale(ex),
16150 lim_exp,
16151 &mut act,
16152 m_e * n_ff_exp,
16153 )?;
16154 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16155 e.qmatvec_expert_q8(
16156 &dev.down,
16157 down_start..down_start + dl.len,
16158 &aq2,
16159 &ad2,
16160 m_e,
16161 m.down_exps.in_f,
16162 m.down_exps.out_f,
16163 dl.qtype,
16164 dl.row_bytes,
16165 )?
16166 } else {
16167 let gate = m.qmatvec_view(
16168 e,
16169 &dev.gate,
16170 gate_start..gate_start + gl.len,
16171 &gv,
16172 m_e,
16173 m.gate_exps.in_f,
16174 m.gate_exps.out_f,
16175 gl.qtype,
16176 gl.row_bytes,
16177 )?;
16178 let up = m.qmatvec_view(
16179 e,
16180 &dev.up,
16181 up_start..up_start + ul.len,
16182 &gv,
16183 m_e,
16184 m.up_exps.in_f,
16185 m.up_exps.out_f,
16186 ul.qtype,
16187 ul.row_bytes,
16188 )?;
16189 let mut act = e.uninit(m_e * n_ff_exp)?;
16190 Self::ffn_act_lim(
16191 e,
16192 cfg,
16193 &gate,
16194 &up,
16195 m.gate_exps.macro_scale(ex),
16196 m.up_exps.macro_scale(ex),
16197 lim_exp,
16198 &mut act,
16199 m_e * n_ff_exp,
16200 )?;
16201 let actv = act.slice(0..m_e * n_ff_exp);
16202 m.qmatvec_view(
16203 e,
16204 &dev.down,
16205 down_start..down_start + dl.len,
16206 &actv,
16207 m_e,
16208 m.down_exps.in_f,
16209 m.down_exps.out_f,
16210 dl.qtype,
16211 dl.row_bytes,
16212 )?
16213 }
16214 } else if use_cache {
16215 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16216 if grouped_q8 {
16217 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16218 let gate = e.with_moe_cache(max_block, |cache, eng| {
16219 let id = BlockId::new(il, PROJ_GATE, ex as u16);
16220 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
16221 eng.qmatvec_expert_q8(
16222 cache.buf(slot),
16223 0..gl.len,
16224 &zq,
16225 &zd,
16226 m_e,
16227 m.gate_exps.in_f,
16228 m.gate_exps.out_f,
16229 gl.qtype,
16230 gl.row_bytes,
16231 )
16232 })?;
16233 let up = e.with_moe_cache(max_block, |cache, eng| {
16234 let id = BlockId::new(il, PROJ_UP, ex as u16);
16235 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
16236 eng.qmatvec_expert_q8(
16237 cache.buf(slot),
16238 0..ul.len,
16239 &zq,
16240 &zd,
16241 m_e,
16242 m.up_exps.in_f,
16243 m.up_exps.out_f,
16244 ul.qtype,
16245 ul.row_bytes,
16246 )
16247 })?;
16248 let mut act = e.uninit(m_e * n_ff_exp)?;
16249 Self::ffn_act_lim(
16250 e,
16251 cfg,
16252 &gate,
16253 &up,
16254 m.gate_exps.macro_scale(ex),
16255 m.up_exps.macro_scale(ex),
16256 lim_exp,
16257 &mut act,
16258 m_e * n_ff_exp,
16259 )?;
16260 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16261 e.with_moe_cache(max_block, |cache, eng| {
16262 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
16263 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
16264 eng.qmatvec_expert_q8(
16265 cache.buf(slot),
16266 0..dl.len,
16267 &aq2,
16268 &ad2,
16269 m_e,
16270 m.down_exps.in_f,
16271 m.down_exps.out_f,
16272 dl.qtype,
16273 dl.row_bytes,
16274 )
16275 })?
16276 } else {
16277 let gate = e.with_moe_cache(max_block, |cache, eng| {
16278 let id = BlockId::new(il, PROJ_GATE, ex as u16);
16279 let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
16280 m.qmatvec_view(
16281 eng,
16282 cache.buf(slot),
16283 0..gl.len,
16284 &gv,
16285 m_e,
16286 m.gate_exps.in_f,
16287 m.gate_exps.out_f,
16288 gl.qtype,
16289 gl.row_bytes,
16290 )
16291 })?;
16292 let up = e.with_moe_cache(max_block, |cache, eng| {
16293 let id = BlockId::new(il, PROJ_UP, ex as u16);
16294 let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
16295 m.qmatvec_view(
16296 eng,
16297 cache.buf(slot),
16298 0..ul.len,
16299 &gv,
16300 m_e,
16301 m.up_exps.in_f,
16302 m.up_exps.out_f,
16303 ul.qtype,
16304 ul.row_bytes,
16305 )
16306 })?;
16307 let mut act = e.uninit(m_e * n_ff_exp)?;
16308 Self::ffn_act_lim(
16309 e,
16310 cfg,
16311 &gate,
16312 &up,
16313 m.gate_exps.macro_scale(ex),
16314 m.up_exps.macro_scale(ex),
16315 lim_exp,
16316 &mut act,
16317 m_e * n_ff_exp,
16318 )?;
16319 let actv = act.slice(0..m_e * n_ff_exp);
16320 e.with_moe_cache(max_block, |cache, eng| {
16321 let id = BlockId::new(il, PROJ_DOWN, ex as u16);
16322 let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
16323 m.qmatvec_view(
16324 eng,
16325 cache.buf(slot),
16326 0..dl.len,
16327 &actv,
16328 m_e,
16329 m.down_exps.in_f,
16330 m.down_exps.out_f,
16331 dl.qtype,
16332 dl.row_bytes,
16333 )
16334 })?
16335 }
16336 } else {
16337 let sg = scratch_g.as_mut().unwrap();
16338 let su = scratch_u.as_mut().unwrap();
16339 let sd = scratch_d.as_mut().unwrap();
16340 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
16341 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
16342 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
16343 if grouped_q8 {
16344 let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16345 let gate = e.qmatvec_expert_q8(
16346 sg,
16347 0..gl.len,
16348 &zq,
16349 &zd,
16350 m_e,
16351 m.gate_exps.in_f,
16352 m.gate_exps.out_f,
16353 gl.qtype,
16354 gl.row_bytes,
16355 )?;
16356 let up = e.qmatvec_expert_q8(
16357 su,
16358 0..ul.len,
16359 &zq,
16360 &zd,
16361 m_e,
16362 m.up_exps.in_f,
16363 m.up_exps.out_f,
16364 ul.qtype,
16365 ul.row_bytes,
16366 )?;
16367 let mut act = e.uninit(m_e * n_ff_exp)?;
16368 Self::ffn_act_lim(
16369 e,
16370 cfg,
16371 &gate,
16372 &up,
16373 m.gate_exps.macro_scale(ex),
16374 m.up_exps.macro_scale(ex),
16375 lim_exp,
16376 &mut act,
16377 m_e * n_ff_exp,
16378 )?;
16379 let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16380 e.qmatvec_expert_q8(
16381 sd,
16382 0..dl.len,
16383 &aq2,
16384 &ad2,
16385 m_e,
16386 m.down_exps.in_f,
16387 m.down_exps.out_f,
16388 dl.qtype,
16389 dl.row_bytes,
16390 )?
16391 } else {
16392 let gate = m.qmatvec_view(
16393 e,
16394 sg,
16395 0..gl.len,
16396 &gv,
16397 m_e,
16398 m.gate_exps.in_f,
16399 m.gate_exps.out_f,
16400 gl.qtype,
16401 gl.row_bytes,
16402 )?;
16403 let up = m.qmatvec_view(
16404 e,
16405 su,
16406 0..ul.len,
16407 &gv,
16408 m_e,
16409 m.up_exps.in_f,
16410 m.up_exps.out_f,
16411 ul.qtype,
16412 ul.row_bytes,
16413 )?;
16414 let mut act = e.uninit(m_e * n_ff_exp)?;
16415 Self::ffn_act_lim(
16416 e,
16417 cfg,
16418 &gate,
16419 &up,
16420 m.gate_exps.macro_scale(ex),
16421 m.up_exps.macro_scale(ex),
16422 lim_exp,
16423 &mut act,
16424 m_e * n_ff_exp,
16425 )?;
16426 let actv = act.slice(0..m_e * n_ff_exp);
16427 m.qmatvec_view(
16428 e,
16429 sd,
16430 0..dl.len,
16431 &actv,
16432 m_e,
16433 m.down_exps.in_f,
16434 m.down_exps.out_f,
16435 dl.qtype,
16436 dl.row_bytes,
16437 )?
16438 }
16439 };
16440
16441 e.scatter_slot(
16443 &y,
16444 &tok_idx_d,
16445 &slot_idx_d,
16446 &weight_d,
16447 &mut slot_buf,
16448 &mut wbuf,
16449 n_embd,
16450 n_used,
16451 m_e,
16452 )?;
16453 }
16454
16455 let mut moe_out = e.zeros(t * n_embd)?;
16457 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
16458
16459 if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
16461 m_dist.sort_unstable();
16462 let active = m_dist.len();
16463 let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
16464 let median = m_dist[active / 2];
16465 let max_m = *m_dist.last().unwrap();
16466 let min_m = m_dist[0];
16467 let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
16468 println!(
16469 "moe-grouped il={il} t={t} active={active}/{n_expert} \
16470 m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
16471 above_gemm_threshold(>=16)={above16}/{active}"
16472 );
16473 }
16474
16475 Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
16476 Ok(moe_out)
16477 }
16478
16479 pub(crate) fn moe_ffn_lockstep(
16486 &self,
16487 e: &Engine,
16488 m: &MoeWeights,
16489 zbatch: &CudaSlice<f32>,
16490 mrows: usize,
16491 il: u16,
16492 max_block: usize,
16493 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16494 use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16495 let cfg = &self.cfg;
16496 let moe = cfg.moe.as_ref().unwrap();
16497 let n_embd = cfg.n_embd as usize;
16498 let n_expert = moe.expert_count as usize;
16499 let n_used = moe.expert_used_count as usize;
16500 let n_ff_exp = moe.expert_ff_length as usize;
16501 let lim_exp = cfg.clamp_exp_at(il as u32);
16503 let lim_shexp = cfg.clamp_shexp_at(il as u32);
16504
16505 let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
16506 if let Some(sig) = cfg.sigmoid_router() {
16507 Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
16508 }
16509 let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
16510 Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
16511 } else {
16512 Self::moe_route_cfg(
16513 e,
16514 &logits,
16515 mrows,
16516 n_expert,
16517 n_used,
16518 m.active_experts.as_deref(),
16519 )?
16520 };
16521 Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
16522
16523 let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
16525 Ok((0..n_expert)
16526 .map(|ex| {
16527 [PROJ_GATE, PROJ_UP, PROJ_DOWN]
16528 .into_iter()
16529 .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
16530 })
16531 .collect())
16532 })?;
16533
16534 struct Group {
16535 rows: Vec<i32>,
16536 slots: Vec<i32>,
16537 weights: Vec<f32>,
16538 }
16539 let mut groups: std::collections::HashMap<usize, Group> = Default::default();
16540 let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
16541 let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
16542 Default::default();
16543 for row in 0..mrows {
16544 for j in 0..n_used {
16545 let ex = sel_all[row * n_used + j] as usize;
16546 let w = w_all[row * n_used + j];
16547 if resident_expert[ex] {
16548 let group = groups.entry(ex).or_insert_with(|| Group {
16549 rows: Vec::new(),
16550 slots: Vec::new(),
16551 weights: Vec::new(),
16552 });
16553 group.rows.push(row as i32);
16554 group.slots.push(j as i32);
16555 group.weights.push(w);
16556 } else {
16557 crate::cpu_experts::record_incomplete_gpu_residency(0);
16558 cpu_rows[row].push((ex, w));
16559 cpu_by_expert.entry(ex).or_default().push((row, w));
16560 }
16561 }
16562 }
16563
16564 let host_rows = e.dtoh(zbatch)?;
16570 let rows_ok = crate::cpu_experts::rows_supported();
16571 enum CpuPart {
16572 Single { row: usize },
16573 Rows { rows: Vec<usize> },
16574 }
16575 let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
16576 let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
16577 if rows_ok {
16578 let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
16579 .into_iter()
16580 .filter(|(_, rows)| rows.len() >= 2)
16581 .collect();
16582 shared.sort_by_key(|(ex, _)| *ex);
16583 for (ex, mut row_weights) in shared {
16584 row_weights.sort_by_key(|(row, _)| *row);
16585 let inputs: Vec<(&[f32], f32)> = row_weights
16586 .iter()
16587 .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
16588 .collect();
16589 let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
16590 .map_err(std::io::Error::other)?;
16591 for &(row, _) in &row_weights {
16592 rows_served.insert((row, ex));
16593 }
16594 tickets.push((
16595 CpuPart::Rows {
16596 rows: row_weights.iter().map(|&(row, _)| row).collect(),
16597 },
16598 crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
16599 ));
16600 }
16601 }
16602 for (row, selected) in cpu_rows.iter().enumerate() {
16603 let leftover: Vec<(usize, f32)> = selected
16604 .iter()
16605 .copied()
16606 .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
16607 .collect();
16608 if leftover.is_empty() {
16609 continue;
16610 }
16611 let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
16612 let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
16613 .map_err(std::io::Error::other)?;
16614 tickets.push((
16615 CpuPart::Single { row },
16616 crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
16617 ));
16618 }
16619
16620 let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
16621 let mut wbuf = e.zeros(mrows * n_used)?;
16622 let mut order: Vec<usize> = groups.keys().copied().collect();
16623 order.sort_by(|&a, &b| {
16624 groups[&b]
16625 .rows
16626 .len()
16627 .cmp(&groups[&a].rows.len())
16628 .then(a.cmp(&b))
16629 });
16630 for &ex in &order {
16631 let group = &groups[&ex];
16632 let m_e = group.rows.len();
16633 let gl = m.gate_exps.expert_layout(ex);
16634 let ul = m.up_exps.expert_layout(ex);
16635 let dl = m.down_exps.expert_layout(ex);
16636 let row_idx_d = e.htod_i32(&group.rows)?;
16637 let slot_idx_d = e.htod_i32(&group.slots)?;
16638 let dmac = m.down_exps.macro_scale(ex);
16639 let weight_d = if dmac == 1.0 {
16640 e.htod(&group.weights)?
16641 } else {
16642 let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
16643 e.htod(&scaled)?
16644 };
16645 let mut gathered = e.zeros(m_e * n_embd)?;
16646 e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
16647 let gv = gathered.slice(0..m_e * n_embd);
16648 let gate = e.with_moe_cache(max_block, |c, eng| {
16649 let slot = c
16650 .resident(BlockId::new(il, PROJ_GATE, ex as u16))
16651 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16652 m.qmatvec_view(
16653 eng,
16654 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16655 0..gl.len,
16656 &gv,
16657 m_e,
16658 m.gate_exps.in_f,
16659 m.gate_exps.out_f,
16660 gl.qtype,
16661 gl.row_bytes,
16662 )
16663 })?;
16664 let up = e.with_moe_cache(max_block, |c, eng| {
16665 let slot = c
16666 .resident(BlockId::new(il, PROJ_UP, ex as u16))
16667 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16668 m.qmatvec_view(
16669 eng,
16670 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16671 0..ul.len,
16672 &gv,
16673 m_e,
16674 m.up_exps.in_f,
16675 m.up_exps.out_f,
16676 ul.qtype,
16677 ul.row_bytes,
16678 )
16679 })?;
16680 let mut act = e.zeros(m_e * n_ff_exp)?;
16681 Self::ffn_act_lim(
16682 e,
16683 cfg,
16684 &gate,
16685 &up,
16686 m.gate_exps.macro_scale(ex),
16687 m.up_exps.macro_scale(ex),
16688 lim_exp,
16689 &mut act,
16690 m_e * n_ff_exp,
16691 )?;
16692 let actv = act.slice(0..m_e * n_ff_exp);
16693 let y = e.with_moe_cache(max_block, |c, eng| {
16694 let slot = c
16695 .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
16696 .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16697 m.qmatvec_view(
16698 eng,
16699 c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16700 0..dl.len,
16701 &actv,
16702 m_e,
16703 m.down_exps.in_f,
16704 m.down_exps.out_f,
16705 dl.qtype,
16706 dl.row_bytes,
16707 )
16708 })?;
16709 e.scatter_slot(
16710 &y,
16711 &row_idx_d,
16712 &slot_idx_d,
16713 &weight_d,
16714 &mut slot_buf,
16715 &mut wbuf,
16716 n_embd,
16717 n_used,
16718 m_e,
16719 )?;
16720 }
16721 let mut moe_out = e.zeros(mrows * n_embd)?;
16722 e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
16723
16724 let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
16726 for (part, ticket) in tickets {
16727 let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
16728 let mut add_row = |row: usize, chunk: &[f32]| {
16729 let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
16730 for (accumulator, value) in sum.iter_mut().zip(chunk) {
16731 *accumulator += value;
16732 }
16733 };
16734 match part {
16735 CpuPart::Single { row } => add_row(row, &cpu_output),
16736 CpuPart::Rows { rows } => {
16737 for (slot, row) in rows.into_iter().enumerate() {
16738 add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
16739 }
16740 }
16741 }
16742 }
16743 for (row, sum) in row_sums.into_iter().enumerate() {
16744 let Some(sum) = sum else { continue };
16745 let cpu_output = e.htod(&sum)?;
16746 let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
16747 e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
16748 }
16749
16750 if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
16751 (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
16752 {
16753 let n_ff_sh = gate_shexp.out_features();
16754 let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
16755 let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
16756 let mut sa = e.zeros(mrows * n_ff_sh)?;
16757 Self::ffn_act_lim(
16758 e,
16759 cfg,
16760 &sg_gate,
16761 &sg_up,
16762 1.0,
16763 1.0,
16764 lim_shexp,
16765 &mut sa,
16766 mrows * n_ff_sh,
16767 )?;
16768 let sh = e.matmul(down_shexp, &sa, mrows)?;
16769 let g = match &m.gate_inp_shexp {
16772 Some(gate_inp_shexp) => {
16773 e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
16774 }
16775 None => e.htod(&vec![1.0f32; mrows])?,
16776 };
16777 e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
16778 }
16779
16780 Ok(moe_out)
16781 }
16782}
16783
16784impl HybridModel {
16790 pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
16804 let g = self
16805 .cfg
16806 .gemma4
16807 .as_ref()
16808 .expect("gemma4_rope_dims on a non-gemma4 config");
16809 if g.swa_pattern[il] {
16810 g.rope_dims_swa as usize
16811 } else {
16812 g.rope_dims_global as usize
16813 }
16814 }
16815
16816 pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
16817 let g = self.cfg.gemma4.as_ref().unwrap();
16818 let swa = g.swa_pattern[il];
16819 let hd = if swa {
16820 g.key_length_swa
16821 } else {
16822 g.key_length_global
16823 } as usize;
16824 (
16828 hd,
16829 g.head_count_kv[il] as usize,
16830 self.cfg.n_head as usize,
16831 if swa {
16832 g.rope_base_swa
16833 } else {
16834 g.rope_base_global
16835 },
16836 1.0,
16837 swa,
16838 )
16839 }
16840
16841 pub(crate) fn gemma4_suppress(
16845 &self,
16846 e: &Engine,
16847 ld: &mut CudaSlice<f32>,
16848 t: usize,
16849 ) -> Result<(), Box<dyn std::error::Error>> {
16850 if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
16851 #[cfg(debug_assertions)]
16856 crate::debug_assert_tensor_stream_device(
16857 ids,
16858 &e.stream(),
16859 "gemma4_suppress.suppress_d",
16860 );
16861 e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
16862 }
16863 Ok(())
16864 }
16865
16866 #[allow(clippy::too_many_arguments)]
16871 fn gemma_fa_one_program() -> bool {
16880 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16881 *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
16882 }
16883
16884 #[allow(clippy::too_many_arguments)] fn gemma4_attn_prime(
16886 &self,
16887 e: &Engine,
16888 fa: &crate::hybrid::FullAttnLayer,
16889 il: usize,
16890 h: &CudaSlice<f32>,
16891 pos_d: &CudaSlice<i32>,
16892 t: usize,
16893 cache: Option<&mut Cache>,
16894 island: Option<&CudaSlice<i32>>,
16895 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16896 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
16897 let eps = self.cfg.rms_eps;
16898 let aux = self.gemma4_aux.as_ref().unwrap();
16899 let ones = aux.ones(e);
16900 #[cfg(debug_assertions)]
16901 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
16902
16903 e.mmq_act_begin();
16906 let q0 = e.matmul(&fa.wq, h, t)?; if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
16908 let v = e.dtoh(&q0)?;
16909 let nan = v.iter().filter(|x| x.is_nan()).count();
16910 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
16911 eprintln!(
16912 "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
16913 v.len()
16914 );
16915 }
16916 let k0 = e.matmul(&fa.wk, h, t)?; let v0 = if swa {
16920 e.matmul(&fa.wv, h, t)?
16921 } else {
16922 e.clone_dtod(&k0)?
16923 };
16924 if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
16925 for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
16926 let v = e.dtoh(buf)?;
16927 let nan = v.iter().filter(|x| x.is_nan()).count();
16928 let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
16929 eprintln!(
16930 "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
16931 v.len()
16932 );
16933 }
16934 }
16935
16936 let mut q = e.uninit(t * nh * hd)?;
16937 let mut k = e.uninit(t * nkv * hd)?;
16938 let mut v = e.uninit(t * nkv * hd)?;
16940 static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16944 let emit = island.is_none()
16947 && t >= 16
16948 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
16949 && *EMIT.get_or_init(|| {
16950 std::env::var("MEMRA_FA_EMIT")
16951 .map(|s| s != "0")
16952 .unwrap_or(true)
16953 });
16954 let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
16955 let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
16956 let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
16957 let v_f16 = emit
16960 && crate::fa_f16pv_on()
16961 && match hd {
16962 512 => true,
16963 256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
16964 _ => false,
16965 };
16966 if emit {
16967 e.rms_norm_qkv_w4b(
16968 &q0,
16969 &k0,
16970 &v0,
16971 fa.q_norm.float_data(),
16972 fa.k_norm.float_data(),
16973 ones,
16974 &mut q,
16975 &mut k,
16976 &mut v,
16977 &mut vb,
16978 hd,
16979 nh * t,
16980 nkv * t,
16981 eps,
16982 v_f16,
16983 )?;
16984 } else {
16985 e.rms_norm_qkv(
16986 &q0,
16987 &k0,
16988 &v0,
16989 fa.q_norm.float_data(),
16990 fa.k_norm.float_data(),
16991 ones,
16992 &mut q,
16993 &mut k,
16994 &mut v,
16995 hd,
16996 nh * t,
16997 nkv * t,
16998 eps,
16999 )?;
17000 }
17001
17002 let ff = if swa {
17003 None
17004 } else {
17005 Some(
17006 aux.rope_freqs(e)
17007 .expect("gemma4 global rope needs rope_freqs.weight"),
17008 )
17009 };
17010 #[cfg(debug_assertions)]
17011 if let Some(ff) = ff {
17012 crate::debug_assert_tensor_stream_device(
17013 ff,
17014 &e.stream(),
17015 "gemma4_attn_prime.rope_freqs",
17016 );
17017 }
17018 if emit {
17019 e.rope_neox2_bf16e(
17020 &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
17021 )?;
17022 } else {
17023 e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
17024 }
17025
17026 if let Some(cache) = cache {
17027 let kvl = cache.kv[il].as_mut().unwrap();
17028 assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
17029 e.append_kv_quantized_rows(
17030 &k,
17031 &v,
17032 &mut kvl.k,
17033 &mut kvl.v,
17034 kvl.len,
17035 t,
17036 kvl.kv_dim_k,
17037 kvl.kv_dim_v,
17038 kvl.k_tok_bytes,
17039 kvl.v_tok_bytes,
17040 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
17041 )?;
17042 kvl.len += t;
17043 }
17044 let mut attn = e.zeros(t * nh * hd)?;
17045 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
17049 if let Some(span) = island {
17050 let w = if swa && t > win { win } else { 0 };
17055 e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
17056 } else if swa && (t > win || Self::gemma_fa_one_program()) {
17057 if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
17058 if emit {
17059 e.fa_prefill_w_pre(
17060 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
17061 )?;
17062 } else {
17063 e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
17064 }
17065 } else {
17066 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
17067 }
17068 } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
17069 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17070 } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
17071 if emit {
17072 e.fa_prefill_hd512_pre(
17073 &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
17074 )?;
17075 } else {
17076 e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17077 }
17078 } else {
17079 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17080 }
17081 e.matmul(&fa.wo, &attn, t)
17082 }
17083
17084 fn gemma4_attn(
17086 &self,
17087 e: &Engine,
17088 fa: &crate::hybrid::FullAttnLayer,
17089 il: usize,
17090 h: &CudaSlice<f32>,
17091 pos_d: &CudaSlice<i32>,
17092 t: usize,
17093 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17094 self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
17095 }
17096
17097 fn gemma4_moe_q8(
17102 &self,
17103 e: &Engine,
17104 m: &crate::hybrid::MoeWeights,
17105 bits: &crate::hybrid::Gemma4MoeBits,
17106 mq: &(CudaSlice<i8>, CudaSlice<f32>),
17107 router_in: &CudaSlice<f32>,
17108 t: usize,
17109 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17110 let cfg = &self.cfg;
17111 let moe = cfg.moe.as_ref().unwrap();
17112 let n_embd = cfg.n_embd as usize;
17113 let n_expert = moe.expert_count as usize;
17114 let n_used = moe.expert_used_count as usize;
17115 let n_ff_exp = moe.expert_ff_length as usize;
17116 let logits = if crate::router_kernel_on() {
17120 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
17121 } else {
17122 e.matmul(&m.gate_inp, router_in, t)?
17123 };
17124 let dev = m.dev_exps.as_ref().unwrap();
17125 let (sel_d, w_d) =
17126 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
17127 let (zq, zd) = mq;
17128 if t == 1 {
17129 let selv = sel_d.slice(0..n_used);
17130 let wv = w_d.slice(0..n_used);
17131 let act = e.moe_gate_up_gelu8_dev_q8(
17132 &dev.ptr_row,
17133 &selv,
17134 zq,
17135 zd,
17136 n_embd,
17137 n_ff_exp,
17138 n_used,
17139 n_expert,
17140 m.gate_exps.qtype,
17141 m.up_exps.qtype,
17142 m.gate_exps.row_bytes,
17143 m.up_exps.row_bytes,
17144 )?;
17145 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
17146 let mut moe_out = e.uninit(n_embd)?;
17147 e.moe_down8_fma_dev_q8(
17148 &dev.ptr_row,
17149 &selv,
17150 &wv,
17151 &aq2,
17152 &ad2,
17153 &mut moe_out.slice_mut(0..n_embd),
17154 n_ff_exp,
17155 n_embd,
17156 n_used,
17157 n_expert,
17158 m.down_exps.qtype,
17159 m.down_exps.row_bytes,
17160 )?;
17161 return Ok(moe_out);
17162 }
17163 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
17164 let act = if csr {
17165 e.moe_gate_up_gelu8_dev_q8_csr(
17166 &dev.ptr_row,
17167 &sel_d,
17168 zq,
17169 zd,
17170 t * n_used,
17171 n_embd,
17172 n_ff_exp,
17173 n_used,
17174 n_expert,
17175 m.gate_exps.qtype,
17176 m.up_exps.qtype,
17177 m.gate_exps.row_bytes,
17178 m.up_exps.row_bytes,
17179 )?
17180 } else {
17181 e.moe_gate_up_gelu8_dev_q8_rows(
17182 &dev.ptr_row,
17183 &sel_d,
17184 zq,
17185 zd,
17186 t,
17187 n_embd,
17188 n_ff_exp,
17189 n_used,
17190 n_expert,
17191 m.gate_exps.qtype,
17192 m.up_exps.qtype,
17193 m.gate_exps.row_bytes,
17194 m.up_exps.row_bytes,
17195 )?
17196 };
17197 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
17198 let mut moe_out = e.uninit(t * n_embd)?;
17199 e.moe_down8_fma_dev_q8_rows_g(
17202 &dev.ptr_row,
17203 &sel_d,
17204 &w_d,
17205 &aq2,
17206 &ad2,
17207 &mut moe_out,
17208 t,
17209 n_ff_exp,
17210 n_embd,
17211 n_used,
17212 n_expert,
17213 m.down_exps.qtype,
17214 m.down_exps.row_bytes,
17215 )?;
17216 Ok(moe_out)
17217 }
17218
17219 fn gemma4_moe(
17223 &self,
17224 e: &Engine,
17225 m: &crate::hybrid::MoeWeights,
17226 bits: &crate::hybrid::Gemma4MoeBits,
17227 moe_in: &CudaSlice<f32>,
17228 router_in: &CudaSlice<f32>,
17229 t: usize,
17230 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17231 let cfg = &self.cfg;
17232 let moe = cfg.moe.as_ref().unwrap();
17233 let n_embd = cfg.n_embd as usize;
17234 let n_expert = moe.expert_count as usize;
17235 let n_used = moe.expert_used_count as usize;
17236 let n_ff_exp = moe.expert_ff_length as usize;
17237
17238 let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
17242 e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
17243 } else {
17244 e.matmul(&m.gate_inp, router_in, t)?
17245 };
17246
17247 if t < PRIME_MIN_T
17252 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17253 && expert_dp4a_supported(m.gate_exps.qtype)
17254 && expert_dp4a_supported(m.up_exps.qtype)
17255 && expert_dp4a_supported(m.down_exps.qtype)
17256 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
17257 {
17258 let dev = m.dev_exps.as_ref().unwrap();
17259 let (sel_d, w_d) =
17260 e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
17261 if t == 1 {
17262 let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
17263 let selv = sel_d.slice(0..n_used);
17264 let wv = w_d.slice(0..n_used);
17265 let act = e.moe_gate_up_gelu8_dev_q8(
17266 &dev.ptr_row,
17267 &selv,
17268 &zq,
17269 &zd,
17270 n_embd,
17271 n_ff_exp,
17272 n_used,
17273 n_expert,
17274 m.gate_exps.qtype,
17275 m.up_exps.qtype,
17276 m.gate_exps.row_bytes,
17277 m.up_exps.row_bytes,
17278 )?;
17279 let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
17280 let mut moe_out = e.uninit(n_embd)?;
17281 e.moe_down8_fma_dev_q8(
17282 &dev.ptr_row,
17283 &selv,
17284 &wv,
17285 &aq2,
17286 &ad2,
17287 &mut moe_out.slice_mut(0..n_embd),
17288 n_ff_exp,
17289 n_embd,
17290 n_used,
17291 n_expert,
17292 m.down_exps.qtype,
17293 m.down_exps.row_bytes,
17294 )?;
17295 return Ok(moe_out);
17296 }
17297 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
17302 let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
17303 let act = if csr {
17304 e.moe_gate_up_gelu8_dev_q8_csr(
17305 &dev.ptr_row,
17306 &sel_d,
17307 &zq,
17308 &zd,
17309 t * n_used,
17310 n_embd,
17311 n_ff_exp,
17312 n_used,
17313 n_expert,
17314 m.gate_exps.qtype,
17315 m.up_exps.qtype,
17316 m.gate_exps.row_bytes,
17317 m.up_exps.row_bytes,
17318 )?
17319 } else {
17320 e.moe_gate_up_gelu8_dev_q8_rows(
17321 &dev.ptr_row,
17322 &sel_d,
17323 &zq,
17324 &zd,
17325 t,
17326 n_embd,
17327 n_ff_exp,
17328 n_used,
17329 n_expert,
17330 m.gate_exps.qtype,
17331 m.up_exps.qtype,
17332 m.gate_exps.row_bytes,
17333 m.up_exps.row_bytes,
17334 )?
17335 };
17336 let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
17337 let mut moe_out = e.uninit(t * n_embd)?;
17338 e.moe_down8_fma_dev_q8_rows_g(
17339 &dev.ptr_row,
17340 &sel_d,
17341 &w_d,
17342 &aq2,
17343 &ad2,
17344 &mut moe_out,
17345 t,
17346 n_ff_exp,
17347 n_embd,
17348 n_used,
17349 n_expert,
17350 m.down_exps.qtype,
17351 m.down_exps.row_bytes,
17352 )?;
17353 return Ok(moe_out);
17354 }
17355
17356 let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
17357 for (i, &sx) in sel_all.iter().enumerate() {
17358 w_all[i] *= bits.per_expert_scale[sx as usize];
17359 }
17360
17361 if t >= PRIME_MIN_T
17365 && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17366 && expert_dp4a_supported(m.gate_exps.qtype)
17367 && expert_dp4a_supported(m.up_exps.qtype)
17368 && expert_dp4a_supported(m.down_exps.qtype)
17369 && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
17370 {
17371 let dev = m.dev_exps.as_ref().unwrap();
17372 let n_pairs = t * n_used;
17373 let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
17374 let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
17375 let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
17376 let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
17377 let pt = e.htod_i32(&pair_tok)?;
17378 let pw = e.htod(&w_all)?;
17379 let toff = e.htod_i32(&tok_off)?;
17380 let tids = e.htod_i32(&tok_ids)?;
17381 let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
17382 for p in 0..n_pairs {
17383 by_ex[pair_ex[p] as usize].push(p as i32);
17384 }
17385 let mut ex_ids: Vec<i32> = Vec::new();
17386 let mut ex_off: Vec<i32> = vec![0];
17387 let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
17388 for (ex, list) in by_ex.iter().enumerate() {
17389 if list.is_empty() {
17390 continue;
17391 }
17392 ex_ids.push(ex as i32);
17393 ex_pairs.extend_from_slice(list);
17394 ex_off.push(ex_pairs.len() as i32);
17395 }
17396 let n_active = ex_ids.len();
17397 let exi = e.htod_i32(&ex_ids)?;
17398 let exo = e.htod_i32(&ex_off)?;
17399 let exp_d = e.htod_i32(&ex_pairs)?;
17400 if crate::moe_f16g_gemma_on()
17408 && f16g_proj_ok(m.gate_exps.qtype, n_embd)
17409 && f16g_proj_ok(m.up_exps.qtype, n_embd)
17410 && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
17411 {
17412 let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
17413 let csr_tok_d = e.htod_i32(&csr_tok)?;
17414 let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
17415 let g_csr = e.moe_f16_grouped(
17416 &dev.ptr_row,
17417 0,
17418 n_expert,
17419 &exi,
17420 &ex_off,
17421 &exo,
17422 &z_f16,
17423 &z_s,
17424 n_embd,
17425 n_ff_exp,
17426 n_active,
17427 n_pairs,
17428 m.gate_exps.qtype,
17429 m.gate_exps.row_bytes,
17430 )?;
17431 let u_csr = e.moe_f16_grouped(
17432 &dev.ptr_row,
17433 1,
17434 n_expert,
17435 &exi,
17436 &ex_off,
17437 &exo,
17438 &z_f16,
17439 &z_s,
17440 n_embd,
17441 n_ff_exp,
17442 n_active,
17443 n_pairs,
17444 m.up_exps.qtype,
17445 m.up_exps.row_bytes,
17446 )?;
17447 let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
17448 let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
17449 let d_csr = e.moe_f16_grouped(
17450 &dev.ptr_row,
17451 2,
17452 n_expert,
17453 &exi,
17454 &ex_off,
17455 &exo,
17456 &a_f16,
17457 &a_s,
17458 n_ff_exp,
17459 n_embd,
17460 n_active,
17461 n_pairs,
17462 m.down_exps.qtype,
17463 m.down_exps.row_bytes,
17464 )?;
17465 let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
17466 let mut moe_out = e.uninit(t * n_embd)?;
17467 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
17468 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
17469 let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
17470 let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
17471 eprintln!(
17472 "[f16g-debug] post-permute bad={} post-scatter bad={}",
17473 scan(&yd),
17474 scan(&mo)
17475 );
17476 }
17477 return Ok(moe_out);
17478 }
17479 let mma = n_embd.is_multiple_of(256)
17482 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
17483 let (gate, up) = if mma {
17484 let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
17485 (
17486 e.mmq_iq_experts(
17487 &dev.ptr_row,
17488 0,
17489 n_expert,
17490 &exi,
17491 &exo,
17492 &exp_d,
17493 &pt,
17494 &z_scr,
17495 n_embd,
17496 n_ff_exp,
17497 n_active,
17498 n_pairs,
17499 t,
17500 m.gate_exps.qtype,
17501 m.gate_exps.row_bytes,
17502 )?,
17503 e.mmq_iq_experts(
17504 &dev.ptr_row,
17505 1,
17506 n_expert,
17507 &exi,
17508 &exo,
17509 &exp_d,
17510 &pt,
17511 &z_scr,
17512 n_embd,
17513 n_ff_exp,
17514 n_active,
17515 n_pairs,
17516 t,
17517 m.up_exps.qtype,
17518 m.up_exps.row_bytes,
17519 )?,
17520 )
17521 } else {
17522 let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
17523 (
17524 e.moe_pairs_matvec_q8_dec(
17525 &dev.ptr_row,
17526 0,
17527 &exi,
17528 &exo,
17529 &exp_d,
17530 &pt,
17531 &zq,
17532 &zd,
17533 n_embd,
17534 n_ff_exp,
17535 n_expert,
17536 n_active,
17537 n_pairs,
17538 m.gate_exps.qtype,
17539 m.gate_exps.row_bytes,
17540 )?,
17541 e.moe_pairs_matvec_q8_dec(
17542 &dev.ptr_row,
17543 1,
17544 &exi,
17545 &exo,
17546 &exp_d,
17547 &pt,
17548 &zq,
17549 &zd,
17550 n_embd,
17551 n_ff_exp,
17552 n_expert,
17553 n_active,
17554 n_pairs,
17555 m.up_exps.qtype,
17556 m.up_exps.row_bytes,
17557 )?,
17558 )
17559 };
17560 let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
17561 let pself = e.htod_i32(&pair_self)?;
17562 let y_down = if mma {
17574 let in_pad = n_ff_exp.div_ceil(256) * 256;
17575 let a_scr = if crate::moe_fuse_actq_on() {
17576 e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
17577 } else {
17578 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
17579 e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
17580 };
17581 e.mmq_iq_experts(
17582 &dev.ptr_row,
17583 2,
17584 n_expert,
17585 &exi,
17586 &exo,
17587 &exp_d,
17588 &pself,
17589 &a_scr,
17590 in_pad,
17591 n_embd,
17592 n_active,
17593 n_pairs,
17594 n_pairs,
17595 m.down_exps.qtype,
17596 m.down_exps.row_bytes,
17597 )?
17598 } else {
17599 let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
17600 let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
17601 e.moe_pairs_matvec_q8_dec(
17602 &dev.ptr_row,
17603 2,
17604 &exi,
17605 &exo,
17606 &exp_d,
17607 &pself,
17608 &aq2,
17609 &ad2,
17610 n_ff_exp,
17611 n_embd,
17612 n_expert,
17613 n_active,
17614 n_pairs,
17615 m.down_exps.qtype,
17616 m.down_exps.row_bytes,
17617 )?
17618 };
17619 let mut moe_out = e.uninit(t * n_embd)?;
17620 e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
17621 return Ok(moe_out);
17622 }
17623
17624 let g_len = m.gate_exps.expert_stride;
17625 let u_len = m.up_exps.expert_stride;
17626 let d_len = m.down_exps.expert_stride;
17627 let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
17631 let (mut sg, mut su, mut sd) = if dev.is_some() {
17632 (None, None, None)
17633 } else {
17634 (
17635 Some(e.alloc_u8_uninit(g_len)?),
17636 Some(e.alloc_u8_uninit(u_len)?),
17637 Some(e.alloc_u8_uninit(d_len)?),
17638 )
17639 };
17640 let mut moe_out = e.zeros(t * n_embd)?;
17641 for tok in 0..t {
17642 let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
17643 let w = &w_all[tok * n_used..(tok + 1) * n_used];
17644 let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
17645 for (j, &ex) in sel.iter().enumerate() {
17646 let ex = ex as usize;
17647 let gate = match dev {
17648 Some(d) => m.qmatvec_view(
17649 e,
17650 &d.gate,
17651 ex * g_len..(ex + 1) * g_len,
17652 &zt,
17653 1,
17654 m.gate_exps.in_f,
17655 m.gate_exps.out_f,
17656 m.gate_exps.qtype,
17657 m.gate_exps.row_bytes,
17658 )?,
17659 None => {
17660 let sg = sg.as_mut().unwrap();
17661 e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
17662 m.qmatvec_view(
17663 e,
17664 sg,
17665 0..g_len,
17666 &zt,
17667 1,
17668 m.gate_exps.in_f,
17669 m.gate_exps.out_f,
17670 m.gate_exps.qtype,
17671 m.gate_exps.row_bytes,
17672 )?
17673 }
17674 };
17675 let up = match dev {
17676 Some(d) => m.qmatvec_view(
17677 e,
17678 &d.up,
17679 ex * u_len..(ex + 1) * u_len,
17680 &zt,
17681 1,
17682 m.up_exps.in_f,
17683 m.up_exps.out_f,
17684 m.up_exps.qtype,
17685 m.up_exps.row_bytes,
17686 )?,
17687 None => {
17688 let su = su.as_mut().unwrap();
17689 e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
17690 m.qmatvec_view(
17691 e,
17692 su,
17693 0..u_len,
17694 &zt,
17695 1,
17696 m.up_exps.in_f,
17697 m.up_exps.out_f,
17698 m.up_exps.qtype,
17699 m.up_exps.row_bytes,
17700 )?
17701 }
17702 };
17703 let mut act = e.uninit(n_ff_exp)?;
17704 e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
17705 let actv = act.slice(0..n_ff_exp);
17706 let y = match dev {
17707 Some(d) => m.qmatvec_view(
17708 e,
17709 &d.down,
17710 ex * d_len..(ex + 1) * d_len,
17711 &actv,
17712 1,
17713 m.down_exps.in_f,
17714 m.down_exps.out_f,
17715 m.down_exps.qtype,
17716 m.down_exps.row_bytes,
17717 )?,
17718 None => {
17719 let sd = sd.as_mut().unwrap();
17720 e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
17721 m.qmatvec_view(
17722 e,
17723 sd,
17724 0..d_len,
17725 &actv,
17726 1,
17727 m.down_exps.in_f,
17728 m.down_exps.out_f,
17729 m.down_exps.qtype,
17730 m.down_exps.row_bytes,
17731 )?
17732 }
17733 };
17734 let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
17735 e.axpy_into(&y, w[j], &mut dst, n_embd)?;
17736 }
17737 }
17738 Ok(moe_out)
17739 }
17740
17741 fn gemma4_layer(
17743 &self,
17744 e: &Engine,
17745 il: usize,
17746 layer: &crate::hybrid::HybridLayer,
17747 x: &CudaSlice<f32>,
17748 pos_d: &CudaSlice<i32>,
17749 t: usize,
17750 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17751 let n_embd = self.cfg.n_embd as usize;
17752 let eps = self.cfg.rms_eps;
17753
17754 let mut h = e.zeros(t * n_embd)?;
17755 e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
17756 let Mixer::Full(fa) = &layer.mixer else {
17757 panic!("gemma4 layer {il} not full-attn")
17758 };
17759 let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
17760 let mut cur = e.zeros(t * n_embd)?;
17762 e.rms_norm(
17763 &o,
17764 layer.post_attn_norm.float_data(),
17765 &mut cur,
17766 n_embd,
17767 t,
17768 eps,
17769 )?;
17770 self.gemma4_layer_tail_add(e, layer, &cur, x, t)
17771 }
17772
17773 fn gemma4_layer_tail_add(
17777 &self,
17778 e: &Engine,
17779 layer: &crate::hybrid::HybridLayer,
17780 cur: &CudaSlice<f32>,
17781 x: &CudaSlice<f32>,
17782 t: usize,
17783 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17784 Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
17785 }
17786
17787 #[allow(clippy::type_complexity)] fn gemma4_layer_tail_add_n(
17791 &self,
17792 e: &Engine,
17793 layer: &crate::hybrid::HybridLayer,
17794 cur: &CudaSlice<f32>,
17795 x: &CudaSlice<f32>,
17796 t: usize,
17797 next_norm: Option<&CudaSlice<f32>>,
17798 ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
17799 let n_embd = self.cfg.n_embd as usize;
17800 let bits = layer.gemma4.as_ref().unwrap();
17801 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
17802 let mut xn = e.uninit(t * n_embd)?;
17803 match next_norm {
17804 Some(w) => {
17805 let mut hn = e.uninit(t * n_embd)?;
17806 e.add_scale_rms_norm(
17807 &sn,
17808 &attn_out,
17809 bits.layer_scale,
17810 w,
17811 &mut xn,
17812 &mut hn,
17813 n_embd,
17814 t,
17815 self.cfg.rms_eps,
17816 )?;
17817 Ok((xn, Some(hn)))
17818 }
17819 None => {
17820 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
17821 Ok((xn, None))
17822 }
17823 }
17824 }
17825
17826 fn gemma4_layer_tail_core(
17829 &self,
17830 e: &Engine,
17831 layer: &crate::hybrid::HybridLayer,
17832 cur: &CudaSlice<f32>,
17833 x: &CudaSlice<f32>,
17834 t: usize,
17835 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17836 self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
17837 }
17838
17839 #[allow(clippy::too_many_arguments)] fn gemma4_layer_tail_core_pn(
17847 &self,
17848 e: &Engine,
17849 layer: &crate::hybrid::HybridLayer,
17850 cur: &CudaSlice<f32>,
17851 x: &CudaSlice<f32>,
17852 t: usize,
17853 pre_norm: Option<&CudaSlice<f32>>,
17854 defer_post_norm: bool,
17855 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17856 let n_embd = self.cfg.n_embd as usize;
17857 let eps = self.cfg.rms_eps;
17858 let bits = layer.gemma4.as_ref().unwrap();
17859
17860 let Some(mbits) = bits.moe_bits.as_ref() else {
17863 let crate::hybrid::Ffn::Dense {
17864 ffn_gate,
17865 ffn_up,
17866 ffn_down,
17867 } = &layer.ffn
17868 else {
17869 panic!("gemma4 dense layer without Dense ffn")
17870 };
17871 let mut attn_out = e.uninit(t * n_embd)?;
17872 let mut zsh = e.uninit(t * n_embd)?;
17873 let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
17876 match pre_norm {
17877 Some(wa) if t == 1 => {
17878 zpair = Some(e.rms_pre_add_rms_norm_q8z(
17879 cur,
17880 wa,
17881 x,
17882 bits.ffn_norm.float_data(),
17883 &mut attn_out,
17884 &mut zsh,
17885 n_embd,
17886 t,
17887 eps,
17888 )?);
17889 }
17890 Some(wa) => e.rms_pre_add_rms_norm(
17891 cur,
17892 wa,
17893 x,
17894 bits.ffn_norm.float_data(),
17895 &mut attn_out,
17896 &mut zsh,
17897 n_embd,
17898 t,
17899 eps,
17900 )?,
17901 None => e.add_rms_norm(
17902 cur,
17903 x,
17904 bits.ffn_norm.float_data(),
17905 &mut attn_out,
17906 &mut zsh,
17907 n_embd,
17908 t,
17909 eps,
17910 )?,
17911 }
17912 let n_ff = ffn_gate.out_features();
17913 let (gate, up) = if t == 1 {
17919 let (zq, zd) = match zpair {
17920 Some(p) => p,
17921 None => e.quantize_q8_1(&zsh, 1, n_embd)?,
17922 };
17923 match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
17924 Some(p) => p,
17925 None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
17927 Some(p) => p,
17928 None => (
17929 e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
17930 e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
17931 ),
17932 },
17933 }
17934 } else {
17935 static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17940 let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
17941 let fused = if f2b {
17942 let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
17943 e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
17944 } else {
17945 None
17946 };
17947 match fused {
17948 Some(p) => p,
17949 None => {
17950 e.mmq_act_begin();
17952 (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
17953 }
17954 }
17955 };
17956 let mut act = e.uninit(t * n_ff)?;
17957 let f0 = if e.uses_q8_1_fast(ffn_down) {
17960 let upv = e.view(&up, t * n_ff);
17961 let up_all = upv.slice(0..t * n_ff);
17962 let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
17963 e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
17964 } else {
17965 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
17966 e.matmul(ffn_down, &act, t)?
17967 };
17968 if defer_post_norm {
17969 return Ok((f0, attn_out));
17970 }
17971 let mut sn = e.uninit(t * n_embd)?;
17972 e.rms_norm(
17973 &f0,
17974 bits.post_ffw_norm.float_data(),
17975 &mut sn,
17976 n_embd,
17977 t,
17978 eps,
17979 )?;
17980 return Ok((sn, attn_out));
17981 };
17982
17983 assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
17984 let mut attn_out = e.uninit(t * n_embd)?;
17989 let mut router_in = e.uninit(t * n_embd)?;
17990 let fast_moe = match &layer.ffn {
17991 crate::hybrid::Ffn::Moe(m) => {
17992 m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17993 && expert_dp4a_supported(m.gate_exps.qtype)
17994 && expert_dp4a_supported(m.up_exps.qtype)
17995 && expert_dp4a_supported(m.down_exps.qtype)
17996 && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
17997 }
17998 _ => false,
17999 };
18000 let q8z = t < PRIME_MIN_T && fast_moe;
18001 let (zsh_f32, zsh_q8, moe_q8) = if q8z {
18002 let (z0, m2) = e.add_rms_norm3_q8z(
18003 cur,
18004 x,
18005 bits.ffn_norm.float_data(),
18006 &mbits.router_scale_pre,
18007 mbits.pre_ffw_norm_2.float_data(),
18008 &mut attn_out,
18009 &mut router_in,
18010 n_embd,
18011 t,
18012 eps,
18013 )?;
18014 (None, Some(z0), Some(m2))
18015 } else {
18016 let mut zsh = e.uninit(t * n_embd)?;
18017 let mut moe_in = e.uninit(t * n_embd)?;
18018 e.add_rms_norm3(
18019 cur,
18020 x,
18021 bits.ffn_norm.float_data(),
18022 &mbits.router_scale_pre,
18023 mbits.pre_ffw_norm_2.float_data(),
18024 &mut attn_out,
18025 &mut zsh,
18026 &mut router_in,
18027 &mut moe_in,
18028 n_embd,
18029 t,
18030 eps,
18031 )?;
18032 (Some((zsh, moe_in)), None, None)
18033 };
18034 let attn_out2 = attn_out;
18035 #[allow(unused_variables)]
18036 let attn_out = &attn_out2;
18037 let n_ff = mbits.shared_gate.out_features();
18038 let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
18039 if t == 1 {
18040 match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
18041 Some(p) => p,
18042 None => match e.matmul_nvfp4_fused2(
18043 &mbits.shared_gate,
18044 &mbits.shared_up,
18045 zq,
18046 zd,
18047 1,
18048 )? {
18049 Some(p) => p,
18050 None => {
18051 let h0 = e.zeros(0)?;
18052 (
18053 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
18054 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
18055 )
18056 }
18057 },
18058 }
18059 } else {
18060 let h0 = e.zeros(0)?;
18062 (
18063 e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
18064 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
18065 )
18066 }
18067 } else {
18068 let (zsh, _) = zsh_f32.as_ref().unwrap();
18069 (
18070 e.matmul(&mbits.shared_gate, zsh, t)?,
18071 e.matmul(&mbits.shared_up, zsh, t)?,
18072 )
18073 };
18074 let mut act = e.uninit(t * n_ff)?;
18075 e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
18076 let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
18077 let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
18078 panic!("gemma4 layer not MoE")
18079 };
18080 let moe0 = match (&moe_q8, &zsh_f32) {
18081 (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
18082 (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
18083 _ => unreachable!(),
18084 };
18085 let mut mlp = e.uninit(t * n_embd)?;
18087 let mut moe = e.uninit(t * n_embd)?;
18088 e.rms_norm2x(
18089 &mlp0,
18090 &moe0,
18091 mbits.post_ffw_norm_1.float_data(),
18092 mbits.post_ffw_norm_2.float_data(),
18093 &mut mlp,
18094 &mut moe,
18095 n_embd,
18096 t,
18097 eps,
18098 )?;
18099
18100 let mut sum = e.uninit(t * n_embd)?;
18103 let mut sn = e.uninit(t * n_embd)?;
18104 e.add_rms_norm(
18105 &mlp,
18106 &moe,
18107 bits.post_ffw_norm.float_data(),
18108 &mut sum,
18109 &mut sn,
18110 n_embd,
18111 t,
18112 eps,
18113 )?;
18114 Ok((sn, attn_out2))
18115 }
18116
18117 #[allow(clippy::type_complexity)] pub(crate) fn gemma4_layer_tail_add_nq_pn(
18128 &self,
18129 e: &Engine,
18130 layer: &crate::hybrid::HybridLayer,
18131 o: &CudaSlice<f32>,
18132 x: &CudaSlice<f32>,
18133 t: usize,
18134 next_norm: Option<&CudaSlice<f32>>,
18135 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
18136 {
18137 let n_embd = self.cfg.n_embd as usize;
18138 let eps = self.cfg.rms_eps;
18139 let bits = layer.gemma4.as_ref().unwrap();
18140 if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
18141 let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
18142 e,
18143 layer,
18144 o,
18145 x,
18146 t,
18147 Some(layer.post_attn_norm.float_data()),
18148 true,
18149 )?;
18150 let mut xn = e.uninit(t * n_embd)?;
18151 return match next_norm {
18152 Some(w) => {
18153 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
18154 &f0,
18155 bits.post_ffw_norm.float_data(),
18156 &attn_out,
18157 bits.layer_scale,
18158 w,
18159 &mut xn,
18160 n_embd,
18161 t,
18162 eps,
18163 )?;
18164 Ok((xn, Some(pair)))
18165 }
18166 None => {
18167 let mut sn = e.uninit(t * n_embd)?;
18168 e.rms_norm(
18169 &f0,
18170 bits.post_ffw_norm.float_data(),
18171 &mut sn,
18172 n_embd,
18173 t,
18174 eps,
18175 )?;
18176 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
18177 Ok((xn, None))
18178 }
18179 };
18180 }
18181 let mut cur = e.uninit(t * n_embd)?;
18182 e.rms_norm(
18183 o,
18184 layer.post_attn_norm.float_data(),
18185 &mut cur,
18186 n_embd,
18187 t,
18188 eps,
18189 )?;
18190 self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
18191 }
18192
18193 #[allow(clippy::type_complexity)] pub(crate) fn gemma4_layer_tail_add_nq(
18195 &self,
18196 e: &Engine,
18197 layer: &crate::hybrid::HybridLayer,
18198 cur: &CudaSlice<f32>,
18199 x: &CudaSlice<f32>,
18200 t: usize,
18201 next_norm: Option<&CudaSlice<f32>>,
18202 ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
18203 {
18204 let n_embd = self.cfg.n_embd as usize;
18205 let bits = layer.gemma4.as_ref().unwrap();
18206 let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
18207 let mut xn = e.uninit(t * n_embd)?;
18208 match next_norm {
18209 Some(w) => {
18210 let pair = e.add_scale_rms_norm_q8_1(
18211 &sn,
18212 &attn_out,
18213 bits.layer_scale,
18214 w,
18215 &mut xn,
18216 n_embd,
18217 t,
18218 self.cfg.rms_eps,
18219 )?;
18220 Ok((xn, Some(pair)))
18221 }
18222 None => {
18223 e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
18224 Ok((xn, None))
18225 }
18226 }
18227 }
18228
18229 fn gemma4_forward(
18232 &self,
18233 e: &Engine,
18234 tokens: &[u32],
18235 last_only: bool,
18236 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
18237 if self.is_gemma4_e4b() {
18240 return self.gemma4_e4b_forward(e, tokens, last_only);
18241 }
18242 let n_embd = self.cfg.n_embd as usize;
18243 let t = tokens.len();
18244 let pos: Vec<i32> = (0..t as i32).collect();
18245 let pos_d = e.htod_i32(&pos)?;
18246
18247 let mut x = self.embed(e, tokens)?;
18248 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
18249 let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
18252 let stat =
18253 |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
18254 let h = e.dtoh(x)?;
18255 let bad = h.iter().filter(|v| !v.is_finite()).count();
18256 let mx = h
18257 .iter()
18258 .filter(|v| v.is_finite())
18259 .fold(0.0f32, |m, v| m.max(v.abs()));
18260 eprintln!(
18261 "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
18262 &h[..3]
18263 );
18264 Ok(())
18265 };
18266 if probe {
18267 stat(e, &x, "embed")?;
18268 }
18269 for (il, layer) in self.layers.iter().enumerate() {
18270 x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
18271 if probe {
18272 stat(e, &x, &format!("L{il}"))?;
18273 }
18274 }
18275 let mut hn = e.zeros(t * n_embd)?;
18276 e.rms_norm(
18277 &x,
18278 self.output_norm.float_data(),
18279 &mut hn,
18280 n_embd,
18281 t,
18282 self.cfg.rms_eps,
18283 )?;
18284 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
18285 let n_vocab = self.output.out_features();
18286 let logits = if last_only {
18287 let hv = e.view(&hn, t * n_embd);
18288 let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
18289 let mut hlast = e.zeros(n_embd)?;
18290 e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
18291 let mut ld = e.matmul(&self.output, &hlast, 1)?;
18292 e.softcap(&mut ld, cap, n_vocab)?;
18293 self.gemma4_suppress(e, &mut ld, 1)?;
18294 e.dtoh(&ld)?
18295 } else {
18296 let mut ld = e.matmul(&self.output, &hn, t)?;
18297 e.softcap(&mut ld, cap, t * n_vocab)?;
18298 self.gemma4_suppress(e, &mut ld, t)?;
18299 e.dtoh(&ld)?
18300 };
18301 Ok(logits)
18302 }
18303
18304 #[allow(clippy::type_complexity)] pub(crate) fn gemma4_prime(
18310 &self,
18311 e: &Engine,
18312 tokens: &[u32],
18313 cache: &mut Cache,
18314 overlay: Option<&crate::vision::EmbedOverlay>,
18315 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18316 if cache.pos != 0 {
18321 return Err(
18322 "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
18323 — prime the full prompt in one call or decode tokenwise"
18324 .into(),
18325 );
18326 }
18327 let n_embd = self.cfg.n_embd as usize;
18328 let eps = self.cfg.rms_eps;
18329 let t = tokens.len();
18330 let pos: Vec<i32> = (0..t as i32).collect();
18331 let pos_d = e.htod_i32(&pos)?;
18332 let mut x = self.embed(e, tokens)?;
18333 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
18334 let island: Option<CudaSlice<i32>> = match overlay {
18341 Some(ov) => {
18342 let mut span_id = vec![-1i32; t];
18343 for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
18344 if pos + n_rows > t {
18345 return Err(format!(
18346 "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
18347 pos + n_rows
18348 )
18349 .into());
18350 }
18351 let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
18352 e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
18353 for s in span_id.iter_mut().skip(pos).take(n_rows) {
18354 *s = i as i32;
18355 }
18356 }
18357 if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
18361 None
18362 } else {
18363 Some(e.htod_i32(&span_id)?)
18364 }
18365 }
18366 None => None,
18367 };
18368 for (il, layer) in self.layers.iter().enumerate() {
18369 let mut h = e.zeros(t * n_embd)?;
18370 e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
18371 let Mixer::Full(fa) = &layer.mixer else {
18372 panic!("gemma4 layer not full-attn")
18373 };
18374 let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
18375 if trace {
18376 let v = e.dtoh(&h)?;
18377 let nan = v.iter().filter(|x| x.is_nan()).count();
18378 eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
18379 }
18380 let o =
18381 self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
18382 if trace {
18383 let v = e.dtoh(&o)?;
18384 let nan = v.iter().filter(|x| x.is_nan()).count();
18385 eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
18386 }
18387 let mut cur = e.zeros(t * n_embd)?;
18388 e.rms_norm(
18389 &o,
18390 layer.post_attn_norm.float_data(),
18391 &mut cur,
18392 n_embd,
18393 t,
18394 eps,
18395 )?;
18396 x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
18397 self.dflash_tap(e, cache, il, &x, t)?;
18398 if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
18400 let h = e.dtoh(&x)?;
18401 let nan = h.iter().filter(|v| v.is_nan()).count();
18402 let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
18403 eprintln!(
18404 "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
18405 h.len()
18406 );
18407 if nan > 0 {
18408 return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
18409 }
18410 }
18411 }
18412 cache.pos += t;
18413 let hiddens = e.clone_dtod(&x)?;
18414 let xv = e.view(&x, t * n_embd);
18415 let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
18416 let mut h_seed = e.zeros(n_embd)?;
18417 e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
18418 let mut hn = e.uninit(n_embd)?;
18419 e.rms_norm(
18420 &h_seed,
18421 self.output_norm.float_data(),
18422 &mut hn,
18423 n_embd,
18424 1,
18425 eps,
18426 )?;
18427 let mut ld = e.matmul(&self.output, &hn, 1)?;
18428 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
18429 e.softcap(&mut ld, cap, self.output.out_features())?;
18430 self.gemma4_suppress(e, &mut ld, 1)?;
18431 let logits = e.dtoh(&ld)?;
18432 Ok((logits, h_seed, hiddens))
18433 }
18434
18435 #[allow(clippy::too_many_arguments)] fn gemma4_decode_attn(
18441 &self,
18442 e: &Engine,
18443 fa: &crate::hybrid::FullAttnLayer,
18444 il: usize,
18445 hq: &CudaSlice<i8>,
18446 hdq: &CudaSlice<f32>,
18447 pos_d: &CudaSlice<i32>,
18448 cache: &mut Cache,
18449 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18450 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
18451 let eps = self.cfg.rms_eps;
18452 let aux = self.gemma4_aux.as_ref().unwrap();
18453 let ones = aux.ones(e);
18454 #[cfg(debug_assertions)]
18455 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
18456 let (hq, hdq) = (hq, hdq);
18457 let h0 = e.zeros(0)?;
18458 let h = &h0;
18459 let (q0, k0, v0) = if swa {
18460 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
18461 Some(t3) => t3,
18462 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
18465 Some((q0, k0)) => {
18466 let v0 = e.matmul_pre(&fa.wv, hq, hdq, h, 1)?;
18467 (q0, k0, v0)
18468 }
18469 None => (
18470 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
18471 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
18472 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
18473 ),
18474 },
18475 }
18476 } else {
18477 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
18478 Some(p) => p,
18479 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
18480 Some(p) => p,
18481 None => (
18482 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
18483 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
18484 ),
18485 },
18486 };
18487 let v0 = e.clone_dtod(&k0)?;
18488 (q0, k0, v0)
18489 };
18490 let mut q = e.uninit(nh * hd)?;
18491 let mut k = e.uninit(nkv * hd)?;
18492 let mut v = e.uninit(nkv * hd)?;
18493 let ff = if swa {
18496 None
18497 } else {
18498 Some(
18499 aux.rope_freqs(e)
18500 .expect("gemma4 global rope needs rope_freqs.weight"),
18501 )
18502 };
18503 #[cfg(debug_assertions)]
18504 if let Some(ff) = ff {
18505 crate::debug_assert_tensor_stream_device(
18506 ff,
18507 &e.stream(),
18508 "gemma4_decode_attn.rope_freqs",
18509 );
18510 }
18511 let kvl = cache.kv[il].as_mut().unwrap();
18512 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
18513 if crate::Engine::qkv_append_on() {
18514 e.rms_norm_qkv_rope_append(
18518 &q0,
18519 &k0,
18520 &v0,
18521 fa.q_norm.float_data(),
18522 fa.k_norm.float_data(),
18523 ones,
18524 &mut q,
18525 &mut k,
18526 &mut v,
18527 hd,
18528 self.gemma4_rope_dims(il),
18529 nh,
18530 nkv,
18531 pos_d,
18532 nh,
18533 nkv,
18534 base,
18535 1.0,
18536 ff,
18537 eps,
18538 &mut kvl.k,
18539 &mut kvl.v,
18540 kvl.len,
18541 kvl.k_tok_bytes,
18542 kvl.v_tok_bytes,
18543 kv_fp8,
18544 )?;
18545 } else {
18546 e.rms_norm_qkv_rope(
18547 &q0,
18548 &k0,
18549 &v0,
18550 fa.q_norm.float_data(),
18551 fa.k_norm.float_data(),
18552 ones,
18553 &mut q,
18554 &mut k,
18555 &mut v,
18556 hd,
18557 self.gemma4_rope_dims(il),
18558 nh,
18559 nkv,
18560 pos_d,
18561 nh,
18562 nkv,
18563 base,
18564 1.0,
18565 ff,
18566 eps,
18567 )?;
18568 e.append_kv_quantized(
18569 &k,
18570 &v,
18571 &mut kvl.k,
18572 &mut kvl.v,
18573 kvl.len,
18574 kvl.kv_dim_k,
18575 kvl.kv_dim_v,
18576 kvl.k_tok_bytes,
18577 kvl.v_tok_bytes,
18578 kv_fp8,
18579 )?;
18580 }
18581 kvl.len += 1;
18582 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
18586 let mut attn = e.uninit(nh * hd)?;
18587 if !swa
18589 && hd == 512
18590 && kvl.len >= crate::fa512_min_tkv()
18591 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
18592 {
18593 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
18594 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
18595 let base = kvl.len as i32;
18597 e.i32_set_k(&mut kvl.len_d, base)?;
18598 e.fa_decode_rows(
18599 &q,
18600 &kp,
18601 &vp,
18602 &mut attn,
18603 hd,
18604 nh,
18605 nkv,
18606 kvl.len - 1,
18607 1,
18608 scale,
18609 kvl.k_tok_bytes,
18610 kvl.v_tok_bytes,
18611 Some((&kvl.len_d, -1)),
18612 false,
18613 false,
18614 None,
18615 )?;
18616 return e.matmul(&fa.wo, &attn, 1);
18617 }
18618 if swa
18620 && kvl.len > win
18621 && hd == 256
18622 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
18623 {
18624 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
18625 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
18626 let base = kvl.len as i32;
18627 e.i32_set_k(&mut kvl.len_d, base)?;
18628 e.fa_decode_rows_w(
18629 &q,
18630 &kp,
18631 &vp,
18632 &mut attn,
18633 hd,
18634 nh,
18635 nkv,
18636 &kvl.len_d,
18637 -1,
18638 1,
18639 scale,
18640 win,
18641 kvl.k_tok_bytes,
18642 kvl.v_tok_bytes,
18643 None,
18644 )?;
18645 return e.matmul(&fa.wo, &attn, 1);
18646 }
18647 let (off_tok, t_kv) = if swa && kvl.len > win {
18648 (kvl.len - win, win)
18649 } else {
18650 (0, kvl.len)
18651 };
18652 let k_view = e.view_u8_range(
18653 &kvl.k,
18654 off_tok * kvl.k_tok_bytes,
18655 (off_tok + t_kv) * kvl.k_tok_bytes,
18656 );
18657 let v_view = e.view_u8_range(
18658 &kvl.v,
18659 off_tok * kvl.v_tok_bytes,
18660 (off_tok + t_kv) * kvl.v_tok_bytes,
18661 );
18662 e.fa_decode_kvmod(
18663 &q,
18664 &k_view,
18665 &v_view,
18666 &mut attn,
18667 hd,
18668 nh,
18669 nkv,
18670 t_kv,
18671 scale,
18672 kvl.k_tok_bytes,
18673 kvl.v_tok_bytes,
18674 swa && crate::Engine::wkv_on(),
18675 )?;
18676 e.matmul(&fa.wo, &attn, 1)
18677 }
18678
18679 #[allow(clippy::too_many_arguments)]
18686 pub fn gemma4_decode_step_dc(
18687 &self,
18688 e: &Engine,
18689 token_d: &CudaSlice<u32>,
18690 pos_d: &mut CudaSlice<i32>,
18691 embd_gpu: &CudaSlice<u8>,
18692 embd_qt: i32,
18693 embd_rb: usize,
18694 cache: &mut Cache,
18695 n_vocab: usize,
18696 cap_bucket_max: Option<(usize, usize)>,
18697 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
18698 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
18699 self.gemma4_decode_step_dc_into(
18700 e,
18701 token_d,
18702 pos_d,
18703 embd_gpu,
18704 embd_qt,
18705 embd_rb,
18706 cache,
18707 n_vocab,
18708 cap_bucket_max,
18709 &mut tok_out,
18710 )?;
18711 Ok(tok_out)
18712 }
18713
18714 #[allow(clippy::too_many_arguments)]
18717 pub fn gemma4_decode_step_dc_into(
18718 &self,
18719 e: &Engine,
18720 token_d: &CudaSlice<u32>,
18721 pos_d: &mut CudaSlice<i32>,
18722 embd_gpu: &CudaSlice<u8>,
18723 embd_qt: i32,
18724 embd_rb: usize,
18725 cache: &mut Cache,
18726 n_vocab: usize,
18727 cap_bucket_max: Option<(usize, usize)>,
18728 tok_out: &mut CudaSlice<u32>,
18729 ) -> Result<(), Box<dyn std::error::Error>> {
18730 let n_embd = self.cfg.n_embd as usize;
18731 let eps = self.cfg.rms_eps;
18732 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
18733 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
18734 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
18735 let n_layers = self.layers.len();
18736 for (il, layer) in self.layers.iter().enumerate() {
18737 let (hq, hdq) = match h_carry.take() {
18738 Some(p) => p,
18739 None => {
18740 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
18741 }
18742 };
18743 let Mixer::Full(fa) = &layer.mixer else {
18744 panic!("gemma4 layer {il} not full-attn")
18745 };
18746 let o =
18747 self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
18748 let next_norm = if il + 1 < n_layers {
18749 Some(self.layers[il + 1].attn_norm.float_data())
18750 } else {
18751 None
18752 };
18753 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
18754 x = xn;
18755 h_carry = hn;
18756 }
18757 let mut hn = e.uninit(n_embd)?;
18758 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
18759 let mut logits = e.matmul(&self.output, &hn, 1)?;
18760 self.gemma4_suppress(e, &mut logits, 1)?; e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
18762 e.inc_seqlen(pos_d)?;
18763 if cap_bucket_max.is_none() {
18764 cache.pos += 1;
18765 }
18766 Ok(())
18767 }
18768
18769 pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
18776 let n_embd = self.cfg.n_embd as usize;
18777 let n_vocab = self.output.out_features();
18778 let n_layers = self.layers.len();
18779 let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
18780 for il in 0..n_layers {
18781 let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
18782 qmax = qmax.max(nh * hd);
18783 kvmax = kvmax.max(nkv * hd);
18784 if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
18785 ffmax = ffmax.max(ffn_gate.out_features());
18786 }
18787 }
18788 Ok(G4DcSlots {
18789 x: e.uninit(n_embd)?,
18790 xn: e.uninit(n_embd)?,
18791 cur: e.uninit(n_embd)?,
18792 hq: e.alloc_i8_uninit(n_embd)?,
18793 hd_: e.uninit(n_embd / 32)?,
18794 q0: e.uninit(qmax)?,
18795 k0: e.uninit(kvmax)?,
18796 v0: e.uninit(kvmax)?,
18797 q: e.uninit(qmax)?,
18798 k: e.uninit(kvmax)?,
18799 v: e.uninit(kvmax)?,
18800 attn: e.uninit(qmax)?,
18801 o: e.uninit(n_embd)?,
18802 attn_out: e.uninit(n_embd)?,
18803 zsh: e.uninit(n_embd)?,
18804 zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
18807 zd: e.uninit(n_embd.max(qmax) / 32)?,
18808 gate: e.uninit(ffmax)?,
18809 up: e.uninit(ffmax)?,
18810 act: e.uninit(ffmax)?,
18811 actq: e.alloc_i8_uninit(ffmax)?,
18812 actd: e.uninit(ffmax / 32)?,
18813 f0: e.uninit(n_embd)?,
18814 sn: e.uninit(n_embd)?,
18815 hn: e.uninit(n_embd)?,
18816 logits: e.uninit(n_vocab)?,
18817 })
18818 }
18819
18820 fn g4_matvec_m1_into(
18823 &self,
18824 e: &Engine,
18825 w: &crate::model::GpuTensor,
18826 aq: &CudaSlice<i8>,
18827 ad: &CudaSlice<f32>,
18828 y: &mut CudaSlice<f32>,
18829 ) -> Result<(), Box<dyn std::error::Error>> {
18830 use crate::model::GpuTensor;
18831 let (bytes, qtype, row_bytes, scale, rp) = match w {
18832 GpuTensor::Quant {
18833 bytes,
18834 qtype,
18835 row_bytes,
18836 scale,
18837 rp,
18838 ..
18839 } => (bytes, *qtype, *row_bytes, *scale, *rp),
18840 _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
18841 };
18842 let (mbytes, mrp) = match w {
18843 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
18844 _ => (bytes, rp),
18845 };
18846 e.qmatvec_mmvq_into(
18847 mbytes,
18848 aq,
18849 ad,
18850 1,
18851 w.in_features(),
18852 w.out_features(),
18853 qtype,
18854 row_bytes,
18855 scale,
18856 mrp,
18857 y,
18858 )
18859 }
18860
18861 #[allow(clippy::too_many_arguments)]
18865 pub fn gemma4_decode_step_dc_slotted(
18866 &self,
18867 e: &Engine,
18868 token_d: &CudaSlice<u32>,
18869 pos_d: &mut CudaSlice<i32>,
18870 embd_gpu: &CudaSlice<u8>,
18871 embd_qt: i32,
18872 embd_rb: usize,
18873 cache: &mut Cache,
18874 n_vocab: usize,
18875 cap_bucket_max: Option<(usize, usize)>,
18876 sl: &mut G4DcSlots,
18877 tok_out: &mut CudaSlice<u32>,
18878 ring: Option<(&mut CudaSlice<u32>, usize)>,
18879 ) -> Result<(), Box<dyn std::error::Error>> {
18880 let n_embd = self.cfg.n_embd as usize;
18881 let eps = self.cfg.rms_eps;
18882 e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
18883 e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
18884 let n_layers = self.layers.len();
18885 let mut has_carry = false;
18886 for il in 0..n_layers {
18887 if !has_carry {
18888 e.rms_norm_q8_1_into(
18889 &sl.x,
18890 self.layers[il].attn_norm.float_data(),
18891 n_embd,
18892 1,
18893 eps,
18894 &mut sl.hq,
18895 &mut sl.hd_,
18896 )?;
18897 }
18898 has_carry = true;
18899 let layer = &self.layers[il];
18900 let Mixer::Full(fa) = &layer.mixer else {
18901 panic!("gemma4 layer {il} not full-attn")
18902 };
18903 self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
18904 if !Engine::g4_pnfold_on() {
18907 e.rms_norm(
18908 &sl.o,
18909 layer.post_attn_norm.float_data(),
18910 &mut sl.cur,
18911 n_embd,
18912 1,
18913 eps,
18914 )?;
18915 }
18916 let next_norm = if il + 1 < n_layers {
18917 Some(self.layers[il + 1].attn_norm.float_data())
18918 } else {
18919 None
18920 };
18921 self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
18922 std::mem::swap(&mut sl.x, &mut sl.xn);
18923 }
18924 e.rms_norm(
18925 &sl.x,
18926 self.output_norm.float_data(),
18927 &mut sl.hn,
18928 n_embd,
18929 1,
18930 eps,
18931 )?;
18932 e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
18933 {
18935 let (zq, zd) = (&sl.zq, &sl.zd);
18936 let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
18937 let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
18938 self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
18939 }
18940 self.gemma4_suppress(e, &mut sl.logits, 1)?;
18941 e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
18942 if let Some((ring, base)) = ring {
18943 e.plain_tok_ring(tok_out, pos_d, base, ring)?;
18947 }
18948 e.inc_seqlen(pos_d)?;
18949 if cap_bucket_max.is_none() {
18950 cache.pos += 1;
18951 }
18952 Ok(())
18953 }
18954
18955 #[allow(clippy::too_many_arguments)]
18957 fn gemma4_decode_attn_dc_slotted(
18958 &self,
18959 e: &Engine,
18960 fa: &crate::hybrid::FullAttnLayer,
18961 il: usize,
18962 pos_d: &CudaSlice<i32>,
18963 cache: &mut Cache,
18964 cap_bucket_max: Option<(usize, usize)>,
18965 sl: &mut G4DcSlots,
18966 ) -> Result<(), Box<dyn std::error::Error>> {
18967 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
18968 let eps = self.cfg.rms_eps;
18969 let aux = self.gemma4_aux.as_ref().unwrap();
18970 let ones = aux.ones(e);
18971 #[cfg(debug_assertions)]
18972 crate::debug_assert_tensor_stream_device(
18973 ones,
18974 &e.stream(),
18975 "gemma4_decode_attn_dc_slotted.ones",
18976 );
18977 {
18978 let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
18979 let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
18980 if swa {
18981 if !e.matmul_q4_fused3_into(
18982 &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
18983 )? {
18984 if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
18988 {
18989 self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
18990 } else {
18991 return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
18992 }
18993 }
18994 } else {
18995 if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
18996 && !e
18997 .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
18998 {
18999 return Err("slotted step: fused2 unavailable".into());
19000 }
19001 let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
19002 e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
19003 }
19004 }
19005 let ff = if swa {
19008 None
19009 } else {
19010 Some(
19011 aux.rope_freqs(e)
19012 .expect("gemma4 global rope needs rope_freqs.weight"),
19013 )
19014 };
19015 #[cfg(debug_assertions)]
19016 if let Some(ff) = ff {
19017 crate::debug_assert_tensor_stream_device(
19018 ff,
19019 &e.stream(),
19020 "gemma4_decode_attn_dc_slotted.rope_freqs",
19021 );
19022 }
19023 let kvl = cache.kv[il].as_mut().unwrap();
19024 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
19025 if crate::Engine::qkv_append_on() {
19026 e.rms_norm_qkv_rope_append_dc(
19028 &sl.q0,
19029 &sl.k0,
19030 &sl.v0,
19031 fa.q_norm.float_data(),
19032 fa.k_norm.float_data(),
19033 ones,
19034 &mut sl.q,
19035 &mut sl.k,
19036 &mut sl.v,
19037 hd,
19038 self.gemma4_rope_dims(il),
19039 nh,
19040 nkv,
19041 pos_d,
19042 nh,
19043 nkv,
19044 base,
19045 1.0,
19046 ff,
19047 eps,
19048 &mut kvl.k,
19049 &mut kvl.v,
19050 &kvl.len_d,
19051 kvl.k_tok_bytes,
19052 kvl.v_tok_bytes,
19053 kv_fp8,
19054 )?;
19055 } else {
19056 e.rms_norm_qkv_rope(
19057 &sl.q0,
19058 &sl.k0,
19059 &sl.v0,
19060 fa.q_norm.float_data(),
19061 fa.k_norm.float_data(),
19062 ones,
19063 &mut sl.q,
19064 &mut sl.k,
19065 &mut sl.v,
19066 hd,
19067 self.gemma4_rope_dims(il),
19068 nh,
19069 nkv,
19070 pos_d,
19071 nh,
19072 nkv,
19073 base,
19074 1.0,
19075 ff,
19076 eps,
19077 )?;
19078 e.append_kv_quantized_dc(
19079 &sl.k,
19080 &sl.v,
19081 &mut kvl.k,
19082 &mut kvl.v,
19083 &kvl.len_d,
19084 kvl.kv_dim_k,
19085 kvl.kv_dim_v,
19086 kvl.k_tok_bytes,
19087 kvl.v_tok_bytes,
19088 kv_fp8,
19089 )?;
19090 }
19091 e.inc_seqlen(&mut kvl.len_d)?;
19092 let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
19093 let k_view = e.view_u8(&kvl.k, kvl.k.len());
19094 let v_view = e.view_u8(&kvl.v, kvl.v.len());
19095 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
19096 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19097 let mut fa_q8 = false;
19101 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
19102 e.fa_decode_rows(
19103 &sl.q,
19104 &k_view,
19105 &v_view,
19106 &mut sl.attn,
19107 hd,
19108 nh,
19109 nkv,
19110 b_glob - 1,
19111 1,
19112 scale,
19113 kvl.k_tok_bytes,
19114 kvl.v_tok_bytes,
19115 Some((&kvl.len_d, -1)),
19116 false,
19117 false,
19118 Some((&mut sl.zq, &mut sl.zd)),
19119 )?;
19120 fa_q8 = true;
19121 } else if swa && b_swa > win && hd == 256 && rows_on {
19122 e.fa_decode_rows_w(
19123 &sl.q,
19124 &k_view,
19125 &v_view,
19126 &mut sl.attn,
19127 hd,
19128 nh,
19129 nkv,
19130 &kvl.len_d,
19131 -1,
19132 1,
19133 scale,
19134 win,
19135 kvl.k_tok_bytes,
19136 kvl.v_tok_bytes,
19137 Some((&mut sl.zq, &mut sl.zd)),
19138 )?;
19139 fa_q8 = true;
19140 } else {
19141 let b = if swa { b_swa } else { b_glob };
19142 e.fa_decode_dc(
19143 &sl.q,
19144 &k_view,
19145 &v_view,
19146 &mut sl.attn,
19147 hd,
19148 nh,
19149 nkv,
19150 &kvl.len_d,
19151 b,
19152 scale,
19153 kvl.k_tok_bytes,
19154 kvl.v_tok_bytes,
19155 swa && crate::Engine::wkv_on(),
19156 )?;
19157 }
19158 if !fa_q8 {
19159 let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
19160 e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
19161 }
19162 {
19163 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
19164 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
19165 self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
19166 }
19167 Ok(())
19168 }
19169
19170 fn gemma4_layer_tail_slotted(
19173 &self,
19174 e: &Engine,
19175 layer: &crate::hybrid::HybridLayer,
19176 next_norm: Option<&CudaSlice<f32>>,
19177 sl: &mut G4DcSlots,
19178 ) -> Result<(), Box<dyn std::error::Error>> {
19179 let n_embd = self.cfg.n_embd as usize;
19180 let eps = self.cfg.rms_eps;
19181 let bits = layer.gemma4.as_ref().unwrap();
19182 let crate::hybrid::Ffn::Dense {
19183 ffn_gate,
19184 ffn_up,
19185 ffn_down,
19186 } = &layer.ffn
19187 else {
19188 return Err("slotted tail: dense ffn only".into());
19189 };
19190 let pnfold = Engine::g4_pnfold_on();
19191 if pnfold {
19192 let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
19195 let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
19196 e.rms_pre_add_rms_norm_q8z_into(
19197 or,
19198 layer.post_attn_norm.float_data(),
19199 xr,
19200 bits.ffn_norm.float_data(),
19201 &mut sl.attn_out,
19202 &mut sl.zsh,
19203 n_embd,
19204 1,
19205 eps,
19206 &mut sl.zq,
19207 &mut sl.zd,
19208 )?;
19209 } else {
19210 e.add_rms_norm(
19211 &sl.cur,
19212 &sl.x,
19213 bits.ffn_norm.float_data(),
19214 &mut sl.attn_out,
19215 &mut sl.zsh,
19216 n_embd,
19217 1,
19218 eps,
19219 )?;
19220 }
19221 let n_ff = ffn_gate.out_features();
19222 if !pnfold {
19223 let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
19224 e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
19225 }
19226 {
19227 let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
19228 let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
19229 if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
19230 && !e.matmul_nvfp4_fused2_into(
19231 ffn_gate,
19232 ffn_up,
19233 zq,
19234 zd,
19235 &mut sl.gate,
19236 &mut sl.up,
19237 )?
19238 {
19239 return Err("slotted tail: ffn fused2 unavailable".into());
19240 }
19241 }
19242 debug_assert!(e.uses_q8_1_fast(ffn_down));
19243 {
19244 let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
19245 let upv = e.view(upr, n_ff);
19246 let up_all = upv.slice(0..n_ff);
19247 let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
19248 e.gelu_tanh_mul_q8_1_into(
19249 gr,
19250 &up_all,
19251 &mut sl.act,
19252 n_ff,
19253 1,
19254 &mut sl.actq,
19255 &mut sl.actd,
19256 )?;
19257 }
19258 {
19259 let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
19260 let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
19261 self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
19262 }
19263 if pnfold {
19264 if let Some(w) = next_norm {
19267 let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
19268 let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
19269 e.rms_pre_add_scale_rms_norm_q8_1_into(
19270 f0r,
19271 bits.post_ffw_norm.float_data(),
19272 aor,
19273 bits.layer_scale,
19274 w,
19275 &mut sl.xn,
19276 n_embd,
19277 1,
19278 eps,
19279 &mut sl.hq,
19280 &mut sl.hd_,
19281 )?;
19282 return Ok(());
19283 }
19284 }
19285 e.rms_norm(
19286 &sl.f0,
19287 bits.post_ffw_norm.float_data(),
19288 &mut sl.sn,
19289 n_embd,
19290 1,
19291 eps,
19292 )?;
19293 match next_norm {
19294 Some(w) => {
19295 e.add_scale_rms_norm_q8_1_into(
19296 &sl.sn,
19297 &sl.attn_out,
19298 bits.layer_scale,
19299 w,
19300 &mut sl.xn,
19301 n_embd,
19302 1,
19303 eps,
19304 &mut sl.hq,
19305 &mut sl.hd_,
19306 )?;
19307 }
19308 None => {
19309 e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
19310 }
19311 }
19312 Ok(())
19313 }
19314
19315 #[allow(clippy::too_many_arguments)]
19317 fn gemma4_decode_attn_dc(
19318 &self,
19319 e: &Engine,
19320 fa: &crate::hybrid::FullAttnLayer,
19321 il: usize,
19322 hq: &CudaSlice<i8>,
19323 hdq: &CudaSlice<f32>,
19324 pos_d: &CudaSlice<i32>,
19325 cache: &mut Cache,
19326 cap_bucket_max: Option<(usize, usize)>,
19327 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19328 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
19329 let eps = self.cfg.rms_eps;
19330 let aux = self.gemma4_aux.as_ref().unwrap();
19331 let ones = aux.ones(e);
19332 #[cfg(debug_assertions)]
19333 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
19334 let (q0, k0, v0) = if swa {
19335 match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
19336 Some(t3) => t3,
19337 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
19339 Some((q0, k0)) => {
19340 let h0 = e.zeros(0)?;
19341 let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
19342 (q0, k0, v0)
19343 }
19344 None => {
19345 let h0 = e.zeros(0)?;
19346 (
19347 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
19348 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
19349 e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
19350 )
19351 }
19352 },
19353 }
19354 } else {
19355 let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
19356 Some(p) => p,
19357 None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
19358 Some(p) => p,
19359 None => {
19360 let h0 = e.zeros(0)?;
19361 (
19362 e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
19363 e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
19364 )
19365 }
19366 },
19367 };
19368 let v0 = e.clone_dtod(&k0)?;
19369 (q0, k0, v0)
19370 };
19371 let mut q = e.uninit(nh * hd)?;
19372 let mut k = e.uninit(nkv * hd)?;
19373 let mut v = e.uninit(nkv * hd)?;
19374 let ff = if swa {
19376 None
19377 } else {
19378 Some(
19379 aux.rope_freqs(e)
19380 .expect("gemma4 global rope needs rope_freqs.weight"),
19381 )
19382 };
19383 #[cfg(debug_assertions)]
19384 if let Some(ff) = ff {
19385 crate::debug_assert_tensor_stream_device(
19386 ff,
19387 &e.stream(),
19388 "gemma4_decode_attn_dc.rope_freqs",
19389 );
19390 }
19391 let kvl = cache.kv[il].as_mut().unwrap();
19392 let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
19393 if crate::Engine::qkv_append_on() {
19394 e.rms_norm_qkv_rope_append_dc(
19396 &q0,
19397 &k0,
19398 &v0,
19399 fa.q_norm.float_data(),
19400 fa.k_norm.float_data(),
19401 ones,
19402 &mut q,
19403 &mut k,
19404 &mut v,
19405 hd,
19406 self.gemma4_rope_dims(il),
19407 nh,
19408 nkv,
19409 pos_d,
19410 nh,
19411 nkv,
19412 base,
19413 1.0,
19414 ff,
19415 eps,
19416 &mut kvl.k,
19417 &mut kvl.v,
19418 &kvl.len_d,
19419 kvl.k_tok_bytes,
19420 kvl.v_tok_bytes,
19421 kv_fp8,
19422 )?;
19423 } else {
19424 e.rms_norm_qkv_rope(
19425 &q0,
19426 &k0,
19427 &v0,
19428 fa.q_norm.float_data(),
19429 fa.k_norm.float_data(),
19430 ones,
19431 &mut q,
19432 &mut k,
19433 &mut v,
19434 hd,
19435 self.gemma4_rope_dims(il),
19436 nh,
19437 nkv,
19438 pos_d,
19439 nh,
19440 nkv,
19441 base,
19442 1.0,
19443 ff,
19444 eps,
19445 )?;
19446 e.append_kv_quantized_dc(
19447 &k,
19448 &v,
19449 &mut kvl.k,
19450 &mut kvl.v,
19451 &kvl.len_d,
19452 kvl.kv_dim_k,
19453 kvl.kv_dim_v,
19454 kvl.k_tok_bytes,
19455 kvl.v_tok_bytes,
19456 kv_fp8,
19457 )?;
19458 }
19459 e.inc_seqlen(&mut kvl.len_d)?;
19460 let mut attn = e.uninit(nh * hd)?;
19461 let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
19464 match cap_bucket_max {
19469 None => {
19470 kvl.len += 1;
19474 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19475 if !swa
19476 && hd == 512
19477 && kvl.len >= crate::fa512_min_tkv()
19478 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
19479 {
19480 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
19483 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
19484 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19485 e.fa_decode_rows(
19486 &q,
19487 &kp,
19488 &vp,
19489 &mut attn,
19490 hd,
19491 nh,
19492 nkv,
19493 kvl.len - 1,
19494 1,
19495 scale,
19496 kvl.k_tok_bytes,
19497 kvl.v_tok_bytes,
19498 Some((&kvl.len_d, -1)),
19499 false,
19500 false,
19501 Some((&mut aq8, &mut ad8)),
19502 )?;
19503 fa_q8 = Some((aq8, ad8));
19504 } else if swa
19505 && kvl.len > win
19506 && hd == 256
19507 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
19508 {
19509 let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
19511 let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
19512 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19513 e.fa_decode_rows_w(
19514 &q,
19515 &kp,
19516 &vp,
19517 &mut attn,
19518 hd,
19519 nh,
19520 nkv,
19521 &kvl.len_d,
19522 -1,
19523 1,
19524 scale,
19525 win,
19526 kvl.k_tok_bytes,
19527 kvl.v_tok_bytes,
19528 Some((&mut aq8, &mut ad8)),
19529 )?;
19530 fa_q8 = Some((aq8, ad8));
19531 } else {
19532 let (off_tok, t_kv) = if swa && kvl.len > win {
19533 (kvl.len - win, win)
19534 } else {
19535 (0, kvl.len)
19536 };
19537 let k_view = e.view_u8_range(
19538 &kvl.k,
19539 off_tok * kvl.k_tok_bytes,
19540 (off_tok + t_kv) * kvl.k_tok_bytes,
19541 );
19542 let v_view = e.view_u8_range(
19543 &kvl.v,
19544 off_tok * kvl.v_tok_bytes,
19545 (off_tok + t_kv) * kvl.v_tok_bytes,
19546 );
19547 e.fa_decode_kvmod(
19548 &q,
19549 &k_view,
19550 &v_view,
19551 &mut attn,
19552 hd,
19553 nh,
19554 nkv,
19555 t_kv,
19556 scale,
19557 kvl.k_tok_bytes,
19558 kvl.v_tok_bytes,
19559 swa && crate::Engine::wkv_on(),
19560 )?;
19561 }
19562 }
19563 Some((b_swa, b_glob)) => {
19564 let k_view = e.view_u8(&kvl.k, kvl.k.len());
19570 let v_view = e.view_u8(&kvl.v, kvl.v.len());
19571 let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
19572 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19573 if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
19574 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19575 e.fa_decode_rows(
19576 &q,
19577 &k_view,
19578 &v_view,
19579 &mut attn,
19580 hd,
19581 nh,
19582 nkv,
19583 b_glob - 1,
19584 1,
19585 scale,
19586 kvl.k_tok_bytes,
19587 kvl.v_tok_bytes,
19588 Some((&kvl.len_d, -1)),
19589 false,
19590 false,
19591 Some((&mut aq8, &mut ad8)),
19592 )?;
19593 fa_q8 = Some((aq8, ad8));
19594 } else if swa && b_swa > win && hd == 256 && rows_on {
19595 let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19596 e.fa_decode_rows_w(
19597 &q,
19598 &k_view,
19599 &v_view,
19600 &mut attn,
19601 hd,
19602 nh,
19603 nkv,
19604 &kvl.len_d,
19605 -1,
19606 1,
19607 scale,
19608 win,
19609 kvl.k_tok_bytes,
19610 kvl.v_tok_bytes,
19611 Some((&mut aq8, &mut ad8)),
19612 )?;
19613 fa_q8 = Some((aq8, ad8));
19614 } else {
19615 let b = if swa { b_swa } else { b_glob };
19616 e.fa_decode_dc(
19617 &q,
19618 &k_view,
19619 &v_view,
19620 &mut attn,
19621 hd,
19622 nh,
19623 nkv,
19624 &kvl.len_d,
19625 b,
19626 scale,
19627 kvl.k_tok_bytes,
19628 kvl.v_tok_bytes,
19629 swa && crate::Engine::wkv_on(),
19630 )?;
19631 }
19632 }
19633 }
19634 if let Some((aq8, ad8)) = fa_q8 {
19637 let mut y = e.uninit(fa.wo.out_features())?;
19638 self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
19639 return Ok(y);
19640 }
19641 e.matmul(&fa.wo, &attn, 1)
19642 }
19643
19644 #[allow(clippy::too_many_arguments)]
19649 #[allow(clippy::map_entry)] pub fn gemma4_generate_graph(
19652 &self,
19653 e: &Engine,
19654 prompt_pos: usize,
19655 first_token: u32,
19656 cache: &mut Cache,
19657 max_new: usize,
19658 eos: &[u32],
19659 mut on_token: impl FnMut(u32) -> bool,
19660 ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
19661 if self.is_gemma4_e4b() {
19662 return Err(
19663 "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
19664 .into(),
19665 );
19666 }
19667 use crate::decode::StopReason;
19668 let n_vocab = self.output.out_features();
19669 let n_embd = self.cfg.n_embd as usize;
19670 let embd_gpu = self
19671 .embd_gpu
19672 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
19673 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
19674 for kvl in cache.kv.iter_mut().flatten() {
19675 e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
19676 }
19677 let mut token_d = e.stream().clone_htod(&[first_token])?;
19678 let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
19679 let g4 = self.cfg.gemma4.as_ref().unwrap();
19680 let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
19681 let nkv_s = g4
19683 .head_count_kv
19684 .iter()
19685 .zip(g4.swa_pattern.iter())
19686 .find(|p| *p.1)
19687 .map(|p| *p.0 as usize)
19688 .unwrap_or(8);
19689 let nkv_g = g4
19690 .head_count_kv
19691 .iter()
19692 .zip(g4.swa_pattern.iter())
19693 .find(|p| !*p.1)
19694 .map(|p| *p.0 as usize)
19695 .unwrap_or(2);
19696 #[allow(clippy::type_complexity)]
19697 let mut graphs: std::collections::HashMap<
19699 ((bool, usize), (bool, usize), bool, bool),
19700 (
19701 cudarc::driver::CudaGraph,
19702 Vec<Box<dyn std::any::Any + Send>>,
19703 ),
19704 > = Default::default();
19705 let mut slots = self.g4_dc_slots(e)?;
19708 const RING: usize = 64;
19711 const DRAIN: usize = 1;
19717 let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
19718 let ring_base = prompt_pos;
19719 let mut out = Vec::with_capacity(max_new);
19720 let mut reason = StopReason::MaxNew;
19721 let mut next = first_token;
19722 let mut captures = 0usize;
19723 for _ in 0..max_new {
19724 out.push(next);
19725 if eos.contains(&next) {
19726 reason = StopReason::Eos;
19727 break;
19728 }
19729 if !on_token(next) {
19730 reason = StopReason::Callback;
19731 break;
19732 }
19733 let t_kv = cache.pos + 1;
19734 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19742 let f512 = crate::fa512_min_tkv();
19743 let key_s = if t_kv > win {
19744 (true, usize::MAX)
19745 } else {
19746 e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
19747 };
19748 let (key_g, rung_end) = if t_kv >= f512 {
19749 let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
19752 ((true, end), end)
19753 } else {
19754 (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
19755 };
19756 let key = (key_s, key_g, t_kv >= f512, t_kv > win);
19757 if !graphs.contains_key(&key) {
19758 let bucket_max = (t_kv, rung_end);
19759 let snap = cache.snapshot(e)?;
19761 let pos_save = e.dtoh_i32_one(&pos_d)?;
19762 let len_save: Vec<Option<i32>> = cache
19763 .kv
19764 .iter()
19765 .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
19766 .collect();
19767 let tok_save = e.dtoh_u32_one(&token_d)?;
19768 let graph = {
19773 let tok_ref = &mut token_d;
19774 let pos_ref = &mut pos_d;
19775 let cache_ref = &mut *cache;
19776 let slots_ref = &mut slots;
19777 let ring_ref = &mut ring;
19778 e.capture_graph_retained_flags(
19779 cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
19780 |e| {
19781 let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
19783 let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
19784 let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
19785 self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
19786 cache_ref, n_vocab, Some(bucket_max),
19787 sl, tok_ref, Some((rg, ring_base)))
19788 })?
19789 };
19790 cache.rollback(e, &snap, 0)?;
19791 e.set_i32_one(&mut pos_d, pos_save)?;
19792 for (il, ls) in len_save.iter().enumerate() {
19793 if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
19794 e.set_i32_one(&mut kvl.len_d, *v)?;
19795 }
19796 }
19797 e.set_u32_one(&mut token_d, tok_save)?;
19798 if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
19799 && let Ok(c) = crate::graph_update::node_census(&graph.0)
19800 {
19801 eprintln!("[graph-census] {c:?}");
19802 }
19803 graphs.insert(key, graph);
19804 captures += 1;
19805 }
19806 let mut chunk = 1usize;
19811 let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
19812 .ok()
19813 .and_then(|v| v.parse().ok())
19814 .unwrap_or(DRAIN);
19815 while chunk < drain_cap && out.len() + chunk < max_new {
19816 let t_next = cache.pos + 1 + chunk;
19817 let key_s2 = if t_next > win {
19818 (true, usize::MAX)
19819 } else {
19820 e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
19821 };
19822 let key_g2 = if t_next >= f512 {
19823 (true, (t_next + 1).next_power_of_two().max(f512 * 2))
19824 } else {
19825 e.fa_bucket_key(t_next, hd_g, nkv_g, false)
19826 };
19827 if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
19828 break;
19829 }
19830 chunk += 1;
19831 }
19832 let g = &graphs.get(&key).unwrap().0;
19833 for _ in 0..chunk {
19834 g.launch()?;
19835 }
19836 e.stream().synchronize()?;
19837 let ringh = e.dtoh_u32(&ring)?;
19838 for j in 0..chunk {
19839 let pos_j = cache.pos + j;
19840 let tok_j = ringh[(pos_j - ring_base) % RING];
19841 cache.pos += 0; if j + 1 == chunk {
19843 next = tok_j;
19844 } else {
19845 out.push(tok_j);
19846 if eos.contains(&tok_j) || !on_token(tok_j) {
19847 reason = if eos.contains(&tok_j) {
19848 StopReason::Eos
19849 } else {
19850 StopReason::Callback
19851 };
19852 let keep = cache.pos + j + 1;
19854 e.set_i32_one(&mut pos_d, keep as i32)?;
19855 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
19856 e.set_i32_one(&mut kvl.len_d, keep as i32)?;
19857 kvl.len = keep;
19858 }
19859 cache.pos = keep;
19860 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
19861 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
19862 }
19863 return Ok((out, reason));
19864 }
19865 }
19866 }
19867 cache.pos += chunk;
19868 for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
19869 kvl.len += chunk;
19870 }
19871 }
19872 if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
19873 eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
19874 }
19875 Ok((out, reason))
19876 }
19877
19878 pub(crate) fn gemma4_decode_step_t(
19884 &self,
19885 e: &Engine,
19886 tokens: &[u32],
19887 pos0: usize,
19888 cache: &mut Cache,
19889 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
19890 Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
19891 }
19892
19893 pub(crate) fn gemma4_decode_step_t_am(
19897 &self,
19898 e: &Engine,
19899 tokens: &[u32],
19900 pos0: usize,
19901 cache: &mut Cache,
19902 ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19903 let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
19904 let t = tokens.len();
19905 let n_vocab = self.output.out_features();
19906 let mut toks = e.stream().alloc_zeros::<u32>(t)?;
19907 for i in 0..t {
19908 e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
19909 }
19910 Ok((e.dtoh_u32(&toks)?, hn))
19911 }
19912
19913 pub(crate) fn gemma4_decode_step_t_am_dev(
19916 &self,
19917 e: &Engine,
19918 tok_d: &CudaSlice<u32>,
19919 t: usize,
19920 pos0: usize,
19921 cache: &mut Cache,
19922 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19923 let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
19924 let n_vocab = self.output.out_features();
19925 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
19926 for i in 0..t {
19927 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
19928 }
19929 Ok((vam, hn))
19930 }
19931
19932 pub(crate) fn gemma4_decode_step_t_h(
19935 &self,
19936 e: &Engine,
19937 tokens: &[u32],
19938 pos0: usize,
19939 cache: &mut Cache,
19940 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19941 let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
19942 let t = tokens.len();
19943 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
19944 e.softcap(&mut ld, cap, t * self.output.out_features())?;
19945 Ok((e.dtoh(&ld)?, hn))
19946 }
19947
19948 pub(crate) fn verify_stream_scratch(
19951 &self,
19952 e: &Engine,
19953 cap: usize,
19954 ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
19955 Ok(VerifyStreamScratch {
19956 pos_d: e.htod_i32(&vec![0i32; cap])?,
19957 row_ctrs: (0..cap)
19958 .map(|_| e.htod_i32(&[0]))
19959 .collect::<Result<_, _>>()?,
19960 })
19961 }
19962
19963 #[allow(clippy::too_many_arguments)] pub(crate) fn gemma4_verify_t_am_stream(
19972 &self,
19973 e: &Engine,
19974 tok_d: &CudaSlice<u32>,
19975 t: usize,
19976 ctr: &CudaSlice<i32>,
19977 hint: usize,
19978 cache: &mut Cache,
19979 scr: &mut VerifyStreamScratch,
19980 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19981 let n_embd = self.cfg.n_embd as usize;
19982 let eps = self.cfg.rms_eps;
19983 assert!(t <= scr.row_ctrs.len() && t <= 64);
19984 e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
19985 for i in 0..t {
19986 e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
19987 }
19988 let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
19989 let embd_gpu = self
19990 .embd_gpu
19991 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
19992 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
19993 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
19994 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
19995 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
19996 let n_layers = self.layers.len();
19997 for (il, layer) in self.layers.iter().enumerate() {
19998 let (hq, hdq) = match h_carry.take() {
19999 Some(p) => p,
20000 None => {
20001 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
20002 }
20003 };
20004 let Mixer::Full(fa) = &layer.mixer else {
20005 panic!("gemma4 layer {il} not full-attn")
20006 };
20007 let o = self
20008 .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
20009 let next_norm = if il + 1 < n_layers {
20010 Some(self.layers[il + 1].attn_norm.float_data())
20011 } else {
20012 None
20013 };
20014 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
20015 x = xn;
20016 h_carry = hn;
20017 self.dflash_tap(e, cache, il, &x, t)?;
20018 }
20019 let mut hn = e.uninit(t * n_embd)?;
20020 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
20021 let ld = e.matmul(&self.output, &hn, t)?;
20022 let n_vocab = self.output.out_features();
20023 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
20024 for i in 0..t {
20025 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
20026 }
20027 Ok((vam, hn))
20028 }
20029
20030 pub(crate) fn dflash_tap(
20037 &self,
20038 e: &Engine,
20039 cache: &mut Cache,
20040 il: usize,
20041 x: &CudaSlice<f32>,
20042 t: usize,
20043 ) -> Result<(), Box<dyn std::error::Error>> {
20044 let Some(taps) = cache.dflash_taps.as_mut() else {
20045 return Ok(());
20046 };
20047 let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
20048 return Ok(());
20049 };
20050 let h = taps.hidden;
20051 let n_taps = taps.layer_ids.len();
20052 let base = taps.base;
20053 debug_assert!(
20054 base + t <= taps.t,
20055 "tap window {base}+{t} exceeds sink {}",
20056 taps.t
20057 );
20058 let xv = e.view(x, t * h);
20059 for r in 0..t {
20060 let row = xv.slice(r * h..(r + 1) * h);
20061 e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
20062 }
20063 Ok(())
20064 }
20065
20066 fn gemma4_verify_trunk(
20067 &self,
20068 e: &Engine,
20069 tokens: &[u32],
20070 pos0: usize,
20071 cache: &mut Cache,
20072 tok_dev: Option<&CudaSlice<u32>>,
20073 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20074 let n_embd = self.cfg.n_embd as usize;
20075 let eps = self.cfg.rms_eps;
20076 let t = tokens.len();
20077 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
20078 let pos_d = e.htod_i32(&pos)?;
20079 let mut x = match tok_dev {
20080 Some(td) => {
20081 let embd_gpu = self
20082 .embd_gpu
20083 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
20084 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
20085 e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
20086 }
20087 None => e.htod(&self.embd.gather(n_embd, tokens))?,
20088 };
20089 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
20090 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20091 let n_layers = self.layers.len();
20092 for (il, layer) in self.layers.iter().enumerate() {
20093 let (hq, hdq) = match h_carry.take() {
20094 Some(p) => p,
20095 None => {
20096 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
20097 }
20098 };
20099 let Mixer::Full(fa) = &layer.mixer else {
20100 panic!("gemma4 layer {il} not full-attn")
20101 };
20102 let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
20103 let next_norm = if il + 1 < n_layers {
20104 Some(self.layers[il + 1].attn_norm.float_data())
20105 } else {
20106 None
20107 };
20108 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
20109 x = xn;
20110 h_carry = hn;
20111 self.dflash_tap(e, cache, il, &x, t)?;
20112 }
20113 let mut hn = e.uninit(t * n_embd)?;
20114 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
20115 let mut ld = e.matmul(&self.output, &hn, t)?;
20116 self.gemma4_suppress(e, &mut ld, t)?; cache.pos += t;
20118 Ok((ld, hn))
20119 }
20120
20121 #[allow(clippy::too_many_arguments)]
20129 fn gemma4_verify_attn_stream(
20130 &self,
20131 e: &Engine,
20132 fa: &crate::hybrid::FullAttnLayer,
20133 il: usize,
20134 hq: &CudaSlice<i8>,
20135 hdq: &CudaSlice<f32>,
20136 pos_d: &CudaSlice<i32>,
20137 t: usize,
20138 cache: &mut Cache,
20139 hint: usize,
20140 row_ctrs: &[CudaSlice<i32>],
20141 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20142 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20143 let eps = self.cfg.rms_eps;
20144 let aux = self.gemma4_aux.as_ref().unwrap();
20145 let ones = aux.ones(e);
20146 #[cfg(debug_assertions)]
20147 crate::debug_assert_tensor_stream_device(
20148 ones,
20149 &e.stream(),
20150 "gemma4_verify_attn_stream.ones",
20151 );
20152 let h0 = e.zeros(0)?;
20153 let h = &h0;
20154 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20157 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
20158 let fused_qkv = if f2b {
20159 if swa {
20160 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
20161 .map(|(a, b, c)| (a, b, Some(c)))
20162 } else {
20163 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
20164 .map(|(a, b)| (a, b, None))
20165 }
20166 } else {
20167 None
20168 };
20169 let (q0, k0, v0) = match fused_qkv {
20170 Some((a, b, cv)) => {
20171 let v = match cv {
20172 Some(c) => c,
20173 None => e.clone_dtod(&b)?,
20174 };
20175 (a, b, v)
20176 }
20177 None => {
20178 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
20179 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
20180 let v0 = if swa {
20181 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
20182 } else {
20183 e.clone_dtod(&k0)?
20184 };
20185 (q0, k0, v0)
20186 }
20187 };
20188 let mut q = e.uninit(t * nh * hd)?;
20189 let mut k = e.uninit(t * nkv * hd)?;
20190 let mut v = e.uninit(t * nkv * hd)?;
20191 let ff = if swa {
20194 None
20195 } else {
20196 Some(
20197 aux.rope_freqs(e)
20198 .expect("gemma4 global rope needs rope_freqs.weight"),
20199 )
20200 };
20201 #[cfg(debug_assertions)]
20202 if let Some(ff) = ff {
20203 crate::debug_assert_tensor_stream_device(
20204 ff,
20205 &e.stream(),
20206 "gemma4_verify_attn_stream.rope_freqs",
20207 );
20208 }
20209 e.rms_norm_qkv_rope(
20210 &q0,
20211 &k0,
20212 &v0,
20213 fa.q_norm.float_data(),
20214 fa.k_norm.float_data(),
20215 ones,
20216 &mut q,
20217 &mut k,
20218 &mut v,
20219 hd,
20220 self.gemma4_rope_dims(il),
20221 nh * t,
20222 nkv * t,
20223 pos_d,
20224 nh,
20225 nkv,
20226 base,
20227 1.0,
20228 ff,
20229 eps,
20230 )?;
20231 let kvl = cache.kv[il].as_mut().unwrap();
20232 e.append_kv_quantized_rows_dc(
20234 &k,
20235 &v,
20236 &mut kvl.k,
20237 &mut kvl.v,
20238 &kvl.len_d,
20239 t,
20240 kvl.kv_dim_k,
20241 kvl.kv_dim_v,
20242 kvl.k_tok_bytes,
20243 kvl.v_tok_bytes,
20244 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
20245 )?;
20246 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20249 let mut attn = e.uninit(t * nh * hd)?;
20250 let k_view = e.view_u8(&kvl.k, kvl.k.len());
20251 let v_view = e.view_u8(&kvl.v, kvl.v.len());
20252 if swa && hint + 1 >= win {
20255 e.fa_decode_rows_w(
20258 &q,
20259 &k_view,
20260 &v_view,
20261 &mut attn,
20262 hd,
20263 nh,
20264 nkv,
20265 &kvl.len_d,
20266 0,
20267 t,
20268 scale,
20269 win,
20270 kvl.k_tok_bytes,
20271 kvl.v_tok_bytes,
20272 None,
20273 )?;
20274 } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
20275 let bucket = (hint + t + 2)
20288 .next_power_of_two()
20289 .min(crate::fa512_min_tkv().saturating_sub(1));
20290 let qv = e.view(&q, t * nh * hd);
20291 #[allow(clippy::needless_range_loop)]
20292 for i in 0..t {
20294 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
20295 let mut q_one = e.uninit(nh * hd)?;
20296 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
20297 let mut a_one = e.uninit(nh * hd)?;
20298 e.fa_decode_dc(
20299 &q_one,
20300 &k_view,
20301 &v_view,
20302 &mut a_one,
20303 hd,
20304 nh,
20305 nkv,
20306 &row_ctrs[i],
20307 bucket,
20308 scale,
20309 kvl.k_tok_bytes,
20310 kvl.v_tok_bytes,
20311 false,
20312 )?;
20313 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
20314 }
20315 } else if hd == 512 {
20316 e.fa_decode_rows(
20319 &q,
20320 &k_view,
20321 &v_view,
20322 &mut attn,
20323 hd,
20324 nh,
20325 nkv,
20326 hint,
20327 t,
20328 scale,
20329 kvl.k_tok_bytes,
20330 kvl.v_tok_bytes,
20331 Some((&kvl.len_d, 0)),
20332 false,
20333 false,
20334 None,
20335 )?;
20336 } else {
20337 e.fa_decode_rows_dc(
20339 &q,
20340 &k_view,
20341 &v_view,
20342 &mut attn,
20343 hd,
20344 nh,
20345 nkv,
20346 &kvl.len_d,
20347 hint + t,
20348 t,
20349 scale,
20350 kvl.k_tok_bytes,
20351 kvl.v_tok_bytes,
20352 0,
20353 swa && crate::Engine::wkv_on(),
20354 )?;
20355 }
20356 e.matmul(&fa.wo, &attn, t)
20357 }
20358
20359 #[allow(clippy::too_many_arguments)] fn gemma4_verify_attn(
20361 &self,
20362 e: &Engine,
20363 fa: &crate::hybrid::FullAttnLayer,
20364 il: usize,
20365 hq: &CudaSlice<i8>,
20366 hdq: &CudaSlice<f32>,
20367 pos_d: &CudaSlice<i32>,
20368 t: usize,
20369 cache: &mut Cache,
20370 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20371 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20372 let eps = self.cfg.rms_eps;
20373 let aux = self.gemma4_aux.as_ref().unwrap();
20374 let ones = aux.ones(e);
20375 #[cfg(debug_assertions)]
20376 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
20377 let n_embd = self.cfg.n_embd as usize;
20378 let _ = n_embd;
20379
20380 let h0 = e.zeros(0)?;
20381 let h = &h0;
20382 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20385 let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
20386 let fused_qkv = if f2b {
20387 if swa {
20388 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
20389 .map(|(a, b, c)| (a, b, Some(c)))
20390 } else {
20391 e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
20392 .map(|(a, b)| (a, b, None))
20393 }
20394 } else {
20395 None
20396 };
20397 let (q0, k0, v0) = match fused_qkv {
20398 Some((a, b, cv)) => {
20399 let v = match cv {
20400 Some(c) => c,
20401 None => e.clone_dtod(&b)?,
20402 };
20403 (a, b, v)
20404 }
20405 None => {
20406 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
20407 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
20408 let v0 = if swa {
20409 e.matmul_pre(&fa.wv, hq, hdq, h, t)?
20410 } else {
20411 e.clone_dtod(&k0)?
20412 };
20413 (q0, k0, v0)
20414 }
20415 };
20416 let mut q = e.uninit(t * nh * hd)?;
20417 let mut k = e.uninit(t * nkv * hd)?;
20418 let mut v = e.uninit(t * nkv * hd)?;
20419 let ff = if swa {
20422 None
20423 } else {
20424 Some(
20425 aux.rope_freqs(e)
20426 .expect("gemma4 global rope needs rope_freqs.weight"),
20427 )
20428 };
20429 #[cfg(debug_assertions)]
20430 if let Some(ff) = ff {
20431 crate::debug_assert_tensor_stream_device(
20432 ff,
20433 &e.stream(),
20434 "gemma4_verify_attn.rope_freqs",
20435 );
20436 }
20437 e.rms_norm_qkv_rope(
20438 &q0,
20439 &k0,
20440 &v0,
20441 fa.q_norm.float_data(),
20442 fa.k_norm.float_data(),
20443 ones,
20444 &mut q,
20445 &mut k,
20446 &mut v,
20447 hd,
20448 self.gemma4_rope_dims(il),
20449 nh * t,
20450 nkv * t,
20451 pos_d,
20452 nh,
20453 nkv,
20454 base,
20455 1.0,
20456 ff,
20457 eps,
20458 )?;
20459 let kvl = cache.kv[il].as_mut().unwrap();
20460 let base_len = kvl.len;
20461 e.append_kv_quantized_rows(
20462 &k,
20463 &v,
20464 &mut kvl.k,
20465 &mut kvl.v,
20466 base_len,
20467 t,
20468 kvl.kv_dim_k,
20469 kvl.kv_dim_v,
20470 kvl.k_tok_bytes,
20471 kvl.v_tok_bytes,
20472 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
20473 )?;
20474 kvl.len += t;
20475 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20476 let mut attn = e.uninit(t * nh * hd)?;
20477 let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
20480 || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
20483 if rows_ok && (!swa || base_len + t <= win) {
20484 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
20485 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
20486 if hd == 512 {
20487 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20489 e.fa_decode_rows(
20490 &q,
20491 &k_view,
20492 &v_view,
20493 &mut attn,
20494 hd,
20495 nh,
20496 nkv,
20497 base_len,
20498 t,
20499 scale,
20500 kvl.k_tok_bytes,
20501 kvl.v_tok_bytes,
20502 Some((&kvl.len_d, 0)),
20503 false,
20504 swa && crate::Engine::wkv_on(),
20505 None,
20506 )?;
20507 } else {
20508 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20512 e.fa_decode_rows_dc(
20513 &q,
20514 &k_view,
20515 &v_view,
20516 &mut attn,
20517 hd,
20518 nh,
20519 nkv,
20520 &kvl.len_d,
20521 base_len + t,
20522 t,
20523 scale,
20524 kvl.k_tok_bytes,
20525 kvl.v_tok_bytes,
20526 0,
20527 swa && crate::Engine::wkv_on(),
20528 )?;
20529 }
20530 return e.matmul(&fa.wo, &attn, t);
20531 }
20532 if hd == 256
20540 && swa
20541 && base_len + 1 >= win
20542 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20543 {
20544 let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
20545 let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
20546 e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20547 e.fa_decode_rows_w(
20548 &q,
20549 &k_view,
20550 &v_view,
20551 &mut attn,
20552 hd,
20553 nh,
20554 nkv,
20555 &kvl.len_d,
20556 0,
20557 t,
20558 scale,
20559 win,
20560 kvl.k_tok_bytes,
20561 kvl.v_tok_bytes,
20562 None,
20563 )?;
20564 return e.matmul(&fa.wo, &attn, t);
20565 }
20566 for i in 0..t {
20567 let avail = base_len + i + 1;
20568 let (off_tok, t_kv) = if swa && avail > win {
20569 (avail - win, win)
20570 } else {
20571 (0, avail)
20572 };
20573 let k_view = e.view_u8_range(
20574 &kvl.k,
20575 off_tok * kvl.k_tok_bytes,
20576 (off_tok + t_kv) * kvl.k_tok_bytes,
20577 );
20578 let v_view = e.view_u8_range(
20579 &kvl.v,
20580 off_tok * kvl.v_tok_bytes,
20581 (off_tok + t_kv) * kvl.v_tok_bytes,
20582 );
20583 let qi = e.view(&q, t * nh * hd);
20584 let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
20585 let mut q_one = e.uninit(nh * hd)?;
20586 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
20587 let mut a_one = e.uninit(nh * hd)?;
20588 if swa
20592 && avail > win
20593 && hd == 256
20594 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20595 {
20596 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
20597 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
20598 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
20599 e.fa_decode_rows_w(
20600 &q_one,
20601 &kp,
20602 &vp,
20603 &mut a_one,
20604 hd,
20605 nh,
20606 nkv,
20607 &kvl.len_d,
20608 0,
20609 1,
20610 scale,
20611 win,
20612 kvl.k_tok_bytes,
20613 kvl.v_tok_bytes,
20614 None,
20615 )?;
20616 } else if !swa
20617 && hd == 512
20618 && avail >= crate::fa512_min_tkv()
20619 && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20620 {
20621 let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
20622 let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
20623 e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
20624 e.fa_decode_rows(
20625 &q_one,
20626 &kp,
20627 &vp,
20628 &mut a_one,
20629 hd,
20630 nh,
20631 nkv,
20632 avail - 1,
20633 1,
20634 scale,
20635 kvl.k_tok_bytes,
20636 kvl.v_tok_bytes,
20637 Some((&kvl.len_d, 0)),
20638 false,
20639 false,
20640 None,
20641 )?;
20642 } else {
20643 e.fa_decode_kvmod(
20644 &q_one,
20645 &k_view,
20646 &v_view,
20647 &mut a_one,
20648 hd,
20649 nh,
20650 nkv,
20651 t_kv,
20652 scale,
20653 kvl.k_tok_bytes,
20654 kvl.v_tok_bytes,
20655 swa && crate::Engine::wkv_on(),
20656 )?;
20657 }
20658 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
20659 }
20660 e.matmul(&fa.wo, &attn, t)
20661 }
20662
20663 pub(crate) fn gemma4_decode_step_h(
20666 &self,
20667 e: &Engine,
20668 token: u32,
20669 cache: &mut Cache,
20670 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20671 if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
20676 let rt = crate::pp::Pp2Rt::get(e)?;
20677 let _walk = rt.acquire_walk("gemma4_decode_step_h_pp2")?;
20678 return self.gemma4_decode_step_h_pp2(e, token, cache, split);
20679 }
20680 if crate::pp::pp_cuts(self.layers.len()).is_some() {
20681 crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
20682 }
20683 let n_embd = self.cfg.n_embd as usize;
20684 let eps = self.cfg.rms_eps;
20685 let pos_d = e.htod_i32(&[cache.pos as i32])?;
20686 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
20687 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20688 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20691 let n_layers = self.layers.len();
20692 for (il, layer) in self.layers.iter().enumerate() {
20693 let (hq, hdq) = match h_carry.take() {
20694 Some(p) => p,
20695 None => {
20696 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
20697 }
20698 };
20699 let Mixer::Full(fa) = &layer.mixer else {
20700 panic!("gemma4 layer {il} not full-attn")
20701 };
20702 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
20703 let next_norm = if il + 1 < n_layers {
20704 Some(self.layers[il + 1].attn_norm.float_data())
20705 } else {
20706 None
20707 };
20708 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20709 x = xn;
20710 h_carry = hn;
20711 }
20712 let mut hn = e.uninit(n_embd)?;
20713 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20714 let h_seed = e.clone_dtod(&x)?;
20715 let mut ld = e.matmul(&self.output, &hn, 1)?;
20716 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20717 e.softcap(&mut ld, cap, self.output.out_features())?; self.gemma4_suppress(e, &mut ld, 1)?;
20719 let logits = e.dtoh(&ld)?;
20720 cache.pos += 1;
20721 Ok((logits, h_seed))
20722 }
20723
20724 fn gemma4_decode_layers(
20732 &self,
20733 e: &Engine,
20734 mut x: CudaSlice<f32>,
20735 lo: usize,
20736 hi: usize,
20737 pos_d: &CudaSlice<i32>,
20738 cache: &mut Cache,
20739 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20740 let n_embd = self.cfg.n_embd as usize;
20741 let eps = self.cfg.rms_eps;
20742 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20743 for il in lo..hi {
20744 let layer = &self.layers[il];
20745 let (hq, hdq) = match h_carry.take() {
20746 Some(p) => p,
20747 None => {
20749 e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
20750 }
20751 };
20752 let Mixer::Full(fa) = &layer.mixer else {
20753 panic!("gemma4 layer {il} not full-attn")
20754 };
20755 let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
20756 let next_norm = if il + 1 < hi {
20757 Some(self.layers[il + 1].attn_norm.float_data())
20758 } else {
20759 None
20760 };
20761 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20762 x = xn;
20763 h_carry = hn;
20764 }
20765 Ok(x)
20766 }
20767
20768 fn gemma4_decode_step_h_pp2(
20776 &self,
20777 e: &Engine,
20778 token: u32,
20779 cache: &mut Cache,
20780 split: usize,
20781 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20782 if crate::pp::pp2_streams_off() {
20783 return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
20784 }
20785 let rt = crate::pp::Pp2Rt::get(e)?;
20786 let e0 = rt.engine(0, e);
20787 let e1 = rt.engine(1, e);
20788 let n_embd = self.cfg.n_embd as usize;
20789 let eps = self.cfg.rms_eps;
20790 let pos = cache.pos as i32;
20791
20792 let slot = {
20794 let _st0 = rt.enter(0);
20795 let pos_d = e0.htod_i32(&[pos])?;
20796 #[cfg(debug_assertions)]
20797 crate::debug_assert_tensor_stream_device(
20798 &pos_d,
20799 &e0.stream(),
20800 "gemma4_decode_step_h_pp2.stage0.pos_d",
20801 );
20802 let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
20803 e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20804 let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
20805 rt.tx(0, &x, n_embd)?
20806 };
20807
20808 let _st1 = rt.enter(1);
20810 let pos_d = e1.htod_i32(&[pos])?;
20811 #[cfg(debug_assertions)]
20812 crate::debug_assert_tensor_stream_device(
20813 &pos_d,
20814 &e1.stream(),
20815 "gemma4_decode_step_h_pp2.stage1.pos_d",
20816 );
20817 let x = rt.rx(0, slot, n_embd)?;
20818 let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
20819
20820 let mut hn = e1.uninit(n_embd)?;
20821 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20822 let h_seed = e1.clone_dtod(&x)?;
20823 let mut ld = e1.matmul(&self.output, &hn, 1)?;
20824 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20825 e1.softcap(&mut ld, cap, self.output.out_features())?;
20826 self.gemma4_suppress(e1, &mut ld, 1)?;
20827 let logits = e1.dtoh(&ld)?;
20828 cache.pos += 1;
20829 Ok((logits, h_seed))
20830 }
20831
20832 fn gemma4_decode_step_h_pp2_samestream(
20835 &self,
20836 e: &Engine,
20837 token: u32,
20838 cache: &mut Cache,
20839 split: usize,
20840 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20841 let n_embd = self.cfg.n_embd as usize;
20842 let eps = self.cfg.rms_eps;
20843 let pos_d = e.htod_i32(&[cache.pos as i32])?;
20844
20845 let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
20847 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20848 let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
20849
20850 let boundary_tx = e.clone_dtod(&x)?;
20852 let boundary_rx = e.clone_dtod(&boundary_tx)?;
20853
20854 let x =
20856 self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
20857
20858 let mut hn = e.uninit(n_embd)?;
20859 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20860 let h_seed = e.clone_dtod(&x)?;
20861 let mut ld = e.matmul(&self.output, &hn, 1)?;
20862 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20863 e.softcap(&mut ld, cap, self.output.out_features())?;
20864 self.gemma4_suppress(e, &mut ld, 1)?;
20865 let logits = e.dtoh(&ld)?;
20866 cache.pos += 1;
20867 Ok((logits, h_seed))
20868 }
20869}
20870
20871impl HybridModel {
20890 pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
20893 let geometry = self
20894 .cfg
20895 .layer_geometry(il as u32)
20896 .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
20897 debug_assert_eq!(
20898 geometry.attention_gate,
20899 memra_gguf::config::AttentionGateKind::SeparateHead
20900 );
20901 geometry
20902 }
20903
20904 #[allow(clippy::too_many_arguments)]
20964 fn step35_attn_pre_wo(
20965 &self,
20966 e: &Engine,
20967 fa: &FullAttnLayer,
20968 mut g3: Vec<CudaSlice<f32>>,
20969 hg: Option<&CudaSlice<f32>>,
20970 gt_pre: Option<&CudaSlice<f32>>,
20971 pos_d: &CudaSlice<i32>,
20972 t: usize,
20973 cache: Option<&mut Cache>,
20974 il: usize,
20975 seq_end: usize,
20976 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20977 let geometry = self.cfg.full_attention_geometry_at(il as u32);
20978 let hd = geometry.head_dim_k as usize;
20979 let nkv = geometry.n_head_kv as usize;
20980 let nh = geometry.n_head as usize;
20981 let rbase = geometry.rope_base;
20982 let scale = geometry.attention_scale();
20983 let swa = geometry.window.is_some();
20984 let eps = self.cfg.rms_eps;
20985 let win = geometry.window.unwrap_or(0) as usize;
20986 let n_rot = geometry.n_rot as usize;
20987
20988 let v = g3.pop().unwrap();
20989 let k0 = g3.pop().unwrap();
20990 let q0 = g3.pop().unwrap();
20991
20992 let mut q = e.uninit(t * nh * hd)?;
20996 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
20997 let mut k = e.uninit(t * nkv * hd)?;
20998 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
20999 let ff = if geometry.rope_factors {
21000 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
21001 } else {
21002 None
21003 };
21004 #[cfg(debug_assertions)]
21005 if let Some(ff) = ff {
21006 crate::debug_assert_tensor_stream_device(
21007 ff,
21008 &e.stream(),
21009 "step35_attn_pre_wo.rope_freqs",
21010 );
21011 }
21012 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
21013
21014 let mut attn = e.uninit(t * nh * hd)?;
21015 match cache {
21016 Some(cache) => {
21017 let base_len = cache.kv[il].as_ref().unwrap().len;
21018 let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
21020 let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
21021 let off = if swa {
21022 let raw = base_len.saturating_sub(win - 1);
21023 if legacy_tkv || legacy_calllocal {
21024 raw
21025 } else {
21026 raw & !31usize
21027 }
21028 } else {
21029 0
21030 };
21031 {
21032 let kvl = cache.kv[il].as_mut().unwrap();
21033 assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
21034 let write_row = e.prepare_kv_append(kvl, off, t)?;
21035 e.append_kv_quantized_rows(
21036 &k,
21037 &v,
21038 &mut kvl.k,
21039 &mut kvl.v,
21040 write_row,
21041 t,
21042 kvl.kv_dim_k,
21043 kvl.kv_dim_v,
21044 kvl.k_tok_bytes,
21045 kvl.v_tok_bytes,
21046 crate::Engine::kv_fp8_on(),
21047 )?;
21048 kvl.len += t;
21049 let new_len = kvl.len as i32;
21050 e.set_i32_one(&mut kvl.len_d, new_len)?;
21051 }
21052 let kvl = cache.kv[il].as_ref().unwrap();
21053 let t_kv = base_len + t - off;
21076 let physical = kvl.physical_rows(off, off + t_kv)?;
21077 let k_view = e.view_u8_range(
21078 &kvl.k,
21079 physical.start * kvl.k_tok_bytes,
21080 physical.end * kvl.k_tok_bytes,
21081 );
21082 let v_view = e.view_u8_range(
21083 &kvl.v,
21084 physical.start * kvl.v_tok_bytes,
21085 physical.end * kvl.v_tok_bytes,
21086 );
21087 let swa_naive = if legacy_tkv {
21099 t_kv > win
21100 } else {
21101 seq_end > win
21102 };
21103 if swa && swa_naive {
21104 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
21117 e.sdpa_naive_w_quantized_view(
21118 &q,
21119 &k_view,
21120 &v_view,
21121 &mut attn,
21122 hd,
21123 nh,
21124 nkv,
21125 t,
21126 t_kv,
21127 scale,
21128 true,
21129 win,
21130 kvl.k_tok_bytes,
21131 kvl.v_tok_bytes,
21132 )?;
21133 } else {
21134 e.fa_prefill_view_ws_w_hd128(
21135 &q,
21136 &k_view,
21137 &v_view,
21138 &mut attn,
21139 hd,
21140 nh,
21141 nkv,
21142 t,
21143 t_kv,
21144 scale,
21145 true,
21146 win,
21147 kvl.k_tok_bytes,
21148 kvl.v_tok_bytes,
21149 )?;
21150 }
21151 } else if std::env::var("MEMRA_NOFA").is_ok() {
21152 e.sdpa_naive_quantized_view(
21153 &q,
21154 &k_view,
21155 &v_view,
21156 &mut attn,
21157 hd,
21158 nh,
21159 nkv,
21160 t,
21161 t_kv,
21162 scale,
21163 true,
21164 kvl.k_tok_bytes,
21165 kvl.v_tok_bytes,
21166 )?;
21167 } else {
21168 e.fa_prefill_view_ws(
21173 &q,
21174 &k_view,
21175 &v_view,
21176 &mut attn,
21177 hd,
21178 nh,
21179 nkv,
21180 t,
21181 t_kv,
21182 scale,
21183 true,
21184 kvl.k_tok_bytes,
21185 kvl.v_tok_bytes,
21186 crate::Engine::kv_fp8_on(),
21187 )?;
21188 }
21189 }
21190 None => {
21191 debug_assert_eq!(
21196 seq_end, t,
21197 "step35 cacheless prefill is monolithic (seq_end == t)"
21198 );
21199 if swa && seq_end > win {
21200 e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
21201 } else if std::env::var("MEMRA_NOFA").is_ok() {
21202 e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
21203 } else {
21204 e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
21205 }
21206 }
21207 }
21208
21209 let gw = fa
21212 .attn_gate
21213 .as_ref()
21214 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
21215 let gt_owned = if gt_pre.is_none() {
21216 Some(e.matmul(
21217 gw,
21218 hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
21219 t,
21220 )?)
21221 } else {
21222 None
21223 };
21224 let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
21225 let mut ag = e.uninit(t * nh * hd)?;
21226 e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
21227 Ok(ag)
21228 }
21229
21230 pub(crate) fn step35_attn(
21233 &self,
21234 e: &Engine,
21235 fa: &FullAttnLayer,
21236 h: &CudaSlice<f32>,
21237 pos_d: &CudaSlice<i32>,
21238 t: usize,
21239 il: usize,
21240 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21241 let g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
21242 Some(g3) => g3,
21243 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
21244 };
21245 let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
21247 self.full_attn_o(e, fa, &ag, t)
21248 }
21249
21250 #[allow(clippy::too_many_arguments)]
21257 pub(crate) fn step35_attn_prime(
21258 &self,
21259 e: &Engine,
21260 fa: &FullAttnLayer,
21261 h: &CudaSlice<f32>,
21262 hx: Option<&CudaSlice<u8>>,
21263 pos_d: &CudaSlice<i32>,
21264 t: usize,
21265 cache: &mut Cache,
21266 il: usize,
21267 seq_end: usize,
21268 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21269 if step_tp_prefill_enabled()? && fa.step_tp_qkv.is_some() {
21270 if hx.is_some() {
21271 return Err(
21272 "rank-local Step prefill preserves BF16 activations and refuses the q8_1 \
21273 pre-quantized prime path"
21274 .into(),
21275 );
21276 }
21277 return self.step35_tp_prefill_attn_resident(e, fa, il, h, pos_d, t, cache, seq_end);
21278 }
21279 let g3 = if fa.step_tp_qkv.is_some() {
21280 if hx.is_some() {
21281 return Err(
21282 "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
21283 pre-quantized prime path"
21284 .into(),
21285 );
21286 }
21287 self.full_attn_tp_qkv(e, fa, h, t)?
21288 .expect("Step Q/K/V TP disappeared after the presence check")
21289 } else {
21290 match hx {
21291 Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
21292 None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
21293 }
21294 };
21295 let ag =
21296 self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
21297 self.full_attn_o(e, fa, &ag, t)
21298 }
21299
21300 fn ensure_step_tp_kv_cache(
21301 &self,
21302 e: &Engine,
21303 fa: &FullAttnLayer,
21304 il: usize,
21305 cache: &mut Cache,
21306 ) -> Result<bool, Box<dyn std::error::Error>> {
21307 let tp = fa
21308 .step_tp_qkv
21309 .as_ref()
21310 .ok_or("Step TP cache hydration lost its resident projections")?;
21311 let geometry = self.cfg.full_attention_geometry_at(il as u32);
21312 let window = geometry.window.map(|window| window as usize);
21313 let ranks = tp.runtime.devices().len();
21314 let head_dim = geometry.head_dim_k as usize;
21315 let kv_heads = geometry.n_head_kv as usize;
21316 let max_ctx = cache.max_ctx;
21317
21318 if cache.tp_kv[il].is_some() {
21319 return Ok(false);
21320 }
21321 let local = cache.kv[il]
21322 .as_ref()
21323 .ok_or_else(|| format!("Step TP layer {il} has no owning-stage KV cache"))?;
21324 if local.kv_dim_k != kv_heads * head_dim || local.kv_dim_v != kv_heads * head_dim {
21325 return Err(format!(
21326 "Step TP layer {il} local KV geometry k={} v={} != {}",
21327 local.kv_dim_k,
21328 local.kv_dim_v,
21329 kv_heads * head_dim
21330 )
21331 .into());
21332 }
21333 let resident_start = window
21334 .map(|window| local.len.saturating_sub(window.saturating_sub(1)) & !31usize)
21335 .unwrap_or(0);
21336 let resident_rows = local.len - resident_start;
21337 let physical = local.physical_rows(resident_start, local.len)?;
21338 let k_rows = if resident_rows == 0 {
21339 Vec::new()
21340 } else {
21341 e.dtoh_u8_view(&e.view_u8_range(
21342 &local.k,
21343 physical.start * local.k_tok_bytes,
21344 physical.end * local.k_tok_bytes,
21345 ))?
21346 };
21347 let v_rows = if resident_rows == 0 {
21348 Vec::new()
21349 } else {
21350 e.dtoh_u8_view(&e.view_u8_range(
21351 &local.v,
21352 physical.start * local.v_tok_bytes,
21353 physical.end * local.v_tok_bytes,
21354 ))?
21355 };
21356 let mut distributed = match window {
21357 Some(window) => tp.runtime.allocate_tp_swa_kv_cache(
21358 kv_heads * head_dim,
21359 kv_heads * head_dim,
21360 max_ctx,
21361 window,
21362 )?,
21363 None => tp.runtime.allocate_tp_kv_cache(
21364 kv_heads * head_dim,
21365 kv_heads * head_dim,
21366 max_ctx,
21367 )?,
21368 };
21369 if distributed.k_tok_bytes() * ranks != local.k_tok_bytes
21370 || distributed.v_tok_bytes() * ranks != local.v_tok_bytes
21371 {
21372 return Err(format!(
21373 "Step TP layer {il} distributed/local KV token bytes disagree: \
21374 k={}x{ranks}/{} v={}x{ranks}/{}",
21375 distributed.k_tok_bytes(),
21376 local.k_tok_bytes,
21377 distributed.v_tok_bytes(),
21378 local.v_tok_bytes,
21379 )
21380 .into());
21381 }
21382 tp.runtime.hydrate_tp_kv_cache_from(
21383 &mut distributed,
21384 local.len,
21385 resident_start,
21386 &k_rows,
21387 &v_rows,
21388 )?;
21389 cache.tp_kv[il] = Some(distributed);
21390 Ok(true)
21391 }
21392
21393 #[allow(clippy::too_many_arguments)]
21394 #[allow(clippy::manual_is_multiple_of)] fn step35_tp_prefill_attn_resident(
21396 &self,
21397 e: &Engine,
21398 fa: &FullAttnLayer,
21399 il: usize,
21400 h: &CudaSlice<f32>,
21401 pos_d: &CudaSlice<i32>,
21402 tokens: usize,
21403 cache: &mut Cache,
21404 seq_end: usize,
21405 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21406 let tp = fa
21407 .step_tp_qkv
21408 .as_ref()
21409 .ok_or("Step TP prefill lost its resident projections")?;
21410 let attention = tp
21411 .attention
21412 .as_ref()
21413 .ok_or("Step TP prefill lost its resident attention auxiliaries")?;
21414 let ranks = tp.runtime.devices().len();
21415 if !step_tp_prefill_shape(
21416 true,
21417 tokens,
21418 ranks,
21419 tp.runtime.native_p2p(),
21420 true,
21421 crate::Engine::kv_fp8_on(),
21422 ) {
21423 return Err(format!(
21424 "rank-local Step prefill requires tokens>={PRIME_MIN_T}, TP2/TP4 native P2P, \
21425 rank-local attention, and q8_0/q5_1 KV; got tokens={tokens} ranks={ranks} \
21426 native_p2p={} fp8_kv={}",
21427 tp.runtime.native_p2p(),
21428 crate::Engine::kv_fp8_on(),
21429 )
21430 .into());
21431 }
21432 for seam in [
21433 "MEMRA_STEP35_SWA_TKV",
21434 "MEMRA_PRIME_CALLLOCAL",
21435 "MEMRA_PRIME_F32CHUNK0",
21436 ] {
21437 if std::env::var(seam).as_deref() == Ok("1") {
21438 return Err(format!(
21439 "rank-local Step prefill has not qualified the legacy seam {seam}=1"
21440 )
21441 .into());
21442 }
21443 }
21444
21445 let geometry = self.cfg.full_attention_geometry_at(il as u32);
21446 let window = geometry.window.map(|window| window as usize);
21447 let head_dim = geometry.head_dim_k as usize;
21448 let heads = geometry.n_head as usize;
21449 let kv_heads = geometry.n_head_kv as usize;
21450 if heads % ranks != 0 || kv_heads % ranks != 0 {
21451 return Err(format!(
21452 "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
21453 )
21454 .into());
21455 }
21456 let local_heads = heads / ranks;
21457 let local_kv_heads = kv_heads / ranks;
21458 let local_kv_dim = local_kv_heads * head_dim;
21459 let hidden = self.cfg.n_embd as usize;
21460 let expected_input = tokens
21461 .checked_mul(hidden)
21462 .ok_or("Step TP prefill input size overflow")?;
21463 if h.len() < expected_input {
21464 return Err(format!(
21465 "Step TP prefill input {} is shorter than {tokens}x{hidden}",
21466 h.len()
21467 )
21468 .into());
21469 }
21470 let positions = e.dtoh_i32(pos_d)?;
21471 if positions.len() != tokens {
21472 return Err(format!(
21473 "rank-local Step prefill positions {} != tokens {tokens}",
21474 positions.len()
21475 )
21476 .into());
21477 }
21478
21479 let mut active_input = e.uninit(expected_input)?;
21480 e.copy_view_into(
21481 &mut active_input,
21482 0,
21483 &h.slice(0..expected_input),
21484 expected_input,
21485 )?;
21486 let mut input = tp.runtime.allocate_replicated_device_rows(tokens, hidden)?;
21487 e.stream().synchronize()?;
21492 tp.runtime
21493 .refresh_replicated_device_rows_from_root(&mut input, &active_input)?;
21494 let q_raw = tp
21495 .runtime
21496 .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &input)?;
21497 let k_raw = tp
21498 .runtime
21499 .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &input)?;
21500 let v_raw = tp
21501 .runtime
21502 .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &input)?;
21503 let mut q = Vec::with_capacity(ranks);
21504 let mut k = Vec::with_capacity(ranks);
21505 for rank in 0..ranks {
21506 let engine = tp
21507 .runtime
21508 .rank_engine(rank)
21509 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21510 let _main = engine.gpu.enter_main()?;
21511 let mut q_rank = engine.uninit(tokens * local_heads * head_dim)?;
21512 engine.rms_norm(
21513 &q_raw[rank],
21514 &attention.q_norm[rank],
21515 &mut q_rank,
21516 head_dim,
21517 tokens * local_heads,
21518 self.cfg.rms_eps,
21519 )?;
21520 let mut k_rank = engine.uninit(tokens * local_kv_dim)?;
21521 engine.rms_norm(
21522 &k_raw[rank],
21523 &attention.k_norm[rank],
21524 &mut k_rank,
21525 head_dim,
21526 tokens * local_kv_heads,
21527 self.cfg.rms_eps,
21528 )?;
21529 let position = engine.htod_i32(&positions)?;
21530 let rope_freqs = if geometry.rope_factors {
21531 self.step35_aux
21532 .as_ref()
21533 .and_then(|aux| aux.rope_freqs(engine))
21534 } else {
21535 None
21536 };
21537 engine.rope_neox2(
21538 &mut q_rank,
21539 &mut k_rank,
21540 &position,
21541 head_dim,
21542 geometry.n_rot as usize,
21543 local_heads,
21544 local_kv_heads,
21545 tokens,
21546 geometry.rope_base,
21547 1.0,
21548 rope_freqs,
21549 )?;
21550 q.push(q_rank);
21551 k.push(k_rank);
21552 }
21553
21554 let gate_weight = fa
21555 .attn_gate
21556 .as_ref()
21557 .ok_or("step35 layer is missing attn_gate.weight")?;
21558 let gate = e.dtoh(&e.matmul(gate_weight, h, tokens)?)?;
21559 if gate.len() != tokens * heads {
21560 return Err(format!(
21561 "Step TP layer {il} gate output {} != {tokens}x{heads}",
21562 gate.len()
21563 )
21564 .into());
21565 }
21566
21567 let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
21568 let base_len = cache.kv[il]
21569 .as_ref()
21570 .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
21571 .len;
21572 let distributed = cache.tp_kv[il]
21573 .as_ref()
21574 .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
21575 if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
21576 return Err(format!(
21577 "Step TP layer {il} cache lengths diverged before prefill: \
21578 local={base_len} distributed={}/{}",
21579 distributed.committed_len(),
21580 distributed.staged_len()
21581 )
21582 .into());
21583 }
21584 let target_len = base_len
21585 .checked_add(tokens)
21586 .ok_or("Step TP prefill cache length overflow")?;
21587 if target_len > cache.max_ctx {
21588 return Err(format!(
21589 "Step TP layer {il} prefill exceeds cache: {base_len}+{tokens}>{}",
21590 cache.max_ctx
21591 )
21592 .into());
21593 }
21594 if seq_end < target_len {
21595 return Err(format!(
21596 "Step TP layer {il} request end {seq_end} precedes chunk end {target_len}"
21597 )
21598 .into());
21599 }
21600
21601 let transaction = cache.tp_kv[il]
21602 .as_mut()
21603 .expect("distributed cache checked above")
21604 .begin_transaction()?;
21605 if let Err(error) = tp.runtime.append_tp_kv_transaction(
21606 cache.tp_kv[il]
21607 .as_mut()
21608 .expect("distributed cache checked above"),
21609 transaction,
21610 &k,
21611 &v_raw,
21612 tokens,
21613 ) {
21614 let _ = tp.runtime.rollback_tp_kv_transaction(
21615 cache.tp_kv[il]
21616 .as_mut()
21617 .expect("distributed cache checked above"),
21618 transaction,
21619 );
21620 return Err(error);
21621 }
21622
21623 let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21624 let distributed = cache.tp_kv[il]
21625 .as_ref()
21626 .expect("distributed cache checked above");
21627 let staged_len = distributed.staged_len();
21628 let view_start = window
21629 .map(|window| base_len.saturating_sub(window.saturating_sub(1)) & !31usize)
21630 .unwrap_or(0);
21631 let physical = distributed.physical_range(view_start, staged_len)?;
21632 let t_kv = staged_len - view_start;
21633 let swa_naive = window.is_some_and(|window| seq_end > window);
21634 let mut gated = Vec::with_capacity(ranks);
21635 #[allow(clippy::needless_range_loop)]
21636 for rank in 0..ranks {
21638 let engine = tp
21639 .runtime
21640 .rank_engine(rank)
21641 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21642 let _main = engine.gpu.enter_main()?;
21643 let rank_cache = distributed
21644 .rank(rank)
21645 .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
21646 let k_view = engine.view_u8_range(
21647 rank_cache.k(),
21648 physical.start * distributed.k_tok_bytes(),
21649 physical.end * distributed.k_tok_bytes(),
21650 );
21651 let v_view = engine.view_u8_range(
21652 rank_cache.v(),
21653 physical.start * distributed.v_tok_bytes(),
21654 physical.end * distributed.v_tok_bytes(),
21655 );
21656 let mut attention_out = engine.uninit(tokens * local_heads * head_dim)?;
21657 if swa_naive {
21658 let window = window.expect("SWA predicate requires a window");
21659 if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
21660 engine.sdpa_naive_w_quantized_view(
21661 &q[rank],
21662 &k_view,
21663 &v_view,
21664 &mut attention_out,
21665 head_dim,
21666 local_heads,
21667 local_kv_heads,
21668 tokens,
21669 t_kv,
21670 geometry.attention_scale(),
21671 true,
21672 window,
21673 distributed.k_tok_bytes(),
21674 distributed.v_tok_bytes(),
21675 )?;
21676 } else {
21677 engine.fa_prefill_view_ws_w_hd128(
21678 &q[rank],
21679 &k_view,
21680 &v_view,
21681 &mut attention_out,
21682 head_dim,
21683 local_heads,
21684 local_kv_heads,
21685 tokens,
21686 t_kv,
21687 geometry.attention_scale(),
21688 true,
21689 window,
21690 distributed.k_tok_bytes(),
21691 distributed.v_tok_bytes(),
21692 )?;
21693 }
21694 } else if std::env::var("MEMRA_NOFA").is_ok() {
21695 engine.sdpa_naive_quantized_view(
21696 &q[rank],
21697 &k_view,
21698 &v_view,
21699 &mut attention_out,
21700 head_dim,
21701 local_heads,
21702 local_kv_heads,
21703 tokens,
21704 t_kv,
21705 geometry.attention_scale(),
21706 true,
21707 distributed.k_tok_bytes(),
21708 distributed.v_tok_bytes(),
21709 )?;
21710 } else {
21711 engine.fa_prefill_view_ws(
21712 &q[rank],
21713 &k_view,
21714 &v_view,
21715 &mut attention_out,
21716 head_dim,
21717 local_heads,
21718 local_kv_heads,
21719 tokens,
21720 t_kv,
21721 geometry.attention_scale(),
21722 true,
21723 distributed.k_tok_bytes(),
21724 distributed.v_tok_bytes(),
21725 false,
21726 )?;
21727 }
21728
21729 let gate_start = rank * local_heads;
21730 let mut gate_rank = Vec::with_capacity(tokens * local_heads);
21731 for token in 0..tokens {
21732 let start = token * heads + gate_start;
21733 gate_rank.extend_from_slice(&gate[start..start + local_heads]);
21734 }
21735 let gate_rank = engine.htod(&gate_rank)?;
21736 let mut gated_rank = engine.uninit(tokens * local_heads * head_dim)?;
21737 engine.attn_head_gate(
21738 &attention_out,
21739 &gate_rank,
21740 &mut gated_rank,
21741 None,
21742 head_dim,
21743 local_heads,
21744 tokens,
21745 )?;
21746 gated.push(gated_rank);
21747 }
21748 for rank in 1..ranks {
21749 let engine = tp
21750 .runtime
21751 .rank_engine(rank)
21752 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21753 let _main = engine.gpu.enter_main()?;
21754 engine.stream().synchronize()?;
21755 }
21756
21757 let (output, k_shadow, v_shadow) = if tp.runtime.bulk_p2p() {
21758 let output = tp
21759 .runtime
21760 .step_bf16_row_parallel_resident_root_device(&tp.o, &gated, tokens)?;
21761 let k_shadow =
21762 tp.runtime
21763 .gather_native_column_shards_device(&k, tokens, local_kv_dim)?;
21764 let v_shadow =
21765 tp.runtime
21766 .gather_native_column_shards_device(&v_raw, tokens, local_kv_dim)?;
21767 let root = tp
21768 .runtime
21769 .rank_engine(0)
21770 .ok_or("Step TP prefill lost its root engine")?;
21771 let _main = root.gpu.enter_main()?;
21772 root.stream().synchronize()?;
21773 (output, k_shadow, v_shadow)
21774 } else {
21775 let attention = tp.runtime.gather_native_column_shards(
21776 &gated,
21777 tokens,
21778 local_heads * head_dim,
21779 )?;
21780 let output = tp
21781 .runtime
21782 .step_bf16_row_parallel_resident_native(&tp.o, &attention, tokens)?;
21783 let k_shadow = tp
21784 .runtime
21785 .gather_native_column_shards(&k, tokens, local_kv_dim)?;
21786 let v_shadow =
21787 tp.runtime
21788 .gather_native_column_shards(&v_raw, tokens, local_kv_dim)?;
21789 (e.htod(&output)?, e.htod(&k_shadow)?, e.htod(&v_shadow)?)
21790 };
21791 let local = cache.kv[il]
21792 .as_mut()
21793 .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
21794 if local.len != base_len {
21795 return Err(format!(
21796 "Step TP layer {il} local cache changed during prefill: \
21797 len={} base={base_len}",
21798 local.len
21799 )
21800 .into());
21801 }
21802 let retain_from = window
21803 .map(|window| {
21804 let staged_retain = staged_len.saturating_sub(window) & !31usize;
21805 let rollback_retain =
21806 base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
21807 staged_retain.min(rollback_retain)
21808 })
21809 .unwrap_or(0);
21810 let write_row = e.prepare_kv_append(local, retain_from, tokens)?;
21811 e.append_kv_quantized_rows(
21812 &k_shadow,
21813 &v_shadow,
21814 &mut local.k,
21815 &mut local.v,
21816 write_row,
21817 tokens,
21818 local.kv_dim_k,
21819 local.kv_dim_v,
21820 local.k_tok_bytes,
21821 local.v_tok_bytes,
21822 false,
21823 )?;
21824 local.len = staged_len;
21825 e.set_i32_one(&mut local.len_d, staged_len as i32)?;
21826 Ok(output)
21827 })();
21828
21829 let output = match staged {
21830 Ok(output) => output,
21831 Err(error) => {
21832 let _ = tp.runtime.rollback_tp_kv_transaction(
21833 cache.tp_kv[il]
21834 .as_mut()
21835 .expect("distributed cache checked above"),
21836 transaction,
21837 );
21838 if let Some(local) = cache.kv[il].as_mut() {
21839 local.len = base_len;
21840 let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
21841 }
21842 return Err(error);
21843 }
21844 };
21845 if let Err(error) = tp.runtime.commit_tp_kv_transaction(
21846 cache.tp_kv[il]
21847 .as_mut()
21848 .expect("distributed cache checked above"),
21849 transaction,
21850 tokens,
21851 ) {
21852 let _ = tp.runtime.rollback_tp_kv_transaction(
21853 cache.tp_kv[il]
21854 .as_mut()
21855 .expect("distributed cache checked above"),
21856 transaction,
21857 );
21858 let local = cache.kv[il].as_mut().expect("local cache checked above");
21859 local.len = base_len;
21860 e.set_i32_one(&mut local.len_d, base_len as i32)?;
21861 return Err(error);
21862 }
21863
21864 let committed = cache.tp_kv[il]
21865 .as_ref()
21866 .expect("distributed cache checked above")
21867 .committed_len();
21868 let local_len = cache.kv[il]
21869 .as_ref()
21870 .expect("local cache checked above")
21871 .len;
21872 if committed != local_len {
21873 return Err(format!(
21874 "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
21875 )
21876 .into());
21877 }
21878 eprintln!(
21879 "[step-tp-prefill-attn] execute layer={} devices={:?} tokens={tokens} \
21880 qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
21881 kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
21882 attention_scope={} input_path=root-device-replicated gate_tensor_parallel=false \
21883 gate_shards=host-canonical o_tensor_parallel=true local_cache_shadow=true \
21884 cache_commit=chunk transport={} native_p2p=true bulk_p2p={} \
21885 output={} performance_claim=false",
21886 tp.layer,
21887 tp.devices,
21888 hydrated,
21889 if window.is_some() {
21890 "rank-local-swa-ring"
21891 } else {
21892 "rank-local-global"
21893 },
21894 tp.runtime.transport_label(),
21895 tp.runtime.bulk_p2p(),
21896 if tp.runtime.bulk_p2p() {
21897 "root-device"
21898 } else {
21899 "root-readback"
21900 },
21901 );
21902 Ok(output)
21903 }
21904
21905 fn step35_tp_decode_attn_resident(
21906 &self,
21907 e: &Engine,
21908 fa: &FullAttnLayer,
21909 il: usize,
21910 h: &CudaSlice<f32>,
21911 pos_d: &CudaSlice<i32>,
21912 cache: &mut Cache,
21913 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21914 static ATTN_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21918 static ATTN_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21919 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
21920 let started = timing.then(std::time::Instant::now);
21921 let result = if crate::tp::step_tp_decode_v2_enabled()? {
21922 self.step35_tp_decode_attn_resident_v2(e, fa, il, h, pos_d, cache)
21923 } else {
21924 self.step35_tp_decode_attn_resident_inner(e, fa, il, h, pos_d, cache)
21925 };
21926 if let Some(started) = started {
21927 use std::sync::atomic::Ordering;
21928 let ns = ATTN_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
21929 + started.elapsed().as_nanos() as u64;
21930 let calls = ATTN_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
21931 if calls.is_multiple_of(430) {
21932 eprintln!(
21933 "[step-tp-attn-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
21934 ns as f64 / 1.0e6,
21935 ns as f64 / calls as f64 / 1.0e3,
21936 );
21937 }
21938 }
21939 result
21940 }
21941
21942 #[allow(clippy::too_many_arguments)]
21943 fn step35_tp_decode_attn_resident_inner(
21944 &self,
21945 e: &Engine,
21946 fa: &FullAttnLayer,
21947 il: usize,
21948 h: &CudaSlice<f32>,
21949 pos_d: &CudaSlice<i32>,
21950 cache: &mut Cache,
21951 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21952 static T_POS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21957 static T_QKV: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21958 static T_NORMROPE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21959 static T_GATE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21960 static T_APPEND: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21961 static T_ATTN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21962 static T_OPROJ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21963 static T_SHADOW: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21964 static T_PHASE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21965 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
21966 #[allow(clippy::manual_is_multiple_of)] fn lap(
21968 runtime: &crate::tp::TpE4m3HostBounce,
21969 e: &Engine,
21970 timer: &std::sync::atomic::AtomicU64,
21971 started: &mut Option<std::time::Instant>,
21972 ) -> Result<(), Box<dyn std::error::Error>> {
21973 let Some(start) = started.as_mut() else {
21974 return Ok(());
21975 };
21976 for rank in 0..runtime.devices().len() {
21977 if let Some(engine) = runtime.rank_engine(rank) {
21978 let _main = engine.gpu.enter_main()?;
21979 engine.stream().synchronize()?;
21980 }
21981 }
21982 e.stream().synchronize()?;
21983 timer.fetch_add(
21984 start.elapsed().as_nanos() as u64,
21985 std::sync::atomic::Ordering::Relaxed,
21986 );
21987 *start = std::time::Instant::now();
21988 Ok(())
21989 }
21990 let tp = fa
21991 .step_tp_qkv
21992 .as_ref()
21993 .ok_or("Step TP decode lost its resident projections")?;
21994 let attention = tp
21995 .attention
21996 .as_ref()
21997 .ok_or("Step TP decode lost its resident attention auxiliaries")?;
21998 if !tp.runtime.native_p2p() {
21999 return Err("rank-local Step attention requires native P2P".into());
22000 }
22001 if crate::Engine::kv_fp8_on() {
22002 return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
22003 }
22004
22005 let geometry = self.step35_geom(il);
22006 let window = geometry.window.map(|window| window as usize);
22007 let ranks = tp.runtime.devices().len();
22008 let head_dim = geometry.head_dim_k as usize;
22009 let heads = geometry.n_head as usize;
22010 let kv_heads = geometry.n_head_kv as usize;
22011 if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
22012 return Err(format!(
22013 "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
22014 )
22015 .into());
22016 }
22017 let local_heads = heads / ranks;
22018 let local_kv_heads = kv_heads / ranks;
22019 let local_kv_dim = local_kv_heads * head_dim;
22020 let max_ctx = cache.max_ctx;
22021
22022 let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
22023
22024 let base_len = cache.kv[il]
22025 .as_ref()
22026 .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
22027 .len;
22028 let distributed = cache.tp_kv[il]
22029 .as_ref()
22030 .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
22031 if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
22032 return Err(format!(
22033 "Step TP layer {il} cache lengths diverged before decode: \
22034 local={base_len} distributed={}/{}",
22035 distributed.committed_len(),
22036 distributed.staged_len()
22037 )
22038 .into());
22039 }
22040
22041 let mut lap_start = timing.then(std::time::Instant::now);
22042 let positions = e.dtoh_i32(pos_d)?;
22043 if positions.len() != 1 {
22044 return Err(format!(
22045 "rank-local Step decode requires one position, got {}",
22046 positions.len()
22047 )
22048 .into());
22049 }
22050 lap(&tp.runtime, e, &T_POS, &mut lap_start)?;
22051 let (q_raw, k_raw, v_raw, input_path) = if let Some(decode_input) =
22052 attention.decode_input.as_ref()
22053 {
22054 let mut decode_input = decode_input
22055 .lock()
22056 .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
22057 e.stream().synchronize()?;
22061 tp.runtime
22062 .refresh_replicated_device_rows_from_root(&mut decode_input, h)?;
22063 let q_raw = tp
22064 .runtime
22065 .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &decode_input)?;
22066 let k_raw = tp
22067 .runtime
22068 .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &decode_input)?;
22069 let v_raw = tp
22070 .runtime
22071 .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &decode_input)?;
22072 (q_raw, k_raw, v_raw, "root-device-replicated")
22073 } else {
22074 let activation = e.dtoh(h)?;
22075 let q_raw =
22076 tp.runtime
22077 .bf16_column_parallel_resident_device_shards(&tp.q, &activation, 1)?;
22078 let k_raw =
22079 tp.runtime
22080 .bf16_column_parallel_resident_device_shards(&tp.k, &activation, 1)?;
22081 let v_raw =
22082 tp.runtime
22083 .bf16_column_parallel_resident_device_shards(&tp.v, &activation, 1)?;
22084 (q_raw, k_raw, v_raw, "host-replicated")
22085 };
22086 lap(&tp.runtime, e, &T_QKV, &mut lap_start)?;
22087 let mut q = Vec::with_capacity(ranks);
22088 let mut k = Vec::with_capacity(ranks);
22089 for rank in 0..ranks {
22090 let engine = tp
22091 .runtime
22092 .rank_engine(rank)
22093 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
22094 let _main = engine.gpu.enter_main()?;
22095 let mut q_rank = engine.uninit(local_heads * head_dim)?;
22096 engine.rms_norm(
22097 &q_raw[rank],
22098 &attention.q_norm[rank],
22099 &mut q_rank,
22100 head_dim,
22101 local_heads,
22102 self.cfg.rms_eps,
22103 )?;
22104 let mut k_rank = engine.uninit(local_kv_dim)?;
22105 engine.rms_norm(
22106 &k_raw[rank],
22107 &attention.k_norm[rank],
22108 &mut k_rank,
22109 head_dim,
22110 local_kv_heads,
22111 self.cfg.rms_eps,
22112 )?;
22113 let position = engine.htod_i32(&positions)?;
22114 let rope_freqs = if geometry.rope_factors {
22115 self.step35_aux
22116 .as_ref()
22117 .and_then(|aux| aux.rope_freqs(engine))
22118 } else {
22119 None
22120 };
22121 engine.rope_neox2(
22122 &mut q_rank,
22123 &mut k_rank,
22124 &position,
22125 head_dim,
22126 geometry.n_rot as usize,
22127 local_heads,
22128 local_kv_heads,
22129 1,
22130 geometry.rope_base,
22131 1.0,
22132 rope_freqs,
22133 )?;
22134 q.push(q_rank);
22135 k.push(k_rank);
22136 }
22137 lap(&tp.runtime, e, &T_NORMROPE, &mut lap_start)?;
22138
22139 let gate_weight = fa
22140 .attn_gate
22141 .as_ref()
22142 .ok_or("step35 layer is missing attn_gate.weight")?;
22143 let gate = e.matmul(gate_weight, h, 1)?;
22144 let gate = e.dtoh(&gate)?;
22145 if gate.len() != heads {
22146 return Err(format!("Step TP layer {il} gate output {} != {heads}", gate.len()).into());
22147 }
22148 lap(&tp.runtime, e, &T_GATE, &mut lap_start)?;
22149
22150 let transaction = cache.tp_kv[il]
22151 .as_mut()
22152 .expect("distributed cache checked above")
22153 .begin_transaction()?;
22154 if let Err(error) = tp.runtime.append_tp_kv_transaction(
22155 cache.tp_kv[il]
22156 .as_mut()
22157 .expect("distributed cache checked above"),
22158 transaction,
22159 &k,
22160 &v_raw,
22161 1,
22162 ) {
22163 let _ = tp.runtime.rollback_tp_kv_transaction(
22164 cache.tp_kv[il]
22165 .as_mut()
22166 .expect("distributed cache checked above"),
22167 transaction,
22168 );
22169 return Err(error);
22170 }
22171 lap(&tp.runtime, e, &T_APPEND, &mut lap_start)?;
22172
22173 let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22174 let distributed = cache.tp_kv[il]
22175 .as_ref()
22176 .expect("distributed cache checked above");
22177 let staged_len = distributed.staged_len();
22178 let view_start = window
22179 .map(|window| staged_len.saturating_sub(window))
22180 .unwrap_or(0);
22181 let physical = distributed.physical_range(view_start, staged_len)?;
22182 let t_kv = staged_len - view_start;
22183 let mut gated = Vec::with_capacity(ranks);
22184 #[allow(clippy::needless_range_loop)]
22185 for rank in 0..ranks {
22187 let engine = tp
22188 .runtime
22189 .rank_engine(rank)
22190 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
22191 let _main = engine.gpu.enter_main()?;
22192 let rank_cache = distributed
22193 .rank(rank)
22194 .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
22195 let k_view = engine.view_u8_range(
22196 rank_cache.k(),
22197 physical.start * distributed.k_tok_bytes(),
22198 physical.end * distributed.k_tok_bytes(),
22199 );
22200 let v_view = engine.view_u8_range(
22201 rank_cache.v(),
22202 physical.start * distributed.v_tok_bytes(),
22203 physical.end * distributed.v_tok_bytes(),
22204 );
22205 let mut attention_out = engine.uninit(local_heads * head_dim)?;
22206 engine.fa_decode_kvmod(
22207 &q[rank],
22208 &k_view,
22209 &v_view,
22210 &mut attention_out,
22211 head_dim,
22212 local_heads,
22213 local_kv_heads,
22214 t_kv,
22215 geometry.attention_scale(),
22216 distributed.k_tok_bytes(),
22217 distributed.v_tok_bytes(),
22218 false,
22219 )?;
22220 let gate_start = rank * local_heads;
22221 let gate_rank = engine.htod(&gate[gate_start..gate_start + local_heads])?;
22222 let mut gated_rank = engine.uninit(local_heads * head_dim)?;
22223 engine.attn_head_gate(
22224 &attention_out,
22225 &gate_rank,
22226 &mut gated_rank,
22227 None,
22228 head_dim,
22229 local_heads,
22230 1,
22231 )?;
22232 gated.push(gated_rank);
22233 }
22234 lap(&tp.runtime, e, &T_ATTN, &mut lap_start)?;
22235
22236 let gathered =
22237 tp.runtime
22238 .gather_native_column_shards(&gated, 1, local_heads * head_dim)?;
22239 let output = tp
22240 .runtime
22241 .step_bf16_row_parallel_resident_native(&tp.o, &gathered, 1)?;
22242 let output = e.htod(&output)?;
22243 lap(&tp.runtime, e, &T_OPROJ, &mut lap_start)?;
22244
22245 let k_shadow = tp
22246 .runtime
22247 .gather_native_column_shards(&k, 1, local_kv_dim)?;
22248 let v_shadow = tp
22249 .runtime
22250 .gather_native_column_shards(&v_raw, 1, local_kv_dim)?;
22251 let k_shadow = e.htod(&k_shadow)?;
22252 let v_shadow = e.htod(&v_shadow)?;
22253 let local = cache.kv[il]
22254 .as_mut()
22255 .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
22256 if local.len != base_len || base_len + 1 > max_ctx {
22257 return Err(format!(
22258 "Step TP layer {il} local cache changed during decode: \
22259 len={} base={base_len} max={max_ctx}",
22260 local.len
22261 )
22262 .into());
22263 }
22264 let retain_from = window
22265 .map(|window| {
22266 let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
22267 let rollback_retain =
22268 base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
22269 staged_retain.min(rollback_retain)
22270 })
22271 .unwrap_or(0);
22272 let write_row = e.prepare_kv_append(local, retain_from, 1)?;
22273 e.append_kv_quantized(
22274 &k_shadow,
22275 &v_shadow,
22276 &mut local.k,
22277 &mut local.v,
22278 write_row,
22279 local.kv_dim_k,
22280 local.kv_dim_v,
22281 local.k_tok_bytes,
22282 local.v_tok_bytes,
22283 false,
22284 )?;
22285 local.len = base_len + 1;
22286 e.set_i32_one(&mut local.len_d, local.len as i32)?;
22287 Ok(output)
22288 })();
22289
22290 let output = match staged {
22291 Ok(output) => output,
22292 Err(error) => {
22293 let _ = tp.runtime.rollback_tp_kv_transaction(
22294 cache.tp_kv[il]
22295 .as_mut()
22296 .expect("distributed cache checked above"),
22297 transaction,
22298 );
22299 if let Some(local) = cache.kv[il].as_mut() {
22300 local.len = base_len;
22301 let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
22302 }
22303 return Err(error);
22304 }
22305 };
22306 if let Err(error) = tp.runtime.commit_tp_kv_transaction(
22307 cache.tp_kv[il]
22308 .as_mut()
22309 .expect("distributed cache checked above"),
22310 transaction,
22311 1,
22312 ) {
22313 let _ = tp.runtime.rollback_tp_kv_transaction(
22314 cache.tp_kv[il]
22315 .as_mut()
22316 .expect("distributed cache checked above"),
22317 transaction,
22318 );
22319 let local = cache.kv[il].as_mut().expect("local cache checked above");
22320 local.len = base_len;
22321 e.set_i32_one(&mut local.len_d, base_len as i32)?;
22322 return Err(error);
22323 }
22324
22325 let committed = cache.tp_kv[il]
22326 .as_ref()
22327 .expect("distributed cache checked above")
22328 .committed_len();
22329 let local_len = cache.kv[il]
22330 .as_ref()
22331 .expect("local cache checked above")
22332 .len;
22333 if committed != local_len {
22334 return Err(format!(
22335 "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
22336 )
22337 .into());
22338 }
22339 lap(&tp.runtime, e, &T_SHADOW, &mut lap_start)?;
22340 if timing {
22341 use std::sync::atomic::Ordering;
22342 let calls = T_PHASE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
22343 if calls.is_multiple_of(430) {
22344 let avg = |t: &std::sync::atomic::AtomicU64| {
22345 t.load(Ordering::Relaxed) as f64 / calls as f64 / 1.0e3
22346 };
22347 eprintln!(
22348 "[step-tp-attn-phase] calls={calls} avg_us pos={:.1} qkv={:.1} \
22349 normrope={:.1} gate={:.1} append={:.1} attn={:.1} oproj={:.1} shadow={:.1}",
22350 avg(&T_POS),
22351 avg(&T_QKV),
22352 avg(&T_NORMROPE),
22353 avg(&T_GATE),
22354 avg(&T_APPEND),
22355 avg(&T_ATTN),
22356 avg(&T_OPROJ),
22357 avg(&T_SHADOW),
22358 );
22359 }
22360 }
22361 eprintln!(
22362 "[step-tp-attn] execute layer={} devices={:?} tokens=1 \
22363 qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
22364 kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
22365 attention_scope={} input_path={} kv_physical_rows={} \
22366 gate_tensor_parallel=false gate_shards=host-canonical o_tensor_parallel=true \
22367 local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
22368 bulk_p2p={} output=root-readback performance_claim=false",
22369 tp.layer,
22370 tp.devices,
22371 hydrated,
22372 if window.is_some() {
22373 "rank-local-swa-ring"
22374 } else {
22375 "rank-local-global"
22376 },
22377 input_path,
22378 cache.tp_kv[il]
22379 .as_ref()
22380 .expect("distributed cache checked above")
22381 .physical_capacity(),
22382 tp.runtime.transport_label(),
22383 tp.runtime.bulk_p2p(),
22384 );
22385 Ok(output)
22386 }
22387
22388 #[allow(clippy::too_many_arguments)]
22395 pub(crate) fn step35_verify_qkv_precompute(
22400 &self,
22401 e: &Engine,
22402 il: usize,
22403 h_t: &CudaSlice<f32>,
22404 t: usize,
22405 ) -> Result<bool, Box<dyn std::error::Error>> {
22406 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22407 return Ok(false);
22408 };
22409 let Some(tp) = fa.step_tp_qkv.as_ref() else {
22410 return Ok(false);
22411 };
22412 let Some(attention) = tp.attention.as_ref() else {
22413 return Ok(false);
22414 };
22415 if !tp.runtime.native_p2p() || !crate::tp::step_tp_qkv_fused_enabled()? {
22416 return Ok(false);
22417 }
22418 let geometry = self.step35_geom(il);
22419 let heads = geometry.n_head as usize;
22420 let ws_index = tp
22421 .runtime
22422 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22423 let gate_shards = attention
22424 .gate_shards_bf16
22425 .as_deref()
22426 .map(crate::tp::StepTpGateShards::Bf16);
22427 tp.runtime.decode_v2_input_qkv_tcol(
22428 ws_index,
22429 e,
22430 h_t,
22431 t,
22432 &tp.q,
22433 &tp.k,
22434 &tp.v,
22435 gate_shards,
22436 )?;
22437 Ok(true)
22438 }
22439
22440 pub(crate) fn step35_verify_oproj_tcol(
22445 &self,
22446 e: &Engine,
22447 il: usize,
22448 t: usize,
22449 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22450 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22451 return Err("tcol o_proj join expects full attention".into());
22452 };
22453 let tp = fa
22454 .step_tp_qkv
22455 .as_ref()
22456 .ok_or("tcol o_proj join lost its resident projections")?;
22457 let heads = self.step35_geom(il).n_head as usize;
22458 let ws_index = tp
22459 .runtime
22460 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22461 tp.runtime.decode_v2_oproj_tcol(ws_index, e, &tp.o, t)
22462 }
22463
22464 #[allow(dead_code)] pub(crate) fn step35_spec_fa2_precheck(
22472 &self,
22473 cache: &Cache,
22474 il: usize,
22475 pos0: usize,
22476 ) -> Result<bool, Box<dyn std::error::Error>> {
22477 fn nope(clause: &str, il: usize, pos0: usize) -> bool {
22480 static DBG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22481 static SEEN: std::sync::Mutex<Vec<&'static str>> = std::sync::Mutex::new(Vec::new());
22482 if *DBG.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1")) {
22483 let mut seen = SEEN.lock().unwrap();
22484 if !seen.contains(&clause) {
22485 seen.push(Box::leak(clause.to_string().into_boxed_str()));
22487 eprintln!("[spec-fa2] precheck FAIL clause={clause} il={il} pos0={pos0}");
22488 }
22489 }
22490 false
22491 }
22492 static ONLY: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
22494 if let Some(only) =
22495 ONLY.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_LAYER").ok()?.parse().ok())
22496 && *only != il
22497 {
22498 return Ok(false);
22499 }
22500 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22501 return Ok(nope("mixer", il, pos0));
22502 };
22503 let Some(tp) = fa.step_tp_qkv.as_ref() else {
22504 return Ok(nope("step_tp", il, pos0));
22505 };
22506 let Some(attention) = tp.attention.as_ref() else {
22507 return Ok(nope("attention", il, pos0));
22508 };
22509 if !tp.runtime.native_p2p()
22510 || crate::Engine::kv_fp8_on()
22511 || !crate::tp::step_tp_dcw_enabled()?
22512 || !crate::tp::step_tp_qkv_fused_enabled()?
22513 || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22514 {
22515 return Ok(nope("runtime-doors", il, pos0));
22516 }
22517 let geometry = self.step35_geom(il);
22518 let head_dim = geometry.head_dim_k as usize;
22519 if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22520 return Ok(nope("fa-class", il, pos0));
22521 }
22522 let Some(distributed) = cache.tp_kv[il].as_ref() else {
22523 return Ok(nope("tp-kv", il, pos0));
22524 };
22525 if distributed.staged_len() != pos0 {
22526 return Ok(nope("staged-len", il, pos0));
22527 }
22528 let (_, would_rebase) = distributed.peek_append_ring(2)?;
22531 if would_rebase {
22532 return Ok(nope("rebase", il, pos0));
22533 }
22534 let window = geometry.window.map(|w| w as usize);
22535 if let Some(w) = window
22543 && pos0 + 2 > w
22544 {
22545 return Ok(nope("swa-capped", il, pos0));
22546 }
22547 let (t0, t1) = (pos0 + 1, pos0 + 2);
22551 if t0 < 96 {
22552 return Ok(nope("dcw-floor", il, pos0));
22553 }
22554 if std::env::var("MEMRA_NO_FA_VEC").is_ok() || t0 < crate::fa_vec_min_tkv() {
22555 return Ok(nope("vec-floor", il, pos0));
22556 }
22557 let ranks = tp.runtime.devices().len();
22563 let local_kv_heads = (geometry.n_head_kv as usize / ranks).max(1);
22564 let sp0 = crate::fa_split_keys_pub(t0, local_kv_heads);
22565 let sp1 = crate::fa_split_keys_pub(t1, local_kv_heads);
22566 if sp0 != sp1 {
22567 return Ok(nope("partition-sp", il, pos0));
22568 }
22569 let (ns0, ns1) = (t0.div_ceil(sp0), t1.div_ceil(sp1));
22570 if ns0 != ns1 {
22571 return Ok(nope("partition-ns", il, pos0));
22572 }
22573 if t0.div_ceil(ns0) != t1.div_ceil(ns1) {
22574 return Ok(nope("partition-per", il, pos0));
22575 }
22576 Ok(true)
22577 }
22578
22579 pub(crate) fn step35_fa_rows_precheck(
22584 &self,
22585 cache: &Cache,
22586 il: usize,
22587 pos0: usize,
22588 t: usize,
22589 ) -> Result<bool, Box<dyn std::error::Error>> {
22590 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22591 return Ok(false);
22592 };
22593 let Some(tp) = fa.step_tp_qkv.as_ref() else {
22594 return Ok(false);
22595 };
22596 let Some(attention) = tp.attention.as_ref() else {
22597 return Ok(false);
22598 };
22599 if !tp.runtime.native_p2p()
22600 || crate::Engine::kv_fp8_on()
22601 || !crate::tp::step_tp_dcw_enabled()?
22602 || !crate::tp::step_tp_qkv_fused_enabled()?
22603 || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22604 {
22605 return Ok(false);
22606 }
22607 let geometry = self.step35_geom(il);
22608 let head_dim = geometry.head_dim_k as usize;
22609 if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22610 return Ok(false);
22611 }
22612 if crate::fa_sm_count() < 128
22613 || std::env::var("MEMRA_FA_SPLIT").is_ok()
22614 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
22615 || std::env::var("MEMRA_FA_SP16").is_ok()
22616 || std::env::var("MEMRA_NO_FA_VEC").is_ok()
22617 {
22618 return Ok(false);
22619 }
22620 let Some(distributed) = cache.tp_kv[il].as_ref() else {
22621 return Ok(false);
22622 };
22623 if distributed.staged_len() != pos0 {
22624 return Ok(false);
22625 }
22626 let (_, would_rebase) = distributed.peek_append_ring(t)?;
22627 if would_rebase {
22628 return Ok(false);
22629 }
22630 let window = geometry.window.map(|w| w as usize);
22633 let t0 = window.map(|w| (pos0 + 1).min(w)).unwrap_or(pos0 + 1);
22634 if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
22635 return Ok(false);
22636 }
22637 Ok(true)
22638 }
22639
22640 pub(crate) fn step35_verify_fa_rows_join(
22644 &self,
22645 e: &Engine,
22646 il: usize,
22647 cache: &Cache,
22648 pos0: usize,
22649 t: usize,
22650 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22651 use cudarc::driver::DevicePtr;
22652 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22653 return Err("fa rows join expects full attention".into());
22654 };
22655 let tp = fa
22656 .step_tp_qkv
22657 .as_ref()
22658 .ok_or("fa rows join lost its resident projections")?;
22659 let geometry = self.step35_geom(il);
22660 let heads = geometry.n_head as usize;
22661 let head_dim = geometry.head_dim_k as usize;
22662 let window = geometry.window.map(|w| w as usize);
22663 let distributed = cache.tp_kv[il]
22664 .as_ref()
22665 .ok_or("fa rows join lost its distributed KV cache")?;
22666 let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
22667 let ladder = |t_kv: usize| -> usize {
22669 if t_kv <= 2048 {
22670 16
22671 } else if t_kv <= 16384 {
22672 64
22673 } else {
22674 128
22675 }
22676 };
22677 let mut max_ns = 1usize;
22678 for r in 0..t {
22679 let t_kv = window
22680 .map(|w| (pos0 + r + 1).min(w))
22681 .unwrap_or(pos0 + r + 1);
22682 max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
22683 }
22684 let ranks = tp.runtime.devices().len();
22689 let mut tables = Vec::with_capacity(ranks);
22690 for rank in 0..ranks {
22691 let engine = tp
22692 .runtime
22693 .rank_engine(rank)
22694 .ok_or("fa rows join lost a rank engine")?;
22695 let rank_cache = distributed
22696 .rank(rank)
22697 .ok_or("fa rows join lost a KV cache rank")?;
22698 let _main = engine.gpu.enter_main()?;
22699 let s = engine.stream();
22700 let (kp, _g0) = rank_cache.k().device_ptr(&s);
22701 let (vp, _g1) = rank_cache.v().device_ptr(&s);
22702 let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
22703 let bp = match rank_cache.base_d() {
22704 Some(b) => {
22705 let (p, _g) = b.device_ptr(&s);
22706 p
22707 }
22708 None => 0u64,
22709 };
22710 let mut host = Vec::with_capacity(t * 6);
22711 for r in 0..t {
22712 host.extend_from_slice(&[kp, vp, lp, bp, 0u64, (t - 1 - r) as u64]);
22713 }
22714 tables.push(engine.stream().clone_htod(&host)?);
22715 }
22716 let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
22717 let ws_index = tp
22718 .runtime
22719 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22720 tp.runtime.decode_v2_fa_rows_join(
22721 ws_index,
22722 e,
22723 &tp.o,
22724 &tabs,
22725 t,
22726 head_dim,
22727 window.unwrap_or(0),
22728 max_ns,
22729 geometry.attention_scale(),
22730 k_tok_bytes,
22731 v_tok_bytes,
22732 )
22733 }
22734
22735 pub(crate) fn step35_batch_fa_rows_precheck(
22739 &self,
22740 caches: &[&mut Cache],
22741 row_to_cache: impl Fn(usize) -> usize,
22742 positions: &[i32],
22743 il: usize,
22744 ) -> Result<bool, Box<dyn std::error::Error>> {
22745 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22746 return Ok(false);
22747 };
22748 let Some(tp) = fa.step_tp_qkv.as_ref() else {
22749 return Ok(false);
22750 };
22751 let Some(attention) = tp.attention.as_ref() else {
22752 return Ok(false);
22753 };
22754 if !tp.runtime.native_p2p()
22755 || crate::Engine::kv_fp8_on()
22756 || !crate::tp::step_tp_dcw_enabled()?
22757 || !crate::tp::step_tp_qkv_fused_enabled()?
22758 || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22759 {
22760 return Ok(false);
22761 }
22762 let geometry = self.step35_geom(il);
22763 let head_dim = geometry.head_dim_k as usize;
22764 if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22765 return Ok(false);
22766 }
22767 if crate::fa_sm_count() < 128
22768 || std::env::var("MEMRA_FA_SPLIT").is_ok()
22769 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
22770 || std::env::var("MEMRA_FA_SP16").is_ok()
22771 || std::env::var("MEMRA_NO_FA_VEC").is_ok()
22772 {
22773 return Ok(false);
22774 }
22775 let window = geometry.window.map(|w| w as usize);
22776 for (r, &pos) in positions.iter().enumerate() {
22777 let cache = &caches[row_to_cache(r)];
22778 let Some(distributed) = cache.tp_kv[il].as_ref() else {
22779 return Ok(false);
22780 };
22781 if distributed.staged_len() != pos as usize {
22782 return Ok(false);
22783 }
22784 if distributed.peek_append_ring(1)?.1 {
22785 return Ok(false);
22786 }
22787 let t0 = window
22788 .map(|w| (pos as usize + 1).min(w))
22789 .unwrap_or(pos as usize + 1);
22790 if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
22791 return Ok(false);
22792 }
22793 }
22794 Ok(true)
22795 }
22796
22797 pub(crate) fn step35_verify_rope_fa_pass(
22803 &self,
22804 e: &Engine,
22805 il: usize,
22806 cache: &Cache,
22807 pos0: usize,
22808 t: usize,
22809 stage_pos: bool,
22810 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22811 use cudarc::driver::DevicePtr;
22812 if !crate::tp::fuse_rope_append_on() {
22813 return Ok(None);
22814 }
22815 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22816 return Ok(None);
22817 };
22818 let Some(tp) = fa.step_tp_qkv.as_ref() else {
22819 return Ok(None);
22820 };
22821 let Some(attention) = tp.attention.as_ref() else {
22822 return Ok(None);
22823 };
22824 let geometry = self.step35_geom(il);
22825 let head_dim = geometry.head_dim_k as usize;
22826 if head_dim != 128 {
22827 return Ok(None);
22828 }
22829 let heads = geometry.n_head as usize;
22830 let window = geometry.window.map(|w| w as usize);
22831 let ranks = tp.runtime.devices().len();
22832 let Some(distributed) = cache.tp_kv[il].as_ref() else {
22833 return Ok(None);
22834 };
22835 if distributed.kv_dim_k() != distributed.kv_dim_v() {
22836 return Ok(None);
22837 }
22838 {
22839 let rank0 = distributed.rank(0).ok_or("verify rope pass lost rank 0")?;
22840 if rank0.base_d().is_none()
22841 && distributed.staged_len() + t > distributed.physical_capacity()
22842 {
22843 return Ok(None);
22844 }
22845 }
22846 let mut rope_freqs = Vec::with_capacity(ranks);
22847 for rank in 0..ranks {
22848 let engine = tp
22849 .runtime
22850 .rank_engine(rank)
22851 .ok_or("verify rope pass lost a rank engine")?;
22852 rope_freqs.push(if geometry.rope_factors {
22853 match self
22854 .step35_aux
22855 .as_ref()
22856 .and_then(|aux| aux.rope_freqs(engine))
22857 {
22858 Some(f) => Some(f),
22859 None => return Ok(None),
22860 }
22861 } else {
22862 None
22863 });
22864 }
22865 let ladder = |t_kv: usize| -> usize {
22866 if t_kv <= 2048 {
22867 16
22868 } else if t_kv <= 16384 {
22869 64
22870 } else {
22871 128
22872 }
22873 };
22874 let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
22875 let mut max_ns = 1usize;
22876 let mut positions = Vec::with_capacity(t);
22877 for r in 0..t {
22878 positions.push((pos0 + r) as i32);
22879 let t_kv = window
22880 .map(|w| (pos0 + r + 1).min(w))
22881 .unwrap_or(pos0 + r + 1);
22882 max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
22883 }
22884 let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
22885 let mut tab_keys = vec![0u64; ranks];
22886 for rank in 0..ranks {
22887 let engine = tp
22888 .runtime
22889 .rank_engine(rank)
22890 .ok_or("verify rope pass lost a rank engine")?;
22891 let rank_cache = distributed
22892 .rank(rank)
22893 .ok_or("verify rope pass lost a KV cache rank")?;
22894 let _main = engine.gpu.enter_main()?;
22895 let s = engine.stream();
22896 let (kp, _g0) = rank_cache.k().device_ptr(&s);
22897 let (vp, _g1) = rank_cache.v().device_ptr(&s);
22898 let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
22899 let bp = match rank_cache.base_d() {
22900 Some(b) => {
22901 let (p, _g) = b.device_ptr(&s);
22902 p
22903 }
22904 None => 0u64,
22905 };
22906 tab_keys[rank] = kp
22907 .rotate_left(17)
22908 .wrapping_add(bp)
22909 .wrapping_add((il as u64) << 32)
22910 .wrapping_add(t as u64)
22911 .wrapping_add(1 << 63);
22912 for _r in 0..t {
22913 session_parts[rank].push([kp, vp, lp, bp]);
22914 }
22915 }
22916 let ws_index = tp
22917 .runtime
22918 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22919 tp.runtime
22920 .decode_v2_rope_fa_rows(
22921 ws_index,
22922 e,
22923 &tp.o,
22924 &session_parts,
22925 &tab_keys,
22926 &positions,
22927 stage_pos,
22928 true,
22929 &attention.q_norm,
22930 &attention.k_norm,
22931 &rope_freqs,
22932 t,
22933 head_dim,
22934 geometry.n_rot as usize,
22935 window.unwrap_or(0),
22936 max_ns,
22937 geometry.attention_scale(),
22938 k_tok_bytes,
22939 v_tok_bytes,
22940 self.cfg.rms_eps,
22941 geometry.rope_base,
22942 )
22943 .map(Some)
22944 }
22945
22946 #[allow(clippy::too_many_arguments)]
22951 pub(crate) fn step35_batch_rope_fa_pass(
22952 &self,
22953 e: &Engine,
22954 il: usize,
22955 caches: &[&mut Cache],
22956 row_to_cache: impl Fn(usize) -> usize,
22957 positions: &[i32],
22958 t: usize,
22959 stage_pos: bool,
22960 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22961 use cudarc::driver::DevicePtr;
22962 if !crate::tp::fuse_rope_append_on() {
22963 return Ok(None);
22964 }
22965 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22966 return Ok(None);
22967 };
22968 let Some(tp) = fa.step_tp_qkv.as_ref() else {
22969 return Ok(None);
22970 };
22971 let Some(attention) = tp.attention.as_ref() else {
22972 return Ok(None);
22973 };
22974 let geometry = self.step35_geom(il);
22975 let head_dim = geometry.head_dim_k as usize;
22976 if head_dim != 128 {
22977 return Ok(None);
22978 }
22979 let heads = geometry.n_head as usize;
22980 let window = geometry.window.map(|w| w as usize);
22981 let ranks = tp.runtime.devices().len();
22982 for r in 0..t {
22985 let cache = &caches[row_to_cache(r)];
22986 let Some(distributed) = cache.tp_kv[il].as_ref() else {
22987 return Ok(None);
22988 };
22989 if distributed.kv_dim_k() != distributed.kv_dim_v() {
22990 return Ok(None);
22991 }
22992 let rank0 = distributed.rank(0).ok_or("rope fa pass lost rank 0")?;
22993 if rank0.base_d().is_none()
22994 && distributed.staged_len() + t > distributed.physical_capacity()
22995 {
22996 return Ok(None);
22997 }
22998 }
22999 let mut rope_freqs = Vec::with_capacity(ranks);
23000 for rank in 0..ranks {
23001 let engine = tp
23002 .runtime
23003 .rank_engine(rank)
23004 .ok_or("rope fa pass lost a rank engine")?;
23005 rope_freqs.push(if geometry.rope_factors {
23006 match self
23007 .step35_aux
23008 .as_ref()
23009 .and_then(|aux| aux.rope_freqs(engine))
23010 {
23011 Some(f) => Some(f),
23012 None => return Ok(None),
23013 }
23014 } else {
23015 None
23016 });
23017 }
23018 let ladder = |t_kv: usize| -> usize {
23019 if t_kv <= 2048 {
23020 16
23021 } else if t_kv <= 16384 {
23022 64
23023 } else {
23024 128
23025 }
23026 };
23027 let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
23028 let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
23029 let mut tab_keys = vec![0u64; ranks];
23030 for (r, &pos) in positions.iter().enumerate().take(t) {
23031 let cache = &caches[row_to_cache(r)];
23032 let distributed = cache.tp_kv[il]
23033 .as_ref()
23034 .ok_or("rope fa pass lost a distributed KV cache")?;
23035 (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23036 let t_kv = window
23037 .map(|w| (pos as usize + 1).min(w))
23038 .unwrap_or(pos as usize + 1);
23039 max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23040 for rank in 0..ranks {
23041 let engine = tp
23042 .runtime
23043 .rank_engine(rank)
23044 .ok_or("rope fa pass lost a rank engine")?;
23045 let rank_cache = distributed
23046 .rank(rank)
23047 .ok_or("rope fa pass lost a KV cache rank")?;
23048 let _main = engine.gpu.enter_main()?;
23049 let s = engine.stream();
23050 let (kp, _g0) = rank_cache.k().device_ptr(&s);
23051 let (vp, _g1) = rank_cache.v().device_ptr(&s);
23052 let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23053 let bp = match rank_cache.base_d() {
23054 Some(b) => {
23055 let (p, _g) = b.device_ptr(&s);
23056 p
23057 }
23058 None => 0u64,
23059 };
23060 tab_keys[rank] = tab_keys[rank]
23061 .rotate_left(9)
23062 .wrapping_add(kp)
23063 .wrapping_add(bp)
23064 .wrapping_add(il as u64);
23065 session_parts[rank].push([kp, vp, lp, bp]);
23066 }
23067 }
23068 let ws_index = tp
23069 .runtime
23070 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23071 tp.runtime
23072 .decode_v2_rope_fa_rows(
23073 ws_index,
23074 e,
23075 &tp.o,
23076 &session_parts,
23077 &tab_keys,
23078 positions,
23079 stage_pos,
23080 false,
23081 &attention.q_norm,
23082 &attention.k_norm,
23083 &rope_freqs,
23084 t,
23085 head_dim,
23086 geometry.n_rot as usize,
23087 window.unwrap_or(0),
23088 max_ns,
23089 geometry.attention_scale(),
23090 k_tok_bytes,
23091 v_tok_bytes,
23092 self.cfg.rms_eps,
23093 geometry.rope_base,
23094 )
23095 .map(Some)
23096 }
23097
23098 #[allow(clippy::too_many_arguments)]
23102 pub(crate) fn step35_batch_fa_rows_join(
23103 &self,
23104 e: &Engine,
23105 il: usize,
23106 caches: &[&mut Cache],
23107 row_to_cache: impl Fn(usize) -> usize,
23108 positions: &[i32],
23109 t: usize,
23110 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23111 use cudarc::driver::DevicePtr;
23112 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23113 return Err("batch fa rows join expects full attention".into());
23114 };
23115 let tp = fa
23116 .step_tp_qkv
23117 .as_ref()
23118 .ok_or("batch fa rows join lost its resident projections")?;
23119 let geometry = self.step35_geom(il);
23120 let heads = geometry.n_head as usize;
23121 let head_dim = geometry.head_dim_k as usize;
23122 let window = geometry.window.map(|w| w as usize);
23123 let ladder = |t_kv: usize| -> usize {
23124 if t_kv <= 2048 {
23125 16
23126 } else if t_kv <= 16384 {
23127 64
23128 } else {
23129 128
23130 }
23131 };
23132 let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
23133 for (r, &pos) in positions.iter().enumerate() {
23134 let cache = &caches[row_to_cache(r)];
23135 let distributed = cache.tp_kv[il]
23136 .as_ref()
23137 .ok_or("batch fa rows join lost a distributed KV cache")?;
23138 (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23139 let t_kv = window
23140 .map(|w| (pos as usize + 1).min(w))
23141 .unwrap_or(pos as usize + 1);
23142 max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23143 }
23144 let ranks = tp.runtime.devices().len();
23148 let mut tables = Vec::with_capacity(ranks);
23149 for rank in 0..ranks {
23150 let engine = tp
23151 .runtime
23152 .rank_engine(rank)
23153 .ok_or("batch fa rows join lost a rank engine")?;
23154 let _main = engine.gpu.enter_main()?;
23155 let s = engine.stream();
23156 let mut host = Vec::with_capacity(t * 6);
23157 for r in 0..t {
23158 let cache = &caches[row_to_cache(r)];
23159 let distributed = cache.tp_kv[il]
23160 .as_ref()
23161 .ok_or("batch fa rows join lost a distributed KV cache")?;
23162 let rank_cache = distributed
23163 .rank(rank)
23164 .ok_or("batch fa rows join lost a KV cache rank")?;
23165 let (kp, _g0) = rank_cache.k().device_ptr(&s);
23166 let (vp, _g1) = rank_cache.v().device_ptr(&s);
23167 let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23168 let bp = match rank_cache.base_d() {
23169 Some(b) => {
23170 let (p, _g) = b.device_ptr(&s);
23171 p
23172 }
23173 None => 0u64,
23174 };
23175 host.extend_from_slice(&[kp, vp, lp, bp, 0u64, 0u64]);
23176 }
23177 tables.push(engine.stream().clone_htod(&host)?);
23178 }
23179 let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
23180 let ws_index = tp
23181 .runtime
23182 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23183 tp.runtime.decode_v2_fa_rows_join(
23184 ws_index,
23185 e,
23186 &tp.o,
23187 &tabs,
23188 t,
23189 head_dim,
23190 window.unwrap_or(0),
23191 max_ns,
23192 geometry.attention_scale(),
23193 k_tok_bytes,
23194 v_tok_bytes,
23195 )
23196 }
23197
23198 #[allow(dead_code)] pub(crate) fn step35_verify_spec_fa2_join(
23203 &self,
23204 e: &Engine,
23205 il: usize,
23206 cache: &Cache,
23207 pos0: usize,
23208 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23209 let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23210 return Err("spec fa2 join expects full attention".into());
23211 };
23212 let tp = fa
23213 .step_tp_qkv
23214 .as_ref()
23215 .ok_or("spec fa2 join lost its resident projections")?;
23216 let geometry = self.step35_geom(il);
23217 let heads = geometry.n_head as usize;
23218 let head_dim = geometry.head_dim_k as usize;
23219 let window = geometry.window.map(|w| w as usize);
23220 let bucket = window.map(|w| (pos0 + 2).min(w)).unwrap_or(pos0 + 2);
23223 let distributed = cache.tp_kv[il]
23224 .as_ref()
23225 .ok_or("spec fa2 join lost its distributed KV cache")?;
23226 let ws_index = tp
23227 .runtime
23228 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23229 tp.runtime.decode_v2_spec_fa2_join(
23230 ws_index,
23231 e,
23232 &tp.o,
23233 distributed,
23234 head_dim,
23235 window.unwrap_or(0),
23236 bucket,
23237 geometry.attention_scale(),
23238 )
23239 }
23240
23241 pub(crate) fn step35_tp_decode_attn_resident_v2(
23242 &self,
23243 e: &Engine,
23244 fa: &FullAttnLayer,
23245 il: usize,
23246 h: &CudaSlice<f32>,
23247 pos_d: &CudaSlice<i32>,
23248 cache: &mut Cache,
23249 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23250 let tp = fa
23251 .step_tp_qkv
23252 .as_ref()
23253 .ok_or("Step TP decode lost its resident projections")?;
23254 let attention = tp
23255 .attention
23256 .as_ref()
23257 .ok_or("Step TP decode lost its resident attention auxiliaries")?;
23258 if !tp.runtime.native_p2p() {
23259 return Err("rank-local Step attention requires native P2P".into());
23260 }
23261 if crate::Engine::kv_fp8_on() {
23262 return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
23263 }
23264
23265 let geometry = self.cfg.full_attention_geometry_at(il as u32);
23266 let window = geometry.window.map(|window| window as usize);
23267 let ranks = tp.runtime.devices().len();
23268 let head_dim = geometry.head_dim_k as usize;
23269 let heads = geometry.n_head as usize;
23270 let kv_heads = geometry.n_head_kv as usize;
23271 if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
23272 return Err(format!(
23273 "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
23274 )
23275 .into());
23276 }
23277 let local_heads = heads / ranks;
23278 let local_kv_heads = kv_heads / ranks;
23279 let max_ctx = cache.max_ctx;
23280
23281 let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
23282
23283 let base_len = cache.kv[il]
23284 .as_ref()
23285 .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
23286 .len;
23287 {
23288 let distributed = cache.tp_kv[il]
23289 .as_ref()
23290 .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
23291 if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
23292 return Err(format!(
23293 "Step TP layer {il} cache lengths diverged before decode: \
23294 local={base_len} distributed={}/{}",
23295 distributed.committed_len(),
23296 distributed.staged_len()
23297 )
23298 .into());
23299 }
23300 }
23301 if pos_d.len() != 1 {
23302 return Err(format!(
23303 "rank-local Step decode requires one position, got {}",
23304 pos_d.len()
23305 )
23306 .into());
23307 }
23308
23309 let decode_input = attention
23310 .decode_input
23311 .as_ref()
23312 .ok_or("Step TP decode v2 requires the replicated decode input")?;
23313 let mut decode_input = decode_input
23314 .lock()
23315 .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
23316
23317 let has_gate = fa.attn_gate.is_some();
23318 let use_gate_shards = has_gate
23323 && (attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some())
23324 && crate::tp::step_tp_qkv_fused_enabled()?;
23325 let gate_raw = if !has_gate || use_gate_shards {
23326 None
23327 } else {
23328 let gate_weight = fa
23329 .attn_gate
23330 .as_ref()
23331 .ok_or("step35 layer is missing attn_gate.weight")?;
23332 let gate_raw = e.matmul(gate_weight, h, 1)?;
23333 if gate_raw.len() != heads {
23334 return Err(format!(
23335 "Step TP layer {il} gate output {} != {heads}",
23336 gate_raw.len()
23337 )
23338 .into());
23339 }
23340 Some(gate_raw)
23341 };
23342
23343 let ws_index = tp
23344 .runtime
23345 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23346 let mut ws_guard = tp
23347 .runtime
23348 .decode_v2_workspace()
23349 .lock()
23350 .map_err(|_| "Step TP decode v2 workspace lock is poisoned")?;
23351 let ws = ws_guard
23352 .get_mut(ws_index)
23353 .ok_or("Step TP decode v2 workspace missing after ensure")?;
23354
23355 let mut rope_freqs = Vec::with_capacity(ranks);
23356 for rank in 0..ranks {
23357 let engine = tp
23358 .runtime
23359 .rank_engine(rank)
23360 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23361 rope_freqs.push(if geometry.rope_factors {
23362 self.step35_aux
23363 .as_ref()
23364 .and_then(|aux| aux.rope_freqs(engine))
23365 } else {
23366 None
23367 });
23368 }
23369 let staged_next = base_len + 1;
23376 let t_kv_eff = window
23377 .map(|window| staged_next.min(window))
23378 .unwrap_or(staged_next);
23379 let dcw = crate::tp::step_tp_dcw_enabled()?
23380 && (use_gate_shards || (!has_gate && crate::tp::step_tp_qkv_fused_enabled()?))
23381 && t_kv_eff >= 96
23382 && {
23383 let (write_row, would_rebase) = cache.tp_kv[il]
23384 .as_ref()
23385 .expect("distributed cache checked above")
23386 .peek_append_ring(1)?;
23387 if !would_rebase {
23388 let base = (base_len - write_row) as i32;
23390 let distributed = cache.tp_kv[il]
23391 .as_mut()
23392 .expect("distributed cache checked above");
23393 for rank in 0..ranks {
23394 let engine = tp.runtime.rank_engine(rank).ok_or_else(|| {
23395 format!("Step TP layer {il} has no engine for rank {rank}")
23396 })?;
23397 let _main = engine.gpu.enter_main()?;
23398 let rank_cache = distributed.rank_mut(rank).ok_or_else(|| {
23399 format!("Step TP layer {il} has no KV cache rank {rank}")
23400 })?;
23401 if rank_cache.base_d().is_none() {
23402 rank_cache.arm_base_d(engine.htod_i32(&[base])?);
23403 }
23404 }
23405 }
23406 !would_rebase
23407 };
23408 let fuse_rope = dcw
23409 && crate::tp::fuse_rope_append_on()
23410 && head_dim == 128
23411 && cache.tp_kv[il]
23412 .as_ref()
23413 .map(|d| d.kv_dim_k() == d.kv_dim_v() && d.kv_dim_k() == local_kv_heads * head_dim)
23414 .unwrap_or(false);
23415
23416 let tcol_col = crate::tp::take_verify_tcol();
23417 let fa2_col = crate::tp::take_spec_fa2_defer();
23424 tp.runtime.decode_v2_input_qkv(
23425 ws,
23426 e,
23427 h,
23428 pos_d,
23429 gate_raw.as_ref(),
23430 if !use_gate_shards {
23431 None
23432 } else if let Some(shards) = attention.gate_shards.as_deref() {
23433 Some(crate::tp::StepTpGateShards::F32(shards))
23434 } else {
23435 attention
23436 .gate_shards_bf16
23437 .as_deref()
23438 .map(crate::tp::StepTpGateShards::Bf16)
23439 },
23440 &mut decode_input,
23441 &tp.q,
23442 &tp.k,
23443 &tp.v,
23444 &attention.q_norm,
23445 &attention.k_norm,
23446 head_dim,
23447 geometry.n_rot as usize,
23448 geometry.rope_base,
23449 &rope_freqs,
23450 self.cfg.rms_eps,
23451 has_gate,
23452 fuse_rope,
23453 tcol_col,
23454 )?;
23455
23456 let transaction = cache.tp_kv[il]
23457 .as_mut()
23458 .expect("distributed cache checked above")
23459 .begin_transaction()?;
23460 let append_result = tp.runtime.append_tp_kv_transaction_inner(
23461 cache.tp_kv[il]
23462 .as_mut()
23463 .expect("distributed cache checked above"),
23464 transaction,
23465 &ws.k,
23466 &ws.v_raw,
23467 1,
23468 dcw,
23469 );
23470 if let Err(error) = append_result {
23471 let _ = tp.runtime.rollback_tp_kv_transaction(
23472 cache.tp_kv[il]
23473 .as_mut()
23474 .expect("distributed cache checked above"),
23475 transaction,
23476 );
23477 return Err(error);
23478 }
23479
23480 let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23481 let (staged_len, physical, k_tok_bytes_c, v_tok_bytes_c, capacity) = {
23484 let distributed = cache.tp_kv[il]
23485 .as_ref()
23486 .expect("distributed cache checked above");
23487 let staged_len = distributed.staged_len();
23488 let view_start = window
23489 .map(|window| staged_len.saturating_sub(window))
23490 .unwrap_or(0);
23491 (
23492 staged_len,
23493 distributed.physical_range(view_start, staged_len)?,
23494 distributed.k_tok_bytes(),
23495 distributed.v_tok_bytes(),
23496 distributed.physical_capacity(),
23497 )
23498 };
23499 let view_start = window
23500 .map(|window| staged_len.saturating_sub(window))
23501 .unwrap_or(0);
23502 let t_kv = staged_len - view_start;
23503 for rank in 0..ranks {
23504 let engine = tp
23505 .runtime
23506 .rank_engine(rank)
23507 .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23508 let _main = engine.gpu.enter_main()?;
23509 if dcw {
23510 {
23514 let distributed_mut = cache.tp_kv[il]
23515 .as_mut()
23516 .expect("distributed cache checked above");
23517 let (kv_dim_k, kv_dim_v) =
23518 (distributed_mut.kv_dim_k(), distributed_mut.kv_dim_v());
23519 let (k_tok_bytes, v_tok_bytes) =
23520 (distributed_mut.k_tok_bytes(), distributed_mut.v_tok_bytes());
23521 let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
23522 format!("Step TP layer {il} has no KV cache rank {rank}")
23523 })?;
23524 let (k_plane, v_plane, len_d, base_d) =
23525 rank_cache.planes_and_counters_mut();
23526 if fuse_rope {
23527 let same_dev = engine.ctx().ordinal() == e.ctx().ordinal();
23530 let crate::tp::StepTpDecodeV2Ws {
23531 q_raw,
23532 k_raw,
23533 v_raw,
23534 q,
23535 k,
23536 pos,
23537 pos_stage,
23538 fuse_ctr,
23539 ..
23540 } = &mut *ws;
23541 let pos_ref: &CudaSlice<i32> = if same_dev {
23545 pos_stage
23546 .as_ref()
23547 .ok_or("step TP decode v2 pos stage not armed")?
23548 } else {
23549 &pos[rank]
23550 };
23551 engine.qk_norm_rope_append_inc_dcw(
23552 &q_raw[rank],
23553 &k_raw[rank],
23554 &v_raw[rank],
23555 &attention.q_norm[rank],
23556 &attention.k_norm[rank],
23557 &mut q[rank],
23558 &mut k[rank],
23559 pos_ref,
23560 k_plane,
23561 v_plane,
23562 len_d,
23563 base_d,
23564 &mut fuse_ctr[rank],
23565 kv_dim_k,
23566 kv_dim_v,
23567 k_tok_bytes,
23568 v_tok_bytes,
23569 head_dim,
23570 geometry.n_rot as usize,
23571 local_heads,
23572 local_kv_heads,
23573 self.cfg.rms_eps,
23574 geometry.rope_base,
23575 1.0,
23576 rope_freqs[rank],
23577 )?;
23578 } else {
23579 engine.append_kv_quantized_dcw(
23580 &ws.k[rank],
23581 &ws.v_raw[rank],
23582 k_plane,
23583 v_plane,
23584 len_d,
23585 base_d,
23586 kv_dim_k,
23587 kv_dim_v,
23588 k_tok_bytes,
23589 v_tok_bytes,
23590 )?;
23591 }
23592 if !fuse_rope {
23593 let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
23594 format!("Step TP layer {il} has no KV cache rank {rank}")
23595 })?;
23596 engine.inc_i32(rank_cache.len_d_mut())?;
23597 }
23598 }
23599 let distributed = cache.tp_kv[il]
23600 .as_ref()
23601 .expect("distributed cache checked above");
23602 let rank_cache = distributed
23603 .rank(rank)
23604 .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23605 let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * k_tok_bytes_c);
23606 let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * v_tok_bytes_c);
23607 if fa2_col.is_some() {
23608 continue;
23611 }
23612 {
23613 let crate::tp::StepTpDecodeV2Ws { q, gate, gated, .. } = &mut *ws;
23616 engine.fa_decode_dcw(
23617 &q[rank],
23618 &k_ring,
23619 &v_ring,
23620 &mut gated[rank],
23621 head_dim,
23622 local_heads,
23623 local_kv_heads,
23624 rank_cache.len_d(),
23625 rank_cache.base_d(),
23626 window.unwrap_or(0),
23627 t_kv,
23628 geometry.attention_scale(),
23629 k_tok_bytes_c,
23630 v_tok_bytes_c,
23631 has_gate.then_some(&gate[rank]),
23632 )?;
23633 }
23634 continue;
23635 }
23636 let distributed = cache.tp_kv[il]
23637 .as_ref()
23638 .expect("distributed cache checked above");
23639 let rank_cache = distributed
23640 .rank(rank)
23641 .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23642 let k_view = engine.view_u8_range(
23643 rank_cache.k(),
23644 physical.start * k_tok_bytes_c,
23645 physical.end * k_tok_bytes_c,
23646 );
23647 let v_view = engine.view_u8_range(
23648 rank_cache.v(),
23649 physical.start * v_tok_bytes_c,
23650 physical.end * v_tok_bytes_c,
23651 );
23652 if has_gate {
23653 engine.fa_decode_kvmod(
23654 &ws.q[rank],
23655 &k_view,
23656 &v_view,
23657 &mut ws.attn_out[rank],
23658 head_dim,
23659 local_heads,
23660 local_kv_heads,
23661 t_kv,
23662 geometry.attention_scale(),
23663 k_tok_bytes_c,
23664 v_tok_bytes_c,
23665 false,
23666 )?;
23667 engine.attn_head_gate(
23668 &ws.attn_out[rank],
23669 &ws.gate[rank],
23670 &mut ws.gated[rank],
23671 None,
23672 head_dim,
23673 local_heads,
23674 1,
23675 )?;
23676 } else {
23677 engine.fa_decode_kvmod(
23678 &ws.q[rank],
23679 &k_view,
23680 &v_view,
23681 &mut ws.gated[rank],
23682 head_dim,
23683 local_heads,
23684 local_kv_heads,
23685 t_kv,
23686 geometry.attention_scale(),
23687 k_tok_bytes_c,
23688 v_tok_bytes_c,
23689 false,
23690 )?;
23691 }
23692 }
23693
23694 let output = if let Some(col) = fa2_col.filter(|_| dcw) {
23701 tp.runtime.decode_v2_stash_fa2(ws, e, col)?;
23705 crate::tp::set_spec_fa2_stashed();
23706 e.uninit(ws.o_out)?
23707 } else if let Some(col) = crate::tp::take_tcol_oproj_defer() {
23708 if tp.runtime.decode_v2_oproj_tcol_eligible(ws, &tp.o) {
23709 tp.runtime.decode_v2_stash_gated(ws, e, col)?;
23710 crate::tp::set_tcol_oproj_stashed();
23711 e.uninit(ws.o_out)?
23712 } else {
23713 tp.runtime.decode_v2_finish(ws, e, &tp.o)?
23714 }
23715 } else {
23716 tp.runtime.decode_v2_finish(ws, e, &tp.o)?
23717 };
23718
23719 let local = cache.kv[il]
23723 .as_mut()
23724 .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
23725 if local.len != base_len || base_len + 1 > max_ctx {
23726 return Err(format!(
23727 "Step TP layer {il} local cache changed during decode: \
23728 len={} base={base_len} max={max_ctx}",
23729 local.len
23730 )
23731 .into());
23732 }
23733 if crate::tp::no_local_shadow_on() {
23734 local.len = base_len + 1;
23737 if !crate::tp::len_mirror_lazy_on() {
23741 e.set_i32_one(&mut local.len_d, local.len as i32)?;
23742 }
23743 } else {
23744 let retain_from = window
23745 .map(|window| {
23746 let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
23747 let rollback_retain =
23748 base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
23749 staged_retain.min(rollback_retain)
23750 })
23751 .unwrap_or(0);
23752 let write_row = e.prepare_kv_append(local, retain_from, 1)?;
23753 e.append_kv_quantized(
23754 &ws.k_shadow,
23755 &ws.v_shadow,
23756 &mut local.k,
23757 &mut local.v,
23758 write_row,
23759 local.kv_dim_k,
23760 local.kv_dim_v,
23761 local.k_tok_bytes,
23762 local.v_tok_bytes,
23763 false,
23764 )?;
23765 local.len = base_len + 1;
23766 e.set_i32_one(&mut local.len_d, local.len as i32)?;
23767 }
23768 Ok(output)
23769 })();
23770
23771 let output = match staged {
23772 Ok(output) => output,
23773 Err(error) => {
23774 let _ = tp.runtime.rollback_tp_kv_transaction(
23775 cache.tp_kv[il]
23776 .as_mut()
23777 .expect("distributed cache checked above"),
23778 transaction,
23779 );
23780 if let Some(local) = cache.kv[il].as_mut() {
23781 local.len = base_len;
23782 let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
23783 }
23784 return Err(error);
23785 }
23786 };
23787 let lazy_commit = fuse_rope && crate::tp::len_mirror_lazy_on();
23792 if lazy_commit {
23793 if let Err(error) = tp.runtime.commit_tp_kv_transaction_external(
23794 cache.tp_kv[il]
23795 .as_mut()
23796 .expect("distributed cache checked above"),
23797 transaction,
23798 1,
23799 ) {
23800 let _ = tp.runtime.rollback_tp_kv_transaction(
23801 cache.tp_kv[il]
23802 .as_mut()
23803 .expect("distributed cache checked above"),
23804 transaction,
23805 );
23806 let local = cache.kv[il].as_mut().expect("local cache checked above");
23807 local.len = base_len;
23808 e.set_i32_one(&mut local.len_d, base_len as i32)?;
23809 return Err(error);
23810 }
23811 } else if let Err(error) = tp.runtime.commit_tp_kv_transaction(
23812 cache.tp_kv[il]
23813 .as_mut()
23814 .expect("distributed cache checked above"),
23815 transaction,
23816 1,
23817 ) {
23818 let _ = tp.runtime.rollback_tp_kv_transaction(
23819 cache.tp_kv[il]
23820 .as_mut()
23821 .expect("distributed cache checked above"),
23822 transaction,
23823 );
23824 let local = cache.kv[il].as_mut().expect("local cache checked above");
23825 local.len = base_len;
23826 e.set_i32_one(&mut local.len_d, base_len as i32)?;
23827 return Err(error);
23828 }
23829
23830 let committed = cache.tp_kv[il]
23831 .as_ref()
23832 .expect("distributed cache checked above")
23833 .committed_len();
23834 let local_len = cache.kv[il]
23835 .as_ref()
23836 .expect("local cache checked above")
23837 .len;
23838 if committed != local_len {
23839 return Err(format!(
23840 "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
23841 )
23842 .into());
23843 }
23844 static V2_LOGGED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
23845 if !V2_LOGGED.swap(true, std::sync::atomic::Ordering::Relaxed) {
23846 eprintln!(
23847 "[step-tp-attn-v2] execute layer={} devices={:?} tokens=1 driver=v2 \
23848 qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
23849 kv_cache_distributed=true kv_cache_hydrated={hydrated} \
23850 attention_tensor_parallel=true attention_scope={} \
23851 input_path=root-device-replicated gate={} gate_tensor_parallel={} \
23852 gate_shards={} o_tensor_parallel=true o_reduce=root-device \
23853 local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
23854 bulk_p2p={} workspace=persistent ordering=evented output=e-device \
23855 performance_claim=false (logged once; every decode layer runs this driver)",
23856 tp.layer,
23857 tp.devices,
23858 if window.is_some() {
23859 "rank-local-swa-ring"
23860 } else {
23861 "rank-local-global"
23862 },
23863 has_gate,
23864 use_gate_shards,
23865 if use_gate_shards {
23866 "device-staged"
23867 } else if has_gate {
23868 "root-staged"
23869 } else {
23870 "none"
23871 },
23872 tp.runtime.transport_label(),
23873 tp.runtime.bulk_p2p(),
23874 );
23875 }
23876 Ok(output)
23877 }
23878
23879 #[allow(clippy::too_many_arguments)]
23889 pub(crate) fn step35_decode_attn(
23890 &self,
23891 e: &Engine,
23892 fa: &FullAttnLayer,
23893 il: usize,
23894 h: &CudaSlice<f32>,
23895 pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
23896 pos_d: &CudaSlice<i32>,
23897 cache: &mut Cache,
23898 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23899 if fa
23900 .step_tp_qkv
23901 .as_ref()
23902 .is_some_and(|tp| tp.attention.is_some())
23903 {
23904 if pre_q.is_some() {
23905 return Err(
23906 "rank-local Step attention preserves BF16 activations and refuses the q8_1 \
23907 pre-quantized decode path"
23908 .into(),
23909 );
23910 }
23911 return self.step35_tp_decode_attn_resident(e, fa, il, h, pos_d, cache);
23912 }
23913
23914 let geometry = self.step35_geom(il);
23915 let hd = geometry.head_dim_k as usize;
23916 let nkv = geometry.n_head_kv as usize;
23917 let nh = geometry.n_head as usize;
23918 let rbase = geometry.rope_base;
23919 let scale = geometry.attention_scale();
23920 let swa = geometry.window.is_some();
23921 let eps = self.cfg.rms_eps;
23922 let win = geometry.window.unwrap_or(0) as usize;
23923 let n_rot = geometry.n_rot as usize;
23924 let n_embd = self.cfg.n_embd as usize;
23925 let gw = fa
23926 .attn_gate
23927 .as_ref()
23928 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
23929
23930 let tp_qkv = if fa.step_tp_qkv.is_some() {
23931 if pre_q.is_some() {
23932 return Err(
23933 "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
23934 pre-quantized decode path"
23935 .into(),
23936 );
23937 }
23938 self.full_attn_tp_qkv(e, fa, h, 1)?
23939 } else {
23940 None
23941 };
23942
23943 let (q0, k0, v0, gt) = match tp_qkv {
23944 Some(mut g3) => {
23945 let v = g3.pop().unwrap();
23946 let k = g3.pop().unwrap();
23947 let q = g3.pop().unwrap();
23948 let gt = e.matmul(gw, h, 1)?;
23949 (q, k, v, gt)
23950 }
23951 None => match pre_q {
23952 Some((hq, hdq)) => {
23953 debug_assert!(
23954 e.uses_q8_1_fast(gw),
23955 "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
23956 (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
23957 );
23958 let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
23959 Some(t3) => t3,
23960 None => (
23961 e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
23962 e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
23963 e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
23964 ),
23965 };
23966 let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
23967 (a, b, c, gt)
23968 }
23969 None => {
23970 if e.uses_q8_1_fast(&fa.wq)
23971 && e.uses_q8_1_fast(&fa.wk)
23972 && e.uses_q8_1_fast(&fa.wv)
23973 && e.uses_q8_1_fast(gw)
23974 {
23975 let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
23976 let (a, b, c) =
23977 match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
23978 Some(t3) => t3,
23979 None => (
23980 e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
23981 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
23982 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
23983 ),
23984 };
23985 let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
23986 (a, b, c, gt)
23987 } else {
23988 (
23989 e.matmul(&fa.wq, h, 1)?,
23990 e.matmul(&fa.wk, h, 1)?,
23991 e.matmul(&fa.wv, h, 1)?,
23992 e.matmul(gw, h, 1)?,
23993 )
23994 }
23995 }
23996 },
23997 };
23998
23999 let mut q = e.uninit(nh * hd)?;
24000 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
24001 let mut k = e.uninit(nkv * hd)?;
24002 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
24003 let ff = if swa {
24004 None
24005 } else {
24006 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
24007 };
24008 #[cfg(debug_assertions)]
24009 if let Some(ff) = ff {
24010 crate::debug_assert_tensor_stream_device(
24011 ff,
24012 &e.stream(),
24013 "step35_decode_attn.rope_freqs",
24014 );
24015 }
24016 e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
24017
24018 if std::env::var("MEMRA_NOFA").is_ok() {
24019 return Err(
24020 "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
24021 cache; unset MEMRA_NOFA to use fa_decode"
24022 .into(),
24023 );
24024 }
24025 let kvl = cache.kv[il].as_mut().unwrap();
24026 let next_len = kvl.len + 1;
24027 let (off, t_kv) = if swa && next_len > win {
24028 (next_len - win, win)
24029 } else {
24030 (0, next_len)
24031 };
24032 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
24033 e.append_kv_quantized(
24034 &k,
24035 &v0,
24036 &mut kvl.k,
24037 &mut kvl.v,
24038 write_row,
24039 kvl.kv_dim_k,
24040 kvl.kv_dim_v,
24041 kvl.k_tok_bytes,
24042 kvl.v_tok_bytes,
24043 crate::Engine::kv_fp8_on(),
24044 )?;
24045 kvl.len = next_len;
24046 let physical = kvl.physical_rows(off, off + t_kv)?;
24047 let k_view = e.view_u8_range(
24048 &kvl.k,
24049 physical.start * kvl.k_tok_bytes,
24050 physical.end * kvl.k_tok_bytes,
24051 );
24052 let v_view = e.view_u8_range(
24053 &kvl.v,
24054 physical.start * kvl.v_tok_bytes,
24055 physical.end * kvl.v_tok_bytes,
24056 );
24057 let mut attn = e.uninit(nh * hd)?;
24058 e.fa_decode_kvmod(
24059 &q,
24060 &k_view,
24061 &v_view,
24062 &mut attn,
24063 hd,
24064 nh,
24065 nkv,
24066 t_kv,
24067 scale,
24068 kvl.k_tok_bytes,
24069 kvl.v_tok_bytes,
24070 crate::Engine::kv_fp8_on(),
24071 )?;
24072
24073 let mut ag = e.uninit(nh * hd)?;
24074 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, 1)?;
24075 self.full_attn_o(e, fa, &ag, 1)
24076 }
24077}
24078
24079impl HybridModel {
24088 pub fn is_gemma4_e4b(&self) -> bool {
24089 self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
24090 }
24091
24092 fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
24096 let g = self.cfg.gemma4.as_ref().unwrap();
24097 let swa = g.swa_pattern[il];
24098 let hd = if swa {
24099 g.key_length_swa
24100 } else {
24101 g.key_length_global
24102 } as usize;
24103 let Mixer::Full(fa) = &self.layers[il].mixer else {
24104 panic!("e4b layer {il} not full-attn")
24105 };
24106 let nh = fa.wq.out_features() / hd;
24107 let nkv = fa.wk.out_features() / hd;
24108 (
24109 hd,
24110 nkv,
24111 nh,
24112 if swa {
24113 g.rope_base_swa
24114 } else {
24115 g.rope_base_global
24116 },
24117 1.0,
24118 swa,
24119 )
24120 }
24121
24122 fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
24124 self.layers[il]
24125 .gemma4
24126 .as_ref()
24127 .and_then(|b| b.e4b.as_ref())
24128 .and_then(|e4| e4.kv_share.map(|t| t as usize))
24129 }
24130
24131 fn gemma4_e4b_inp_pl(
24136 &self,
24137 e: &Engine,
24138 tokens: &[u32],
24139 x_scaled: &CudaSlice<f32>,
24140 t: usize,
24141 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24142 let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
24143 self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
24144 }
24145
24146 fn gemma4_e4b_inp_pl_dev(
24148 &self,
24149 e: &Engine,
24150 tok_d: &CudaSlice<u32>,
24151 x_scaled: &CudaSlice<f32>,
24152 t: usize,
24153 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24154 let aux = self.gemma4_aux.as_ref().unwrap();
24155 let m = aux.e4b.as_ref().unwrap();
24156 let n_embd = self.cfg.n_embd as usize;
24157 let n_layer = self.layers.len();
24158 let width = m.n_epl * n_layer;
24159 let tbl = m.tok_tbl_gpu.get_or_init(|| {
24160 e.upload_u8(&m.tok_embd_bytes)
24161 .expect("e4b per-layer token table upload")
24162 });
24163 let mut a =
24164 e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
24165 e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
24166 let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
24167 e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
24168 let mut pn = e.uninit(t * width)?;
24169 e.rms_norm(
24170 &p,
24171 m.proj_norm.float_data(),
24172 &mut pn,
24173 m.n_epl,
24174 t * n_layer,
24175 self.cfg.rms_eps,
24176 )?;
24177 let mut out = e.uninit(t * width)?;
24178 e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
24179 Ok(out)
24180 }
24181
24182 #[allow(clippy::too_many_arguments)]
24187 fn gemma4_e4b_attn(
24188 &self,
24189 e: &Engine,
24190 il: usize,
24191 hq: &CudaSlice<i8>,
24192 hdq: &CudaSlice<f32>,
24193 pos_d: &CudaSlice<i32>,
24194 t: usize,
24195 cache: &mut Cache,
24196 dc_bucket: Option<usize>,
24197 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24198 let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
24199 let eps = self.cfg.rms_eps;
24200 let aux = self.gemma4_aux.as_ref().unwrap();
24201 let ones = aux.ones(e);
24202 #[cfg(debug_assertions)]
24203 crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
24204 let Mixer::Full(fa) = &self.layers[il].mixer else {
24205 unreachable!()
24206 };
24207 let h0 = e.zeros(0)?;
24211 let h = &h0;
24212
24213 let ff = if swa {
24214 None
24215 } else {
24216 Some(
24217 aux.rope_freqs(e)
24218 .expect("e4b global rope needs rope_freqs.weight"),
24219 )
24220 };
24221 #[cfg(debug_assertions)]
24222 if let Some(ff) = ff {
24223 crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
24224 }
24225 let share = self.gemma4_e4b_kv_target(il);
24226 let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
24228 let mut q;
24229 if let Some(_tgt) = share {
24230 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
24231 q = e.uninit(t * nh * hd)?;
24232 let mut kdummy = e.uninit(1)?;
24235 let mut vdummy = e.uninit(1)?;
24236 e.rms_norm_qkv_rope(
24237 &q0,
24238 &q0,
24239 &q0,
24240 fa.q_norm.float_data(),
24241 fa.q_norm.float_data(),
24242 ones,
24243 &mut q,
24244 &mut kdummy,
24245 &mut vdummy,
24246 hd,
24247 self.gemma4_rope_dims(il),
24248 nh * t,
24249 0,
24250 pos_d,
24251 nh,
24252 1,
24253 base,
24254 1.0,
24255 ff,
24256 eps,
24257 )?;
24258 } else {
24259 let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
24263 let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
24264 q = e.uninit(t * nh * hd)?;
24265 let mut k = e.uninit(t * nkv * hd)?;
24266 let mut v = e.uninit(t * nkv * hd)?;
24267 if t == 1 && cat.is_some() {
24268 #[allow(clippy::unnecessary_unwrap)]
24269 let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
24271 e.rms_norm_qkv_rope_cat(
24272 &qkv0,
24273 fa.q_norm.float_data(),
24274 fa.k_norm.float_data(),
24275 ones,
24276 &mut q,
24277 &mut k,
24278 &mut v,
24279 hd,
24280 self.gemma4_rope_dims(il),
24281 nh,
24282 nkv,
24283 pos_d,
24284 nh,
24285 nkv,
24286 base,
24287 1.0,
24288 ff,
24289 eps,
24290 )?;
24291 } else {
24292 let (q0, k0, v0) = match if t == 1 {
24293 e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
24294 } else {
24295 static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24298 if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
24299 e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
24300 } else {
24301 None
24302 }
24303 } {
24304 Some(triple) => triple,
24305 None => (
24306 e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
24307 e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
24308 e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
24309 ), };
24311 e.rms_norm_qkv_rope(
24314 &q0,
24315 &k0,
24316 &v0,
24317 fa.q_norm.float_data(),
24318 fa.k_norm.float_data(),
24319 ones,
24320 &mut q,
24321 &mut k,
24322 &mut v,
24323 hd,
24324 self.gemma4_rope_dims(il),
24325 nh * t,
24326 nkv * t,
24327 pos_d,
24328 nh,
24329 nkv,
24330 base,
24331 1.0,
24332 ff,
24333 eps,
24334 )?;
24335 }
24336 let kvl = cache.kv[il].as_mut().unwrap();
24337 let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24341 if dc_bucket.is_some() {
24342 debug_assert!(t == 1);
24347 e.append_kv_quantized_row_dc_inc(
24349 &k,
24350 &v,
24351 &mut kvl.k,
24352 &mut kvl.v,
24353 &mut kvl.len_d,
24354 kvl.kv_dim_k,
24355 kvl.kv_dim_v,
24356 kvl.k_tok_bytes,
24357 kvl.v_tok_bytes,
24358 cls,
24359 )?;
24360 } else {
24361 e.append_kv_quantized_rows(
24362 &k,
24363 &v,
24364 &mut kvl.k,
24365 &mut kvl.v,
24366 kvl.len,
24367 t,
24368 kvl.kv_dim_k,
24369 kvl.kv_dim_v,
24370 kvl.k_tok_bytes,
24371 kvl.v_tok_bytes,
24372 cls,
24373 )?;
24374 kvl.len += t;
24375 }
24376 kv_f32 = Some((k, v));
24377 }
24378 let kvl_idx = share.unwrap_or(il);
24381 let kvl = cache.kv[kvl_idx].as_ref().unwrap();
24382 let base_len = kvl.len - t; let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
24384 let mut attn = e.uninit(t * nh * hd)?;
24385 if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
24397 if let Some((kf, vf)) = &kv_f32 {
24398 if hd == 256 && t <= win {
24399 e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24400 return e.matmul(&fa.wo, &attn, t);
24401 }
24402 if hd == 256 && swa && t > win {
24403 e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
24404 return e.matmul(&fa.wo, &attn, t);
24405 }
24406 if hd == 512 && !swa {
24407 e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24408 return e.matmul(&fa.wo, &attn, t);
24409 }
24410 } else if share.is_some() {
24411 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24412 let k_view = e.view_u8(&kvl.k, kvl.k.len());
24413 let v_view = e.view_u8(&kvl.v, kvl.v.len());
24414 if hd == 256 && (!swa || t <= win) {
24415 e.fa_prefill_view(
24417 &q,
24418 &k_view,
24419 &v_view,
24420 &mut attn,
24421 hd,
24422 nh,
24423 nkv,
24424 t,
24425 t,
24426 scale,
24427 true,
24428 kvl.k_tok_bytes,
24429 kvl.v_tok_bytes,
24430 g,
24431 )?;
24432 return e.matmul(&fa.wo, &attn, t);
24433 }
24434 let kv_dim = nkv * hd;
24437 let mut kf = e.uninit(t * kv_dim)?;
24438 let mut vf = e.uninit(t * kv_dim)?;
24439 e.fa_dequant_kv_view_f32(
24440 &k_view,
24441 &v_view,
24442 &mut kf,
24443 &mut vf,
24444 kv_dim,
24445 kv_dim,
24446 t,
24447 kvl.k_tok_bytes,
24448 kvl.v_tok_bytes,
24449 g,
24450 )?;
24451 if hd == 512 {
24452 e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24453 } else {
24454 e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
24455 }
24456 return e.matmul(&fa.wo, &attn, t);
24457 }
24458 }
24459 if let Some(bucket) = dc_bucket {
24460 assert!(t == 1);
24465 let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
24471 bucket.min(crate::fa512_min_tkv().saturating_sub(1))
24472 } else {
24473 bucket
24474 };
24475 let k_view = e.view_u8(&kvl.k, kvl.k.len());
24476 let v_view = e.view_u8(&kvl.v, kvl.v.len());
24477 let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24478 if crate::Engine::wpf_level() >= 1 {
24486 e.prefetch_weight_l2(&fa.wo)?;
24487 }
24488 if e.uses_q8_1_fast(&fa.wo) {
24491 let mut oq = e.alloc_i8_uninit(nh * hd)?;
24492 let mut od = e.zeros(nh * hd / 32)?;
24493 e.fa_decode_dc_q8(
24494 &q,
24495 &k_view,
24496 &v_view,
24497 &mut attn,
24498 hd,
24499 nh,
24500 nkv,
24501 &kvl.len_d,
24502 bucket,
24503 scale,
24504 kvl.k_tok_bytes,
24505 kvl.v_tok_bytes,
24506 g,
24507 Some((&mut oq, &mut od)),
24508 )?;
24509 return e.matmul_pre(&fa.wo, &oq, &od, &attn, t);
24510 }
24511 e.fa_decode_dc(
24512 &q,
24513 &k_view,
24514 &v_view,
24515 &mut attn,
24516 hd,
24517 nh,
24518 nkv,
24519 &kvl.len_d,
24520 bucket,
24521 scale,
24522 kvl.k_tok_bytes,
24523 kvl.v_tok_bytes,
24524 g,
24525 )?;
24526 return e.matmul(&fa.wo, &attn, t);
24527 }
24528 for i in 0..t {
24529 let avail = base_len + i + 1;
24530 let (off_tok, t_kv) = if swa && avail > win {
24531 (avail - win, win)
24532 } else {
24533 (0, avail)
24534 };
24535 let k_view = e.view_u8_range(
24536 &kvl.k,
24537 off_tok * kvl.k_tok_bytes,
24538 (off_tok + t_kv) * kvl.k_tok_bytes,
24539 );
24540 let v_view = e.view_u8_range(
24541 &kvl.v,
24542 off_tok * kvl.v_tok_bytes,
24543 (off_tok + t_kv) * kvl.v_tok_bytes,
24544 );
24545 let qv = e.view(&q, t * nh * hd);
24546 let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
24547 let mut q_one = e.uninit(nh * hd)?;
24548 e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
24549 let mut a_one = e.uninit(nh * hd)?;
24550 e.fa_decode_kvmod(
24554 &q_one,
24555 &k_view,
24556 &v_view,
24557 &mut a_one,
24558 hd,
24559 nh,
24560 nkv,
24561 t_kv,
24562 scale,
24563 kvl.k_tok_bytes,
24564 kvl.v_tok_bytes,
24565 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
24566 )?;
24567 e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
24568 }
24569 e.matmul(&fa.wo, &attn, t)
24570 }
24571
24572 fn gemma4_e4b_trunk(
24577 &self,
24578 e: &Engine,
24579 tokens: &[u32],
24580 pos0: usize,
24581 cache: &mut Cache,
24582 head_last: bool,
24583 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24584 let n_embd = self.cfg.n_embd as usize;
24585 let t = tokens.len();
24586 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
24587 let pos_d = e.htod_i32(&pos)?;
24588 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
24589 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
24590 let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
24591 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
24592 }
24593
24594 #[allow(clippy::too_many_arguments)] fn gemma4_e4b_trunk_core(
24599 &self,
24600 e: &Engine,
24601 x_in: CudaSlice<f32>,
24602 inp_pl: CudaSlice<f32>,
24603 pos_d: &CudaSlice<i32>,
24604 t: usize,
24605 cache: &mut Cache,
24606 dc_bucket: Option<usize>,
24607 cap_logits: bool,
24608 head_last: bool,
24609 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24610 let n_embd = self.cfg.n_embd as usize;
24611 let eps = self.cfg.rms_eps;
24612 let n_layer = self.layers.len();
24613 let mut x = x_in;
24614 let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
24615 let n_epl = aux_e4b.n_epl;
24616
24617 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
24623 for il in 0..n_layer {
24624 let layer = &self.layers[il];
24625 let (hq, hdq) = match h_carry.take() {
24626 Some(p) => p,
24627 None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
24628 };
24629 let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
24630 let bits = layer.gemma4.as_ref().unwrap();
24633 let e4b = bits.e4b.as_ref().expect("e4b layer bits");
24634 let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
24645 let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
24646 e,
24647 layer,
24648 &o,
24649 &x,
24650 t,
24651 Some(layer.post_attn_norm.float_data()),
24652 fuse_exit,
24653 )?;
24654 let mut resid = e.uninit(t * n_embd)?;
24655 let g = if fuse_exit {
24661 let (rq, rd) = e.rms_pre_add_q8_1(
24663 &sn,
24664 bits.post_ffw_norm.float_data(),
24665 &attn_out,
24666 &mut resid,
24667 n_embd,
24668 t,
24669 self.cfg.rms_eps,
24670 )?;
24671 e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
24672 } else {
24673 e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
24674 e.matmul(&e4b.inp_gate, &resid, t)?
24675 };
24676 let mut act = e.uninit(t * n_epl)?;
24677 let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
24678 let ipv = e.view(&inp_pl, n_epl * n_layer);
24679 let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
24680 let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
24681 e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
24682 } else {
24683 let mut inp_this = e.uninit(t * n_epl)?;
24684 e.copy_rows_strided(
24685 &inp_pl,
24686 &mut inp_this,
24687 n_epl,
24688 t,
24689 n_epl * n_layer,
24690 il * n_epl,
24691 )?;
24692 e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
24693 e.matmul(&e4b.proj, &act, t)?
24694 };
24695 let next_norm = if il + 1 < n_layer {
24698 self.layers[il + 1].attn_norm.float_data()
24699 } else {
24700 self.output_norm.float_data()
24701 };
24702 let mut xn = e.uninit(t * n_embd)?;
24703 let pair = e.rms_pre_add_scale_rms_norm_q8_1(
24704 &y,
24705 e4b.post_norm.float_data(),
24706 &resid,
24707 bits.layer_scale,
24708 next_norm,
24709 &mut xn,
24710 n_embd,
24711 t,
24712 eps,
24713 )?;
24714 h_carry = Some(pair);
24715 x = xn;
24716 }
24717 let (oq, odq) = h_carry.take().unwrap();
24721 let h0 = e.zeros(0)?;
24722 let hm = if head_last { 1 } else { t };
24723 let (hq, hd) = if head_last && t > 1 {
24724 let mut q1 = e.uninit_i8(n_embd)?;
24725 e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
24726 let nb = n_embd / 32;
24727 let mut d1 = e.uninit(nb)?;
24728 e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
24729 (q1, d1)
24730 } else {
24731 (oq, odq)
24732 };
24733 let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
24734 if cap_logits {
24738 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
24739 e.softcap(&mut ld, cap, hm * self.output.out_features())?;
24740 }
24741 self.gemma4_suppress(e, &mut ld, hm)?; Ok((ld, x))
24743 }
24744
24745 pub fn gemma4_e4b_decode_step_t_am_dev(
24752 &self,
24753 e: &Engine,
24754 tok_d: &CudaSlice<u32>,
24755 t: usize,
24756 pos0: usize,
24757 cache: &mut Cache,
24758 ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24759 let n_embd = self.cfg.n_embd as usize;
24760 let eps = self.cfg.rms_eps;
24761 let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
24762 let pos_d = e.htod_i32(&pos)?;
24763 let embd_gpu = self
24764 .embd_gpu
24765 .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
24766 let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
24767 let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
24768 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
24769 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
24770 let (ld, xp) =
24771 self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
24772 let n_vocab = self.output.out_features();
24775 let mut vam = e.stream().alloc_zeros::<u32>(t)?;
24776 for i in 0..t {
24777 e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
24778 }
24779 let mut hn = e.uninit(t * n_embd)?;
24780 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
24781 cache.pos += t;
24782 Ok((vam, hn))
24783 }
24784
24785 pub(crate) fn gemma4_e4b_decode_step_t_h(
24788 &self,
24789 e: &Engine,
24790 tokens: &[u32],
24791 pos0: usize,
24792 cache: &mut Cache,
24793 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24794 let n_embd = self.cfg.n_embd as usize;
24795 let eps = self.cfg.rms_eps;
24796 let t = tokens.len();
24797 let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
24798 let mut hn = e.uninit(t * n_embd)?;
24799 e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
24800 cache.pos += t;
24801 Ok((e.dtoh(&ld)?, hn))
24802 }
24803
24804 #[allow(clippy::too_many_arguments)] pub fn gemma4_e4b_decode_step_dcg(
24811 &self,
24812 e: &Engine,
24813 token_d: &mut CudaSlice<u32>,
24814 pos_d: &mut CudaSlice<i32>,
24815 embd_gpu: &CudaSlice<u8>,
24816 embd_qt: i32,
24817 embd_rb: usize,
24818 cache: &mut Cache,
24819 n_vocab: usize,
24820 bucket: usize,
24821 ) -> Result<(), Box<dyn std::error::Error>> {
24822 let n_embd = self.cfg.n_embd as usize;
24823 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
24824 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
24825 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
24826 let (ld, _x) =
24827 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
24828 e.argmax_token_device_into(&ld, token_d, n_vocab)?;
24829 e.inc_seqlen(pos_d)?;
24830 Ok(())
24831 }
24832
24833 #[allow(clippy::too_many_arguments)]
24841 pub fn gemma4_e4b_decode_step_dc(
24842 &self,
24843 e: &Engine,
24844 token_d: &CudaSlice<u32>,
24845 pos_d: &mut CudaSlice<i32>,
24846 embd_gpu: &CudaSlice<u8>,
24847 embd_qt: i32,
24848 embd_rb: usize,
24849 cache: &mut Cache,
24850 n_vocab: usize,
24851 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
24852 let n_embd = self.cfg.n_embd as usize;
24853 let eps = self.cfg.rms_eps;
24854 let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
24855 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
24856 let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
24857 let (ld, _x) =
24858 self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
24859 let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
24860 e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
24861 e.inc_seqlen(pos_d)?;
24862 cache.pos += 1;
24863 let _ = eps;
24864 Ok(tok_out)
24865 }
24866
24867 pub(crate) fn gemma4_e4b_decode_step_h(
24870 &self,
24871 e: &Engine,
24872 token: u32,
24873 cache: &mut Cache,
24874 ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24875 let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
24876 let logits = e.dtoh(&ld)?;
24877 cache.pos += 1;
24878 Ok((logits, x))
24879 }
24880
24881 #[allow(clippy::type_complexity)] pub(crate) fn gemma4_e4b_prime(
24886 &self,
24887 e: &Engine,
24888 tokens: &[u32],
24889 cache: &mut Cache,
24890 ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24891 if cache.pos != 0 {
24894 return Err(
24895 "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
24896 call or decode tokenwise"
24897 .into(),
24898 );
24899 }
24900 let n_embd = self.cfg.n_embd as usize;
24901 let t = tokens.len();
24902 let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
24903 cache.pos += t;
24904 let last = e.dtoh(&ld)?; let xv = e.view(&x, t * n_embd);
24906 let row = xv.slice((t - 1) * n_embd..t * n_embd);
24907 let mut h_seed = e.uninit(n_embd)?;
24908 e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
24909 Ok((last, h_seed, x))
24910 }
24911
24912 pub(crate) fn gemma4_e4b_forward(
24914 &self,
24915 e: &Engine,
24916 tokens: &[u32],
24917 last_only: bool,
24918 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
24919 let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
24920 let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
24921 e.dtoh(&ld) }
24923}
24924
24925#[cfg(test)]
24926mod prime_chunk_schedule_tests {
24927 use super::{
24928 PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, PrimePpSignal, PrimePpStageChannels, PrimePpWaveCredits,
24929 PrimePpWaveSlot, active_matrix_values, align_prime_ranges_to_gdn,
24930 dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges, fixed_prime_chunk_ranges_for_ring,
24931 move_prime_cache_layers, parse_step_ep_grouped_prefill, parse_step_tp_prefill,
24932 prime_cache_stage_for_layer, recv_prime_pp_signal, restore_prime_cache_layers,
24933 step_grouped_decode_shape, step_grouped_prefill_shape, step_tp_prefill_shape,
24934 validate_step_prime_batch_modes,
24935 };
24936
24937 fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
24938 ranges.iter().map(|(start, end)| end - start).collect()
24939 }
24940
24941 #[allow(clippy::manual_clamp)] fn auto_chunk(t: usize) -> usize {
24943 t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
24944 }
24945
24946 #[test]
24947 fn ppn_prime_cache_partition_moves_and_restores_every_layer() {
24948 let round_trip = |fence: &[usize], layers: usize| {
24949 let original: Vec<Option<usize>> = (0..layers).map(Some).collect();
24950 let mut parent = original.clone();
24951 let mut stages: Vec<Vec<Option<usize>>> =
24952 (0..fence.len() - 1).map(|_| vec![None; layers]).collect();
24953
24954 move_prime_cache_layers(&mut parent, &mut stages, fence);
24955 assert!(parent.iter().all(Option::is_none));
24956 for layer in 0..layers {
24957 let owner = prime_cache_stage_for_layer(fence, layer);
24958 for (stage, values) in stages.iter().enumerate() {
24959 assert_eq!(values[layer], (stage == owner).then_some(layer));
24960 }
24961 }
24962
24963 restore_prime_cache_layers(&mut parent, &mut stages, fence);
24964 assert_eq!(parent, original);
24965 assert!(stages.iter().flatten().all(Option::is_none));
24966 };
24967
24968 round_trip(&[0, 5, 8], 10);
24970 round_trip(&[0, 2, 5, 8], 10);
24971 round_trip(&[0, 1, 3, 6, 8], 10);
24972 }
24973
24974 #[test]
24975 fn ppn_prime_wave_credit_requires_the_exact_oldest_wave_and_slot() {
24976 let mut credits = PrimePpWaveCredits::default();
24977 let wave0 = PrimePpWaveSlot { wave: 0, slot: 1 };
24978 let wave1 = PrimePpWaveSlot { wave: 1, slot: 0 };
24979 credits.record_send(wave0).unwrap();
24980 assert_eq!(credits.release_required(), None);
24981 credits.record_send(wave1).unwrap();
24982 assert_eq!(credits.release_required(), Some(wave0));
24983
24984 assert!(
24985 credits
24986 .record_release(PrimePpWaveSlot { wave: 0, slot: 0 })
24987 .unwrap_err()
24988 .contains("does not match oldest pending")
24989 );
24990 assert_eq!(credits.release_required(), Some(wave0));
24991 credits.record_release(wave0).unwrap();
24992 credits
24993 .record_send(PrimePpWaveSlot { wave: 2, slot: 1 })
24994 .unwrap();
24995 assert!(
24996 credits
24997 .record_send(PrimePpWaveSlot { wave: 4, slot: 0 })
24998 .unwrap_err()
24999 .contains("while wave 3 was next")
25000 );
25001 assert!(
25002 credits
25003 .record_send(PrimePpWaveSlot { wave: 3, slot: 1 })
25004 .unwrap_err()
25005 .contains("reused slot 1")
25006 );
25007 }
25008
25009 #[test]
25010 fn ppn_prime_wave_signal_reports_order_error_injected_error_and_closure() {
25011 let expected = PrimePpWaveSlot { wave: 2, slot: 1 };
25012
25013 let (sender, receiver) = std::sync::mpsc::channel();
25014 sender.send(PrimePpSignal::Slot(expected)).unwrap();
25015 assert_eq!(
25016 recv_prime_pp_signal(&receiver, expected, true, "test").unwrap(),
25017 expected
25018 );
25019
25020 let (sender, receiver) = std::sync::mpsc::channel();
25021 sender
25022 .send(PrimePpSignal::Slot(PrimePpWaveSlot { wave: 3, slot: 1 }))
25023 .unwrap();
25024 assert!(
25025 recv_prime_pp_signal(&receiver, expected, true, "test")
25026 .unwrap_err()
25027 .contains("expected wave/slot")
25028 );
25029
25030 let (sender, receiver) = std::sync::mpsc::channel();
25031 sender
25032 .send(PrimePpSignal::Error("injected stage failure".into()))
25033 .unwrap();
25034 assert_eq!(
25035 recv_prime_pp_signal(&receiver, expected, true, "test").unwrap_err(),
25036 "injected stage failure"
25037 );
25038
25039 let (upstream_sender, upstream_receiver) = std::sync::mpsc::channel();
25040 let (outgoing_sender, outgoing_receiver) = std::sync::mpsc::channel();
25041 let (_release_sender, released_downstream) = std::sync::mpsc::channel();
25042 PrimePpStageChannels {
25043 incoming: None,
25044 release_upstream: Some(upstream_sender),
25045 outgoing: outgoing_sender,
25046 released_downstream,
25047 }
25048 .notify_failure("injected worker error");
25049 assert_eq!(
25050 recv_prime_pp_signal(&upstream_receiver, expected, false, "test").unwrap_err(),
25051 "injected worker error"
25052 );
25053 assert_eq!(
25054 recv_prime_pp_signal(&outgoing_receiver, expected, false, "test").unwrap_err(),
25055 "injected worker error"
25056 );
25057
25058 let (sender, receiver) = std::sync::mpsc::channel::<PrimePpSignal>();
25059 drop(sender);
25060 assert!(
25061 recv_prime_pp_signal(&receiver, expected, true, "test")
25062 .unwrap_err()
25063 .contains("channel closed while waiting for wave 2")
25064 );
25065 }
25066
25067 #[test]
25072 fn auto_prime_ranges_align_to_the_gdn_grid() {
25073 let c = 32usize; let assert_covers = |ranges: &[(usize, usize)], t: usize| {
25075 assert_eq!(ranges.first().map(|&(s, _)| s), Some(0));
25076 assert_eq!(ranges.last().map(|&(_, e)| e), Some(t));
25077 for w in ranges.windows(2) {
25078 assert_eq!(w[0].1, w[1].0, "ranges must stay contiguous");
25079 }
25080 assert!(ranges.iter().all(|&(s, e)| e > s), "no empty range");
25081 };
25082
25083 let t = 9510usize;
25086 let fill = auto_chunk(t);
25087 let fixed = fixed_prime_chunk_ranges(t, fill);
25088 assert!(
25089 fixed[..fixed.len() - 1].iter().any(|&(_, e)| e % c != 0),
25090 "broken arm vanished: fixed auto boundaries all landed on-grid"
25091 );
25092 let dynamic = dynamic_prime_chunk_ranges(t, fill, &fixed);
25093 assert!(
25094 dynamic[..dynamic.len() - 1]
25095 .iter()
25096 .any(|&(_, e)| e % c != 0),
25097 "broken arm vanished: dynamic auto boundaries all landed on-grid"
25098 );
25099
25100 for ranges in [&fixed, &dynamic] {
25101 let aligned = align_prime_ranges_to_gdn(ranges, t, c);
25102 assert_covers(&aligned, t);
25103 for &(_, e) in &aligned[..aligned.len() - 1] {
25104 assert_eq!(e % c, 0, "internal boundary {e} off the {c}-grid");
25105 }
25106 for (&(_, a), &(_, b)) in aligned.iter().zip(ranges.iter()) {
25108 assert!(a <= b && b - a < c);
25109 }
25110 }
25111
25112 let tight = vec![(0usize, 33usize), (33, 40), (40, 200)];
25115 let aligned = align_prime_ranges_to_gdn(&tight, 200, c);
25116 assert_covers(&aligned, 200);
25117 assert_eq!(aligned, vec![(0, 32), (32, 200)]);
25118
25119 assert_eq!(align_prime_ranges_to_gdn(&[(0, 200)], 200, c), [(0, 200)]);
25121 assert_eq!(align_prime_ranges_to_gdn(&tight, 200, 0), tight.as_slice());
25122 let on_grid = vec![(0usize, 128usize), (128, 256), (256, 300)];
25123 assert_eq!(
25124 align_prime_ranges_to_gdn(&on_grid, 300, c),
25125 on_grid.as_slice()
25126 );
25127 }
25128
25129 #[test]
25130 fn active_matrix_prefix_scopes_reused_prime_slabs() {
25131 assert_eq!(
25132 active_matrix_values(40 * 4096, 29, 4096, "activation").unwrap(),
25133 29 * 4096
25134 );
25135 assert_eq!(
25136 active_matrix_values(29 * 4096, 29, 4096, "activation").unwrap(),
25137 29 * 4096
25138 );
25139 assert_eq!(
25140 active_matrix_values(29 * 4096, 24, 4096, "activation").unwrap(),
25141 24 * 4096
25142 );
25143 assert!(active_matrix_values(28 * 4096, 29, 4096, "activation").is_err());
25144 assert!(active_matrix_values(usize::MAX, usize::MAX, 2, "activation").is_err());
25145 }
25146
25147 #[test]
25148 fn step_tp_prefill_batch_refuses_before_scheduler_fallback() {
25149 assert!(validate_step_prime_batch_modes(false, false).is_ok());
25150
25151 let grouped_without_tp = validate_step_prime_batch_modes(false, true).unwrap_err();
25152 assert!(grouped_without_tp.contains("requires MEMRA_STEP_TP_PREFILL=1"));
25153
25154 for grouped in [false, true] {
25155 let err = validate_step_prime_batch_modes(true, grouped).unwrap_err();
25156 assert!(err.contains("did not clear the live-server performance gate"));
25157 assert!(err.contains("per-session grouped prefill"));
25158 }
25159 }
25160
25161 #[test]
25162 fn step_grouped_path_is_eager_single_token_only() {
25163 assert!(step_grouped_decode_shape(false, 1));
25164 assert!(!step_grouped_decode_shape(true, 1));
25165 assert!(!step_grouped_decode_shape(false, 2));
25166 assert!(!step_grouped_decode_shape(true, 2));
25167 }
25168
25169 #[test]
25170 fn step_grouped_prefill_door_is_strict_and_capacity_bounded() {
25171 assert!(!parse_step_ep_grouped_prefill(None).unwrap());
25172 assert!(!parse_step_ep_grouped_prefill(Some("")).unwrap());
25173 assert!(!parse_step_ep_grouped_prefill(Some("0")).unwrap());
25174 assert!(parse_step_ep_grouped_prefill(Some("1")).unwrap());
25175 assert!(parse_step_ep_grouped_prefill(Some("true")).is_err());
25176 assert!(parse_step_ep_grouped_prefill(Some("2")).is_err());
25177
25178 assert!(step_grouped_prefill_shape(true, true, PRIME_MIN_T));
25179 assert!(step_grouped_prefill_shape(
25180 true,
25181 true,
25182 crate::cache::PRIME_CHUNK_MAX_TOKENS,
25183 ));
25184 assert!(!step_grouped_prefill_shape(true, true, PRIME_MIN_T - 1,));
25185 assert!(!step_grouped_prefill_shape(
25186 true,
25187 true,
25188 crate::cache::PRIME_CHUNK_MAX_TOKENS + 1,
25189 ));
25190 assert!(!step_grouped_prefill_shape(false, true, PRIME_MIN_T));
25191 assert!(!step_grouped_prefill_shape(true, false, PRIME_MIN_T));
25192 }
25193
25194 #[test]
25195 fn step_tp_prefill_door_is_strict_and_default_off() {
25196 assert!(!parse_step_tp_prefill(None).unwrap());
25197 assert!(!parse_step_tp_prefill(Some("")).unwrap());
25198 assert!(!parse_step_tp_prefill(Some("0")).unwrap());
25199 assert!(parse_step_tp_prefill(Some("1")).unwrap());
25200 assert!(parse_step_tp_prefill(Some("true")).is_err());
25201 assert!(parse_step_tp_prefill(Some("2")).is_err());
25202 }
25203
25204 #[test]
25205 fn step_tp_prefill_requires_a_qualified_even_rank_shape() {
25206 assert!(step_tp_prefill_shape(
25207 true,
25208 PRIME_MIN_T,
25209 4,
25210 true,
25211 true,
25212 false,
25213 ));
25214 assert!(!step_tp_prefill_shape(
25215 false,
25216 PRIME_MIN_T,
25217 4,
25218 true,
25219 true,
25220 false,
25221 ));
25222 assert!(!step_tp_prefill_shape(
25223 true,
25224 PRIME_MIN_T - 1,
25225 4,
25226 true,
25227 true,
25228 false,
25229 ));
25230 assert!(step_tp_prefill_shape(
25232 true,
25233 PRIME_MIN_T,
25234 2,
25235 true,
25236 true,
25237 false
25238 ));
25239 assert!(!step_tp_prefill_shape(
25240 true,
25241 PRIME_MIN_T,
25242 1,
25243 true,
25244 true,
25245 false
25246 ));
25247 assert!(!step_tp_prefill_shape(
25248 true,
25249 PRIME_MIN_T,
25250 3,
25251 true,
25252 true,
25253 false
25254 ));
25255 assert!(!step_tp_prefill_shape(
25256 true,
25257 PRIME_MIN_T,
25258 4,
25259 false,
25260 true,
25261 false,
25262 ));
25263 assert!(!step_tp_prefill_shape(
25264 true,
25265 PRIME_MIN_T,
25266 4,
25267 true,
25268 false,
25269 false,
25270 ));
25271 assert!(!step_tp_prefill_shape(
25272 true,
25273 PRIME_MIN_T,
25274 4,
25275 true,
25276 true,
25277 true,
25278 ));
25279 }
25280
25281 #[test]
25282 fn fixed_schedule_retains_measured_geometry() {
25283 assert_eq!(
25284 sizes(&fixed_prime_chunk_ranges(461, 128)),
25285 vec![128, 128, 128, 77]
25286 );
25287 assert_eq!(
25288 sizes(&fixed_prime_chunk_ranges(1833, 230)),
25289 vec![230, 230, 230, 230, 230, 230, 230, 223]
25290 );
25291 assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
25292 let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
25293 assert_eq!(capped, vec![4096, 4088, 16]);
25294 assert!(capped.iter().all(|&rows| rows <= 4096));
25295 assert_eq!(
25296 sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
25297 vec![4100],
25298 "flag-off schedule remains byte-for-byte the legacy monolithic tail",
25299 );
25300 }
25301
25302 #[test]
25303 fn dynamic_schedule_matches_registered_shapes() {
25304 let cases = [
25305 (461, vec![64, 141, 132, 124]),
25306 (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
25307 (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
25308 ];
25309 for (t, expected) in cases {
25310 let chunk = auto_chunk(t);
25311 let fixed = fixed_prime_chunk_ranges(t, chunk);
25312 assert_eq!(
25313 sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
25314 expected
25315 );
25316 }
25317 }
25318
25319 #[test]
25320 fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
25321 for t in 256..=8192 {
25322 let chunk = auto_chunk(t);
25323 let fixed = fixed_prime_chunk_ranges(t, chunk);
25324 let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
25325 assert_eq!(dynamic.len(), fixed.len(), "T={t}");
25326 assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
25327 assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
25328 for pair in dynamic.windows(2) {
25329 assert_eq!(pair[0].1, pair[1].0, "T={t}");
25330 }
25331 assert!(
25332 dynamic
25333 .iter()
25334 .all(|(start, end)| end - start >= PRIME_MIN_T),
25335 "T={t} sizes={:?}",
25336 sizes(&dynamic)
25337 );
25338 if dynamic.len() >= 3 {
25339 let chunk_sizes = sizes(&dynamic);
25340 assert!(
25341 chunk_sizes[0] < chunk_sizes[1],
25342 "T={t} sizes={chunk_sizes:?}"
25343 );
25344 assert!(
25345 chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
25346 "T={t} sizes={chunk_sizes:?}"
25347 );
25348 }
25349 }
25350 }
25351}
25352
25353#[cfg(test)]
25354mod page_prefetch_tests {
25355 use super::{
25356 grouped_worker_prefetch_position, page_prefetch_positions,
25357 page_prefetch_window_from_values, worker_prefetch_positions,
25358 };
25359
25360 #[test]
25361 fn page_prefetch_window_keeps_existing_opt_in_default() {
25362 assert_eq!(page_prefetch_window_from_values(false, None), 0);
25363 assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
25364 assert_eq!(page_prefetch_window_from_values(true, None), 1);
25365 assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
25366 assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
25367 assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
25368 }
25369
25370 #[test]
25371 fn rolling_page_prefetch_advises_each_future_expert_once() {
25372 let advised: Vec<_> = (0..7)
25373 .flat_map(|position| page_prefetch_positions(position, 7, 3))
25374 .collect();
25375 assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
25376
25377 let one_ahead: Vec<_> = (0..4)
25378 .flat_map(|position| page_prefetch_positions(position, 4, 1))
25379 .collect();
25380 assert_eq!(one_ahead, vec![1, 2, 3]);
25381 assert!(page_prefetch_positions(0, 4, 0).is_empty());
25382 }
25383
25384 #[test]
25385 fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
25386 assert_eq!(grouped_worker_prefetch_position(0, None), None);
25387 let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
25388 .chain(
25389 (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
25390 )
25391 .collect();
25392 assert_eq!(positions, vec![0, 1, 2, 3]);
25393 assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
25394 }
25395
25396 #[test]
25397 fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
25398 let queued: Vec<_> = (0..8)
25399 .flat_map(|position| worker_prefetch_positions(position, 8, 5))
25400 .collect();
25401 assert_eq!(queued, (0..8).collect::<Vec<_>>());
25402
25403 let one_at_a_time: Vec<_> = (0..4)
25404 .flat_map(|position| worker_prefetch_positions(position, 4, 1))
25405 .collect();
25406 assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
25407 assert!(worker_prefetch_positions(0, 4, 0).is_empty());
25408 }
25409}
25410
25411pub struct G4DcSlots {
25412 x: CudaSlice<f32>,
25413 xn: CudaSlice<f32>,
25414 cur: CudaSlice<f32>,
25415 hq: CudaSlice<i8>,
25416 hd_: CudaSlice<f32>,
25417 q0: CudaSlice<f32>,
25418 k0: CudaSlice<f32>,
25419 v0: CudaSlice<f32>,
25420 q: CudaSlice<f32>,
25421 k: CudaSlice<f32>,
25422 v: CudaSlice<f32>,
25423 attn: CudaSlice<f32>,
25424 o: CudaSlice<f32>,
25425 attn_out: CudaSlice<f32>,
25426 zsh: CudaSlice<f32>,
25427 zq: CudaSlice<i8>,
25428 zd: CudaSlice<f32>,
25429 gate: CudaSlice<f32>,
25430 up: CudaSlice<f32>,
25431 act: CudaSlice<f32>,
25432 actq: CudaSlice<i8>,
25433 actd: CudaSlice<f32>,
25434 f0: CudaSlice<f32>,
25435 sn: CudaSlice<f32>,
25436 hn: CudaSlice<f32>,
25437 logits: CudaSlice<f32>,
25438}
25439
25440pub struct Step35TokenGraphState {
25445 pub graphs: Vec<(usize, crate::tp::TokenGraph)>,
25447 pub token_d: cudarc::driver::CudaSlice<u32>,
25448 pub pos_d: cudarc::driver::CudaSlice<i32>,
25449 pub logits_stage: cudarc::driver::CudaSlice<f32>,
25450 pub x: cudarc::driver::CudaSlice<f32>,
25455 pub x1: cudarc::driver::CudaSlice<f32>,
25456 pub mixed_stage: cudarc::driver::CudaSlice<f32>,
25457 pub sh_stage: cudarc::driver::CudaSlice<f32>,
25458 pub k_shadow_stage: cudarc::driver::CudaSlice<f32>,
25459 pub v_shadow_stage: cudarc::driver::CudaSlice<f32>,
25460 pub router_logits: cudarc::driver::CudaSlice<f32>,
25463 pub shexp_gate: cudarc::driver::CudaSlice<f32>,
25464 pub shexp_up: cudarc::driver::CudaSlice<f32>,
25465 pub shexp_act: cudarc::driver::CudaSlice<f32>,
25466 pub gate_sig: cudarc::driver::CudaSlice<f32>,
25467 pub dense_z: cudarc::driver::CudaSlice<f32>,
25468 pub dense_gate: cudarc::driver::CudaSlice<f32>,
25469 pub dense_up: cudarc::driver::CudaSlice<f32>,
25470 pub dense_act: cudarc::driver::CudaSlice<f32>,
25471 pub hn: cudarc::driver::CudaSlice<f32>,
25472 pub probe_mixed: cudarc::driver::CudaSlice<f32>,
25475 pub probe_x: cudarc::driver::CudaSlice<f32>,
25476 pub token_hist: cudarc::driver::CudaSlice<u32>,
25479 pub hist_idx: cudarc::driver::CudaSlice<i32>,
25480}
25481
25482impl HybridModel {
25483 #[allow(clippy::type_complexity)] pub(crate) fn step35_token_graph_step(
25495 &self,
25496 e: &Engine,
25497 token: u32,
25498 cache: &mut Cache,
25499 ) -> Result<Option<(Vec<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
25500 if !self.uses_sliding_gated_moe_program()
25501 || !crate::tp::step_tp_graph_enabled()?
25502 || !crate::tp::step_tp_dcw_enabled()?
25503 || !crate::tp::step_tp_qkv_fused_enabled()?
25504 || !crate::tp::step_tp_dev_router_enabled()?
25505 || !crate::tp::step_nvfp4_dev_routes_enabled()?
25506 {
25507 return Ok(None);
25508 }
25509 if !crate::spec::graph_launch_headroom_ok(e) {
25515 static NOTED: std::sync::Once = std::sync::Once::new();
25516 NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
25517 return Ok(None);
25518 }
25519 let n_embd = self.cfg.n_embd as usize;
25520 let n_vocab = self.cfg.n_vocab as usize;
25521 let n_layers = self.layers.len();
25522 let pos = cache.pos;
25523 let staged_next = pos + 1;
25524 if staged_next < 96 {
25525 return Ok(None); }
25527
25528 for il in 0..n_layers {
25531 let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
25532 return Ok(None); };
25534 if tp_kv.peek_append_ring(1)?.1 {
25535 return Ok(None);
25536 }
25537 }
25538
25539 let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
25542 if !fa_vec {
25543 return Ok(None);
25544 }
25545 let sp = crate::fa_split_keys(staged_next, 8);
25546 let bucket_max = (n_splits * sp).max(staged_next);
25547
25548 let mut state_guard = self
25549 .step35_token_graph
25550 .lock()
25551 .map_err(|_| "step35 token graph lock is poisoned")?;
25552 if state_guard.is_none() {
25553 let _main = e.gpu.enter_main()?;
25554 let n_expert = self
25555 .cfg
25556 .moe
25557 .as_ref()
25558 .map(|m| m.expert_count as usize)
25559 .unwrap_or(0);
25560 let n_ff_sh = self
25561 .layers
25562 .iter()
25563 .find_map(|l| match &l.ffn {
25564 crate::hybrid::Ffn::Moe(m) => m.gate_shexp.as_ref().map(|g| g.out_features()),
25565 _ => None,
25566 })
25567 .unwrap_or(0);
25568 let n_ff_dense = self
25569 .layers
25570 .iter()
25571 .find_map(|l| match &l.ffn {
25572 crate::hybrid::Ffn::Dense { ffn_gate, .. } => Some(ffn_gate.out_features()),
25573 _ => None,
25574 })
25575 .unwrap_or(0);
25576 *state_guard = Some(Step35TokenGraphState {
25577 graphs: Vec::new(),
25578 token_d: e.stream().clone_htod(&[0u32])?,
25579 pos_d: e.htod_i32(&[pos as i32])?,
25580 logits_stage: e.htod(&vec![0.0f32; n_vocab])?,
25581 x: e.htod(&vec![0.0f32; n_embd])?,
25582 x1: e.htod(&vec![0.0f32; n_embd])?,
25583 mixed_stage: e.htod(&vec![0.0f32; n_embd])?,
25584 sh_stage: e.htod(&vec![0.0f32; n_embd])?,
25585 k_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
25586 v_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
25587 router_logits: e.htod(&vec![0.0f32; n_expert.max(1)])?,
25588 shexp_gate: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25589 shexp_up: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25590 shexp_act: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25591 gate_sig: e.htod(&[1.0f32; 1])?,
25592 dense_z: e.htod(&vec![0.0f32; n_embd])?,
25593 dense_gate: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25594 dense_up: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25595 dense_act: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25596 hn: e.htod(&vec![0.0f32; n_embd])?,
25597 probe_mixed: e.htod(&vec![0.0f32; n_embd])?,
25598 probe_x: e.htod(&vec![0.0f32; n_embd])?,
25599 token_hist: e.stream().clone_htod(&[0u32; 16])?,
25600 hist_idx: e.htod_i32(&[0])?,
25601 });
25602 }
25603 let state = state_guard.as_mut().expect("state armed above");
25604 {
25608 let _main = e.gpu.enter_main()?;
25609 let Step35TokenGraphState {
25610 logits_stage,
25611 token_d,
25612 ..
25613 } = &mut *state;
25614 e.argmax_token_device_into(logits_stage, token_d, n_vocab)?;
25615 }
25616
25617 if state.graphs.is_empty() {
25622 self.step35_token_graph_build(e, cache, state, bucket_max)?;
25625 }
25626 {
25627 let (b, g) = state.graphs.first_mut().expect("graph built above");
25628 if *b != bucket_max {
25629 g.retarget_bucket(bucket_max)?;
25630 *b = bucket_max;
25631 }
25632 }
25633 let graph = state
25634 .graphs
25635 .first()
25636 .map(|(_, g)| g)
25637 .expect("graph built above");
25638
25639 let tg_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
25640 let t_fence = tg_timing.then(std::time::Instant::now);
25641 {
25646 let fa0 = match &self.layers[0].mixer {
25647 Mixer::Full(fa) => fa,
25648 _ => return Err("step35 token graph expects full-attention layers".into()),
25649 };
25650 let tp0 = fa0
25651 .step_tp_qkv
25652 .as_ref()
25653 .ok_or("step35 token graph lost its TP state")?;
25654 for rank in 0..tp0.runtime.devices().len() {
25655 let engine = tp0
25656 .runtime
25657 .rank_engine(rank)
25658 .ok_or("step35 token graph lost a rank engine")?;
25659 let _main = engine.gpu.enter_main()?;
25660 engine.stream().synchronize()?;
25661 }
25662 }
25663
25664 {
25666 let _main = e.gpu.enter_main()?;
25667 e.set_u32_one(&mut state.token_d, token)?;
25668 e.set_i32_one(&mut state.pos_d, pos as i32)?;
25669 }
25670 let t_launch = tg_timing.then(std::time::Instant::now);
25671 graph.launch(e)?;
25672 let t_book = tg_timing.then(std::time::Instant::now);
25673 for il in 0..n_layers {
25678 let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
25679 let transaction = tp_kv.begin_transaction()?;
25680 let fa = match &self.layers[il].mixer {
25681 Mixer::Full(fa) => fa,
25682 _ => return Err("step35 token graph expects full-attention layers".into()),
25683 };
25684 let tp = fa
25685 .step_tp_qkv
25686 .as_ref()
25687 .ok_or("step35 token graph lost its TP state")?;
25688 let empty: [CudaSlice<f32>; 0] = [];
25691 tp.runtime.append_tp_kv_transaction_inner(
25692 tp_kv,
25693 transaction,
25694 &empty,
25695 &empty,
25696 1,
25697 true,
25698 )?;
25699 tp.runtime
25700 .commit_tp_kv_transaction_external(tp_kv, transaction, 1)?;
25701 if let Some(local) = cache.kv[il].as_mut() {
25703 local.len = pos + 1;
25704 let _main = e.gpu.enter_main()?;
25705 e.set_i32_one(&mut local.len_d, (pos + 1) as i32)?;
25706 }
25707 }
25708 cache.pos = pos + 1;
25709 let t_sync = tg_timing.then(std::time::Instant::now);
25710 let (logits, h_seed) = {
25711 let _main = e.gpu.enter_main()?;
25712 e.stream().synchronize()?;
25713 (e.dtoh(&state.logits_stage)?, e.clone_dtod(&state.x)?)
25714 };
25715 if let (Some(f), Some(l), Some(b), Some(sy)) = (t_fence, t_launch, t_book, t_sync) {
25716 use std::sync::atomic::{AtomicU64, Ordering};
25717 static NS: [AtomicU64; 5] = [
25718 AtomicU64::new(0),
25719 AtomicU64::new(0),
25720 AtomicU64::new(0),
25721 AtomicU64::new(0),
25722 AtomicU64::new(0),
25723 ];
25724 static CALLS: AtomicU64 = AtomicU64::new(0);
25725 let now = std::time::Instant::now();
25726 NS[0].fetch_add((l - f).as_nanos() as u64, Ordering::Relaxed); NS[1].fetch_add((b - l).as_nanos() as u64, Ordering::Relaxed); NS[2].fetch_add((sy - b).as_nanos() as u64, Ordering::Relaxed); NS[3].fetch_add((now - sy).as_nanos() as u64, Ordering::Relaxed); NS[4].fetch_add((now - f).as_nanos() as u64, Ordering::Relaxed); let calls = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
25732 if calls.is_multiple_of(100) {
25733 let avg = |i: usize| NS[i].load(Ordering::Relaxed) as f64 / calls as f64 / 1e3;
25734 eprintln!(
25735 "[tg-timing] calls={calls} fence_us={:.0} launch_us={:.0} book_us={:.0} \
25736 syncdtoh_us={:.0} total_us={:.0}",
25737 avg(0),
25738 avg(1),
25739 avg(2),
25740 avg(3),
25741 avg(4)
25742 );
25743 }
25744 }
25745 if std::env::var("MEMRA_TG_PROBE_LAYER").is_ok() {
25747 use std::io::Write;
25748 let (pm, px) = {
25749 let _main = e.gpu.enter_main()?;
25750 (e.dtoh(&state.probe_mixed)?, e.dtoh(&state.probe_x)?)
25751 };
25752 for (path, data) in [
25753 ("/root/tg-probe-mixed.bin", &pm),
25754 ("/root/tg-probe-x.bin", &px),
25755 ] {
25756 let mut fo = std::fs::OpenOptions::new()
25757 .create(true)
25758 .append(true)
25759 .open(path)?;
25760 for v in data {
25761 fo.write_all(&v.to_le_bytes())?;
25762 }
25763 }
25764 }
25765 if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
25768 let hh = {
25769 let _main = e.gpu.enter_main()?;
25770 e.dtoh(&state.hn)?
25771 };
25772 use std::io::Write;
25773 let mut fo = std::fs::OpenOptions::new()
25774 .create(true)
25775 .append(true)
25776 .open(path)?;
25777 for v in &hh {
25778 fo.write_all(&v.to_le_bytes())?;
25779 }
25780 }
25781 if std::env::var("MEMRA_STEP_TP_GRAPH_DEBUG").as_deref() == Ok("1") {
25784 for il in [0usize, 1, 44] {
25785 let tp_kv = cache.tp_kv[il].as_ref().expect("eligibility checked above");
25786 let host_len = tp_kv.staged_len();
25787 let fa = match &self.layers[il].mixer {
25788 Mixer::Full(fa) => fa,
25789 _ => continue,
25790 };
25791 let tp = fa
25792 .step_tp_qkv
25793 .as_ref()
25794 .ok_or("step35 token graph lost its TP state")?;
25795 for rank in 0..tp.runtime.devices().len() {
25796 let engine = tp
25797 .runtime
25798 .rank_engine(rank)
25799 .ok_or("step35 token graph lost a rank engine")?;
25800 let rank_cache = tp_kv.rank(rank).ok_or("debug rank cache missing")?;
25801 let _main = engine.gpu.enter_main()?;
25802 engine.stream().synchronize()?;
25803 let len_d = engine.dtoh_i32_one(rank_cache.len_d())?;
25804 let base_d = match rank_cache.base_d() {
25805 Some(b) => engine.dtoh_i32_one(b)?,
25806 None => -1,
25807 };
25808 eprintln!(
25809 "[graph-debug] pos={pos} il={il} rank={rank} host_len={host_len} \
25810 len_d={len_d} base_d={base_d}"
25811 );
25812 }
25813 }
25814 }
25815 Ok(Some((logits, h_seed)))
25816 }
25817
25818 pub(crate) fn head_split_matvec(
25824 &self,
25825 e: &Engine,
25826 hn: &CudaSlice<f32>,
25827 ) -> Result<Option<Vec<f32>>, Box<dyn std::error::Error>> {
25828 if self.head_split_fill_device(e, hn)?.is_none() {
25829 return Ok(None);
25830 }
25831 let guard = HEAD_SPLIT_WS
25832 .lock()
25833 .map_err(|_| "head split lock is poisoned")?;
25834 let ws = guard.as_ref().expect("filled above");
25835 let _main = e.gpu.enter_main()?;
25836 Ok(Some(e.dtoh(&ws.logits_e)?))
25837 }
25838
25839 fn head_split_fill_device(
25843 &self,
25844 e: &Engine,
25845 hn: &CudaSlice<f32>,
25846 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
25847 use cudarc::driver::DevicePtr;
25848 let crate::model::GpuTensor::FloatBf16 { data: head, .. } = &self.output else {
25849 return Ok(None);
25850 };
25851 let Some(rank1) = self.layers.first().and_then(|l| match &l.mixer {
25852 Mixer::Full(fa) => fa
25853 .step_tp_qkv
25854 .as_ref()
25855 .and_then(|tp| tp.runtime.rank_engine(1)),
25856 _ => None,
25857 }) else {
25858 return Ok(None);
25859 };
25860 let n_embd = self.cfg.n_embd as usize;
25861 let n_vocab = self.cfg.n_vocab as usize;
25862 let half = n_vocab / 2;
25863 let mut guard = HEAD_SPLIT_WS
25864 .lock()
25865 .map_err(|_| "head split lock is poisoned")?;
25866 let pin = {
25867 let _main = e.gpu.enter_main()?;
25868 let stream = e.stream();
25869 let (ptr, _g) = head.device_ptr(&stream);
25870 ptr
25871 };
25872 if guard.as_ref().is_none_or(|ws| ws.pin != pin) {
25873 let hi_rows = n_vocab - half;
25875 let (w1, hn1, y1, ev_done) = {
25876 let _r1 = rank1.gpu.enter_main()?;
25877 (
25878 rank1.alloc_u8_uninit(hi_rows * n_embd * 2)?,
25879 rank1.htod(&vec![0.0f32; n_embd])?,
25880 rank1.htod(&vec![0.0f32; hi_rows])?,
25881 rank1.ctx().new_event(None)?,
25882 )
25883 };
25884 {
25885 use cudarc::driver::sys;
25886 let src = pin + (half * n_embd * 2) as u64;
25887 let dst = {
25888 let _r1 = rank1.gpu.enter_main()?;
25889 let rstream = rank1.stream();
25890 let (d, _g) = w1.device_ptr(&rstream);
25891 d
25892 };
25893 let _r1 = rank1.gpu.enter_main()?;
25894 let r = unsafe {
25895 sys::cuMemcpyAsync(
25896 dst as sys::CUdeviceptr,
25897 src as sys::CUdeviceptr,
25898 hi_rows * n_embd * 2,
25899 rank1.stream().cu_stream() as sys::CUstream,
25900 )
25901 };
25902 if r != sys::CUresult::CUDA_SUCCESS {
25903 return Err(format!("head split replica upload: {r:?}").into());
25904 }
25905 rank1.stream().synchronize()?;
25906 }
25907 let (logits_e, ev_hn) = {
25908 let _main = e.gpu.enter_main()?;
25909 (e.htod(&vec![0.0f32; n_vocab])?, e.ctx().new_event(None)?)
25910 };
25911 let (raw_hn1, raw_y1) = {
25912 let _r1 = rank1.gpu.enter_main()?;
25913 let rstream = rank1.stream();
25914 let (a, _g0) = hn1.device_ptr(&rstream);
25915 let (b, _g1) = y1.device_ptr(&rstream);
25916 (a, b)
25917 };
25918 let raw_logits_hi = {
25919 let _main = e.gpu.enter_main()?;
25920 let stream = e.stream();
25921 let (l, _g) = logits_e.device_ptr(&stream);
25922 l + (half * 4) as u64
25923 };
25924 *guard = Some(HeadSplit {
25925 pin,
25926 w1,
25927 hn1,
25928 y1,
25929 logits_e,
25930 ev_hn,
25931 ev_done,
25932 raw_hn1,
25933 raw_y1,
25934 raw_logits_hi,
25935 samp: None,
25936 });
25937 }
25938 let ws = guard.as_mut().expect("armed above");
25939 let hi_rows = n_vocab - half;
25940 let raw_hn = {
25942 let _main = e.gpu.enter_main()?;
25943 let stream = e.stream();
25944 let (h, _g) = hn.device_ptr(&stream);
25945 ws.ev_hn.record(&stream)?;
25946 h
25947 };
25948 {
25949 let _r1 = rank1.gpu.enter_main()?;
25950 rank1.stream().wait(&ws.ev_hn)?;
25951 crate::tp::raw_copy_bytes(ws.raw_hn1, raw_hn, n_embd * 4, rank1)?;
25952 let HeadSplit { w1, hn1, y1, .. } = &mut *ws;
25953 rank1.matvec_bf16_into(w1, hn1, y1, n_embd, hi_rows)?;
25954 crate::tp::raw_copy_bytes(ws.raw_logits_hi, ws.raw_y1, hi_rows * 4, rank1)?;
25955 ws.ev_done.record(&rank1.stream())?;
25956 }
25957 {
25958 let _main = e.gpu.enter_main()?;
25959 let head_lo = head.slice(0..half * n_embd * 2);
25960 let HeadSplit { logits_e, .. } = &mut *ws;
25961 e.matvec_bf16_view_into(&head_lo, hn, logits_e, n_embd, half)?;
25963 e.stream().wait(&ws.ev_done)?;
25964 Ok(Some(()))
25965 }
25966 }
25967
25968 pub(crate) fn head_split_argmax_device(
25973 &self,
25974 e: &Engine,
25975 hn: &CudaSlice<f32>,
25976 token_d: &mut CudaSlice<u32>,
25977 ) -> Result<bool, Box<dyn std::error::Error>> {
25978 if self.head_split_fill_device(e, hn)?.is_none() {
25979 return Ok(false);
25980 }
25981 let n_vocab = self.cfg.n_vocab as usize;
25982 let guard = HEAD_SPLIT_WS
25983 .lock()
25984 .map_err(|_| "head split lock is poisoned")?;
25985 let ws = guard.as_ref().expect("filled above");
25986 let _main = e.gpu.enter_main()?;
25987 e.argmax_token_device_into(&ws.logits_e, token_d, n_vocab)?;
25988 Ok(true)
25989 }
25990
25991 pub(crate) fn head_split_sample_device(
25997 &self,
25998 e: &Engine,
25999 hn: &CudaSlice<f32>,
26000 token_d: &mut CudaSlice<u32>,
26001 samp: &crate::decode_batch::DevSamp,
26002 ctr: u32,
26003 ) -> Result<bool, Box<dyn std::error::Error>> {
26004 if self.head_split_fill_device(e, hn)?.is_none() {
26005 return Ok(false);
26006 }
26007 let n_vocab = self.cfg.n_vocab as usize;
26008 let guard = HEAD_SPLIT_WS
26009 .lock()
26010 .map_err(|_| "head split lock is poisoned")?;
26011 let mut guard = guard;
26012 let ws = guard.as_mut().expect("filled above");
26013 let _main = e.gpu.enter_main()?;
26014 if ws.samp.is_none() {
26015 ws.samp = Some(SampScratch {
26016 pb: e.zeros(n_vocab)?,
26017 th: e.zeros(1)?,
26018 z: e.zeros(1)?,
26019 mx: e.zeros(1)?,
26020 rows: e.htod_i32(&[0i32])?,
26021 });
26022 }
26023 let filtered = samp.top_k > 0 || samp.top_p < 1.0 || samp.min_p > 0.0;
26024 let HeadSplit {
26025 logits_e,
26026 samp: scratch,
26027 ..
26028 } = &mut *ws;
26029 let sc = scratch.as_mut().expect("armed above");
26030 if filtered {
26031 e.filter_stats(
26032 logits_e, n_vocab, &sc.rows, &mut sc.th, &mut sc.z, &mut sc.mx, n_vocab, 1,
26033 samp.temp, samp.top_k, samp.top_p, samp.min_p,
26034 )?;
26035 let SampScratch { pb, th, mx, .. } = sc;
26036 e.gumbel_perturb_filtered_col(
26037 logits_e, 0, pb, n_vocab, samp.seed, ctr, samp.temp, mx, th, 0,
26038 )?;
26039 } else {
26040 e.gumbel_perturb_col(logits_e, 0, &mut sc.pb, n_vocab, samp.seed, ctr, samp.temp)?;
26041 }
26042 e.argmax_token_device_col(&sc.pb, 0, n_vocab, token_d, 0)?;
26043 Ok(true)
26044 }
26045
26046 pub(crate) fn head_split_logits_dtoh(
26049 &self,
26050 e: &Engine,
26051 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
26052 let guard = HEAD_SPLIT_WS
26053 .lock()
26054 .map_err(|_| "head split lock is poisoned")?;
26055 let ws = guard.as_ref().ok_or("head split logits not armed")?;
26056 let _main = e.gpu.enter_main()?;
26057 e.dtoh(&ws.logits_e)
26058 }
26059
26060 #[allow(clippy::type_complexity)] pub fn step35_token_graph_chunk(
26070 &self,
26071 e: &Engine,
26072 token: u32,
26073 k_target: usize,
26074 cache: &mut Cache,
26075 ) -> Result<Option<(Vec<u32>, Vec<f32>)>, Box<dyn std::error::Error>> {
26076 if !self.uses_sliding_gated_moe_program()
26077 || !crate::tp::step_tp_graph_enabled()?
26078 || !crate::tp::step_tp_dcw_enabled()?
26079 || !crate::tp::step_tp_qkv_fused_enabled()?
26080 || !crate::tp::step_tp_dev_router_enabled()?
26081 || !crate::tp::step_nvfp4_dev_routes_enabled()?
26082 {
26083 return Ok(None);
26084 }
26085 if !crate::spec::graph_launch_headroom_ok(e) {
26088 static NOTED: std::sync::Once = std::sync::Once::new();
26089 NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
26090 return Ok(None);
26091 }
26092 let n_layers = self.layers.len();
26093 let pos = cache.pos;
26094 let staged_next = pos + 1;
26095 if staged_next < 96 {
26096 return Ok(None);
26097 }
26098 let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
26101 if !fa_vec {
26102 return Ok(None);
26103 }
26104 let sp = crate::fa_split_keys(staged_next, 8);
26105 let bucket_max = (n_splits * sp).max(staged_next);
26106 let to_boundary = bucket_max.saturating_sub(staged_next) + 1;
26107 let mut k = k_target.min(to_boundary).min(16);
26108 if k < 2 {
26109 return Ok(None);
26110 }
26111 for il in 0..n_layers {
26113 let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
26114 return Ok(None);
26115 };
26116 while k >= 2 && tp_kv.peek_append_ring(k)?.1 {
26117 k -= 1;
26118 }
26119 if k < 2 {
26120 return Ok(None);
26121 }
26122 }
26123
26124 let mut state_guard = self
26125 .step35_token_graph
26126 .lock()
26127 .map_err(|_| "step35 token graph lock is poisoned")?;
26128 let Some(state) = state_guard.as_mut() else {
26129 return Ok(None); };
26131 if state.graphs.is_empty() {
26132 return Ok(None);
26133 }
26134 {
26135 let (b, g) = state.graphs.first_mut().expect("checked above");
26136 if *b != bucket_max {
26137 g.retarget_bucket(bucket_max)?;
26138 *b = bucket_max;
26139 }
26140 }
26141 let graph = state.graphs.first().map(|(_, g)| g).expect("checked above");
26142
26143 {
26145 let fa0 = match &self.layers[0].mixer {
26146 Mixer::Full(fa) => fa,
26147 _ => return Err("step35 token graph expects full-attention layers".into()),
26148 };
26149 let tp0 = fa0
26150 .step_tp_qkv
26151 .as_ref()
26152 .ok_or("step35 token graph lost its TP state")?;
26153 for rank in 0..tp0.runtime.devices().len() {
26154 let engine = tp0
26155 .runtime
26156 .rank_engine(rank)
26157 .ok_or("step35 token graph lost a rank engine")?;
26158 let _main = engine.gpu.enter_main()?;
26159 engine.stream().synchronize()?;
26160 }
26161 }
26162
26163 {
26166 let _main = e.gpu.enter_main()?;
26167 e.set_u32_one(&mut state.token_d, token)?;
26168 e.set_i32_one(&mut state.pos_d, pos as i32)?;
26169 e.set_i32_one(&mut state.hist_idx, 0)?;
26170 }
26171 for _ in 0..k {
26172 graph.launch(e)?;
26173 }
26174 for il in 0..n_layers {
26176 let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
26177 let transaction = tp_kv.begin_transaction()?;
26178 let fa = match &self.layers[il].mixer {
26179 Mixer::Full(fa) => fa,
26180 _ => return Err("step35 token graph expects full-attention layers".into()),
26181 };
26182 let tp = fa
26183 .step_tp_qkv
26184 .as_ref()
26185 .ok_or("step35 token graph lost its TP state")?;
26186 let empty: [CudaSlice<f32>; 0] = [];
26187 tp.runtime.append_tp_kv_transaction_inner(
26188 tp_kv,
26189 transaction,
26190 &empty,
26191 &empty,
26192 k,
26193 true,
26194 )?;
26195 tp.runtime
26196 .commit_tp_kv_transaction_external(tp_kv, transaction, k)?;
26197 if let Some(local) = cache.kv[il].as_mut() {
26198 local.len = pos + k;
26199 let _main = e.gpu.enter_main()?;
26200 e.set_i32_one(&mut local.len_d, (pos + k) as i32)?;
26201 }
26202 }
26203 cache.pos = pos + k;
26204 let (hist, logits) = {
26205 let _main = e.gpu.enter_main()?;
26206 e.stream().synchronize()?;
26207 (e.dtoh_u32(&state.token_hist)?, e.dtoh(&state.logits_stage)?)
26208 };
26209 Ok(Some((hist[..k].to_vec(), logits)))
26210 }
26211}
26212
26213impl HybridModel {
26214 #[allow(clippy::too_many_arguments)]
26219 fn step35_token_graph_build(
26220 &self,
26221 e: &Engine,
26222 cache: &mut Cache,
26223 state: &mut Step35TokenGraphState,
26224 bucket_max: usize,
26225 ) -> Result<(), Box<dyn std::error::Error>> {
26226 use cudarc::driver::DevicePtr;
26227 let n_embd = self.cfg.n_embd as usize;
26228 let eps = self.cfg.rms_eps;
26229 let n_layers = self.layers.len();
26230 let started = std::time::Instant::now();
26231 if !crate::router_kernel_on() {
26232 return Err(
26233 "step35 token graph requires the router kernel (MEMRA_ROUTER_KERNEL=0)".into(),
26234 );
26235 }
26236 if !Engine::bf16_mmv_on() || !n_embd.is_multiple_of(8) {
26237 return Err("step35 token graph requires MEMRA_BF16_MMV bf16-resident matvecs".into());
26238 }
26239
26240 let embd_gpu = self
26242 .embd_gpu_try(e)
26243 .ok_or("step35 token graph could not upload the device embed table")?;
26244 let embd_qtype = match self.embd.ggml_type {
26245 memra_gguf::GgmlType::BF16 => crate::QT_BF16,
26246 memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
26247 other => return Err(format!("token graph embed dtype {other:?} unhandled").into()),
26248 };
26249 let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
26250
26251 let (p_mixed, p_kshadow, p_vshadow) = {
26253 let _main = e.gpu.enter_main()?;
26254 let stream = e.stream();
26255 let (a, _g) = state.mixed_stage.device_ptr(&stream);
26256 let (b, _g) = state.k_shadow_stage.device_ptr(&stream);
26257 let (c, _g) = state.v_shadow_stage.device_ptr(&stream);
26258 (a, b, c)
26259 };
26260
26261 crate::tp::token_graph_build_begin()?;
26262 let mut group_id: u32 = 0;
26263 for il in 0..n_layers {
26264 let layer = &self.layers[il];
26265 let fa = match &layer.mixer {
26266 Mixer::Full(fa) => fa,
26267 _ => return Err("step35 token graph expects full-attention layers".into()),
26268 };
26269 let tp = fa
26270 .step_tp_qkv
26271 .as_ref()
26272 .ok_or("step35 token graph lost its TP state")?;
26273 let attention = tp
26274 .attention
26275 .as_ref()
26276 .ok_or("step35 token graph lost its attention aux")?;
26277 let geometry = self.step35_geom(il);
26278 let window = geometry.window.map(|w| w as usize);
26279 let head_dim = geometry.head_dim_k as usize;
26280 let heads = geometry.n_head as usize;
26281 let kv_heads = geometry.n_head_kv as usize;
26282 let ranks = tp.runtime.devices().len();
26283 let local_heads = heads / ranks;
26284 let local_kv_heads = kv_heads / ranks;
26285 let layer_bucket = window.map(|w| bucket_max.min(w)).unwrap_or(bucket_max);
26286 let use_gate_shards =
26287 attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some();
26288 if !use_gate_shards {
26289 return Err("step35 token graph requires the fused gate shards".into());
26290 }
26291
26292 let ws_index = tp
26293 .runtime
26294 .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
26295 let ws_mutex = tp.runtime.decode_v2_workspace();
26296 let mut ws_guard = ws_mutex
26297 .lock()
26298 .map_err(|_| "step TP decode v2 workspace lock is poisoned")?;
26299 let ws = ws_guard
26300 .get_mut(ws_index)
26301 .ok_or("step TP decode v2 workspace missing after ensure")?;
26302 tp.runtime
26303 .decode_v2_arm_token_mirrors(ws, p_mixed, (p_kshadow, p_vshadow))?;
26304 let mut rope_freqs = Vec::with_capacity(ranks);
26305 for rank in 0..ranks {
26306 let engine = tp
26307 .runtime
26308 .rank_engine(rank)
26309 .ok_or("step35 token graph lost a rank engine")?;
26310 rope_freqs.push(if geometry.rope_factors {
26311 self.step35_aux
26312 .as_ref()
26313 .and_then(|aux| aux.rope_freqs(engine))
26314 } else {
26315 None
26316 });
26317 }
26318 let gate_shards_arg = if let Some(shards) = attention.gate_shards.as_deref() {
26319 Some(crate::tp::StepTpGateShards::F32(shards))
26320 } else {
26321 attention
26322 .gate_shards_bf16
26323 .as_deref()
26324 .map(crate::tp::StepTpGateShards::Bf16)
26325 };
26326
26327 let decode_input = attention
26329 .decode_input
26330 .as_ref()
26331 .ok_or("step35 token graph requires the replicated decode input")?;
26332 let mut decode_input = decode_input
26333 .lock()
26334 .map_err(|_| "replicated decode input lock is poisoned")?;
26335 if ws.h_stage.is_none() {
26337 return Err(
26338 "step35 token graph requires the stage flow armed (run eager dcw first)".into(),
26339 );
26340 }
26341 {
26342 let state_x = &mut state.x;
26343 let token_d = &state.token_d;
26344 let pos_d = &state.pos_d;
26345 crate::tp::graph_section(e, None, || {
26346 let _main = e.gpu.enter_main()?;
26347 if il == 0 {
26348 e.embed_gather_device_into(
26349 embd_gpu,
26350 token_d,
26351 state_x,
26352 n_embd,
26353 embd_qtype,
26354 embd_row_bytes,
26355 )?;
26356 }
26357 {
26358 let h_stage = ws.h_stage.as_mut().expect("stage armed checked above");
26359 e.rms_norm(
26360 state_x,
26361 layer.attn_norm.float_data(),
26362 h_stage,
26363 n_embd,
26364 1,
26365 eps,
26366 )?;
26367 }
26368 {
26369 let pos_stage = ws.pos_stage.as_mut().expect("stage armed above");
26370 let mut dst = pos_stage.slice_mut(0..1);
26371 e.stream().memcpy_dtod(&pos_d.slice(0..1), &mut dst)?;
26372 }
26373 Ok(())
26374 })?;
26375 }
26376
26377 group_id += 1;
26379 for rank in 0..ranks {
26380 let engine = tp
26381 .runtime
26382 .rank_engine(rank)
26383 .ok_or("step35 token graph lost a rank engine")?;
26384 {
26385 let ceiling = window
26390 .map(|w| cache.max_ctx.min(w))
26391 .unwrap_or(cache.max_ctx);
26392 let _main = engine.gpu.enter_main()?;
26393 engine.fa_dcw_pool_ensure(
26394 head_dim,
26395 local_heads,
26396 local_kv_heads,
26397 ceiling.min(2048),
26398 )?;
26399 engine.fa_dcw_pool_ensure(head_dim, local_heads, local_kv_heads, ceiling)?;
26400 engine.fa_dcw_pool_ensure(
26401 head_dim,
26402 local_heads,
26403 local_kv_heads,
26404 layer_bucket,
26405 )?;
26406 }
26407 let runtime = &tp.runtime;
26408 let q_norm = &attention.q_norm;
26409 let k_norm = &attention.k_norm;
26410 let gate_ref = gate_shards_arg.as_ref();
26411 crate::tp::graph_section(engine, Some(group_id), || {
26412 runtime.decode_v2_input_qkv_rank(
26413 ws,
26414 &state.pos_d,
26415 &mut decode_input,
26416 &tp.q,
26417 &tp.k,
26418 &tp.v,
26419 q_norm,
26420 k_norm,
26421 head_dim,
26422 geometry.n_rot as usize,
26423 geometry.rope_base,
26424 &rope_freqs,
26425 eps,
26426 gate_ref,
26427 true,
26428 true,
26429 false,
26430 rank,
26431 None,
26432 )?;
26433 let distributed = cache.tp_kv[il]
26436 .as_mut()
26437 .ok_or("step35 token graph lost a TP cache")?;
26438 let (kv_dim_k, kv_dim_v) = (distributed.kv_dim_k(), distributed.kv_dim_v());
26439 let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
26440 let capacity = distributed.physical_capacity();
26441 {
26442 let rank_cache = distributed
26443 .rank_mut(rank)
26444 .ok_or("step35 token graph lost a rank cache")?;
26445 let (k_plane, v_plane, len_d, base_d) =
26446 rank_cache.planes_and_counters_mut();
26447 engine.append_kv_quantized_dcw(
26448 &ws.k[rank],
26449 &ws.v_raw[rank],
26450 k_plane,
26451 v_plane,
26452 len_d,
26453 base_d,
26454 kv_dim_k,
26455 kv_dim_v,
26456 ktb,
26457 vtb,
26458 )?;
26459 }
26460 {
26461 let rank_cache = distributed
26462 .rank_mut(rank)
26463 .ok_or("step35 token graph lost a rank cache")?;
26464 engine.inc_i32(rank_cache.len_d_mut())?;
26465 }
26466 let rank_cache = distributed
26467 .rank(rank)
26468 .ok_or("step35 token graph lost a rank cache")?;
26469 let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * ktb);
26470 let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * vtb);
26471 engine.fa_decode_dcw(
26476 &ws.q[rank],
26477 &k_ring,
26478 &v_ring,
26479 &mut ws.attn_out[rank],
26480 head_dim,
26481 local_heads,
26482 local_kv_heads,
26483 rank_cache.len_d(),
26484 rank_cache.base_d(),
26485 window.unwrap_or(0),
26486 layer_bucket,
26487 geometry.attention_scale(),
26488 ktb,
26489 vtb,
26490 None,
26491 )?;
26492 engine.attn_head_gate(
26493 &ws.attn_out[rank],
26494 &ws.gate[rank],
26495 &mut ws.gated[rank],
26496 None,
26497 head_dim,
26498 local_heads,
26499 1,
26500 )?;
26501 runtime.decode_v2_finish_rank_partial(ws, &tp.o, true, rank)?;
26502 Ok(())
26503 })?;
26504 }
26505
26506 {
26508 let root = tp
26509 .runtime
26510 .rank_engine(0)
26511 .ok_or("step35 token graph lost the root engine")?;
26512 let runtime = &tp.runtime;
26513 crate::tp::graph_section(root, None, || runtime.decode_v2_finish_root_fused(ws))?;
26514 }
26515 drop(ws_guard);
26516 drop(decode_input);
26517
26518 let probe_layer: Option<usize> = std::env::var("MEMRA_TG_PROBE_LAYER")
26519 .ok()
26520 .and_then(|v| v.parse().ok());
26521 if probe_layer == Some(il) {
26522 let Step35TokenGraphState {
26523 mixed_stage,
26524 probe_mixed,
26525 ..
26526 } = &mut *state;
26527 crate::tp::graph_section(e, None, || {
26528 let _main = e.gpu.enter_main()?;
26529 let mut dst = probe_mixed.slice_mut(0..n_embd);
26530 e.stream()
26531 .memcpy_dtod(&mixed_stage.slice(0..n_embd), &mut dst)?;
26532 Ok(())
26533 })?;
26534 }
26535
26536 match &layer.ffn {
26538 crate::hybrid::Ffn::Dense {
26539 ffn_gate,
26540 ffn_up,
26541 ffn_down,
26542 } => {
26543 let n_ff = ffn_gate.out_features();
26544 let lim = self.cfg.clamp_shexp_at(il as u32);
26545 if lim.is_some() {
26549 return Err("step35 token graph dense FFN with clamp unsupported".into());
26550 }
26551 let (wg_d, wu_d, wd_d) = match (ffn_gate, ffn_up, ffn_down) {
26552 (
26553 crate::model::GpuTensor::FloatBf16 { data: wg, .. },
26554 crate::model::GpuTensor::FloatBf16 { data: wu, .. },
26555 crate::model::GpuTensor::FloatBf16 { data: wd, .. },
26556 ) => (wg, wu, wd),
26557 _ => {
26558 return Err(
26559 "step35 token graph dense FFN requires bf16-resident weights"
26560 .into(),
26561 );
26562 }
26563 };
26564 crate::tp::graph_section(e, None, || {
26565 let _main = e.gpu.enter_main()?;
26566 let Step35TokenGraphState {
26567 x,
26568 x1,
26569 mixed_stage,
26570 dense_z,
26571 dense_gate,
26572 dense_up,
26573 dense_act,
26574 sh_stage,
26575 ..
26576 } = &mut *state;
26577 e.add_rms_norm(
26578 x,
26579 mixed_stage,
26580 layer.post_attn_norm.float_data(),
26581 x1,
26582 dense_z,
26583 n_embd,
26584 1,
26585 eps,
26586 )?;
26587 e.matvec_bf16_into(wg_d, dense_z, dense_gate, n_embd, n_ff)?;
26591 e.matvec_bf16_into(wu_d, dense_z, dense_up, n_embd, n_ff)?;
26592 Self::ffn_act_lim(
26593 e, &self.cfg, dense_gate, dense_up, 1.0, 1.0, lim, dense_act, n_ff,
26594 )?;
26595 e.matvec_bf16_into(wd_d, dense_act, sh_stage, n_ff, n_embd)?;
26596 e.add(x1, sh_stage, x, n_embd)?;
26597 Ok(())
26598 })?;
26599 }
26600 crate::hybrid::Ffn::Moe(m) => {
26601 let moe = self
26602 .cfg
26603 .moe
26604 .as_ref()
26605 .ok_or("step35 token graph needs moe cfg")?;
26606 let n_expert = moe.expert_count as usize;
26607 let n_used = moe.expert_used_count as usize;
26608 let sigmoid = self
26609 .cfg
26610 .sigmoid_router()
26611 .ok_or("step35 token graph needs the sigmoid router")?;
26612 let step_tp = m
26613 .step_tp
26614 .as_ref()
26615 .ok_or("step35 token graph needs TP experts")?;
26616 let bank = match &step_tp.experts {
26617 crate::hybrid::StepTpExpertBank::Nvfp4(bank) => bank,
26618 _ => return Err("step35 token graph needs the NVFP4 bank".into()),
26619 };
26620 let routes_ws_mutex = bank.device_workspace_handle();
26621 let mut routes_guard = routes_ws_mutex
26622 .lock()
26623 .map_err(|_| "routes workspace lock is poisoned")?;
26624 let routes_ws = routes_guard
26625 .as_mut()
26626 .ok_or("step35 token graph requires the routes workspace warmed")?;
26627 routes_ws.arm_stages(e, bank.input_width, n_used)?;
26628 step_tp.runtime.routes_arm_raw(bank, routes_ws)?;
26629 let p_z = {
26630 let root = step_tp
26631 .runtime
26632 .rank_engine(0)
26633 .ok_or("routes root engine missing")?;
26634 let _main = root.gpu.enter_main()?;
26635 let stream = root.stream();
26636 let in_stage = routes_ws
26637 .in_stage_handle()
26638 .ok_or("routes in stage not armed")?;
26639 let (a, _g) = in_stage.device_ptr(&stream);
26640 a
26641 };
26642 let local_out = bank.expert_width / ranks;
26643
26644 crate::tp::graph_section(e, None, || {
26646 let _main = e.gpu.enter_main()?;
26647 {
26648 let in_stage = routes_ws
26649 .in_stage_mut()
26650 .ok_or("routes in stage not armed")?;
26651 let Step35TokenGraphState {
26652 x, x1, mixed_stage, ..
26653 } = &mut *state;
26654 e.add_rms_norm(
26655 x,
26656 mixed_stage,
26657 layer.post_attn_norm.float_data(),
26658 x1,
26659 in_stage,
26660 n_embd,
26661 1,
26662 eps,
26663 )?;
26664 }
26665 {
26666 let z_ref = routes_ws
26667 .in_stage_handle()
26668 .ok_or("routes in stage not armed")?;
26669 e.router_gemv_into(
26670 m.gate_inp.float_data(),
26671 z_ref,
26672 &mut state.router_logits,
26673 n_embd,
26674 n_expert,
26675 1,
26676 )?;
26677 }
26678 let (sel_e, w_e) = routes_ws
26679 .dev_route_e_mut()
26680 .ok_or("routes staging not armed")?;
26681 e.moe_router_sigmoid_topk_into(
26682 &state.router_logits,
26683 1,
26684 n_expert,
26685 n_used,
26686 m.active_count(),
26687 &m.exp_probs_b_dev,
26688 &m.active_experts_dev,
26689 sigmoid.0,
26690 sigmoid.1,
26691 sel_e,
26692 w_e,
26693 )?;
26694 Ok(())
26695 })?;
26696
26697 group_id += 1;
26699 for rank in 0..ranks {
26700 let engine = step_tp
26701 .runtime
26702 .rank_engine(rank)
26703 .ok_or("routes rank engine missing")?;
26704 let runtime = &step_tp.runtime;
26705 crate::tp::graph_section(engine, Some(group_id), || {
26706 runtime.routes_rank_section(
26707 bank,
26708 routes_ws,
26709 p_z,
26710 local_out,
26711 n_used,
26712 step_tp.activation_limit,
26713 rank,
26714 )
26715 })?;
26716 }
26717
26718 {
26720 let root = step_tp
26721 .runtime
26722 .rank_engine(0)
26723 .ok_or("routes root engine missing")?;
26724 let runtime = &step_tp.runtime;
26725 crate::tp::graph_section(root, None, || {
26726 runtime.routes_root_section(bank, routes_ws)
26727 })?;
26728 }
26729
26730 let lim_sh = self.cfg.clamp_shexp_at(il as u32);
26734 let (wg_sh, wu_sh, wd_sh) = match (&m.gate_shexp, &m.up_shexp, &m.down_shexp) {
26735 (
26736 Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
26737 Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
26738 Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
26739 ) => (wg, wu, wd),
26740 _ => {
26741 return Err(
26742 "step35 token graph shexp requires bf16-resident weights".into()
26743 );
26744 }
26745 };
26746 let n_ff_sh = m
26747 .gate_shexp
26748 .as_ref()
26749 .expect("matched Some above")
26750 .out_features();
26751 let gate_inp_shexp = m.gate_inp_shexp.as_ref();
26754 crate::tp::graph_section(e, None, || {
26755 let _main = e.gpu.enter_main()?;
26756 let (z_ref, out_stage) = routes_ws
26757 .in_and_out_stages_mut()
26758 .ok_or("routes stages not armed")?;
26759 let Step35TokenGraphState {
26760 x,
26761 x1,
26762 sh_stage,
26763 shexp_gate,
26764 shexp_up,
26765 shexp_act,
26766 gate_sig,
26767 ..
26768 } = &mut *state;
26769 e.matvec_bf16_dual_into(
26770 wg_sh, wu_sh, z_ref, shexp_gate, shexp_up, n_embd, n_ff_sh,
26771 )?;
26772 Self::ffn_act_lim(
26773 e, &self.cfg, shexp_gate, shexp_up, 1.0, 1.0, lim_sh, shexp_act,
26774 n_ff_sh,
26775 )?;
26776 e.matvec_bf16_into(wd_sh, shexp_act, sh_stage, n_ff_sh, n_embd)?;
26777 if let Some(gate_w) = gate_inp_shexp {
26778 e.sigmoid_dot_rows_into(
26779 z_ref,
26780 gate_w.float_data(),
26781 gate_sig,
26782 n_embd,
26783 1,
26784 )?;
26785 }
26786 e.add_scaled_rows(sh_stage, gate_sig, out_stage, n_embd, 1)?;
26787 e.add(x1, out_stage, x, n_embd)?;
26788 Ok(())
26789 })?;
26790 }
26791 }
26792 if probe_layer == Some(il) {
26793 let Step35TokenGraphState { x, probe_x, .. } = &mut *state;
26794 crate::tp::graph_section(e, None, || {
26795 let _main = e.gpu.enter_main()?;
26796 let mut dst = probe_x.slice_mut(0..n_embd);
26797 e.stream().memcpy_dtod(&x.slice(0..n_embd), &mut dst)?;
26798 Ok(())
26799 })?;
26800 }
26801 }
26802
26803 let head = match &self.output {
26805 crate::model::GpuTensor::FloatBf16 { data, .. } => data,
26806 _ => return Err("step35 token graph head requires the bf16-resident output".into()),
26807 };
26808 crate::tp::graph_section(e, None, || {
26809 let _main = e.gpu.enter_main()?;
26810 let Step35TokenGraphState {
26811 x,
26812 hn,
26813 logits_stage,
26814 token_d,
26815 pos_d,
26816 token_hist,
26817 hist_idx,
26818 ..
26819 } = &mut *state;
26820 e.rms_norm(x, self.output_norm.float_data(), hn, n_embd, 1, eps)?;
26821 e.matvec_bf16_into(head, hn, logits_stage, n_embd, self.cfg.n_vocab as usize)?;
26822 e.argmax_token_device_into(logits_stage, token_d, self.cfg.n_vocab as usize)?;
26828 e.u32_hist_append(token_d, token_hist, hist_idx)?;
26829 e.inc_i32(pos_d)?;
26830 Ok(())
26831 })?;
26832
26833 let graph = crate::tp::token_graph_build_finish()?;
26834 state.graphs.push((bucket_max, graph));
26835 eprintln!(
26836 "[step35-token-graph] built bucket={bucket_max} layers={n_layers} \
26837 build_ms={:.0} performance_claim=false",
26838 started.elapsed().as_secs_f64() * 1e3
26839 );
26840 Ok(())
26841 }
26842}