metrale_model_layers/weight_map/
quantize_fp8_bs.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Load-time BF16 to FP8 E4M3 quantization with one FP32 scale per 128x128 block.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - `FP8_BLOCK` equals `FP8_BLOCK` in `kernels/gb10/common/moe_fp8_grouped_gemm.cu`.
8//!
9//! The FP8 grouped MoE GEMM reads block scales `[N/128, K/128]`, not per-row
10//! ones, and cannot tell the two apart by shape; `Fp8Weight::scale_format`
11//! records which one a weight carries.
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::GpuBackend;
15use metrale_gpu_runtime::kernel_args::KernelLaunch;
16
17use super::{DenseWeight, Fp8Weight, WeightQuantFormat};
18
19/// 2026-09-25: Elements per block scale, on both axes.
20const FP8_BLOCK: usize = 128;
21
22/// 2026-09-25: Quantize an `[n, k]` BF16 dense weight to block-scaled FP8 E4M3
23/// on the GPU, as an `Fp8Weight` tagged `Fp8BlockScaled` with FP32 scales
24/// `[ceil(n/128), ceil(k/128)]`.
25///
26/// It does not synchronize `stream` and does not free the BF16 source; the
27/// caller must synchronize before freeing it.
28pub fn quantize_to_fp8_blockscaled(
29    bf16_weight: &DenseWeight,
30    n: usize,
31    k: usize,
32    gpu: &dyn GpuBackend,
33    quantize_kernel: metrale_gpu_runtime::gpu::KernelHandle,
34    stream: u64,
35) -> Result<Fp8Weight> {
36    anyhow::ensure!(
37        n > 0 && k > 0,
38        "quantize_to_fp8_blockscaled: empty [{n},{k}]"
39    );
40
41    let n_blocks = n.div_ceil(FP8_BLOCK);
42    let k_blocks = k.div_ceil(FP8_BLOCK);
43
44    let weight_buf = gpu.alloc(n * k)?;
45    let scale_buf = gpu.alloc(n_blocks * k_blocks * 4)?;
46
47    KernelLaunch::new(gpu, quantize_kernel)
48        .grid([k_blocks as u32, n_blocks as u32, 1])
49        .block([256, 1, 1])
50        .arg_ptr(bf16_weight.weight)
51        .arg_ptr(weight_buf)
52        .arg_ptr(scale_buf)
53        .arg_u32(n as u32)
54        .arg_u32(k as u32)
55        .launch(stream)?;
56
57    Ok(Fp8Weight {
58        weight: weight_buf,
59        row_scale: scale_buf,
60        n: n as u32,
61        k: k as u32,
62        scale_format: WeightQuantFormat::Fp8BlockScaled,
63    })
64}
metrale_model_layers/weight_map/quantize_fp8_bs.rs

metrale_model_layers/weight_map/
quantize_fp8_bs.rs
