spark_model/weight_loader/dflash_loader.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! DFlash drafter weight loader.
4//!
5//! Loads `z-lab/Qwen3.6-{27B,35B-A3B}-DFlash`-style drafter checkpoints into
6//! the typed [`DflashWeights`] structure consumed by
7//! [`crate::layers::BlockDiffusionDraftHead`]. The drafter is a small
8//! Qwen3-architecture transformer (8 layers, hidden=2048, GQA 32:4) with
9//! these distinctive parts vs. a vanilla Qwen3:
10//!
11//! * `model.fc` — `[len(target_layer_ids) * target_hidden, draft_hidden]`
12//! BF16 projection that maps the stack of captured target hidden states
13//! into the drafter's input space.
14//! * `model.hidden_norm` — RMSNorm applied to the projected target context
15//! before mixing with token embeddings.
16//! * `lm_head` — drafter ships its own (NOT tied to target's), so
17//! `tie_word_embeddings=false`.
18//! * Optional `d2t` — draft-vocab → target-vocab id remap (absent when
19//! drafter shares vocab with target, as in Qwen3.6-35B-A3B-DFlash where
20//! both = 248320).
21//! * Special `mask_token_id` (`248070` for Qwen3.6-DFlash) used for the γ
22//! "to-be-predicted" positions in block diffusion.
23//!
24//! Under TP the drafter is **not sharded** — it's small (~1–2 GB BF16),
25//! every rank loads the full set. Mirrors the existing MTP-under-TP pattern
26//! (`MTP loads ALL experts on every rank — no EP all_reduce needed`).
27
28use anyhow::{Context, Result};
29use spark_runtime::gpu::GpuBackend;
30use spark_runtime::weights::WeightStore;
31
32use crate::weight_map::{DenseWeight, dense};
33
34mod config;
35pub use config::*;
36
37/// Raw weight bundle for the DFlash drafter, post-load.
38///
39/// Verified against `z-lab/Qwen3.6-35B-A3B-DFlash` (commit 42d3b34, May 2026):
40/// the checkpoint ships 91 BF16 tensors — `fc.weight`, `hidden_norm.weight`,
41/// `norm.weight`, plus 11 weights per drafter layer × 8 layers. **No
42/// `embed_tokens` or `lm_head` are in the checkpoint** — the drafter shares
43/// the target's embedding and LM head at construction time. This matches the
44/// vLLM PR #40898 flow: when those keys are absent, vLLM's `AutoWeightsLoader`
45/// adds them to `skip_substrs`, leaving the runtime to slot in the target's
46/// pointers.
47#[allow(dead_code)]
48pub struct DflashWeights {
49 pub config: DflashConfig,
50
51 /// `[draft_hidden, len(target_layer_ids) * target_hidden]`.
52 /// Qwen3.6-35B-A3B-DFlash: `[2048, 10240]`.
53 pub fc: DenseWeight,
54 /// `[draft_hidden]` — RMSNorm applied to the projected target context
55 /// before mixing with token embeddings.
56 pub hidden_norm: DenseWeight,
57 /// `[draft_hidden]` — final RMSNorm before LM head.
58 pub norm: DenseWeight,
59
60 pub layers: Vec<DflashLayerWeights>,
61
62 /// Present iff the drafter has a draft-id → target-id mapping (i.e.
63 /// `draft_vocab_size != target_vocab_size`). Absent for
64 /// Qwen3.6-35B-A3B-DFlash (both vocabs = 248320).
65 pub draft_id_to_target_id: Option<Vec<i64>>,
66
67 // ── DSpark heads (optional; RadixArk Qwen3.8-27B-DSpark ships all 4) ──
68 /// Markov head `markov_w1`: `[vocab, markov_rank]` BF16 embedding table
69 /// (prev-token → latent). Present iff `config.markov_rank > 0` and the
70 /// checkpoint carries the tensor.
71 pub markov_w1: Option<DenseWeight>,
72 /// Markov head `markov_w2`: `[vocab, markov_rank]` BF16
73 /// (`nn.Linear(rank, vocab, bias=False).weight`, i.e. `[N, K]` for the
74 /// GEMV convention). Projects the latent back to a full-vocab bias.
75 pub markov_w2: Option<DenseWeight>,
76 /// Confidence head weight: `[1, hidden(+rank)]` BF16.
77 pub confidence_proj: Option<DenseWeight>,
78 /// Confidence head bias: `[1]` BF16.
79 pub confidence_bias: Option<DenseWeight>,
80
81 // ── DFlash2 candidate selector (None on DFlash1/DSpark drafters) ──
82 /// `candidate_selector.predecessor_codebook` `[vocab, selector_rank]` BF16.
83 pub selector_pred: Option<DenseWeight>,
84 /// `candidate_selector.successor_codebook` `[vocab, selector_rank]` BF16.
85 pub selector_succ: Option<DenseWeight>,
86 /// `candidate_selector.hidden_projection.weight` `[selector_rank, hidden]`.
87 pub selector_hidden_proj: Option<DenseWeight>,
88}
89
90/// Per-drafter-layer raw weights (BF16). Same shape across all 8 layers.
91#[allow(dead_code)]
92pub struct DflashLayerWeights {
93 pub input_layernorm: DenseWeight,
94 pub post_attention_layernorm: DenseWeight,
95 pub q_proj: DenseWeight,
96 pub k_proj: DenseWeight,
97 pub v_proj: DenseWeight,
98 pub o_proj: DenseWeight,
99 pub q_norm: DenseWeight,
100 pub k_norm: DenseWeight,
101 pub gate_proj: DenseWeight,
102 pub up_proj: DenseWeight,
103 pub down_proj: DenseWeight,
104
105 // ── DFlash2 grouped dynamic causal convs (None on DFlash1 drafters) ──
106 /// `attention_conv.base_kernel` `[2, kernel_size, hidden]` BF16 —
107 /// static tap weights (index 0 = prepare/pre-sublayer, 1 = finish/post).
108 pub attention_conv_base: Option<DenseWeight>,
109 /// `attention_conv.kernel_projection.weight`
110 /// `[2 * kernel_size * groups, hidden]` BF16 — dynamic tap generator.
111 pub attention_conv_proj: Option<DenseWeight>,
112 /// `mlp_conv.base_kernel`, same shape as attention_conv_base.
113 pub mlp_conv_base: Option<DenseWeight>,
114 /// `mlp_conv.kernel_projection.weight`, same shape as attention_conv_proj.
115 pub mlp_conv_proj: Option<DenseWeight>,
116}
117
118/// Probe a [`WeightStore`] for the presence of DFlash drafter weights.
119/// Returns true if the store contains the unique `fc.weight` tensor that
120/// DFlash drafters ship — a lightweight detection that doesn't load any
121/// data. Both bare-key and `model.`-prefixed layouts are accepted; the
122/// canonical `z-lab/Qwen3.6-{27B,35B-A3B}-DFlash` checkpoints ship the
123/// bare layout (verified against commit 42d3b34, May 2026).
124pub fn store_has_dflash_weights(store: &WeightStore) -> bool {
125 store.contains("fc.weight") || store.contains("model.fc.weight")
126}
127
128/// Parse a DFlash drafter's `config.json` into a [`DflashConfig`]. Used by
129/// `main.rs` after fetching the drafter's HF metadata to size the runtime
130/// `BlockDiffusionDraftHead` (layer count, head_dim, vocab_size, the
131/// `target_layer_ids` capture indices).
132pub fn parse_dflash_config(json: &str) -> Result<DflashConfig> {
133 serde_json::from_str(json).context("Parsing DFlash drafter config.json")
134}
135
136/// Load DFlash drafter weights from a separate [`WeightStore`] pointing at
137/// the drafter checkpoint.
138///
139/// The drafter ships its weights at the **root** of the safetensors file
140/// (no `model.` prefix), in the same naming convention as a vanilla Qwen3
141/// transformer minus `embed_tokens` and `lm_head`. Atlas's runtime fills
142/// those two from the *target* model's embedding / LM head at construction
143/// time — exactly mirroring vLLM's "absent in checkpoint → skip_substrs →
144/// share with parent" flow.
145///
146/// The probed key list (verified against `z-lab/Qwen3.6-35B-A3B-DFlash`):
147///
148/// ```text
149/// fc.weight [H, 5*H_target]
150/// hidden_norm.weight [H]
151/// norm.weight [H]
152/// layers.{0..L-1}.input_layernorm.weight [H]
153/// layers.{0..L-1}.post_attention_layernorm.weight [H]
154/// layers.{0..L-1}.self_attn.q_proj.weight [Q*Hd, H]
155/// layers.{0..L-1}.self_attn.k_proj.weight [Kv*Hd, H]
156/// layers.{0..L-1}.self_attn.v_proj.weight [Kv*Hd, H]
157/// layers.{0..L-1}.self_attn.o_proj.weight [H, Q*Hd]
158/// layers.{0..L-1}.self_attn.q_norm.weight [Hd]
159/// layers.{0..L-1}.self_attn.k_norm.weight [Hd]
160/// layers.{0..L-1}.mlp.gate_proj.weight [I, H]
161/// layers.{0..L-1}.mlp.up_proj.weight [I, H]
162/// layers.{0..L-1}.mlp.down_proj.weight [H, I]
163/// ```
164///
165/// where `H=2048`, `H_target=2048`, `Q=32`, `Kv=4`, `Hd=128`, `I=6144`,
166/// `L=8` for Qwen3.6-35B-A3B-DFlash.
167///
168/// Under TP the drafter is replicated, not sharded — `tp_size>1` produces
169/// the same per-rank result as `tp_size=1`. Memory cost: ~948 MB BF16
170/// per rank, trivially below the 119 GB GB10 budget.
171pub fn load_dflash_weights(
172 drafter_store: &WeightStore,
173 drafter_config: &DflashConfig,
174 _gpu: &dyn GpuBackend,
175 _tp_size: usize,
176) -> Result<Option<DflashWeights>> {
177 if !store_has_dflash_weights(drafter_store) {
178 tracing::debug!("DFlash drafter store has no `fc.weight` — skipping");
179 return Ok(None);
180 }
181
182 // Detect bare vs. `model.`-prefixed layout. `z-lab` checkpoints use
183 // bare; we accept either to be robust against a hypothetical re-upload
184 // that uses the prefixed layout.
185 let prefix = if drafter_store.contains("model.fc.weight") {
186 "model."
187 } else {
188 ""
189 };
190
191 let fc = dense(drafter_store, &format!("{prefix}fc.weight"))
192 .context("DFlash drafter: load fc.weight")?;
193 let hidden_norm = dense(drafter_store, &format!("{prefix}hidden_norm.weight"))
194 .context("DFlash drafter: load hidden_norm.weight")?;
195 let norm = dense(drafter_store, &format!("{prefix}norm.weight"))
196 .context("DFlash drafter: load norm.weight")?;
197
198 // DFlash2 detection: conv/selector dims declared in dflash_config AND the
199 // layer-0 conv tensor present. All four families load per layer or none.
200 let dflash2_conv = drafter_config
201 .dflash_config
202 .as_ref()
203 .map(|c| c.conv_kernel_size > 0 && c.conv_group_size > 0)
204 .unwrap_or(false)
205 && drafter_store.contains(&format!("{prefix}layers.0.attention_conv.base_kernel"));
206
207 let layer_count = drafter_config.num_hidden_layers;
208 let mut layers = Vec::with_capacity(layer_count);
209 for i in 0..layer_count {
210 let lp = format!("{prefix}layers.{i}");
211 let (attention_conv_base, attention_conv_proj, mlp_conv_base, mlp_conv_proj) =
212 if dflash2_conv {
213 (
214 Some(dense(
215 drafter_store,
216 &format!("{lp}.attention_conv.base_kernel"),
217 )?),
218 Some(dense(
219 drafter_store,
220 &format!("{lp}.attention_conv.kernel_projection.weight"),
221 )?),
222 Some(dense(drafter_store, &format!("{lp}.mlp_conv.base_kernel"))?),
223 Some(dense(
224 drafter_store,
225 &format!("{lp}.mlp_conv.kernel_projection.weight"),
226 )?),
227 )
228 } else {
229 (None, None, None, None)
230 };
231 let layer = DflashLayerWeights {
232 input_layernorm: dense(drafter_store, &format!("{lp}.input_layernorm.weight"))?,
233 post_attention_layernorm: dense(
234 drafter_store,
235 &format!("{lp}.post_attention_layernorm.weight"),
236 )?,
237 q_proj: dense(drafter_store, &format!("{lp}.self_attn.q_proj.weight"))?,
238 k_proj: dense(drafter_store, &format!("{lp}.self_attn.k_proj.weight"))?,
239 v_proj: dense(drafter_store, &format!("{lp}.self_attn.v_proj.weight"))?,
240 o_proj: dense(drafter_store, &format!("{lp}.self_attn.o_proj.weight"))?,
241 q_norm: dense(drafter_store, &format!("{lp}.self_attn.q_norm.weight"))?,
242 k_norm: dense(drafter_store, &format!("{lp}.self_attn.k_norm.weight"))?,
243 gate_proj: dense(drafter_store, &format!("{lp}.mlp.gate_proj.weight"))?,
244 up_proj: dense(drafter_store, &format!("{lp}.mlp.up_proj.weight"))?,
245 down_proj: dense(drafter_store, &format!("{lp}.mlp.down_proj.weight"))?,
246 attention_conv_base,
247 attention_conv_proj,
248 mlp_conv_base,
249 mlp_conv_proj,
250 };
251 layers.push(layer);
252 }
253
254 // DFlash2 candidate selector. NOTE: the two codebooks ship with NO
255 // `.weight` suffix (raw nn.Parameter-style keys — the z-lab loader
256 // key-maps them; verified against the incoai safetensors header).
257 let selector_key = format!("{prefix}candidate_selector.predecessor_codebook");
258 let (selector_pred, selector_succ, selector_hidden_proj) = if drafter_config
259 .dflash_config
260 .as_ref()
261 .map(|c| c.selector_rank > 0 && c.selector_top_k > 0)
262 .unwrap_or(false)
263 && drafter_store.contains(&selector_key)
264 {
265 (
266 Some(
267 dense(drafter_store, &selector_key)
268 .context("DFlash2: load candidate_selector.predecessor_codebook")?,
269 ),
270 Some(
271 dense(
272 drafter_store,
273 &format!("{prefix}candidate_selector.successor_codebook"),
274 )
275 .context("DFlash2: load candidate_selector.successor_codebook")?,
276 ),
277 Some(
278 dense(
279 drafter_store,
280 &format!("{prefix}candidate_selector.hidden_projection.weight"),
281 )
282 .context("DFlash2: load candidate_selector.hidden_projection.weight")?,
283 ),
284 )
285 } else {
286 (None, None, None)
287 };
288 if dflash2_conv || selector_pred.is_some() {
289 tracing::info!(
290 "DFlash2 heads loaded: convs={} (k={}, group={}), selector={} (rank={}, top_k={})",
291 dflash2_conv,
292 drafter_config
293 .dflash_config
294 .as_ref()
295 .map(|c| c.conv_kernel_size)
296 .unwrap_or(0),
297 drafter_config
298 .dflash_config
299 .as_ref()
300 .map(|c| c.conv_group_size)
301 .unwrap_or(0),
302 selector_pred.is_some(),
303 drafter_config
304 .dflash_config
305 .as_ref()
306 .map(|c| c.selector_rank)
307 .unwrap_or(0),
308 drafter_config
309 .dflash_config
310 .as_ref()
311 .map(|c| c.selector_top_k)
312 .unwrap_or(0),
313 );
314 }
315
316 // `d2t` (draft-id → target-id) is absent from Qwen3.6-DFlash because
317 // both vocabs are 248320. If a future drafter ships a smaller vocab
318 // (vLLM supports this via `draft_vocab_size`), the int64 mapping table
319 // would land here. Probing first to keep this loader compatible.
320 let draft_id_to_target_id = if drafter_store.contains(&format!("{prefix}d2t"))
321 || drafter_store.contains(&format!("{prefix}draft_id_to_target_id"))
322 {
323 // Mapping is loaded into device memory by upstream paths — for now
324 // we just record presence. Phase 2.5 will copy it to a host Vec<i64>
325 // when the head needs it for logit remapping.
326 tracing::warn!(
327 "DFlash drafter has draft-id→target-id mapping; remapping path is not yet wired (Phase 2.5 follow-up)"
328 );
329 Some(Vec::new())
330 } else {
331 None
332 };
333
334 // ── DSpark heads (optional) ──────────────────────────────────────
335 // Tensor names verified against RadixArk/Qwen3.8-27B-DSpark
336 // (model.safetensors, 62 tensors): `markov_head.markov_w1.weight`
337 // [248320, 256], `markov_head.markov_w2.weight` [248320, 256],
338 // `confidence_head.proj.weight` [1, 5376], `confidence_head.proj.bias`
339 // [1] — all BF16, bare layout (same prefix convention as fc.weight).
340 let markov_key = format!("{prefix}markov_head.markov_w1.weight");
341 let (markov_w1, markov_w2) =
342 if drafter_config.markov_rank > 0 && drafter_store.contains(&markov_key) {
343 if let Some(kind) = drafter_config.markov_head_type.as_deref()
344 && kind != "vanilla"
345 {
346 anyhow::bail!(
347 "DSpark drafter declares markov_head_type={kind:?}; only \"vanilla\" \
348 (low-rank bigram bias) is defined by the reference implementation"
349 );
350 }
351 let w1 = dense(drafter_store, &markov_key)
352 .context("DSpark drafter: load markov_head.markov_w1.weight")?;
353 let w2 = dense(
354 drafter_store,
355 &format!("{prefix}markov_head.markov_w2.weight"),
356 )
357 .context("DSpark drafter: load markov_head.markov_w2.weight")?;
358 (Some(w1), Some(w2))
359 } else {
360 if drafter_config.markov_rank > 0 {
361 tracing::warn!(
362 "DSpark drafter config declares markov_rank={} but the checkpoint has \
363 no {markov_key} — running as plain DFlash (Markov bias disabled)",
364 drafter_config.markov_rank,
365 );
366 }
367 (None, None)
368 };
369 let conf_key = format!("{prefix}confidence_head.proj.weight");
370 let (confidence_proj, confidence_bias) =
371 if drafter_config.enable_confidence_head && drafter_store.contains(&conf_key) {
372 let w = dense(drafter_store, &conf_key)
373 .context("DSpark drafter: load confidence_head.proj.weight")?;
374 let b = dense(drafter_store, &format!("{prefix}confidence_head.proj.bias"))
375 .context("DSpark drafter: load confidence_head.proj.bias")?;
376 (Some(w), Some(b))
377 } else {
378 (None, None)
379 };
380 if markov_w1.is_some() || confidence_proj.is_some() {
381 tracing::info!(
382 "DSpark heads loaded: markov={} (rank={}), confidence={} (with_markov={})",
383 markov_w1.is_some(),
384 drafter_config.markov_rank,
385 confidence_proj.is_some(),
386 drafter_config.confidence_head_with_markov,
387 );
388 }
389
390 tracing::info!(
391 "DFlash drafter loaded: {} layers, hidden={}, vocab={}, γ={}, target_layers={:?}",
392 layers.len(),
393 drafter_config.hidden_size,
394 drafter_config.vocab_size,
395 drafter_config.block_size,
396 drafter_config
397 .dflash_config
398 .as_ref()
399 .map(|c| c.target_layer_ids.as_slice())
400 .unwrap_or(&[]),
401 );
402
403 Ok(Some(DflashWeights {
404 config: drafter_config.clone(),
405 fc,
406 hidden_norm,
407 norm,
408 layers,
409 draft_id_to_target_id,
410 markov_w1,
411 markov_w2,
412 confidence_proj,
413 confidence_bias,
414 selector_pred,
415 selector_succ,
416 selector_hidden_proj,
417 }))
418}
419
420#[cfg(test)]
421#[path = "dflash_loader/loader_tests.rs"]
422mod loader_tests;