metrale_model_arch/weight_loader/deepseek_v4/
mtp.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Loads the DeepSeek-V4 multi-token-prediction (MTP) draft module
4//! from the `mtp.0.*` tensors.
5//!
6//! The body is a V4 transformer layer built by `assemble_layer` from the
7//! `mtp.0` prefix. The MTP-specific pieces are the input combiner
8//! (`enorm`/`hnorm`, `e_proj`/`h_proj`), the final `norm` and the module's own
9//! HC head. The token embedding and the LM head are not loaded here: the
10//! proposer, `crate::deepseek_v4_mtp::DeepseekV4MtpHead`, takes the target
11//! model's at build time.
12//!
13//! Owner: model-arch weight loader.
14//! Invariants: none beyond the types.
15
16use anyhow::Result;
17use metrale_cache::kv_cache::KvCacheDtype;
18use metrale_config::ModelConfig;
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
20use metrale_model_weights::weights::WeightStore;
21
22use metrale_model_layers::layer::TransformerLayer;
23use metrale_model_layers::layers::qwen3_attention::HcHeadWeights;
24use metrale_model_layers::weight_map::DenseWeight;
25use metrale_model_layers::weight_map::quant_helpers::dense_auto;
26
27/// 2026-09-25: A loaded DeepSeek-V4 MTP draft module: the V4 layer body plus the
28/// MTP input combiner, final norm and HC head.
29#[allow(dead_code)]
30pub struct DeepseekV4MtpModule {
31    /// 2026-09-25: The V4 layer body, built by `assemble_layer` from the `mtp.0` prefix.
32    pub body: Box<dyn TransformerLayer>,
33    /// 2026-09-25: RMSNorm weight applied to the token embedding before `e_proj`.
34    pub enorm: DenseWeight,
35    /// 2026-09-25: RMSNorm weight applied to the target's hidden state before `h_proj`.
36    pub hnorm: DenseWeight,
37    /// 2026-09-25: `[hidden, hidden]` projection of the normed embedding; summed with
38    /// the `h_proj` branch.
39    pub e_proj: DenseWeight,
40    /// 2026-09-25: `[hidden, hidden]` projection of the normed hidden state; summed
41    /// with the `e_proj` branch.
42    pub h_proj: DenseWeight,
43    /// 2026-09-25: Final RMSNorm weight, applied before the target's LM head.
44    pub norm: DenseWeight,
45    /// 2026-09-25: The module's own HC head (`mtp.0.hc_head_*`). The body is built
46    /// with `layer_idx = num_hidden_layers`, so it is not the last model layer
47    /// and never collapses the HC streams; the proposer does that with these
48    /// weights. The body holds a clone of the same pointers. `Some` in every
49    /// returned module: `assemble_layer` refuses `hc_mult == 0`.
50    pub hc_head: Option<HcHeadWeights>,
51}
52
53/// 2026-09-25: Loads the DeepSeek-V4 MTP draft module.
54///
55/// Returns `Ok(None)`, loading nothing, when `num_mtp_modules == 0` or the
56/// store has no `mtp.0.enorm.weight`. Any other missing `mtp.0.*` tensor is an
57/// error.
58pub fn load_v4_mtp_module(
59    store: &WeightStore,
60    config: &ModelConfig,
61    gpu: &dyn GpuBackend,
62    layer_kv_dtypes: &[KvCacheDtype],
63) -> Result<Option<DeepseekV4MtpModule>> {
64    if config.num_mtp_modules == 0 {
65        return Ok(None);
66    }
67    if !store.contains("mtp.0.enorm.weight") {
68        tracing::info!(
69            "DeepSeek-V4: num_mtp_modules={} but no mtp.0.* tensors in checkpoint — MTP disabled",
70            config.num_mtp_modules
71        );
72        return Ok(None);
73    }
74
75    let prefix = "mtp.0";
76    let ap = format!("{prefix}.attn");
77    let null = DenseWeight {
78        weight: DevicePtr::NULL,
79    };
80
81    // 2026-09-25: The body loads the same tensors as the `o_lora_rank > 0` branch
82    // of `load_all_layers`, whatever `o_lora_rank` is, and its absorption views
83    // stay NULL.
84    let input_norm = dense_auto(store, &format!("{prefix}.attn_norm.weight"), gpu)?;
85    let post_attn_norm = dense_auto(store, &format!("{prefix}.ffn_norm.weight"), gpu)?;
86    let wq_a = dense_auto(store, &format!("{ap}.wq_a.weight"), gpu)?;
87    let wq_b = dense_auto(store, &format!("{ap}.wq_b.weight"), gpu)?;
88    let q_a_norm = dense_auto(store, &format!("{ap}.q_norm.weight"), gpu)?;
89    let wkv_a = dense_auto(store, &format!("{ap}.wkv.weight"), gpu)?;
90    let kv_a_norm = dense_auto(store, &format!("{ap}.kv_norm.weight"), gpu)?;
91    let wo_a = dense_auto(store, &format!("{ap}.wo_a.weight"), gpu)?;
92    let wo_b = dense_auto(store, &format!("{ap}.wo_b.weight"), gpu)?;
93
94    // 2026-09-25: The module's own HC head (`mtp.0.hc_head_*`), separate from the
95    // model-level head that `load_all_layers` gives the main layers.
96    let hc_head = if config.hc_mult > 0 {
97        let hc = config.hc_mult;
98        let hc_dim = hc * config.hidden_size;
99        let head_fn = super::assemble::load_hc_f32(
100            store,
101            &[format!("{prefix}.hc_head_fn")],
102            hc * hc_dim,
103            gpu,
104        )?;
105        let head_base =
106            super::assemble::load_hc_f32(store, &[format!("{prefix}.hc_head_base")], hc, gpu)?;
107        let head_scale =
108            super::assemble::load_hc_f32(store, &[format!("{prefix}.hc_head_scale")], 1, gpu)?;
109        Some(HcHeadWeights {
110            hc_fn: head_fn,
111            hc_base: head_base,
112            hc_scale: head_scale,
113            // 2026-09-25: `None` selects the Sinkhorn head mixer.
114            lowrank: None,
115        })
116    } else {
117        None
118    };
119
120    let mut yarn_inv_freq = DevicePtr::NULL;
121    yarn_inv_freq = super::compute::ensure_yarn_inv_freq(&mut yarn_inv_freq, config, gpu)?;
122
123    // 2026-09-25: `layer_idx = num_hidden_layers` makes the body neither the
124    // first nor the last model layer, so it neither expands nor collapses the HC
125    // streams; the proposer does both.
126    let body = super::assemble::assemble_layer(
127        config.num_hidden_layers,
128        prefix,
129        true,
130        input_norm,
131        post_attn_norm,
132        wq_a,
133        None,
134        wq_b,
135        None,
136        q_a_norm,
137        wkv_a,
138        None,
139        null,
140        kv_a_norm,
141        wo_b,
142        None,
143        null,
144        null,
145        null,
146        null,
147        null,
148        null,
149        yarn_inv_freq,
150        wo_a,
151        hc_head.clone(),
152        store,
153        config,
154        gpu,
155        layer_kv_dtypes,
156    )?;
157
158    // 2026-09-25: The norms are loaded unchanged: the proposer applies them with
159    // `rms_norm_vanilla`, which scales by `w`, not `1 + w`. `dense_auto` returns every
160    // tensor here as BF16, dequantizing FP8 or packed NVFP4.
161    let enorm = dense_auto(store, &format!("{prefix}.enorm.weight"), gpu)?;
162    let hnorm = dense_auto(store, &format!("{prefix}.hnorm.weight"), gpu)?;
163    let norm = dense_auto(store, &format!("{prefix}.norm.weight"), gpu)?;
164    let e_proj = dense_auto(store, &format!("{prefix}.e_proj.weight"), gpu)?;
165    let h_proj = dense_auto(store, &format!("{prefix}.h_proj.weight"), gpu)?;
166
167    tracing::info!(
168        "DeepSeek-V4 MTP module loaded: reused V4 body (MLA + mHC + 256-expert MoE) \
169         + combiner (enorm/hnorm + e_proj/h_proj) + final norm"
170    );
171
172    Ok(Some(DeepseekV4MtpModule {
173        body,
174        enorm,
175        hnorm,
176        e_proj,
177        h_proj,
178        norm,
179        hc_head,
180    }))
181}
metrale_model_arch/weight_loader/deepseek_v4/mtp.rs

metrale_model_arch/weight_loader/deepseek_v4/
mtp.rs
