update

Author: dsikka

src/compressed_tensors/quantization/utils/helpers.py

@@ -28,8 +28,9 @@
 )
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
 from compressed_tensors.quantization.utils.mxfp4_utils import (
-    maybe_convert_from_mxfp4_scale,
-    maybe_convert_to_mxfp4_scales,
+    generate_mxfp4_scales,
+    maybe_convert_from_mxfp4_exp,
+    should_generatre_mxfp4_scales,
 )
 from compressed_tensors.utils import deprecated
 from loguru import logger
@@ -92,8 +93,10 @@ def calculate_qparams(
     # 1. Generate scale and zero-point
     if quantization_args.symmetric:
         max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals))
-        # scales = max_val_pos / (float(bit_range) / 2)
-        scales = maybe_convert_to_mxfp4_scales(max_val_pos)
+        if should_generatre_mxfp4_scales(args=quantization_args):
+            scales = generate_mxfp4_scales(x=max_val_pos)
+        else:
+            scales = max_val_pos / (float(bit_range) / 2)
         zero_points = torch.zeros(scales.shape, device=device, dtype=min_vals.dtype)
     else:
         if (
@@ -117,10 +120,10 @@ def calculate_qparams(
             scales, dtype=quantization_args.scale_dtype
         )
 
-    # Optionally remove exponent
-    scales = maybe_convert_from_mxfp4_scale(quantization_args, scales)
+    # 4. Optionally remove exponent
+    scales = maybe_convert_from_mxfp4_exp(quantization_args, scales)
 
-    # 4. Update any 0s with small values to
+    # 5. Update any 0s with small values to
     # prevent div by 0
     eps = _get_dtype_eps(
         dtype=quantization_args.scale_dtype
@@ -133,7 +136,7 @@ def calculate_qparams(
         scales,
     )
 
-    # 5. Round the zp to zp_dtype
+    # 6. Round the zp to zp_dtype
     zero_points = round_to_quantized_type_dtype(
         zero_points, dtype=quantization_args.zp_dtype, cast_to_original_dtype=False
     )

src/compressed_tensors/quantization/utils/mxfp4_utils.py

@@ -21,17 +21,21 @@
 
 
 __all__ = [
-    "maybe_convert_from_mxfp4_scale",
+    "maybe_convert_from_mxfp4_exp",
     "generate_mxfp4_scales",
     "round_to_power_2",
-    "maybe_convert_to_mxfp4_scales",
+    "should_generatre_mxfp4_scales",
 ]
 
 # Reference: https://github.com/vllm-project/vllm/blob/main/tests/quantization/reference_mxfp4.py # noqa: E501
 
 
-def maybe_convert_from_mxfp4_scale(
-    args: QuantizationArgs, scale: torch.Tensor, dtype: torch.dtype = torch.bfloat16
+def should_generatre_mxfp4_scales(args: QuantizationArgs):
+    return args.num_bits == 4 and args.type == "float" and args.group_size == 32
+
+
+def maybe_convert_from_mxfp4_exp(
+    args: QuantizationArgs, scale: torch.Tensor
 ) -> torch.Tensor:
     """
     Converts mxfp4 scales. Scales are powers of 2, with the
@@ -41,30 +45,14 @@ def maybe_convert_from_mxfp4_scale(
     :param scale: uint8 exponent scale
     :param dtype: dense dtype
     """
-    is_mxfp4 = args.num_bits == 4 and args.type == "float" and args.group_size == 32
-    if is_mxfp4:
+    original_dtype = scale.dtype
+    if should_generatre_mxfp4_scales(args):
         scale_exp = scale.to(torch.int32) - 127
         scale = 2.00 ** (scale_exp.to(torch.float))
-        return scale.to(dtype)
+        return scale.to(original_dtype)
     return scale
 
 
-def maybe_convert_to_mxfp4_scales(
-    args: QuantizationArgs, scales: torch.Tensor
-) -> torch.Tensor:
-    """
-    Conver the scales to be mxfp4 compatible scales, if quant args are FP4 with group_size 32.
-    If not, return original scales
-
-    :param args: quantization args
-    :param scales: scales to update
-    """
-    is_mxfp4 = args.num_bits == 4 and args.type == "float" and args.group_size == 32
-    if is_mxfp4:
-        return generate_mxfp4_scales(x=scales)
-    return scales
-
-
 def round_to_power_2(x: torch.Tensor) -> torch.Tensor:
     """
     Round values to the closest power of 2.
@@ -99,28 +87,17 @@ def round_to_power_2(x: torch.Tensor) -> torch.Tensor:
     return block_max_uint.to(torch.uint16).view(torch.bfloat16)
 
 
-def generate_mxfp4_scales(x: torch.Tensor, clamp: bool = False) -> torch.Tensor:
+def generate_mxfp4_scales(x: torch.Tensor) -> torch.Tensor:
     """
     Generate mxfp4 scales. The scales require the following steps
     1. Round to the closest power of 2
     2. Convert to exponent
-    3. Optionally, store in uint8
 
     Called when calculating qparams using observers.
 
     :param x: tensor to round to closest power of 2
-    :returns uint8 scales as exponents
+    :returns scales as exponents
     """
     # Round to closest power of 2
     scale_power_2 = round_to_power_2(x)
-    # Convert to exponent
-    scale_exp = 127 + torch.floor(torch.log2(scale_power_2)).to(torch.int32) - 2
-    # Clamp and store in uint8, as expected by mxfp4
-    if clamp:
-        scale_exp = torch.clamp(
-            scale_exp,
-            max=torch.iinfo(torch.uint8).max,
-            min=torch.iinfo(torch.uint8).min,
-        )
-        return scale_exp.to(torch.uint8)
-    return scale_exp
+    return 127 + torch.floor(torch.log2(scale_power_2)) - 2

tests/test_quantization/test_utils/test_mxfp4_utils.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 import torch
+from compressed_tensors.quantization import round_to_quantized_type_dtype
 from compressed_tensors.quantization.utils import (
     generate_mxfp4_scales,
-    maybe_convert_from_mxfp4_scale,
+    maybe_convert_from_mxfp4_exp,
     round_to_power_2,
 )
 
@@ -77,7 +78,6 @@ def test_mxfp4_scales_e2e():
     max_vals = torch.max(max_vals, torch.zeros_like(max_vals))
     block_max = torch.max(torch.abs(min_vals), torch.abs(max_vals))
 
-    scales_generated = generate_mxfp4_scales(block_max, clamp=True)
     args = QuantizationArgs(
         num_bits=4,
         type=QuantizationType.FLOAT,
@@ -86,7 +86,11 @@ def test_mxfp4_scales_e2e():
         scale_dtype=torch.uint8,
         zp_dtype=torch.uint8,
     )
-    converted_ct = maybe_convert_from_mxfp4_scale(args=args, scale=scales_generated)
+
+    scales = generate_mxfp4_scales(block_max)
+    scales = round_to_quantized_type_dtype(scales, dtype=args.scale_dtype)
+
+    converted_ct = maybe_convert_from_mxfp4_exp(args=args, scale=scales)
 
     scales_exp = torch.log2(converted_ct)
     block_max_exp = torch.floor(torch.log2(round_to_power_2(block_max))) - 2