metrale_model_arch/weight_loader/
qwen35.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen35WeightLoader`, the loader for the Qwen3.5/3.6 MoE family and
4//! `holo3_1_moe`: layers (`load_layers`), embedding, final norm, LM head, MTP weights and
5//! the vision tower.
6//!
7//! Owner: model-arch weight loader (Qwen3.5).
8//! Invariants: none beyond the types.
9
10pub(crate) mod load_layers;
11
12use anyhow::{Context, Result};
13use metrale_cache::kv_cache::KvCacheDtype;
14use metrale_config::ModelConfig;
15use metrale_gpu_runtime::gpu::GpuBackend;
16use metrale_model_weights::weights::{WeightDtype, WeightStore};
17
18use super::ModelWeightLoader;
19use metrale_model_layers::layer::TransformerLayer;
20use metrale_model_layers::weight_map::loaders_moe::load_mtp;
21use metrale_model_layers::weight_map::quant_helpers::dense_auto;
22use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, detect_nvfp4_variant};
23
24pub struct Qwen35WeightLoader;
25
26fn vision_dense_auto(
27    store: &WeightStore,
28    prefix: &str,
29    gpu: &dyn GpuBackend,
30) -> Result<DenseWeight> {
31    let name = format!("{prefix}.weight");
32    let w = store.get(&name)?;
33    match w.dtype {
34        WeightDtype::BF16 | WeightDtype::FP8E4M3 => {
35            metrale_model_layers::weight_map::dense_auto_fp8_or_bf16(store, prefix, gpu)
36        }
37        WeightDtype::FP32 => metrale_model_layers::weight_map::dense_f32_safe(store, &name, gpu),
38        other => anyhow::bail!("vision_dense_auto: unsupported dtype {other:?} for {name}"),
39    }
40}
41
42fn vision_tensor_dense_auto(
43    store: &WeightStore,
44    name: &str,
45    gpu: &dyn GpuBackend,
46) -> Result<DenseWeight> {
47    let w = store.get(name)?;
48    match w.dtype {
49        WeightDtype::BF16 => Ok(DenseWeight { weight: w.ptr }),
50        WeightDtype::FP32 => metrale_model_layers::weight_map::dense_f32_safe(store, name, gpu),
51        other => anyhow::bail!("vision_tensor_dense_auto: unsupported dtype {other:?} for {name}"),
52    }
53}
54
55impl ModelWeightLoader for Qwen35WeightLoader {
56    /// 2026-10-02: Its MoE layers build the transposed NVFP4 prefill tables only when the
57    /// serve's memory plan says they fit (`load_layers.rs`).
58    fn reads_expert_table_plan(&self) -> bool {
59        true
60    }
61
62    fn supports_tp(&self) -> bool {
63        // 2026-09-25: The native-FP8 and NVFP4 attention arms shard Q/K/V/O, and the BF16
64        // dense and NVFP4 linear-attention builders shard by head. The BF16-dequant attention
65        // arm and the native-FP8 linear-attention builder refuse TP > 1.
66        true
67    }
68
69    fn load_layers(
70        &self,
71        store: &WeightStore,
72        config: &ModelConfig,
73        gpu: &dyn GpuBackend,
74        layer_kv_dtypes: &[KvCacheDtype],
75    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
76        load_layers::load_layers(self, store, config, gpu, layer_kv_dtypes)
77    }
78
79    fn load_embedding(
80        &self,
81        store: &WeightStore,
82        config: &ModelConfig,
83        gpu: &dyn GpuBackend,
84    ) -> Result<DenseWeight> {
85        let prefix = &config.weight_prefix;
86        dense_auto(store, &format!("{prefix}.embed_tokens.weight"), gpu)
87    }
88
89    fn load_final_norm(
90        &self,
91        store: &WeightStore,
92        config: &ModelConfig,
93        gpu: &dyn GpuBackend,
94    ) -> Result<DenseWeight> {
95        let prefix = &config.weight_prefix;
96        dense_auto(store, &format!("{prefix}.norm.weight"), gpu)
97    }
98
99    fn load_lm_head(
100        &self,
101        store: &WeightStore,
102        config: &ModelConfig,
103        gpu: &dyn GpuBackend,
104    ) -> Result<DenseWeight> {
105        // 2026-09-25: The head is the first of `lm_head.weight`, `language_model.lm_head.weight`
106        // and `model.lm_head.weight` present. An FP8E4M3 head is dequantized to BF16
107        // (`dense_auto_fp8_or_bf16`, which rejects any dtype but BF16 and FP8E4M3); every other
108        // dtype is passed to `dense` unchecked, so an NVFP4 head's packed bytes load as they are.
109        for prefix in ["lm_head", "language_model.lm_head", "model.lm_head"] {
110            let key = format!("{prefix}.weight");
111            if !store.contains(&key) {
112                continue;
113            }
114            let is_fp8 = store
115                .get(&key)
116                .map(|w| w.dtype == WeightDtype::FP8E4M3)
117                .unwrap_or(false);
118            return if is_fp8 {
119                metrale_model_layers::weight_map::dense_auto_fp8_or_bf16(store, prefix, gpu)
120            } else {
121                metrale_model_layers::weight_map::dense(store, &key)
122            };
123        }
124        // 2026-09-25: No head tensor: the head is the embedding table (tied embeddings).
125        let prefix = &config.weight_prefix;
126        metrale_model_layers::weight_map::dense(store, &format!("{prefix}.embed_tokens.weight"))
127    }
128
129    fn load_mtp_weights(
130        &self,
131        store: &WeightStore,
132        config: &ModelConfig,
133        gpu: &dyn GpuBackend,
134    ) -> Result<Option<MtpWeights>> {
135        if !store.contains("mtp.fc.weight") {
136            tracing::info!("No MTP weights found — speculative decoding disabled");
137            return Ok(None);
138        }
139        let variant = detect_nvfp4_variant(store, config);
140        tracing::info!(
141            "Loading MTP weights ({} experts, variant={:?})...",
142            config.num_experts,
143            variant
144        );
145        let mtp = load_mtp(store, config.num_experts, gpu, variant)?;
146        tracing::info!(
147            "MTP weights loaded: fc=[2048,4096], {} experts, attn layer",
148            mtp.experts.len(),
149        );
150        Ok(Some(mtp))
151    }
152
153    /// 2026-09-25: Loads the ViT tower from `model.visual` or `model.language_model.visual`.
154    /// Returns `None` when `config.vision` is `None` or neither prefix holds
155    /// `patch_embed.proj.weight`. Reads `depth` blocks, one deepstack merger per
156    /// `deepstack_visual_indexes` entry, and the final merger. Linear weights may be BF16,
157    /// FP32 or FP8E4M3 (dequantized to BF16); biases, norms and embeddings BF16 or FP32.
158    fn load_vision_encoder(
159        &self,
160        store: &WeightStore,
161        config: &ModelConfig,
162        gpu: &dyn GpuBackend,
163    ) -> Result<Option<metrale_model_layers::layers::VisionTower>> {
164        let vcfg = match &config.vision {
165            Some(v) => v.clone(),
166            None => return Ok(None),
167        };
168        // 2026-09-25: The flat `model.visual.*` and the nested `model.language_model.visual.*`
169        // layouts are both accepted; the first that holds `patch_embed.proj.weight` wins.
170        let vp = if store.contains("model.visual.patch_embed.proj.weight") {
171            "model.visual"
172        } else if store.contains("model.language_model.visual.patch_embed.proj.weight") {
173            "model.language_model.visual"
174        } else {
175            tracing::warn!(
176                "Vision encoder tensors absent under both `model.visual.*` and \
177                 `model.language_model.visual.*`; skipping vision tower (text-only mode)"
178            );
179            return Ok(None);
180        };
181
182        let patch_embed_w =
183            vision_tensor_dense_auto(store, &format!("{vp}.patch_embed.proj.weight"), gpu)?;
184        let patch_embed_b =
185            vision_tensor_dense_auto(store, &format!("{vp}.patch_embed.proj.bias"), gpu)?;
186        let pos_embed = vision_tensor_dense_auto(store, &format!("{vp}.pos_embed.weight"), gpu)?;
187        let pos_embed_shape = store.get(&format!("{vp}.pos_embed.weight"))?.shape.clone();
188        let num_position_embeddings = pos_embed_shape
189            .first()
190            .copied()
191            .context("pos_embed shape missing rows")?;
192
193        let mut blocks = Vec::with_capacity(vcfg.depth);
194        for i in 0..vcfg.depth {
195            let bp = format!("{vp}.blocks.{i}");
196            blocks.push(metrale_model_layers::layers::ViTBlock {
197                norm1_w: vision_tensor_dense_auto(store, &format!("{bp}.norm1.weight"), gpu)?
198                    .weight,
199                norm1_b: vision_tensor_dense_auto(store, &format!("{bp}.norm1.bias"), gpu)?.weight,
200                qkv_w: vision_dense_auto(store, &format!("{bp}.attn.qkv"), gpu)?.weight,
201                qkv_b: vision_tensor_dense_auto(store, &format!("{bp}.attn.qkv.bias"), gpu)?.weight,
202                proj_w: vision_dense_auto(store, &format!("{bp}.attn.proj"), gpu)?.weight,
203                proj_b: vision_tensor_dense_auto(store, &format!("{bp}.attn.proj.bias"), gpu)?
204                    .weight,
205                norm2_w: vision_tensor_dense_auto(store, &format!("{bp}.norm2.weight"), gpu)?
206                    .weight,
207                norm2_b: vision_tensor_dense_auto(store, &format!("{bp}.norm2.bias"), gpu)?.weight,
208                fc1_w: vision_dense_auto(store, &format!("{bp}.mlp.linear_fc1"), gpu)?.weight,
209                fc1_b: vision_tensor_dense_auto(store, &format!("{bp}.mlp.linear_fc1.bias"), gpu)?
210                    .weight,
211                fc2_w: vision_dense_auto(store, &format!("{bp}.mlp.linear_fc2"), gpu)?.weight,
212                fc2_b: vision_tensor_dense_auto(store, &format!("{bp}.mlp.linear_fc2.bias"), gpu)?
213                    .weight,
214            });
215        }
216
217        let mut deepstack = Vec::with_capacity(vcfg.deepstack_visual_indexes.len());
218        for i in 0..vcfg.deepstack_visual_indexes.len() {
219            let mp = format!("{vp}.deepstack_merger_list.{i}");
220            deepstack.push(metrale_model_layers::layers::MergerLayer {
221                norm_w: vision_tensor_dense_auto(store, &format!("{mp}.norm.weight"), gpu)?.weight,
222                norm_b: vision_tensor_dense_auto(store, &format!("{mp}.norm.bias"), gpu)?.weight,
223                fc1_w: vision_dense_auto(store, &format!("{mp}.linear_fc1"), gpu)?.weight,
224                fc1_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc1.bias"), gpu)?
225                    .weight,
226                fc2_w: vision_dense_auto(store, &format!("{mp}.linear_fc2"), gpu)?.weight,
227                fc2_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc2.bias"), gpu)?
228                    .weight,
229            });
230        }
231
232        let mp = format!("{vp}.merger");
233        let merger = metrale_model_layers::layers::MergerLayer {
234            norm_w: vision_tensor_dense_auto(store, &format!("{mp}.norm.weight"), gpu)?.weight,
235            norm_b: vision_tensor_dense_auto(store, &format!("{mp}.norm.bias"), gpu)?.weight,
236            fc1_w: vision_dense_auto(store, &format!("{mp}.linear_fc1"), gpu)?.weight,
237            fc1_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc1.bias"), gpu)?.weight,
238            fc2_w: vision_dense_auto(store, &format!("{mp}.linear_fc2"), gpu)?.weight,
239            fc2_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc2.bias"), gpu)?.weight,
240        };
241
242        let deepstack_indexes = vcfg.deepstack_visual_indexes.clone();
243        let ve = metrale_model_layers::layers::VisionEncoder::new(
244            patch_embed_w.weight,
245            patch_embed_b.weight,
246            pos_embed.weight,
247            num_position_embeddings,
248            blocks,
249            deepstack,
250            deepstack_indexes,
251            merger,
252            vcfg.hidden_size,
253            vcfg.num_heads,
254            vcfg.spatial_merge_size,
255            vcfg.out_hidden_size,
256            vcfg.intermediate_size,
257            vcfg.patch_size,
258            vcfg.max_pixels,
259            gpu,
260        )?;
261        tracing::info!(
262            "Qwen3.6 vision encoder loaded: depth={}, hidden={}, heads={}, FP8-blocks>=4",
263            vcfg.depth,
264            vcfg.hidden_size,
265            vcfg.num_heads,
266        );
267        Ok(Some(metrale_model_layers::layers::VisionTower::qwen(ve)))
268    }
269}
metrale_model_arch/weight_loader/qwen35.rs

metrale_model_arch/weight_loader/
qwen35.rs
