metrale_model_layers/weight_map/
model_a.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Basic store accessors: raw pointers, scalars, KV scales, ModelOpt NVFP4, native MXFP4, and dtype-converting dense loads.
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
16/// 2026-09-25: Whole-model weights built by `ModelWeights::from_store`.
17pub struct ModelWeights {
18    pub embed_tokens: DenseWeight,
19    pub final_norm: DenseWeight,
20    /// 2026-09-25: `lm_head.weight`, or `embed_tokens` when the checkpoint has none.
21    pub lm_head: DenseWeight,
22    pub layers: Vec<LayerWeights>,
23}
24
25/// 2026-09-25: The device pointer of tensor `name`.
26pub(crate) fn ptr(store: &WeightStore, name: &str) -> Result<DevicePtr> {
27    Ok(store.get(name)?.ptr)
28}
29
30/// 2026-09-25: Read a one-element FP32 tensor to the host; any other dtype or size is an error.
31pub(crate) fn scalar_f32(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result<f32> {
32    let w = store.get(name)?;
33    ensure!(
34        w.dtype == WeightDtype::FP32,
35        "Expected FP32 for {name}, got {:?}",
36        w.dtype
37    );
38    ensure!(
39        w.num_elements() == 1,
40        "Expected scalar for {name}, got {} elements",
41        w.num_elements()
42    );
43    let mut buf = [0u8; 4];
44    gpu.copy_d2h(w.ptr, &mut buf)?;
45    Ok(f32::from_le_bytes(buf))
46}
47
48/// 2026-09-28: Read a one-element BF16 or FP32 tensor to the host as f32; any
49/// other dtype or element count is an error. Checkpoints ship FP8 KV scales in
50/// both widths (unsloth: BF16 `[1]`; ModelOpt: FP32 `[]`), and an F16 export is
51/// already BF16 in the store.
52pub fn scalar_bf16_or_f32(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result<f32> {
53    let w = store.get(name)?;
54    ensure!(
55        w.num_elements() == 1,
56        "Expected a scalar for {name}, got shape {:?}",
57        w.shape
58    );
59    match w.dtype {
60        WeightDtype::BF16 => {
61            let mut buf = [0u8; 2];
62            gpu.copy_d2h(w.ptr, &mut buf)?;
63            Ok(f32::from_bits(u32::from(u16::from_le_bytes(buf)) << 16))
64        }
65        WeightDtype::FP32 => {
66            let mut buf = [0u8; 4];
67            gpu.copy_d2h(w.ptr, &mut buf)?;
68            Ok(f32::from_le_bytes(buf))
69        }
70        other => bail!("Expected BF16 or FP32 for {name}, got {other:?}"),
71    }
72}
73
74/// 2026-09-28: The checkpoint's FP8 KV-cache scales `(k_scale, v_scale)` for the
75/// attention layer at `attn_prefix`, or `None` when it ships none. The keys come
76/// from `WeightStore::kv_scale_keys`, the resolver serve's census also uses.
77///
78/// Each scale is the dequant multiplier (`amax / 448`): the KV write divides by
79/// it and the FP8 attention kernels multiply by it. A scale that is not a
80/// finite positive number is an error, since the write computes `1 / scale`.
81pub fn load_checkpoint_kv_scales(
82    store: &WeightStore,
83    attn_prefix: &str,
84    gpu: &dyn GpuBackend,
85) -> Result<Option<(f32, f32)>> {
86    let Some(keys) = store.kv_scale_keys(attn_prefix)? else {
87        return Ok(None);
88    };
89    let read = |key: &str| -> Result<f32> {
90        let v = scalar_bf16_or_f32(store, key, gpu)
91            .with_context(|| format!("FP8 KV scale {key} ({:?})", keys.spelling))?;
92        ensure!(
93            v.is_finite() && v > 0.0,
94            "FP8 KV scale {key} is {v}; it must be finite and positive"
95        );
96        Ok(v)
97    };
98    let (k, v) = (read(&keys.k)?, read(&keys.v)?);
99    tracing::debug!(
100        "Loaded k_scale={k:.6} from {}, v_scale={v:.6} from {}",
101        keys.k,
102        keys.v
103    );
104    Ok(Some((k, v)))
105}
106
107/// 2026-09-28: [`load_checkpoint_kv_scales`], with `(1.0, 1.0)` for a layer the
108/// checkpoint ships no scales for. Serve logs how many layers that is
109/// (`resolve_kv_cache_config`), from the same resolver.
110pub fn load_kv_scales(
111    store: &WeightStore,
112    attn_prefix: &str,
113    gpu: &dyn GpuBackend,
114) -> Result<(f32, f32)> {
115    Ok(
116        load_checkpoint_kv_scales(store, attn_prefix, gpu)?.unwrap_or_else(|| {
117            tracing::debug!("No FP8 KV scales for {attn_prefix} in the checkpoint, using 1.0");
118            (1.0, 1.0)
119        }),
120    )
121}
122
123/// 2026-09-25: Build a `QuantizedWeight` from ModelOpt NVFP4 keys: `weight`,
124/// `weight_scale`, the FP32 scalar `weight_scale_2` (read to the host), and an
125/// optional `input_scale`.
126pub fn quantized(
127    store: &WeightStore,
128    prefix: &str,
129    gpu: &dyn GpuBackend,
130) -> Result<QuantizedWeight> {
131    let input_scale_key = format!("{prefix}.input_scale");
132    Ok(QuantizedWeight {
133        weight: ptr(store, &format!("{prefix}.weight"))?,
134        weight_scale: ptr(store, &format!("{prefix}.weight_scale"))?,
135        weight_scale_2: scalar_f32(store, &format!("{prefix}.weight_scale_2"), gpu)?,
136        input_scale: if store.contains(&input_scale_key) {
137            ptr(store, &input_scale_key)?
138        } else {
139            DevicePtr::NULL
140        },
141        weight_scale_2_vec: DevicePtr::NULL,
142        act: metrale_config::Nvfp4Act::Unstamped,
143    })
144}
145
146/// 2026-09-25: [`quantized_mxfp4_e8m0_pair`] on `{prefix}.weight` and `{prefix}.scale`.
147pub fn quantized_mxfp4_e8m0(store: &WeightStore, prefix: &str) -> Result<QuantizedWeight> {
148    quantized_mxfp4_e8m0_pair(
149        store,
150        &format!("{prefix}.weight"),
151        &format!("{prefix}.scale"),
152    )
153}
154
155/// 2026-09-25: Hand out a native MXFP4 weight's store pointers unchanged: E2M1
156/// nibbles `[n, k/2]` in `weight_key` and one E8M0 scale per 32 weights in
157/// `scale_key`. Any other group size inferred from the element counts is an
158/// error. `weight_scale_2` is 1.0 because the format has no per-tensor scale.
159pub fn quantized_mxfp4_e8m0_pair(
160    store: &WeightStore,
161    weight_key: &str,
162    scale_key: &str,
163) -> Result<QuantizedWeight> {
164    let w = store.get(weight_key)?;
165    ensure!(
166        w.shape.len() == 2,
167        "{weight_key}: expected packed rank-2 matrix"
168    );
169    let n = w.shape[0];
170    let k_packed = w.shape[1];
171    let total_nibbles = n * k_packed * 2;
172    let scale_t = store.get(scale_key)?;
173    let num_groups = scale_t.num_elements();
174    ensure!(
175        num_groups > 0 && total_nibbles.is_multiple_of(num_groups),
176        "{weight_key}: MXFP4 weight nibbles {total_nibbles} not divisible by E8M0 scale groups {num_groups}"
177    );
178    let block = total_nibbles / num_groups;
179    ensure!(
180        block == 32,
181        "{weight_key}: native MXFP4 expects GROUP_SIZE=32, inferred {block} (scale groups {num_groups}) \
182         — refusing to land a non-MX checkpoint on the transcode-free path"
183    );
184    Ok(QuantizedWeight {
185        weight: ptr(store, weight_key)?,
186        weight_scale: ptr(store, scale_key)?,
187        weight_scale_2: 1.0,
188        input_scale: DevicePtr::NULL,
189        weight_scale_2_vec: DevicePtr::NULL,
190        act: metrale_config::Nvfp4Act::Unstamped,
191    })
192}
193
194pub fn dense(store: &WeightStore, name: &str) -> Result<DenseWeight> {
195    let w = store.get(name)?;
196    Ok(DenseWeight { weight: w.ptr })
197}
198
199/// 2026-09-25: Load `{prefix}.weight` as BF16: a BF16 tensor is the store
200/// pointer, FP8E4M3 goes through `dequant_fp8_blockscaled_to_bf16`, and any
201/// other dtype is an error.
202pub fn dense_auto_fp8_or_bf16(
203    store: &WeightStore,
204    prefix: &str,
205    gpu: &dyn GpuBackend,
206) -> Result<DenseWeight> {
207    let w = store.get(&format!("{prefix}.weight"))?;
208    match w.dtype {
209        WeightDtype::BF16 => Ok(DenseWeight { weight: w.ptr }),
210        WeightDtype::FP8E4M3 => dequant_fp8_blockscaled_to_bf16(store, prefix, gpu),
211        other => anyhow::bail!(
212            "dense_auto_fp8_or_bf16: unsupported dtype {:?} for {prefix}.weight",
213            other
214        ),
215    }
216}
217
218/// 2026-09-25: Load `name` as BF16: an FP32 tensor is truncated into a new
219/// buffer, any other dtype is returned as the store pointer unchanged.
220pub fn dense_f32_safe(
221    store: &WeightStore,
222    name: &str,
223    gpu: &dyn GpuBackend,
224) -> Result<DenseWeight> {
225    let w = store.get(name)?;
226    if w.dtype == WeightDtype::FP32 {
227        // 2026-09-25: `f32_to_bf16_trunc` keeps the high 16 bits of each f32
228        // (round toward zero). It runs on the default stream and is not
229        // synchronized; a reader on another stream must synchronize it first.
230        let n = w.num_elements();
231        let ptr = gpu.alloc(n * 2)?;
232        let trunc = gpu.kernel("quantize_nvfp4", "f32_to_bf16_trunc")?;
233        let blocks = (n.div_ceil(256) as u32).max(1);
234        metrale_gpu_runtime::kernel_args::KernelLaunch::new(gpu, trunc)
235            .grid([blocks, 1, 1])
236            .block([256, 1, 1])
237            .arg_ptr(w.ptr)
238            .arg_ptr(ptr)
239            .arg_u32(n as u32)
240            .launch(gpu.default_stream())?;
241        Ok(DenseWeight { weight: ptr })
242    } else {
243        Ok(DenseWeight { weight: w.ptr })
244    }
245}
246
247/// 2026-09-25: Load `name` as FP32: an FP32 tensor is the store pointer, BF16 is
248/// widened on the host into a new buffer, and any other dtype is an error.
249/// `load_ssm` uses it for `A_log` and `dt_bias`.
250pub fn dense_keep_f32(
251    store: &WeightStore,
252    name: &str,
253    gpu: &dyn GpuBackend,
254) -> Result<DenseWeight> {
255    let w = store.get(name)?;
256    match w.dtype {
257        WeightDtype::FP32 => {
258            // 2026-09-25: Already FP32: the store's pointer, unconverted.
259            Ok(DenseWeight { weight: w.ptr })
260        }
261        WeightDtype::BF16 => {
262            tracing::info!(
263                "dense_keep_f32: promoting {name} from BF16 to FP32 ({:?})",
264                w.shape
265            );
266            let n = w.num_elements();
267            let mut bf16_buf = vec![0u8; n * 2];
268            gpu.copy_d2h(w.ptr, &mut bf16_buf)?;
269            let f32_buf: Vec<u8> = bf16_buf
270                .chunks_exact(2)
271                .flat_map(|c| {
272                    let bits = u16::from_le_bytes([c[0], c[1]]);
273                    let f32_bits = (bits as u32) << 16;
274                    f32_bits.to_le_bytes()
275                })
276                .collect();
277            let ptr = gpu.alloc(f32_buf.len())?;
278            gpu.copy_h2d(&f32_buf, ptr)?;
279            Ok(DenseWeight { weight: ptr })
280        }
281        other => {
282            bail!("dense_keep_f32: unsupported dtype {:?} for {name}", other);
283        }
284    }
285}
286
287/// 2026-09-25: Load a BF16 tensor into a new F32 device buffer, widened on the
288/// host; any other dtype is an error. The Nemotron-H Mamba loader uses it for
289/// `conv1d.bias`, `A_log`, `D` and `dt_bias`.
290pub(crate) fn dense_bf16_as_f32(
291    store: &WeightStore,
292    name: &str,
293    gpu: &dyn GpuBackend,
294) -> Result<DenseWeight> {
295    let w = store.get(name)?;
296    ensure!(
297        w.dtype == WeightDtype::BF16,
298        "Expected BF16 for {name}, got {:?}",
299        w.dtype
300    );
301    let n = w.num_elements();
302    let mut bf16_buf = vec![0u8; n * 2];
303    gpu.copy_d2h(w.ptr, &mut bf16_buf)?;
304    let f32_buf: Vec<u8> = bf16_buf
305        .chunks_exact(2)
306        .flat_map(|c| {
307            let bits = u16::from_le_bytes([c[0], c[1]]);
308            let f32_bits = (bits as u32) << 16;
309            f32_bits.to_le_bytes()
310        })
311        .collect();
312    let ptr = gpu.alloc(f32_buf.len())?;
313    gpu.copy_h2d(&f32_buf, ptr)?;
314    Ok(DenseWeight { weight: ptr })
315}
316
317/// 2026-09-25: Load an FP32 tensor into a new BF16 device buffer, converted on
318/// the host with `f32_to_bf16`; any other dtype is an error.
319pub fn dense_f32_as_bf16(
320    store: &WeightStore,
321    name: &str,
322    gpu: &dyn GpuBackend,
323) -> Result<DenseWeight> {
324    let w = store.get(name)?;
325    ensure!(
326        w.dtype == WeightDtype::FP32,
327        "Expected FP32 for {name}, got {:?}",
328        w.dtype
329    );
330    let n = w.num_elements();
331    let mut f32_buf = vec![0u8; n * 4];
332    gpu.copy_d2h(w.ptr, &mut f32_buf)?;
333    let bf16_buf: Vec<u8> = f32_buf
334        .chunks_exact(4)
335        .flat_map(|c| {
336            let val = f32::from_le_bytes([c[0], c[1], c[2], c[3]]);
337            f32_to_bf16(val).to_le_bytes()
338        })
339        .collect();
340    let ptr = gpu.alloc(bf16_buf.len())?;
341    gpu.copy_h2d(&bf16_buf, ptr)?;
342    Ok(DenseWeight { weight: ptr })
343}
344
345#[cfg(test)]
346#[path = "kv_scales_tests.rs"]
347mod kv_scales_tests;
metrale_model_layers/weight_map/model_a.rs

metrale_model_layers/weight_map/
model_a.rs
