metrale_model_layers/weight_map/
moe.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Weight structs for MoE layers, whole transformer layers and the MTP head.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use anyhow::{Context, Result, bail, ensure};
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12use metrale_model_weights::weights::{WeightDtype, WeightStore};
13
14use super::*;
15
16/// 2026-09-25: Weights of one MoE block.
17pub struct MoeWeights {
18    /// 2026-09-25: Router gate, BF16, one row of `hidden_size` per router logit.
19    pub gate: DenseWeight,
20    pub shared_expert: ExpertWeight,
21    /// 2026-09-25: Shared-expert gate `[1, hidden_size]` BF16: the shared
22    /// expert's output is scaled by `sigmoid(dot(input, gate))`, or by 1 when
23    /// this pointer is NULL.
24    pub shared_expert_gate: DenseWeight,
25    /// 2026-09-25: One entry per expert id; experts this rank does not hold are
26    /// `ExpertWeight::null()`.
27    pub experts: Vec<ExpertWeight>,
28    /// 2026-09-25: RMSNorm weight applied to the router input before the gate
29    /// GEMV (`MoeLayer::router_input`). The Gemma-4 loader sets it to the BF16
30    /// vector `router.scale * hidden_size^-0.5`; `None` feeds the gate the
31    /// input unchanged.
32    pub router_pre_norm: Option<DenseWeight>,
33    /// 2026-09-25: Per-expert F32 bias `[num_experts]` added to the router
34    /// scores for top-k selection.
35    pub correction_bias: Option<DenseWeight>,
36}
37
38impl MoeWeights {
39    /// 2026-09-25: `num_experts` experts with NULL pointers, for tests.
40    #[cfg(test)]
41    pub fn empty(num_experts: usize) -> Self {
42        let null_dense = DenseWeight {
43            weight: DevicePtr::NULL,
44        };
45        let null_quant = QuantizedWeight {
46            weight: DevicePtr::NULL,
47            weight_scale: DevicePtr::NULL,
48            weight_scale_2: 1.0,
49            input_scale: DevicePtr::NULL,
50            weight_scale_2_vec: DevicePtr::NULL,
51            act: metrale_config::Nvfp4Act::Unstamped,
52        };
53        let null_expert = ExpertWeight {
54            gate_proj: null_quant,
55            up_proj: null_quant,
56            down_proj: null_quant,
57        };
58        Self {
59            gate: null_dense,
60            shared_expert: null_expert,
61            shared_expert_gate: null_dense,
62            experts: vec![null_expert; num_experts],
63            router_pre_norm: None,
64            correction_bias: None,
65        }
66    }
67}
68
69/// 2026-09-25: All weights of one attention or linear-attention layer.
70pub enum LayerWeights {
71    FullAttention {
72        input_norm: DenseWeight,
73        attn: AttentionWeights,
74        post_attn_norm: DenseWeight,
75        moe: MoeWeights,
76    },
77    LinearAttention {
78        input_norm: DenseWeight,
79        ssm: SsmWeights,
80        post_attn_norm: DenseWeight,
81        moe: MoeWeights,
82    },
83}
84
85/// 2026-09-25: BF16 gate, up and down projections of one expert.
86#[derive(Debug, Clone, Copy)]
87pub struct DenseExpertWeight {
88    pub gate_proj: DenseWeight,
89    pub up_proj: DenseWeight,
90    pub down_proj: DenseWeight,
91}
92
93/// 2026-09-25: MTP (multi-token prediction) head weights: one decoder layer
94/// plus the concat projection `fc`.
95///
96/// Every `DenseWeight` here is BF16 (`load_mtp` dequantizes FP8 or NVFP4
97/// projections). `MtpHead::new` quantizes the projections when it builds the head.
98pub struct MtpWeights {
99    /// 2026-09-25: RMSNorm weight for the token embedding before the concat.
100    pub pre_fc_norm_embedding: DenseWeight,
101    /// 2026-09-25: RMSNorm weight for the target hidden state before the concat.
102    pub pre_fc_norm_hidden: DenseWeight,
103    /// 2026-09-25: Concat projection `[hidden_size, 2*hidden_size]` BF16.
104    pub fc: DenseWeight,
105    /// 2026-09-25: RMSNorm weight before attention.
106    pub input_layernorm: DenseWeight,
107    /// 2026-09-25: Attention projections, BF16.
108    pub q_proj: DenseWeight,
109    pub k_proj: DenseWeight,
110    pub v_proj: DenseWeight,
111    pub o_proj: DenseWeight,
112    pub q_norm: DenseWeight,
113    pub k_norm: DenseWeight,
114    /// 2026-09-25: RMSNorm weight after attention (fused with the residual add).
115    pub post_attn_layernorm: DenseWeight,
116    /// 2026-09-25: MoE router gate, BF16. NULL when `dense_ffn` is `Some`.
117    pub moe_gate: DenseWeight,
118    /// 2026-09-25: Shared expert, BF16. NULL fields when `dense_ffn` is `Some`.
119    pub shared_expert: DenseExpertWeight,
120    /// 2026-09-25: Shared-expert gate `[1, hidden_size]` BF16. NULL when `dense_ffn` is `Some`.
121    pub shared_expert_gate: DenseWeight,
122    /// 2026-09-25: One BF16 entry per expert. Empty when `dense_ffn` is `Some`.
123    pub experts: Vec<DenseExpertWeight>,
124    /// 2026-09-25: Dense FFN (`gate_proj`, `up_proj`, `down_proj`) of an MTP
125    /// head without a router. When `Some`, `MtpHead::new` builds no MoE and
126    /// the MoE fields above are unused.
127    pub dense_ffn: Option<DenseExpertWeight>,
128    /// 2026-09-25: Final RMSNorm weight.
129    pub norm: DenseWeight,
130    /// 2026-09-25: The routed and shared experts as the checkpoint's own FP8
131    /// block-scaled tables (`mtp.layers.0.mlp.experts.{e}.*_proj.{weight,weight_scale_inv}`).
132    /// When `Some`, `experts` and `shared_expert` above are fresh BF16 dequants
133    /// of the same tensors; a head that builds its MoE from these tables calls
134    /// [`Self::release_bf16_expert_dequants`]. `None` when the experts are not
135    /// FP8 on disk, are stacked, or the variant is `Bf16Raw`.
136    pub fp8_experts: Option<MtpFp8Experts>,
137}
138
139/// 2026-09-25: Native FP8 block-scaled MTP MoE tables: `experts[e]` is routed
140/// expert `e` (all `num_experts` are loaded), plus the shared expert. The FP8
141/// bytes belong to the weight store; only the FP32 block scales are owned
142/// allocations.
143pub struct MtpFp8Experts {
144    pub experts: Vec<Fp8ExpertWeight>,
145    pub shared_expert: Fp8ExpertWeight,
146}
147
148impl MtpWeights {
149    /// 2026-09-25: Free the BF16 dequants of the routed and shared experts once
150    /// an FP8 `MoeLayer` serves them. Returns an error when
151    /// [`Self::fp8_experts`] is `None`. `load_mtp` attaches the tables only for
152    /// unstacked FP8 experts and a variant other than `Bf16Raw`; every expert
153    /// entry is then a fresh dequant allocation, never a store pointer. Leaves
154    /// `experts` empty and `shared_expert` NULL.
155    pub fn release_bf16_expert_dequants(&mut self, gpu: &dyn GpuBackend) -> Result<()> {
156        ensure!(
157            self.fp8_experts.is_some(),
158            "release_bf16_expert_dequants: no FP8 tables — the BF16 experts are the only copy"
159        );
160        let mut released = 0usize;
161        for de in self.experts.drain(..) {
162            for w in [de.gate_proj, de.up_proj, de.down_proj] {
163                gpu.free(w.weight)?;
164                released += 1;
165            }
166        }
167        let null = DenseWeight {
168            weight: DevicePtr::NULL,
169        };
170        let shared = std::mem::replace(
171            &mut self.shared_expert,
172            DenseExpertWeight {
173                gate_proj: null,
174                up_proj: null,
175                down_proj: null,
176            },
177        );
178        for w in [shared.gate_proj, shared.up_proj, shared.down_proj] {
179            if !w.weight.is_null() {
180                gpu.free(w.weight)?;
181                released += 1;
182            }
183        }
184        tracing::info!(
185            "MTP head: released {released} BF16 expert dequants (native FP8 tables serve the MoE)"
186        );
187        Ok(())
188    }
189}
metrale_model_layers/weight_map/moe.rs

metrale_model_layers/weight_map/
moe.rs
