metrale_model_layers/weight_map/
quant_helpers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: FP8-to-BF16 dequantization, dtype-dispatching dense loads, and the compressed-tensors NVFP4 loader.
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: Convert FP8 E4M3 bytes to little-endian BF16 bytes on the host,
17/// multiplying each value by `scale`.
18pub(super) fn dequant_fp8_bytes_to_bf16(fp8_buf: &[u8], scale: f32) -> Vec<u8> {
19    fp8_buf
20        .iter()
21        .flat_map(|&byte| {
22            let val = fp8_e4m3_to_f32(byte) * scale;
23            f32_to_bf16(val).to_le_bytes()
24        })
25        .collect()
26}
27
28/// 2026-09-25: Dequantize a block-scaled FP8 E4M3 weight `{prefix}.weight`
29/// `[N, K]` into a new BF16 device buffer:
30/// `bf16[i,j] = fp8[i,j] * scale[i/block_n, j/block_k]`, with the block sizes
31/// inferred from the scale's shape.
32///
33/// A `weight_scale_inv`, or else a 2-D `weight_scale`, must be BF16 or FP32 and
34/// is applied on the device by `dequant_fp8_blockscaled_bf16`. Otherwise a 1-D
35/// `weight_scale` or a `.scale` (FP32, BF16 or E8M0) is applied on the host.
36/// Having none of them is an error.
37pub fn dequant_fp8_blockscaled_to_bf16(
38    store: &WeightStore,
39    prefix: &str,
40    gpu: &dyn GpuBackend,
41) -> Result<DenseWeight> {
42    use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
43
44    let w = store.get(&format!("{prefix}.weight"))?;
45    ensure!(
46        w.dtype == WeightDtype::FP8E4M3,
47        "Expected FP8E4M3 for {prefix}.weight, got {:?}",
48        w.dtype,
49    );
50    ensure!(
51        w.shape.len() == 2,
52        "Expected 2D weight for {prefix}, got {:?}",
53        w.shape
54    );
55    let n = w.shape[0];
56    let k = w.shape[1];
57    let total = n * k;
58    let byte_size = w.byte_size();
59    ensure!(
60        total == byte_size,
61        "FP8 size mismatch: total={total} byte_size={byte_size}"
62    );
63
64    // 2026-09-25: The kernel reads a 2-D BF16 or FP32 scale only, so E8M0
65    // `.scale` and 1-D `weight_scale` fall through to the host path. Keep this a
66    // fall-through: a weight with only `.scale` has no `weight_scale*` key.
67    let gpu_scale = store
68        .get(&format!("{prefix}.weight_scale_inv"))
69        .ok()
70        .or_else(|| {
71            store
72                .get(&format!("{prefix}.weight_scale"))
73                .ok()
74                .filter(|s| s.shape.len() == 2)
75        });
76    if let Some(s) = gpu_scale {
77        ensure!(
78            s.dtype == WeightDtype::BF16 || s.dtype == WeightDtype::FP32,
79            "Expected BF16 or FP32 for {prefix} GPU block scale, got {:?}",
80            s.dtype,
81        );
82        let sn = s.shape[0];
83        let sk = s.shape[1];
84        let block_n = (n / sn) as u32;
85        let block_k = (k / sk) as u32;
86        let scale_is_f32 = s.dtype == WeightDtype::FP32;
87
88        let out = gpu.alloc(total * 2)?;
89
90        let stream = gpu.default_stream();
91        let kernel = gpu.kernel(
92            "dequant_fp8_blockscaled_bf16",
93            "dequant_fp8_blockscaled_bf16",
94        )?;
95        KernelLaunch::new(gpu, kernel)
96            .grid([div_ceil(k as u32, 64), div_ceil(n as u32, 4), 1])
97            .block([64, 4, 1])
98            .arg_ptr(w.ptr)
99            .arg_ptr(s.ptr)
100            .arg_ptr(out)
101            .arg_u32(n as u32)
102            .arg_u32(k as u32)
103            .arg_u32(block_n)
104            .arg_u32(block_k)
105            .arg_u32(sk as u32)
106            .arg_u32(scale_is_f32 as u32)
107            .launch(stream)?;
108        // 2026-09-25: Not synchronized: `out` is written on the default stream,
109        // and a reader on another stream must synchronize it first.
110
111        tracing::debug!(
112            "GPU-dequanted FP8 blockscaled {prefix}: [{n}, {k}] block=[{block_n}, {block_k}] → BF16",
113        );
114        return Ok(DenseWeight { weight: out });
115    }
116
117    // 2026-09-25: Host path: download the FP8 weight and the scale, dequantize,
118    // upload the BF16 result.
119    enum ScaleDtype {
120        Fp32,
121        Bf16,
122        E8M0,
123    }
124    let mut fp8_buf = vec![0u8; total];
125    gpu.copy_d2h(w.ptr, &mut fp8_buf).with_context(|| {
126        format!(
127            "D2H failed for {prefix}.weight: ptr={}, size={total}",
128            w.ptr.0
129        )
130    })?;
131
132    let (scale_buf, _sn, sk, block_n, block_k, scale_dtype) = if let Ok(s) =
133        store.get(&format!("{prefix}.weight_scale"))
134    {
135        ensure!(
136            s.dtype == WeightDtype::BF16 || s.dtype == WeightDtype::FP32,
137            "Expected BF16 or FP32 2-D block scale for {prefix}.weight_scale, got {:?}",
138            s.dtype,
139        );
140        let rank = s.shape.len();
141        let (sn, sk) = if rank == 2 {
142            (s.shape[0], s.shape[1])
143        } else if rank == 1 {
144            (s.shape[0], 1)
145        } else {
146            bail!(
147                "Expected 1-D or 2-D scale for {prefix}.weight_scale, got shape {:?}",
148                s.shape
149            );
150        };
151        let block_n = if sn > 1 { n / sn } else { n };
152        let block_k = if sk > 1 { k / sk } else { k };
153        let scale_is_f32 = s.dtype == WeightDtype::FP32;
154        let scale_bytes_per = if scale_is_f32 { 4 } else { 2 };
155        let mut buf = vec![0u8; sn * sk * scale_bytes_per];
156        gpu.copy_d2h(s.ptr, &mut buf).with_context(|| {
157            format!(
158                "D2H failed for {prefix}.weight_scale: ptr={}, size={}",
159                s.ptr.0,
160                sn * sk * scale_bytes_per
161            )
162        })?;
163        let sd = if scale_is_f32 {
164            ScaleDtype::Fp32
165        } else {
166            ScaleDtype::Bf16
167        };
168        (buf, sn, sk, block_n, block_k, sd)
169    } else if let Ok(s) = store.get(&format!("{prefix}.scale")) {
170        let rank = s.shape.len();
171        let (sn, sk) = if rank == 2 {
172            (s.shape[0], s.shape[1])
173        } else if rank == 1 {
174            (s.shape[0], 1)
175        } else {
176            bail!(
177                "Expected 1-D or 2-D scale for {prefix}.scale, got shape {:?}",
178                s.shape
179            );
180        };
181        let block_n = if sn > 1 { n / sn } else { n };
182        let block_k = if sk > 1 { k / sk } else { k };
183        let sd = match s.dtype {
184            WeightDtype::FP32 => ScaleDtype::Fp32,
185            WeightDtype::BF16 => ScaleDtype::Bf16,
186            WeightDtype::FP8E8M0 => ScaleDtype::E8M0,
187            other => bail!(
188                "Expected FP32, BF16, or FP8E8M0 for {prefix}.scale, got {:?}",
189                other,
190            ),
191        };
192        let scale_bytes_per = s.dtype.byte_size();
193        let mut buf = vec![0u8; sn * sk * scale_bytes_per];
194        gpu.copy_d2h(s.ptr, &mut buf).with_context(|| {
195            format!(
196                "D2H failed for {prefix}.scale: ptr={}, size={}",
197                s.ptr.0,
198                sn * sk * scale_bytes_per
199            )
200        })?;
201        (buf, sn, sk, block_n, block_k, sd)
202    } else {
203        bail!(
204            "FP8 tensor {prefix}: no .weight_scale_inv, .weight_scale, or .scale found for dequant"
205        );
206    };
207
208    let mut bf16_out = vec![0u8; total * 2];
209    for row in 0..n {
210        let scale_row = row / block_n;
211        for col in 0..k {
212            let scale_col = col / block_k;
213            let scale_idx = scale_row * sk + scale_col;
214            let scale_f32 = match scale_dtype {
215                ScaleDtype::E8M0 => fp8_e8m0_to_f32(scale_buf[scale_idx]),
216                ScaleDtype::Fp32 => {
217                    let b = [
218                        scale_buf[scale_idx * 4],
219                        scale_buf[scale_idx * 4 + 1],
220                        scale_buf[scale_idx * 4 + 2],
221                        scale_buf[scale_idx * 4 + 3],
222                    ];
223                    f32::from_le_bytes(b)
224                }
225                ScaleDtype::Bf16 => {
226                    let b = [scale_buf[scale_idx * 2], scale_buf[scale_idx * 2 + 1]];
227                    bf16_bytes_to_f32(b)
228                }
229            };
230
231            let fp8_byte = fp8_buf[row * k + col];
232            let val = fp8_e4m3_to_f32(fp8_byte) * scale_f32;
233            let bf16_val = f32_to_bf16(val);
234            let out_idx = (row * k + col) * 2;
235            let [lo, hi] = bf16_val.to_le_bytes();
236            bf16_out[out_idx] = lo;
237            bf16_out[out_idx + 1] = hi;
238        }
239    }
240
241    let out = gpu.alloc(bf16_out.len())?;
242    gpu.copy_h2d(&bf16_out, out)?;
243    tracing::debug!(
244        "CPU-dequanted FP8 blockscaled {prefix}: [{n}, {k}] block=[{block_n}, {block_k}] → BF16",
245    );
246    Ok(DenseWeight { weight: out })
247}
248
249/// 2026-09-25: Widen little-endian BF16 bytes to f32.
250pub(super) fn bf16_bytes_to_f32(bytes: [u8; 2]) -> f32 {
251    let bits = u16::from_le_bytes(bytes);
252    f32::from_bits((bits as u32) << 16)
253}
254
255/// 2026-09-25: Load `name` as a BF16 `DenseWeight`, converting by dtype:
256/// - BF16: the store pointer itself;
257/// - FP32: `dense_f32_safe`;
258/// - FP8E4M3: `dequant_fp8_blockscaled_to_bf16` when the weight has a
259///   `weight_scale_inv`, a multi-element `weight_scale` or a `.scale`, else the
260///   per-tensor `dequant_fp8_to_bf16`;
261/// - UInt8 (packed NVFP4, `[n, k/2]`): `dequant_nvfp4_to_bf16`.
262///
263/// Any other dtype is an error. FP8 and NVFP4 names must end in `.weight`.
264pub fn dense_auto(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result<DenseWeight> {
265    let w = store.get(name)?;
266    match w.dtype {
267        WeightDtype::BF16 => Ok(DenseWeight { weight: w.ptr }),
268        WeightDtype::FP32 => dense_f32_safe(store, name, gpu),
269        WeightDtype::FP8E4M3 => {
270            let prefix = name
271                .strip_suffix(".weight")
272                .ok_or_else(|| anyhow::anyhow!("FP8 tensor {name} doesn't end with .weight"))?;
273            // 2026-09-25: A one-element `weight_scale` is a per-tensor scale; every
274            // other scale form is per block.
275            let has_blockscale = store.contains(&format!("{prefix}.weight_scale_inv"));
276            let has_per_row_scale = store
277                .get(&format!("{prefix}.weight_scale"))
278                .map(|s| s.num_elements() > 1)
279                .unwrap_or(false);
280            let has_v4_scale = store.contains(&format!("{prefix}.scale"));
281            if has_blockscale || has_per_row_scale || has_v4_scale {
282                dequant_fp8_blockscaled_to_bf16(store, prefix, gpu)
283            } else {
284                dequant_fp8_to_bf16(store, prefix, gpu)
285            }
286        }
287        WeightDtype::UInt8 => {
288            let prefix = name
289                .strip_suffix(".weight")
290                .ok_or_else(|| anyhow::anyhow!("NVFP4 tensor {name} doesn't end with .weight"))?;
291            if w.shape.len() != 2 {
292                anyhow::bail!(
293                    "dense_auto: packed NVFP4 {name} must be 2-D, got {:?}",
294                    w.shape
295                );
296            }
297            crate::weight_map::dequant_nvfp4_to_bf16(store, prefix, w.shape[0], w.shape[1] * 2, gpu)
298        }
299        other => anyhow::bail!("dense_auto: unsupported dtype {:?} for {name}", other),
300    }
301}
302
303/// 2026-09-25: Build a `QuantizedWeight` from compressed-tensors NVFP4 keys:
304/// `weight_packed`, `weight_scale`, `weight_global_scale` and an optional
305/// `input_global_scale`. compressed-tensors stores the global scale as the
306/// reciprocal of `weight_scale_2`, so `weight_scale_2 = 1 / weight_global_scale`.
307pub fn quantized_v2(
308    store: &WeightStore,
309    prefix: &str,
310    gpu: &dyn GpuBackend,
311) -> Result<QuantizedWeight> {
312    let raw_global_scale = scalar_f32(store, &format!("{prefix}.weight_global_scale"), gpu)?;
313    // 2026-09-25: A zero or non-finite global scale is a load error, because its
314    // reciprocal would be stored as `weight_scale_2`.
315    if !raw_global_scale.is_finite() || raw_global_scale.abs() < f32::MIN_POSITIVE {
316        anyhow::bail!(
317            "{prefix}.weight_global_scale is non-finite or zero ({raw_global_scale}); \
318             checkpoint likely corrupted"
319        );
320    }
321    Ok(QuantizedWeight {
322        weight: ptr(store, &format!("{prefix}.weight_packed"))?,
323        weight_scale: ptr(store, &format!("{prefix}.weight_scale"))?,
324        weight_scale_2: 1.0 / raw_global_scale,
325        // 2026-09-25: Weight-only checkpoints have no `input_global_scale`; then NULL.
326        input_scale: ptr(store, &format!("{prefix}.input_global_scale")).unwrap_or(DevicePtr::NULL),
327        weight_scale_2_vec: DevicePtr::NULL,
328        act: metrale_config::Nvfp4Act::Unstamped,
329    })
330}
metrale_model_layers/weight_map/quant_helpers.rs

metrale_model_layers/weight_map/
quant_helpers.rs
