metrale_model_layers/weight_map/
model_b.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `ModelWeights::from_store`, the whole-model loader for ModelOpt NVFP4 checkpoints with attention and linear-attention layers.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use anyhow::{Context, Result, bail, ensure};
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12use metrale_model_weights::weights::{WeightDtype, WeightStore};
13
14use super::*;
15
16impl ModelWeights {
17    /// 2026-09-25: Build typed weight references from a flat `WeightStore`,
18    /// loading every layer as `Nvfp4Variant::Standard`.
19    ///
20    /// `layer_types[i]` must be `FullAttention` or `LinearAttention`: `SparseAttention`
21    /// is an error, and `SlidingAttention` or `Moe` panics.
22    pub fn from_store(
23        store: &WeightStore,
24        layer_types: &[metrale_config::LayerType],
25        num_experts: usize,
26        gpu: &dyn GpuBackend,
27        config: &metrale_config::ModelConfig,
28    ) -> Result<Self> {
29        let embed_tokens = dense(store, "model.embed_tokens.weight")?;
30        let final_norm = dense(store, "model.norm.weight")?;
31
32        let lm_head = if store.contains("lm_head.weight") {
33            dense(store, "lm_head.weight")?
34        } else {
35            embed_tokens
36        };
37
38        let mut layers = Vec::with_capacity(layer_types.len());
39        for (i, lt) in layer_types.iter().enumerate() {
40            let lp = config.layer_prefix(i);
41            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
42            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
43            // 2026-09-25: Placeholder kernel handles (0): a key that `quantized_any`
44            // routes to load-time quantization gets no usable kernel here.
45            let dummy_qctx = QuantizeCtx {
46                absmax_k: metrale_gpu_runtime::gpu::KernelHandle(0),
47                quantize_k: metrale_gpu_runtime::gpu::KernelHandle(0),
48                stream: 0,
49            };
50            let moe = load_moe(
51                store,
52                &lp,
53                num_experts,
54                gpu,
55                config,
56                Nvfp4Variant::Standard,
57                dummy_qctx,
58            )?;
59
60            match lt {
61                metrale_config::LayerType::FullAttention => {
62                    let attn = load_attention(
63                        store,
64                        &lp,
65                        gpu,
66                        Nvfp4Variant::Standard,
67                        dummy_qctx,
68                        config,
69                    )?;
70                    layers.push(LayerWeights::FullAttention {
71                        input_norm,
72                        attn,
73                        post_attn_norm,
74                        moe,
75                    });
76                }
77                metrale_config::LayerType::LinearAttention => {
78                    let ssm =
79                        load_ssm(store, &lp, gpu, Nvfp4Variant::Standard, dummy_qctx, config)?;
80                    layers.push(LayerWeights::LinearAttention {
81                        input_norm,
82                        ssm,
83                        post_attn_norm,
84                        moe,
85                    });
86                }
87                metrale_config::LayerType::SlidingAttention => {
88                    unreachable!("unexpected SlidingAttention in this loader")
89                }
90                metrale_config::LayerType::Moe => {
91                    unreachable!("Qwen3 has no standalone MoE layers")
92                }
93                // 2026-09-25: `LayerWeights` has no sparse variant, and a sparse layer
94                // bound as dense attention would attend over the whole cache.
95                metrale_config::LayerType::SparseAttention => anyhow::bail!(
96                    "layer {i}: SparseAttention has no weight-map variant in this loader"
97                ),
98            }
99
100            if (i + 1) % 12 == 0 {
101                tracing::info!("Mapped weights for layers 0..{}", i + 1);
102            }
103        }
104
105        tracing::info!(
106            "Weight map: {} layers ({} attention, {} SSM)",
107            layers.len(),
108            layers
109                .iter()
110                .filter(|l| matches!(l, LayerWeights::FullAttention { .. }))
111                .count(),
112            layers
113                .iter()
114                .filter(|l| matches!(l, LayerWeights::LinearAttention { .. }))
115                .count(),
116        );
117
118        Ok(Self {
119            embed_tokens,
120            final_norm,
121            lm_head,
122            layers,
123        })
124    }
125}
metrale_model_layers/weight_map/model_b.rs

metrale_model_layers/weight_map/
model_b.rs
