metrale_core/
mxfp4_e8m0.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Native MXFP4 host unpack: packed E2M1 nibbles + E8M0 scales.
4//!
5//! Owner: metrale-core.
6//! Invariants:
7//! - E8M0 bytes 0 and 255 decode to 0.0; every other byte `b` to `2^(b-127)`.
8//! - Within a byte, the even flat index is the low nibble.
9//!
10//! The Kimi-K3 host path (`metrale_model_weights::kimi_k3_host::mxfp4`) calls
11//! this rather than keeping its own unpack. The device scale decode
12//! `mx_block_scale<true>` (`kernels/gb10/common/mx_block_scale.cuh`) must
13//! stay byte-exact with [`fp8_e8m0_to_f32`].
14//!
15//! The Kimi-K3 GPU GEMM compiles `deepseek-v4-flash/nvfp4/moe_w4a16_grouped_gemm.cu`
16//! through the `[sources] use` list of `kernels/gb10/kimi-k3/*/KERNEL.toml`;
17//! do not copy that file into kimi-k3.
18
19use anyhow::{Result, ensure};
20
21/// 2026-09-25: E2M1 nibble → f32, indexed by the 4-bit code (bit 3 is the sign).
22pub const E2M1: [f32; 16] = [
23    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,
24];
25
26/// 2026-09-25: Native MXFP4 group size; the GPU lander
27/// `weight_map::quantized_mxfp4_e8m0_pair` refuses any other.
28pub const GROUP_SIZE: usize = 32;
29
30/// 2026-09-25: FP8 E8M0 → f32 (unsigned exponent, bias 127). exp=0 and exp=255 → 0.0.
31const FP8_E8M0_LUT: [f32; 256] = {
32    let mut table = [0.0f32; 256];
33    let mut i: u32 = 0;
34    while i < 256 {
35        let exp = i as u8;
36        table[i as usize] = if exp == 0 || exp == 255 {
37            0.0f32
38        } else {
39            f32::from_bits((exp as u32) << 23)
40        };
41        i += 1;
42    }
43    table
44};
45
46/// 2026-09-25: Convert one E8M0 scale byte to f32 with a single table load.
47#[inline(always)]
48pub fn fp8_e8m0_to_f32(bits: u8) -> f32 {
49    FP8_E8M0_LUT[bits as usize]
50}
51
52/// 2026-09-25: Unpack packed E2M1 `[n, k/2]` + E8M0 scales to f32 `[n, k]`.
53///
54/// The block size is inferred as `n*k / scales.len()`; unlike the GPU lander
55/// this does not require 32. Even flat index = low nibble, odd = high nibble.
56/// Errors when `n*k` is odd, `packed` is not `n*k/2` bytes, or `scales` is
57/// empty or does not divide `n*k`. Panics if `n*k` overflows `usize`.
58pub fn dequant_nvfp4_e8m0_to_f32(
59    packed: &[u8],
60    scales: &[u8],
61    n: usize,
62    k: usize,
63) -> Result<Vec<f32>> {
64    let total = n.checked_mul(k).expect("mxfp4 n*k");
65    ensure!(
66        total.is_multiple_of(2),
67        "MXFP4 E8M0: n*k={total} is odd (need even nibble count)"
68    );
69    let packed_bytes = total / 2;
70    ensure!(
71        packed.len() == packed_bytes,
72        "MXFP4 E8M0: packed {} B, expected {packed_bytes} for [{n},{k}]",
73        packed.len()
74    );
75    let num_groups = scales.len();
76    ensure!(
77        num_groups > 0 && total.is_multiple_of(num_groups),
78        "MXFP4 E8M0: weight elems {total} not divisible by E8M0 scale groups {num_groups}"
79    );
80    let block = total / num_groups;
81    let mut out = vec![0.0f32; total];
82    for (group, &sb) in scales.iter().enumerate() {
83        let block_scale = fp8_e8m0_to_f32(sb);
84        for elem in 0..block {
85            let flat_idx = group * block + elem;
86            let byte_idx = flat_idx / 2;
87            let nibble = if flat_idx.is_multiple_of(2) {
88                packed[byte_idx] & 0x0F
89            } else {
90                (packed[byte_idx] >> 4) & 0x0F
91            };
92            out[flat_idx] = E2M1[nibble as usize] * block_scale;
93        }
94    }
95    Ok(out)
96}
97
98/// 2026-09-25: Same unpack, then `numeric::f32_to_bf16` per element.
99pub fn dequant_nvfp4_e8m0_to_bf16(
100    packed: &[u8],
101    scales: &[u8],
102    n: usize,
103    k: usize,
104) -> Result<Vec<u16>> {
105    let f = dequant_nvfp4_e8m0_to_f32(packed, scales, n, k)?;
106    Ok(f.into_iter().map(crate::numeric::f32_to_bf16).collect())
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn e8m0_pow2_and_sentinels() {
115        assert_eq!(fp8_e8m0_to_f32(0), 0.0);
116        assert_eq!(fp8_e8m0_to_f32(255), 0.0);
117        assert_eq!(fp8_e8m0_to_f32(127), 1.0);
118        assert_eq!(fp8_e8m0_to_f32(128), 2.0);
119        assert_eq!(fp8_e8m0_to_f32(126), 0.5);
120    }
121
122    #[test]
123    fn matches_dsv4_nibble_order_and_lut() {
124        // 2026-09-25: Low nibble first: 0x12 holds low=2 (1.0) then high=1 (0.5), at scale 2^0.
125        let packed = [0x12u8];
126        let scales = [127u8];
127        let got = dequant_nvfp4_e8m0_to_f32(&packed, &scales, 1, 2).unwrap();
128        assert_eq!(got, vec![1.0, 0.5]);
129    }
130}
metrale_core/mxfp4_e8m0.rs

metrale_core/
mxfp4_e8m0.rs
