spark_model/weight_map/
quantize_fns.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `weight_map.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::{Context, Result, bail, ensure};
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9use spark_runtime::weights::{WeightDtype, WeightStore};
10
11use super::*;
12
13/// Quantize a BF16 dense weight to FP8 E4M3 on GPU.
14///
15/// Allocates FP8 weight buffer + per-row scale buffer, runs GPU quantization
16/// kernel. Called once at model load time (not on the hot path).
17pub fn quantize_to_fp8(
18    bf16_weight: &DenseWeight,
19    n: usize,
20    k: usize,
21    gpu: &dyn GpuBackend,
22    quantize_kernel: spark_runtime::gpu::KernelHandle,
23    stream: u64,
24) -> Result<Fp8DenseWeight> {
25    use spark_runtime::kernel_args::KernelLaunch;
26
27    // Allocate FP8 weight buffer: N * K bytes
28    let fp8_buf = gpu.alloc(n * k)?;
29    // Allocate per-row scale buffer: N * 4 bytes (f32)
30    let scale_buf = gpu.alloc(n * 4)?;
31
32    // Launch quantization kernel: Grid=(N), Block=(256)
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    // Synchronize to ensure quantization completes before using the buffers
44    gpu.synchronize(stream)?;
45
46    Ok(Fp8DenseWeight {
47        weight: fp8_buf,
48        row_scale: scale_buf,
49    })
50}
51
52/// Load an FP8 E4M3 checkpoint weight with per-row f32 scales.
53///
54/// Expects two tensors in the store:
55///   - `{name}.weight`: FP8E4M3 [N, K] (1 byte per element)
56///   - `{name}.weight_scale`: f32 `[N]` per-row dequant scale
57///
58/// Both are already on GPU from safetensors mmap — no conversion needed.
59/// Returns an [`Fp8Weight`] ready for the `w8a16_gemv` LUT kernel.
60pub fn load_fp8_weight(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result<Fp8Weight> {
61    let w = store.get(&format!("{name}.weight"))?;
62    ensure!(
63        w.dtype == WeightDtype::FP8E4M3,
64        "Expected FP8E4M3 for {name}.weight, got {:?}",
65        w.dtype,
66    );
67    ensure!(
68        w.shape.len() == 2,
69        "Expected 2D weight for {name}, got {:?}",
70        w.shape
71    );
72    let n = w.shape[0];
73    let k = w.shape[1];
74
75    // FP8 weight bytes: already on GPU from WeightStore, 1 byte per element
76    let weight_ptr = w.ptr;
77
78    // Per-row scale: try `.weight_scale` (per-row f32 [N])
79    let scale_key = format!("{name}.weight_scale");
80    let s = store.get(&scale_key).with_context(|| {
81        format!("Missing per-row scale tensor {scale_key} for FP8 weight {name}")
82    })?;
83    ensure!(
84        s.shape.len() == 1 && s.shape[0] == n,
85        "Expected [{n}] shape for {scale_key}, got {:?}",
86        s.shape,
87    );
88
89    // Scale tensor may be BF16 or f32 on disk. If BF16, convert to f32 on CPU.
90    let row_scale_ptr = if s.dtype == WeightDtype::FP32 {
91        // Already f32 on GPU — use directly
92        s.ptr
93    } else if s.dtype == WeightDtype::BF16 {
94        // BF16 → f32 conversion on CPU, upload to GPU
95        let mut bf16_buf = vec![0u8; n * 2];
96        gpu.copy_d2h(s.ptr, &mut bf16_buf)?;
97        let mut f32_buf = vec![0u8; n * 4];
98        for i in 0..n {
99            let bf16_bytes = [bf16_buf[i * 2], bf16_buf[i * 2 + 1]];
100            let val = bf16_bytes_to_f32(bf16_bytes);
101            let f32_bytes = val.to_le_bytes();
102            f32_buf[i * 4..i * 4 + 4].copy_from_slice(&f32_bytes);
103        }
104        let f32_ptr = gpu.alloc(n * 4)?;
105        gpu.copy_h2d(&f32_buf, f32_ptr)?;
106        f32_ptr
107    } else {
108        anyhow::bail!(
109            "Unsupported dtype {:?} for {scale_key}, expected FP32 or BF16",
110            s.dtype,
111        );
112    };
113
114    Ok(Fp8Weight {
115        weight: weight_ptr,
116        row_scale: row_scale_ptr,
117        n: n as u32,
118        k: k as u32,
119        // `load_fp8_weight` reads `.weight_scale` which is shape `[N]` f32.
120        // That's the per-row F32 layout, consumed by `w8a16_gemv` /
121        // `w8a16_gemm`. Tag accordingly so kernel asserts don't panic.
122        scale_format: WeightQuantFormat::Fp8PerRow,
123    })
124}
spark_model/weight_map/quantize_fns.rs

spark_model/weight_map/
quantize_fns.rs
