metrale_model_arch/mistral_loader/loader_impl/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `ModelWeightLoader` for `MistralWeightLoader`: every layer runs the same sequence of load steps over one `ctx::MistralLayerCtx`.
4//!
5//! The steps, in the order `load_layers_inner` runs them:
6//! - `phase_lora_qkv`: wq_a, wq_b, wkv_a, wkv_b and their norms, NVFP4 copies, TP shard.
7//! - `phase_per_head`: W_UK_T (per-head transpose), W_UV and wq_b_rope.
8//! - `phase_qk_absorbed`: W_QK_absorbed, computed on the CPU.
9//! - `phase_block_diag`: block-diagonal W_UK and W_UV.
10//! - `phase_o_proj`: the output projection and its NVFP4 copy.
11//! - `yarn`: the YaRN inv_freq table, computed for the first layer and shared.
12//! - `phase_assemble`: `MlaWeights`, the MoE FFN and the `TransformerLayer`.
13//!
14//! Owner: model-arch (Mistral loader).
15//! Invariants:
16//! - Layers are built in index order, and the first error ends the load.
17//! - Every layer gets the same YaRN inv_freq pointer.
18
19use anyhow::{Context, Result};
20use metrale_cache::kv_cache::KvCacheDtype;
21use metrale_config::ModelConfig;
22use metrale_gpu_runtime::gpu::GpuBackend;
23use metrale_model_weights::weights::WeightStore;
24
25use super::MistralWeightLoader;
26use crate::weight_loader::ModelWeightLoader;
27use metrale_model_layers::layer::TransformerLayer;
28use metrale_model_layers::layers::VisionTower;
29use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, dense};
30
31// 2026-09-25: `ctx` and the three steps that read no tensor names are
32// `pub(crate)` because the LongCat loader runs them too. The steps that read
33// Mistral tensor names stay private.
34pub(crate) mod ctx;
35mod phase_assemble;
36pub(crate) mod phase_block_diag;
37mod phase_lora_qkv;
38mod phase_o_proj;
39pub(crate) mod phase_per_head;
40pub(crate) mod phase_qk_absorbed;
41mod yarn;
42
43impl ModelWeightLoader for MistralWeightLoader {
44    fn supports_tp(&self) -> bool {
45        // 2026-09-25: Under TP, `phase_lora_qkv` shards wq_b and wkv_b
46        // column-parallel on the head axis; wq_a, wkv_a and the two latent
47        // norms stay replicated. The per-head steps loop over
48        // `num_key_value_heads`, which the server has already divided by the
49        // TP size, so they build this rank's heads only.
50        true
51    }
52
53    fn load_layers(
54        &self,
55        store: &WeightStore,
56        config: &ModelConfig,
57        gpu: &dyn GpuBackend,
58        layer_kv_dtypes: &[KvCacheDtype],
59    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
60        self.load_layers_inner(store, config, gpu, layer_kv_dtypes)
61    }
62
63    fn load_embedding(
64        &self,
65        store: &WeightStore,
66        _config: &ModelConfig,
67        _gpu: &dyn GpuBackend,
68    ) -> Result<DenseWeight> {
69        dense(store, "tok_embeddings.weight")
70            .or_else(|_| dense(store, "model.embed_tokens.weight"))
71            .context("Mistral: embedding not found")
72    }
73
74    fn load_final_norm(
75        &self,
76        store: &WeightStore,
77        _config: &ModelConfig,
78        _gpu: &dyn GpuBackend,
79    ) -> Result<DenseWeight> {
80        dense(store, "norm.weight")
81            .or_else(|_| dense(store, "model.norm.weight"))
82            .context("Mistral: final norm not found")
83    }
84
85    fn load_lm_head(
86        &self,
87        store: &WeightStore,
88        config: &ModelConfig,
89        _gpu: &dyn GpuBackend,
90    ) -> Result<DenseWeight> {
91        if store.contains("output.weight") {
92            dense(store, "output.weight")
93        } else if store.contains("lm_head.weight") {
94            dense(store, "lm_head.weight")
95        } else if config.tie_word_embeddings {
96            dense(store, "tok_embeddings.weight")
97                .or_else(|_| dense(store, "model.embed_tokens.weight"))
98                .context("Mistral: tied embedding lm_head not found")
99        } else {
100            anyhow::bail!("Mistral: lm_head/output weight not found")
101        }
102    }
103
104    fn load_mtp_weights(
105        &self,
106        _store: &WeightStore,
107        _config: &ModelConfig,
108        _gpu: &dyn GpuBackend,
109    ) -> Result<Option<MtpWeights>> {
110        Ok(None)
111    }
112
113    fn load_vision_encoder(
114        &self,
115        _store: &WeightStore,
116        _config: &ModelConfig,
117        _gpu: &dyn GpuBackend,
118    ) -> Result<Option<VisionTower>> {
119        Ok(None)
120    }
121}
122
123impl MistralWeightLoader {
124    pub(crate) fn load_layers_inner(
125        &self,
126        store: &WeightStore,
127        config: &ModelConfig,
128        gpu: &dyn GpuBackend,
129        layer_kv_dtypes: &[KvCacheDtype],
130    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
131        let n = config.num_hidden_layers;
132        let q_lora = config.q_lora_rank;
133        let kv_lora = config.kv_lora_rank;
134        let nope = config.qk_nope_head_dim;
135        let rope = config.qk_rope_head_dim;
136        let v_dim = config.v_head_dim;
137
138        tracing::info!(
139            "Mistral MLA→GQA: expanding LoRA on GPU (q_lora={q_lora}, kv_lora={kv_lora}, \
140             nope={nope}, rope={rope}, v_dim={v_dim})"
141        );
142
143        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
144        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
145        let stream = gpu.default_stream();
146
147        let mut layers: Vec<Box<dyn TransformerLayer>> = Vec::with_capacity(n);
148        let mut yarn_inv_freq_shared = metrale_gpu_runtime::gpu::DevicePtr::NULL;
149
150        for i in 0..n {
151            let mut ctx =
152                ctx::MistralLayerCtx::new(store, config, gpu, absmax_k, quantize_k, stream, i);
153            phase_lora_qkv::load_lora_qkv(&mut ctx)?;
154            phase_per_head::build_per_head_views(&mut ctx)?;
155            phase_qk_absorbed::build_w_qk_absorbed(&mut ctx)?;
156            phase_block_diag::build_block_diagonals(&mut ctx)?;
157            phase_o_proj::load_o_proj(&mut ctx)?;
158            let yarn_inv_freq =
159                ctx::ensure_yarn_inv_freq(&mut yarn_inv_freq_shared, config, rope, gpu)?;
160            let layer = phase_assemble::assemble_layer(ctx, yarn_inv_freq, layer_kv_dtypes)?;
161            layers.push(layer);
162
163            if (i + 1) % 6 == 0 || i == n - 1 {
164                let free = gpu.free_memory().unwrap_or(0);
165                tracing::info!("L{}/{n} done — {:.1} GB free", i + 1, free as f64 / 1e9);
166            }
167        }
168        Ok(layers)
169    }
170}
metrale_model_arch/mistral_loader/loader_impl/mod.rs

metrale_model_arch/mistral_loader/loader_impl/
mod.rs
