metrale_model_layers/weight_map/
loaders_moe.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: MoE and MTP-head weight loaders: the `mlp` MoE block, the Mistral `w1/w2/w3` MoE block, and the MTP head.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - `MtpWeights::fp8_experts` is `Some` only when the experts are unstacked,
8//!   the variant is not `Bf16Raw`, and expert 0 is FP8 with `weight_scale_inv`.
9
10#![allow(unused_imports)]
11
12use anyhow::{Context, Result, bail, ensure};
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14use metrale_model_weights::weights::{WeightDtype, WeightStore};
15
16use super::*;
17
18pub(super) fn load_moe_inner(
19    store: &WeightStore,
20    layer_prefix: &str,
21    num_experts: usize,
22    gpu: &dyn GpuBackend,
23    config: &metrale_config::ModelConfig,
24    variant: Nvfp4Variant,
25    qctx: QuantizeCtx,
26    skip_routed_experts: bool,
27) -> Result<MoeWeights> {
28    let p = format!("{layer_prefix}.mlp");
29    let inter = config.moe_intermediate_size;
30    let h = config.hidden_size;
31
32    let gate = dense(store, &format!("{p}.gate.weight"))?;
33    let shared_expert_gate = dense(store, &format!("{p}.shared_expert_gate.weight"))?;
34
35    // 2026-09-25: The shared expert is loaded even when `skip_routed_experts` is set.
36    let shared_expert = ExpertWeight {
37        gate_proj: quantized_any(
38            store,
39            &format!("{p}.shared_expert.gate_proj"),
40            inter,
41            h,
42            gpu,
43            variant,
44            qctx,
45        )?,
46        up_proj: quantized_any(
47            store,
48            &format!("{p}.shared_expert.up_proj"),
49            inter,
50            h,
51            gpu,
52            variant,
53            qctx,
54        )?,
55        down_proj: quantized_any(
56            store,
57            &format!("{p}.shared_expert.down_proj"),
58            h,
59            inter,
60            gpu,
61            variant,
62            qctx,
63        )?,
64    };
65
66    let mut experts = Vec::with_capacity(num_experts);
67    for e in 0..num_experts {
68        if skip_routed_experts || !config.is_local_expert(e) {
69            experts.push(ExpertWeight::null());
70        } else {
71            experts.push(ExpertWeight {
72                gate_proj: quantized_any(
73                    store,
74                    &format!("{p}.experts.{e}.gate_proj"),
75                    inter,
76                    h,
77                    gpu,
78                    variant,
79                    qctx,
80                )?,
81                up_proj: quantized_any(
82                    store,
83                    &format!("{p}.experts.{e}.up_proj"),
84                    inter,
85                    h,
86                    gpu,
87                    variant,
88                    qctx,
89                )?,
90                down_proj: quantized_any(
91                    store,
92                    &format!("{p}.experts.{e}.down_proj"),
93                    h,
94                    inter,
95                    gpu,
96                    variant,
97                    qctx,
98                )?,
99            });
100        }
101    }
102
103    Ok(MoeWeights {
104        gate,
105        shared_expert,
106        shared_expert_gate,
107        experts,
108        router_pre_norm: None,
109        correction_bias: None,
110    })
111}
112
113/// 2026-09-25: Load a Mistral MoE block (`w1/w2/w3` naming: w1 = gate_proj,
114/// w2 = down_proj, w3 = up_proj) through `quantized_v2`.
115pub fn load_moe_mistral(
116    store: &WeightStore,
117    layer: usize,
118    num_experts: usize,
119    gpu: &dyn GpuBackend,
120    config: &metrale_config::ModelConfig,
121) -> Result<MoeWeights> {
122    // 2026-09-25: Keys are `layers.{i}.…`; only the gate also accepts `model.layers.{i}.gate.weight`.
123    let p = format!("layers.{layer}");
124
125    let gate = dense(store, &format!("{p}.gate.weight"))
126        .or_else(|_| dense(store, &format!("model.layers.{layer}.gate.weight")))
127        .context("Mistral: MoE gate weight not found")?;
128
129    // 2026-09-25: The shared expert is loaded only when `shared_experts.w1.weight_packed`
130    // exists; `shared_expert_gate` is always NULL.
131    let se_prefix = format!("{p}.shared_experts");
132    let shared_expert = if store.contains(&format!("{se_prefix}.w1.weight_packed")) {
133        ExpertWeight {
134            gate_proj: quantized_v2(store, &format!("{se_prefix}.w1"), gpu)
135                .context("Mistral: shared expert w1")?,
136            up_proj: quantized_v2(store, &format!("{se_prefix}.w3"), gpu)
137                .context("Mistral: shared expert w3")?,
138            down_proj: quantized_v2(store, &format!("{se_prefix}.w2"), gpu)
139                .context("Mistral: shared expert w2")?,
140        }
141    } else {
142        tracing::warn!("L{layer}: no shared expert weights, using NULL");
143        ExpertWeight::null()
144    };
145    let shared_expert_gate = DenseWeight {
146        weight: metrale_gpu_runtime::gpu::DevicePtr::NULL,
147    };
148
149    let mut experts = Vec::with_capacity(num_experts);
150    for e in 0..num_experts {
151        if config.is_local_expert(e) {
152            let ep = format!("{p}.experts.{e}");
153            let gate_proj = quantized_v2(store, &format!("{ep}.w1"), gpu)
154                .with_context(|| format!("Mistral expert {e}: w1 not found"))?;
155            let up_proj = quantized_v2(store, &format!("{ep}.w3"), gpu)
156                .with_context(|| format!("Mistral expert {e}: w3 not found"))?;
157            let down_proj = quantized_v2(store, &format!("{ep}.w2"), gpu)
158                .with_context(|| format!("Mistral expert {e}: w2 not found"))?;
159            experts.push(ExpertWeight {
160                gate_proj,
161                up_proj,
162                down_proj,
163            });
164        } else {
165            experts.push(ExpertWeight::null());
166        }
167    }
168
169    Ok(MoeWeights {
170        gate,
171        shared_expert,
172        shared_expert_gate,
173        experts,
174        router_pre_norm: None,
175        correction_bias: None,
176    })
177}
178
179/// 2026-09-25: Load the MTP head (`mtp.*` keys) as BF16 `DenseWeight`s.
180///
181/// Projections go through `dense_auto` (FP8 and packed NVFP4 are dequantized),
182/// except under `Bf16Raw`, which hands out the store pointers. Two FFN shapes:
183/// - **MoE**: router `mtp.layers.0.mlp.gate.weight`, per-expert or stacked
184///   expert tensors, and a shared expert.
185/// - **Dense**: `mtp.layers.0.mlp.{gate,up,down}_proj.weight` with no
186///   `.gate.weight` router, returned in `dense_ffn`.
187pub fn load_mtp(
188    store: &WeightStore,
189    num_experts: usize,
190    gpu: &dyn GpuBackend,
191    variant: Nvfp4Variant,
192) -> Result<MtpWeights> {
193    let p = "mtp.layers.0.self_attn";
194    let mlp = "mtp.layers.0.mlp";
195
196    let load = |name: &str| -> Result<DenseWeight> {
197        match variant {
198            Nvfp4Variant::Fp8Dequanted => dense_auto(store, name, gpu),
199            Nvfp4Variant::Bf16Raw => dense(store, name),
200            // 2026-09-25: A BF16 tensor comes back as the store pointer, like
201            // `dense`; a packed NVFP4 (U8) head is dequantized to BF16.
202            _ => dense_auto(store, name, gpu),
203        }
204    };
205
206    // 2026-09-25: Return before any MoE-shaped load, so a dense head does not
207    // fail on its missing `shared_expert.*` tensors.
208    let dense_gate_proj = format!("{mlp}.gate_proj.weight");
209    let moe_router = format!("{mlp}.gate.weight");
210    if store.contains(&dense_gate_proj) && !store.contains(&moe_router) {
211        let dense_ffn = DenseExpertWeight {
212            gate_proj: load(&dense_gate_proj)?,
213            up_proj: load(&format!("{mlp}.up_proj.weight"))?,
214            down_proj: load(&format!("{mlp}.down_proj.weight"))?,
215        };
216        let null = DenseWeight {
217            weight: DevicePtr::NULL,
218        };
219        return Ok(MtpWeights {
220            pre_fc_norm_embedding: dense(store, "mtp.pre_fc_norm_embedding.weight")?,
221            pre_fc_norm_hidden: dense(store, "mtp.pre_fc_norm_hidden.weight")?,
222            fc: load("mtp.fc.weight")?,
223            input_layernorm: dense(store, "mtp.layers.0.input_layernorm.weight")?,
224            q_proj: load(&format!("{p}.q_proj.weight"))?,
225            k_proj: load(&format!("{p}.k_proj.weight"))?,
226            v_proj: load(&format!("{p}.v_proj.weight"))?,
227            o_proj: load(&format!("{p}.o_proj.weight"))?,
228            q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
229            k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
230            post_attn_layernorm: dense(store, "mtp.layers.0.post_attention_layernorm.weight")?,
231            moe_gate: null,
232            shared_expert: DenseExpertWeight {
233                gate_proj: null,
234                up_proj: null,
235                down_proj: null,
236            },
237            shared_expert_gate: null,
238            experts: Vec::new(),
239            dense_ffn: Some(dense_ffn),
240            norm: dense(store, "mtp.norm.weight")?,
241            fp8_experts: None,
242        });
243    }
244
245    let shared_expert = DenseExpertWeight {
246        gate_proj: load(&format!("{mlp}.shared_expert.gate_proj.weight"))?,
247        up_proj: load(&format!("{mlp}.shared_expert.up_proj.weight"))?,
248        down_proj: load(&format!("{mlp}.shared_expert.down_proj.weight"))?,
249    };
250
251    // 2026-09-25: Two expert layouts: per expert
252    // (`experts.{e}.{gate,up,down}_proj.weight`), or stacked
253    // (`experts.gate_up_proj` `[E, 2*I, H]` and `experts.down_proj` `[E, H, I]`,
254    // both present), which `load_mtp_experts_stacked` slices without copying.
255    let stacked_gate_up = format!("{mlp}.experts.gate_up_proj");
256    let stacked_down = format!("{mlp}.experts.down_proj");
257    let stacked = store.contains(&stacked_gate_up) && store.contains(&stacked_down);
258    let experts = if stacked {
259        load_mtp_experts_stacked(store, mlp, num_experts)?
260    } else {
261        let mut v = Vec::with_capacity(num_experts);
262        for e in 0..num_experts {
263            v.push(DenseExpertWeight {
264                gate_proj: load(&format!("{mlp}.experts.{e}.gate_proj.weight"))?,
265                up_proj: load(&format!("{mlp}.experts.{e}.up_proj.weight"))?,
266                down_proj: load(&format!("{mlp}.experts.{e}.down_proj.weight"))?,
267            });
268        }
269        v
270    };
271    // 2026-09-25: Not under `Bf16Raw`: there `load` hands out store pointers,
272    // which `release_bf16_expert_dequants` would free. Not for stacked experts,
273    // whose entries alias the stacked tensors.
274    let fp8_experts =
275        if !stacked && variant != Nvfp4Variant::Bf16Raw && mtp_experts_native_fp8(store, mlp) {
276            Some(load_mtp_fp8_experts(store, mlp, num_experts, gpu)?)
277        } else {
278            None
279        };
280
281    Ok(MtpWeights {
282        pre_fc_norm_embedding: dense(store, "mtp.pre_fc_norm_embedding.weight")?,
283        pre_fc_norm_hidden: dense(store, "mtp.pre_fc_norm_hidden.weight")?,
284        fc: load("mtp.fc.weight")?,
285        input_layernorm: dense(store, "mtp.layers.0.input_layernorm.weight")?,
286        q_proj: load(&format!("{p}.q_proj.weight"))?,
287        k_proj: load(&format!("{p}.k_proj.weight"))?,
288        v_proj: load(&format!("{p}.v_proj.weight"))?,
289        o_proj: load(&format!("{p}.o_proj.weight"))?,
290        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
291        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
292        post_attn_layernorm: dense(store, "mtp.layers.0.post_attention_layernorm.weight")?,
293        moe_gate: dense(store, &format!("{mlp}.gate.weight"))?,
294        shared_expert,
295        shared_expert_gate: dense(store, &format!("{mlp}.shared_expert_gate.weight"))?,
296        experts,
297        dense_ffn: None,
298        norm: dense(store, "mtp.norm.weight")?,
299        fp8_experts,
300    })
301}
302
303/// 2026-09-25: True when expert 0's `gate_proj.weight` is FP8 E4M3 with a
304/// `weight_scale_inv` sibling. Only expert 0 is probed; an expert in another
305/// format then fails the FP8 dtype check in `load_fp8_block_scaled_as_fp8weight`.
306pub(crate) fn mtp_experts_native_fp8(store: &WeightStore, mlp: &str) -> bool {
307    let weight = format!("{mlp}.experts.0.gate_proj.weight");
308    store
309        .get(&weight)
310        .map(|w| w.dtype == WeightDtype::FP8E4M3)
311        .unwrap_or(false)
312        && store.contains(&format!("{mlp}.experts.0.gate_proj.weight_scale_inv"))
313}
314
315/// 2026-09-25: The MTP routed and shared experts as FP8 block-scaled tables.
316/// The FP8 bytes are the store's; the FP32 block scales are new allocations.
317fn load_mtp_fp8_experts(
318    store: &WeightStore,
319    mlp: &str,
320    num_experts: usize,
321    gpu: &dyn GpuBackend,
322) -> Result<MtpFp8Experts> {
323    let proj = |prefix: &str, name: &str| -> Result<Fp8Weight> {
324        load_fp8_block_scaled_as_fp8weight(store, &format!("{prefix}.{name}"), gpu)
325    };
326    let load = |prefix: &str| -> Result<Fp8ExpertWeight> {
327        Ok(Fp8ExpertWeight {
328            gate_proj: proj(prefix, "gate_proj")?,
329            up_proj: proj(prefix, "up_proj")?,
330            down_proj: proj(prefix, "down_proj")?,
331        })
332    };
333    let mut experts = Vec::with_capacity(num_experts);
334    for e in 0..num_experts {
335        experts.push(
336            load(&format!("{mlp}.experts.{e}"))
337                .with_context(|| format!("MTP expert {e}: native FP8 tables"))?,
338        );
339    }
340    let shared_expert =
341        load(&format!("{mlp}.shared_expert")).context("MTP shared expert: native FP8 tables")?;
342    Ok(MtpFp8Experts {
343        experts,
344        shared_expert,
345    })
346}
347
348#[cfg(test)]
349mod mtp_fp8_detect_tests {
350    use super::*;
351    use metrale_model_weights::weights::WeightTensor;
352    use std::collections::HashMap;
353
354    const MLP: &str = "mtp.layers.0.mlp";
355
356    fn store(entries: Vec<(String, WeightDtype)>) -> WeightStore {
357        let mut m = HashMap::new();
358        for (name, dtype) in entries {
359            m.insert(
360                name,
361                WeightTensor {
362                    ptr: DevicePtr::NULL,
363                    shape: vec![512, 2048],
364                    dtype,
365                },
366            );
367        }
368        WeightStore::from_map(m)
369    }
370
371    fn gate0(suffix: &str) -> String {
372        format!("{MLP}.experts.0.gate_proj.{suffix}")
373    }
374
375    #[test]
376    fn fp8_weight_with_block_scale_is_native_fp8() {
377        let s = store(vec![
378            (gate0("weight"), WeightDtype::FP8E4M3),
379            (gate0("weight_scale_inv"), WeightDtype::BF16),
380        ]);
381        assert!(mtp_experts_native_fp8(&s, MLP));
382    }
383
384    #[test]
385    fn bf16_experts_are_not_native_fp8() {
386        let s = store(vec![(gate0("weight"), WeightDtype::BF16)]);
387        assert!(!mtp_experts_native_fp8(&s, MLP));
388    }
389
390    #[test]
391    fn fp8_without_a_block_scale_is_refused() {
392        let s = store(vec![(gate0("weight"), WeightDtype::FP8E4M3)]);
393        assert!(!mtp_experts_native_fp8(&s, MLP));
394    }
395
396    #[test]
397    fn absent_experts_are_not_native_fp8() {
398        assert!(!mtp_experts_native_fp8(&store(vec![]), MLP));
399    }
400}
metrale_model_layers/weight_map/loaders_moe.rs

metrale_model_layers/weight_map/
loaders_moe.rs
