metrale_core/

numeric.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! 2026-09-25: Host-side numeric conversions: the FP8 E4M3 decode table, the f32 -> BF16 cast and BF16 widening.
//!
//! Owner: metrale-core.
//! Invariants:
//! - No CUDA and no `cfg`: the module compiles under every feature combination.
//! - `FP8_E4M3_LUT` decodes the two NaN bytes (`0x7F`, `0xFF`) to signed zero.
//! - `E4M3_LUT_GMOE` in `kernels/gb10/common/moe_fp8_grouped_gemm.cu` must agree with
//!   `FP8_E4M3_LUT` element for element; it does, including the signed zeros.

/// 2026-09-25: FP8 E4M3 -> f32 lookup table, one entry per byte value.
///
/// OCP FP8 E4M3FN: sign(1) | exponent(4) | mantissa(3), bias 7. There are no
/// infinities; `0x7F` / `0xFF` are NaN and max finite is +/-448.0 (exp 15,
/// mant 6).
///
/// NaN entries decode to signed zero: one bad byte then cannot turn a whole
/// dequantised tensor into NaN.
///
/// Built at compile time, so a decode is one indexed load with no branches.
#[allow(clippy::if_same_then_else)]
pub static FP8_E4M3_LUT: [f32; 256] = {
    let mut table = [0.0f32; 256];
    let mut i: u32 = 0;
    while i < 256 {
        let bits = i as u8;
        let sign = (bits >> 7) & 1;
        let exp = (bits >> 3) & 0x0F;
        let mantissa = bits & 0x07;

        let val = if exp == 0 && mantissa == 0 {
            0.0f32
        } else if exp == 0x0F && mantissa == 0x07 {
            0.0f32
        } else if exp == 0 {
            // 2026-09-25: Subnormal: 2^(-6) * (mantissa / 8).
            (mantissa as f32) * (0.015625f32 / 8.0)
        } else {
            // 2026-09-25: Normal: 2^(exp-7) * (1 + mantissa/8), assembled directly in
            // f32 bits — f32 exponent = fp8_exp - 7 + 127 = fp8_exp + 120,
            // f32 mantissa = fp8_mant << 20 (3 bits left-aligned into 23).
            let f32_exp = (exp as u32 + 120) << 23;
            let f32_mant = (mantissa as u32) << 20;
            f32::from_bits(f32_exp | f32_mant)
        };

        table[i as usize] = if sign == 1 { -val } else { val };
        i += 1;
    }
    table
};

/// 2026-09-25: Decode one FP8 E4M3 byte to f32 with a single table load.
#[inline(always)]
pub fn fp8_e4m3_to_f32(bits: u8) -> f32 {
    FP8_E4M3_LUT[bits as usize]
}

/// 2026-09-25: Convert f32 to BF16 with IEEE-754 round-to-nearest-even.
///
/// Truncation (`bits >> 16`) is biased toward zero; the tests below pin the
/// cases where the two differ.
///
/// NaN maps to the canonical quiet-NaN pattern with the sign preserved.
///
/// `METRALE_DISABLE_RNE` is a bisect escape hatch that reverts to
/// truncation. It is a presence check, not a value check: `=0` disables RNE
/// just as `=1` does.
///
/// The escape hatch is read once per process. Weight loaders call this once
/// per element (`weight_map::quant_helpers`, `weight_map::fp8_lut`), and a
/// `std::env::var` per call would allocate a `String` and take the
/// process-wide environment lock each time.
///
/// Caching is safe because the test that varies the variable re-execs the
/// test binary as a child process with it set
/// (`disable_rne_presence_uses_truncation` below), so each process resolves
/// it once at its own start. Do not convert that test to an in-process
/// `set_var`.
#[inline(always)]
pub fn f32_to_bf16(val: f32) -> u16 {
    static DISABLE_RNE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    if *DISABLE_RNE.get_or_init(|| std::env::var("METRALE_DISABLE_RNE").is_ok()) {
        return (val.to_bits() >> 16) as u16;
    }
    let bits = val.to_bits();
    if val.is_nan() {
        let sign = ((bits >> 16) & 0x8000) as u16;
        return sign | 0x7FC0;
    }
    let lsb = (bits >> 16) & 1;
    let rounding_bias = 0x7FFFu32 + lsb;
    (bits.wrapping_add(rounding_bias) >> 16) as u16
}

/// 2026-09-25: Widen little-endian BF16 bytes to f32. Exact: BF16 is the
/// top 16 bits of an f32, so this is a shift, never a rounding.
#[inline(always)]
pub fn bf16_bytes_to_f32(bytes: [u8; 2]) -> f32 {
    let bits = u16::from_le_bytes(bytes);
    f32::from_bits((bits as u32) << 16)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fp8_lut_reference_values() {
        assert_eq!(fp8_e4m3_to_f32(0x00).to_bits(), 0x0000_0000);
        assert_eq!(fp8_e4m3_to_f32(0x80).to_bits(), 0x8000_0000);
        assert_eq!(fp8_e4m3_to_f32(0x38), 1.0);
        assert_eq!(fp8_e4m3_to_f32(0xB8), -1.0);
        assert_eq!(fp8_e4m3_to_f32(0x3C), 1.5);
        assert_eq!(fp8_e4m3_to_f32(0x7E), 448.0);
        assert_eq!(fp8_e4m3_to_f32(0xFE), -448.0);
        assert_eq!(fp8_e4m3_to_f32(0x7F).to_bits(), 0x0000_0000);
        assert_eq!(fp8_e4m3_to_f32(0xFF).to_bits(), 0x8000_0000);

        // 2026-09-25: Subnormals: 2^(-6) * mant/8.
        let eps = 1e-10;
        assert!((fp8_e4m3_to_f32(0x01) - 0.001953125).abs() < eps);
        assert!((fp8_e4m3_to_f32(0x07) - 0.013671875).abs() < eps);
    }

    #[test]
    #[allow(clippy::if_same_then_else)]
    fn fp8_lut_matches_ocp_values_and_metrale_nan_policy_for_all_bytes() {
        // 2026-09-25: Re-derived from the OCP finite-value definition with float
        // math, independently of the table's bit assembly. The two OCP NaN
        // encodings map to signed zero, as `E4M3_LUT_GMOE` does on the device.
        for i in 0u16..256 {
            let bits = i as u8;
            let sign = (bits >> 7) & 1;
            let exp = (bits >> 3) & 0x0F;
            let mant = bits & 0x07;

            let magnitude = if exp == 0x0F && mant == 0x07 {
                0.0f32
            } else if exp == 0 && mant == 0 {
                0.0f32
            } else if exp == 0 {
                (mant as f32 / 8.0) * 2.0f32.powi(-6)
            } else {
                (1.0 + mant as f32 / 8.0) * 2.0f32.powi(exp as i32 - 7)
            };
            let expected = if sign == 1 { -magnitude } else { magnitude };
            let actual = fp8_e4m3_to_f32(bits);
            assert_eq!(
                actual.to_bits(),
                expected.to_bits(),
                "LUT mismatch at {i:#04x}: expected {expected:?}, got {actual:?}"
            );
        }
    }

    /// 2026-09-25: The assertions that separate round-to-nearest-even from
    /// truncation toward zero. Truncation fails every "round up" case here.
    #[test]
    fn f32_to_bf16_is_rne_byte_exact() {
        fn convert(bits: u32) -> u16 {
            f32_to_bf16(f32::from_bits(bits))
        }

        // 2026-09-25: Below half-ULP: round down. Truncation agrees.
        assert_eq!(convert(0x3F80_0800), 0x3F80, "1.0 + below-half-ULP -> 1.0");
        // 2026-09-25: Exactly half-ULP, LSB=0: tie -> round to even (down). Does
        // not distinguish RNE from truncation; kept for the tie coverage.
        assert_eq!(
            convert(0x3F80_8000),
            0x3F80,
            "1.0 + exact-half-ULP, LSB=0 -> 1.0 (even)"
        );
        // 2026-09-25: Above half-ULP: round up. Truncation would give 0x3F80.
        assert_eq!(
            convert(0x3F80_8001),
            0x3F81,
            "1.0 + above-half-ULP -> next bf16 (truncation would give 0x3F80)"
        );
        // 2026-09-25: Exactly half-ULP, LSB=1: tie -> round to even (up).
        // Truncation would give 0x3F81.
        assert_eq!(
            convert(0x3F81_8000),
            0x3F82,
            "1.0078125 + exact-half-ULP, LSB=1 -> 1.015625"
        );
        // 2026-09-25: Negative values round the magnitude the same way.
        assert_eq!(convert(0xBF80_8001), 0xBF81, "negative round up");
        assert_eq!(convert(0x0000_0000), 0x0000, "+0.0");
        assert_eq!(convert(0x8000_0000), 0x8000, "-0.0");
        // 2026-09-25: Smallest f32 subnormal (2^-149) is below half-ULP and
        // rounds down to 0.
        assert_eq!(convert(0x0000_0001), 0x0000, "tiny subnormal -> 0");
        assert_eq!(convert(0x7F80_0000), 0x7F80, "+inf");
        assert_eq!(convert(0xFF80_0000), 0xFF80, "-inf");
        // 2026-09-25: Max-finite f32 rounds up to +inf in bf16.
        assert_eq!(
            convert(0x7F7F_FFFF),
            0x7F80,
            "max-finite f32 rounds to +inf bf16"
        );
        // 2026-09-25: NaN -> canonical quiet NaN, sign preserved; a signalling
        // NaN is quieted rather than passed through.
        assert_eq!(convert(0x7FC0_0000), 0x7FC0, "qnan +");
        assert_eq!(convert(0xFFC0_0000), 0xFFC0, "qnan -");
        assert_eq!(convert(0x7F80_0001), 0x7FC0, "snan + -> qnan +");
    }

    /// 2026-09-25: Byte-exact round-to-nearest-even results for common values,
    /// including cases where truncation gives a different result.
    #[test]
    fn f32_to_bf16_matches_pytorch() {
        let cases: &[(u32, u16, &str)] = &[
            (0x3F80_0000, 0x3F80, "1.0"),
            (0x4000_0000, 0x4000, "2.0"),
            (0xC000_0000, 0xC000, "-2.0"),
            (0x3FC0_0000, 0x3FC0, "1.5"),
            (
                0x3DCC_CCCD,
                0x3DCD,
                "0.1 -> RNE rounds UP to 0x3DCD (trunc=0x3DCC)",
            ),
            (0x3F4C_CCCD, 0x3F4D, "0.8 -> RNE rounds UP to 0x3F4D"),
            (0x40C9_0FDB, 0x40C9, "pi -> truncates (next bit < half)"),
            (0x402D_F854, 0x402E, "e -> RNE rounds UP (next bit > half)"),
            (0x4490_0000, 0x4490, "1152.0"),
            (0x3727_C5AC, 0x3728, "1e-5 -> RNE rounds UP"),
        ];
        for (f32_bits, want, desc) in cases {
            let got = f32_to_bf16(f32::from_bits(*f32_bits));
            assert_eq!(
                got, *want,
                "f32={f32_bits:#010x} ({desc}): want bf16={want:#06x}, got {got:#06x}"
            );
        }
    }

    #[test]
    fn disable_rne_presence_uses_truncation() {
        let this = format!(
            "{}::disable_rne_presence_uses_truncation",
            module_path!()
                .split_once("::")
                .map_or(module_path!(), |(_, p)| p)
        );
        const CHILD_MARKER: &str = "METRALE_NUMERIC_RNE_CHILD";

        if std::env::var_os(CHILD_MARKER).is_some() {
            assert_eq!(
                f32_to_bf16(f32::from_bits(0x3F80_8001)),
                0x3F80,
                "the escape hatch must truncate an above-half-ULP value"
            );
            return;
        }

        for value in ["0", "1"] {
            let output = std::process::Command::new(std::env::current_exe().unwrap())
                .args(["--exact", this.as_str()])
                .env(CHILD_MARKER, "1")
                .env("METRALE_DISABLE_RNE", value)
                .output()
                .unwrap();
            assert!(
                output.status.success(),
                "METRALE_DISABLE_RNE={value} child failed:\n{}",
                String::from_utf8_lossy(&output.stdout)
            );
            assert!(
                String::from_utf8_lossy(&output.stdout).contains("1 passed"),
                "the child must run exactly this test: {}",
                String::from_utf8_lossy(&output.stdout)
            );
        }
    }

    #[test]
    fn bf16_widening_is_byte_exact_for_every_pattern() {
        for bits in 0u32..=0xFFFF {
            let bf16 = bits as u16;
            let widened = bf16_bytes_to_f32(bf16.to_le_bytes());
            assert_eq!(
                widened.to_bits(),
                bits << 16,
                "widening moved bits for bf16 {bf16:#06x}"
            );
        }
    }

    #[test]
    fn bf16_narrowing_preserves_every_non_nan_bf16_value() {
        for bits in 0u32..=0xFFFF {
            let bf16 = bits as u16;
            let widened = f32::from_bits(bits << 16);
            if widened.is_nan() {
                continue;
            }
            assert_eq!(
                f32_to_bf16(widened),
                bf16,
                "round trip failed for bf16 {bf16:#06x}"
            );
        }
    }
}