Fix MX block quantization: row-local blocking, OCP shared exponent, element rounding

src/chop/nn/quantizers/_minifloat_mx/__init__.py

@@ -5,13 +5,18 @@
 """
 
 from .meta import MinifloatMeta, MinifloatTensorMeta
-from .fake import extract_minifloat_component, compose_minifloat_component
+from .fake import (
+    compose_minifloat_component,
+    extract_minifloat_component,
+    quantize_minifloat_value,
+)
 from .minifloat import minifloat_quantizer_sim
 
 __all__ = [
     "MinifloatMeta",
     "MinifloatTensorMeta",
     "extract_minifloat_component",
     "compose_minifloat_component",
+    "quantize_minifloat_value",
     "minifloat_quantizer_sim",
 ]

src/chop/nn/quantizers/_minifloat_mx/fake.py

@@ -8,89 +8,125 @@
 from .meta import MinifloatMeta
 
 
-def extract_minifloat_component(x: Tensor, minifloat_meta: MinifloatMeta) -> Tensor:
-    """
-    Extract minifloat representation from float tensor.
-
-    Args:
-        x: Input float tensor
-        minifloat_meta: Minifloat format specification
-
-    Returns:
-        Tensor of uint16 containing minifloat representation
-    """
+def _minifloat_fields(
+    x: Tensor, minifloat_meta: MinifloatMeta
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Return (sign, biased exponent, fraction) for one minifloat rounding."""
     y_exp_bits = minifloat_meta.exp_bits
     y_frac_bits = minifloat_meta.frac_bits
     always_finite = minifloat_meta.is_finite
     round_mode = minifloat_meta.round_mode
 
-    y_exp_bias = (1 << (y_exp_bits - 1)) - 1
+    y_exp_bias = 1 if y_exp_bits == 1 else (1 << (y_exp_bits - 1)) - 1
     y_exp_max = (1 << y_exp_bits) - 1 if always_finite else (1 << y_exp_bits) - 2
     y_exp_max_biased = y_exp_max - y_exp_bias
-    y_exp_min = 0
-    y_exp_min_biased = y_exp_min - y_exp_bias
+    y_exp_min_biased = 1 - y_exp_bias
     y_frac_max = (1 << y_frac_bits) - 1
+    y_frac_levels = 1 << y_frac_bits
 
     x = x.to(torch.float32)
-    y_sign = x < 0
-    x_int32 = x.abs().view(torch.int32)
-    flush_to_zero = (x_int32 & 0x7F800000) == 0
-    x_normal = torch.where(flush_to_zero, 0.0, x)
-
-    x_frac, x_exp = x_normal.abs().frexp()
-    x_frac = x_frac * 2
-    x_exp = x_exp - 1
-
-    if not always_finite:
-        x_is_inf = x.isinf()
-        x_is_nan = x.isnan()
-
-    y_exp = x_exp
-    underflow = y_exp < y_exp_min_biased
-    overflow = y_exp > y_exp_max_biased
-    y_exp = y_exp + y_exp_bias
-
-    y_frac = x_frac.view(torch.int32) & 0x7FFFFF
-
-    if round_mode == "rz":
-        y_frac = y_frac >> (23 - y_frac_bits)
-    else:
-        y_frac = (y_frac >> 8).float()
-        div = 1 << (15 - y_frac_bits)
-        y_frac = y_frac / div
+    y_sign = torch.signbit(x)
+    magnitude = x.abs()
+    nonzero_finite = torch.isfinite(magnitude) & (magnitude != 0)
+
+    def round_magnitude(value: Tensor) -> Tensor:
+        if round_mode == "rn":
+            return value.round()
+        if round_mode == "rz":
+            return value.floor()
         if round_mode == "ru":
-            y_frac = y_frac.ceil()
-        elif round_mode == "rd":
-            y_frac = y_frac.floor()
-        elif round_mode == "rn":
-            y_frac = y_frac.round()
-        else:
-            raise ValueError(f"Unknown rounding mode: {round_mode}")
-        y_frac = y_frac.to(torch.int32)
-
-    y_is_subnormal = (y_exp == y_exp_min) & (y_frac != 0)
-    y_frac = torch.where(y_is_subnormal, (y_frac | (1 << y_frac_bits)) >> 1, y_frac)
-
-    # underflow -> 0
-    y_frac = torch.where(underflow, 0, y_frac)
-    y_exp = torch.where(underflow, 0, y_exp)
-    # overflow -> max
+            return torch.where(y_sign, value.floor(), value.ceil())
+        if round_mode == "rd":
+            return torch.where(y_sign, value.ceil(), value.floor())
+        raise ValueError(f"Unknown rounding mode: {round_mode}")
+
+    min_normal = float(2**y_exp_min_biased)
+    subnormal_step = min_normal / y_frac_levels
+    is_subnormal = nonzero_finite & (magnitude < min_normal)
+
+    safe_magnitude = torch.where(
+        nonzero_finite,
+        magnitude,
+        torch.ones_like(magnitude),
+    )
+    unbiased_exp = torch.floor(torch.log2(safe_magnitude))
+    overflow = (~torch.isfinite(magnitude)) | (
+        nonzero_finite & (unbiased_exp > y_exp_max_biased)
+    )
+    clamped_exp = unbiased_exp.clamp(
+        min=y_exp_min_biased,
+        max=y_exp_max_biased,
+    )
+    normal_fraction = (safe_magnitude / torch.exp2(clamped_exp) - 1.0) * y_frac_levels
+    y_frac = round_magnitude(normal_fraction).to(torch.int32)
+    y_exp = (clamped_exp + y_exp_bias).to(torch.int32)
+
+    carry = y_frac >= y_frac_levels
+    y_frac = torch.where(carry, 0, y_frac)
+    y_exp = torch.where(carry, y_exp + 1, y_exp)
+
+    subnormal_fraction = round_magnitude(magnitude / subnormal_step).to(torch.int32)
+    subnormal_carry = subnormal_fraction >= y_frac_levels
+    y_frac = torch.where(
+        is_subnormal,
+        torch.where(subnormal_carry, 0, subnormal_fraction),
+        y_frac,
+    )
+    y_exp = torch.where(
+        is_subnormal,
+        torch.where(subnormal_carry, 1, 0),
+        y_exp,
+    )
+
+    overflow = overflow | (y_exp > y_exp_max)
     y_frac = torch.where(overflow, y_frac_max, y_frac)
     y_exp = torch.where(overflow, y_exp_max, y_exp)
-    # flush to zero
-    y_frac = torch.where(flush_to_zero, 0, y_frac)
-    y_exp = torch.where(flush_to_zero, 0, y_exp)
+    y_frac = torch.where(nonzero_finite | overflow, y_frac, 0)
+    y_exp = torch.where(nonzero_finite | overflow, y_exp, 0)
+    return y_sign, y_exp, y_frac
 
-    if not always_finite:
-        y_frac = torch.where(x_is_inf, 0, y_frac)
-        y_frac = torch.where(x_is_nan, (1 << y_frac_bits) - 1, y_frac)
-        y_exp = torch.where(x_is_inf, y_exp_max, y_exp)
-        y_exp = torch.where(x_is_nan, y_exp_max, y_exp)
 
+def extract_minifloat_component(x: Tensor, minifloat_meta: MinifloatMeta) -> Tensor:
+    """
+    Extract minifloat representation from float tensor.
+
+    Args:
+        x: Input float tensor
+        minifloat_meta: Minifloat format specification
+
+    Returns:
+        Tensor of uint16 containing minifloat representation
+    """
+    y_sign, y_exp, y_frac = _minifloat_fields(x, minifloat_meta)
+    y_frac_bits = minifloat_meta.frac_bits
     y = (y_exp << y_frac_bits) | y_frac
-    y = torch.where(y_sign, y + (1 << (y_exp_bits + y_frac_bits)), y)
-    y = y.to(torch.uint16)
-    return y
+    y = torch.where(y_sign, y + (1 << (minifloat_meta.exp_bits + y_frac_bits)), y)
+    return y.to(torch.uint16)
+
+
+def quantize_minifloat_value(x: Tensor, minifloat_meta: MinifloatMeta) -> Tensor:
+    """Round to the minifloat grid and return the value, skipping the encoding.
+
+    Equivalent to ``compose_minifloat_component(extract_minifloat_component(x))``
+    but without the uint16 round trip. The encoder never emits an inf/nan
+    exponent code, so the decoder's special-value branches cannot apply here.
+    """
+    y_sign, y_exp, y_frac = _minifloat_fields(x, minifloat_meta)
+    y_exp_bias = (
+        1
+        if minifloat_meta.exp_bits == 1
+        else ((1 << (minifloat_meta.exp_bits - 1)) - 1)
+    )
+    frac_levels = 1 << minifloat_meta.frac_bits
+    subnormal_step = float(2 ** (1 - y_exp_bias)) / frac_levels
+    fraction = y_frac.to(torch.float32)
+    magnitude = torch.where(
+        y_exp == 0,
+        fraction * subnormal_step,
+        (1.0 + fraction / frac_levels)
+        * torch.exp2((y_exp - y_exp_bias).to(torch.float32)),
+    )
+    return torch.where(y_sign, -magnitude, magnitude)
 
 
 def compose_minifloat_component(
@@ -115,13 +151,13 @@ def compose_minifloat_component(
 
     x_sign_mask = 1 << (exp_bits + frac_bits)
     x_frac_mask = (1 << frac_bits) - 1
-    x_exp_bias = (1 << (exp_bits - 1)) - 1
+    x_exp_bias = 1 if exp_bits == 1 else (1 << (exp_bits - 1)) - 1
 
     assert elements.dtype == torch.uint16
     elements = elements.to(torch.int32)
     y_sign = (elements & x_sign_mask) << (31 - (exp_bits + frac_bits))
 
-    elements = elements & 0x7FFF
+    elements = elements & (x_sign_mask - 1)
     x_exp = (elements >> frac_bits) & ((1 << exp_bits) - 1)
     x_frac = elements & x_frac_mask
     is_subnormal = (x_exp == 0) & (x_frac != 0)

src/chop/nn/quantizers/_minifloat_mx/meta.py

@@ -15,7 +15,7 @@ class MinifloatMeta:
     Args:
         exp_bits: Number of exponent bits
         frac_bits: Number of fraction bits
-        is_finite: Whether the minifloat type supports inf/nan
+        is_finite: Whether every exponent code represents a finite value
         round_mode: Rounding mode - "rn" (nearest), "rd" (down), "ru" (up), "rz" (truncate)
     """
 

src/chop/nn/quantizers/_minifloat_mx/minifloat.py

@@ -6,7 +6,7 @@
 from torch import Tensor
 
 from .meta import MinifloatMeta, MinifloatTensorMeta
-from .fake import extract_minifloat_component, compose_minifloat_component
+from .fake import quantize_minifloat_value
 
 
 def minifloat_quantizer_sim(
@@ -25,9 +25,5 @@ def minifloat_quantizer_sim(
     Returns:
         Dequantized tensor
     """
-    ori_dtype = tensor.dtype
-    element = extract_minifloat_component(tensor, minifloat_meta)
-
-    return compose_minifloat_component(
-        element, minifloat_meta, output_dtype=output_dtype or ori_dtype
-    )
+    value = quantize_minifloat_value(tensor, minifloat_meta)
+    return value.to(output_dtype or tensor.dtype)

src/chop/nn/quantizers/mxfp/fake.py

@@ -1,49 +1,59 @@
-"""
-Fake MXFP quantization — aligned with Plena-Acc-Sim.
-
-"""
+"""PLENA MXFP quantize-dequantize helpers."""
 
 import torch
 from torch import Tensor
 
-from ..minifloat import minifloat_denorm_quantizer
+from .._minifloat_mx import minifloat_quantizer_sim
 from .meta import MXFPMeta
 
 
 def extract_mxfp_components(
     x: Tensor, mxfp_meta: MXFPMeta, percentile: float = 1.0
 ) -> tuple[Tensor, Tensor]:
-    """Extract MXFP components (Plena-aligned).
+    """Split a blocked tensor into shared scales and minifloat elements.
 
     Returns:
-        scales:   per-block shared exponent in **log2 domain** (float),
-                  shape ``[n_blocks, 1]``. Recompose with ``2 ** scales``.
-        elements: per-element minifloat-quantized values (float, same dtype
-                  as ``x``), shape ``[n_blocks, B]``. Already in scaled
-                  (block-relative) space — multiply by ``2 ** scales`` to
-                  get the dequantized tensor.
+        scales:   per-block shared exponent in log2 domain (float), shape
+                  ``[n_blocks, 1]``. Recompose with ``2 ** scales``.
+        elements: per-element minifloat values in block-relative space,
+                  shape ``[n_blocks, B]``.
     """
     B = mxfp_meta.block_size
-    assert x.numel() % B == 0, (
-        f"Input tensor size {x.numel()} is not divisible by block size {B}."
-    )
+    assert (
+        x.numel() % B == 0
+    ), f"Input tensor size {x.numel()} is not divisible by block size {B}."
     n_blocks = x.numel() // B
 
     x = x.flatten().reshape(n_blocks, B)
+    magnitude = x.abs().to(torch.float32)
+    # quantile() sorts each block; at percentile 1.0 the result is the maximum,
+    # which is also the only path that avoids quantile's input-size limit.
     per_block_max = (
-        x.abs().to(torch.float32).quantile(percentile, dim=1, keepdim=True) + 1e-9
+        magnitude.amax(dim=1, keepdim=True)
+        if percentile == 1.0
+        else magnitude.quantile(percentile, dim=1, keepdim=True)
+    )
+    nonzero_block = per_block_max > 0
+    safe_block_max = torch.where(
+        nonzero_block,
+        per_block_max,
+        torch.ones_like(per_block_max),
     )
-    scales = per_block_max.log2().ceil()
+    # OCP MX shared exponent: place the block maximum at the top of the
+    # element format's exponent range, so the elements use their full range.
+    scales = safe_block_max.log2().floor() - mxfp_meta.element_max_exponent
+    scale_bias = 2 ** (mxfp_meta.scale_exp_bits - 1) - 1
     scales = scales.clamp(
-        min=-(2 ** (mxfp_meta.scale_exp_bits - 1)),
-        max=2 ** (mxfp_meta.scale_exp_bits - 1) - 1,
+        min=-scale_bias,
+        max=2**mxfp_meta.scale_exp_bits - 1 - scale_bias,
     )
+    scales = torch.where(nonzero_block, scales, torch.zeros_like(scales))
 
     q_tensor = x / 2**scales
-    elements = minifloat_denorm_quantizer(
+    elements = minifloat_quantizer_sim(
         q_tensor,
-        width=mxfp_meta.element_frac_bits + mxfp_meta.element_exp_bits + 1,
-        exponent_width=mxfp_meta.element_exp_bits,
+        minifloat_meta=mxfp_meta.element_meta,
+        output_dtype=x.dtype,
     )
 
     return scales, elements
@@ -55,13 +65,11 @@ def compose_mxfp_tensor(
     mxfp_meta: MXFPMeta,
     output_dtype: torch.dtype,
 ) -> Tensor:
-    """Reconstruct the dequantized tensor from Plena-style MXFP components.
+    """Reconstruct the dequantized tensor from MXFP components.
 
-    ``elements`` are the already-block-relative minifloat values; the shared
-    log2 scale just shifts them back into the original block magnitude.
-    ``mxfp_meta`` is unused for the recompose itself (kept in the signature
-    so callers don't need to know whether the impl is Plena- or OCP-style).
+    ``elements`` are block-relative minifloat values; the shared log2 scale
+    shifts them back into the original block magnitude.
     """
-    del mxfp_meta  # only needed by the OCP-style impl
+    del mxfp_meta  # shape and format already fixed by extraction
     dequantized = (elements * 2**shared_scales).flatten().to(output_dtype)
     return dequantized