atlas_core/
mxfp4_e8m0.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! DeepSeek-V4 native MXFP4 host unpack: packed E2M1 nibbles + E8M0 scales.
4//!
5//! This is the SSOT for the CPU loop previously inlined in
6//! `spark_model::weight_map::dequant_nvfp4_e8m0_to_bf16`. K3 official experts
7//! (`weight_packed` + `weight_scale`) call this; they do not grow a second
8//! dequant stack. CUDA `mx_block_scale<true>` must stay byte-exact with
9//! [`fp8_e8m0_to_f32`].
10//!
11//! GPU GEMM: kimi-k3 `{mxfp4,nvfp4}/KERNEL.toml` `[build].extra_cu` points at
12//! `kernels/gb10/deepseek-v4-flash/nvfp4/moe_w4a16_grouped_gemm.cu`
13//! (`moe_w4a16_grouped_gemm_ptrtable_e8m0`). Do not copy it into kimi-k3.
14
15use anyhow::{Result, ensure};
16
17/// E2M1 nibble → f32. Same table as DSV4 `dequant_nvfp4_e8m0_to_bf16`.
18pub const E2M1: [f32; 16] = [
19    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,
20];
21
22/// Native MXFP4 group size on the DSV4 GPU lander (`quantized_mxfp4_e8m0`).
23pub const GROUP_SIZE: usize = 32;
24
25/// FP8 E8M0 → f32 (unsigned exponent, bias 127). exp=0 and exp=255 → 0.0.
26const FP8_E8M0_LUT: [f32; 256] = {
27    let mut table = [0.0f32; 256];
28    let mut i: u32 = 0;
29    while i < 256 {
30        let exp = i as u8;
31        table[i as usize] = if exp == 0 || exp == 255 {
32            0.0f32
33        } else {
34            f32::from_bits((exp as u32) << 23)
35        };
36        i += 1;
37    }
38    table
39};
40
41/// Convert one E8M0 scale byte to f32 (branchless LUT).
42#[inline(always)]
43pub fn fp8_e8m0_to_f32(bits: u8) -> f32 {
44    FP8_E8M0_LUT[bits as usize]
45}
46
47/// Unpack packed E2M1 `[n, k/2]` + E8M0 scales to f32 `[n, k]`.
48///
49/// Block size is `n*k / scales.len()` (DSV4 infers; GPU lander requires 32).
50/// Even flat index = low nibble, odd = high nibble.
51pub fn dequant_nvfp4_e8m0_to_f32(
52    packed: &[u8],
53    scales: &[u8],
54    n: usize,
55    k: usize,
56) -> Result<Vec<f32>> {
57    let total = n.checked_mul(k).expect("mxfp4 n*k");
58    ensure!(
59        total.is_multiple_of(2),
60        "MXFP4 E8M0: n*k={total} is odd (need even nibble count)"
61    );
62    let packed_bytes = total / 2;
63    ensure!(
64        packed.len() == packed_bytes,
65        "MXFP4 E8M0: packed {} B, expected {packed_bytes} for [{n},{k}]",
66        packed.len()
67    );
68    let num_groups = scales.len();
69    ensure!(
70        num_groups > 0 && total.is_multiple_of(num_groups),
71        "MXFP4 E8M0: weight elems {total} not divisible by E8M0 scale groups {num_groups}"
72    );
73    let block = total / num_groups;
74    let mut out = vec![0.0f32; total];
75    for (group, &sb) in scales.iter().enumerate() {
76        let block_scale = fp8_e8m0_to_f32(sb);
77        for elem in 0..block {
78            let flat_idx = group * block + elem;
79            let byte_idx = flat_idx / 2;
80            let nibble = if flat_idx.is_multiple_of(2) {
81                packed[byte_idx] & 0x0F
82            } else {
83                (packed[byte_idx] >> 4) & 0x0F
84            };
85            out[flat_idx] = E2M1[nibble as usize] * block_scale;
86        }
87    }
88    Ok(out)
89}
90
91/// Same unpack, then `f32_to_bf16` — matches DSV4 host upload.
92pub fn dequant_nvfp4_e8m0_to_bf16(
93    packed: &[u8],
94    scales: &[u8],
95    n: usize,
96    k: usize,
97) -> Result<Vec<u16>> {
98    let f = dequant_nvfp4_e8m0_to_f32(packed, scales, n, k)?;
99    Ok(f.into_iter().map(crate::numeric::f32_to_bf16).collect())
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn e8m0_pow2_and_sentinels() {
108        assert_eq!(fp8_e8m0_to_f32(0), 0.0);
109        assert_eq!(fp8_e8m0_to_f32(255), 0.0);
110        assert_eq!(fp8_e8m0_to_f32(127), 1.0);
111        assert_eq!(fp8_e8m0_to_f32(128), 2.0);
112        assert_eq!(fp8_e8m0_to_f32(126), 0.5);
113    }
114
115    #[test]
116    fn matches_dsv4_nibble_order_and_lut() {
117        // Same table as spark_model::weight_map::fp8_lut (DSV4). Low nibble first.
118        let packed = [0x12u8]; // low=2 → 1.0, high=1 → 0.5
119        let scales = [127u8]; // 2^0
120        let got = dequant_nvfp4_e8m0_to_f32(&packed, &scales, 1, 2).unwrap();
121        assert_eq!(got, vec![1.0, 0.5]);
122    }
123}
atlas_core/mxfp4_e8m0.rs

atlas_core/
mxfp4_e8m0.rs
