metrale_model_arch/weight_loader/
qwen3.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen3WeightLoader`, the loader for `qwen3_next` checkpoints: full-attention
4//! and linear-attention (GDN) layers, each with an MoE block, plus embedding, final norm,
5//! LM head and MTP weights.
6//!
7//! Owner: model-arch weight loader (Qwen3-Next).
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_cache::kv_cache::KvCacheDtype;
12use metrale_config::{LayerType, ModelConfig};
13use metrale_gpu_runtime::gpu::GpuBackend;
14use metrale_model_weights::weights::WeightStore;
15
16use super::{ModelWeightLoader, QuantFormat};
17use crate::tp_shard::{TpShardKind, load_qkvo_tp, shard_dense_bf16, shard_fp8_block_scaled};
18use metrale_model_layers::layer::TransformerLayer;
19use metrale_model_layers::layers::{FfnComponent, MoeLayer, Qwen3AttentionLayer, Qwen3SsmLayer};
20use metrale_model_layers::weight_map::loaders_moe::load_mtp;
21use metrale_model_layers::weight_map::{
22    AttentionWeights, DenseWeight, MtpWeights, Nvfp4Variant, QuantizeCtx, QuantizedWeight, dense,
23    detect_nvfp4_variant, load_attention, load_fp8_block_scaled_as_fp8weight, load_kv_scales,
24    load_moe, load_moe_qwen35_fp8_experts, load_moe_skip_experts, load_ssm, quantize_to_nvfp4,
25};
26
27pub struct Qwen3WeightLoader;
28
29impl ModelWeightLoader for Qwen3WeightLoader {
30    fn supports_tp(&self) -> bool {
31        // 2026-09-25: Full-attention layers shard on both arms. Linear-attention layers are
32        // loaded unsharded (`load_ssm`).
33        true
34    }
35
36    fn load_layers(
37        &self,
38        store: &WeightStore,
39        config: &ModelConfig,
40        gpu: &dyn GpuBackend,
41        layer_kv_dtypes: &[KvCacheDtype],
42    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
43        let layer_types = if config.layer_types.is_empty() {
44            (0..config.num_hidden_layers)
45                .map(|i| config.layer_type(i))
46                .collect::<Vec<_>>()
47        } else {
48            config.layer_types.clone()
49        };
50
51        let mut layers: Vec<Box<dyn TransformerLayer>> =
52            Vec::with_capacity(config.num_hidden_layers);
53        let mut attn_idx = 0usize;
54
55        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
56        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
57        let stream = gpu.default_stream();
58        let qctx = QuantizeCtx {
59            absmax_k,
60            quantize_k,
61            stream,
62        };
63
64        let variant = detect_nvfp4_variant(store, config);
65        let quant_format = if variant == Nvfp4Variant::Fp8Dequanted {
66            QuantFormat::Fp8
67        } else {
68            QuantFormat::Nvfp4
69        };
70        let native_fp8 = quant_format == QuantFormat::Fp8;
71        tracing::info!(
72            "Qwen3 weight variant: {:?}, native_fp8: {}",
73            variant,
74            native_fp8
75        );
76
77        let h = config.hidden_size;
78
79        // 2026-09-25: `super::moe_prefill_copies_fit` decides whether the transposed MoE prefill
80        // tables fit in free memory; `METRALE_MOE_PREFILL_COPIES=0` forces them off.
81        let skip_moe_transpose = !super::moe_prefill_copies_fit(config, gpu);
82
83        for (i, lt) in layer_types.iter().enumerate() {
84            let lp = config.layer_prefix(i);
85            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
86            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
87
88            let moe_weights = if native_fp8 {
89                load_moe_skip_experts(store, &lp, config.num_experts, gpu, config, variant, qctx)?
90            } else {
91                load_moe(store, &lp, config.num_experts, gpu, config, variant, qctx)?
92            };
93            // 2026-09-25: `METRALE_BF16_ROUTER=1` keeps the router gate BF16 (`gate_nvfp4` None),
94            // so the gate GEMM runs dense on `weights.gate` (`moe/forward_batched_gate.rs`).
95            let gate_nvfp4 = if std::env::var("METRALE_BF16_ROUTER").as_deref() == Ok("1") {
96                None
97            } else {
98                Some(quantize_to_nvfp4(
99                    &moe_weights.gate,
100                    config.num_experts,
101                    h,
102                    gpu,
103                    absmax_k,
104                    quantize_k,
105                    stream,
106                )?)
107            };
108            let mut moe_layer =
109                MoeLayer::new(moe_weights, config.num_experts, gate_nvfp4, gpu, config)?;
110            if !native_fp8 && !skip_moe_transpose {
111                moe_layer.transpose_for_prefill(gpu, config)?;
112            }
113            if !native_fp8 {
114                moe_layer.predequant_for_prefill(gpu, config, stream)?;
115            }
116
117            // 2026-09-25: Native FP8: install the checkpoint's FP8 routed experts. A load error is
118            // discarded, and the layer then has no routed experts (`load_moe_skip_experts`).
119            if native_fp8
120                && let Ok(fp8_experts) =
121                    load_moe_qwen35_fp8_experts(store, &lp, config.num_experts, gpu, config)
122            {
123                let sp = format!("{lp}.mlp.shared_expert");
124                use metrale_gpu_runtime::gpu::DevicePtr;
125                use metrale_model_layers::weight_map::{Fp8ExpertWeight as FEW, Fp8Weight as FW};
126                let null_fw = FW {
127                    weight: DevicePtr::NULL,
128                    row_scale: DevicePtr::NULL,
129                    n: 0,
130                    k: 0,
131                    // 2026-09-25: Stands in for a shared-expert tensor that failed to load, tagged
132                    // like the block-scaled loader's output.
133                    scale_format:
134                        metrale_model_layers::weight_map::WeightQuantFormat::Fp8BlockScaled,
135                };
136                let sh_gate =
137                    load_fp8_block_scaled_as_fp8weight(store, &format!("{sp}.gate_proj"), gpu);
138                let sh_up =
139                    load_fp8_block_scaled_as_fp8weight(store, &format!("{sp}.up_proj"), gpu);
140                let sh_down =
141                    load_fp8_block_scaled_as_fp8weight(store, &format!("{sp}.down_proj"), gpu);
142                let shared_fp8 = FEW {
143                    gate_proj: sh_gate.unwrap_or(null_fw),
144                    up_proj: sh_up.unwrap_or(null_fw),
145                    down_proj: sh_down.unwrap_or(null_fw),
146                };
147                if let Err(e) = moe_layer.set_fp8_experts(&fp8_experts, shared_fp8, gpu) {
148                    tracing::error!("Layer {i}: FP8 expert tables failed: {e:#}");
149                } else {
150                    tracing::info!("Layer {i}: MoE experts loaded as native FP8");
151                }
152            }
153
154            let ffn = FfnComponent::Moe(moe_layer);
155
156            match lt {
157                LayerType::FullAttention if native_fp8 => {
158                    let p = format!("{lp}.self_attn");
159                    let tp_rank = config.tp_rank;
160                    let tp_size = config.tp_world_size.max(1);
161                    let block_size = 128usize;
162                    let load_fp8 = |name: &str,
163                                    _full_n: usize,
164                                    _full_k: usize,
165                                    kind: TpShardKind|
166                     -> Result<metrale_model_layers::weight_map::Fp8Weight> {
167                        let src =
168                            load_fp8_block_scaled_as_fp8weight(store, &format!("{p}.{name}"), gpu)?;
169                        if tp_size == 1 {
170                            return Ok(src);
171                        }
172                        let sharded =
173                            shard_fp8_block_scaled(&src, kind, tp_rank, tp_size, block_size, gpu)?;
174                        gpu.free(src.weight)?;
175                        gpu.free(src.row_scale)?;
176                        Ok(sharded)
177                    };
178                    let [q_fp8, k_fp8, v_fp8, o_fp8] = load_qkvo_tp(config, load_fp8)?;
179
180                    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu)?;
181                    let dummy = DenseWeight {
182                        weight: metrale_gpu_runtime::gpu::DevicePtr::NULL,
183                    };
184                    let attn = AttentionWeights {
185                        q_proj: dummy,
186                        k_proj: dummy,
187                        v_proj: dummy,
188                        o_proj: QuantizedWeight::null(),
189                        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
190                        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
191                        q_norm_full: None,
192                        k_norm_full: None,
193                        k_scale,
194                        v_scale,
195                    };
196
197                    let layer_kv_dtype = layer_kv_dtypes[attn_idx];
198                    let mut layer = Qwen3AttentionLayer::new(
199                        input_norm,
200                        attn,
201                        post_attn_norm,
202                        ffn,
203                        attn_idx,
204                        None,
205                        None,
206                        None,
207                        gpu,
208                        layer_kv_dtype,
209                        config.fp8_kv_calibration_tokens,
210                        config,
211                    )?;
212                    layer.set_fp8_weights(Some(q_fp8), Some(k_fp8), Some(v_fp8), Some(o_fp8));
213                    if let Err(e) = layer.transpose_fp8_for_prefill(gpu, stream) {
214                        tracing::warn!("Layer {i}: FP8 transpose failed: {e}");
215                    }
216                    layers.push(Box::new(layer));
217                    attn_idx += 1;
218                }
219                LayerType::FullAttention => {
220                    let mut attn = load_attention(store, &lp, gpu, variant, qctx, config)?;
221                    let tp_rank = config.tp_rank;
222                    let tp_size = config.tp_world_size.max(1);
223                    // 2026-09-25: Under TP, shard each projection before Q/K/V are quantized. The
224                    // sizes below come from `config`, whose head counts are already per rank
225                    // (server `serve_phases/topology.rs`).
226                    if tp_size > 1 {
227                        use crate::tp_shard::TpAttentionDims;
228                        let dims = TpAttentionDims::from_config(config);
229                        let (qp, _, _) = shard_dense_bf16(
230                            attn.q_proj.weight,
231                            dims.full_q_n,
232                            dims.h,
233                            TpShardKind::ColumnParallel,
234                            tp_rank,
235                            tp_size,
236                            gpu,
237                        )?;
238                        if qp != attn.q_proj.weight {
239                            gpu.free(attn.q_proj.weight)?;
240                        }
241                        attn.q_proj.weight = qp;
242                        let (kp, _, _) = shard_dense_bf16(
243                            attn.k_proj.weight,
244                            dims.full_kv_n,
245                            dims.h,
246                            TpShardKind::ColumnParallel,
247                            tp_rank,
248                            tp_size,
249                            gpu,
250                        )?;
251                        if kp != attn.k_proj.weight {
252                            gpu.free(attn.k_proj.weight)?;
253                        }
254                        attn.k_proj.weight = kp;
255                        let (vp, _, _) = shard_dense_bf16(
256                            attn.v_proj.weight,
257                            dims.full_kv_n,
258                            dims.h,
259                            TpShardKind::ColumnParallel,
260                            tp_rank,
261                            tp_size,
262                            gpu,
263                        )?;
264                        if vp != attn.v_proj.weight {
265                            gpu.free(attn.v_proj.weight)?;
266                        }
267                        attn.v_proj.weight = vp;
268                        let (op, _, _) = shard_dense_bf16(
269                            attn.o_proj.weight,
270                            dims.h,
271                            dims.full_o_in,
272                            TpShardKind::RowParallel,
273                            tp_rank,
274                            tp_size,
275                            gpu,
276                        )?;
277                        if op != attn.o_proj.weight {
278                            gpu.free(attn.o_proj.weight)?;
279                        }
280                        attn.o_proj.weight = op;
281                    }
282                    let q_nvfp4 = quantize_to_nvfp4(
283                        &attn.q_proj,
284                        config.num_attention_heads * config.head_dim * 2,
285                        h,
286                        gpu,
287                        absmax_k,
288                        quantize_k,
289                        stream,
290                    )?;
291                    let k_nvfp4 = quantize_to_nvfp4(
292                        &attn.k_proj,
293                        config.num_key_value_heads * config.head_dim,
294                        h,
295                        gpu,
296                        absmax_k,
297                        quantize_k,
298                        stream,
299                    )?;
300                    let v_nvfp4 = quantize_to_nvfp4(
301                        &attn.v_proj,
302                        config.num_key_value_heads * config.head_dim,
303                        h,
304                        gpu,
305                        absmax_k,
306                        quantize_k,
307                        stream,
308                    )?;
309                    let layer_kv_dtype = layer_kv_dtypes[attn_idx];
310                    let mut layer = Qwen3AttentionLayer::new(
311                        input_norm,
312                        attn,
313                        post_attn_norm,
314                        ffn,
315                        attn_idx,
316                        Some(q_nvfp4),
317                        Some(k_nvfp4),
318                        Some(v_nvfp4),
319                        gpu,
320                        layer_kv_dtype,
321                        config.fp8_kv_calibration_tokens,
322                        config,
323                    )?;
324                    let qt = q_nvfp4.transpose_for_gemm(
325                        gpu,
326                        config.num_attention_heads * config.head_dim * 2,
327                        h,
328                    )?;
329                    let kt = k_nvfp4.transpose_for_gemm(
330                        gpu,
331                        config.num_key_value_heads * config.head_dim,
332                        h,
333                    )?;
334                    let vt = v_nvfp4.transpose_for_gemm(
335                        gpu,
336                        config.num_key_value_heads * config.head_dim,
337                        h,
338                    )?;
339                    let ot = layer.attn.o_proj.transpose_for_gemm(
340                        gpu,
341                        h,
342                        config.num_attention_heads * config.head_dim,
343                    )?;
344                    layer.set_prefill_weights(Some(qt), Some(kt), Some(vt), Some(ot));
345                    layer.predequant_for_prefill(gpu, config, stream)?;
346                    layers.push(Box::new(layer));
347                    attn_idx += 1;
348                }
349                // 2026-09-25: Linear attention takes one path for every variant: `load_ssm`, then
350                // QKVZ quantized to NVFP4. It is not sharded under TP.
351                LayerType::LinearAttention => {
352                    let ssm = load_ssm(store, &lp, gpu, variant, qctx, config)?;
353                    let qkvz_nvfp4 = quantize_to_nvfp4(
354                        &ssm.in_proj_qkvz,
355                        config.ssm_qkvz_size(),
356                        h,
357                        gpu,
358                        absmax_k,
359                        quantize_k,
360                        stream,
361                    )?;
362                    layers.push(Box::new(Qwen3SsmLayer::new(
363                        input_norm,
364                        ssm,
365                        post_attn_norm,
366                        ffn,
367                        Some(qkvz_nvfp4),
368                        config,
369                        gpu,
370                    )?));
371                }
372                LayerType::SlidingAttention => {
373                    unreachable!("unexpected SlidingAttention in this loader")
374                }
375                LayerType::Moe => unreachable!("Qwen3 has no standalone MoE layers"),
376                // 2026-09-25: GLM-5.3 `deepseek_sparse_attention` needs a DSA indexer and per-query
377                // top-k, which this loader does not build; it is refused, not served as dense
378                // attention.
379                LayerType::SparseAttention => anyhow::bail!(
380                    "layer {i}: SparseAttention needs a DSA indexer and per-query top-k; Qwen3 has neither"
381                ),
382            }
383
384            if (i + 1) % 12 == 0 {
385                tracing::info!("Loaded layers 0..{}", i + 1);
386            }
387        }
388
389        tracing::info!(
390            "Weight loader: {} layers ({} attention, {} SSM)",
391            layers.len(),
392            attn_idx,
393            layers.len() - attn_idx,
394        );
395
396        Ok(layers)
397    }
398
399    fn load_embedding(
400        &self,
401        store: &WeightStore,
402        _config: &ModelConfig,
403        _gpu: &dyn GpuBackend,
404    ) -> Result<DenseWeight> {
405        dense(store, "model.embed_tokens.weight")
406    }
407
408    fn load_final_norm(
409        &self,
410        store: &WeightStore,
411        _config: &ModelConfig,
412        _gpu: &dyn GpuBackend,
413    ) -> Result<DenseWeight> {
414        dense(store, "model.norm.weight")
415    }
416
417    fn load_lm_head(
418        &self,
419        store: &WeightStore,
420        _config: &ModelConfig,
421        _gpu: &dyn GpuBackend,
422    ) -> Result<DenseWeight> {
423        if store.contains("lm_head.weight") {
424            dense(store, "lm_head.weight")
425        } else {
426            dense(store, "model.embed_tokens.weight")
427        }
428    }
429
430    fn load_mtp_weights(
431        &self,
432        store: &WeightStore,
433        config: &ModelConfig,
434        gpu: &dyn GpuBackend,
435    ) -> Result<Option<MtpWeights>> {
436        if !store.contains("mtp.fc.weight") {
437            tracing::info!("No MTP weights found — speculative decoding disabled");
438            return Ok(None);
439        }
440        let variant = detect_nvfp4_variant(store, config);
441        tracing::info!("Loading MTP weights (variant={:?})...", variant);
442        let mtp = load_mtp(store, config.num_experts, gpu, variant)?;
443        tracing::info!(
444            "MTP weights loaded: fc=[2048,4096], {} experts, attn layer",
445            mtp.experts.len(),
446        );
447        Ok(Some(mtp))
448    }
449}
metrale_model_arch/weight_loader/qwen3.rs

metrale_model_arch/weight_loader/
qwen3.rs
