metrale_model_arch/weight_loader/
step3p7.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Step3p7WeightLoader`, the Step 3.7 Flash weight loader, with
4//! its norm-offset and fused-expert helpers.
5//!
6//! Owner: model-arch weight loader.
7//! Invariants: none beyond the types.
8//!
9//! Layers are read under `{weight_prefix}.layers.{i}`, with
10//! `model.language_model` when the prefix is empty. A layer with
11//! `moe.gate.weight` gets a MoE FFN and any other a dense FFN
12//! (`load_layers.rs`). Routed experts are read per expert
13//! (`moe.experts.{e}.*`) when such tensors exist, otherwise from one fused
14//! tensor per projection that `slice_fused_experts` splits by byte offset.
15//! MTP modules, probed at `model.layers.{num_hidden_layers}`, are detected
16//! but not loaded.
17
18mod load_layers;
19
20use anyhow::Result;
21use metrale_cache::kv_cache::KvCacheDtype;
22use metrale_config::ModelConfig;
23use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
24use metrale_model_weights::weights::WeightStore;
25
26use super::ModelWeightLoader;
27use metrale_model_layers::layer::TransformerLayer;
28use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, QuantizedWeight, dense};
29
30pub struct Step3p7WeightLoader;
31
32/// 2026-09-25: Add 1.0 to each of `size` BF16 norm weights in place, so the
33/// standard RMSNorm kernel, `(x / rms) * weight`, computes
34/// `(x / rms) * (weight + 1)`. The loader applies it to every norm it loads.
35fn offset_norm_weights_plus_one(
36    weight: &DenseWeight,
37    size: usize,
38    gpu: &dyn GpuBackend,
39) -> Result<()> {
40    let byte_len = size * 2;
41    let mut buf = vec![0u8; byte_len];
42    gpu.copy_d2h(weight.weight, &mut buf)?;
43
44    for i in 0..size {
45        let bits = u16::from_le_bytes([buf[i * 2], buf[i * 2 + 1]]);
46        let f32_val = f32::from_bits((bits as u32) << 16);
47        let new_val = f32_val + 1.0;
48        // 2026-09-25: Round to nearest even when truncating to BF16.
49        let f32_bits = new_val.to_bits();
50        let new_bits = ((f32_bits + 0x7FFF + ((f32_bits >> 16) & 1)) >> 16) as u16;
51        buf[i * 2] = new_bits as u8;
52        buf[i * 2 + 1] = (new_bits >> 8) as u8;
53    }
54
55    gpu.copy_h2d(&buf, weight.weight)?;
56    Ok(())
57}
58
59/// 2026-09-25: Split one fused NVFP4 projection into `num_experts`
60/// `QuantizedWeight` views of `[n, k]`. Expert `e` starts at `e * n * k / 2`
61/// bytes of the packed weight, `e * n * ceil(k / 16)` bytes of the group
62/// scales and, when present, `e * n * 4` bytes of the input scale; all share
63/// `global_scale_2`. Nothing is copied.
64fn slice_fused_experts(
65    fused_weight: DevicePtr,
66    fused_scale: DevicePtr,
67    fused_input_scale: DevicePtr,
68    global_scale_2: f32,
69    num_experts: usize,
70    n: usize,
71    k: usize,
72) -> Vec<QuantizedWeight> {
73    let group_size = 16usize;
74    let packed_bytes_per_expert = n * k / 2;
75    let scale_bytes_per_expert = n * k.div_ceil(group_size);
76    let input_scale_bytes_per_expert = n * 4;
77
78    (0..num_experts)
79        .map(|e| QuantizedWeight {
80            weight: fused_weight.offset(e * packed_bytes_per_expert),
81            weight_scale: fused_scale.offset(e * scale_bytes_per_expert),
82            weight_scale_2: global_scale_2,
83            input_scale: if fused_input_scale == DevicePtr::NULL {
84                DevicePtr::NULL
85            } else {
86                fused_input_scale.offset(e * input_scale_bytes_per_expert)
87            },
88            weight_scale_2_vec: DevicePtr::NULL,
89            act: metrale_config::Nvfp4Act::Unstamped,
90        })
91        .collect()
92}
93
94/// 2026-09-25: Whether any store tensor name starts with
95/// `{layer_prefix}.moe.experts.`.
96fn has_per_expert_tensors(store: &WeightStore, layer_prefix: &str) -> bool {
97    let pattern = format!("{layer_prefix}.moe.experts.");
98    let found = store.names().any(|k| k.starts_with(&pattern));
99    tracing::debug!("has_per_expert_tensors('{layer_prefix}'): pattern='{pattern}', found={found}");
100    found
101}
102
103/// 2026-09-25: The store pointers of one fused NVFP4 projection:
104/// `(weight, weight_scale, input_scale or NULL, weight_scale_2 as f32)`.
105fn load_fused_nvfp4(
106    store: &WeightStore,
107    prefix: &str,
108    gpu: &dyn GpuBackend,
109) -> Result<(DevicePtr, DevicePtr, DevicePtr, f32)> {
110    let weight = store.get(&format!("{prefix}.weight"))?.ptr;
111    let weight_scale = store.get(&format!("{prefix}.weight_scale"))?.ptr;
112
113    let ws2_key = format!("{prefix}.weight_scale_2");
114    let ws2_ptr = store.get(&ws2_key)?.ptr;
115    let mut ws2_buf = [0u8; 4];
116    gpu.copy_d2h(ws2_ptr, &mut ws2_buf)?;
117    let weight_scale_2 = f32::from_le_bytes(ws2_buf);
118
119    let is_key = format!("{prefix}.input_scale");
120    let input_scale = if store.contains(&is_key) {
121        store.get(&is_key)?.ptr
122    } else {
123        DevicePtr::NULL
124    };
125
126    Ok((weight, weight_scale, input_scale, weight_scale_2))
127}
128
129impl ModelWeightLoader for Step3p7WeightLoader {
130    fn supports_tp(&self) -> bool {
131        false
132    }
133
134    fn load_layers(
135        &self,
136        store: &WeightStore,
137        config: &ModelConfig,
138        gpu: &dyn GpuBackend,
139        layer_kv_dtypes: &[KvCacheDtype],
140    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
141        load_layers::load_layers(store, config, gpu, layer_kv_dtypes)
142    }
143
144    fn load_embedding(
145        &self,
146        store: &WeightStore,
147        config: &ModelConfig,
148        _gpu: &dyn GpuBackend,
149    ) -> Result<DenseWeight> {
150        let prefix = if config.weight_prefix.is_empty() {
151            "model.language_model"
152        } else {
153            &config.weight_prefix
154        };
155        dense(store, &format!("{prefix}.embed_tokens.weight"))
156    }
157
158    fn load_final_norm(
159        &self,
160        store: &WeightStore,
161        config: &ModelConfig,
162        gpu: &dyn GpuBackend,
163    ) -> Result<DenseWeight> {
164        let prefix = if config.weight_prefix.is_empty() {
165            "model.language_model"
166        } else {
167            &config.weight_prefix
168        };
169        let w = dense(store, &format!("{prefix}.norm.weight"))?;
170        offset_norm_weights_plus_one(&w, config.hidden_size, gpu)?;
171        Ok(w)
172    }
173
174    fn load_lm_head(
175        &self,
176        store: &WeightStore,
177        _config: &ModelConfig,
178        _gpu: &dyn GpuBackend,
179    ) -> Result<DenseWeight> {
180        dense(store, "lm_head.weight")
181    }
182
183    fn load_mtp_weights(
184        &self,
185        _store: &WeightStore,
186        _config: &ModelConfig,
187        _gpu: &dyn GpuBackend,
188    ) -> Result<Option<MtpWeights>> {
189        Ok(None)
190    }
191
192    fn load_mtp_weights_multi(
193        &self,
194        store: &WeightStore,
195        config: &ModelConfig,
196        _gpu: &dyn GpuBackend,
197    ) -> Result<Vec<MtpWeights>> {
198        let first_mtp_idx = config.num_hidden_layers;
199        let probe = format!("model.layers.{first_mtp_idx}.input_layernorm.weight");
200        if !store.contains(&probe) {
201            tracing::info!(
202                "step3p7: no MTP module weights found \
203                 (expected at layer {first_mtp_idx}); MTP disabled"
204            );
205            return Ok(Vec::new());
206        }
207
208        tracing::info!(
209            "step3p7: MTP module weights detected at layers {}-{} but MTP loader \
210             not yet implemented. Run with --speculative 0 for non-MTP decode.",
211            first_mtp_idx,
212            first_mtp_idx + config.mtp_num_hidden_layers - 1,
213        );
214        Ok(Vec::new())
215    }
216}
metrale_model_arch/weight_loader/step3p7.rs

metrale_model_arch/weight_loader/
step3p7.rs
