metrale_model_arch/weight_loader/deepseek_v41/
load_layers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: Builds the DeepSeek-V4.1 layers and the `V41Runtime` they share:
5//! the expert arena and cache, the attention, MoE and engram engines, and the
6//! engram hash tables.
7//!
8//! Owner: model-arch weight loader.
9//! Invariants:
10//! - Every returned layer holds a clone of the same `Arc<V41Runtime>`.
11//! - Loading fails unless `hc_mult` is in 1..=4, the expert stacks hold
12//!   `num_experts` experts and `compress_ratios` has an entry for every layer.
13
14use std::sync::{Arc, Mutex};
15
16use anyhow::{Context, Result, ensure};
17use metrale_config::ModelConfig;
18use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
19use metrale_model_weights::weights::WeightStore;
20use metrale_model_weights::weights::expert_stream::{
21    EngramRowReader, ExpertArena, ExpertSliceMap, ExpertSource, ShardFiles,
22};
23
24use crate::attn_v41::{
25    AttnV41, AttnV41Cfg, AttnV41LayerWeights, CompressorWeightsGpu, IndexerWeightsGpu, LayerRole,
26    SharedV41,
27};
28use crate::deepseek_v41_layer::SegState;
29use crate::deepseek_v41_layer::{DeepSeekV41Layer, V41Runtime};
30use crate::engram_v41::{EngramHashTables, EngramHasher, EngramLayerWeights, EngramV41};
31use crate::moe_v41::{MoeV41, MoeV41Cfg, MoeV41LayerWeights, RouterWeights};
32use metrale_model_layers::layer::TransformerLayer;
33use metrale_model_layers::weight_map::DenseWeight;
34
35use super::{
36    DEFAULT_CANDIDATE_BLOCK, DEFAULT_CANDIDATE_SOURCE, DEFAULT_CANDIDATE_TOPK_BLOCKS, bf16_ptr,
37    download_f32, env_usize, f32_ptr, hc_site, resident_mat,
38};
39
40pub fn load_layers(
41    store: &WeightStore,
42    config: &ModelConfig,
43    gpu: &dyn GpuBackend,
44    _layer_kv_dtypes: &[metrale_cache::kv_cache::KvCacheDtype],
45) -> Result<Vec<Box<dyn TransformerLayer>>> {
46    let n_layers = config.num_hidden_layers;
47    let (dim, hc) = (config.hidden_size, config.hc_mult);
48    ensure!(
49        (1..=4).contains(&hc),
50        "deepseek-v4.1: hc_mult {hc} outside 1..=4"
51    );
52    let head_dim = config.head_dim;
53    let nh = config.num_attention_heads;
54
55    let anchor = store
56        .deferred("model.layers.0.ffn.experts_stack.gate")
57        .context("deepseek-v4.1: the routed expert stacks were not deferred by the GGUF loader")?;
58    let model_dir = anchor.path.parent().context("shard path has no parent")?;
59    let files = Arc::new(ShardFiles::open_dir(model_dir)?);
60    let slices = Arc::new(ExpertSliceMap::new(files.clone())?);
61    let rows = EngramRowReader::new(files.clone())?;
62    ensure!(
63        slices.num_experts() == config.num_experts,
64        "expert stacks hold {} experts, config says {}",
65        slices.num_experts(),
66        config.num_experts
67    );
68
69    let max_seq =
70        env_usize("METRALE_DS41_MAX_SEQ", 8192).min(config.max_position_embeddings.max(1));
71    let max_tokens = env_usize("METRALE_DS41_MAX_TOKENS", 2048).min(max_seq);
72    let cache_gib: f64 = std::env::var("METRALE_DS41_EXPERT_CACHE_GIB")
73        .ok()
74        .and_then(|v| v.parse().ok())
75        .filter(|g: &f64| g.is_finite() && *g > 0.0)
76        .unwrap_or(88.0);
77    let cache_bytes = (cache_gib * (1u64 << 30) as f64) as usize;
78    let reader_threads = env_usize("METRALE_DS41_READER_THREADS", 16);
79    let ratios: Vec<usize> = config
80        .compress_ratios
81        .iter()
82        .take(n_layers)
83        .copied()
84        .collect();
85    ensure!(
86        ratios.len() == n_layers,
87        "compress_ratios has {} entries for {n_layers} layers",
88        ratios.len()
89    );
90    let has = |l: usize, t: &str| store.contains(&format!("model.layers.{l}.{t}"));
91    let kv_sources: Vec<usize> = (0..n_layers)
92        .filter(|&l| has(l, "compressor.wkv.weight"))
93        .collect();
94    let index_sources: Vec<usize> = (0..n_layers)
95        .filter(|&l| has(l, "indexer.wq_b.weight"))
96        .collect();
97    ensure!(
98        !kv_sources.is_empty() && !index_sources.is_empty(),
99        "deepseek-v4.1: no compressor / indexer tensors found"
100    );
101    let cand_src = config
102        .candidate_source_layer_id
103        .unwrap_or_else(|| env_usize("METRALE_DS41_CANDIDATE_SOURCE", DEFAULT_CANDIDATE_SOURCE));
104    let cand_topk_blocks = if config.candidate_topk_blocks > 0 {
105        config.candidate_topk_blocks
106    } else {
107        DEFAULT_CANDIDATE_TOPK_BLOCKS
108    };
109    let cand_block = if config.candidate_block_size > 0 {
110        config.candidate_block_size
111    } else {
112        DEFAULT_CANDIDATE_BLOCK
113    };
114    let attn_cfg = AttnV41Cfg {
115        dim,
116        n_heads: nh,
117        head_dim,
118        rope_dim: config.rotary_dim,
119        q_rank: config.q_lora_rank,
120        o_rank: config.o_lora_rank,
121        groups: config.o_groups.max(1),
122        window: config.sliding_window as usize,
123        eps: config.rms_norm_eps as f32,
124        index_heads: config.index_n_heads,
125        index_hd: config.index_head_dim,
126        index_topk: config.index_topk,
127        cand_topk_blocks,
128        cand_block,
129        max_seq,
130        max_tokens,
131        rope_theta: config.rope_theta as f32,
132        compress_rope_theta: config.compress_rope_theta,
133        rope_factor: if config.yarn_factor > 0.0 {
134            config.yarn_factor
135        } else {
136            16.0
137        },
138        orig_seq: if config.yarn_original_max_position_embeddings > 0 {
139            config.yarn_original_max_position_embeddings
140        } else {
141            65536
142        },
143        beta_fast: if config.yarn_beta_fast > 0.0 {
144            config.yarn_beta_fast
145        } else {
146            32.0
147        },
148        beta_slow: if config.yarn_beta_slow > 0.0 {
149            config.yarn_beta_slow
150        } else {
151            1.0
152        },
153    };
154    let moe_cfg = MoeV41Cfg {
155        dim,
156        inter: config.moe_intermediate_size,
157        n_routed: config.num_experts,
158        topk: config.num_experts_per_tok,
159        gate_temp: 1.0,
160        norm_topk_prob: config.norm_topk_prob,
161        route_scale: config.routed_scaling_factor as f32,
162        swiglu_limit: config.swiglu_limit,
163        max_tokens,
164    };
165    tracing::info!(
166        "DeepSeek-V4.1: {n_layers} layers, kv sources {kv_sources:?}, index sources {index_sources:?}, candidate source {cand_src} ({cand_topk_blocks} x {cand_block}), window {}, max_seq {max_seq}, max_tokens {max_tokens}, expert cache {cache_gib} GiB, {reader_threads} readers",
167        attn_cfg.window
168    );
169
170    let token_map = files
171        .header(0)
172        .get_i64_array("deepseek41.engram.token_map")
173        .context("deepseek41.engram.token_map missing from the GGUF metadata")?;
174    let tables = Arc::new(EngramHashTables::from_flat(
175        config.engram_layer_ids.clone(),
176        config.engram_max_ngram_size,
177        config.engram_n_heads,
178        config.engram_pad_token_id,
179        token_map,
180        &config.engram_multipliers,
181        &config.engram_primes,
182        &config.engram_offsets,
183    )?);
184    let cols = tables.n_hash_cols();
185
186    let layout = slices.slot_layout();
187    // 2026-09-25: The expert slots are device memory behind a page-locked
188    // staging ring; `METRALE_DS41_ARENA_DEVICE=0` selects a page-locked arena
189    // the GPU reads in place.
190    let arena_device = !std::env::var("METRALE_DS41_ARENA_DEVICE").is_ok_and(|v| v == "0");
191    let staging_slots = env_usize("METRALE_DS41_STAGING_SLOTS", 16);
192    let (arena, mut lru) =
193        ExpertArena::alloc(gpu, cache_bytes, layout, arena_device, staging_slots)?;
194    tracing::info!(
195        "DeepSeek-V4.1: expert arena {}; MemAvailable now {}",
196        arena.describe(),
197        mem_available()
198    );
199    // 2026-09-25: `METRALE_DS41_READER_POOL=1` reads misses and predicted
200    // experts on persistent reader threads, which `METRALE_DS41_PREFETCH_K`
201    // needs; otherwise each fetch uses scoped threads. The pool is only armed
202    // with the page-locked arena.
203    if std::env::var("METRALE_DS41_READER_POOL").is_ok_and(|v| v == "1") {
204        if arena_device {
205            tracing::warn!(
206                "DeepSeek-V4.1: METRALE_DS41_READER_POOL=1 ignored with the device arena (set METRALE_DS41_ARENA_DEVICE=0 for the pool)"
207            );
208        } else {
209            lru.set_pool(slices.clone(), reader_threads);
210            tracing::info!("DeepSeek-V4.1: expert reader pool of {reader_threads} threads");
211        }
212    }
213    if let Ok(path) = std::env::var("METRALE_DS41_ROUTE_TRACE") {
214        lru.set_trace(&path)?;
215        tracing::info!("DeepSeek-V4.1: expert route trace -> {path}");
216    }
217    tracing::info!(
218        "DeepSeek-V4.1: expert cache {} slots of {:.2} MiB",
219        lru.n_slots(),
220        layout.bytes as f64 / 1048576.0
221    );
222    let attn = AttnV41::new(gpu, attn_cfg.clone())?;
223    let moe = MoeV41::new(gpu, moe_cfg.clone())?;
224    let mut engram = EngramV41::new(
225        gpu,
226        dim,
227        hc,
228        config.engram_head_dim,
229        cols,
230        config.rms_norm_eps as f32,
231        max_tokens,
232    )?;
233    for &l in &config.engram_layer_ids {
234        let lp = format!("model.layers.{l}");
235        let q = download_f32(gpu, store, &format!("{lp}.engram.wq"))?;
236        let k = download_f32(gpu, store, &format!("{lp}.engram.wk"))?;
237        ensure!(
238            q.len() == hc * dim && k.len() == hc * dim,
239            "engram q/k of layer {l}: {} / {} elements, expected {}",
240            q.len(),
241            k.len(),
242            hc * dim
243        );
244        engram.add_layer(
245            gpu,
246            EngramLayerWeights {
247                layer: l,
248                wkv: bf16_ptr(store, &format!("{lp}.engram.wkv"))?,
249                qk: EngramV41::upload_qk(gpu, &q, &k)?,
250                raw: DevicePtr(0),
251                rows: DevicePtr(0),
252                wkv_q2k: super::engram_q2k::engram_wkv_q2k(gpu, &files, l)?,
253            },
254        )?;
255    }
256    let alloc_f32 = |n: usize| gpu.alloc((n * 4).max(16));
257    let rt = Arc::new(V41Runtime {
258        attn: Mutex::new(attn),
259        moe: Mutex::new(moe),
260        engram: Mutex::new(engram),
261        lru: Mutex::new(lru),
262        arena,
263        slices,
264        rows,
265        hasher: Mutex::new(EngramHasher::new(tables.clone(), max_seq)),
266        tables,
267        shared: Mutex::new(SharedV41::default()),
268        step_hashes: Mutex::new(None),
269        pre_prev: alloc_f32(max_tokens * hc)?,
270        mixes_s: alloc_f32(max_tokens * (2 + hc) * hc)?,
271        reader_threads,
272        n_layers,
273        hc_mult: hc,
274        hidden: dim,
275        sinkhorn_iters: config.hc_sinkhorn_iters.max(1),
276        hc_eps: config.hc_eps,
277        norm_eps: config.rms_norm_eps as f32,
278        seg: Mutex::new(SegState::default()),
279        roles: Mutex::new(vec![None; n_layers]),
280        engram_layers: Mutex::new(Vec::new()),
281        pre_a: alloc_f32(max_tokens * hc)?,
282        pre_f: alloc_f32(max_tokens * hc)?,
283        post_s: alloc_f32(max_tokens * hc)?,
284        comb_s: alloc_f32(max_tokens * hc * hc)?,
285        attn_in: gpu.alloc(max_tokens * dim * 2)?,
286        attn_cfg: attn_cfg.clone(),
287        moe_cfg,
288        max_tokens,
289        step_moe: Mutex::new(Default::default()),
290        step_attn_ms: Mutex::new(0.0),
291        step_engram_ms: Mutex::new(0.0),
292        step_start: Mutex::new(None),
293        graph_disabled: std::sync::atomic::AtomicBool::new(false),
294        pred_x: gpu.alloc(3 * dim * 2)?,
295        pred_trace: Mutex::new(match std::env::var("METRALE_DS41_PREDICT_TRACE") {
296            Ok(path) => {
297                tracing::info!("DeepSeek-V4.1: routing prediction trace -> {path}");
298                Some(std::io::BufWriter::new(std::fs::File::create(&path)?))
299            }
300            Err(_) => None,
301        }),
302    });
303
304    let k_hc_expand = gpu.kernel("hyper_connection", "hc_expand")?;
305    let k_hc_post = gpu.kernel("hyper_connection", "hc_post")?;
306    let k_mixes_dot = gpu.kernel("hc_v41", "hc_v41_mixes_dot")?;
307    let k_mixes_finish = gpu.kernel("hc_v41", "hc_v41_mixes_finish")?;
308    let k_collapse = gpu.kernel("hc_v41", "hc_v41_collapse")?;
309    let k_finish_collapse = gpu.kernel("hc_v41", "hc_v41_finish_collapse")?;
310    let k_collapse_wide = gpu.kernel("hc_v41", "hc_v41_collapse_wide")?;
311    let k_post_wide = gpu.kernel("hc_v41", "hc_v41_post_wide")?;
312    // 2026-09-25: Every DeepSeek-V4.1 layer norm scales by `w` through
313    // `rms_norm_vanilla`; the shared `rms_norm` kernel scales by `1 + w`. The
314    // final norm gets the vanilla kernel through `ships_vanilla_norm_weights`.
315    let k_rms_norm = gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")?;
316
317    let mut layers: Vec<Box<dyn TransformerLayer>> = Vec::with_capacity(n_layers);
318    for l in 0..n_layers {
319        let lp = format!("model.layers.{l}");
320        let ratio = ratios[l];
321        let role = LayerRole {
322            ratio,
323            is_kv_source: kv_sources.contains(&l),
324            is_index_source: index_sources.contains(&l),
325            is_candidate_source: cand_src == l,
326            uses_candidates: cand_src < l,
327        };
328        let comp = if role.is_kv_source {
329            let norm = f32_ptr(
330                gpu,
331                store,
332                &format!("{lp}.compressor.norm.weight"),
333                head_dim,
334            )?;
335            if ratio > 1 {
336                Some(CompressorWeightsGpu {
337                    kv: f32_ptr(
338                        gpu,
339                        store,
340                        &format!("{lp}.compressor.wkv.weight"),
341                        head_dim * dim,
342                    )?,
343                    gate: Some(f32_ptr(
344                        gpu,
345                        store,
346                        &format!("{lp}.compressor.wgate.weight"),
347                        head_dim * dim,
348                    )?),
349                    norm,
350                })
351            } else {
352                Some(CompressorWeightsGpu {
353                    kv: bf16_ptr(store, &format!("{lp}.compressor.wkv.weight"))?,
354                    gate: None,
355                    norm,
356                })
357            }
358        } else {
359            None
360        };
361        let idx = if role.is_index_source {
362            Some(IndexerWeightsGpu {
363                wq_b: bf16_ptr(store, &format!("{lp}.indexer.wq_b.weight"))?,
364                weights_proj: bf16_ptr(store, &format!("{lp}.indexer.proj.weight"))?,
365                wk: if role.is_kv_source {
366                    Some(bf16_ptr(store, &format!("{lp}.indexer.wk.weight"))?)
367                } else {
368                    None
369                },
370                k_norm: if role.is_kv_source {
371                    Some(f32_ptr(
372                        gpu,
373                        store,
374                        &format!("{lp}.indexer.k_norm.weight"),
375                        config.index_head_dim,
376                    )?)
377                } else {
378                    None
379                },
380            })
381        } else {
382            None
383        };
384        let attn_w = AttnV41LayerWeights {
385            role,
386            sink: f32_ptr(gpu, store, &format!("{lp}.attn.attn_sink"), nh)?,
387            wq_a: resident_mat(store, &format!("{lp}.attn.wq_a.weight"))?,
388            q_norm: f32_ptr(
389                gpu,
390                store,
391                &format!("{lp}.attn.q_norm.weight"),
392                config.q_lora_rank,
393            )?,
394            wq_b: resident_mat(store, &format!("{lp}.attn.wq_b.weight"))?,
395            wkv: resident_mat(store, &format!("{lp}.attn.wkv.weight"))?,
396            kv_norm: f32_ptr(gpu, store, &format!("{lp}.attn.kv_norm.weight"), head_dim)?,
397            wo_a: resident_mat(store, &format!("{lp}.attn.wo_a.weight"))?,
398            wo_b: resident_mat(store, &format!("{lp}.attn.wo_b.weight"))?,
399            comp,
400            idx,
401        };
402        let gate_bias = download_f32(
403            gpu,
404            store,
405            &format!("{lp}.ffn.gate.e_score_correction_bias"),
406        )?;
407        ensure!(
408            gate_bias.len() == config.num_experts,
409            "layer {l}: correction bias has {} entries",
410            gate_bias.len()
411        );
412        let gate_bias_dev = MoeV41::upload_bias(gpu, &gate_bias)?;
413        let moe_w = MoeV41LayerWeights {
414            layer: l as u32,
415            gate_w: bf16_ptr(store, &format!("{lp}.ffn.gate.weight"))?,
416            gate_bias,
417            gate_bias_dev,
418            shared_w1: resident_mat(store, &format!("{lp}.ffn.shared_experts.w1"))?,
419            shared_w2: resident_mat(store, &format!("{lp}.ffn.shared_experts.w2"))?,
420            shared_w3: resident_mat(store, &format!("{lp}.ffn.shared_experts.w3"))?,
421        };
422        let next_router = if l + 1 < n_layers {
423            let np = format!("model.layers.{}", l + 1);
424            Some(RouterWeights {
425                layer: (l + 1) as u32,
426                gate_w: bf16_ptr(store, &format!("{np}.ffn.gate.weight"))?,
427                gate_bias: download_f32(
428                    gpu,
429                    store,
430                    &format!("{np}.ffn.gate.e_score_correction_bias"),
431                )?,
432            })
433        } else {
434            None
435        };
436        let engram_index = rt.tables.hash_index(l);
437        rt.roles.lock().unwrap()[l] = Some(role);
438        if let Some(hi) = engram_index {
439            rt.engram_layers.lock().unwrap().push((l, hi));
440        }
441        layers.push(Box::new(DeepSeekV41Layer {
442            idx: l,
443            role,
444            rt: rt.clone(),
445            attn_w,
446            moe_w,
447            next_router,
448            engram_index,
449            hc_attn: hc_site(gpu, store, &lp, "attn", config)?,
450            hc_ffn: hc_site(gpu, store, &lp, "ffn", config)?,
451            attn_norm: DenseWeight {
452                weight: bf16_ptr(store, &format!("{lp}.attn_norm.weight"))?,
453            },
454            ffn_norm: DenseWeight {
455                weight: bf16_ptr(store, &format!("{lp}.ffn_norm.weight"))?,
456            },
457            k_hc_expand,
458            k_hc_post,
459            k_mixes_dot,
460            k_mixes_finish,
461            k_collapse,
462            k_finish_collapse,
463            k_collapse_wide,
464            k_post_wide,
465            k_rms_norm,
466        }));
467    }
468    Ok(layers)
469}
470
471/// 2026-09-25: `MemAvailable` from `/proc/meminfo` in GiB, or "unknown".
472fn mem_available() -> String {
473    std::fs::read_to_string("/proc/meminfo")
474        .ok()
475        .and_then(|s| {
476            s.lines()
477                .find(|l| l.starts_with("MemAvailable:"))
478                .and_then(|l| l.split_whitespace().nth(1))
479                .and_then(|kb| kb.parse::<f64>().ok())
480        })
481        .map(|kb| format!("{:.1} GiB", kb / 1048576.0))
482        .unwrap_or_else(|| "unknown".to_string())
483}
metrale_model_arch/weight_loader/deepseek_v41/load_layers.rs

metrale_model_arch/weight_loader/deepseek_v41/
load_layers.rs
