spark_model/quant_format/
fp8_blockscaled.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! FP8 E4M3 block-scaled serialization (DeepSeek-V3 / Qwen3.5-35B-A3B-FP8
4//! convention).
5//!
6//! | field             | tensor name           | dtype          |
7//! | ----------------- | --------------------- | -------------- |
8//! | FP8 payload       | `.weight`             | float8_e4m3    |
9//! | block scales      | `.weight_scale_inv`   | bf16 [N/BS,K/BS] |
10//!
11//! Atlas consumes these via runtime BF16→NVFP4 re-quantization inside
12//! [`crate::weight_map::quantized_from_fp8`]. The `ignore_modules` list,
13//! when present, flags modules that should stay BF16 after dequant
14//! (skipping the NVFP4 step).
15//!
16//! Maps to the existing [`Nvfp4Variant::Fp8Dequanted`] dispatch.
17
18use crate::quant_format::{QuantFormat, module_matches_pattern};
19use crate::weight_map::Nvfp4Variant;
20
21/// FP8 block-scaled checkpoint.
22#[derive(Debug)]
23pub struct Fp8BlockScaledFormat {
24    pub ignore_modules: Vec<String>,
25}
26
27impl Fp8BlockScaledFormat {
28    pub fn new(ignore_modules: Vec<String>) -> Self {
29        Self { ignore_modules }
30    }
31}
32
33impl QuantFormat for Fp8BlockScaledFormat {
34    fn name(&self) -> &'static str {
35        "fp8-blockscaled"
36    }
37
38    fn base_variant(&self) -> Nvfp4Variant {
39        Nvfp4Variant::Fp8Dequanted
40    }
41
42    fn is_ignored(&self, module_path: &str) -> bool {
43        self.ignore_modules
44            .iter()
45            .any(|pat| module_matches_pattern(module_path, pat))
46    }
47}