1#![allow(clippy::needless_range_loop)]
55
56use std::ops::Range;
57use std::sync::Arc;
58use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
59
60use cudarc::driver::CudaSlice;
61
62use crate::Engine;
63use crate::kda::{ConvArm, KdaAttnLayer};
64use crate::model::GpuTensor;
65use memra_kv::{Cache, LatentKvLayer, RecurLayer};
66
67pub const GLM5_TP_ALLOWED_RANKS: [usize; 2] = [2, 4];
78
79pub const GLM5_TP_TRANSPORT_TAG: &str = "glm5-tp-transport";
85
86pub type Glm5TpLayerSpec = crate::tp::StepEpLayerSpec;
87
88pub fn glm5_tp_env_raw() -> Option<String> {
94 std::env::var("MEMRA_GLM5_TP").ok()
95}
96
97pub fn glm5_tp_armed() -> bool {
101 matches!(glm5_tp_env_raw().as_deref(), Some(v) if !v.is_empty() && v != "0")
102}
103
104pub fn parse_glm5_tp_layer_specs(
108 value: Option<&str>,
109 trunk_layers: usize,
110) -> Result<Vec<Glm5TpLayerSpec>, String> {
111 crate::tp::parse_layer_specs_for_trunk("MEMRA_GLM5_TP", value, Some(trunk_layers))
112}
113
114pub fn gate_same_device() -> bool {
120 std::env::var("MEMRA_GLM5_TP_GATE_SAME_DEV").as_deref() == Ok("1")
121}
122
123#[derive(Clone, Copy, PartialEq, Eq, Debug)]
137pub enum GateRed {
138 SwapWo,
139 SwapEpGateUp,
140 SkipPeerCombine,
141 CorruptEpMap,
142}
143
144pub fn gate_red() -> Result<Option<GateRed>, String> {
145 match std::env::var("MEMRA_GLM5_TP_GATE_RED").ok().as_deref() {
146 None | Some("") => Ok(None),
147 Some("swap-wo") => Ok(Some(GateRed::SwapWo)),
148 Some("swap-ep-gateup") => Ok(Some(GateRed::SwapEpGateUp)),
149 Some("skip-peer-combine") => Ok(Some(GateRed::SkipPeerCombine)),
150 Some("corrupt-ep-map") => Ok(Some(GateRed::CorruptEpMap)),
151 Some(other) => Err(format!(
152 "MEMRA_GLM5_TP_GATE_RED={other:?} is not a known red arm \
153 (swap-wo | swap-ep-gateup | skip-peer-combine | corrupt-ep-map)"
154 )),
155 }
156}
157
158pub struct Glm5TpRt {
166 pub peers: Vec<Engine>,
167 pub root_dev: usize,
168 pub peer_devs: Vec<usize>,
169 pub same_device_gate: bool,
173 pub transport: crate::tp_transport::TpTransport,
177 link: Option<crate::tp_transport::PeerPullLink>,
180}
181
182impl Glm5TpRt {
183 pub fn new(devices: &[usize]) -> Result<Self, Box<dyn std::error::Error>> {
184 let root_dev = devices[0];
185 let peer_devs: Vec<usize> = devices[1..].to_vec();
186 for &d in &peer_devs {
187 if d == root_dev || peer_devs.iter().filter(|&&x| x == d).count() > 1 {
188 return Err(format!(
189 "MEMRA_GLM5_TP rank devices must be distinct in serving; got {devices:?} \
190 (the same-device form exists only for the rig gate binary)"
191 )
192 .into());
193 }
194 }
195 let mut peers = Vec::with_capacity(peer_devs.len());
196 for &d in &peer_devs {
197 peers.push(Engine::new(d)?);
198 }
199 Ok(Self {
200 peers,
201 root_dev,
202 peer_devs,
203 same_device_gate: false,
204 transport: crate::tp_transport::TpTransport::HostCanonical,
205 link: None,
206 })
207 }
208
209 pub fn new_gate_same_device(
215 root_dev: usize,
216 ranks: usize,
217 ) -> Result<Self, Box<dyn std::error::Error>> {
218 let mut peers = Vec::with_capacity(ranks - 1);
219 for _ in 1..ranks {
220 peers.push(Engine::new(root_dev)?);
221 }
222 Ok(Self {
223 peers,
224 root_dev,
225 peer_devs: vec![root_dev; ranks - 1],
226 same_device_gate: true,
227 transport: crate::tp_transport::TpTransport::HostCanonical,
228 link: None,
229 })
230 }
231
232 pub fn ranks(&self) -> usize {
234 self.peers.len() + 1
235 }
236
237 pub fn arm_transport(&mut self, root: &Engine) -> Result<(), Box<dyn std::error::Error>> {
242 let (transport, armed_flag) = crate::tp_transport::transport_env()?;
248 let engines: Vec<&Engine> = std::iter::once(root).chain(self.peers.iter()).collect();
249 let link = crate::tp_transport::arm_transport(
250 transport,
251 armed_flag,
252 GLM5_TP_TRANSPORT_TAG,
253 &engines,
254 self.same_device_gate,
255 )?;
256 self.transport = transport;
257 self.link = link;
258 Ok(())
259 }
260
261 pub fn hop<'a>(&'a self, root: &'a Engine) -> crate::tp_transport::Hop<'a> {
265 crate::tp_transport::Hop {
266 engines: std::iter::once(root).chain(self.peers.iter()).collect(),
267 transport: self.transport,
268 link: self.link.as_ref(),
269 }
270 }
271}
272
273pub struct Glm5TpModelView {
281 pub trunk_layers: usize,
282 pub layer_class: Vec<Glm5LayerClass>,
284 pub layer_is_moe: Vec<bool>,
286 pub kda_heads: usize,
287 pub kda_head_dim: usize,
288 pub mla_heads: usize,
289 pub n_routed_experts: usize,
290 pub top_k: usize,
291}
292
293#[derive(Clone, Copy, PartialEq, Eq, Debug)]
294pub enum Glm5LayerClass {
295 Kda,
296 Mla,
297}
298
299pub struct Glm5TpLoadPlan {
303 pub rt: Arc<Glm5TpRt>,
304 pub layers: std::collections::BTreeSet<usize>,
305 pub ep_map: Option<crate::ep_map::EpMap>,
306}
307
308fn load_glm5_ep_map(
315 view: &Glm5TpModelView,
316 layers: &std::collections::BTreeSet<usize>,
317 ranks: usize,
318) -> Result<Option<crate::ep_map::EpMap>, Box<dyn std::error::Error>> {
319 let Some((flag, path)) = crate::ep_map::ep_map_env()? else {
320 return Ok(None);
321 };
322 if path.is_empty() {
323 return Err(format!(
324 "{flag} is set but empty (fail-closed: unset the flag for \
325 the even split; an empty value never silently means default)"
326 )
327 .into());
328 }
329 let text = std::fs::read_to_string(&path)
330 .map_err(|e| format!("{flag}={path}: cannot read the map file ({e}) — refused by name"))?;
331 let map = crate::ep_map::EpMap::parse(&text).map_err(|e| format!("{flag}={path}: {e}"))?;
332 if map.ranks != ranks {
333 return Err(format!(
334 "{flag}={path}: map declares ranks={}, this load is TP-{ranks} \
335 (re-mint the map for the armed rank count)",
336 map.ranks
337 )
338 .into());
339 }
340 if map.n_experts != view.n_routed_experts {
341 return Err(format!(
342 "{flag}={path}: map declares expert_count={}, the model routes {}",
343 map.n_experts, view.n_routed_experts
344 )
345 .into());
346 }
347 if map.entry_rank != 0 {
348 return Err(format!(
349 "{flag}={path}: entry_rank={} but the glm5 TP first-hop card is \
350 rank 0 (root: router + combine + shared expert) — re-mint with \
351 --entry-rank 0 (refused rather than silently remapping ranks)",
352 map.entry_rank
353 )
354 .into());
355 }
356 let ep_layers: Vec<usize> = layers
357 .iter()
358 .copied()
359 .filter(|&il| view.layer_is_moe[il])
360 .collect();
361 map.validate_layer_cover(&ep_layers)
362 .map_err(|e| format!("{flag}={path}: {e}"))?;
363 let digest = {
365 use sha2::{Digest, Sha256};
366 let mut h = Sha256::new();
367 h.update(text.as_bytes());
368 let out = h.finalize();
369 out.iter().map(|b| format!("{b:02x}")).collect::<String>()
370 };
371 eprintln!(
372 "[glm5-tp-preflight] ep-map armed path={path} sha256={digest} layers={} \
373 experts={} ranks={} entry_rank={} performance_claim=false",
374 map.layers.len(),
375 map.n_experts,
376 map.ranks,
377 map.entry_rank,
378 );
379 Ok(Some(map))
380}
381
382pub const GLM5_TP_REFUSED_DOOR_FLAGS: [(&str, &str); 4] = [
392 (
393 "MEMRA_HC_FUSED_PRE",
394 "the fused mHC pre-chain is gated on the unsharded walk only",
395 ),
396 (
397 "MEMRA_HC_DECODE_WS",
398 "the workspace decode walk carries no TP mixer branches",
399 ),
400 (
401 "MEMRA_KDA_FUSED_PROJ",
402 "the fused six-projection door (either operand arm) is gated on full-width \
403 projections, never head shards",
404 ),
405 (
406 "MEMRA_MLA_DECODE_SPLIT",
407 "the absorb/decompress split is gated on the full-head geometry",
408 ),
409];
410
411pub fn refuse_glm5_tp_door_composition(armed: impl Fn(&str) -> bool) -> Result<(), String> {
418 crate::tp::refuse_door_composition("MEMRA_GLM5_TP", &GLM5_TP_REFUSED_DOOR_FLAGS, armed)
419}
420
421pub fn prepare_glm5_tp_load(
424 e: &Engine,
425 view: &Glm5TpModelView,
426) -> Result<Option<Glm5TpLoadPlan>, Box<dyn std::error::Error>> {
427 let raw = glm5_tp_env_raw();
428 let specs = parse_glm5_tp_layer_specs(raw.as_deref(), view.trunk_layers)?;
429 if specs.is_empty() {
430 return Ok(None);
431 }
432
433 if crate::pp::pp_cuts(view.trunk_layers).is_some() {
436 return Err(
437 "MEMRA_GLM5_TP + MEMRA_PP_STAGES>1: the TP x PP composition is unwired and \
438 refuses until its own gate exists (stage 5 of the tp2 lane names it)"
439 .into(),
440 );
441 }
442 if !crate::tp::step_tp_layer_specs()?.is_empty()
443 || !crate::tp::step_ep_layer_specs()?.is_empty()
444 {
445 return Err(
446 "MEMRA_GLM5_TP + MEMRA_STEP_TP/MEMRA_STEP_EP: the step and glm5 parallel \
447 contracts never co-arm"
448 .into(),
449 );
450 }
451 refuse_glm5_tp_door_composition(|flag| std::env::var(flag).as_deref() == Ok("1"))?;
452
453 let devices = specs[0].devices.clone();
456 let ranks = devices.len();
457 if !GLM5_TP_ALLOWED_RANKS.contains(&ranks) {
458 return Err(format!(
459 "MEMRA_GLM5_TP names {ranks} devices per layer; the qualified rank envelope is \
460 {GLM5_TP_ALLOWED_RANKS:?} (TP-3 is a head-padding question, not built — see the \
461 module doc)"
462 )
463 .into());
464 }
465
466 if view.layer_class.len() != view.trunk_layers || view.layer_is_moe.len() != view.trunk_layers {
468 return Err(format!(
469 "glm5-tp preflight: layer class map ({}/{}) does not cover the {}-layer trunk",
470 view.layer_class.len(),
471 view.layer_is_moe.len(),
472 view.trunk_layers
473 )
474 .into());
475 }
476 if !view.kda_heads.is_multiple_of(ranks) || view.kda_heads == 0 {
477 return Err(format!(
478 "glm5-tp: {} KDA heads do not shard across {ranks} ranks",
479 view.kda_heads
480 )
481 .into());
482 }
483 if view.kda_head_dim != crate::kda::KDA_HEAD_DIM {
484 return Err(format!(
485 "glm5-tp: KDA head_dim {} is not the {} the scan kernel is instantiated for",
486 view.kda_head_dim,
487 crate::kda::KDA_HEAD_DIM
488 )
489 .into());
490 }
491 if !view.mla_heads.is_multiple_of(ranks) || view.mla_heads == 0 {
492 return Err(format!(
493 "glm5-tp: {} MLA heads do not shard across {ranks} ranks",
494 view.mla_heads
495 )
496 .into());
497 }
498 if !view.n_routed_experts.is_multiple_of(ranks) || view.n_routed_experts == 0 {
499 return Err(format!(
500 "glm5-tp: {} routed experts do not partition across {ranks} ranks",
501 view.n_routed_experts
502 )
503 .into());
504 }
505 if view.top_k > view.n_routed_experts {
506 return Err("glm5-tp: top_k exceeds the routed expert count".into());
507 }
508
509 for s in &specs {
510 if s.devices != devices {
511 return Err(format!(
512 "MEMRA_GLM5_TP carries ONE runtime group: layer {} names devices {:?}, \
513 the first spec names {:?}",
514 s.layer, s.devices, devices
515 )
516 .into());
517 }
518 if s.layer >= view.trunk_layers {
519 return Err(format!(
520 "MEMRA_GLM5_TP layer {} outside the {}-layer trunk",
521 s.layer, view.trunk_layers
522 )
523 .into());
524 }
525 }
526 let root_dev = e.ctx().ordinal();
527 if devices[0] != root_dev {
528 return Err(format!(
529 "MEMRA_GLM5_TP rank list {:?} must start with the owning device {root_dev} \
530 (the owner-first rank law)",
531 devices
532 )
533 .into());
534 }
535
536 let red = gate_red()?;
538 let same_dev = gate_same_device();
539 if let Some(red) = red {
540 eprintln!("[glm5-tp-preflight] GATE RED ARM armed: {red:?} — outputs MUST diverge");
541 }
542 let mut rt = if same_dev {
543 eprintln!(
544 "[glm5-tp-preflight] GATE same-device emulation: {} peer ranks are additional \
545 contexts on device {root_dev} (spec devices {:?} are logical rank ids)",
546 ranks - 1,
547 &devices[1..],
548 );
549 Glm5TpRt::new_gate_same_device(root_dev, ranks)?
550 } else {
551 Glm5TpRt::new(&devices)?
552 };
553 rt.arm_transport(e)?;
556 let rt = Arc::new(rt);
557 let layers: std::collections::BTreeSet<usize> = specs.iter().map(|s| s.layer).collect();
558 let ep_map = load_glm5_ep_map(view, &layers, ranks)?;
559 let (mut kda_n, mut mla_n, mut moe_n) = (0usize, 0usize, 0usize);
560 for &il in &layers {
561 match view.layer_class[il] {
562 Glm5LayerClass::Kda => kda_n += 1,
563 Glm5LayerClass::Mla => mla_n += 1,
564 }
565 if view.layer_is_moe[il] {
566 moe_n += 1;
567 }
568 }
569 eprintln!(
570 "[glm5-tp-preflight] armed ranks={ranks} devices={devices:?} layers={} \
571 kda_shard={kda_n} mla_shard={mla_n} moe_ep={moe_n} kda_heads_per_rank={} \
572 mla_heads_per_rank={} experts_per_rank={} transport={} \
573 weights_loaded=false performance_claim=false",
574 layers.len(),
575 view.kda_heads / ranks,
576 view.mla_heads / ranks,
577 view.n_routed_experts / ranks,
578 rt.transport.name(),
579 );
580 Ok(Some(Glm5TpLoadPlan { rt, layers, ep_map }))
581}
582
583fn outer_rows(ne: &[u64]) -> (usize, usize) {
588 let outer = *ne.last().expect("tensor has at least one axis") as usize;
591 let inner: usize = ne[..ne.len() - 1].iter().map(|&d| d as usize).product();
592 (outer, inner.max(1))
593}
594
595fn shard_rows(
599 src_engine: &Engine,
600 dst: &Engine,
601 t: &GpuTensor,
602 rows: Range<usize>,
603) -> Result<GpuTensor, Box<dyn std::error::Error>> {
604 match t {
605 GpuTensor::Float { data, ne } => {
606 let (outer, inner) = outer_rows(ne);
607 if rows.end > outer {
608 return Err(format!("shard rows {rows:?} exceed outer axis {outer}").into());
609 }
610 let host = src_engine.dtoh(data)?;
611 let piece = &host[rows.start * inner..rows.end * inner];
612 let mut ne2 = ne.clone();
613 *ne2.last_mut().unwrap() = (rows.end - rows.start) as u64;
614 Ok(GpuTensor::Float {
615 data: dst.htod(piece)?,
616 ne: ne2,
617 })
618 }
619 GpuTensor::FloatBf16 { data, ne } => {
620 let (outer, inner) = outer_rows(ne);
621 if rows.end > outer {
622 return Err(format!("shard rows {rows:?} exceed outer axis {outer}").into());
623 }
624 let host = src_engine.dtoh_u8(data)?;
625 let piece = &host[rows.start * inner * 2..rows.end * inner * 2];
626 let mut ne2 = ne.clone();
627 *ne2.last_mut().unwrap() = (rows.end - rows.start) as u64;
628 Ok(GpuTensor::FloatBf16 {
629 data: dst.htod_bytes(piece)?,
630 ne: ne2,
631 })
632 }
633 GpuTensor::Quant {
634 bytes,
635 qtype,
636 row_bytes,
637 ne,
638 scale,
639 rp,
640 fp8,
641 rp4,
642 blk,
643 f16,
644 #[cfg(memra_cutlass)]
645 cutlass,
646 } => {
647 if *rp {
648 return Err(
649 "glm5-tp shard: rp split-plane mirror layout is unwired — load \
650 the TP-armed tensor with MEMRA_RP=0 (raw layout is bit-identical \
651 by the mirror's own contract)"
652 .into(),
653 );
654 }
655 if fp8.is_some() || rp4.is_some() || blk.is_some() || f16.is_some() {
656 return Err(
657 "glm5-tp shard: a decode/prefill mirror (fp8/rp4/blk/f16) is present on a \
658 TP-armed tensor — mirrors are unwired for shards in v1; disable the \
659 mirror door for this load"
660 .into(),
661 );
662 }
663 #[cfg(memra_cutlass)]
664 if cutlass.is_some() {
665 return Err("glm5-tp shard: cutlass prefill operand unwired for shards".into());
666 }
667 let (outer, inner) = outer_rows(ne);
668 if ne.len() != 2 {
669 return Err("glm5-tp shard: quantized shards are 2D-only in v1".into());
670 }
671 let _ = inner;
672 if rows.end > outer {
673 return Err(format!("shard rows {rows:?} exceed outer axis {outer}").into());
674 }
675 let host = src_engine.dtoh_u8(bytes)?;
676 let piece = &host[rows.start * row_bytes..rows.end * row_bytes];
677 let mut ne2 = ne.clone();
678 *ne2.last_mut().unwrap() = (rows.end - rows.start) as u64;
679 Ok(GpuTensor::Quant {
680 bytes: dst.htod_bytes(piece)?,
681 qtype: *qtype,
682 row_bytes: *row_bytes,
683 ne: ne2,
684 scale: *scale,
685 rp: false,
686 fp8: None,
687 rp4: None,
688 blk: None,
689 f16: None,
690 #[cfg(memra_cutlass)]
691 cutlass: None,
692 })
693 }
694 }
695}
696
697fn replicate(
699 src_engine: &Engine,
700 dst: &Engine,
701 t: &GpuTensor,
702) -> Result<GpuTensor, Box<dyn std::error::Error>> {
703 let (outer, _) = outer_rows(t.ne());
704 shard_rows(src_engine, dst, t, 0..outer)
705}
706
707pub(crate) fn rank_engine<'a>(e: &'a Engine, rt: &'a Glm5TpRt, r: usize) -> &'a Engine {
710 if r == 0 { e } else { &rt.peers[r - 1] }
711}
712
713pub struct Glm5TpKda {
721 pub rt: Arc<Glm5TpRt>,
722 pub peers: Vec<KdaAttnLayer>,
724 pub full_qkv: usize,
726 pub n_embd: usize,
728}
729
730impl Glm5TpKda {
731 pub fn ranks(&self) -> usize {
732 self.peers.len() + 1
733 }
734}
735
736static KDA_MARKED: AtomicBool = AtomicBool::new(false);
737
738pub(crate) fn shard_kda_layer(
742 e: &Engine,
743 rt: &Arc<Glm5TpRt>,
744 la: KdaAttnLayer,
745) -> Result<KdaAttnLayer, Box<dyn std::error::Error>> {
746 if la.tp.is_some() {
747 return Err("shard_kda_layer: layer is already sharded".into());
748 }
749 let ranks = rt.ranks();
750 let heads = la.heads();
751 let head_dim = la.head_dim();
752 let qkv = la.qkv();
753 let kernel = la.conv_kernel();
754 if !heads.is_multiple_of(ranks) {
755 return Err(format!("KDA heads {heads} do not shard across {ranks} ranks").into());
756 }
757 let hl = heads / ranks; let ql = qkv / ranks; let n_embd = la.wo.out_features();
760 if !n_embd.is_multiple_of(ranks) {
761 return Err(format!("KDA wo out {n_embd} does not split across ranks").into());
762 }
763 let hh = n_embd / ranks;
764
765 let mut shard_plan = la.plan;
766 shard_plan.num_heads = hl as u32;
767
768 let conv_host = e.dtoh(&la.conv)?;
771 let conv_rank =
772 |dst: &Engine, r: usize| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
773 let mut piece = Vec::with_capacity(3 * ql * kernel);
774 for p in 0..3 {
775 let a = (p * qkv + r * ql) * kernel;
776 piece.extend_from_slice(&conv_host[a..a + ql * kernel]);
777 }
778 dst.htod(&piece)
779 };
780
781 let wo_rank = |r: usize| -> usize {
784 match gate_red() {
785 Ok(Some(GateRed::SwapWo)) => (r + 1) % ranks,
786 _ => r,
787 }
788 };
789
790 let rank_shard = |dst: &Engine, r: usize| -> Result<KdaAttnLayer, Box<dyn std::error::Error>> {
791 let wr = wo_rank(r);
792 Ok(KdaAttnLayer {
793 plan: shard_plan,
794 wq: shard_rows(e, dst, &la.wq, r * ql..(r + 1) * ql)?,
795 wk: shard_rows(e, dst, &la.wk, r * ql..(r + 1) * ql)?,
796 wv: shard_rows(e, dst, &la.wv, r * ql..(r + 1) * ql)?,
797 f_a: replicate(e, dst, &la.f_a)?,
798 f_b: shard_rows(e, dst, &la.f_b, r * ql..(r + 1) * ql)?,
799 g_a: replicate(e, dst, &la.g_a)?,
800 g_b: shard_rows(e, dst, &la.g_b, r * ql..(r + 1) * ql)?,
801 b_proj: shard_rows(e, dst, &la.b_proj, r * hl..(r + 1) * hl)?,
802 wo: shard_rows(e, dst, &la.wo, wr * hh..(wr + 1) * hh)?,
806 conv: conv_rank(dst, r)?,
807 a_log: shard_rows(e, dst, &la.a_log, r * hl..(r + 1) * hl)?,
808 dt_bias: shard_rows(e, dst, &la.dt_bias, r * ql..(r + 1) * ql)?,
809 o_norm: replicate(e, dst, &la.o_norm)?,
810 tp: None,
811 })
812 };
813
814 let mut root = rank_shard(e, 0)?;
815 let mut peers = Vec::with_capacity(ranks - 1);
816 for r in 1..ranks {
817 peers.push(rank_shard(&rt.peers[r - 1], r)?);
818 }
819 if !KDA_MARKED.swap(true, Ordering::Relaxed) {
820 eprintln!(
821 "[glm5-tp-kda] head shard armed: ranks={ranks} heads_per_rank={hl} \
822 head_dim={head_dim} wo=column-over-gather transport={} performance_claim=false",
823 rt.transport.name(),
824 );
825 }
826 root.tp = Some(Box::new(Glm5TpKda {
827 rt: Arc::clone(rt),
828 peers,
829 full_qkv: qkv,
830 n_embd,
831 }));
832 Ok(root)
833}
834
835fn ensure_kda_tp_state<'c>(
840 e: &Engine,
841 rt: &Glm5TpRt,
842 la_root: &KdaAttnLayer,
843 cache: &'c mut Cache,
844 il: usize,
845) -> Result<&'c mut Vec<RecurLayer>, Box<dyn std::error::Error>> {
846 if cache.glm5_tp_recur.len() <= il {
847 return Err(format!("glm5-tp: cache carries no TP recur slot for layer {il}").into());
848 }
849 if cache.glm5_tp_recur[il].is_none() {
850 let conv_pad = la_root.conv_width() * (la_root.conv_kernel() - 1);
851 let state = la_root.state_width();
852 let mk = |dev: &Engine| -> Result<RecurLayer, Box<dyn std::error::Error>> {
853 Ok(RecurLayer {
854 conv_state: dev.zeros(conv_pad)?,
855 ssm_state: dev.zeros(state)?,
856 ssm_state_alt: dev.zeros(state)?,
857 })
858 };
859 let mut planes = Vec::with_capacity(rt.ranks());
860 planes.push(mk(e)?);
861 for p in &rt.peers {
862 planes.push(mk(p)?);
863 }
864 cache.glm5_tp_recur[il] = Some(planes);
865 }
866 Ok(cache.glm5_tp_recur[il].as_mut().unwrap())
867}
868
869#[allow(clippy::too_many_arguments)] fn kda_tp_core(
886 e: &Engine,
887 la_root: &KdaAttnLayer,
888 x: &CudaSlice<f32>,
889 t: usize,
890 eps: f32,
891 cache: &mut Cache,
892 il: usize,
893 arm: ConvArm,
894 verify_stash: Option<&mut Glm5TpKdaVerifyStash>,
895 mut scan_clock: Option<&mut u64>,
896) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
897 let tp = la_root
898 .tp
899 .as_ref()
900 .ok_or("kda_tp_core called on an unsharded layer")?;
901 let rt = &tp.rt;
902 let ranks = rt.ranks();
903 let ql = la_root.qkv(); let full = tp.full_qkv;
905 let n_embd = tp.n_embd;
906 let hh = n_embd / ranks;
907 let rows_exact = verify_stash.is_some();
908 let mut captured: Vec<Option<(CudaSlice<f32>, crate::kda::KdaRowsStash)>> =
911 (0..ranks).map(|_| None).collect();
912
913 let hop = rt.hop(e);
914 let x_peers = crate::tp_transport::fanout_f32(&hop, x, x.len())?;
918 let states = ensure_kda_tp_state(e, rt, la_root, cache, il)?;
919
920 let mut gated: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
923 for r in (1..ranks).chain(std::iter::once(0)) {
924 let dev = if r == 0 { e } else { &rt.peers[r - 1] };
925 let la = if r == 0 { la_root } else { &tp.peers[r - 1] };
926 let xin = if r == 0 { x } else { &x_peers[r - 1] };
927 let snap = if rows_exact {
930 Some(dev.clone_dtod(&states[r].ssm_state)?)
931 } else {
932 None
933 };
934 let mut rank_stash: Option<crate::kda::KdaRowsStash> = None;
935 let mut rank_scan_ns = 0u64;
936 let out = {
937 let RecurLayer {
938 conv_state,
939 ssm_state,
940 ssm_state_alt,
941 } = &mut states[r];
942 let out = crate::kda::kda_core_gated(
943 dev,
944 la,
945 xin,
946 t,
947 eps,
948 conv_state,
949 ssm_state,
950 ssm_state_alt,
951 arm,
952 if rows_exact {
953 crate::kda::KdaStash::Rows(&mut rank_stash)
954 } else {
955 crate::kda::KdaStash::None
956 },
957 scan_clock.as_deref_mut().map(|_| &mut rank_scan_ns),
958 )?;
959 std::mem::swap(ssm_state, ssm_state_alt);
960 out
961 };
962 if let Some(clock) = scan_clock.as_deref_mut() {
963 *clock += rank_scan_ns;
964 }
965 if rows_exact {
966 let snap = snap.expect("verify arm cloned the snapshot above");
967 let rank_stash = rank_stash
968 .ok_or("kda_core_gated returned without filling the requested rows stash")?;
969 captured[r] = Some((snap, rank_stash));
970 }
971 gated[r] = Some(out);
972 }
973 if let Some(stash_vec) = verify_stash {
974 stash_vec.clear();
975 for c in captured {
976 stash_vec.push(c.expect("every rank captured on the verify arm"));
977 }
978 }
979
980 debug_assert_eq!(full, ranks * ql);
984 let gated_refs: Vec<&CudaSlice<f32>> = gated
985 .iter()
986 .map(|g| g.as_ref().expect("filled above"))
987 .collect();
988 let fulls = crate::tp_transport::gather_parts(&hop, &gated_refs, t, ql)?;
989
990 let mut ys = Vec::with_capacity(ranks);
993 if rows_exact {
994 ys.push(e.matmul_rows_exact(&la_root.wo, &fulls[0], t)?);
995 for r in 1..ranks {
996 ys.push(rt.peers[r - 1].matmul_rows_exact(&tp.peers[r - 1].wo, &fulls[r], t)?);
997 }
998 } else {
999 ys.push(e.matmul(&la_root.wo, &fulls[0], t)?);
1000 for r in 1..ranks {
1001 ys.push(rt.peers[r - 1].matmul(&tp.peers[r - 1].wo, &fulls[r], t)?);
1002 }
1003 }
1004
1005 let y_refs: Vec<&CudaSlice<f32>> = ys.iter().collect();
1007 crate::tp_transport::concat_parts_on_root(&hop, &y_refs, t, hh)
1008}
1009
1010#[allow(clippy::too_many_arguments)] pub(crate) fn kda_tp_cached(
1013 e: &Engine,
1014 la_root: &KdaAttnLayer,
1015 x: &CudaSlice<f32>,
1016 t: usize,
1017 eps: f32,
1018 cache: &mut Cache,
1019 il: usize,
1020 arm: ConvArm,
1021) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1022 kda_tp_core(e, la_root, x, t, eps, cache, il, arm, None, None)
1023}
1024
1025pub type Glm5TpKdaVerifyStash = Vec<(CudaSlice<f32>, crate::kda::KdaRowsStash)>;
1033
1034#[allow(clippy::too_many_arguments)] pub(crate) fn kda_tp_verify_rows(
1040 e: &Engine,
1041 la_root: &KdaAttnLayer,
1042 x: &CudaSlice<f32>,
1043 t: usize,
1044 eps: f32,
1045 cache: &mut Cache,
1046 il: usize,
1047 scan_clock: Option<&mut u64>,
1048) -> Result<(CudaSlice<f32>, Glm5TpKdaVerifyStash), Box<dyn std::error::Error>> {
1049 let mut stash: Glm5TpKdaVerifyStash = Vec::new();
1050 let out = kda_tp_core(
1051 e,
1052 la_root,
1053 x,
1054 t,
1055 eps,
1056 cache,
1057 il,
1058 ConvArm::Prefill,
1059 Some(&mut stash),
1060 scan_clock,
1061 )?;
1062 Ok((out, stash))
1063}
1064
1065pub(crate) fn kda_tp_verify_rollback(
1069 e: &Engine,
1070 la_root: &KdaAttnLayer,
1071 stash: &Glm5TpKdaVerifyStash,
1072 keep: usize,
1073 cache: &mut Cache,
1074 il: usize,
1075) -> Result<(), Box<dyn std::error::Error>> {
1076 let tp = la_root
1077 .tp
1078 .as_ref()
1079 .ok_or("kda_tp_verify_rollback called on an unsharded layer")?;
1080 let rt = &tp.rt;
1081 let ranks = rt.ranks();
1082 if stash.len() != ranks {
1083 return Err(format!(
1084 "glm5-tp verify rollback: stash carries {} ranks, the runtime has {ranks}",
1085 stash.len()
1086 )
1087 .into());
1088 }
1089 let states = cache.glm5_tp_recur[il]
1090 .as_mut()
1091 .ok_or_else(|| format!("glm5-tp verify rollback: layer {il} has no per-rank state"))?;
1092 for r in 0..ranks {
1093 let dev = if r == 0 { e } else { &rt.peers[r - 1] };
1094 let la = if r == 0 { la_root } else { &tp.peers[r - 1] };
1095 let (snap, rows) = &stash[r];
1096 crate::kda::kda_verify_rollback_rows_on(dev, la, snap, rows, keep, &mut states[r], il)?;
1097 }
1098 Ok(())
1099}
1100
1101pub struct Glm5TpMla {
1108 pub rt: Arc<Glm5TpRt>,
1109 pub peers: Vec<crate::hybrid::MlaAttnLayer>,
1111 pub full_heads: usize,
1113 pub n_embd: usize,
1115}
1116
1117impl Glm5TpMla {
1118 pub fn ranks(&self) -> usize {
1119 self.peers.len() + 1
1120 }
1121}
1122
1123static MLA_MARKED: AtomicBool = AtomicBool::new(false);
1124
1125pub(crate) fn shard_mla_layer(
1126 e: &Engine,
1127 rt: &Arc<Glm5TpRt>,
1128 la: crate::hybrid::MlaAttnLayer,
1129) -> Result<crate::hybrid::MlaAttnLayer, Box<dyn std::error::Error>> {
1130 use crate::hybrid::{MlaAttnLayer, MlaIndexer};
1131 if la.tp.is_some() {
1132 return Err("shard_mla_layer: layer is already sharded".into());
1133 }
1134 let ranks = rt.ranks();
1135 let g = la.geom;
1136 let nh = g.n_head;
1137 if !nh.is_multiple_of(ranks) {
1138 return Err(format!("MLA heads {nh} do not shard across {ranks} ranks").into());
1139 }
1140 let hl = nh / ranks;
1141 let head_q = g.d_nope + g.d_rope; let n_embd = la.wo.out_features();
1143 if !n_embd.is_multiple_of(ranks) {
1144 return Err(format!("MLA wo out {n_embd} does not split across ranks").into());
1145 }
1146 let hh = n_embd / ranks;
1147
1148 let mut shard_geom = g;
1149 shard_geom.n_head = hl;
1150
1151 let replicate_indexer =
1152 |dst: &Engine, ix: &MlaIndexer| -> Result<MlaIndexer, Box<dyn std::error::Error>> {
1153 Ok(MlaIndexer {
1154 wq_b: replicate(e, dst, &ix.wq_b)?,
1155 wk: replicate(e, dst, &ix.wk)?,
1156 k_norm_w: replicate(e, dst, &ix.k_norm_w)?,
1157 k_norm_b: replicate(e, dst, &ix.k_norm_b)?,
1158 weights_proj: replicate(e, dst, &ix.weights_proj)?,
1159 kpool_gate: replicate(e, dst, &ix.kpool_gate)?,
1160 kpool_ape: replicate(e, dst, &ix.kpool_ape)?,
1161 geom: ix.geom,
1162 })
1163 };
1164
1165 let wo_rank = |r: usize| -> usize {
1167 match gate_red() {
1168 Ok(Some(GateRed::SwapWo)) => (r + 1) % ranks,
1169 _ => r,
1170 }
1171 };
1172
1173 let rank_shard = |dst: &Engine, r: usize| -> Result<MlaAttnLayer, Box<dyn std::error::Error>> {
1174 let wr = wo_rank(r);
1175 Ok(MlaAttnLayer {
1176 wq_a: replicate(e, dst, &la.wq_a)?,
1177 q_a_norm: replicate(e, dst, &la.q_a_norm)?,
1178 wq_b: shard_rows(e, dst, &la.wq_b, r * hl * head_q..(r + 1) * hl * head_q)?,
1179 wkv_a: replicate(e, dst, &la.wkv_a)?,
1180 kv_a_norm: replicate(e, dst, &la.kv_a_norm)?,
1181 wk_b: shard_rows(e, dst, &la.wk_b, r * hl..(r + 1) * hl)?,
1183 wv_b: shard_rows(e, dst, &la.wv_b, r * hl..(r + 1) * hl)?,
1184 wo: shard_rows(e, dst, &la.wo, wr * hh..(wr + 1) * hh)?,
1186 geom: shard_geom,
1187 index: match &la.index {
1188 Some(ix) => Some(replicate_indexer(dst, ix)?),
1189 None => None,
1190 },
1191 tp: None,
1192 tp_shard: true,
1193 })
1194 };
1195
1196 let mut root = rank_shard(e, 0)?;
1197 let mut peers = Vec::with_capacity(ranks - 1);
1198 for r in 1..ranks {
1199 peers.push(rank_shard(&rt.peers[r - 1], r)?);
1200 }
1201 if !MLA_MARKED.swap(true, Ordering::Relaxed) {
1202 eprintln!(
1203 "[glm5-tp-mla] head shard armed: ranks={ranks} heads_per_rank={hl} kv_rank={} \
1204 latent=replicated indexer=replicated wo=column-over-gather transport={} \
1205 performance_claim=false",
1206 g.kv_rank,
1207 rt.transport.name(),
1208 );
1209 }
1210 root.tp = Some(Box::new(Glm5TpMla {
1211 rt: Arc::clone(rt),
1212 peers,
1213 full_heads: nh,
1214 n_embd,
1215 }));
1216 Ok(root)
1217}
1218
1219pub(crate) fn ensure_mla_peer_latent(
1223 rt: &Glm5TpRt,
1224 canonical: &LatentKvLayer,
1225 cache_slot: &mut Option<Vec<LatentKvLayer>>,
1226) -> Result<(), Box<dyn std::error::Error>> {
1227 if cache_slot.is_some() {
1228 return Ok(());
1229 }
1230 let mut planes = Vec::with_capacity(rt.peers.len());
1231 for dev in &rt.peers {
1232 let rows = dev.zeros(canonical.rows.len())?;
1233 let len_d = dev.htod_i32(&[0])?;
1236 let index_rows = match &canonical.index_rows {
1237 Some(p) => Some(dev.zeros(p.len())?),
1238 None => None,
1239 };
1240 planes.push(LatentKvLayer {
1241 rows,
1242 width: canonical.width,
1243 index_width: canonical.index_width,
1244 len: 0,
1245 len_d,
1246 index_rows,
1247 index_ring_rows: canonical.index_ring_rows,
1248 index_pool_keys: None, index_pools_ready: 0,
1250 index_pool: canonical.index_pool,
1251 });
1252 }
1253 *cache_slot = Some(planes);
1254 Ok(())
1255}
1256
1257pub struct EpRankSlab {
1265 pub gate: CudaSlice<u8>,
1266 pub up: CudaSlice<u8>,
1267 pub down: CudaSlice<u8>,
1268 pub n_experts: usize,
1269}
1270
1271pub struct Glm5EpExps {
1281 pub rt: Arc<Glm5TpRt>,
1282 pub slabs: Vec<EpRankSlab>,
1284 pub owner_of: Vec<u8>,
1286 pub local_of: Vec<u32>,
1288 pub ptr_rows: Vec<CudaSlice<u64>>,
1297}
1298
1299impl Glm5EpExps {
1300 pub fn owner(&self, expert: usize) -> usize {
1302 self.owner_of[expert] as usize
1303 }
1304
1305 pub fn ranks(&self) -> usize {
1306 self.slabs.len()
1307 }
1308}
1309
1310static EP_MARKED: AtomicBool = AtomicBool::new(false);
1311
1312pub static GLM5_EP_PEER_SLOT_DISPATCHES: AtomicU64 = AtomicU64::new(0);
1317
1318pub fn glm5_ep_peer_slot_dispatches() -> u64 {
1319 GLM5_EP_PEER_SLOT_DISPATCHES.load(Ordering::Relaxed)
1320}
1321
1322pub static GLM5_EP_DIET_DISPATCHES: AtomicU64 = AtomicU64::new(0);
1329
1330pub fn glm5_ep_diet_dispatches() -> u64 {
1332 GLM5_EP_DIET_DISPATCHES.load(Ordering::Relaxed)
1333}
1334
1335pub static GLM5_EP_DIET_BULK_RETURNS: AtomicU64 = AtomicU64::new(0);
1339
1340pub fn glm5_ep_diet_bulk_returns() -> u64 {
1342 GLM5_EP_DIET_BULK_RETURNS.load(Ordering::Relaxed)
1343}
1344
1345pub static GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED: AtomicU64 = AtomicU64::new(0);
1349
1350pub fn glm5_ep_diet_peer_roundtrips_avoided() -> u64 {
1352 GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED.load(Ordering::Relaxed)
1353}
1354
1355pub static GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED: AtomicU64 = AtomicU64::new(0);
1361
1362pub fn glm5_ep_diet_fanout_uploads_avoided() -> u64 {
1364 GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED.load(Ordering::Relaxed)
1365}
1366
1367pub static GLM5_EP_GROUPED_PRIME_DISPATCHES: AtomicU64 = AtomicU64::new(0);
1371
1372pub fn glm5_ep_grouped_prime_dispatches() -> u64 {
1374 GLM5_EP_GROUPED_PRIME_DISPATCHES.load(Ordering::Relaxed)
1375}
1376
1377pub(crate) fn arm_moe_ep(
1381 e: &Engine,
1382 rt: &Arc<Glm5TpRt>,
1383 m: &mut crate::hybrid::MoeWeights,
1384 placement: Option<&[u8]>,
1385) -> Result<(), Box<dyn std::error::Error>> {
1386 if m.glm5_ep.is_some() {
1387 return Err("arm_moe_ep: layer is already EP-armed".into());
1388 }
1389 let ranks = rt.ranks();
1390 let n_expert = m.gate_exps.n_expert;
1391 if !n_expert.is_multiple_of(ranks) {
1392 return Err(format!(
1393 "glm5-tp EP: {n_expert} experts do not partition across {ranks} ranks"
1394 )
1395 .into());
1396 }
1397 if m.gate_exps.layouts.is_some() || m.up_exps.layouts.is_some() || m.down_exps.layouts.is_some()
1398 {
1399 return Err("glm5-tp EP: per-expert mixed layouts are unwired for EP shards".into());
1400 }
1401 let owner_of: Vec<u8> = match placement {
1402 Some(owners) => {
1403 if owners.len() != n_expert {
1406 return Err(format!(
1407 "glm5-tp EP: placement row carries {} owners for a {n_expert}-expert bank",
1408 owners.len()
1409 )
1410 .into());
1411 }
1412 if owners.iter().any(|&r| (r as usize) >= ranks) {
1413 return Err(
1414 format!("glm5-tp EP: placement row names a rank outside TP-{ranks}").into(),
1415 );
1416 }
1417 owners.to_vec()
1418 }
1419 None => crate::ep_map::EpMap::even_owners(n_expert, ranks),
1420 };
1421 let mut local_of = vec![0u32; n_expert];
1423 let mut owned: Vec<Vec<usize>> = vec![Vec::new(); ranks];
1424 for ex in 0..n_expert {
1425 let r = owner_of[ex] as usize;
1426 local_of[ex] = owned[r].len() as u32;
1427 owned[r].push(ex);
1428 }
1429 if owned.iter().any(|o| o.is_empty()) {
1430 return Err("glm5-tp EP: placement leaves a rank with zero experts (refused)".into());
1431 }
1432 let slab =
1433 |dev: &Engine, experts: &[usize]| -> Result<EpRankSlab, Box<dyn std::error::Error>> {
1434 let cut = |h: &crate::model::HostExps,
1441 pad: usize|
1442 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1443 let stride = h.expert_stride;
1444 let bytes = h.bytes.as_bytes();
1445 let contiguous = experts.windows(2).all(|w| w[1] == w[0] + 1);
1448 if contiguous {
1449 let a = experts[0] * stride;
1450 let b = (experts[experts.len() - 1] + 1) * stride;
1451 return dev.htod_bytes_padded(&bytes[a..b], pad);
1452 }
1453 let mut staged = Vec::with_capacity(experts.len() * stride);
1456 for &ex in experts {
1457 staged.extend_from_slice(&bytes[ex * stride..(ex + 1) * stride]);
1458 }
1459 dev.htod_bytes_padded(&staged, pad)
1460 };
1461 Ok(EpRankSlab {
1462 gate: cut(&m.gate_exps, 8)?,
1463 up: cut(&m.up_exps, 8)?,
1464 down: cut(&m.down_exps, 144)?,
1465 n_experts: experts.len(),
1466 })
1467 };
1468 let mut slabs = Vec::with_capacity(ranks);
1469 for r in 0..ranks {
1470 slabs.push(slab(rank_engine(e, rt, r), &owned[r])?);
1471 }
1472 if matches!(gate_red(), Ok(Some(GateRed::SwapEpGateUp))) {
1474 let root = &mut slabs[0];
1475 std::mem::swap(&mut root.gate, &mut root.up);
1476 }
1477 if matches!(gate_red(), Ok(Some(GateRed::CorruptEpMap))) {
1480 let n0 = owned[0].len() as u32;
1481 for &ex in &owned[0] {
1482 local_of[ex] = n0 - 1 - local_of[ex];
1483 }
1484 }
1485 let ptr_table = |dev: &Engine,
1490 slab: &EpRankSlab,
1491 rank: u8|
1492 -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
1493 use cudarc::driver::DevicePtr;
1494 let (pg, pu, pd) = {
1495 let s = dev.stream();
1496 let (pg, _g0) = slab.gate.device_ptr(&s);
1497 let (pu, _g1) = slab.up.device_ptr(&s);
1498 let (pd, _g2) = slab.down.device_ptr(&s);
1499 (pg, pu, pd)
1500 };
1501 let mut host = vec![0u64; 3 * n_expert];
1502 for ex in 0..n_expert {
1503 if owner_of[ex] != rank {
1504 continue; }
1506 let local = local_of[ex] as usize;
1507 host[ex] = pg + (local * m.gate_exps.expert_stride) as u64;
1508 host[n_expert + ex] = pu + (local * m.up_exps.expert_stride) as u64;
1509 host[2 * n_expert + ex] = pd + (local * m.down_exps.expert_stride) as u64;
1510 }
1511 dev.htod_u64(&host)
1512 };
1513 let mut ptr_rows = Vec::with_capacity(ranks);
1514 for r in 0..ranks {
1515 ptr_rows.push(ptr_table(rank_engine(e, rt, r), &slabs[r], r as u8)?);
1516 }
1517 if !EP_MARKED.swap(true, Ordering::Relaxed) {
1518 eprintln!(
1519 "[glm5-tp-ep] expert-parallel armed: experts_per_rank={:?} ownership={} \
1520 router=root combine=slot-ordered-fmaf transport={} \
1521 performance_claim=false",
1522 owned.iter().map(Vec::len).collect::<Vec<_>>(),
1523 if placement.is_some() {
1528 "measured-map"
1529 } else {
1530 "even-split"
1531 },
1532 rt.transport.name(),
1533 );
1534 }
1535 m.dev_exps = None;
1538 m.glm5_ep = Some(Glm5EpExps {
1539 rt: Arc::clone(rt),
1540 slabs,
1541 owner_of,
1542 local_of,
1543 ptr_rows,
1544 });
1545 Ok(())
1546}
1547
1548#[cfg(test)]
1549mod tests {
1550 use super::*;
1551
1552 #[test]
1553 fn parse_is_literal_and_fail_closed() {
1554 assert!(parse_glm5_tp_layer_specs(None, 45).unwrap().is_empty());
1556 assert!(parse_glm5_tp_layer_specs(Some(""), 45).unwrap().is_empty());
1557 assert!(parse_glm5_tp_layer_specs(Some("0"), 45).unwrap().is_empty());
1558 let all = parse_glm5_tp_layer_specs(Some("all@0,1"), 45).unwrap();
1560 assert_eq!(all.len(), 45);
1561 assert_eq!(all[0].devices, vec![0, 1]);
1562 let all4 = parse_glm5_tp_layer_specs(Some("all@0,1"), 4).unwrap();
1563 assert_eq!(all4.len(), 4);
1564 let quad = parse_glm5_tp_layer_specs(Some("all@0,1,2,3"), 45).unwrap();
1566 assert_eq!(quad.len(), 45);
1567 assert_eq!(quad[0].devices, vec![0, 1, 2, 3]);
1568 let r = parse_glm5_tp_layer_specs(Some("0-2@0,1;4@0,1"), 45).unwrap();
1570 assert_eq!(
1571 r.iter().map(|s| s.layer).collect::<Vec<_>>(),
1572 vec![0, 1, 2, 4]
1573 );
1574 assert!(parse_glm5_tp_layer_specs(Some("0@0,0"), 45).is_err());
1576 assert!(parse_glm5_tp_layer_specs(Some("0@0,1;0@0,1"), 45).is_err());
1577 assert!(parse_glm5_tp_layer_specs(Some("banana"), 45).is_err());
1578 }
1579
1580 fn fixture_view() -> Glm5TpModelView {
1581 Glm5TpModelView {
1582 trunk_layers: 4,
1583 layer_class: vec![
1584 Glm5LayerClass::Kda,
1585 Glm5LayerClass::Mla,
1586 Glm5LayerClass::Kda,
1587 Glm5LayerClass::Mla,
1588 ],
1589 layer_is_moe: vec![false, true, true, true],
1590 kda_heads: 4,
1591 kda_head_dim: 128,
1592 mla_heads: 4,
1593 n_routed_experts: 4,
1594 top_k: 2,
1595 }
1596 }
1597
1598 #[test]
1602 fn preflight_geometry_laws_are_dimension_derived() {
1603 let v = fixture_view();
1606 for ranks in GLM5_TP_ALLOWED_RANKS {
1607 assert_eq!(v.kda_heads % ranks, 0);
1608 assert_eq!(v.mla_heads % ranks, 0);
1609 assert_eq!(v.n_routed_experts % ranks, 0);
1610 }
1611 let odd = Glm5TpModelView {
1612 kda_heads: 3,
1613 ..fixture_view()
1614 };
1615 assert_ne!(odd.kda_heads % 2, 0);
1616 let bad_dim = Glm5TpModelView {
1617 kda_head_dim: 64,
1618 ..fixture_view()
1619 };
1620 assert_ne!(bad_dim.kda_head_dim, crate::kda::KDA_HEAD_DIM);
1621 let odd_experts = Glm5TpModelView {
1622 n_routed_experts: 5,
1623 ..fixture_view()
1624 };
1625 assert_ne!(odd_experts.n_routed_experts % 2, 0);
1626 assert!(!GLM5_TP_ALLOWED_RANKS.contains(&3));
1628 }
1629
1630 #[test]
1631 fn armed_check_counts_parse_errors_as_armed() {
1632 for (v, armed) in [
1637 ("", false),
1638 ("0", false),
1639 ("all@0,1", true),
1640 ("all@0,1,2,3", true),
1641 ("junk", true),
1642 ] {
1643 let is_armed = !v.is_empty() && v != "0";
1644 assert_eq!(is_armed, armed);
1645 }
1646 }
1647
1648 #[test]
1649 fn every_refused_door_composition_bites_by_name() {
1650 for (flag, _) in GLM5_TP_REFUSED_DOOR_FLAGS {
1656 let err = refuse_glm5_tp_door_composition(|f| f == flag)
1657 .expect_err("an armed door must refuse");
1658 assert!(err.contains("MEMRA_GLM5_TP"), "{err}");
1659 assert!(err.contains(flag), "{err}");
1660 assert!(err.contains("unproven composition"), "{err}");
1661 }
1662 refuse_glm5_tp_door_composition(|_| false).expect("cold doors must pass");
1664 refuse_glm5_tp_door_composition(|f| f == "MEMRA_GLM5_VERIFY_BATCH")
1668 .expect("verify-batch is refused via the spec co-refusal, not here");
1669 }
1670}