atlas_core/
dtype.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3use serde::{Deserialize, Serialize};
4
5/// Supported quantization and precision types for Atlas kernels.
6///
7/// Each variant maps to a specific bit-width and numeric format used by
8/// SM121 tensor cores or CUDA ALU paths.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
10pub enum DType {
11    /// 4-bit floating point (1 sign + 2 exponent + 1 mantissa)
12    /// Used for NVFP4 weight storage. Values: {0, 0.5, 1, 1.5, 2, 3, 4, 6}
13    E2M1,
14
15    /// 8-bit floating point (1 sign + 4 exponent + 3 mantissa)
16    /// Used for NVFP4 block scales and FP8 quantization
17    FP8E4M3,
18
19    /// 8-bit floating point (1 sign + 5 exponent + 2 mantissa)
20    /// Used for some FP8 quantization schemes
21    FP8E5M2,
22
23    /// 16-bit brain floating point
24    BF16,
25
26    /// 16-bit IEEE floating point
27    FP16,
28
29    /// 32-bit IEEE floating point
30    FP32,
31}
32
33impl DType {
34    /// Size in bytes per element. E2M1 is sub-byte (4 bits) but we report
35    /// the packed size (2 elements per byte).
36    pub const fn element_size_bits(&self) -> usize {
37        match self {
38            DType::E2M1 => 4,
39            DType::FP8E4M3 | DType::FP8E5M2 => 8,
40            DType::BF16 | DType::FP16 => 16,
41            DType::FP32 => 32,
42        }
43    }
44
45    /// Number of elements that pack into a 32-bit word.
46    pub const fn elements_per_u32(&self) -> usize {
47        32 / self.element_size_bits()
48    }
49}
50
51/// Quantization configuration for a weight tensor.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct QuantConfig {
54    /// Weight storage type (e.g., E2M1 for NVFP4)
55    pub weight_type: DType,
56
57    /// Scale factor type (e.g., FP8E4M3 for NVFP4 block scales)
58    pub scale_type: DType,
59
60    /// Number of weights per scale factor (block size)
61    pub group_size: usize,
62
63    /// Global scale factor type (FP32 for NVFP4)
64    pub global_scale_type: DType,
65}
66
67impl QuantConfig {
68    /// NVFP4: E2M1 weights with FP8 block scales (group_size=16)
69    pub fn nvfp4() -> Self {
70        Self {
71            weight_type: DType::E2M1,
72            scale_type: DType::FP8E4M3,
73            group_size: 16,
74            global_scale_type: DType::FP32,
75        }
76    }
77
78    /// FP8 per-tensor quantization
79    pub fn fp8() -> Self {
80        Self {
81            weight_type: DType::FP8E4M3,
82            scale_type: DType::FP32,
83            group_size: 0, // per-tensor
84            global_scale_type: DType::FP32,
85        }
86    }
87}
atlas_core/dtype.rs

atlas_core/
dtype.rs
