metrale_model_arch/weight_loader/dflash_loader/
config.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The drafter `config.json` schema: [`DflashConfig`] with the
4//! nested [`DflashRopeScaling`] and [`DflashSubConfig`] blocks, and their
5//! serde defaults.
6//!
7//! Owner: model-arch weight loader.
8//! Invariants: none beyond the types.
9
10use serde::Deserialize;
11
12/// 2026-09-25: The fields of a drafter's `config.json` that the engine reads,
13/// under their `config.json` names. Unknown keys are ignored.
14#[derive(Debug, Clone, Deserialize)]
15pub struct DflashConfig {
16    pub hidden_size: usize,
17    pub num_hidden_layers: usize,
18    pub intermediate_size: usize,
19    pub num_attention_heads: usize,
20    pub num_key_value_heads: usize,
21    pub head_dim: usize,
22    pub vocab_size: usize,
23    #[serde(default)]
24    pub draft_vocab_size: Option<usize>,
25    #[serde(default)]
26    pub tie_word_embeddings: bool,
27    /// 2026-09-25: Top-level block size γ; 16 when absent. Read it through
28    /// [`DflashConfig::effective_block_size`].
29    #[serde(default = "default_block_size")]
30    pub block_size: usize,
31    #[serde(default)]
32    pub dflash_config: Option<DflashSubConfig>,
33    /// 2026-09-25: Drafter RoPE θ; 10,000,000 when absent.
34    #[serde(default = "default_rope_theta")]
35    pub rope_theta: f32,
36    /// 2026-09-25: The `rope_scaling` block, also read under the key
37    /// `rope_parameters`. `None` means plain RoPE. With `rope_type == "yarn"`
38    /// the head builds a YaRN inv_freq table; any other block falls back to
39    /// plain RoPE with a warning.
40    #[serde(default, alias = "rope_parameters")]
41    pub rope_scaling: Option<DflashRopeScaling>,
42    /// 2026-09-25: DSpark Markov head rank, a top-level key; 0 when absent,
43    /// which loads no Markov head.
44    #[serde(default)]
45    pub markov_rank: usize,
46    /// 2026-09-25: DSpark Markov head type. When the Markov head loads, any
47    /// value other than `"vanilla"` fails the load.
48    #[serde(default)]
49    pub markov_head_type: Option<String>,
50    /// 2026-09-25: Load the DSpark confidence head (`confidence_head.proj.*`)
51    /// when the checkpoint has it; false when absent.
52    #[serde(default)]
53    pub enable_confidence_head: bool,
54    /// 2026-09-25: When true the confidence head reads the hidden state and
55    /// then the Markov embedding (`hidden_size + markov_rank` inputs); when
56    /// false, the hidden state only. True when absent.
57    #[serde(default = "default_true")]
58    pub confidence_head_with_markov: bool,
59}
60
61fn default_true() -> bool {
62    true
63}
64
65fn default_rope_theta() -> f32 {
66    10_000_000.0
67}
68
69/// 2026-09-25: The `rope_scaling` fields the drafter head reads.
70#[derive(Debug, Clone, Deserialize)]
71pub struct DflashRopeScaling {
72    /// 2026-09-25: Only `"yarn"` is recognised; any other value, or none, falls
73    /// back to plain RoPE with a warning when the head is built.
74    #[serde(default)]
75    pub rope_type: Option<String>,
76    #[serde(default)]
77    pub factor: Option<f32>,
78    #[serde(default)]
79    pub beta_fast: Option<f32>,
80    #[serde(default)]
81    pub beta_slow: Option<f32>,
82    #[serde(default)]
83    pub original_max_position_embeddings: Option<f32>,
84}
85
86fn default_block_size() -> usize {
87    16
88}
89
90/// 2026-09-25: The nested `dflash_config` block of the drafter's `config.json`.
91#[derive(Debug, Clone, Deserialize)]
92pub struct DflashSubConfig {
93    /// 2026-09-25: Token id that fills the draft block after its anchor token.
94    pub mask_token_id: u32,
95    /// 2026-09-25: Target-model layers whose hidden states the drafter reads.
96    /// The head refuses to build when the list is empty.
97    pub target_layer_ids: Vec<usize>,
98    /// 2026-09-25: `"dspark"` marks a drafter whose row j predicts position
99    /// j + 1; the head then rotates the draft vector right by one, unless
100    /// `METRALE_DSPARK_SHIFT` overrides it.
101    #[serde(default)]
102    pub projector_type: Option<String>,
103
104    /// 2026-09-25: DFlash2 conv kernel size; 0 when absent. With
105    /// `conv_group_size`, a non-zero value lets the conv tensors load.
106    #[serde(default)]
107    pub conv_kernel_size: usize,
108    /// 2026-09-25: DFlash2 channels per conv group; 0 when absent.
109    #[serde(default)]
110    pub conv_group_size: usize,
111    /// 2026-09-25: DFlash2 selector codebook rank; 0 when absent. With
112    /// `selector_top_k`, a non-zero value lets the selector tensors load.
113    #[serde(default)]
114    pub selector_rank: usize,
115    /// 2026-09-25: DFlash2 candidates kept per position; 0 when absent.
116    #[serde(default)]
117    pub selector_top_k: usize,
118    /// 2026-09-25: The drafter's trained block size, when the checkpoint states
119    /// it here. `None` falls back to the top-level field. Read it through
120    /// [`DflashConfig::effective_block_size`].
121    #[serde(default)]
122    pub block_size: Option<usize>,
123}
124
125impl DflashConfig {
126    /// 2026-09-25: Resolved block size γ: `dflash_config.block_size` when set,
127    /// else the top-level field. The sub-config comes first because serde
128    /// fills the top-level default of 16 for a checkpoint that never states
129    /// it.
130    pub fn effective_block_size(&self) -> usize {
131        self.dflash_config
132            .as_ref()
133            .and_then(|c| c.block_size)
134            .unwrap_or(self.block_size)
135    }
136}
metrale_model_arch/weight_loader/dflash_loader/config.rs

metrale_model_arch/weight_loader/dflash_loader/
config.rs
