metrale_model_layers/weight_map/
expert.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-layer weight structs (attention, SSM, MoE experts) and
4//! `QuantWeight`, the one enum layer dispatch matches over NVFP4, FP8, dense
5//! BF16 and packed Q2 weights.
6//!
7//! Owner: model-layers (weight loading).
8//! Invariants: none beyond the types.
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
18/// 2026-09-25: FP8 expert weight: gate/up/down projections as `Fp8Weight`s.
19#[derive(Debug, Clone, Copy)]
20pub struct Fp8ExpertWeight {
21    pub gate_proj: Fp8Weight,
22    pub up_proj: Fp8Weight,
23    pub down_proj: Fp8Weight,
24}
25
26/// 2026-09-25: Full-attention layer weights.
27#[derive(Debug, Clone, Copy)]
28pub struct AttentionWeights {
29    pub q_proj: DenseWeight,
30    pub k_proj: DenseWeight,
31    pub v_proj: DenseWeight,
32    pub o_proj: QuantizedWeight,
33    /// 2026-09-25: Per-head Q RMSNorm weight, `[head_dim]`.
34    pub q_norm: DenseWeight,
35    /// 2026-09-25: Per-head K RMSNorm weight, `[head_dim]`.
36    pub k_norm: DenseWeight,
37    /// 2026-09-25: Full-width Q RMSNorm weight, `[num_heads * head_dim]`: one
38    /// norm over the whole Q projection before RoPE. The MiniMax loader sets it;
39    /// when `Some`, the attention forward applies it instead of the per-head
40    /// `q_norm`.
41    pub q_norm_full: Option<DenseWeight>,
42    /// 2026-09-25: Full-width K RMSNorm weight, `[num_kv_heads * head_dim]`.
43    pub k_norm_full: Option<DenseWeight>,
44    /// 2026-09-25: K scale for the FP8 KV cache, used when no FP8 calibration is
45    /// loaded.
46    pub k_scale: f32,
47    /// 2026-09-25: V scale for the FP8 KV cache, as `k_scale`.
48    pub v_scale: f32,
49}
50
51/// 2026-09-25: Linear-attention (Gated DeltaNet) layer weights.
52#[derive(Debug, Clone, Copy)]
53pub struct SsmWeights {
54    pub in_proj_qkvz: DenseWeight,
55    pub in_proj_ba: DenseWeight,
56    pub conv1d: DenseWeight,
57    pub a_log: DenseWeight,
58    pub dt_bias: DenseWeight,
59    pub norm: DenseWeight,
60    pub out_proj: QuantizedWeight,
61}
62
63/// 2026-09-25: One MoE expert's NVFP4 gate/up/down projections.
64#[derive(Debug, Clone, Copy)]
65pub struct ExpertWeight {
66    pub gate_proj: QuantizedWeight,
67    pub up_proj: QuantizedWeight,
68    pub down_proj: QuantizedWeight,
69}
70
71impl ExpertWeight {
72    /// 2026-09-25: Null expert, every pointer NULL. Loaders use it for experts
73    /// this rank does not own under expert parallelism, and the MoE kernels
74    /// skip a NULL expert.
75    pub fn null() -> Self {
76        Self {
77            gate_proj: QuantizedWeight::null(),
78            up_proj: QuantizedWeight::null(),
79            down_proj: QuantizedWeight::null(),
80        }
81    }
82}
83
84/// 2026-09-25: A weight in any supported format. `ops::quant_gemv` /
85/// `ops::quant_gemm` match on it to pick the kernel.
86#[derive(Debug, Clone, Copy)]
87pub enum QuantWeight {
88    /// 2026-09-25: NVFP4 E2M1: packed nibbles + FP8 group scales + f32 global
89    /// scale. Kernel: `w4a16_gemv` / `w4a16_gemm`.
90    Nvfp4(QuantizedWeight),
91
92    /// 2026-09-25: FP8 E4M3 weights and their scale buffer. Kernel: `w8a16_gemv` /
93    /// `w8a16_gemm`.
94    Fp8(Fp8Weight),
95
96    /// 2026-09-25: BF16 dense (unquantized). Kernel: `dense_gemv` / `dense_gemm`.
97    Dense(DenseWeight),
98
99    /// 2026-09-25: Keep-packed Q2_0: raw `block_q2_0` bytes kept 2-bit on the
100    /// device. `quant_gemv` / `quant_gemm` refuse it: decode runs
101    /// `q2_0_gemv_vec` and prefill dequantizes to BF16 first.
102    PackedQ2(PackedQ2Weight),
103}
104
105impl QuantWeight {
106    /// 2026-09-25: An NVFP4 weight with NULL pointers.
107    pub fn null() -> Self {
108        Self::Nvfp4(QuantizedWeight::null())
109    }
110
111    /// 2026-09-25: Whether this weight's data pointer is NULL.
112    pub fn is_null(&self) -> bool {
113        match self {
114            Self::Nvfp4(w) => w.is_null(),
115            Self::Fp8(w) => w.weight.is_null(),
116            Self::Dense(w) => w.weight.is_null(),
117            Self::PackedQ2(w) => w.is_null(),
118        }
119    }
120
121    /// 2026-09-25: The keep-packed Q2_0 weight, if this is that variant.
122    pub fn as_packed_q2(&self) -> Option<&PackedQ2Weight> {
123        match self {
124            Self::PackedQ2(w) => Some(w),
125            _ => None,
126        }
127    }
128
129    /// 2026-09-25: The NVFP4 weight, if this is that variant.
130    pub fn as_nvfp4(&self) -> Option<&QuantizedWeight> {
131        match self {
132            Self::Nvfp4(w) => Some(w),
133            _ => None,
134        }
135    }
136
137    /// 2026-09-25: The FP8 weight, if this is that variant.
138    pub fn as_fp8(&self) -> Option<&Fp8Weight> {
139        match self {
140            Self::Fp8(w) => Some(w),
141            _ => None,
142        }
143    }
144
145    /// 2026-09-25: The dense weight, if this is that variant.
146    pub fn as_dense(&self) -> Option<&DenseWeight> {
147        match self {
148            Self::Dense(w) => Some(w),
149            _ => None,
150        }
151    }
152}
153
154impl From<QuantizedWeight> for QuantWeight {
155    fn from(w: QuantizedWeight) -> Self {
156        Self::Nvfp4(w)
157    }
158}
159
160impl From<Fp8Weight> for QuantWeight {
161    fn from(w: Fp8Weight) -> Self {
162        Self::Fp8(w)
163    }
164}
165
166impl From<DenseWeight> for QuantWeight {
167    fn from(w: DenseWeight) -> Self {
168        Self::Dense(w)
169    }
170}
171
172impl From<PackedQ2Weight> for QuantWeight {
173    fn from(w: PackedQ2Weight) -> Self {
174        Self::PackedQ2(w)
175    }
176}
177
178/// 2026-09-25: Per-expert weights in any supported format. `ExpertWeight`
179/// (NVFP4) and `Fp8ExpertWeight` convert into it.
180#[derive(Debug, Clone, Copy)]
181pub struct QuantExpertWeight {
182    pub gate_proj: QuantWeight,
183    pub up_proj: QuantWeight,
184    pub down_proj: QuantWeight,
185}
186
187impl QuantExpertWeight {
188    pub fn null() -> Self {
189        Self {
190            gate_proj: QuantWeight::null(),
191            up_proj: QuantWeight::null(),
192            down_proj: QuantWeight::null(),
193        }
194    }
195}
196
197impl From<ExpertWeight> for QuantExpertWeight {
198    fn from(w: ExpertWeight) -> Self {
199        Self {
200            gate_proj: QuantWeight::Nvfp4(w.gate_proj),
201            up_proj: QuantWeight::Nvfp4(w.up_proj),
202            down_proj: QuantWeight::Nvfp4(w.down_proj),
203        }
204    }
205}
206
207impl From<Fp8ExpertWeight> for QuantExpertWeight {
208    fn from(w: Fp8ExpertWeight) -> Self {
209        Self {
210            gate_proj: QuantWeight::Fp8(w.gate_proj),
211            up_proj: QuantWeight::Fp8(w.up_proj),
212            down_proj: QuantWeight::Fp8(w.down_proj),
213        }
214    }
215}
metrale_model_layers/weight_map/expert.rs

metrale_model_layers/weight_map/
expert.rs
