metrale_model_arch/glm5next_mlp/
weights.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Device weight handles for one GLM-5.3 MLP site, already sharded for this rank.
4//!
5//! Owner: model-arch (GLM-5.3).
6//! Invariants: none beyond the types; `build.rs` fills them.
7
8use metrale_gpu_runtime::gpu::DevicePtr;
9
10/// 2026-09-25: One NVFP4 projection: packed `e2m1` pairs, per-16-element `e4m3` block scales,
11/// and one global `f32` scale.
12///
13/// `scale_2` is held on the host and passed to the kernels by value (`arg_f32`), not as a
14/// device pointer.
15#[derive(Debug, Clone, Copy)]
16pub struct Nvfp4Proj {
17    /// 2026-09-25: `[out, in/2]` U8, two `e2m1` codes per byte.
18    pub packed: DevicePtr,
19    /// 2026-09-25: `[out, in/16]` F8_E4M3 block scales.
20    pub scale: DevicePtr,
21    /// 2026-09-25: The single global F32 scale.
22    pub scale_2: f32,
23}
24
25/// 2026-09-25: A BF16 SwiGLU MLP: the dense layers `0..first_k_dense_replace`, and the shared
26/// expert of every routed layer.
27///
28/// TP: `gate_proj`/`up_proj` (stored `[inter, hidden]`) are sliced by row, `down_proj` (stored
29/// `[hidden, inter]`) by column, so the output is a partial sum when `tp_world_size > 1`.
30#[derive(Debug, Clone, Copy)]
31pub struct Glm5NextDenseMlpWeights {
32    /// 2026-09-25: `[local_inter, hidden]` BF16.
33    pub gate_proj: DevicePtr,
34    /// 2026-09-25: `[local_inter, hidden]` BF16.
35    pub up_proj: DevicePtr,
36    /// 2026-09-25: `[hidden, local_inter]` BF16.
37    pub down_proj: DevicePtr,
38}
39
40/// 2026-09-25: One routed NVFP4 expert, owned whole by one EP rank.
41#[derive(Debug, Clone, Copy)]
42pub struct Glm5NextExpertWeights {
43    pub gate_proj: Nvfp4Proj,
44    pub up_proj: Nvfp4Proj,
45    pub down_proj: Nvfp4Proj,
46}
47
48/// 2026-09-25: Device pointer tables for one projection over the full expert set.
49///
50/// Indexed by global expert id, so the grouped kernels go from the router's on-device `ids` to
51/// weights without a host read. Experts another EP rank owns have null `packed`/`scale`
52/// pointers; the kernels write nothing for those slots, and the caller zeroes the output first.
53#[derive(Debug, Clone, Copy)]
54pub struct Glm5NextExpertPtrTable {
55    /// 2026-09-25: `[num_experts]` U64 device pointers to each expert's packed NVFP4 weight.
56    pub packed_ptrs: DevicePtr,
57    /// 2026-09-25: `[num_experts]` U64 device pointers to each expert's block scales.
58    pub scale_ptrs: DevicePtr,
59    /// 2026-09-25: `[num_experts]` F32 per-expert global scale.
60    pub scale2_vals: DevicePtr,
61}
62
63/// 2026-09-25: The three projections' pointer tables for one routed site.
64#[derive(Debug, Clone, Copy)]
65pub struct Glm5NextMoePtrTables {
66    pub gate: Glm5NextExpertPtrTable,
67    pub up: Glm5NextExpertPtrTable,
68    pub down: Glm5NextExpertPtrTable,
69}
70
71/// 2026-09-25: A routed MoE site's weights for this rank.
72pub struct Glm5NextMoeWeights {
73    /// 2026-09-25: `[num_experts, hidden]` BF16 router, the full matrix on every rank (see the
74    /// module header). Run through the FP32-output GEMM.
75    pub router: DevicePtr,
76    /// 2026-09-25: `[num_experts]` F32 selection bias (`gate.e_score_correction_bias`). It
77    /// steers selection only; the emitted weight is the chosen expert's unbiased score.
78    pub router_bias: DevicePtr,
79    /// 2026-09-25: The shared expert: dense BF16, TP-sharded, added to every token unscaled.
80    pub shared: Glm5NextDenseMlpWeights,
81    /// 2026-09-25: `local_experts` entries in ascending global id, indexed by
82    /// `Glm5NextMlpConfig::local_slot(global_id)`, never by the global id.
83    pub experts: Vec<Glm5NextExpertWeights>,
84    /// 2026-09-25: Global-id-indexed pointer tables over the same experts, for the
85    /// device-dispatched forward. Null entries mark remote ids.
86    pub ptrs: Glm5NextMoePtrTables,
87}
metrale_model_arch/glm5next_mlp/weights.rs

metrale_model_arch/glm5next_mlp/
weights.rs
