Use OCP-conformant shared exponent and unify the MXFP paths

memory_mapping/rand_gen.py

@@ -1,7 +1,7 @@
 import torch
 import os
 import numpy as np
-from plena_utils.torch_fp_conversion import pack_fp_to_bin
+from plena_quant.mxfp.utils import pack_fp_to_bin
 from plena_utils.debugger import set_excepthook
 from plena_utils.logger import set_logging_verbosity, get_logger
 from plena_quant.quantizer.hardware_quantizer import _mx_fp_quantize_hardware

plena_quant/common/minifloat.py

@@ -4,7 +4,6 @@
 from torch import Tensor
 
 from .utils import ste_clamp, ste_round
-from .hardware_utils import hardware_round
 
 
 def _minifloat_denorm_quantize_hardware(
@@ -84,9 +83,7 @@ def _minifloat_ieee_quantize_hardware(x: Tensor, width: int, exponent_width: int
     exponent_bias: the value of the exponent bias. If None, the default bias will be (2**exponent_bits - 1) >> 1.
 
     ---
-    For example:
-    a minifloat(bits=8, exponent_bits=4, mantissa_bits=3) number,
-    1 0111 011, is equal to (-1)**1 * 2**(7-15) * (1+3/8) = -0.00537109375
+    For example, E4M3 bits ``1 0111 011`` represent ``-1.375``.
 
     ---
 
@@ -96,39 +93,100 @@ def _minifloat_ieee_quantize_hardware(x: Tensor, width: int, exponent_width: int
 
     # set default bias
     if exponent_bias in (None, "none", "None"):
-        exponent_bias = 2 ** (exponent_width - 1) - 1
-    # upper and lower bound of shifted exponent
-    exponent_max = 2**exponent_width - 2 - exponent_bias
-    exponent_min = -exponent_bias
-    # upper and lower bound of shifted minifloat mantissa
-    shift = 2**mantissa_bits
-    shifted_mantissa_max = 2**mantissa_bits - 1
-    shifted_mantissa_min = 0
-
-    sign = torch.sign(x + 1e-9)
+        exponent_bias = (
+            1 if exponent_width == 1 else 2 ** (exponent_width - 1) - 1
+        )
+    max_exponent_code = (
+        2**exponent_width - 1
+        if exponent_width == 1
+        else 2**exponent_width - 2
+    )
+    exponent_max = max_exponent_code - exponent_bias
+    exponent_min = 1 - exponent_bias
+    fraction_levels = 2**mantissa_bits
+    fraction_max = fraction_levels - 1
 
     value = torch.abs(x)
-    # clip the exponent before calculating mantissa
-    exponent = torch.floor(torch.log2(value + 1e-9))
-    overflow = exponent > exponent_max
-    exponent = ste_clamp(exponent, exponent_min, exponent_max)
-
-    mantissa = value / 2**exponent
-
-    shift = 2**mantissa_bits
-    # fmt: off
-    # if the clipped exponent is zero, the minifloat is in a subnormal form
-    # this `is_normal` also help the grad keeps 1 if input x is 0, or the zero-initialized value will be trapped in 0
-    if isinstance(exponent_bias, (int, float)):
-        exponent_bias = torch.tensor([exponent_bias], dtype=exponent.dtype, device=exponent.device)
-    is_normal = (~torch.isclose(exponent, -exponent_bias))
-
-    shifted_mantissa = is_normal*ste_clamp(hardware_round((mantissa - 1)*shift), shifted_mantissa_min, shifted_mantissa_max) +\
-        (~is_normal)*ste_clamp(hardware_round(mantissa*shift), shifted_mantissa_min, shifted_mantissa_max)
-    shifted_mantissa[overflow] = shifted_mantissa_max
-    mantissa = is_normal*(1.0+shifted_mantissa/shift) + (~is_normal)*(shifted_mantissa/shift)
-    # this `is_close_to_0` helps the grad keeps 1 if input x is 0, or the zero-initialized value will be trapped in 0
-    is_close_to_0 = torch.isclose(value, torch.tensor([0.0], dtype=value.dtype, device=value.device))
-    minifloat_ieee_x = (~is_close_to_0)*(sign * (2**exponent) * mantissa) + is_close_to_0*x
-    # fmt: on
-    return minifloat_ieee_x, exponent, sign * mantissa
+    finite_nonzero = torch.isfinite(value) & (value != 0)
+    safe_value = torch.where(finite_nonzero, value, torch.ones_like(value))
+    raw_exponent = torch.floor(torch.log2(safe_value))
+    min_normal = float(2**exponent_min)
+    is_subnormal = finite_nonzero & (value < min_normal)
+    exponent = ste_clamp(raw_exponent, exponent_min, exponent_max)
+
+    normal_fraction = ste_round(
+        (safe_value / 2**exponent - 1.0) * fraction_levels
+    )
+    normal_carry = normal_fraction >= fraction_levels
+    normal_fraction = torch.where(
+        normal_carry,
+        torch.zeros_like(normal_fraction),
+        normal_fraction,
+    )
+    normal_fraction = ste_clamp(normal_fraction, 0, fraction_max)
+    normal_exponent = exponent + normal_carry.to(exponent.dtype)
+
+    subnormal_fraction = ste_round(value / min_normal * fraction_levels)
+    subnormal_carry = subnormal_fraction >= fraction_levels
+    subnormal_fraction = torch.where(
+        subnormal_carry,
+        torch.zeros_like(subnormal_fraction),
+        subnormal_fraction,
+    )
+    subnormal_fraction = ste_clamp(subnormal_fraction, 0, fraction_max)
+
+    overflow = (~torch.isfinite(value)) | (
+        finite_nonzero
+        & ((raw_exponent > exponent_max) | (normal_exponent > exponent_max))
+    )
+    quantized_exponent = torch.where(
+        is_subnormal,
+        torch.full_like(exponent, exponent_min),
+        normal_exponent,
+    )
+    encoded_exponent = torch.where(
+        is_subnormal & ~subnormal_carry,
+        torch.full_like(exponent, -exponent_bias),
+        quantized_exponent,
+    )
+    fraction = torch.where(
+        is_subnormal,
+        subnormal_fraction,
+        normal_fraction,
+    )
+    mantissa = torch.where(
+        is_subnormal & ~subnormal_carry,
+        fraction / fraction_levels,
+        1.0 + fraction / fraction_levels,
+    )
+
+    encoded_exponent = torch.where(
+        overflow,
+        torch.full_like(encoded_exponent, exponent_max),
+        encoded_exponent,
+    )
+    mantissa = torch.where(
+        overflow,
+        torch.full_like(mantissa, 1.0 + fraction_max / fraction_levels),
+        mantissa,
+    )
+    encoded_exponent = torch.where(
+        finite_nonzero | overflow,
+        encoded_exponent,
+        torch.full_like(encoded_exponent, -exponent_bias),
+    )
+    mantissa = torch.where(
+        finite_nonzero | overflow,
+        mantissa,
+        torch.zeros_like(mantissa),
+    )
+
+    negative = torch.signbit(x)
+    signed_mantissa = torch.where(negative, -mantissa, mantissa)
+    minifloat_ieee_x = signed_mantissa * 2**torch.where(
+        overflow,
+        torch.full_like(quantized_exponent, exponent_max),
+        quantized_exponent,
+    )
+    minifloat_ieee_x = torch.where(value == 0, x, minifloat_ieee_x)
+    return minifloat_ieee_x, encoded_exponent, signed_mantissa

plena_quant/mxfp/quantizer.py

@@ -17,8 +17,8 @@ def _mx_fp_quantize_hardware(
     skip_first_dim: bool = False,
 ):
     """
-    - Convert IEEE FP32/64 to Block Minifloat (BM) which is also called as MXFP, where an exponent bias is shared over all elements in a block
-    - `2**-bias_shared x [(-1)^s1 x 2^exponent1 x mantissa1, (-1)^s2 x 2^exponent2 x mantissa2, ...]`
+    - Convert IEEE FP32/64 to MXFP with one shared exponent per block.
+    - Dequantized values are ``2**shared_exponent * element``.
     - See https://openreview.net/forum?id=6zaTwpNSsQ2
 
     ---
@@ -31,7 +31,14 @@ def _mx_fp_quantize_hardware(
     - `exponent_bias_width`: the number of bits of the shared exponent bias
     - `block_size`: a list of integers where each integer is the block size on that dimension. See function `block`.
 
+    The shared exponent follows the OCP MX convention: the block maximum is
+    placed at the top of the element format's finite exponent range. This is
+    a software accuracy convention; the datapath consumes the resulting E8M0
+    scale code without selecting a scale-placement rule.
+
     """
+    if isinstance(block_size, int):
+        block_size = [block_size]
     if len(block_size) == 1:
         block_size = [1, block_size[0]]
     else:
@@ -54,9 +61,33 @@ def _mx_fp_quantize_hardware(
     ).permute(0, 1, 3, 2, 4)
     px = px.reshape(-1, block_size[0] * block_size[1])
 
-    per_block_max = px.abs().max(dim=-1, keepdim=True).values + 1e-9
+    per_block_max = px.abs().max(dim=-1, keepdim=True).values
+    nonzero_block = per_block_max > 0
+    safe_block_max = torch.where(
+        nonzero_block,
+        per_block_max,
+        torch.ones_like(per_block_max),
+    )
+    scale_bias = 2 ** (exponent_bias_width - 1) - 1
+    scale_max = 2**exponent_bias_width - 1 - scale_bias
+    element_exponent_bias = (
+        1 if exponent_width == 1 else 2 ** (exponent_width - 1) - 1
+    )
+    element_max_exponent_code = (
+        2**exponent_width - 1
+        if exponent_width == 1
+        else 2**exponent_width - 2
+    )
+    element_max_exponent = element_max_exponent_code - element_exponent_bias
     per_block_exponent_bias = ste_clamp(
-        torch.floor(torch.log2(per_block_max)), -(2 ** (exponent_bias_width - 1)), 2 ** (exponent_bias_width - 1) - 1
+        torch.floor(torch.log2(safe_block_max)) - element_max_exponent,
+        -scale_bias,
+        scale_max,
+    )
+    per_block_exponent_bias = torch.where(
+        nonzero_block,
+        per_block_exponent_bias,
+        torch.zeros_like(per_block_exponent_bias),
     )
 
     px = px / 2**per_block_exponent_bias
@@ -69,13 +100,16 @@ def _mx_fp_quantize_hardware(
     per_block_bm_x = per_block_bm_x * 2**per_block_exponent_bias
 
     bm_x = per_block_bm_x.reshape(
-        -1, px_shape[0] // block_size[0], px_shape[1] // block_size[1], block_size[0], block_size[1]
+        -1,
+        px_shape[-2] // block_size[0],
+        px_shape[-1] // block_size[1],
+        block_size[0],
+        block_size[1],
     )
     bm_x = bm_x.permute(0, 1, 3, 2, 4)
     bm_x = bm_x.reshape(-1, px_shape[-2], px_shape[-1])
     bm_x = bm_x[:, : x_shape[-2], : x_shape[-1]]
 
-    bias_bias = 2 ** (exponent_bias_width - 1) - 1
-    per_block_exponent_bias = per_block_exponent_bias + bias_bias
+    per_block_exponent_bias = per_block_exponent_bias + scale_bias
 
     return bm_x, per_block_fp_exp, per_block_fp_mant, per_block_exponent_bias

plena_quant/mxfp/utils.py

@@ -11,19 +11,18 @@ def pack_fp_to_bin(signed_exponent, signed_mantissa, exp_width, man_width):
     signed_exponent = signed_exponent.reshape(-1)
     signed_mantissa = signed_mantissa.reshape(-1)
 
-    sign = signed_mantissa.sign()
-    sign_bit = torch.where(sign < 0, torch.tensor(1), torch.tensor(0))
+    sign_bit = torch.signbit(signed_mantissa).to(torch.int64)
 
-    exponent_bias = (2 ** (exp_width - 1)) - 1
-    exponent_bit = signed_exponent + exponent_bias
+    exponent_bias = 1 if exp_width == 1 else (2 ** (exp_width - 1)) - 1
+    exponent_bit = (signed_exponent + exponent_bias).to(torch.int64)
 
     for item in exponent_bit:
         assert item >= 0 and item <= (2**exp_width - 1), "Exponent out of range!"
 
     mantissa = torch.where(signed_mantissa < 0, -signed_mantissa, signed_mantissa)
     mantissa_bit = torch.where(exponent_bit == 0, mantissa, mantissa - 1)
 
-    mantissa_bit = mantissa_bit * 2 ** (man_width)
+    mantissa_bit = (mantissa_bit * 2 ** (man_width)).round().to(torch.int64)
 
     result = ((sign_bit * 2 ** (exp_width + man_width)) + exponent_bit * 2 ** (man_width) + mantissa_bit).int()
 
@@ -44,54 +43,32 @@ def split_bin(bits: int | BitArray, exp_width: int, mant_width: int):
     exponent = int(exponent_bits.bin, 2)
     mantissa = int(mantissa_bits.bin, 2)
     # Bias and reconstruct
-    bias = (1 << (exp_width - 1)) - 1
-    exponent_val = exponent - bias
-    exponent_min = 0 - bias
-
-    if exponent_val == exponent_min:
+    bias = 1 if exp_width == 1 else (1 << (exp_width - 1)) - 1
+    if exponent == 0:
+        exponent_val = 1 - bias
         mantissa_val = mantissa / (1 << mant_width)
     else:
+        exponent_val = exponent - bias
         mantissa_val = 1.0 + (mantissa / (1 << mant_width))
 
     return exponent_val, -mantissa_val if sign else mantissa_val
 
 
 def bin_2_fp(bits: int | BitArray | torch.Tensor | list, exp_width: int, mant_width: int):
-    if isinstance(bits, torch.Tensor) or isinstance(bits, list):
-        # Handle tensor input - process each element
-        results = []
-        if isinstance(bits, torch.Tensor):
-            _tensor = True
-            bits = bits.reshape(-1)
-            bits_list = bits.tolist()
-            for bit_val in bits_list:
-                exp_val, mant_val = split_bin(bit_val, exp_width, mant_width)
-                results.append(mant_val * 2**exp_val)
-        else:
-            results = []
-
-        if isinstance(bits_list, list):
-            # Handle multi-element tensor
-            results = []
-            for bit_val in bits_list:
-                exp_val, mant_val = split_bin(bit_val, exp_width, mant_width)
-                results.append(mant_val * 2**exp_val)
-            return torch.tensor(results)
-        else:
-            # Handle single-element tensor
-            exp_val, mant_val = split_bin(bits_list, exp_width, mant_width)
-            return torch.tensor(mant_val)
-    elif isinstance(bits, list):
-        # Handle list input
-        results = []
-        for bit_val in bits:
-            exp_val, mant_val = split_bin(bit_val, exp_width, mant_width)
-            results.append(mant_val * 2**exp_val)
-        return results
-    else:
-        # Handle single int or BitArray
-        exp_val, mant_val = split_bin(bits, exp_width, mant_width)
-        return mant_val * 2**exp_val
+    if isinstance(bits, torch.Tensor):
+        shape = bits.shape
+        values = [
+            bin_2_fp(value, exp_width, mant_width)
+            for value in bits.reshape(-1).tolist()
+        ]
+        return torch.tensor(values, device=bits.device).reshape(shape)
+    if isinstance(bits, list):
+        return [
+            bin_2_fp(value, exp_width, mant_width)
+            for value in bits
+        ]
+    exp_val, mant_val = split_bin(bits, exp_width, mant_width)
+    return mant_val * 2**exp_val
 
 
 def fp_2_bin(fp: torch.Tensor | float | list, exp_width: int, mant_width: int):