spark_model/layers/glm5next_mlp/
weights.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The GLM MLP weight contract, already sharded for this rank.
4
5use spark_runtime::gpu::DevicePtr;
6
7/// One NVFP4 projection as ModelOpt stores it: packed `e2m1` pairs, per-block `e4m3` scales,
8/// and one global `f32` scale.
9///
10/// 🪤 The three travel together and are meaningless apart. `weight_scale_2` is a **scalar read
11/// off the device at load** and passed by value — `w4a16_gemm` takes it as `arg_f32`, not as a
12/// pointer, so uploading it and passing the address silently reinterprets a pointer as a float.
13#[derive(Debug, Clone, Copy)]
14pub struct Nvfp4Proj {
15    /// `[out, in/2]` U8 — two `e2m1` codes per byte.
16    pub packed: DevicePtr,
17    /// `[out, in/16]` F8_E4M3 block scales.
18    pub scale: DevicePtr,
19    /// The single global F32 scale.
20    pub scale_2: f32,
21}
22
23/// A dense SwiGLU MLP in BF16 — layers `0..first_k_dense_replace`, and the shared expert of
24/// every routed layer. Same shape, same kernels, different widths.
25///
26/// TP: `gate_proj`/`up_proj` are column-parallel (**row**-sliced, since each is stored `[inter,
27/// hidden]`), `down_proj` is row-parallel (**column**-sliced on `[hidden, inter]`). The output
28/// is therefore a partial sum whenever `tp_world_size > 1`.
29#[derive(Debug, Clone, Copy)]
30pub struct Glm5NextDenseMlpWeights {
31    /// `[local_inter, hidden]` BF16.
32    pub gate_proj: DevicePtr,
33    /// `[local_inter, hidden]` BF16.
34    pub up_proj: DevicePtr,
35    /// `[hidden, local_inter]` BF16 — row-parallel.
36    pub down_proj: DevicePtr,
37}
38
39/// One routed expert. NVFP4, owned **whole** by one EP rank — never split further.
40#[derive(Debug, Clone, Copy)]
41pub struct Glm5NextExpertWeights {
42    pub gate_proj: Nvfp4Proj,
43    pub up_proj: Nvfp4Proj,
44    pub down_proj: Nvfp4Proj,
45}
46
47/// Device-side pointer tables for ONE projection across the **full** expert set.
48///
49/// Indexed by GLOBAL expert id, so the grouped kernel can go straight from the router's
50/// on-device `ids` to weights with no host round trip. Experts another EP rank owns carry a
51/// **null** `packed`/`scale` pointer; the kernel writes nothing for those slots and the
52/// caller's pre-zeroed output row stands.
53#[derive(Debug, Clone, Copy)]
54pub struct Glm5NextExpertPtrTable {
55    /// `[num_experts]` U64 device pointers to each expert's packed NVFP4 weight.
56    pub packed_ptrs: DevicePtr,
57    /// `[num_experts]` U64 device pointers to each expert's block scales.
58    pub scale_ptrs: DevicePtr,
59    /// `[num_experts]` F32 per-expert `weight_scale_2`.
60    pub scale2_vals: DevicePtr,
61}
62
63/// 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/// A routed MoE site's weights for this rank.
72pub struct Glm5NextMoeWeights {
73    /// `[num_experts, hidden]` BF16 router. 🪤 **REPLICATED, and it must stay that way** — see
74    /// the module header. Consumed through the FP32-out GEMM.
75    pub router: DevicePtr,
76    /// `[num_experts]` F32 selection bias (`gate.e_score_correction_bias`).
77    ///
78    /// 🪤 Steers SELECTION ONLY. The emitted weight is the chosen expert's *unbiased* score.
79    pub router_bias: DevicePtr,
80    /// The shared expert — dense BF16, TP-sharded, added to every token unscaled.
81    pub shared: Glm5NextDenseMlpWeights,
82    /// Exactly `local_experts` entries, indexed by **local slot**, in ascending global id.
83    ///
84    /// 🪤 Indexed by `Glm5NextMlpConfig::local_slot(global_id)`, never by the global id.
85    pub experts: Vec<Glm5NextExpertWeights>,
86    /// Global-id-indexed device pointer tables over the same experts, for the grouped
87    /// device-dispatch forward. Null entries mark remote ids.
88    pub ptrs: Glm5NextMoePtrTables,
89}
spark_model/layers/glm5next_mlp/weights.rs

spark_model/layers/glm5next_mlp/
weights.rs
