spark_model/weight_map/
quantize_fp8_bs.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Runtime BF16 → FP8 (E4M3) weight quantization with 128×128 block scales.
4//!
5//! The third runtime quantizer, and the one that exists for MoE experts.
6//! Atlas already had:
7//!   - BF16 → NVFP4 (`quantize_to_nvfp4`) — 4 bits, what `Bf16Raw` models get
8//!   - BF16 → FP8 per-ROW (`quantize_bf16_to_fp8`) — for the DFlash head
9//!
10//! Neither fits routed experts on a plain-BF16 checkpoint. NVFP4 costs real
11//! output quality, and the FP8 grouped MoE GEMM
12//! (`moe_fp8_grouped_gemm.cu`) reads BLOCK scales `[N/128, K/128]` FP32,
13//! not per-row ones. Feeding it a per-row buffer is not a shape error the
14//! kernel can detect — it would index a shorter array and silently dequant
15//! with the wrong scale, which is why `Fp8Weight::scale_format` is tagged
16//! and asserted at dispatch.
17//!
18//! Why block-scaled FP8 rather than BF16 for experts: on LongCat-Flash-Lite
19//! the routed experts are 63.0 GB of the 70.2 GB resident. BF16 does not fit
20//! (+47 GB against a 97.3 GB budget at 0.80 util); FP8 does (+15.75 GB). And
21//! because MoE reads only top-12 of 256 experts per token, the DECODE cost is
22//! +0.74 GB/token — less than either of the dense BF16 levers.
23
24use anyhow::Result;
25use spark_runtime::gpu::GpuBackend;
26use spark_runtime::kernel_args::KernelLaunch;
27
28use super::{DenseWeight, Fp8Weight, WeightQuantFormat};
29
30/// Elements per block scale, both axes. Must match `FP8_BLOCK` in
31/// `moe_fp8_grouped_gemm.cu` — the consumer hardcodes 128.
32const FP8_BLOCK: usize = 128;
33
34/// Quantize an `[n, k]` BF16 dense weight to block-scaled FP8 E4M3 on GPU.
35///
36/// Returns a `Fp8Weight` tagged `Fp8BlockScaled`, laid out exactly as the
37/// on-disk Qwen FP8 releases are after widening, so every consumer that
38/// already accepts those accepts this with no change.
39///
40/// Called once per projection at load time, never on the hot path. The BF16
41/// source is the caller's to free — this does not take ownership.
42pub fn quantize_to_fp8_blockscaled(
43    bf16_weight: &DenseWeight,
44    n: usize,
45    k: usize,
46    gpu: &dyn GpuBackend,
47    quantize_kernel: spark_runtime::gpu::KernelHandle,
48    stream: u64,
49) -> Result<Fp8Weight> {
50    anyhow::ensure!(
51        n > 0 && k > 0,
52        "quantize_to_fp8_blockscaled: empty [{n},{k}]"
53    );
54
55    let n_blocks = n.div_ceil(FP8_BLOCK);
56    let k_blocks = k.div_ceil(FP8_BLOCK);
57
58    // One byte per weight, one f32 per [128,128] tile.
59    let weight_buf = gpu.alloc(n * k)?;
60    let scale_buf = gpu.alloc(n_blocks * k_blocks * 4)?;
61
62    KernelLaunch::new(gpu, quantize_kernel)
63        .grid([k_blocks as u32, n_blocks as u32, 1])
64        .block([256, 1, 1])
65        .arg_ptr(bf16_weight.weight)
66        .arg_ptr(weight_buf)
67        .arg_ptr(scale_buf)
68        .arg_u32(n as u32)
69        .arg_u32(k as u32)
70        .launch(stream)?;
71
72    Ok(Fp8Weight {
73        weight: weight_buf,
74        row_scale: scale_buf,
75        n: n as u32,
76        k: k as u32,
77        scale_format: WeightQuantFormat::Fp8BlockScaled,
78    })
79}
spark_model/weight_map/quantize_fp8_bs.rs

spark_model/weight_map/
quantize_fp8_bs.rs
