spark_model/weight_map/
model_b.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `weight_map.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::{Context, Result, bail, ensure};
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9use spark_runtime::weights::{WeightDtype, WeightStore};
10
11use super::*;
12
13impl ModelWeights {
14    /// Build typed weight references from a flat WeightStore.
15    ///
16    /// `layer_types` maps layer index → FullAttention or LinearAttention.
17    /// `num_experts` is 512 for Qwen3-Next.
18    pub fn from_store(
19        store: &WeightStore,
20        layer_types: &[atlas_core::config::LayerType],
21        num_experts: usize,
22        gpu: &dyn GpuBackend,
23        config: &atlas_core::config::ModelConfig,
24    ) -> Result<Self> {
25        let embed_tokens = dense(store, "model.embed_tokens.weight")?;
26        let final_norm = dense(store, "model.norm.weight")?;
27
28        // LM head may be tied to embed_tokens.
29        let lm_head = if store.contains("lm_head.weight") {
30            dense(store, "lm_head.weight")?
31        } else {
32            embed_tokens
33        };
34
35        let mut layers = Vec::with_capacity(layer_types.len());
36        for (i, lt) in layer_types.iter().enumerate() {
37            let lp = config.layer_prefix(i);
38            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
39            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
40            // ModelWeights::from_store is only used for Standard NVFP4 (test/legacy path).
41            let dummy_qctx = QuantizeCtx {
42                absmax_k: spark_runtime::gpu::KernelHandle(0),
43                quantize_k: spark_runtime::gpu::KernelHandle(0),
44                stream: 0,
45            };
46            let moe = load_moe(
47                store,
48                &lp,
49                num_experts,
50                gpu,
51                config,
52                Nvfp4Variant::Standard,
53                dummy_qctx,
54            )?;
55
56            match lt {
57                atlas_core::config::LayerType::FullAttention => {
58                    let attn = load_attention(
59                        store,
60                        &lp,
61                        gpu,
62                        Nvfp4Variant::Standard,
63                        dummy_qctx,
64                        config,
65                    )?;
66                    layers.push(LayerWeights::FullAttention {
67                        input_norm,
68                        attn,
69                        post_attn_norm,
70                        moe,
71                    });
72                }
73                atlas_core::config::LayerType::LinearAttention => {
74                    let ssm =
75                        load_ssm(store, &lp, gpu, Nvfp4Variant::Standard, dummy_qctx, config)?;
76                    layers.push(LayerWeights::LinearAttention {
77                        input_norm,
78                        ssm,
79                        post_attn_norm,
80                        moe,
81                    });
82                }
83                atlas_core::config::LayerType::SlidingAttention => {
84                    unreachable!("unexpected SlidingAttention in this loader")
85                }
86                atlas_core::config::LayerType::Moe => {
87                    unreachable!("Qwen3 has no standalone MoE layers")
88                }
89                // GLM-5.3's `deepseek_sparse_attention`. `LayerWeights` has no sparse
90                // variant, so bail rather than fall through to `FullAttention` — a
91                // sparse layer bound as dense attends over the whole cache and produces
92                // plausible output, which is the worst failure mode available.
93                atlas_core::config::LayerType::SparseAttention => anyhow::bail!(
94                    "layer {i}: SparseAttention has no weight-map variant in this loader"
95                ),
96            }
97
98            if (i + 1) % 12 == 0 {
99                tracing::info!("Mapped weights for layers 0..{}", i + 1);
100            }
101        }
102
103        tracing::info!(
104            "Weight map: {} layers ({} attention, {} SSM)",
105            layers.len(),
106            layers
107                .iter()
108                .filter(|l| matches!(l, LayerWeights::FullAttention { .. }))
109                .count(),
110            layers
111                .iter()
112                .filter(|l| matches!(l, LayerWeights::LinearAttention { .. }))
113                .count(),
114        );
115
116        Ok(Self {
117            embed_tokens,
118            final_norm,
119            lm_head,
120            layers,
121        })
122    }
123}
124
125// ── Nemotron-H weight types and loaders ──
spark_model/weight_map/model_b.rs

spark_model/weight_map/
model_b.rs
