spark_model/weight_loader/
deepseek_v4.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! DeepSeek-V4 weight loader (MLA + MoE).
4//!
5//! Implements full layer loading for DeepSeek-V4-Flash, reusing the same
6//! MLA attention pattern as Mistral Small 4 with DeepSeek weight naming.
7
8mod assemble;
9mod attn_sink;
10mod compute;
11mod csa_ape;
12mod load_layers;
13// MTP draft-module loader for nvidia/DeepSeek-V4-Flash-NVFP4.
14mod mtp;
15pub(crate) use mtp::{DeepseekV4MtpModule, load_v4_mtp_module};
16
17use anyhow::{Context, Result};
18use atlas_core::config::ModelConfig;
19use spark_runtime::gpu::GpuBackend;
20use spark_runtime::kv_cache::KvCacheDtype;
21use spark_runtime::weights::WeightStore;
22
23use super::ModelWeightLoader;
24use crate::layer::TransformerLayer;
25use crate::weight_map::{DenseWeight, MtpWeights, dense, dense_auto};
26
27pub struct DeepSeekV4WeightLoader;
28
29impl ModelWeightLoader for DeepSeekV4WeightLoader {
30    fn supports_tp(&self) -> bool {
31        // DeepSeek-V4 uses num_key_value_heads=1 (MQA), which makes
32        // head-parallel TP sharding impossible — 1 is not divisible by
33        // any tp_size > 1.  Multi-spark deployments MUST use pure EP
34        // (tp-size 1, ep-size 2/4/...) instead.
35        false
36    }
37
38    fn load_layers(
39        &self,
40        store: &WeightStore,
41        config: &ModelConfig,
42        gpu: &dyn GpuBackend,
43        layer_kv_dtypes: &[KvCacheDtype],
44    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
45        load_layers::load_all_layers(store, config, gpu, layer_kv_dtypes)
46    }
47
48    fn load_embedding(
49        &self,
50        store: &WeightStore,
51        _config: &ModelConfig,
52        _gpu: &dyn GpuBackend,
53    ) -> Result<DenseWeight> {
54        // RedHatAI re-quant uses flattened naming; try it first, then standard HF names.
55        if let Ok(w) = dense(store, "embed.weight") {
56            return Ok(w);
57        }
58        if let Ok(w) = dense(store, "model.embed_tokens.weight") {
59            return Ok(w);
60        }
61        dense(store, "embed_tokens.weight")
62            .context("DeepSeek-V4: no embedding tensor found (tried embed.weight, model.embed_tokens.weight, embed_tokens.weight)")
63    }
64
65    fn load_final_norm(
66        &self,
67        store: &WeightStore,
68        _config: &ModelConfig,
69        _gpu: &dyn GpuBackend,
70    ) -> Result<DenseWeight> {
71        // DeepSeek-V4 ships HF-vanilla RMSNorm weights (scale = weight). Load them
72        // EXACTLY; the model dispatches `rms_norm_vanilla` (see
73        // `crate::ships_vanilla_norm_weights`).
74        if let Ok(w) = dense_auto(store, "norm.weight", _gpu) {
75            return Ok(w);
76        }
77        if let Ok(w) = dense_auto(store, "model.norm.weight", _gpu) {
78            return Ok(w);
79        }
80        dense_auto(store, "final_norm.weight", _gpu)
81            .context("DeepSeek-V4: no final norm tensor found (tried norm.weight, model.norm.weight, final_norm.weight)")
82    }
83
84    fn load_lm_head(
85        &self,
86        store: &WeightStore,
87        config: &ModelConfig,
88        _gpu: &dyn GpuBackend,
89    ) -> Result<DenseWeight> {
90        // Try standard HF name first
91        if store.contains("lm_head.weight") {
92            return dense(store, "lm_head.weight");
93        }
94        // RedHatAI / consolidated checkpoints
95        if store.contains("output.weight") {
96            return dense(store, "output.weight");
97        }
98        if store.contains("head.weight") {
99            return dense(store, "head.weight");
100        }
101        // Tied embeddings: either config says so, or no separate head exists
102        if config.tie_word_embeddings
103            || store.contains("embed.weight")
104            || store.contains("model.embed_tokens.weight")
105        {
106            // Tied: reuse the embedding tensor. Inline the same dense lookups as
107            // load_embedding (load_lm_head has no `gpu`, and they don't need it).
108            if let Ok(w) = dense(store, "embed.weight") {
109                return Ok(w);
110            }
111            if let Ok(w) = dense(store, "model.embed_tokens.weight") {
112                return Ok(w);
113            }
114            return dense(store, "embed_tokens.weight")
115                .context("DeepSeek-V4: tied lm_head — no embedding tensor found");
116        }
117        anyhow::bail!(
118            "DeepSeek-V4: lm_head not found (tried lm_head.weight, output.weight, head.weight, and tied embeddings)"
119        )
120    }
121
122    fn load_mtp_weights(
123        &self,
124        _store: &WeightStore,
125        _config: &ModelConfig,
126        _gpu: &dyn GpuBackend,
127    ) -> Result<Option<MtpWeights>> {
128        Ok(None)
129    }
130}
spark_model/weight_loader/deepseek_v4.rs

spark_model/weight_loader/
deepseek_v4.rs
