metrale_cache/kv_dequant/
luts.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The dequantization codebooks and the scale group size.
4//!
5//! Owner: cache.
6//! Invariants: each codebook holds the same values as the `lut_init` array of
7//! the decode kernel named on it (the kernel keeps the turbo codebooks as
8//! `__half`).
9
10/// 2026-09-25: Elements per scale group in every grouped layout; the
11/// `NVFP4_GROUP_SIZE` of `kernels/gb10/common/paged_decode_attn_nvfp4.cu`.
12pub const NVFP4_GROUP_SIZE: usize = 16;
13
14/// 2026-09-25: E2M1 4-bit codebook (NVFP4), as in
15/// `kernels/gb10/common/paged_decode_attn_nvfp4.cu`.
16pub const NVFP4_E2M1_LUT: [f32; 16] = [
17    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,
18];
19
20/// 2026-09-25: Turbo4 16-level Lloyd-Max codebook, as in
21/// `kernels/gb10/common/paged_decode_attn_turbo4.cu`.
22pub const TURBO4_LUT: [f32; 16] = [
23    -2.7326, -2.0690, -1.6180, -1.2562, -0.9423, -0.6568, -0.3880, -0.1284, 0.1284, 0.3880, 0.6568,
24    0.9423, 1.2562, 1.6180, 2.0690, 2.7326,
25];
26
27/// 2026-09-25: Turbo3 8-level Lloyd-Max codebook, as in
28/// `kernels/gb10/common/paged_decode_attn_turbo3.cu`.
29pub const TURBO3_LUT: [f32; 8] = [
30    -2.1520, -1.3440, -0.7560, -0.2451, 0.2451, 0.7560, 1.3440, 2.1520,
31];
32
33/// 2026-09-25: FP8 E4M3 byte to f32, all 256 bytes, computed at compile
34/// time. `0x7F` and `0xFF` are NaN.
35const E4M3_LUT: [f32; 256] = {
36    let mut lut = [0.0f32; 256];
37    let mut byte = 0u32;
38    while byte < 256 {
39        let sign_bit = (byte >> 7) & 1;
40        let exp = ((byte >> 3) & 0xF) as i32;
41        let mant = byte & 0x7;
42        let s: f32 = if sign_bit == 0 { 1.0 } else { -1.0 };
43        lut[byte as usize] = if exp == 0 {
44            if mant == 0 {
45                s * 0.0
46            } else {
47                s * (mant as f32) * exp2(-9)
48            }
49        } else if exp == 0xF && mant == 0x7 {
50            f32::NAN
51        } else {
52            s * exp2(exp - 7) * (1.0 + (mant as f32) / 8.0)
53        };
54        byte += 1;
55    }
56    lut
57};
58
59/// 2026-09-25: `2^e` for the exponents the table uses (`-9..=8`), because
60/// `f32::powi` is not a const fn. Each multiply or divide by 2 is exact in
61/// that range, so the result is bit-identical to `powi`; a test checks it.
62const fn exp2(e: i32) -> f32 {
63    let mut v = 1.0f32;
64    let mut i = 0i32;
65    if e >= 0 {
66        while i < e {
67            v *= 2.0;
68            i += 1;
69        }
70    } else {
71        while i < -e {
72            v /= 2.0;
73            i += 1;
74        }
75    }
76    v
77}
78
79/// 2026-09-25: The E4M3 table.
80pub fn e4m3_lut() -> &'static [f32; 256] {
81    &E4M3_LUT
82}
83
84#[cfg(test)]
85mod lut_tests {
86    use super::*;
87
88    /// 2026-09-25: Every entry is bit-identical to the same formula computed
89    /// with `f32::powi` (NaN entries must both be NaN).
90    #[test]
91    fn the_const_table_is_bit_identical_to_the_powi_formula() {
92        for byte in 0..256u32 {
93            let sign_bit = (byte >> 7) & 1;
94            let exp = ((byte >> 3) & 0xF) as i32;
95            let mant = byte & 0x7;
96            let s: f32 = if sign_bit == 0 { 1.0 } else { -1.0 };
97            let expected: f32 = if exp == 0 {
98                if mant == 0 {
99                    s * 0.0
100                } else {
101                    s * (mant as f32) * 2.0f32.powi(-9)
102                }
103            } else if exp == 0xF && mant == 0x7 {
104                f32::NAN
105            } else {
106                s * 2.0f32.powi(exp - 7) * (1.0 + (mant as f32) / 8.0)
107            };
108            let got = E4M3_LUT[byte as usize];
109            if expected.is_nan() {
110                assert!(got.is_nan(), "byte {byte}: expected NaN, got {got}");
111            } else {
112                assert_eq!(
113                    got.to_bits(),
114                    expected.to_bits(),
115                    "byte {byte}: {got} != {expected}"
116                );
117            }
118        }
119    }
120
121    #[test]
122    fn exp2_matches_powi_across_the_domain() {
123        for e in -9..=8i32 {
124            assert_eq!(exp2(e).to_bits(), 2.0f32.powi(e).to_bits(), "2^{e}");
125        }
126    }
127}
metrale_cache/kv_dequant/luts.rs

metrale_cache/kv_dequant/
luts.rs
