metrale_model_arch/weight_loader/
qwen3_vl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen3VLWeightLoader`: text layers that are all full attention
4//! with an ungated Q projection and a MoE FFN without shared experts, plus the
5//! vision encoder. No MTP weights.
6//!
7//! Owner: model-arch weight loader.
8//! Invariants: none beyond the types.
9
10use anyhow::{Context, Result};
11use metrale_cache::kv_cache::KvCacheDtype;
12use metrale_config::ModelConfig;
13use metrale_gpu_runtime::gpu::GpuBackend;
14use metrale_model_weights::weights::WeightStore;
15
16use super::{ModelWeightLoader, WeightFormat};
17use crate::tp_shard::{TpShardKind, load_qkvo_tp, shard_quantized_nvfp4};
18use metrale_model_layers::layer::TransformerLayer;
19use metrale_model_layers::layers::vision_encoder::{MergerLayer, ViTBlock};
20use metrale_model_layers::layers::{
21    FfnComponent, MoeLayer, Qwen3AttentionLayer, VisionEncoder, VisionTower,
22};
23use metrale_model_layers::weight_map::{
24    AttentionWeights, DenseWeight, MtpWeights, dense, detect_nvfp4_variant, load_kv_scales,
25    load_moe_no_shared, quantize_to_nvfp4, quantized_auto,
26};
27
28pub struct Qwen3VLWeightLoader;
29
30impl ModelWeightLoader for Qwen3VLWeightLoader {
31    fn supports_tp(&self) -> bool {
32        // 2026-09-25: Q/K/V are sharded column-parallel and O row-parallel
33        // (`load_qkvo_tp`, `shard_quantized_nvfp4`); q_norm/k_norm, the MoE and
34        // the vision encoder are loaded whole on every rank.
35        true
36    }
37
38    fn load_layers(
39        &self,
40        store: &WeightStore,
41        config: &ModelConfig,
42        gpu: &dyn GpuBackend,
43        layer_kv_dtypes: &[KvCacheDtype],
44    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
45        let mut layers: Vec<Box<dyn TransformerLayer>> =
46            Vec::with_capacity(config.num_hidden_layers);
47
48        let variant = detect_nvfp4_variant(store, config);
49        let weight_format = WeightFormat::detect(store, config);
50        tracing::info!(
51            "Weight format: {:?}, NVFP4 variant: {:?}",
52            weight_format,
53            variant
54        );
55
56        // 2026-09-25: Whether to build the MoE prefill copies, decided once for
57        // every layer: `moe_prefill_copies_fit` reads free memory, which
58        // shrinks as layers load.
59        let moe_prefill_copies = super::moe_prefill_copies_fit(config, gpu);
60        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
61        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
62        let stream = gpu.default_stream();
63        let h = config.hidden_size;
64
65        for i in 0..config.num_hidden_layers {
66            let lp = config.layer_prefix(i);
67            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
68            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
69
70            let moe_weights =
71                load_moe_no_shared(store, &lp, config.num_experts, gpu, config, variant)?;
72            let gate_nvfp4 = quantize_to_nvfp4(
73                &moe_weights.gate,
74                config.num_experts,
75                h,
76                gpu,
77                absmax_k,
78                quantize_k,
79                stream,
80            )?;
81            let mut moe_layer = MoeLayer::new(
82                moe_weights,
83                config.num_experts,
84                Some(gate_nvfp4),
85                gpu,
86                config,
87            )?;
88            if moe_prefill_copies {
89                moe_layer.transpose_for_prefill(gpu, config)?;
90                moe_layer.predequant_for_prefill(gpu, config, stream)?;
91            }
92            let ffn = FfnComponent::Moe(moe_layer);
93
94            // 2026-09-25: Each projection is loaded with `quantized_auto` and,
95            // for TP > 1, sliced by `load_qkvo_tp`.
96            let p = format!("{lp}.self_attn");
97            let tp_rank = config.tp_rank;
98            let tp_size = config.tp_world_size.max(1);
99            let group_size = 16usize;
100            let load_proj =
101                |name: &str,
102                 full_n: usize,
103                 full_k: usize,
104                 kind: TpShardKind|
105                 -> Result<metrale_model_layers::weight_map::QuantizedWeight> {
106                    let src = quantized_auto(store, &format!("{p}.{name}"), gpu, variant)?;
107                    if tp_size == 1 {
108                        return Ok(src);
109                    }
110                    let sharded = shard_quantized_nvfp4(
111                        &src, full_n, full_k, kind, tp_rank, tp_size, group_size, gpu,
112                    )?;
113                    gpu.free(src.weight)?;
114                    gpu.free(src.weight_scale)?;
115                    Ok(sharded)
116                };
117            let [q, k, v, o] = load_qkvo_tp(config, load_proj)?;
118            let dummy = DenseWeight {
119                weight: metrale_gpu_runtime::gpu::DevicePtr::NULL,
120            };
121            let (k_scale, v_scale) = load_kv_scales(store, &p, gpu)?;
122            let attn = AttentionWeights {
123                q_proj: dummy,
124                k_proj: dummy,
125                v_proj: dummy,
126                o_proj: o,
127                q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
128                k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
129                q_norm_full: None,
130                k_norm_full: None,
131                k_scale,
132                v_scale,
133            };
134
135            layers.push(Box::new(Qwen3AttentionLayer::new_ungated(
136                input_norm,
137                attn,
138                post_attn_norm,
139                ffn,
140                i,
141                Some(q),
142                Some(k),
143                Some(v),
144                gpu,
145                layer_kv_dtypes[i],
146                config.fp8_kv_calibration_tokens,
147                config,
148            )?));
149
150            if (i + 1) % 10 == 0 {
151                tracing::info!("Loaded layers 0..{}", i + 1);
152            }
153        }
154
155        tracing::info!(
156            "Qwen3-VL weight loader: {} layers (all attention, ungated)",
157            layers.len(),
158        );
159
160        Ok(layers)
161    }
162
163    fn load_embedding(
164        &self,
165        store: &WeightStore,
166        config: &ModelConfig,
167        _gpu: &dyn GpuBackend,
168    ) -> Result<DenseWeight> {
169        let prefix = &config.weight_prefix;
170        dense(store, &format!("{prefix}.embed_tokens.weight"))
171    }
172
173    fn load_final_norm(
174        &self,
175        store: &WeightStore,
176        config: &ModelConfig,
177        _gpu: &dyn GpuBackend,
178    ) -> Result<DenseWeight> {
179        let prefix = &config.weight_prefix;
180        dense(store, &format!("{prefix}.norm.weight"))
181    }
182
183    fn load_lm_head(
184        &self,
185        store: &WeightStore,
186        config: &ModelConfig,
187        _gpu: &dyn GpuBackend,
188    ) -> Result<DenseWeight> {
189        for pattern in &[
190            "lm_head.weight",
191            "language_model.lm_head.weight",
192            "model.lm_head.weight",
193        ] {
194            if store.contains(pattern) {
195                return dense(store, pattern);
196            }
197        }
198        let prefix = &config.weight_prefix;
199        dense(store, &format!("{prefix}.embed_tokens.weight"))
200    }
201
202    fn load_mtp_weights(
203        &self,
204        _store: &WeightStore,
205        _config: &ModelConfig,
206        _gpu: &dyn GpuBackend,
207    ) -> Result<Option<MtpWeights>> {
208        Ok(None)
209    }
210
211    fn load_vision_encoder(
212        &self,
213        store: &WeightStore,
214        config: &ModelConfig,
215        gpu: &dyn GpuBackend,
216    ) -> Result<Option<VisionTower>> {
217        let vcfg = match &config.vision {
218            Some(v) => v.clone(),
219            None => return Ok(None),
220        };
221        let vp = "model.visual";
222
223        let patch_embed_w = dense(store, &format!("{vp}.patch_embed.proj.weight"))?;
224        let patch_embed_b = dense(store, &format!("{vp}.patch_embed.proj.bias"))?;
225        let pos_embed = dense(store, &format!("{vp}.pos_embed.weight"))?;
226        let pos_embed_shape = store.get(&format!("{vp}.pos_embed.weight"))?.shape.clone();
227        let num_position_embeddings = pos_embed_shape
228            .first()
229            .copied()
230            .context("pos_embed shape missing rows")?;
231
232        let mut blocks = Vec::with_capacity(vcfg.depth);
233        for i in 0..vcfg.depth {
234            let bp = format!("{vp}.blocks.{i}");
235            blocks.push(ViTBlock {
236                norm1_w: dense(store, &format!("{bp}.norm1.weight"))?.weight,
237                norm1_b: dense(store, &format!("{bp}.norm1.bias"))?.weight,
238                qkv_w: dense(store, &format!("{bp}.attn.qkv.weight"))?.weight,
239                qkv_b: dense(store, &format!("{bp}.attn.qkv.bias"))?.weight,
240                proj_w: dense(store, &format!("{bp}.attn.proj.weight"))?.weight,
241                proj_b: dense(store, &format!("{bp}.attn.proj.bias"))?.weight,
242                norm2_w: dense(store, &format!("{bp}.norm2.weight"))?.weight,
243                norm2_b: dense(store, &format!("{bp}.norm2.bias"))?.weight,
244                fc1_w: dense(store, &format!("{bp}.mlp.linear_fc1.weight"))?.weight,
245                fc1_b: dense(store, &format!("{bp}.mlp.linear_fc1.bias"))?.weight,
246                fc2_w: dense(store, &format!("{bp}.mlp.linear_fc2.weight"))?.weight,
247                fc2_b: dense(store, &format!("{bp}.mlp.linear_fc2.bias"))?.weight,
248            });
249        }
250
251        let mut deepstack = Vec::with_capacity(vcfg.deepstack_visual_indexes.len());
252        for i in 0..vcfg.deepstack_visual_indexes.len() {
253            let mp = format!("{vp}.deepstack_merger_list.{i}");
254            deepstack.push(MergerLayer {
255                norm_w: dense(store, &format!("{mp}.norm.weight"))?.weight,
256                norm_b: dense(store, &format!("{mp}.norm.bias"))?.weight,
257                fc1_w: dense(store, &format!("{mp}.linear_fc1.weight"))?.weight,
258                fc1_b: dense(store, &format!("{mp}.linear_fc1.bias"))?.weight,
259                fc2_w: dense(store, &format!("{mp}.linear_fc2.weight"))?.weight,
260                fc2_b: dense(store, &format!("{mp}.linear_fc2.bias"))?.weight,
261            });
262        }
263
264        let mp = format!("{vp}.merger");
265        let merger = MergerLayer {
266            norm_w: dense(store, &format!("{mp}.norm.weight"))?.weight,
267            norm_b: dense(store, &format!("{mp}.norm.bias"))?.weight,
268            fc1_w: dense(store, &format!("{mp}.linear_fc1.weight"))?.weight,
269            fc1_b: dense(store, &format!("{mp}.linear_fc1.bias"))?.weight,
270            fc2_w: dense(store, &format!("{mp}.linear_fc2.weight"))?.weight,
271            fc2_b: dense(store, &format!("{mp}.linear_fc2.bias"))?.weight,
272        };
273
274        let deepstack_indexes = vcfg.deepstack_visual_indexes.clone();
275        let ve = VisionEncoder::new(
276            patch_embed_w.weight,
277            patch_embed_b.weight,
278            pos_embed.weight,
279            num_position_embeddings,
280            blocks,
281            deepstack,
282            deepstack_indexes,
283            merger,
284            vcfg.hidden_size,
285            vcfg.num_heads,
286            vcfg.spatial_merge_size,
287            vcfg.out_hidden_size,
288            vcfg.intermediate_size,
289            vcfg.patch_size,
290            vcfg.max_pixels,
291            gpu,
292        )?;
293        tracing::info!(
294            "Vision encoder loaded: depth={}, hidden={}, heads={}, deepstack={:?}",
295            vcfg.depth,
296            vcfg.hidden_size,
297            vcfg.num_heads,
298            vcfg.deepstack_visual_indexes,
299        );
300        Ok(Some(VisionTower::qwen(ve)))
301    }
302}
metrale_model_arch/weight_loader/qwen3_vl.rs

metrale_model_arch/weight_loader/
qwen3_vl.rs
