metrale_model_layers/weight_map/
quantize_fns.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: FP8 E4M3 with per-row scales: load-time quantization from BF16 and the checkpoint loader.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - Every returned FP8 weight carries one FP32 scale per row.
8
9#![allow(unused_imports)]
10
11use anyhow::{Context, Result, bail, ensure};
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13use metrale_model_weights::weights::{WeightDtype, WeightStore};
14
15use super::*;
16
17/// 2026-09-25: Quantize a BF16 dense weight `[n, k]` to FP8 E4M3 (`n*k` bytes)
18/// with one FP32 scale per row. It synchronizes `stream`, so it belongs at load
19/// time, not in a forward pass.
20pub fn quantize_to_fp8(
21    bf16_weight: &DenseWeight,
22    n: usize,
23    k: usize,
24    gpu: &dyn GpuBackend,
25    quantize_kernel: metrale_gpu_runtime::gpu::KernelHandle,
26    stream: u64,
27) -> Result<Fp8DenseWeight> {
28    use metrale_gpu_runtime::kernel_args::KernelLaunch;
29
30    let fp8_buf = gpu.alloc(n * k)?;
31    let scale_buf = gpu.alloc(n * 4)?;
32
33    KernelLaunch::new(gpu, quantize_kernel)
34        .grid([n as u32, 1, 1])
35        .block([256, 1, 1])
36        .arg_ptr(bf16_weight.weight)
37        .arg_ptr(fp8_buf)
38        .arg_ptr(scale_buf)
39        .arg_u32(n as u32)
40        .arg_u32(k as u32)
41        .launch(stream)?;
42
43    gpu.synchronize(stream)?;
44
45    Ok(Fp8DenseWeight {
46        weight: fp8_buf,
47        row_scale: scale_buf,
48    })
49}
50
51/// 2026-09-25: Load an FP8 E4M3 checkpoint weight with per-row scales as an
52/// [`Fp8Weight`] tagged `Fp8PerRow`.
53///
54/// Expects `{name}.weight` FP8E4M3 `[N, K]` and `{name}.weight_scale` `[N]`.
55/// The weight is the store's pointer. An FP32 scale is the store's pointer
56/// too; a BF16 scale is widened on the host into a new buffer; any other dtype
57/// is an error.
58pub fn load_fp8_weight(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result<Fp8Weight> {
59    let w = store.get(&format!("{name}.weight"))?;
60    ensure!(
61        w.dtype == WeightDtype::FP8E4M3,
62        "Expected FP8E4M3 for {name}.weight, got {:?}",
63        w.dtype,
64    );
65    ensure!(
66        w.shape.len() == 2,
67        "Expected 2D weight for {name}, got {:?}",
68        w.shape
69    );
70    let n = w.shape[0];
71    let k = w.shape[1];
72
73    let weight_ptr = w.ptr;
74
75    let scale_key = format!("{name}.weight_scale");
76    let s = store.get(&scale_key).with_context(|| {
77        format!("Missing per-row scale tensor {scale_key} for FP8 weight {name}")
78    })?;
79    ensure!(
80        s.shape.len() == 1 && s.shape[0] == n,
81        "Expected [{n}] shape for {scale_key}, got {:?}",
82        s.shape,
83    );
84
85    let row_scale_ptr = if s.dtype == WeightDtype::FP32 {
86        s.ptr
87    } else if s.dtype == WeightDtype::BF16 {
88        let mut bf16_buf = vec![0u8; n * 2];
89        gpu.copy_d2h(s.ptr, &mut bf16_buf)?;
90        let mut f32_buf = vec![0u8; n * 4];
91        for i in 0..n {
92            let bf16_bytes = [bf16_buf[i * 2], bf16_buf[i * 2 + 1]];
93            let val = bf16_bytes_to_f32(bf16_bytes);
94            let f32_bytes = val.to_le_bytes();
95            f32_buf[i * 4..i * 4 + 4].copy_from_slice(&f32_bytes);
96        }
97        let f32_ptr = gpu.alloc(n * 4)?;
98        gpu.copy_h2d(&f32_buf, f32_ptr)?;
99        f32_ptr
100    } else {
101        anyhow::bail!(
102            "Unsupported dtype {:?} for {scale_key}, expected FP32 or BF16",
103            s.dtype,
104        );
105    };
106
107    Ok(Fp8Weight {
108        weight: weight_ptr,
109        row_scale: row_scale_ptr,
110        n: n as u32,
111        k: k as u32,
112        scale_format: WeightQuantFormat::Fp8PerRow,
113    })
114}
metrale_model_layers/weight_map/quantize_fns.rs

metrale_model_layers/weight_map/
quantize_fns.rs
