metrale_model_layers/weight_map/
loaders_mtp.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The stacked MTP expert slicer and the `Nvfp4Variant` on-disk format enum.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - The stacked expert pointers alias store tensors; no loader frees them.
8
9#![allow(unused_imports)]
10
11use anyhow::{Context, Result, bail, ensure};
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13use metrale_model_weights::weights::{WeightDtype, WeightStore};
14
15use super::*;
16
17/// 2026-09-25: Slice a stacked MTP expert layout into per-expert
18/// `DenseExpertWeight`s by pointer offset, without copying.
19///
20/// Expects two BF16 tensors with `E == num_experts`:
21///   `{mlp}.experts.gate_up_proj` `[E, 2*I, H]`: the first `I` rows of axis 1
22///   are gate, the next `I` rows up;
23///   `{mlp}.experts.down_proj` `[E, H, I]`.
24///
25/// The returned pointers alias the store's tensors and must not be passed
26/// to `gpu.free()`.
27pub(super) fn load_mtp_experts_stacked(
28    store: &WeightStore,
29    mlp: &str,
30    num_experts: usize,
31) -> Result<Vec<DenseExpertWeight>> {
32    let gate_up = store.get(&format!("{mlp}.experts.gate_up_proj"))?;
33    let down = store.get(&format!("{mlp}.experts.down_proj"))?;
34
35    ensure!(
36        gate_up.shape.len() == 3,
37        "MTP stacked experts.gate_up_proj: expected 3D [E,2I,H], got {:?}",
38        gate_up.shape
39    );
40    ensure!(
41        down.shape.len() == 3,
42        "MTP stacked experts.down_proj: expected 3D [E,H,I], got {:?}",
43        down.shape
44    );
45    ensure!(
46        gate_up.shape[0] == num_experts,
47        "MTP stacked experts.gate_up_proj: expert dim {} != num_experts {num_experts}",
48        gate_up.shape[0]
49    );
50    ensure!(
51        down.shape[0] == num_experts,
52        "MTP stacked experts.down_proj: expert dim {} != num_experts {num_experts}",
53        down.shape[0]
54    );
55
56    let two_inter = gate_up.shape[1];
57    let hidden = gate_up.shape[2];
58    ensure!(
59        two_inter % 2 == 0,
60        "MTP stacked experts.gate_up_proj: 2nd dim must be even (gate+up fused), got {two_inter}"
61    );
62    let intermediate = two_inter / 2;
63
64    ensure!(
65        down.shape[1] == hidden,
66        "MTP stacked: gate_up_proj.hidden ({hidden}) != down_proj.hidden ({})",
67        down.shape[1]
68    );
69    ensure!(
70        down.shape[2] == intermediate,
71        "MTP stacked: down_proj.intermediate ({}) != gate_up_proj/2 ({intermediate})",
72        down.shape[2]
73    );
74
75    // 2026-09-25: BF16 only: these tensors are handed out without conversion.
76    ensure!(
77        matches!(gate_up.dtype, WeightDtype::BF16),
78        "MTP stacked experts.gate_up_proj: expected BF16, got {:?}",
79        gate_up.dtype
80    );
81    ensure!(
82        matches!(down.dtype, WeightDtype::BF16),
83        "MTP stacked experts.down_proj: expected BF16, got {:?}",
84        down.dtype
85    );
86
87    let elt = WeightDtype::BF16.byte_size();
88    let half_bytes = intermediate * hidden * elt;
89    let gate_up_stride = two_inter * hidden * elt;
90    let down_stride = hidden * intermediate * elt;
91
92    let mut experts = Vec::with_capacity(num_experts);
93    for e in 0..num_experts {
94        let base_gu = gate_up.ptr.offset(e * gate_up_stride);
95        experts.push(DenseExpertWeight {
96            gate_proj: DenseWeight { weight: base_gu },
97            up_proj: DenseWeight {
98                weight: base_gu.offset(half_bytes),
99            },
100            down_proj: DenseWeight {
101                weight: down.ptr.offset(e * down_stride),
102            },
103        });
104    }
105    Ok(experts)
106}
107
108/// 2026-09-25: On-disk weight format, chosen by `detect_nvfp4_variant`.
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub enum Nvfp4Variant {
111    /// 2026-09-25: ModelOpt NVFP4: `weight`, `weight_scale`, `weight_scale_2`,
112    /// optional `input_scale` (read by `quantized`).
113    Standard,
114    /// 2026-09-25: compressed-tensors NVFP4: `weight_packed`, `weight_scale`,
115    /// `weight_global_scale`, optional `input_global_scale` (read by `quantized_v2`).
116    CompressedTensors,
117    /// 2026-09-25: Block-scaled FP8 E4M3 (`weight` plus its block scale).
118    /// `quantized_any` turns such a weight into NVFP4 through
119    /// a BF16 dequant (`quantized_from_fp8`); loaders that serve FP8 natively
120    /// read it with `load_fp8_block_scaled_as_fp8weight` instead.
121    Fp8Dequanted,
122    /// 2026-09-25: No usable NVFP4 or FP8 metadata: the BF16 `.weight` is
123    /// quantized to NVFP4 at load, and detection logs a warning.
124    Bf16Raw,
125}
metrale_model_layers/weight_map/loaders_mtp.rs

metrale_model_layers/weight_map/
loaders_mtp.rs
