spark_model/weight_loader/
qwen35.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3pub(crate) mod load_layers;
4
5use anyhow::{Context, Result};
6use atlas_core::config::ModelConfig;
7use spark_runtime::gpu::GpuBackend;
8use spark_runtime::kv_cache::KvCacheDtype;
9use spark_runtime::weights::{WeightDtype, WeightStore};
10
11use super::ModelWeightLoader;
12use crate::layer::TransformerLayer;
13use crate::weight_map::{DenseWeight, MtpWeights, dense_auto, detect_nvfp4_variant, load_mtp};
14
15pub struct Qwen35WeightLoader;
16
17fn vision_dense_auto(
18    store: &WeightStore,
19    prefix: &str,
20    gpu: &dyn GpuBackend,
21) -> Result<DenseWeight> {
22    let name = format!("{prefix}.weight");
23    let w = store.get(&name)?;
24    match w.dtype {
25        WeightDtype::BF16 | WeightDtype::FP8E4M3 => {
26            crate::weight_map::dense_auto_fp8_or_bf16(store, prefix, gpu)
27        }
28        WeightDtype::FP32 => crate::weight_map::dense_f32_safe(store, &name, gpu),
29        other => anyhow::bail!("vision_dense_auto: unsupported dtype {other:?} for {name}"),
30    }
31}
32
33fn vision_tensor_dense_auto(
34    store: &WeightStore,
35    name: &str,
36    gpu: &dyn GpuBackend,
37) -> Result<DenseWeight> {
38    let w = store.get(name)?;
39    match w.dtype {
40        WeightDtype::BF16 => Ok(DenseWeight { weight: w.ptr }),
41        WeightDtype::FP32 => crate::weight_map::dense_f32_safe(store, name, gpu),
42        other => anyhow::bail!("vision_tensor_dense_auto: unsupported dtype {other:?} for {name}"),
43    }
44}
45
46impl ModelWeightLoader for Qwen35WeightLoader {
47    fn supports_tp(&self) -> bool {
48        // FullAttention layers are TP-sharded across all 3 quant paths
49        // (FP8 native, NVFP4-from-disk, BF16 → NVFP4). LinearAttention
50        // (GDN SSM) layers are now TP-sharded head-parallel for the BF16
51        // and NVFP4 paths (GDN HeadParallel): linear_num_key/value_heads
52        // are divided per rank in topology.rs, each rank owns a contiguous
53        // head range, and out_proj is row-parallel with one all-reduce.
54        // Native block-scaled FP8 SSM still requires TP=1 (per-128-row
55        // scale slicing deferred) — build_linear_attention_fp8 errors
56        // clearly when tp_size > 1.
57        true
58    }
59
60    fn load_layers(
61        &self,
62        store: &WeightStore,
63        config: &ModelConfig,
64        gpu: &dyn GpuBackend,
65        layer_kv_dtypes: &[KvCacheDtype],
66    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
67        load_layers::load_layers(self, store, config, gpu, layer_kv_dtypes)
68    }
69
70    fn load_embedding(
71        &self,
72        store: &WeightStore,
73        config: &ModelConfig,
74        gpu: &dyn GpuBackend,
75    ) -> Result<DenseWeight> {
76        let prefix = &config.weight_prefix;
77        dense_auto(store, &format!("{prefix}.embed_tokens.weight"), gpu)
78    }
79
80    fn load_final_norm(
81        &self,
82        store: &WeightStore,
83        config: &ModelConfig,
84        gpu: &dyn GpuBackend,
85    ) -> Result<DenseWeight> {
86        let prefix = &config.weight_prefix;
87        dense_auto(store, &format!("{prefix}.norm.weight"), gpu)
88    }
89
90    fn load_lm_head(
91        &self,
92        store: &WeightStore,
93        config: &ModelConfig,
94        gpu: &dyn GpuBackend,
95    ) -> Result<DenseWeight> {
96        // lm_head location varies by quantizer:
97        //   Sehyo: "lm_head.weight"
98        //   Kbenkhaled: "language_model.lm_head.weight"
99        //
100        // Dequant FP8 ONLY; hand every other dtype through untouched.
101        // `dense` does no dtype check, which is correct for a BF16 head and for
102        // a Standard-NVFP4 head (`weight` U8-packed + weight_scale_2, e.g.
103        // nvidia/Qwen3.6-*-NVFP4) that the consumer unpacks itself — routing
104        // those through `dense_auto_fp8_or_bf16` hard-errors on `unsupported
105        // dtype UInt8`. But it is WRONG for FP8: mixed-precision checkpoints
106        // (unsloth Qwen3.6-*-NVFP4, 2026-07-10) keep lm_head as FP8 E4M3 +
107        // per-row `weight_scale`, and feeding those bytes to a BF16 GEMM reads
108        // 2x the allocation on the largest tensor in the model →
109        // CUDA_ERROR_ILLEGAL_ADDRESS at the first sync after build.
110        for prefix in ["lm_head", "language_model.lm_head", "model.lm_head"] {
111            let key = format!("{prefix}.weight");
112            if !store.contains(&key) {
113                continue;
114            }
115            let is_fp8 = store
116                .get(&key)
117                .map(|w| w.dtype == WeightDtype::FP8E4M3)
118                .unwrap_or(false);
119            return if is_fp8 {
120                crate::weight_map::dense_auto_fp8_or_bf16(store, prefix, gpu)
121            } else {
122                crate::weight_map::dense(store, &key)
123            };
124        }
125        // Tied embeddings: the head IS the embedding table.
126        let prefix = &config.weight_prefix;
127        crate::weight_map::dense(store, &format!("{prefix}.embed_tokens.weight"))
128    }
129
130    fn load_mtp_weights(
131        &self,
132        store: &WeightStore,
133        config: &ModelConfig,
134        gpu: &dyn GpuBackend,
135    ) -> Result<Option<MtpWeights>> {
136        if !store.contains("mtp.fc.weight") {
137            tracing::info!("No MTP weights found — speculative decoding disabled");
138            return Ok(None);
139        }
140        let variant = detect_nvfp4_variant(store, config);
141        tracing::info!(
142            "Loading MTP weights ({} experts, variant={:?})...",
143            config.num_experts,
144            variant
145        );
146        let mtp = load_mtp(store, config.num_experts, gpu, variant)?;
147        tracing::info!(
148            "MTP weights loaded: fc=[2048,4096], {} experts, attn layer",
149            mtp.experts.len(),
150        );
151        Ok(Some(mtp))
152    }
153
154    /// Load the Qwen3.6 ViT tower. Returns `None` when `config.vision` is
155    /// `None` (Qwen3.5 text-only). Otherwise matches the Qwen3-VL shape
156    /// exactly (27 blocks, `model.visual.*` prefix, optional deepstack
157    /// merger list + final merger) but auto-dequants FP8 per-channel
158    /// weights to BF16 for blocks 4+. Blocks 0-3 are exempted in the
159    /// checkpoint's `modules_to_not_convert` list and stay BF16 on disk.
160    fn load_vision_encoder(
161        &self,
162        store: &WeightStore,
163        config: &ModelConfig,
164        gpu: &dyn GpuBackend,
165    ) -> Result<Option<crate::layers::VisionEncoder>> {
166        let vcfg = match &config.vision {
167            Some(v) => v.clone(),
168            None => return Ok(None),
169        };
170        // AEON-7's v2 NVFP4 re-quant (and other multimodal-preserved
171        // checkpoints quantized via AutoModelForImageTextToText) keeps
172        // the canonical nested layout `model.language_model.visual.*`
173        // instead of the flat `model.visual.*` form. Probe the canonical
174        // tensor under both prefixes; first hit wins.
175        let vp = if store.contains("model.visual.patch_embed.proj.weight") {
176            "model.visual"
177        } else if store.contains("model.language_model.visual.patch_embed.proj.weight") {
178            "model.language_model.visual"
179        } else {
180            tracing::warn!(
181                "Vision encoder tensors absent under both `model.visual.*` and \
182                 `model.language_model.visual.*`; skipping vision tower (text-only mode)"
183            );
184            return Ok(None);
185        };
186
187        let patch_embed_w =
188            vision_tensor_dense_auto(store, &format!("{vp}.patch_embed.proj.weight"), gpu)?;
189        let patch_embed_b =
190            vision_tensor_dense_auto(store, &format!("{vp}.patch_embed.proj.bias"), gpu)?;
191        let pos_embed = vision_tensor_dense_auto(store, &format!("{vp}.pos_embed.weight"), gpu)?;
192        let pos_embed_shape = store.get(&format!("{vp}.pos_embed.weight"))?.shape.clone();
193        let num_position_embeddings = pos_embed_shape
194            .first()
195            .copied()
196            .context("pos_embed shape missing rows")?;
197
198        let mut blocks = Vec::with_capacity(vcfg.depth);
199        for i in 0..vcfg.depth {
200            let bp = format!("{vp}.blocks.{i}");
201            blocks.push(crate::layers::ViTBlock {
202                norm1_w: vision_tensor_dense_auto(store, &format!("{bp}.norm1.weight"), gpu)?
203                    .weight,
204                norm1_b: vision_tensor_dense_auto(store, &format!("{bp}.norm1.bias"), gpu)?.weight,
205                qkv_w: vision_dense_auto(store, &format!("{bp}.attn.qkv"), gpu)?.weight,
206                qkv_b: vision_tensor_dense_auto(store, &format!("{bp}.attn.qkv.bias"), gpu)?.weight,
207                proj_w: vision_dense_auto(store, &format!("{bp}.attn.proj"), gpu)?.weight,
208                proj_b: vision_tensor_dense_auto(store, &format!("{bp}.attn.proj.bias"), gpu)?
209                    .weight,
210                norm2_w: vision_tensor_dense_auto(store, &format!("{bp}.norm2.weight"), gpu)?
211                    .weight,
212                norm2_b: vision_tensor_dense_auto(store, &format!("{bp}.norm2.bias"), gpu)?.weight,
213                fc1_w: vision_dense_auto(store, &format!("{bp}.mlp.linear_fc1"), gpu)?.weight,
214                fc1_b: vision_tensor_dense_auto(store, &format!("{bp}.mlp.linear_fc1.bias"), gpu)?
215                    .weight,
216                fc2_w: vision_dense_auto(store, &format!("{bp}.mlp.linear_fc2"), gpu)?.weight,
217                fc2_b: vision_tensor_dense_auto(store, &format!("{bp}.mlp.linear_fc2.bias"), gpu)?
218                    .weight,
219            });
220        }
221
222        let mut deepstack = Vec::with_capacity(vcfg.deepstack_visual_indexes.len());
223        for i in 0..vcfg.deepstack_visual_indexes.len() {
224            let mp = format!("{vp}.deepstack_merger_list.{i}");
225            deepstack.push(crate::layers::MergerLayer {
226                norm_w: vision_tensor_dense_auto(store, &format!("{mp}.norm.weight"), gpu)?.weight,
227                norm_b: vision_tensor_dense_auto(store, &format!("{mp}.norm.bias"), gpu)?.weight,
228                fc1_w: vision_dense_auto(store, &format!("{mp}.linear_fc1"), gpu)?.weight,
229                fc1_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc1.bias"), gpu)?
230                    .weight,
231                fc2_w: vision_dense_auto(store, &format!("{mp}.linear_fc2"), gpu)?.weight,
232                fc2_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc2.bias"), gpu)?
233                    .weight,
234            });
235        }
236
237        let mp = format!("{vp}.merger");
238        let merger = crate::layers::MergerLayer {
239            norm_w: vision_tensor_dense_auto(store, &format!("{mp}.norm.weight"), gpu)?.weight,
240            norm_b: vision_tensor_dense_auto(store, &format!("{mp}.norm.bias"), gpu)?.weight,
241            fc1_w: vision_dense_auto(store, &format!("{mp}.linear_fc1"), gpu)?.weight,
242            fc1_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc1.bias"), gpu)?.weight,
243            fc2_w: vision_dense_auto(store, &format!("{mp}.linear_fc2"), gpu)?.weight,
244            fc2_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc2.bias"), gpu)?.weight,
245        };
246
247        let deepstack_indexes = vcfg.deepstack_visual_indexes.clone();
248        let ve = crate::layers::VisionEncoder::new(
249            patch_embed_w.weight,
250            patch_embed_b.weight,
251            pos_embed.weight,
252            num_position_embeddings,
253            blocks,
254            deepstack,
255            deepstack_indexes,
256            merger,
257            vcfg.hidden_size,
258            vcfg.num_heads,
259            vcfg.spatial_merge_size,
260            vcfg.out_hidden_size,
261            vcfg.intermediate_size,
262            vcfg.patch_size,
263            vcfg.max_pixels,
264            gpu,
265        )?;
266        tracing::info!(
267            "Qwen3.6 vision encoder loaded: depth={}, hidden={}, heads={}, FP8-blocks>=4",
268            vcfg.depth,
269            vcfg.hidden_size,
270            vcfg.num_heads,
271        );
272        Ok(Some(ve))
273    }
274}
spark_model/weight_loader/qwen35.rs

spark_model/weight_loader/
qwen35.rs
