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