metrale_model_arch/glm5next_mlp/forward/
workspace.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The byte sizes of the GLM-5.3 MLP scratch buffers and
4//! `Glm5NextMlpWorkspace::new`, which allocates them.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants:
8//! - `mlp_ws_bytes` lists the sizes `Glm5NextMlpWorkspace::new` allocates, in its order.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::GpuBackend;
12
13use super::Glm5NextMlpWorkspace;
14use crate::glm5next_mlp::Glm5NextMlpConfig;
15
16/// 2026-09-25: Byte size of each `Glm5NextMlpWorkspace` buffer, in the order `new` allocates
17/// them: a_gate, a_up, a_act, logits, ids, wts, expert_out, shared_out, u_eid, u_slot,
18/// sorted_token_ids, sorted_expert_ids, expert_offsets, token_to_perm. The loader logs their
19/// sum. `new` computes the same sizes on its own, so the two must change together.
20pub fn mlp_ws_bytes(cfg: &Glm5NextMlpConfig, max_rows: usize) -> [usize; 14] {
21    let rows = max_rows.max(1);
22    let max_inter = cfg
23        .local_dense_intermediate
24        .max(cfg.moe_intermediate)
25        .max(cfg.local_shared_intermediate)
26        .max(1);
27    let act_elems = (rows * max_inter)
28        .max(rows * cfg.top_k * cfg.moe_intermediate)
29        .max(1);
30    [
31        act_elems * 2,
32        act_elems * 2,
33        act_elems * 2,
34        rows * cfg.num_experts * 4,
35        rows * cfg.top_k * 4,
36        rows * cfg.top_k * 4,
37        rows * cfg.top_k * cfg.hidden * 2,
38        rows * cfg.hidden * 2,
39        rows * cfg.top_k * 4,
40        rows * cfg.top_k * rows * 4,
41        rows * cfg.top_k * 4,
42        rows * cfg.top_k * 4,
43        (cfg.num_experts + 1) * 4,
44        rows * cfg.top_k * 4,
45    ]
46}
47
48/// 2026-09-25: Total device bytes of one MLP workspace of `max_rows` rows.
49pub fn mlp_ws_total_bytes(cfg: &Glm5NextMlpConfig, max_rows: usize) -> usize {
50    mlp_ws_bytes(cfg, max_rows).iter().sum()
51}
52
53impl Glm5NextMlpWorkspace {
54    pub fn new(gpu: &dyn GpuBackend, cfg: &Glm5NextMlpConfig, max_rows: usize) -> Result<Self> {
55        let rows = max_rows.max(1);
56        let max_inter = cfg
57            .local_dense_intermediate
58            .max(cfg.moe_intermediate)
59            .max(cfg.local_shared_intermediate)
60            .max(1);
61        // 2026-09-25: The activation buffers hold the widest of a dense or shared-expert pass
62        // (`rows * max_inter`) and a routed pass over every (row, slot)
63        // (`rows * top_k * moe_intermediate`).
64        let act_elems = (rows * max_inter)
65            .max(rows * cfg.top_k * cfg.moe_intermediate)
66            .max(1);
67        Ok(Self {
68            a_gate: gpu.alloc(act_elems * 2)?,
69            a_up: gpu.alloc(act_elems * 2)?,
70            a_act: gpu.alloc(act_elems * 2)?,
71            logits: gpu.alloc(rows * cfg.num_experts * 4)?,
72            ids: gpu.alloc(rows * cfg.top_k * 4)?,
73            wts: gpu.alloc(rows * cfg.top_k * 4)?,
74            expert_out: gpu.alloc(rows * cfg.top_k * cfg.hidden * 2)?,
75            shared_out: gpu.alloc(rows * cfg.hidden * 2)?,
76            u_eid: gpu.alloc(rows * cfg.top_k * 4)?,
77            u_slot: gpu.alloc(rows * cfg.top_k * rows * 4)?,
78            // 2026-09-25: The grouped-GEMM routing tables are allocated whatever the grouped
79            // levers say.
80            sorted_token_ids: gpu.alloc(rows * cfg.top_k * 4)?,
81            sorted_expert_ids: gpu.alloc(rows * cfg.top_k * 4)?,
82            expert_offsets: gpu.alloc((cfg.num_experts + 1) * 4)?,
83            token_to_perm: gpu.alloc(rows * cfg.top_k * 4)?,
84            max_inter,
85            max_rows: rows,
86            max_total_expanded: rows * cfg.top_k,
87        })
88    }
89}
metrale_model_arch/glm5next_mlp/forward/workspace.rs

metrale_model_arch/glm5next_mlp/forward/
workspace.rs
