Skip to main content

memra_engine/
glm5_tp.rs

1//! glm5_next (GLM-5.3-Flash) TP-N seam — `MEMRA_GLM5_TP` (lane/glm5-tp2, 2026-08-31;
2//! rank-widened to TP-4 by lane/glm5-composition, 2026-09-01).
3//!
4//! WHAT THIS IS. A correctness-first tensor-parallel execution program for the glm5_next
5//! hybrid trunk, per the lane's shard map (`research/glm53-flash-bringup-20260827/
6//! tp2-20260831/SHARD-MAP.md`). Per layer class:
7//!
8//!   * KDA (34 layers): head-sharded, `heads / ranks` per rank. Each rank runs the UNCHANGED
9//!     `kda_core_gated` program on its shard (per-head kernels: conv, L2 norm, gate, scan,
10//!     gated rmsnorm are all head-independent), the gated `[t, qkv/ranks]` parts are
11//!     gathered, and each rank's COLUMN-parallel `wo` slice (out rows over the FULL gathered
12//!     input) computes its slice of the output with the same plain matvec kernel — joins are
13//!     pure data movement, never a partial-sum reduction, which is what makes model-level
14//!     TP-vs-plain BYTE identity the bar instead of a tolerance band.
15//!   * MLA/DSA (11 layers): head-sharded per-head operands (`wq_b`, `wk_b`, `wv_b`);
16//!     REPLICATED per-token shared work (`wq_a`/`q_a_norm`, `wkv_a`/`kv_a_norm`, the whole
17//!     indexer + k-pool selection) — every rank computes identical bytes from identical
18//!     inputs, so the latent + indexer planes are replicated per rank and no per-token
19//!     cross-rank hop exists in the latent chain. `wo` is column-parallel over the gathered
20//!     attention parts, exactly like KDA.
21//!   * MoE (sparse-FFN layers): EP-N, whole experts, contiguous slices (even split: rank =
22//!     expert / (n_expert/ranks)). The router stays root-computed (host sigmoid top-k,
23//!     unchanged); each owner extracts its slots' UNWEIGHTED down rows with the same
24//!     fused-epilogue kernels at n_used=1, and root re-applies the slot-ordered fmaf
25//!     accumulation chain — the same rounded-operation sequence as the plain
26//!     `moe_down8_fma_q8` walk. Shared expert, dense MLPs, router, mHC, norms, embed and
27//!     lm_head stay ROOT-OWNED (the `MEMRA_STEP_TP` owner-stage precedent).
28//!
29//! TRANSPORT is a SEPARATE, SWAPPABLE AXIS (`MEMRA_GLM5_TP_TRANSPORT`,
30//! lane/glm5-tp-transport 2026-09-01). Because every cross-rank hop above is pure movement,
31//! the transport arm cannot change a bit — so this module names the hop SHAPES and
32//! `tp_transport` owns the bytes. `host-canonical` (the default, and what every banked
33//! glm5 TP number was measured on) bounces each hop through host with a full stream drain
34//! per leg; `peer-pull` issues a consumer-side device peer copy per hop with event ordering
35//! and no host boundary. The join-diet doors are an orthogonal axis (they cut hop COUNT; the
36//! transport cuts hop COST) and compose.
37//!
38//! FAIL-CLOSED SURFACE. The preflight refuses before any TP CUDA state exists: non-glm5
39//! plans, rank counts outside the qualified set (2 and 4 — see [`GLM5_TP_ALLOWED_RANKS`]),
40//! head/expert counts that do not divide, duplicate devices (serving parse), co-armed
41//! `MEMRA_PP_STAGES>1`, `MEMRA_STEP_TP`/`MEMRA_STEP_EP`. A sharded layer POISONS every plain
42//! path: `kda_core`, `mla_attn_cached` and the batched walks refuse a TP-armed layer by
43//! name. The memra-server worker refuses the flag outright (serving wiring is the named
44//! box-lane increment, not v1).
45//!
46//! Engagement markers: `[glm5-tp-preflight]`, `[glm5-tp-kda]`, `[glm5-tp-mla]`,
47//! `[glm5-tp-ep]`, `[glm5-tp-transport]` — every marker carries `performance_claim=false`,
48//! and the first four name the LIVE transport rather than a hardcoded string (the
49//! tp2-battery greps `transport=` on all four seams, and a hardcoded value would have made a
50//! transport A/B unreadable from the boot log).
51
52// lane/clippy-zero-restore-20260901: perf-gated TP2 host code (fresh lane receipts);
53// index loops stay in their gated shape — iterator reshapes are not bit-neutral by inspection.
54#![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
67/// The qualified rank envelope: TP-2 (the v1 seam, box-battery-gated) and TP-4.
68///
69/// GEOMETRY (lane/glm5-tp-transport, 2026-09-01). The DSA indexer is REPLICATED per rank by
70/// this seam's own shard map (`shard_mla_layer`'s `replicate_indexer`), so its 32 heads
71/// impose no divisibility constraint at any rank count. TP-4 needs NO padding on the
72/// glm5_next geometry: 64/4 = 16 KDA heads, 64/4 = 16 MLA heads, 288/4 = 72 experts,
73/// 4096/4 = 1024 `wo` out rows, KDA `head_dim` 128 is rank-count independent. TP-3 remains
74/// refused: the only real obstruction is the 64 attention/KDA heads (a HEAD-PADDING
75/// question, 64 -> 66, RESEARCH.md §1.5d — not built). See `tp-transport-20260901/LANE.md`
76/// "TP-4 divisibility".
77pub const GLM5_TP_ALLOWED_RANKS: [usize; 2] = [2, 4];
78
79/// This family's receipt marker for the GENERAL transport seam (`tp_transport`, generalized
80/// lane/glm5-extract2). The tag is the caller's so lane/glm5-tp-transport's and
81/// lane/glm5-composition's banked gate and box receipts keep their exact bytes while a second
82/// family gets its own marker — the same rule phase 1 set for `[glm5-phase]` on the shared
83/// spec timers.
84pub const GLM5_TP_TRANSPORT_TAG: &str = "glm5-tp-transport";
85
86pub type Glm5TpLayerSpec = crate::tp::StepEpLayerSpec;
87
88// ------------------------------------------------------------------------------------------
89// Flag
90// ------------------------------------------------------------------------------------------
91
92/// Raw `MEMRA_GLM5_TP` value. Empty / unset / `"0"` = seam off.
93pub fn glm5_tp_env_raw() -> Option<String> {
94    std::env::var("MEMRA_GLM5_TP").ok()
95}
96
97/// Cheap armed check for co-refusal sites (server boot, spec doors). Parse errors count as
98/// ARMED so a misspelled spec still refuses the co-armed program instead of racing the
99/// loader's own refusal.
100pub fn glm5_tp_armed() -> bool {
101    matches!(glm5_tp_env_raw().as_deref(), Some(v) if !v.is_empty() && v != "0")
102}
103
104/// Parse the shared `LAYER[-LAYER]@DEVICE,DEVICE[,...][;...]` grammar for the glm5 door.
105/// `trunk_layers` is the loaded model's trunk length (the `all` shorthand expands against
106/// it — the model contract owns that number, never a constant in the parser).
107pub 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
114/// Gate-harness knob, never a serving flag: `MEMRA_GLM5_TP_GATE_SAME_DEV=1` builds every
115/// peer rank as an ADDITIONAL CUDA CONTEXT ON THE ROOT DEVICE (the one-card rig gate's
116/// emulation; the ppN same-device-stages precedent). The spec's non-root device ids become
117/// logical rank ids. The serving worker refuses `MEMRA_GLM5_TP` outright, so this can never
118/// leak into serving.
119pub fn gate_same_device() -> bool {
120    std::env::var("MEMRA_GLM5_TP_GATE_SAME_DEV").as_deref() == Ok("1")
121}
122
123/// Gate-harness RED-arm knob, never a serving flag (`MEMRA_GLM5_TP_GATE_RED`):
124///   * `swap-wo` — each rank's column `wo` slice takes the NEXT rank's out rows (a broken
125///     shard map); the gate run MUST diverge from plain.
126///   * `swap-ep-gateup` — the root EP slab's gate and up projections swap (wrong expert
127///     weights); MUST diverge.
128///   * `skip-peer-combine` — the EP combine drops every peer-owned slot; MUST diverge,
129///     which is also the non-vacuity proof that the peer ranks contribute real work.
130///   * `corrupt-ep-map` — the placement's local-slot table for rank 0 is reversed after
131///     the slabs are built (owner table and slab bytes disagree — a corrupted map row);
132///     MUST diverge. This is the red that proves the MEASURED-placement indirection is
133///     load-bearing, not decorative.
134///
135/// Unknown values refuse at load.
136#[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
158// ------------------------------------------------------------------------------------------
159// Runtime
160// ------------------------------------------------------------------------------------------
161
162/// The TP-N rank runtime. Rank 0 (root) executes on the model's own engine — the PP-owner
163/// context, exactly like the step seam's owner-first rank law. Ranks `1..ranks` each own a
164/// full peer Engine, in `MEMRA_GLM5_TP` device order (`peers[i]` = rank `i + 1`).
165pub struct Glm5TpRt {
166    pub peers: Vec<Engine>,
167    pub root_dev: usize,
168    pub peer_devs: Vec<usize>,
169    /// True only when built through [`Glm5TpRt::new_gate_same_device`] — the one-card rig
170    /// gate's multi-context emulation (the ppN same-device gate precedent). The env-driven
171    /// serving parse can never reach this: the grammar refuses duplicate devices.
172    pub same_device_gate: bool,
173    /// Which transport every cross-rank hop of this runtime moves its bytes with
174    /// (`MEMRA_GLM5_TP_TRANSPORT`, default `host-canonical`). Frozen at
175    /// [`Glm5TpRt::arm_transport`] time, announced once, and named in every gate log.
176    pub transport: crate::tp_transport::TpTransport,
177    /// The peer-pull ordering primitives — `Some` only on the peer-pull arm, and only after
178    /// its byte-integrity ladder passed.
179    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    /// Same-device multi-context runtime for the ONE-CARD rig gate (exactness only). Every
210    /// peer rank is an additional CUDA context on the root device: the whole shard/join
211    /// walk — shard loads, replicated compute, gathers, canonical combines — executes
212    /// exactly as on N cards, minus real peer transport (which the pro6000 batteries
213    /// qualify on the box card class separately).
214    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    /// Rank count of this runtime (root + peers).
233    pub fn ranks(&self) -> usize {
234        self.peers.len() + 1
235    }
236
237    /// Freeze the transport for this runtime: read the flag, grant peer access (real groups
238    /// only), and run the byte-integrity pull ladder over every ordered rank pair. Called
239    /// from the preflight BEFORE any layer is sharded, so a bad fabric refuses the load
240    /// rather than corrupting a shard.
241    pub fn arm_transport(&mut self, root: &Engine) -> Result<(), Box<dyn std::error::Error>> {
242        // The seam is general; the TAG and the flag name are the FAMILY's (lane/glm5-extract2,
243        // the phase-1 caller-owned-tag pattern) — `[glm5-tp-transport]` bytes stay exactly as
244        // lane/glm5-tp-transport and lane/glm5-composition banked them, and `armed_flag` is
245        // whichever of `MEMRA_TP_TRANSPORT` / `MEMRA_GLM5_TP_TRANSPORT` the operator actually
246        // set, so a peer-access or ladder refusal names their flag.
247        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    /// Build the per-hop transport handle. Every cross-rank movement in the glm5 TP walk
262    /// goes through one of `tp_transport`'s named hop shapes with this handle, which is
263    /// what makes the arm swap a ONE-PLACE change and the movement census automatic.
264    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
273// ------------------------------------------------------------------------------------------
274// Preflight
275// ------------------------------------------------------------------------------------------
276
277/// What the loader tells the preflight about the model, extracted from the plan/config
278/// BEFORE any TP CUDA state exists. Structural laws are dimension-derived (they hold for
279/// the mini fixture and the real artifact alike): the laws ARE the geometry checks.
280pub struct Glm5TpModelView {
281    pub trunk_layers: usize,
282    /// Per-layer mixer class, `trunk_layers` entries.
283    pub layer_class: Vec<Glm5LayerClass>,
284    /// Per-layer "has routed-expert FFN" flag (dense-prefix layers are false).
285    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
299/// The armed load plan: the runtime plus the layer set the spec selected, plus the
300/// measured expert-placement map when `MEMRA_EP_MAP` (or its glm5 alias) is armed (validated at
301/// preflight, before any TP CUDA state — absent flag = the even split, byte-unchanged).
302pub 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
308/// Load + validate the placement map against the model view and the armed layer set.
309/// The env seam is the general `ep_map::ep_map_env()` (`MEMRA_EP_MAP`, glm5 alias
310/// honored); every refusal names the flag that ARMED the load. `Some("")` REFUSES (a
311/// set-but-empty flag is an operator error, never a silent even split). Fail-closed on
312/// every axis: unreadable file, malformed text, rank/expert-count mismatch, layer-cover
313/// mismatch. Returns `None` only when both names are UNSET.
314fn 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    // Receipt anchor: the map bytes that armed this load, named by digest.
364    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
382/// Decode-diet doors that never co-arm with `MEMRA_GLM5_TP` in v1 (merge-forward
383/// 2026-08-31): every TP-x-door pair is UNPROVEN. The TP byte/band gates ran with every
384/// door cold, and each door's own gate ran on the unsharded walk, so v1 refuses by name
385/// rather than silently picking an arm; a pair unlocks only with its own composition gate
386/// (the `MEMRA_GLM5_TP` row in docs/FLAGS.md carries the matrix). `MEMRA_GLM5_VERIFY_BATCH`
387/// is absent DELIBERATELY: its walk exists only inside glm5 spec sessions — co-refused on
388/// a sharded model unless `MEMRA_GLM5_SPEC_TP=1` arms the GATED composition
389/// (lane/glm5-composition; the spec x TP pair HAS its composition gate, `glm5-tp-gate`
390/// arms S2/Q-S4), whose admission REQUIRES the batched walk by name.
391pub 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
411/// The pure composition law over [`GLM5_TP_REFUSED_DOOR_FLAGS`]: the first armed door
412/// refuses by name, before any TP CUDA state exists. Delegates to the general
413/// [`crate::tp::refuse_door_composition`] pattern (lane/glm5-extract-general) with this
414/// door's own table — error bytes unchanged. `armed` reports whether a flag is set to
415/// `"1"` (env in production; a plain set in the unit test — the module keeps its tests
416/// env-mutation-free).
417pub 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
421/// Fail-closed preflight + runtime construction. Returns `None` when the seam is off.
422/// Every illegal geometry refuses HERE, before any rank engine or shard exists.
423pub 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    // Co-armed programs refuse by name: two parallel/spec programs on one model never
434    // silently coexist (the MEMRA_DSPARK precedent).
435    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    // One device group across the whole spec (one runtime group), root-first; the rank
454    // count comes from the device list and must be in the qualified envelope.
455    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    // Structural geometry laws, all dimension-derived.
467    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    // Validate the gate red-arm spelling at load (fail-closed), and pick the transport.
537    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    // Transport arms HERE — after the rank engines exist, BEFORE any layer is sharded. A
554    // peer-pull ladder failure refuses the load with zero TP shards built (lane/glm5-tp-transport).
555    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
583// ------------------------------------------------------------------------------------------
584// Shard mechanics
585// ------------------------------------------------------------------------------------------
586
587fn outer_rows(ne: &[u64]) -> (usize, usize) {
588    // GGML axis order: ne[0] is the fastest (innermost). The shardable axis is the LAST
589    // (outermost) — out rows on a 2D projection, the head axis on a 3D per-head slab.
590    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
595/// Copy `rows` of `t`'s outermost axis onto `dst` (host bounce; load-time only). Mirror
596/// planes (`rp`/`rp4`/`f16`/`fp8`/`blk`) REFUSE by name: v1 shards carry the raw layout —
597/// a pure byte-permutation difference, bit-identical by the mirrors' own contracts.
598fn 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
697/// Full replica of `t` on `dst` (host bounce). Same mirror refusals as [`shard_rows`].
698fn 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
707/// Rank r's engine within a runtime, given the root engine (rank 0 has no owned Engine in
708/// the runtime — it IS the model's engine).
709pub(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
713// ------------------------------------------------------------------------------------------
714// KDA sidecar
715// ------------------------------------------------------------------------------------------
716
717/// The KDA TP sidecar: the peer ranks' head shards plus the runtime handle. The OUTER
718/// `KdaAttnLayer` that carries this in its `tp` field is the root shard; every shard's
719/// `wo` field holds that rank's COLUMN slice (out rows over the full `qkv` input).
720pub struct Glm5TpKda {
721    pub rt: Arc<Glm5TpRt>,
722    /// `peers[i]` is rank `i + 1`'s shard, resident on `rt.peers[i]`.
723    pub peers: Vec<KdaAttnLayer>,
724    /// Full-width qkv of the UNSHARDED layer (`ranks * shard qkv`) — the gather width.
725    pub full_qkv: usize,
726    /// Full hidden width (`wo` out rows across all ranks).
727    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
738/// Shard one loaded KDA layer: returns the ROOT shard (heads/ranks, wo out-rows
739/// `0..H/ranks`) with the peer shards in its `tp` sidecar. The full layer's tensors are
740/// consumed and dropped — per-layer transient VRAM is one layer, never the model.
741pub(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; // heads per rank
758    let ql = qkv / ranks; // channels per rank
759    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    // Per-rank fused conv slice: plane p occupies rows [p*qkv, (p+1)*qkv) of the fused
769    // [3*qkv, kernel] buffer; rank r takes channel rows [r*ql, (r+1)*ql) of each plane.
770    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    // Gate red arm: a broken shard map hands each rank the NEXT rank's wo out rows
782    // (the two-rank swap generalized to a rotation — still guaranteed wrong on every rank).
783    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            // COLUMN-parallel wo: rank r owns OUT rows [r*hh, (r+1)*hh) over the FULL qkv
803            // input — consumed by the join over the gathered gated tensor, never by
804            // kda_core_gated itself.
805            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
835/// Ensure layer `il`'s per-rank KDA state planes exist (lazily, sized for the SHARD
836/// geometry — the canonical `cache.recur[il]` planes are full-width and stay untouched
837/// as allocated; the TP walk never reads them). Index 0 = root's plane on `e`, index r =
838/// rank r's plane on its peer engine.
839fn 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/// The KDA TP walk, ONE body for both consumers (the #80 review's dedup finding — the
870/// forked verify twin had already drifted to root-first issue order):
871///   * prime/decode ([`kda_tp_cached`]): `verify_stash = None`, plain `wo` matmul —
872///     byte-for-byte the pre-composition walk.
873///   * spec x TP verify rows ([`kda_tp_verify_rows`]): `verify_stash = Some`, per-rank
874///     pre-round ssm snapshot + batched `KdaStash::Rows` capture, `wo` on the ROWS-EXACT
875///     class (the unsharded verify walk's own routing), per-rank scan-ns accumulated into
876///     `scan_clock` so the `[glm5-phase-v]` receipt keeps its sequential-floor share on
877///     the composed shape.
878///
879/// Issue order is PEERS FIRST, ROOT LAST on both arms (v1's order; the twins document it).
880/// THREE cross-rank hop shapes, each a named `tp_transport` shape: fan-out of `x`,
881/// gather of the gated parts, concat of the `wo` parts. On `host-canonical` at two ranks
882/// that is 5 draining `dtoh` + 4 `htod` per layer-call, exactly as v1; on `peer-pull` it is
883/// device peer copies, local copies and 0 host boundaries.
884#[allow(clippy::too_many_arguments)] // mirrors the kda entry contract shape
885fn 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(); // per-rank channels
904    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    // Per-rank verify capture, RANK-indexed regardless of issue order; assembled into the
909    // caller's stash after the loop.
910    let mut captured: Vec<Option<(CudaSlice<f32>, crate::kda::KdaRowsStash)>> =
911        (0..ranks).map(|_| None).collect();
912
913    let hop = rt.hop(e);
914    // HOP 1 — fan-out of the mixer input to every peer rank. `x.len()` and not `t * n_embd`:
915    // the v1 arm moved the WHOLE buffer, and the arms must move identical byte ranges or
916    // the transport A/B stops being a transport A/B.
917    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    // Peer shards first (host-canonical serial walk; overlap is the box arc), root last —
921    // v1's issue order at two ranks, both arms.
922    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        // Verify arm: the pre-round snapshot on the rank's engine, BEFORE the batched
928        // call advances the resident state (the ckpt contract's per-rank twin).
929        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    // HOP 2 — gather the gated parts into the FULL [t, qkv] layout on EVERY rank
981    // (column-parallel wo needs the whole input on each rank). Token-major interleave: row
982    // tok is [rank0 ql | rank1 ql | ...]. `full == ranks * ql` by the shard map.
983    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    // Per-rank column wo slices: each output element is one full-K dot by the SAME kernel
991    // class the consumer's unsharded walk uses — no cross-rank arithmetic in this join.
992    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    // HOP 3 — concat the column parts into the mixer output on ROOT.
1006    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/// The KDA TP walk for one prime/decode call — [`kda_tp_core`] with no verify capture.
1011#[allow(clippy::too_many_arguments)] // mirrors the kda entry contract shape
1012pub(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
1025/// Per-rank rollback material of ONE sharded-KDA verify round (lane/glm5-composition, the
1026/// spec x TP composition): index = rank; each entry is that rank's pre-round ssm snapshot
1027/// (cloned on the rank's engine BEFORE its batched call advanced the resident state) plus
1028/// the batched [`crate::kda::KdaRowsStash`] its `KdaStash::Rows` call filled. Rollback to
1029/// `keep` rows restores every rank through `kda_verify_rollback_rows_on` with the rank's
1030/// own engine/shard/plane tuple — the same two-plane contract as the unsharded stash,
1031/// per rank.
1032pub type Glm5TpKdaVerifyStash = Vec<(CudaSlice<f32>, crate::kda::KdaRowsStash)>;
1033
1034/// The sharded-KDA VERIFY walk (spec x TP composition) — [`kda_tp_core`] with the verify
1035/// capture armed: batched `KdaStash::Rows` per rank, ROWS-EXACT `wo` (the unsharded verify
1036/// walk's own routing), per-rank scan-ns accumulated into `scan_clock`. Returns the mixer
1037/// output plus the rank-indexed rollback stash the ckpt banks.
1038#[allow(clippy::too_many_arguments)] // mirrors the kda verify entry contract shape
1039pub(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
1065/// Roll every rank's sharded-KDA state back to "after row `keep-1`" from a spec x TP verify
1066/// round (the [`Glm5TpKdaVerifyStash`] contract). Full accept never calls this — the
1067/// resident per-rank states ARE the state after the last kept row.
1068pub(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
1101// ------------------------------------------------------------------------------------------
1102// MLA sidecar
1103// ------------------------------------------------------------------------------------------
1104
1105/// The MLA TP sidecar: the peer ranks' head shards (with replicated `wq_a`/`wkv_a`/norms
1106/// and full indexer replicas) plus the runtime handle.
1107pub struct Glm5TpMla {
1108    pub rt: Arc<Glm5TpRt>,
1109    /// `peers[i]` is rank `i + 1`'s shard, resident on `rt.peers[i]`.
1110    pub peers: Vec<crate::hybrid::MlaAttnLayer>,
1111    /// Full head count of the unsharded layer.
1112    pub full_heads: usize,
1113    /// Full hidden width (`wo` out rows across all ranks).
1114    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; // per-head wq_b out rows
1142    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    // Gate red arm: a broken shard map hands each rank the NEXT rank's wo out rows.
1166    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            // 3D per-head slabs: the head axis is outermost.
1182            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            // COLUMN-parallel wo: rank r owns OUT rows over the full N*V input.
1185            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
1219/// Ensure the PEER ranks' replicated latent planes for layer `il` exist, geometry-cloned
1220/// from the canonical (root) plane. The canonical plane IS the root replica — the root path
1221/// is unchanged. `cache_slot` holds one plane per peer rank (`[i]` = rank `i + 1`).
1222pub(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        // Fresh replica starts at len 0 like a fresh canonical plane; the walk appends to
1234        // every replica in the same calls, so the lengths stay in lock-step by construction.
1235        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, // lazily allocated by the core, exactly like the canonical plane
1249            index_pools_ready: 0,
1250            index_pool: canonical.index_pool,
1251        });
1252    }
1253    *cache_slot = Some(planes);
1254    Ok(())
1255}
1256
1257// ------------------------------------------------------------------------------------------
1258// MoE EP sidecar
1259// ------------------------------------------------------------------------------------------
1260
1261/// One rank's expert slab: the rank's owned experts packed in ASCENDING expert-id order
1262/// for every projection, device-resident on that rank. For the even split the packing is
1263/// the contiguous slice — byte-for-byte the pre-map layout.
1264pub struct EpRankSlab {
1265    pub gate: CudaSlice<u8>,
1266    pub up: CudaSlice<u8>,
1267    pub down: CudaSlice<u8>,
1268    pub n_experts: usize,
1269}
1270
1271/// The MoE EP sidecar on `MoeWeights`: per-rank expert slabs, the placement tables,
1272/// and the runtime handle. Router, shared expert, macros and all host metadata stay on
1273/// the unchanged `MoeWeights`.
1274///
1275/// PLACEMENT INDEPENDENCE (the contract the gate's skewed-map arm proves): `owner_of`
1276/// only selects WHICH rank runs the identical per-expert program over identical
1277/// host-canonical input bytes; `local_of` indexes the same expert bytes wherever they
1278/// were packed; the combine stays slot-ordered on root. The map moves bytes, never
1279/// changes arithmetic.
1280pub struct Glm5EpExps {
1281    pub rt: Arc<Glm5TpRt>,
1282    /// `slabs[r]` is rank r's expert slab (`[0]` = root's, on the model's engine).
1283    pub slabs: Vec<EpRankSlab>,
1284    /// `owner_of[expert]` = owning rank (0 = root).
1285    pub owner_of: Vec<u8>,
1286    /// `local_of[expert]` = slot inside the owner's slab (ascending-id packing order).
1287    pub local_of: Vec<u32>,
1288    /// Per-rank grouped-dispatch pointer tables, `[rank]`, each the `DevExps::ptr_row`
1289    /// shape ([3 * n_expert] u64 device pointers: gate | up | down planes, indexed by GLOBAL
1290    /// expert id, resident on the owning rank's device). Owned experts point at
1291    /// `slab_base + local * stride`; non-owned entries are 0 and never dereferenced — the EP
1292    /// grouped-prime CSR is built per rank from `owner_of`, so a foreign id cannot reach the
1293    /// wrong rank's table. Built AFTER the gate-red slab mutations, from the FINAL slab
1294    /// buffers and the FINAL `local_of`, so `swap-ep-gateup` and `corrupt-ep-map` bite the
1295    /// grouped walk exactly as they bite the sequential one.
1296    pub ptr_rows: Vec<CudaSlice<u64>>,
1297}
1298
1299impl Glm5EpExps {
1300    /// Owner rank of `expert` under the armed placement (even split when no map).
1301    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
1312/// Engagement counter: PEER-owned expert slots dispatched by the EP walk (counted before
1313/// any gate-red skip, so a red arm can still assert a peer was ROUTED). Gates read it to
1314/// prove the peer ranks contribute real expert work — a token stream that never routes a
1315/// peer-owned expert makes every EP identity arm vacuous.
1316pub 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
1322// ---- EP dispatch-diet engagement counters (lane/glm5-ep-diet, 2026-08-31) ----------------
1323// The box A/B greps announces and reads these deltas; the rig gate asserts them non-vacuous
1324// on the ON arms and FLAT on the pinned-`=0` arms.
1325
1326/// Layer-calls that took the dieted EP walk (`MEMRA_GLM5_EP_DIET`) instead of the v1
1327/// per-slot host-canonical walk.
1328pub static GLM5_EP_DIET_DISPATCHES: AtomicU64 = AtomicU64::new(0);
1329
1330/// Snapshot of [`GLM5_EP_DIET_DISPATCHES`] — gates take a before/after delta.
1331pub fn glm5_ep_diet_dispatches() -> u64 {
1332    GLM5_EP_DIET_DISPATCHES.load(Ordering::Relaxed)
1333}
1334
1335/// Bulk peer-row block returns performed by the dieted walk (one per (layer-call, peer
1336/// rank) that routed at least one slot owned by that rank; each replaces that call's ENTIRE
1337/// per-slot return dribble for that rank).
1338pub static GLM5_EP_DIET_BULK_RETURNS: AtomicU64 = AtomicU64::new(0);
1339
1340/// Snapshot of [`GLM5_EP_DIET_BULK_RETURNS`].
1341pub fn glm5_ep_diet_bulk_returns() -> u64 {
1342    GLM5_EP_DIET_BULK_RETURNS.load(Ordering::Relaxed)
1343}
1344
1345/// Per-slot synchronous peer round-trips (one peer DtoH + one root pageable HtoD each, the
1346/// v1 walk's dominant hop class) that the dieted walk folded into its bulk returns — one
1347/// count per peer-owned slot bulked.
1348pub static GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED: AtomicU64 = AtomicU64::new(0);
1349
1350/// Snapshot of [`GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED`].
1351pub fn glm5_ep_diet_peer_roundtrips_avoided() -> u64 {
1352    GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED.load(Ordering::Relaxed)
1353}
1354
1355/// Per-token peer z uploads the dieted walk avoided: `t-1` per (fanned layer-call, peer
1356/// rank) (one bulk [t, n_embd] upload replaces t per-token uploads) plus `t` per (layer-call,
1357/// rank) whose routing never touched that rank's experts (the fan-out is skipped entirely —
1358/// the placement-map multiplier: single-rank layer-calls move ZERO activation bytes off
1359/// root).
1360pub static GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED: AtomicU64 = AtomicU64::new(0);
1361
1362/// Snapshot of [`GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED`].
1363pub fn glm5_ep_diet_fanout_uploads_avoided() -> u64 {
1364    GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED.load(Ordering::Relaxed)
1365}
1366
1367/// Layer-calls that took the per-rank grouped-GEMM EP prime (`MEMRA_GLM5_EP_GROUPED_PRIME`).
1368/// Stays 0 whenever the plain grouped-prefill conjuncts do not hold (e.g. non-f16g-eligible
1369/// expert qtypes — the rig fixture's Q8_0 bank always falls closed to the sequential walk).
1370pub static GLM5_EP_GROUPED_PRIME_DISPATCHES: AtomicU64 = AtomicU64::new(0);
1371
1372/// Snapshot of [`GLM5_EP_GROUPED_PRIME_DISPATCHES`].
1373pub fn glm5_ep_grouped_prime_dispatches() -> u64 {
1374    GLM5_EP_GROUPED_PRIME_DISPATCHES.load(Ordering::Relaxed)
1375}
1376
1377/// Arm one MoE layer for EP. `placement` is the layer's validated map row
1378/// (`owners[expert] = rank`) when `MEMRA_EP_MAP` (or its glm5 alias) is armed; `None` = the
1379/// even split, whose ascending-id packing is byte-for-byte the pre-map contiguous slices.
1380pub(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            // The preflight validated the map; re-assert the two structural laws at the
1404            // consumption site so a wiring bug can never hand a foreign row to a layer.
1405            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    // Ascending-id packing per rank + the local-slot table.
1422    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            // Tail-slack pads mirror the resident-slab builder (`build_dev_exps`): 8 B
1435            // alignment slack on gate/up and 144 B on down — the ragged-k grouped GEMM
1436            // walks whole superblocks and may overread past the LAST row (harmless bytes,
1437            // the zero-padded k-range multiplies them away; the slack only prevents the
1438            // OOB fault). Bytes at every in-slab offset are unchanged, so the sequential
1439            // per-slot views read exactly what they read before.
1440            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                // Contiguous ascending run (the even split, and any contiguous map row):
1446                // one direct upload of the existing byte range — no host copy.
1447                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                // General map row: pack the owned experts ascending into one staging
1454                // buffer (load-time only; per-layer transient host = one rank's slab).
1455                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    // Gate red arm: wrong expert weights on the root rank (gate/up swapped).
1473    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    // Gate red arm: a corrupted map row — the local-slot table for rank 0 is reversed
1478    // AFTER the slabs were packed, so the owner table and the slab bytes disagree.
1479    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    // Per-rank grouped-dispatch pointer tables (lane/glm5-ep-diet): the `DevExps::ptr_row`
1486    // shape over each rank's OWN slab, built from the FINAL slab buffers and the FINAL
1487    // `local_of` so both gate reds above flow into the grouped walk too. ~3*n_expert*8 B per
1488    // rank per layer — negligible next to the slabs they index.
1489    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; // non-owned: 0, never dereferenced (rank CSRs filter by owner)
1505            }
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            // ORDER MATTERS and it was wrong once: the ownership string must land on
1524            // `ownership={}` and the transport on `transport={}`. The first gate run's
1525            // receipt-extract printed `[glm5-tp-ep] transport=even-split`, which is how the
1526            // swap was caught — a receipt line is only worth what its argument order is.
1527            if placement.is_some() {
1528                "measured-map"
1529            } else {
1530                "even-split"
1531            },
1532            rt.transport.name(),
1533        );
1534    }
1535    // The root-resident full slab (if the loader built one) is superseded by the EP slices;
1536    // dropping it returns its VRAM and removes the arm that would silently bypass EP.
1537    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        // Off spellings.
1555        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        // The full-model shorthand expands against the CALLER's trunk, not a constant.
1559        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        // The TP-4 device list parses through the same grammar.
1565        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        // Explicit ranges.
1569        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        // Refusals: duplicate devices, duplicate layers, garbage.
1575        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    /// Structural preflight refusals, exercised WITHOUT constructing any CUDA state: every
1599    /// geometry law here fires before `prepare_glm5_tp_load` reaches the runtime build.
1600    /// (The armed happy path needs an Engine and lives in the gate binary.)
1601    #[test]
1602    fn preflight_geometry_laws_are_dimension_derived() {
1603        // The checks below mirror prepare_glm5_tp_load's law order on the view alone, at
1604        // BOTH qualified rank counts.
1605        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        // TP-3 stays outside the qualified envelope (head padding not built).
1627        assert!(!GLM5_TP_ALLOWED_RANKS.contains(&3));
1628    }
1629
1630    #[test]
1631    fn armed_check_counts_parse_errors_as_armed() {
1632        // glm5_tp_armed is a cheap co-refusal predicate: any nonempty non-"0" value counts,
1633        // including a spec the parser would refuse — the co-armed program must not race the
1634        // loader's own refusal.
1635        // (Env-mutation-free: the predicate's contract is pure string classification.)
1636        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        // The merge-forward composition matrix (2026-08-31): each decode-diet door armed
1651        // alone must refuse, naming BOTH flags — a silent pick is the failure mode this
1652        // guards. (Env-mutation-free: the law is pure over the armed predicate; the live
1653        // env read is one closure at the prepare_glm5_tp_load call site, and the tp-gate
1654        // red receipt exercises it end to end.)
1655        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        // All doors cold = no refusal.
1663        refuse_glm5_tp_door_composition(|_| false).expect("cold doors must pass");
1664        // The verify-batch flag is DELIBERATELY not in the matrix (the gated spec x TP
1665        // composition owns that pair — its admission REQUIRES the batched walk); arming
1666        // it alone must not trip this law.
1667        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}