metrale_model_arch/weight_loader/
deepseek_v41.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: DeepSeek-V4.1 loader: assembles
5//! [`DeepSeekV41Layer`](crate::deepseek_v41_layer::DeepSeekV41Layer)s from a
6//! GGUF checkpoint.
7//!
8//! Attention, the compressor and indexer projections, the routers, the shared
9//! experts, the engram projections, the norms, the embedding and the head are
10//! read from the weight store. The projections read through `resident_mat` may
11//! stay raw Q2_K / Q3_K, and a Q6_K head stays raw. The routed expert stacks
12//! must be deferred by the GGUF loader; they are streamed from the shard files
13//! through `expert_stream`, as are the engram rows.
14//!
15//! The loader widens to f32 the tensors the kernels read as f32: the attention
16//! sinks, the `q_norm`/`kv_norm`/compressor/indexer norm weights, the
17//! compressor projections of layers with ratio > 1, and the HC mixes. It forms
18//! the engram `q * k` product on the host.
19//!
20//! The KV and index source layers are the layers whose tensors include
21//! `compressor.wkv.weight` and `indexer.wq_b.weight`. Candidate settings the
22//! config leaves unset fall back to `DEFAULT_CANDIDATE_SOURCE` (after
23//! `METRALE_DS41_CANDIDATE_SOURCE`), `DEFAULT_CANDIDATE_TOPK_BLOCKS` and
24//! `DEFAULT_CANDIDATE_BLOCK`.
25//!
26//! Owner: model-arch weight loader.
27//! Invariants: none beyond the types.
28
29mod engram_q2k;
30pub mod load_layers;
31
32use anyhow::{Context, Result, ensure};
33use metrale_config::ModelConfig;
34use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
35use metrale_model_weights::weights::{WeightDtype, WeightStore};
36
37use crate::weight_loader::ModelWeightLoader;
38use metrale_model_layers::layer::TransformerLayer;
39use metrale_model_layers::layers::ops::ResidentMat;
40use metrale_model_layers::layers::qwen3_attention::HcSiteWeights;
41use metrale_model_layers::weight_map::quant_helpers::dense_auto;
42use metrale_model_layers::weight_map::{DenseWeight, MtpWeights};
43
44pub struct DeepSeekV41WeightLoader;
45
46pub(super) const DEFAULT_CANDIDATE_SOURCE: usize = 20;
47pub(super) const DEFAULT_CANDIDATE_TOPK_BLOCKS: usize = 2048;
48pub(super) const DEFAULT_CANDIDATE_BLOCK: usize = 8;
49
50pub(super) fn env_usize(key: &str, default: usize) -> usize {
51    std::env::var(key)
52        .ok()
53        .and_then(|v| v.parse().ok())
54        .unwrap_or(default)
55}
56
57pub fn bf16_ptr(store: &WeightStore, name: &str) -> Result<DevicePtr> {
58    let t = store.get(name)?;
59    ensure!(
60        t.dtype == WeightDtype::BF16,
61        "{name}: expected bf16, got {:?}",
62        t.dtype
63    );
64    Ok(t.ptr)
65}
66
67/// 2026-09-25: A resident projection as the store holds it: bf16, or raw
68/// Q2_K / Q3_K blocks. Any other dtype is an error.
69pub(super) fn resident_mat(store: &WeightStore, name: &str) -> Result<ResidentMat> {
70    let t = store.get(name)?;
71    match t.dtype {
72        WeightDtype::BF16 => Ok(ResidentMat::Bf16(t.ptr)),
73        WeightDtype::Q2K => Ok(ResidentMat::Q2K(t.ptr)),
74        WeightDtype::Q3K => Ok(ResidentMat::Q3K(t.ptr)),
75        d => anyhow::bail!("{name}: expected bf16, Q2_K or Q3_K, got {d:?}"),
76    }
77}
78
79pub fn download_f32(gpu: &dyn GpuBackend, store: &WeightStore, name: &str) -> Result<Vec<f32>> {
80    let t = store.get(name)?;
81    let n = t.num_elements();
82    match t.dtype {
83        WeightDtype::BF16 => {
84            let mut b = vec![0u8; n * 2];
85            gpu.copy_d2h(t.ptr, &mut b)?;
86            Ok(b.chunks_exact(2)
87                .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
88                .collect())
89        }
90        WeightDtype::FP32 => {
91            let mut b = vec![0u8; n * 4];
92            gpu.copy_d2h(t.ptr, &mut b)?;
93            Ok(b.chunks_exact(4)
94                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
95                .collect())
96        }
97        WeightDtype::Q2K | WeightDtype::Q3K => {
98            use metrale_model_weights::weights::dequant_cpu::{GgmlType, dequant_to_f32};
99            let gt = if t.dtype == WeightDtype::Q2K {
100                GgmlType::Q2K
101            } else {
102                GgmlType::Q3K
103            };
104            let mut b = vec![0u8; t.byte_size()];
105            gpu.copy_d2h(t.ptr, &mut b)?;
106            let mut out = vec![0f32; n];
107            dequant_to_f32(gt, &b, n, &mut out)
108                .with_context(|| format!("{name}: CPU dequant of the resident {gt:?} blocks"))?;
109            Ok(out)
110        }
111        other => anyhow::bail!("{name}: cannot widen {other:?} to f32"),
112    }
113}
114
115pub fn upload_f32(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
116    let bytes: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
117    let p = gpu.alloc(bytes.len().max(4))?;
118    gpu.copy_h2d(&bytes, p)?;
119    Ok(p)
120}
121
122/// 2026-09-25: A new f32 device copy of `name`, which must hold `expect`
123/// elements. The store's tensor is widened on the host (`download_f32`) and
124/// uploaded.
125pub fn f32_ptr(
126    gpu: &dyn GpuBackend,
127    store: &WeightStore,
128    name: &str,
129    expect: usize,
130) -> Result<DevicePtr> {
131    let v = download_f32(gpu, store, name)?;
132    ensure!(
133        v.len() == expect,
134        "{name}: {} elements, expected {expect}",
135        v.len()
136    );
137    upload_f32(gpu, &v)
138}
139
140pub(super) fn hc_site(
141    gpu: &dyn GpuBackend,
142    store: &WeightStore,
143    lp: &str,
144    site: &str,
145    c: &ModelConfig,
146) -> Result<HcSiteWeights> {
147    let hc = c.hc_mult;
148    let mix_hc = (2 + hc) * hc;
149    Ok(HcSiteWeights {
150        hc_fn: f32_ptr(
151            gpu,
152            store,
153            &format!("{lp}.hc_{site}_fn"),
154            mix_hc * hc * c.hidden_size,
155        )?,
156        hc_base: f32_ptr(gpu, store, &format!("{lp}.hc_{site}_base"), mix_hc)?,
157        hc_scale: f32_ptr(gpu, store, &format!("{lp}.hc_{site}_scale"), 3)?,
158        lowrank: None,
159    })
160}
161
162impl ModelWeightLoader for DeepSeekV41WeightLoader {
163    fn supports_tp(&self) -> bool {
164        false
165    }
166
167    fn load_layers(
168        &self,
169        store: &WeightStore,
170        config: &ModelConfig,
171        gpu: &dyn GpuBackend,
172        layer_kv_dtypes: &[metrale_cache::kv_cache::KvCacheDtype],
173    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
174        load_layers::load_layers(store, config, gpu, layer_kv_dtypes)
175    }
176
177    fn load_embedding(
178        &self,
179        store: &WeightStore,
180        _config: &ModelConfig,
181        gpu: &dyn GpuBackend,
182    ) -> Result<DenseWeight> {
183        dense_auto(store, "model.embed_tokens.weight", gpu)
184    }
185
186    fn load_final_norm(
187        &self,
188        store: &WeightStore,
189        _config: &ModelConfig,
190        gpu: &dyn GpuBackend,
191    ) -> Result<DenseWeight> {
192        dense_auto(store, "model.norm.weight", gpu)
193    }
194
195    fn load_lm_head(
196        &self,
197        store: &WeightStore,
198        _config: &ModelConfig,
199        gpu: &dyn GpuBackend,
200    ) -> Result<DenseWeight> {
201        let t = store.get("lm_head.weight")?;
202        if t.dtype == WeightDtype::Q6K {
203            // 2026-09-25: A Q6_K head is returned as the raw blocks, unconverted.
204            return Ok(DenseWeight { weight: t.ptr });
205        }
206        dense_auto(store, "lm_head.weight", gpu)
207    }
208
209    fn load_mtp_weights(
210        &self,
211        _store: &WeightStore,
212        _config: &ModelConfig,
213        _gpu: &dyn GpuBackend,
214    ) -> Result<Option<MtpWeights>> {
215        Ok(None)
216    }
217}
218
219#[cfg(all(test, feature = "cuda"))]
220mod real_file_tests;
metrale_model_arch/weight_loader/deepseek_v41.rs

metrale_model_arch/weight_loader/
deepseek_v41.rs
