spark_model/weight_loader/dflash_loader/
config.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Drafter HF `config.json` schema: [`DflashConfig`] plus the nested
4//! [`DflashRopeScaling`] / [`DflashSubConfig`] blocks and their serde
5//! defaults. Split out of `dflash_loader.rs` to stay under the 500-LoC
6//! cap; field semantics are documented on the types themselves.
7
8use serde::Deserialize;
9
10/// Drafter HF `config.json` (subset Atlas consumes). Mirrors
11/// `z-lab/Qwen3.6-35B-A3B-DFlash/config.json` field names verbatim so
12/// `serde_json::from_str` works directly on the raw file.
13#[derive(Debug, Clone, Deserialize)]
14pub struct DflashConfig {
15    pub hidden_size: usize,
16    pub num_hidden_layers: usize,
17    pub intermediate_size: usize,
18    pub num_attention_heads: usize,
19    pub num_key_value_heads: usize,
20    pub head_dim: usize,
21    pub vocab_size: usize,
22    #[serde(default)]
23    pub draft_vocab_size: Option<usize>,
24    #[serde(default)]
25    pub tie_word_embeddings: bool,
26    /// Block size γ. Qwen3.6-DFlash ships `block_size: 16`.
27    #[serde(default = "default_block_size")]
28    pub block_size: usize,
29    /// DFlash-specific nested config object.
30    #[serde(default)]
31    pub dflash_config: Option<DflashSubConfig>,
32    /// Drafter base RoPE θ. Defaults to 10M (matches Qwen3.6-DFlash).
33    #[serde(default = "default_rope_theta")]
34    pub rope_theta: f32,
35    /// HF-style `rope_scaling` block. `None` ⇒ plain RoPE (the v2 2026-04-27
36    /// Qwen3.6-DFlash drafter ships `rope_scaling: null`). When present and
37    /// `rope_type == "yarn"`, the drafter's YaRN parameters are used to
38    /// build the inv_freq table at construction time.
39    /// `alias = "rope_parameters"`: newer transformers releases (RadixArk
40    /// DSpark, incoai DFlash2) ship the block under that key — without the
41    /// alias Atlas silently drops the scaling (the RadixArk config.json had
42    /// to be hand-patched before this).
43    #[serde(default, alias = "rope_parameters")]
44    pub rope_scaling: Option<DflashRopeScaling>,
45    /// DSpark Markov head rank. `0` (default) ⇒ plain DFlash drafter with no
46    /// Markov head. RadixArk `Qwen3.8-27B-DSpark` ships `markov_rank: 256`
47    /// top-level (SpecForge `DSparkConfig` convention — see the checkpoint's
48    /// `dspark.py`: DSpark fields are declared as top-level config attrs).
49    #[serde(default)]
50    pub markov_rank: usize,
51    /// DSpark Markov head flavor. Only `"vanilla"` (low-rank learned bigram
52    /// bias) is defined by the reference; anything else is rejected at load.
53    #[serde(default)]
54    pub markov_head_type: Option<String>,
55    /// DSpark confidence head (`AcceptRatePredictor`): a `Linear(input, 1)`
56    /// predicting per-draft-position acceptance probability, used for
57    /// adaptive block length. Loaded when present; consumed by the dynamic-K
58    /// scheduling phase (the Markov fixup works without it).
59    #[serde(default)]
60    pub enable_confidence_head: bool,
61    /// When true the confidence head's input is `[hidden ‖ markov_embed]`
62    /// (input_dim = hidden_size + markov_rank); when false, hidden only.
63    #[serde(default = "default_true")]
64    pub confidence_head_with_markov: bool,
65}
66
67fn default_true() -> bool {
68    true
69}
70
71fn default_rope_theta() -> f32 {
72    10_000_000.0
73}
74
75/// Subset of HF `rope_scaling` block consumed by Atlas. Mirrors the field
76/// names in `transformers`' Qwen3 config so `serde_json::from_str` works
77/// directly on the drafter's `config.json`.
78#[derive(Debug, Clone, Deserialize)]
79pub struct DflashRopeScaling {
80    /// Currently only `"yarn"` is recognised; anything else falls back to
81    /// plain RoPE with a warning logged at construction time.
82    #[serde(default)]
83    pub rope_type: Option<String>,
84    #[serde(default)]
85    pub factor: Option<f32>,
86    #[serde(default)]
87    pub beta_fast: Option<f32>,
88    #[serde(default)]
89    pub beta_slow: Option<f32>,
90    #[serde(default)]
91    pub original_max_position_embeddings: Option<f32>,
92}
93
94fn default_block_size() -> usize {
95    16
96}
97
98/// Nested `dflash_config` block in the drafter's `config.json`.
99#[derive(Debug, Clone, Deserialize)]
100pub struct DflashSubConfig {
101    /// Token id used to fill the γ "to-be-predicted" positions during draft
102    /// inference. `248070` for Qwen3.6-DFlash.
103    pub mask_token_id: u32,
104    /// Target-model layer indices to capture intermediate hidden states from.
105    /// `[1, 10, 19, 28, 37]` for Qwen3.6-35B-A3B-DFlash. Order matters:
106    /// shallow-to-deep concatenation is what `fc` expects.
107    pub target_layer_ids: Vec<usize>,
108    /// Draft flavor tag. `"dspark"` marks a SpecForge-lineage drafter, whose
109    /// row convention is SHIFTED (row j's output = token at position j+1; the
110    /// anchor row's output is draft #1) versus the z-lab convention Atlas's
111    /// forward was built for (row j predicts at j, row 0 = echo). The runtime
112    /// keys the draft-vector rotation off this tag. DFlash2 checkpoints have
113    /// no tag and keep the z-lab convention (verified against z-lab
114    /// `dflash/model.py::dflash_generate` — anchor-row output discarded,
115    /// mask rows fill in place).
116    #[serde(default)]
117    pub projector_type: Option<String>,
118
119    // ── DFlash2 fields (incoai/z-lab `DFlash2DraftModel`); absent = DFlash1 ──
120    /// Two-tap dynamic conv kernel size (2 for DFlash2).
121    #[serde(default)]
122    pub conv_kernel_size: usize,
123    /// Channels per conv group (16 for DFlash2 → hidden/16 groups).
124    #[serde(default)]
125    pub conv_group_size: usize,
126    /// Selector codebook rank (256 for DFlash2).
127    #[serde(default)]
128    pub selector_rank: usize,
129    /// Candidates kept per position for the selector walk (16 for DFlash2).
130    #[serde(default)]
131    pub selector_top_k: usize,
132    /// Block size γ as the drafter was TRAINED, when the checkpoint states it
133    /// here. DFlash2 ships `block_size: 8` inside `dflash_config`; DFlash1
134    /// ships it top-level only. `None` = not stated, fall back to the
135    /// top-level field. Read through [`DflashConfig::effective_block_size`].
136    #[serde(default)]
137    pub block_size: Option<usize>,
138}
139
140impl DflashConfig {
141    /// Resolved block size γ: the drafter's own trained value when the
142    /// checkpoint states it, else the top-level field.
143    ///
144    /// The top-level `block_size` defaults to 16, and serde fills that default
145    /// happily for a checkpoint that never mentioned it — so a DFlash2 drafter
146    /// trained at 8 comes up as 16 unless the sub-config is consulted first.
147    /// That is not a cosmetic mismatch: the serve then runs num_drafts=15
148    /// against an 8-block drafter, which measured 0% accept on EVERY verify
149    /// step, and sizes the drafter's per-sequence pools for twice the block it
150    /// will ever use. `--dflash-gamma` still overrides both.
151    pub fn effective_block_size(&self) -> usize {
152        self.dflash_config
153            .as_ref()
154            .and_then(|c| c.block_size)
155            .unwrap_or(self.block_size)
156    }
157}
spark_model/weight_loader/dflash_loader/config.rs

spark_model/weight_loader/dflash_loader/
config.rs
