spark_runtime/kv_dequant/
luts.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The dequantization codebooks, and the group size they are indexed in.
4//!
5//! Split from the dequant routines that read them purely for the file-size
6//! cap: the tables are data, the routines are the loops over that data, and
7//! the seam falls between them cleanly. Every table matches its CUDA-side
8//! counterpart byte-for-byte — see the parent module's table for which
9//! kernel each mirrors.
10
11/// Group size for per-group FP8 scales. Matches `NVFP4_GROUP_SIZE` in the
12/// per-quant attention kernels.
13pub const NVFP4_GROUP_SIZE: usize = 16;
14
15/// E2M1 4-bit codebook (NVFP4). Matches
16/// `kernels/gb10/common/paged_decode_attn_nvfp4.cu:118`.
17pub const NVFP4_E2M1_LUT: [f32; 16] = [
18    0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
19];
20
21/// Turbo4 16-level Lloyd-Max codebook. Matches
22/// `kernels/gb10/common/paged_decode_attn_turbo4.cu:121`.
23pub const TURBO4_LUT: [f32; 16] = [
24    -2.7326, -2.0690, -1.6180, -1.2562, -0.9423, -0.6568, -0.3880, -0.1284, 0.1284, 0.3880, 0.6568,
25    0.9423, 1.2562, 1.6180, 2.0690, 2.7326,
26];
27
28/// Turbo3 8-level Lloyd-Max codebook. Matches
29/// `kernels/gb10/common/paged_decode_attn_turbo3.cu:137`.
30pub const TURBO3_LUT: [f32; 8] = [
31    -2.1520, -1.3440, -0.7560, -0.2451, 0.2451, 0.7560, 1.3440, 2.1520,
32];
33
34/// E4M3 → f32 LUT (256 entries).
35///
36/// A `const`, not a lazily-filled static: the table is pure arithmetic over a
37/// fixed domain, so it is computed at compile time and there is no runtime
38/// state to initialise, synchronise or invalidate.
39const E4M3_LUT: [f32; 256] = {
40    let mut lut = [0.0f32; 256];
41    let mut byte = 0u32;
42    while byte < 256 {
43        let sign_bit = (byte >> 7) & 1;
44        let exp = ((byte >> 3) & 0xF) as i32;
45        let mant = byte & 0x7;
46        let s: f32 = if sign_bit == 0 { 1.0 } else { -1.0 };
47        lut[byte as usize] = if exp == 0 {
48            if mant == 0 {
49                s * 0.0
50            } else {
51                s * (mant as f32) * exp2(-9)
52            }
53        } else if exp == 0xF && mant == 0x7 {
54            f32::NAN
55        } else {
56            s * exp2(exp - 7) * (1.0 + (mant as f32) / 8.0)
57        };
58        byte += 1;
59    }
60    lut
61};
62
63/// `2^e`, for the small integer exponents this table needs (`-9..=8`).
64///
65/// `f32::powi` is not a const fn. Repeated multiply/divide is EXACT here rather
66/// than merely close: powers of two are representable exactly in binary
67/// floating point, so every step is lossless and the result is bit-identical to
68/// `powi`. A test pins that.
69const fn exp2(e: i32) -> f32 {
70    let mut v = 1.0f32;
71    let mut i = 0i32;
72    if e >= 0 {
73        while i < e {
74            v *= 2.0;
75            i += 1;
76        }
77    } else {
78        while i < -e {
79            v /= 2.0;
80            i += 1;
81        }
82    }
83    v
84}
85
86/// Borrow the compile-time table.
87pub fn e4m3_lut() -> &'static [f32; 256] {
88    &E4M3_LUT
89}
90
91#[cfg(test)]
92mod lut_tests {
93    use super::*;
94
95    /// The table moved from a lazily-filled `OnceLock` to a `const`, which
96    /// meant replacing `f32::powi` (not const) with `exp2`. This proves the
97    /// substitution changed no value: every entry must be bit-identical to the
98    /// `powi` formula it replaced.
99    #[test]
100    fn the_const_table_is_bit_identical_to_the_powi_formula() {
101        for byte in 0..256u32 {
102            let sign_bit = (byte >> 7) & 1;
103            let exp = ((byte >> 3) & 0xF) as i32;
104            let mant = byte & 0x7;
105            let s: f32 = if sign_bit == 0 { 1.0 } else { -1.0 };
106            let expected: f32 = if exp == 0 {
107                if mant == 0 {
108                    s * 0.0
109                } else {
110                    s * (mant as f32) * 2.0f32.powi(-9)
111                }
112            } else if exp == 0xF && mant == 0x7 {
113                f32::NAN
114            } else {
115                s * 2.0f32.powi(exp - 7) * (1.0 + (mant as f32) / 8.0)
116            };
117            let got = E4M3_LUT[byte as usize];
118            if expected.is_nan() {
119                assert!(got.is_nan(), "byte {byte}: expected NaN, got {got}");
120            } else {
121                assert_eq!(
122                    got.to_bits(),
123                    expected.to_bits(),
124                    "byte {byte}: {got} != {expected}"
125                );
126            }
127        }
128    }
129
130    #[test]
131    fn exp2_matches_powi_across_the_domain() {
132        for e in -9..=8i32 {
133            assert_eq!(exp2(e).to_bits(), 2.0f32.powi(e).to_bits(), "2^{e}");
134        }
135    }
136}
spark_runtime/kv_dequant/luts.rs

spark_runtime/kv_dequant/
luts.rs
