spark_model/weight_loader/
qwen3.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3use anyhow::Result;
4use atlas_core::config::{LayerType, ModelConfig};
5use spark_runtime::gpu::GpuBackend;
6use spark_runtime::kv_cache::KvCacheDtype;
7use spark_runtime::weights::WeightStore;
8
9use super::{ModelWeightLoader, QuantFormat};
10use crate::layer::TransformerLayer;
11use crate::layers::{FfnComponent, MoeLayer, Qwen3AttentionLayer, Qwen3SsmLayer};
12use crate::tp_shard::{TpShardKind, load_qkvo_tp, shard_dense_bf16, shard_fp8_block_scaled};
13use crate::weight_map::{
14    AttentionWeights, DenseWeight, MtpWeights, Nvfp4Variant, QuantizeCtx, QuantizedWeight, dense,
15    detect_nvfp4_variant, load_attention, load_fp8_block_scaled_as_fp8weight, load_kv_scales,
16    load_moe, load_moe_qwen35_fp8_experts, load_moe_skip_experts, load_mtp, load_ssm,
17    quantize_to_nvfp4,
18};
19
20pub struct Qwen3WeightLoader;
21
22impl ModelWeightLoader for Qwen3WeightLoader {
23    fn supports_tp(&self) -> bool {
24        // Qwen3-Next FullAttention layers (gated) are TP-sharded across
25        // both quant paths (FP8 native, BF16 → NVFP4). LinearAttention
26        // (GDN SSM) layers run full-replica per rank — same trade-off as
27        // qwen35.rs: SSM weight memory not recovered, but functionally
28        // correct and the bulk of compute (attention + MoE) is sharded.
29        true
30    }
31
32    fn load_layers(
33        &self,
34        store: &WeightStore,
35        config: &ModelConfig,
36        gpu: &dyn GpuBackend,
37        layer_kv_dtypes: &[KvCacheDtype],
38    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
39        let layer_types = if config.layer_types.is_empty() {
40            (0..config.num_hidden_layers)
41                .map(|i| config.layer_type(i))
42                .collect::<Vec<_>>()
43        } else {
44            config.layer_types.clone()
45        };
46
47        let mut layers: Vec<Box<dyn TransformerLayer>> =
48            Vec::with_capacity(config.num_hidden_layers);
49        let mut attn_idx = 0usize;
50
51        // Kernels + stream for BF16→NVFP4 runtime quantization of dense weights
52        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
53        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
54        let stream = gpu.default_stream();
55        let qctx = QuantizeCtx {
56            absmax_k,
57            quantize_k,
58            stream,
59        };
60
61        // Detect weight format variant (Standard NVFP4, CompressedTensors, or FP8 block-scaled).
62        let variant = detect_nvfp4_variant(store, config);
63        let quant_format = if variant == Nvfp4Variant::Fp8Dequanted {
64            QuantFormat::Fp8
65        } else {
66            QuantFormat::Nvfp4
67        };
68        let native_fp8 = quant_format == QuantFormat::Fp8;
69        tracing::info!(
70            "Qwen3 weight variant: {:?}, native_fp8: {}",
71            variant,
72            native_fp8
73        );
74
75        let h = config.hidden_size;
76
77        // SSOT: the budget arithmetic and the `ATLAS_MOE_PREFILL_COPIES` lever
78        // live in `super::moe_prefill_copies_fit` — shared with every other MoE
79        // loader instead of one inline copy per family.
80        let skip_moe_transpose = !super::moe_prefill_copies_fit(config, gpu);
81
82        for (i, lt) in layer_types.iter().enumerate() {
83            let lp = config.layer_prefix(i);
84            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
85            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
86
87            // ── MoE weights ──
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            // ATLAS_BF16_ROUTER=1: keep the MoE router/gate in BF16 (skip the
94            // NVFP4 quant) so expert SELECTION is decided by full-precision gate
95            // logits. The bf16moe experiment showed dequanting EXPERTS to BF16
96            // eliminates the empty_path tool-call drift (FP8 flips were the seed)
97            // but halves decode throughput. The router is a tiny num_experts×h
98            // GEMM, so making ONLY it high-precision targets the expert-selection
99            // flips at ~zero throughput cost (experts stay FP8). The forward
100            // (dense_gemv/dense_gemm) already falls back to weights.gate (BF16)
101            // when gate_nvfp4 is None. Explicit opt-in (PCND); default unchanged.
102            let gate_nvfp4 = if std::env::var("ATLAS_BF16_ROUTER").as_deref() == Ok("1") {
103                None
104            } else {
105                Some(quantize_to_nvfp4(
106                    &moe_weights.gate,
107                    config.num_experts,
108                    h,
109                    gpu,
110                    absmax_k,
111                    quantize_k,
112                    stream,
113                )?)
114            };
115            let mut moe_layer =
116                MoeLayer::new(moe_weights, config.num_experts, gate_nvfp4, gpu, config)?;
117            if !native_fp8 && !skip_moe_transpose {
118                moe_layer.transpose_for_prefill(gpu, config)?;
119            }
120            if !native_fp8 {
121                moe_layer.predequant_for_prefill(gpu, config, stream)?;
122            }
123
124            // Native FP8 MoE: load FP8 expert weights for fused batch dispatch
125            if native_fp8
126                && let Ok(fp8_experts) =
127                    load_moe_qwen35_fp8_experts(store, &lp, config.num_experts, gpu, config)
128            {
129                let sp = format!("{lp}.mlp.shared_expert");
130                use crate::weight_map::{Fp8ExpertWeight as FEW, Fp8Weight as FW};
131                use spark_runtime::gpu::DevicePtr;
132                let null_fw = FW {
133                    weight: DevicePtr::NULL,
134                    row_scale: DevicePtr::NULL,
135                    n: 0,
136                    k: 0,
137                    // Placeholder for absent shared-expert tensor: the
138                    // calling site checks `weight == NULL` before
139                    // launching any kernel, so the tag is conventional.
140                    // Match the block-scaled FP8 loader the other
141                    // arms use so the format is consistent.
142                    scale_format: crate::weight_map::WeightQuantFormat::Fp8BlockScaled,
143                };
144                let sh_gate =
145                    load_fp8_block_scaled_as_fp8weight(store, &format!("{sp}.gate_proj"), gpu);
146                let sh_up =
147                    load_fp8_block_scaled_as_fp8weight(store, &format!("{sp}.up_proj"), gpu);
148                let sh_down =
149                    load_fp8_block_scaled_as_fp8weight(store, &format!("{sp}.down_proj"), gpu);
150                let shared_fp8 = FEW {
151                    gate_proj: sh_gate.unwrap_or(null_fw),
152                    up_proj: sh_up.unwrap_or(null_fw),
153                    down_proj: sh_down.unwrap_or(null_fw),
154                };
155                if let Err(e) = moe_layer.set_fp8_experts(&fp8_experts, shared_fp8, gpu) {
156                    tracing::error!("Layer {i}: FP8 expert tables failed: {e:#}");
157                } else {
158                    tracing::info!("Layer {i}: MoE experts loaded as native FP8");
159                }
160            }
161
162            let ffn = FfnComponent::Moe(moe_layer);
163
164            match lt {
165                // ── Native FP8 Attention ──
166                LayerType::FullAttention if native_fp8 => {
167                    let p = format!("{lp}.self_attn");
168                    let tp_rank = config.tp_rank;
169                    let tp_size = config.tp_world_size.max(1);
170                    let block_size = 128usize;
171                    let load_fp8 = |name: &str,
172                                    _full_n: usize,
173                                    _full_k: usize,
174                                    kind: TpShardKind|
175                     -> Result<crate::weight_map::Fp8Weight> {
176                        let src =
177                            load_fp8_block_scaled_as_fp8weight(store, &format!("{p}.{name}"), gpu)?;
178                        if tp_size == 1 {
179                            return Ok(src);
180                        }
181                        let sharded =
182                            shard_fp8_block_scaled(&src, kind, tp_rank, tp_size, block_size, gpu)?;
183                        gpu.free(src.weight)?;
184                        gpu.free(src.row_scale)?;
185                        Ok(sharded)
186                    };
187                    let [q_fp8, k_fp8, v_fp8, o_fp8] = load_qkvo_tp(config, load_fp8)?;
188
189                    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
190                    let dummy = DenseWeight {
191                        weight: spark_runtime::gpu::DevicePtr::NULL,
192                    };
193                    let attn = AttentionWeights {
194                        q_proj: dummy,
195                        k_proj: dummy,
196                        v_proj: dummy,
197                        o_proj: QuantizedWeight::null(),
198                        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
199                        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
200                        q_norm_full: None,
201                        k_norm_full: None,
202                        k_scale,
203                        v_scale,
204                    };
205
206                    let layer_kv_dtype = layer_kv_dtypes[attn_idx];
207                    let mut layer = Qwen3AttentionLayer::new(
208                        input_norm,
209                        attn,
210                        post_attn_norm,
211                        ffn,
212                        attn_idx,
213                        None,
214                        None,
215                        None, // No NVFP4 — w8a16 handles everything
216                        gpu,
217                        layer_kv_dtype,
218                        config.fp8_kv_calibration_tokens,
219                        config,
220                    )?;
221                    layer.set_fp8_weights(Some(q_fp8), Some(k_fp8), Some(v_fp8), Some(o_fp8));
222                    if let Err(e) = layer.transpose_fp8_for_prefill(gpu, stream) {
223                        tracing::warn!("Layer {i}: FP8 transpose failed: {e}");
224                    }
225                    layers.push(Box::new(layer));
226                    attn_idx += 1;
227                }
228                // ── NVFP4 Attention (original path) ──
229                LayerType::FullAttention => {
230                    let mut attn = load_attention(store, &lp, gpu, variant, qctx, config)?;
231                    let tp_rank = config.tp_rank;
232                    let tp_size = config.tp_world_size.max(1);
233                    // TP shard each BF16 projection BEFORE quantization. After
234                    // sharding, dims are TP-LOCAL and config head counts (already
235                    // divided in main.rs) match — so the post-shard quantize
236                    // calls below use the correct local sizes naturally.
237                    if tp_size > 1 {
238                        use crate::tp_shard::TpAttentionDims;
239                        let dims = TpAttentionDims::from_config(config);
240                        let (qp, _, _) = shard_dense_bf16(
241                            attn.q_proj.weight,
242                            dims.full_q_n,
243                            dims.h,
244                            TpShardKind::ColumnParallel,
245                            tp_rank,
246                            tp_size,
247                            gpu,
248                        )?;
249                        if qp != attn.q_proj.weight {
250                            gpu.free(attn.q_proj.weight)?;
251                        }
252                        attn.q_proj.weight = qp;
253                        let (kp, _, _) = shard_dense_bf16(
254                            attn.k_proj.weight,
255                            dims.full_kv_n,
256                            dims.h,
257                            TpShardKind::ColumnParallel,
258                            tp_rank,
259                            tp_size,
260                            gpu,
261                        )?;
262                        if kp != attn.k_proj.weight {
263                            gpu.free(attn.k_proj.weight)?;
264                        }
265                        attn.k_proj.weight = kp;
266                        let (vp, _, _) = shard_dense_bf16(
267                            attn.v_proj.weight,
268                            dims.full_kv_n,
269                            dims.h,
270                            TpShardKind::ColumnParallel,
271                            tp_rank,
272                            tp_size,
273                            gpu,
274                        )?;
275                        if vp != attn.v_proj.weight {
276                            gpu.free(attn.v_proj.weight)?;
277                        }
278                        attn.v_proj.weight = vp;
279                        let (op, _, _) = shard_dense_bf16(
280                            attn.o_proj.weight,
281                            dims.h,
282                            dims.full_o_in,
283                            TpShardKind::RowParallel,
284                            tp_rank,
285                            tp_size,
286                            gpu,
287                        )?;
288                        if op != attn.o_proj.weight {
289                            gpu.free(attn.o_proj.weight)?;
290                        }
291                        attn.o_proj.weight = op;
292                    }
293                    let q_nvfp4 = quantize_to_nvfp4(
294                        &attn.q_proj,
295                        config.num_attention_heads * config.head_dim * 2,
296                        h,
297                        gpu,
298                        absmax_k,
299                        quantize_k,
300                        stream,
301                    )?;
302                    let k_nvfp4 = quantize_to_nvfp4(
303                        &attn.k_proj,
304                        config.num_key_value_heads * config.head_dim,
305                        h,
306                        gpu,
307                        absmax_k,
308                        quantize_k,
309                        stream,
310                    )?;
311                    let v_nvfp4 = quantize_to_nvfp4(
312                        &attn.v_proj,
313                        config.num_key_value_heads * config.head_dim,
314                        h,
315                        gpu,
316                        absmax_k,
317                        quantize_k,
318                        stream,
319                    )?;
320                    let layer_kv_dtype = layer_kv_dtypes[attn_idx];
321                    let mut layer = Qwen3AttentionLayer::new(
322                        input_norm,
323                        attn,
324                        post_attn_norm,
325                        ffn,
326                        attn_idx,
327                        Some(q_nvfp4),
328                        Some(k_nvfp4),
329                        Some(v_nvfp4),
330                        gpu,
331                        layer_kv_dtype,
332                        config.fp8_kv_calibration_tokens,
333                        config,
334                    )?;
335                    let qt = q_nvfp4.transpose_for_gemm(
336                        gpu,
337                        config.num_attention_heads * config.head_dim * 2,
338                        h,
339                    )?;
340                    let kt = k_nvfp4.transpose_for_gemm(
341                        gpu,
342                        config.num_key_value_heads * config.head_dim,
343                        h,
344                    )?;
345                    let vt = v_nvfp4.transpose_for_gemm(
346                        gpu,
347                        config.num_key_value_heads * config.head_dim,
348                        h,
349                    )?;
350                    let ot = layer.attn.o_proj.transpose_for_gemm(
351                        gpu,
352                        h,
353                        config.num_attention_heads * config.head_dim,
354                    )?;
355                    layer.set_prefill_weights(Some(qt), Some(kt), Some(vt), Some(ot));
356                    layer.predequant_for_prefill(gpu, config, stream)?;
357                    layers.push(Box::new(layer));
358                    attn_idx += 1;
359                }
360                // ── SSM (FP8→BF16→NVFP4 conversion, same path for native_fp8 and non-native) ──
361                // Native FP8 SSM decode is disabled upstream (Qwen35 `&& false`) due to
362                // block-scale → per-row-scale precision loss. Instead, dequant FP8→BF16
363                // then quantize BF16→NVFP4. Only qkvz + out_proj need conversion (tiny).
364                //
365                LayerType::LinearAttention => {
366                    let ssm = load_ssm(store, &lp, gpu, variant, qctx, config)?;
367                    let qkvz_nvfp4 = quantize_to_nvfp4(
368                        &ssm.in_proj_qkvz,
369                        config.ssm_qkvz_size(),
370                        h,
371                        gpu,
372                        absmax_k,
373                        quantize_k,
374                        stream,
375                    )?;
376                    layers.push(Box::new(Qwen3SsmLayer::new(
377                        input_norm,
378                        ssm,
379                        post_attn_norm,
380                        ffn,
381                        Some(qkvz_nvfp4),
382                        config,
383                        gpu,
384                    )?));
385                }
386                LayerType::SlidingAttention => {
387                    unreachable!("unexpected SlidingAttention in this loader")
388                }
389                LayerType::Moe => unreachable!("Qwen3 has no standalone MoE layers"),
390                // GLM-5.3's `deepseek_sparse_attention`: a full-rank mixer whose visible key set
391                // is chosen at runtime by an indexer. Hard error, not a silent fallthrough into
392                // the dense-attention arm -- that would attend over the WHOLE cache and look right.
393                LayerType::SparseAttention => anyhow::bail!(
394                    "layer {i}: SparseAttention needs a DSA indexer and per-query top-k; Qwen3 has neither"
395                ),
396            }
397
398            if (i + 1) % 12 == 0 {
399                tracing::info!("Loaded layers 0..{}", i + 1);
400            }
401        }
402
403        tracing::info!(
404            "Weight loader: {} layers ({} attention, {} SSM)",
405            layers.len(),
406            attn_idx,
407            layers.len() - attn_idx,
408        );
409
410        Ok(layers)
411    }
412
413    fn load_embedding(
414        &self,
415        store: &WeightStore,
416        _config: &ModelConfig,
417        _gpu: &dyn GpuBackend,
418    ) -> Result<DenseWeight> {
419        dense(store, "model.embed_tokens.weight")
420    }
421
422    fn load_final_norm(
423        &self,
424        store: &WeightStore,
425        _config: &ModelConfig,
426        _gpu: &dyn GpuBackend,
427    ) -> Result<DenseWeight> {
428        dense(store, "model.norm.weight")
429    }
430
431    fn load_lm_head(
432        &self,
433        store: &WeightStore,
434        _config: &ModelConfig,
435        _gpu: &dyn GpuBackend,
436    ) -> Result<DenseWeight> {
437        if store.contains("lm_head.weight") {
438            dense(store, "lm_head.weight")
439        } else {
440            dense(store, "model.embed_tokens.weight")
441        }
442    }
443
444    fn load_mtp_weights(
445        &self,
446        store: &WeightStore,
447        config: &ModelConfig,
448        gpu: &dyn GpuBackend,
449    ) -> Result<Option<MtpWeights>> {
450        if !store.contains("mtp.fc.weight") {
451            tracing::info!("No MTP weights found — speculative decoding disabled");
452            return Ok(None);
453        }
454        let variant = detect_nvfp4_variant(store, config);
455        tracing::info!("Loading MTP weights (variant={:?})...", variant);
456        let mtp = load_mtp(store, config.num_experts, gpu, variant)?;
457        tracing::info!(
458            "MTP weights loaded: fc=[2048,4096], {} experts, attn layer",
459            mtp.experts.len(),
460        );
461        Ok(Some(mtp))
462    }
463}
spark_model/weight_loader/qwen3.rs

spark_model/weight_loader/
qwen3.rs
