metrale_model_arch/glm5next_mlp/
build.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Bind one GLM-5.3 MLP site for a rank: TP-slice the dense FFN or shared expert,
4//! and build the routed experts' pointer tables over GLOBAL expert ids.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants:
8//! - TP: `gate_proj`/`up_proj` (`[inter, hidden]`) are split by row and `down_proj`
9//!   (`[hidden, inter]`) by column, so the dense output is a partial sum.
10//! - EP: an expert is owned whole by one rank. In each pointer table, an id another rank owns
11//!   keeps a null `packed` pointer.
12//! - The router weight and bias are loaded unsliced on every rank.
13
14use anyhow::{Result, bail};
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
16
17use super::Glm5NextMlpConfig;
18use super::weights::{
19    Glm5NextDenseMlpWeights, Glm5NextExpertPtrTable, Glm5NextExpertWeights, Glm5NextMoePtrTables,
20    Glm5NextMoeWeights, Nvfp4Proj,
21};
22
23/// 2026-09-25: A BF16 or F32 tensor as host `f32`, by layer-relative name.
24pub type LoadFn<'a> = &'a dyn Fn(&str) -> Result<Vec<f32>>;
25/// 2026-09-25: One routed expert's weights, by GLOBAL id. Experts are not sliced, so the
26/// loader's closure (`bind_expert`) can return NVFP4 pointers straight from the weight store.
27pub type ExpertFn<'a> = &'a dyn Fn(usize) -> Result<Glm5NextExpertWeights>;
28
29/// 2026-09-25: Rows `[start, end)` of a `[rows, row_elems]` row-major tensor, for a
30/// column-parallel projection.
31fn row_slice(v: &[f32], row_elems: usize, start: usize, end: usize) -> Vec<f32> {
32    v[start * row_elems..end * row_elems].to_vec()
33}
34
35/// 2026-09-25: Columns `[start, end)` of every row, for the row-parallel `down_proj`.
36fn col_slice(v: &[f32], row_elems: usize, start: usize, end: usize) -> Vec<f32> {
37    v.chunks(row_elems)
38        .flat_map(|r| r[start..end].iter().copied())
39        .collect()
40}
41
42fn up_bf16(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
43    let b: Vec<u8> = v
44        .iter()
45        .flat_map(|x| half::bf16::from_f32(*x).to_le_bytes())
46        .collect();
47    let p = gpu.alloc(b.len().max(1))?;
48    gpu.copy_h2d(&b, p)?;
49    Ok(p)
50}
51
52fn up_f32(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
53    let b: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
54    let p = gpu.alloc(b.len().max(1))?;
55    gpu.copy_h2d(&b, p)?;
56    Ok(p)
57}
58
59/// 2026-09-25: TP-slice and upload one BF16 SwiGLU MLP, a dense layer or a routed layer's
60/// shared expert. `full_inter` is the unsliced width; the rank keeps `full_inter / tp`. Errors
61/// when `full_inter` does not divide over TP or a tensor has the wrong element count.
62pub fn build_dense_mlp(
63    gpu: &dyn GpuBackend,
64    cfg: &Glm5NextMlpConfig,
65    tp_rank: usize,
66    full_inter: usize,
67    prefix: &str,
68    load: LoadFn<'_>,
69) -> Result<Glm5NextDenseMlpWeights> {
70    let tp = cfg.tp_world_size;
71    if !full_inter.is_multiple_of(tp) {
72        bail!("GLM MLP {prefix}: intermediate {full_inter} does not divide over tp {tp}");
73    }
74    let local = full_inter / tp;
75    let lo = tp_rank * local;
76
77    let get = |n: &str| -> Result<Vec<f32>> { load(&format!("{prefix}.{n}")) };
78
79    let expect = |name: &str, v: &[f32], want: usize| -> Result<()> {
80        if v.len() != want {
81            bail!(
82                "GLM MLP {prefix}.{name}: {} elements, expected {want}",
83                v.len()
84            );
85        }
86        Ok(())
87    };
88
89    let gate = get("gate_proj.weight")?;
90    expect("gate_proj.weight", &gate, full_inter * cfg.hidden)?;
91    let up = get("up_proj.weight")?;
92    expect("up_proj.weight", &up, full_inter * cfg.hidden)?;
93    let down = get("down_proj.weight")?;
94    expect("down_proj.weight", &down, cfg.hidden * full_inter)?;
95
96    Ok(Glm5NextDenseMlpWeights {
97        gate_proj: up_bf16(gpu, &row_slice(&gate, cfg.hidden, lo, lo + local))?,
98        up_proj: up_bf16(gpu, &row_slice(&up, cfg.hidden, lo, lo + local))?,
99        down_proj: up_bf16(gpu, &col_slice(&down, full_inter, lo, lo + local))?,
100    })
101}
102
103/// 2026-09-25: One projection's pointer table over all `num_experts` GLOBAL ids. An id another
104/// rank owns keeps a null `packed` pointer, which the expert kernels skip.
105fn build_expert_ptr_table(
106    gpu: &dyn GpuBackend,
107    cfg: &Glm5NextMlpConfig,
108    experts: &[Glm5NextExpertWeights],
109    proj: impl Fn(&Glm5NextExpertWeights) -> Nvfp4Proj,
110) -> Result<Glm5NextExpertPtrTable> {
111    let n = cfg.num_experts;
112    let mut packed = vec![0u8; n * 8];
113    let mut scale = vec![0u8; n * 8];
114    let mut scale2 = vec![0u8; n * 4];
115    for id in 0..n {
116        let Some(local) = cfg.local_slot(id) else {
117            continue;
118        };
119        let p = proj(&experts[local]);
120        packed[id * 8..id * 8 + 8].copy_from_slice(&p.packed.0.to_le_bytes());
121        scale[id * 8..id * 8 + 8].copy_from_slice(&p.scale.0.to_le_bytes());
122        scale2[id * 4..id * 4 + 4].copy_from_slice(&p.scale_2.to_le_bytes());
123    }
124    let packed_ptrs = gpu.alloc(packed.len())?;
125    gpu.copy_h2d(&packed, packed_ptrs)?;
126    let scale_ptrs = gpu.alloc(scale.len())?;
127    gpu.copy_h2d(&scale, scale_ptrs)?;
128    let scale2_vals = gpu.alloc(scale2.len())?;
129    gpu.copy_h2d(&scale2, scale2_vals)?;
130    Ok(Glm5NextExpertPtrTable {
131        packed_ptrs,
132        scale_ptrs,
133        scale2_vals,
134    })
135}
136
137/// 2026-09-25: Bind one routed MoE site for this rank: the router and bias unsliced, the shared
138/// expert TP-sliced, and the `local_experts` routed experts this EP rank owns.
139pub fn build_moe(
140    gpu: &dyn GpuBackend,
141    cfg: &Glm5NextMlpConfig,
142    tp_rank: usize,
143    full_shared_inter: usize,
144    load: LoadFn<'_>,
145    expert: ExpertFn<'_>,
146) -> Result<Glm5NextMoeWeights> {
147    // 2026-09-25: The router weight and bias are loaded whole on every rank, so every rank
148    // selects the same experts.
149    let router = load("mlp.gate.weight")?;
150    if router.len() != cfg.num_experts * cfg.hidden {
151        bail!(
152            "GLM MoE mlp.gate.weight: {} elements, expected {} ({} experts x {} hidden)",
153            router.len(),
154            cfg.num_experts * cfg.hidden,
155            cfg.num_experts,
156            cfg.hidden
157        );
158    }
159    let bias = load("mlp.gate.e_score_correction_bias")?;
160    if bias.len() != cfg.num_experts {
161        bail!(
162            "GLM MoE e_score_correction_bias: {} entries, expected {}",
163            bias.len(),
164            cfg.num_experts
165        );
166    }
167
168    let shared = build_dense_mlp(
169        gpu,
170        cfg,
171        tp_rank,
172        full_shared_inter,
173        "mlp.shared_experts",
174        load,
175    )?;
176
177    // 2026-09-25: Ascending GLOBAL id: slot `i` holds id `local_expert_range().start + i`, the
178    // inverse of `Glm5NextMlpConfig::local_slot`.
179    let mut experts = Vec::with_capacity(cfg.local_experts);
180    for id in cfg.local_expert_range() {
181        experts.push(expert(id)?);
182    }
183
184    let ptrs = Glm5NextMoePtrTables {
185        gate: build_expert_ptr_table(gpu, cfg, &experts, |e| e.gate_proj)?,
186        up: build_expert_ptr_table(gpu, cfg, &experts, |e| e.up_proj)?,
187        down: build_expert_ptr_table(gpu, cfg, &experts, |e| e.down_proj)?,
188    };
189
190    Ok(Glm5NextMoeWeights {
191        router: up_bf16(gpu, &router)?,
192        // 2026-09-25: FP32: `glm5next_router_topk` reads `bias` as `const float*`.
193        router_bias: up_f32(gpu, &bias)?,
194        shared,
195        experts,
196        ptrs,
197    })
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    /// 2026-09-25: `row_slice` takes rows and `col_slice` takes columns, on tensors whose values
205    /// encode their index.
206    #[test]
207    fn dense_slicing_takes_rows_for_gate_and_columns_for_down() {
208        let gate: Vec<f32> = (0..4)
209            .flat_map(|r| (0..3).map(move |c| (r * 10 + c) as f32))
210            .collect();
211        assert_eq!(
212            row_slice(&gate, 3, 2, 4),
213            vec![20., 21., 22., 30., 31., 32.]
214        );
215
216        let down: Vec<f32> = (0..3)
217            .flat_map(|r| (0..4).map(move |c| (r * 10 + c) as f32))
218            .collect();
219        assert_eq!(col_slice(&down, 4, 2, 4), vec![2., 3., 12., 13., 22., 23.]);
220        // 2026-09-25: On a square tensor the two slices have the same length and different
221        // values, so only the values tell them apart.
222        let sq: Vec<f32> = (0..16).map(|i| i as f32).collect();
223        let by_row = row_slice(&sq, 4, 2, 4);
224        let by_col = col_slice(&sq, 4, 2, 4);
225        assert_eq!(by_row.len(), by_col.len());
226        assert_ne!(by_row, by_col);
227    }
228}
metrale_model_arch/glm5next_mlp/build.rs

metrale_model_arch/glm5next_mlp/
build.rs
