metrale_model_arch/weight_loader/
minimax.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: MiniMax M2 weight loader (`minimax_m2`): every layer is full attention plus a
4//! routed MoE.
5//!
6//! Owner: model-arch weight loader (MiniMax).
7//! Invariants:
8//! - Q/K/V are sharded column-parallel and O row-parallel (`load_qkvo_tp`) before being
9//!   quantized to NVFP4 at load; the q/k norms are sharded 1-D (`load_qk_norms_tp`).
10//! - The per-head `q_norm`/`k_norm` slots are null; `q_norm_full`/`k_norm_full` carry the
11//!   full-width norm weights.
12//! - `load_mtp_weights_multi` returns no modules when the checkpoint has no
13//!   `model.layers.{num_hidden_layers}` tensors, and an error when it has them.
14
15use anyhow::Result;
16use metrale_cache::kv_cache::KvCacheDtype;
17use metrale_config::ModelConfig;
18use metrale_gpu_runtime::gpu::GpuBackend;
19use metrale_model_weights::weights::WeightStore;
20
21use super::ModelWeightLoader;
22use crate::tp_shard::{
23    TpShardKind, load_qk_norms_tp, load_qkvo_tp, shard_dense_1d_bf16, shard_dense_bf16,
24};
25use metrale_model_layers::layer::TransformerLayer;
26use metrale_model_layers::layers::{FfnComponent, MoeLayer, Qwen3AttentionLayer};
27use metrale_model_layers::weight_map::quant_helpers::dense_auto;
28use metrale_model_layers::weight_map::ssm_qwen35_more::load_moe_minimax;
29use metrale_model_layers::weight_map::{
30    AttentionWeights, DenseWeight, MtpWeights, QuantizedWeight, dense, detect_nvfp4_variant,
31    load_kv_scales, quantize_to_nvfp4,
32};
33
34pub struct MinimaxM2WeightLoader;
35
36impl ModelWeightLoader for MinimaxM2WeightLoader {
37    fn supports_tp(&self) -> bool {
38        // 2026-09-25: `load_layers` shards Q/K/V/O and the q/k norms by
39        // `config.tp_rank` / `config.tp_world_size`.
40        true
41    }
42
43    fn load_layers(
44        &self,
45        store: &WeightStore,
46        config: &ModelConfig,
47        gpu: &dyn GpuBackend,
48        layer_kv_dtypes: &[KvCacheDtype],
49    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
50        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
51        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
52        let stream = gpu.default_stream();
53        let h = config.hidden_size;
54        let variant = detect_nvfp4_variant(store, config);
55        tracing::info!(
56            "minimax_m2: loading {} layers, variant={:?}, hidden_size={h}",
57            config.num_hidden_layers,
58            variant,
59        );
60
61        // 2026-09-25: The MoE prefill transpose is not done here; `factory::build`
62        // runs it after load (`maybe_run_minimax_m2_moe_transpose`).
63
64        let mut layers: Vec<Box<dyn TransformerLayer>> =
65            Vec::with_capacity(config.num_hidden_layers);
66
67        // 2026-09-25: Null weight for the per-head q_norm/k_norm slots; the
68        // full-width norms go in `q_norm_full`/`k_norm_full`.
69        let dummy_norm = DenseWeight {
70            weight: metrale_gpu_runtime::gpu::DevicePtr::NULL,
71        };
72
73        for i in 0..config.num_hidden_layers {
74            let lp = format!("model.layers.{i}");
75            tracing::debug!("minimax_m2: layer {i}");
76            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
77            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
78
79            // 2026-09-25: MoE through `load_moe_minimax`, with the router gate
80            // also quantized to NVFP4.
81            let moe_weights = load_moe_minimax(
82                store,
83                &lp,
84                config.num_experts,
85                gpu,
86                config,
87                variant,
88                absmax_k,
89                quantize_k,
90                stream,
91            )?;
92            let gate_nvfp4 = quantize_to_nvfp4(
93                &moe_weights.gate,
94                config.num_experts,
95                h,
96                gpu,
97                absmax_k,
98                quantize_k,
99                stream,
100            )?;
101            let mut moe_layer = MoeLayer::new(
102                moe_weights,
103                config.num_experts,
104                Some(gate_nvfp4),
105                gpu,
106                config,
107            )?;
108            // 2026-09-25: `predequant_for_prefill` runs here; the prefill
109            // transpose runs after load (see above).
110            moe_layer.predequant_for_prefill(gpu, config, stream)?;
111            let ffn = FfnComponent::Moe(moe_layer);
112
113            // 2026-09-25: Attention: ungated Q, full-width q/k norms.
114            let p = format!("{lp}.self_attn");
115            // 2026-09-25: Each projection is read through `dense_auto` (a new
116            // BF16 buffer for an FP8 or F32 source), TP-sharded and quantized to
117            // NVFP4. The BF16 buffers and the store's source tensor are then
118            // freed, and the returned `DenseWeight` is null.
119            let tp_rank = config.tp_rank;
120            let tp_size = config.tp_world_size;
121            let load_and_quant = |name: &str,
122                                  full_n: usize,
123                                  full_k: usize,
124                                  kind: TpShardKind|
125             -> Result<(DenseWeight, QuantizedWeight)> {
126                let wkey = format!("{p}.{name}.weight");
127                let scale_key = format!("{p}.{name}.weight_scale_inv");
128                let (src_ptr, src_dtype) = {
129                    let t = store.get(&wkey)?;
130                    (t.ptr, t.dtype)
131                };
132                let src_is_fp8 = src_dtype == metrale_model_weights::weights::WeightDtype::FP8E4M3;
133                let src_is_f32 = src_dtype == metrale_model_weights::weights::WeightDtype::FP32;
134                let scale_ptr = if src_is_fp8 && store.contains(&scale_key) {
135                    Some(store.get(&scale_key)?.ptr)
136                } else {
137                    None
138                };
139                let dense_w = dense_auto(store, &wkey, gpu)?;
140                // 2026-09-25: TP-shard the BF16 weight before NVFP4 quantization.
141                // With tp_size <= 1 the helper returns the input pointer.
142                let (sharded_ptr, local_n, local_k) =
143                    shard_dense_bf16(dense_w.weight, full_n, full_k, kind, tp_rank, tp_size, gpu)?;
144                let sharded = DenseWeight {
145                    weight: sharded_ptr,
146                };
147                let q = quantize_to_nvfp4(
148                    &sharded, local_n, local_k, gpu, absmax_k, quantize_k, stream,
149                )?;
150                if sharded_ptr != dense_w.weight {
151                    gpu.free(sharded_ptr)?;
152                }
153                if src_is_fp8 {
154                    // 2026-09-25: FP8 source: free the BF16 dequant buffer, the
155                    // FP8 source and its block scale. The store keeps the stale
156                    // pointers.
157                    gpu.free(dense_w.weight)?;
158                    gpu.free(src_ptr)?;
159                    if let Some(sp) = scale_ptr {
160                        gpu.free(sp)?;
161                    }
162                } else {
163                    // 2026-09-25: BF16 or F32 source: free the F32-to-BF16
164                    // buffer when one was made, and the source.
165                    if src_is_f32 {
166                        gpu.free(dense_w.weight)?;
167                    }
168                    gpu.free(src_ptr)?;
169                }
170                Ok((
171                    DenseWeight {
172                        weight: metrale_gpu_runtime::gpu::DevicePtr::NULL,
173                    },
174                    q,
175                ))
176            };
177            // 2026-09-25: Q/K/V column-parallel and O row-parallel through
178            // `load_qkvo_tp`, which passes each projection's full dims to
179            // `load_and_quant`.
180            let [
181                (q_dense, q_nvfp4),
182                (k_dense, k_nvfp4),
183                (v_dense, v_nvfp4),
184                (_o_dense, o_nvfp4),
185            ] = load_qkvo_tp(config, |name, full_n, full_k, kind| {
186                load_and_quant(name, full_n, full_k, kind)
187            })?;
188
189            let (k_scale, v_scale) = load_kv_scales(store, &p, gpu)?;
190
191            // 2026-09-25: Full-width q/k norm weights, sharded 1-D to match the
192            // local Q/K outputs; `shard_dense_1d_bf16` returns the input pointer
193            // when tp_size <= 1.
194            let (q_norm_full, k_norm_full) = load_qk_norms_tp(config, |name, full_dim| {
195                let src = dense(store, &format!("{p}.{name}.weight"))?;
196                let (ptr, _) = shard_dense_1d_bf16(src.weight, full_dim, tp_rank, tp_size, gpu)?;
197                Ok::<_, anyhow::Error>(DenseWeight { weight: ptr })
198            })?;
199
200            let attn = AttentionWeights {
201                q_proj: q_dense,
202                k_proj: k_dense,
203                v_proj: v_dense,
204                o_proj: o_nvfp4,
205                // 2026-09-25: Per-head slots null; the full-width norms follow.
206                q_norm: dummy_norm,
207                k_norm: dummy_norm,
208                q_norm_full: Some(q_norm_full),
209                k_norm_full: Some(k_norm_full),
210                k_scale,
211                v_scale,
212            };
213
214            let mut layer = Qwen3AttentionLayer::new_ungated(
215                input_norm,
216                attn,
217                post_attn_norm,
218                ffn,
219                i,
220                Some(q_nvfp4),
221                Some(k_nvfp4),
222                Some(v_nvfp4),
223                gpu,
224                layer_kv_dtypes[i],
225                config.fp8_kv_calibration_tokens,
226                config,
227            )?;
228
229            // 2026-09-25: Transposed copies of the attention NVFP4 weights for
230            // prefill (`set_prefill_weights`).
231            let q_proj_n = config.num_attention_heads * config.head_dim;
232            let kv_proj_n = config.num_key_value_heads * config.head_dim;
233            let qt = q_nvfp4.transpose_for_gemm(gpu, q_proj_n, h)?;
234            let kt = k_nvfp4.transpose_for_gemm(gpu, kv_proj_n, h)?;
235            let vt = v_nvfp4.transpose_for_gemm(gpu, kv_proj_n, h)?;
236            let ot = layer.attn.o_proj.transpose_for_gemm(gpu, h, q_proj_n)?;
237            layer.set_prefill_weights(Some(qt), Some(kt), Some(vt), Some(ot));
238
239            layers.push(Box::new(layer));
240        }
241
242        tracing::info!("minimax_m2: built {} layers", layers.len());
243        Ok(layers)
244    }
245
246    fn load_embedding(
247        &self,
248        store: &WeightStore,
249        _config: &ModelConfig,
250        _gpu: &dyn GpuBackend,
251    ) -> Result<DenseWeight> {
252        dense(store, "model.embed_tokens.weight")
253    }
254
255    fn load_final_norm(
256        &self,
257        store: &WeightStore,
258        _config: &ModelConfig,
259        _gpu: &dyn GpuBackend,
260    ) -> Result<DenseWeight> {
261        dense(store, "model.norm.weight")
262    }
263
264    fn load_lm_head(
265        &self,
266        store: &WeightStore,
267        _config: &ModelConfig,
268        _gpu: &dyn GpuBackend,
269    ) -> Result<DenseWeight> {
270        dense(store, "lm_head.weight")
271    }
272
273    fn load_mtp_weights(
274        &self,
275        _store: &WeightStore,
276        _config: &ModelConfig,
277        _gpu: &dyn GpuBackend,
278    ) -> Result<Option<MtpWeights>> {
279        // 2026-09-25: MiniMax's MTP modules go through `load_mtp_weights_multi`.
280        Ok(None)
281    }
282
283    fn load_mtp_weights_multi(
284        &self,
285        store: &WeightStore,
286        config: &ModelConfig,
287        _gpu: &dyn GpuBackend,
288    ) -> Result<Vec<MtpWeights>> {
289        // 2026-09-25: MTP modules would start at `model.layers.{num_hidden_layers}`.
290        // Without that layer's `input_layernorm.weight` the checkpoint has none:
291        // return no modules.
292        let first_mtp_idx = config.num_hidden_layers;
293        let probe = format!("model.layers.{first_mtp_idx}.input_layernorm.weight");
294        if !store.contains(&probe) {
295            tracing::info!(
296                "minimax_m2: no MTP module weights found in checkpoint \
297                 (expected starting at layer {first_mtp_idx}); MTP disabled"
298            );
299            return Ok(Vec::new());
300        }
301
302        // 2026-09-25: Per-module extraction is not implemented, so a checkpoint
303        // that has the modules fails the load.
304        anyhow::bail!(
305            "minimax_m2: MTP module weights detected at layer {first_mtp_idx} \
306             but the MiniMax loader hasn't implemented per-module extraction \
307             yet. Run with --speculative 0 for non-MTP decode, or await \
308             MiniMax M5 phase-3 (populate load_mtp_weights_multi with the \
309             concrete Mixtral-convention weight keys)."
310        )
311    }
312}
metrale_model_arch/weight_loader/minimax.rs

metrale_model_arch/weight_loader/
minimax.rs
