metrale_model_arch/weight_loader/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The [`ModelWeightLoader`] trait, and one loader per architecture that turns a
4//! flat [`WeightStore`] into typed [`TransformerLayer`]s.
5//!
6//! Owner: model-arch weight loader.
7//! Invariants: none beyond the types.
8
9pub mod deepseek_v4;
10pub mod deepseek_v41;
11pub mod dflash_loader;
12mod gemma4;
13/// 2026-09-25: GLM-5.3 checkpoint tensor classification (`classify`).
14pub mod glm5_next;
15mod laguna;
16mod longcat;
17mod minimax;
18mod nemotron;
19mod nllb;
20mod qwen3;
21mod qwen35;
22mod qwen35_dense;
23mod qwen3_vl;
24#[cfg_attr(test, allow(unreachable_pub))]
25pub(crate) mod qwen4_exp;
26mod step3p7;
27
28pub use deepseek_v4::DeepSeekV4WeightLoader;
29pub use dflash_loader::{
30    DflashConfig, DflashLayerWeights, DflashSubConfig, DflashWeights, load_dflash_weights,
31    store_has_dflash_weights,
32};
33pub mod glm5_next_load;
34pub mod glm5_next_mtp;
35pub mod glm5_next_vision;
36pub use gemma4::Gemma4WeightLoader;
37pub use glm5_next_load::Glm5NextWeightLoader;
38pub(crate) use glm5_next_mtp::Glm5NextMtpModule;
39
40pub use laguna::LagunaWeightLoader;
41pub use longcat::LongcatWeightLoader;
42pub use minimax::MinimaxM2WeightLoader;
43pub use nemotron::NemotronHWeightLoader;
44pub use nllb::NllbWeightLoader;
45pub use qwen3::Qwen3WeightLoader;
46pub use qwen3_vl::Qwen3VLWeightLoader;
47pub use qwen4_exp::Qwen4ExpWeightLoader;
48pub use qwen35::Qwen35WeightLoader;
49pub use qwen35_dense::Qwen35DenseWeightLoader;
50/// 2026-09-25: The native-FP8 dense loader's derived-copy decision table and the
51/// shape arithmetic that prices it.
52pub use qwen35_dense::fp8_residency;
53/// 2026-09-25: That table evaluated from the config before the checkpoint
54/// loads; preflight's headroom check reads it.
55pub use qwen35_dense::predicted_residency;
56pub use step3p7::Step3p7WeightLoader;
57
58use anyhow::Result;
59use metrale_cache::kv_cache::KvCacheDtype;
60use metrale_config::ModelConfig;
61use metrale_gpu_runtime::gpu::GpuBackend;
62use metrale_model_weights::weights::{DeferHook, WeightStore};
63
64use metrale_model_layers::layer::TransformerLayer;
65use metrale_model_layers::layers::VisionTower;
66use metrale_model_layers::weight_map::{
67    DenseWeight, MtpWeights, Nvfp4Variant, detect_nvfp4_variant,
68};
69
70/// 2026-09-25: Can this box hold the transposed MoE prefill copies for every
71/// layer, and does the operator want them? Without them MoE prefill uses the
72/// fallback grouped GEMM. The qwen3, qwen3_vl and gemma4 loaders ask here.
73///
74/// `METRALE_MOE_PREFILL_COPIES=0` forces the fallback. Any other value, or
75/// unset, builds them when their NVFP4 bytes for all `num_hidden_layers` fit in
76/// the measured free memory less 2 GiB; if free memory cannot be read, they are
77/// not built.
78pub(crate) fn moe_prefill_copies_fit(config: &ModelConfig, gpu: &dyn GpuBackend) -> bool {
79    if std::env::var("METRALE_MOE_PREFILL_COPIES").ok().as_deref() == Some("0") {
80        tracing::info!("METRALE_MOE_PREFILL_COPIES=0: MoE prefill uses the fallback grouped GEMM");
81        return false;
82    }
83    let inter = config.moe_intermediate_size;
84    let h = config.hidden_size;
85    // 2026-09-25: NVFP4: one scale byte per 16 elements beside the packed
86    // pairs (half a byte per element).
87    let group_size = 16usize;
88    let gu_bytes = inter * h / 2 + inter * h / group_size;
89    let d_bytes = h * inter / 2 + h * inter / group_size;
90    let per_layer = config.num_experts * (2 * gu_bytes + d_bytes);
91    let total = per_layer * config.num_hidden_layers;
92    let available = gpu.free_memory().unwrap_or(0);
93    let headroom = 2 * 1024 * 1024 * 1024;
94    let fits = total <= available.saturating_sub(headroom);
95    if !fits {
96        tracing::warn!(
97            "Skipping MoE weight transposition ({:.1} GB needed, {:.1} GB available). \
98             Prefill will use fallback grouped GEMM.",
99            total as f64 / (1024.0 * 1024.0 * 1024.0),
100            available as f64 / (1024.0 * 1024.0 * 1024.0),
101        );
102    }
103    fits
104}
105
106/// 2026-09-25: Runtime quantization format of a loader's weights, which the
107/// qwen3 and qwen35 loaders resolve from the checkpoint's NVFP4 variant.
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum QuantFormat {
110    /// 2026-09-25: NVFP4 E2M1; every variant other than `Fp8Dequanted`.
111    Nvfp4,
112    /// 2026-09-25: FP8 E4M3 block-scaled, served natively (`Fp8Dequanted`).
113    Fp8,
114}
115
116impl QuantFormat {
117    /// 2026-09-25: Peak GPU memory as a multiple of on-disk weight bytes, for an
118    /// OOM pre-flight estimate.
119    pub fn peak_memory_multiplier(&self) -> f64 {
120        match self {
121            Self::Nvfp4 => 1.15,
122            Self::Fp8 => 1.5,
123        }
124    }
125}
126
127/// 2026-09-25: On-disk weight format, derived from `detect_nvfp4_variant`.
128#[derive(Debug, Clone, Copy, PartialEq, Eq)]
129pub enum WeightFormat {
130    /// 2026-09-25: ModelOpt (`Standard`) or compressed-tensors NVFP4 on disk,
131    /// and `Bf16Raw` checkpoints, which the loader quantises to NVFP4.
132    Nvfp4,
133    /// 2026-09-25: Block-scaled FP8 E4M3 on disk (`Fp8Dequanted`).
134    Fp8BlockScaled,
135    /// 2026-09-25: BF16 on disk. `WeightFormat::detect` never returns it.
136    Bf16Dense,
137}
138
139impl WeightFormat {
140    /// 2026-09-25: Detect the weight format from a [`WeightStore`] by probing
141    /// tensor names.
142    pub fn detect(store: &WeightStore, config: &ModelConfig) -> Self {
143        match detect_nvfp4_variant(store, config) {
144            Nvfp4Variant::Fp8Dequanted => Self::Fp8BlockScaled,
145            Nvfp4Variant::CompressedTensors | Nvfp4Variant::Standard => Self::Nvfp4,
146            // 2026-09-25: A `Bf16Raw` weight is quantized to NVFP4 at load.
147            Nvfp4Variant::Bf16Raw => Self::Nvfp4,
148        }
149    }
150
151    pub fn is_fp8(&self) -> bool {
152        matches!(self, Self::Fp8BlockScaled)
153    }
154}
155
156/// 2026-09-25: Loads weights from a [`WeightStore`] into typed layer objects.
157pub trait ModelWeightLoader {
158    /// 2026-09-25: Whether this loader's weight slicing honours
159    /// `config.tp_world_size` / `config.tp_rank`. There is no default, so every
160    /// loader declares it.
161    ///
162    /// With `--tp-size > 1`, startup (`serve_phases/topology.rs`) refuses a
163    /// loader that returns `false`. `weight_loader/minimax.rs` is a loader
164    /// that returns `true`.
165    fn supports_tp(&self) -> bool;
166
167    /// 2026-09-25: Load all transformer layers from the weight store.
168    ///
169    /// `layer_kv_dtypes` holds one KV cache dtype per attention layer, indexed
170    /// by the layer's 0-based ordinal among the attention layers.
171    fn load_layers(
172        &self,
173        store: &WeightStore,
174        config: &ModelConfig,
175        gpu: &dyn GpuBackend,
176        layer_kv_dtypes: &[KvCacheDtype],
177    ) -> Result<Vec<Box<dyn TransformerLayer>>>;
178
179    /// 2026-09-25: Drop store tensors this loader has finished with. `factory::build`
180    /// calls it after every `load_*` reader and before the KV budget is
181    /// computed, so the freed memory counts.
182    ///
183    /// The default keeps everything, which is right for a loader that binds
184    /// the store's device pointers zero-copy. A loader that uploads its own
185    /// copies (a TP shard, a host round-trip, a dtype conversion) overrides it
186    /// to free the originals.
187    fn prune_after_load(
188        &self,
189        _store: &mut WeightStore,
190        _config: &ModelConfig,
191        _gpu: &dyn GpuBackend,
192    ) -> Result<()> {
193        Ok(())
194    }
195
196    /// 2026-09-25: Per-(layer, role) weight precision schedule. The default is
197    /// `PrecisionSchedule::default()`, in which every role is `Dtype::Inherit`.
198    fn precision_schedule(
199        &self,
200        _config: &ModelConfig,
201    ) -> crate::precision_schedule::PrecisionSchedule {
202        crate::precision_schedule::PrecisionSchedule::default()
203    }
204
205    fn load_embedding(
206        &self,
207        store: &WeightStore,
208        config: &ModelConfig,
209        gpu: &dyn GpuBackend,
210    ) -> Result<DenseWeight>;
211
212    /// 2026-09-25: Build the n-gram embedding, for an architecture that fuses
213    /// hashed n-gram lookups into the input embedding (the LongCat loader).
214    /// The default `None` leaves the plain `embed_tokens` gather in place.
215    fn load_ngram_embedding(
216        &self,
217        _store: &WeightStore,
218        _config: &ModelConfig,
219        _gpu: &dyn GpuBackend,
220        _max_tokens: usize,
221    ) -> Result<Option<metrale_model_layers::layers::ngram_embed::NgramEmbedding>> {
222        Ok(None)
223    }
224    /// 2026-09-25: Load the final RMSNorm weight used before the LM head.
225    fn load_final_norm(
226        &self,
227        store: &WeightStore,
228        config: &ModelConfig,
229        gpu: &dyn GpuBackend,
230    ) -> Result<DenseWeight>;
231    fn load_lm_head(
232        &self,
233        store: &WeightStore,
234        config: &ModelConfig,
235        gpu: &dyn GpuBackend,
236    ) -> Result<DenseWeight>;
237
238    /// 2026-09-25: Load single-module MTP head weights, or `None`.
239    fn load_mtp_weights(
240        &self,
241        store: &WeightStore,
242        config: &ModelConfig,
243        gpu: &dyn GpuBackend,
244    ) -> Result<Option<MtpWeights>>;
245
246    /// 2026-09-25: Load every MTP module; `factory::build` calls this one.
247    /// The default wraps `load_mtp_weights` (zero or one module). The MiniMax
248    /// and Step3p7 loaders override it.
249    fn load_mtp_weights_multi(
250        &self,
251        store: &WeightStore,
252        config: &ModelConfig,
253        gpu: &dyn GpuBackend,
254    ) -> Result<Vec<MtpWeights>> {
255        Ok(self
256            .load_mtp_weights(store, config, gpu)?
257            .into_iter()
258            .collect())
259    }
260
261    /// 2026-09-25: Per-attention-layer `(num_kv_heads, head_dim)`, indexed like
262    /// `layer_kv_dtypes`, for models whose attention geometry differs per layer
263    /// (the Gemma-4 loader). The default is empty: every layer uses the global
264    /// dims.
265    fn kv_layer_dims(&self, _config: &ModelConfig) -> Vec<(usize, usize)> {
266        Vec::new()
267    }
268
269    /// 2026-09-25: Load DFlash drafter weights from the drafter's own
270    /// `WeightStore`. The default returns `None`, and no loader overrides it:
271    /// `factory::build` loads the drafter with the free function
272    /// `dflash_loader::load_dflash_weights`.
273    fn load_dflash_weights(
274        &self,
275        _drafter_store: &WeightStore,
276        _config: &ModelConfig,
277        _gpu: &dyn GpuBackend,
278        _tp_size: usize,
279    ) -> Result<Option<DflashWeights>> {
280        Ok(None)
281    }
282
283    /// 2026-09-25: Load startup-static PEFT LoRA adapters into the LoRA pool,
284    /// through `metrale_model_layers::lora::load_lora_adapters_multi`.
285    /// `factory::build` calls it before the buffer arena and the KV budget, so
286    /// the pool counts as used memory.
287    fn load_lora_adapters(
288        &self,
289        adapters: &[metrale_model_layers::lora::LoraAdapterInput<'_>],
290        config: &ModelConfig,
291        gpu: &dyn GpuBackend,
292        max_loras: usize,
293        max_lora_rank: usize,
294    ) -> Result<Option<metrale_model_layers::lora::LoraWeights>> {
295        metrale_model_layers::lora::load_lora_adapters_multi(
296            adapters,
297            config,
298            gpu,
299            max_loras,
300            max_lora_rank,
301        )
302        .map(Some)
303    }
304
305    /// 2026-09-25: Will this loader bind a vision encoder for a multimodal
306    /// checkpoint?
307    ///
308    /// Default `true`, so a loader that does not override it never loses a
309    /// weight. On `false` the checkpoint loader skips the tower's tensors
310    /// (`serve_phases/weights.rs`). `factory::build` frees an unbound tower
311    /// afterwards either way, keyed off the bind result, so this only lowers
312    /// peak memory.
313    fn binds_vision_encoder(&self) -> bool {
314        true
315    }
316
317    /// 2026-10-02: Does this loader read the serve's MoE expert-table decision
318    /// (`metrale_model_layers::layers::MoeExpertTables`)? The serve plans and publishes it before
319    /// load only for a loader that does.
320    fn reads_expert_table_plan(&self) -> bool {
321        false
322    }
323
324    /// 2026-09-25: Which checkpoint tensors will this loader read from the host
325    /// at bind time, so the checkpoint loader records their location instead
326    /// of uploading them?
327    ///
328    /// Default `None`: upload everything. `serve_phases/weights.rs` asks once,
329    /// before the load. A deferred tensor is not skipped: it is recorded in
330    /// `WeightStore::deferred`, and the loader that asked reads it. See
331    /// [`metrale_model_weights::weights::DeferHook`] for the contract.
332    fn defer_predicate(&self, _config: &ModelConfig) -> Option<DeferHook> {
333        None
334    }
335
336    /// 2026-09-25: Load vision encoder weights; the default returns `None`.
337    fn load_vision_encoder(
338        &self,
339        _store: &WeightStore,
340        _config: &ModelConfig,
341        _gpu: &dyn GpuBackend,
342    ) -> Result<Option<VisionTower>> {
343        Ok(None)
344    }
345}
metrale_model_arch/weight_loader/mod.rs

metrale_model_arch/weight_loader/
mod.rs
