FourOverSix

Author: Krishna Teja Chitty Venkata

src/compressed_tensors/compressors/nvfp4/base.py

@@ -18,6 +18,7 @@
     QuantizationType,
 )
 from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
+from compressed_tensors.quantization.quant_args import round_to_quantized_type_args
 from compressed_tensors.utils import TensorStateDict, getattr_chain
 
 
@@ -60,6 +61,59 @@ def _compress_scale(
     def _decompress_scale(cls, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
         return scale.to(dtype)
 
+    @classmethod
+    def _adjust_scale_for_four_over_six(
+        cls,
+        weight: torch.Tensor,
+        scale: torch.Tensor,
+        global_scale: torch.Tensor | None,
+        weights: QuantizationArgs,
+    ) -> torch.Tensor:
+        """
+        Pre-adjust per-group scales for Four Over Six: for each group, compare
+        MSE of standard quantization (scale-to-6) vs alternative (scale-to-4,
+        i.e. scale * 1.5). Groups where scale-to-4 wins get their scale
+        multiplied by 1.5 so that the stored scale reflects the actual
+        quantization used.
+        """
+        from compressed_tensors.quantization.utils import calculate_range
+
+        q_min, q_max = calculate_range(weights, weight.device)
+
+        group_size = weights.group_size
+        rows, cols = weight.shape
+        num_groups = cols // group_size
+        w_grouped = weight.reshape(rows, num_groups, group_size)
+
+        if global_scale is not None:
+            eff_scale = scale / global_scale
+        else:
+            eff_scale = scale.clone()
+
+        eff_scale_3d = eff_scale.unsqueeze(-1)
+
+        scaled_a = w_grouped / eff_scale_3d
+        q_a = round_to_quantized_type_args(
+            tensor=scaled_a, args=weights, min=q_min, max=q_max
+        )
+        dq_a = q_a.to(eff_scale.dtype) * eff_scale_3d
+
+        eff_scale_b = eff_scale * 1.5
+        eff_scale_b_3d = eff_scale_b.unsqueeze(-1)
+        scaled_b = w_grouped / eff_scale_b_3d
+        q_b = round_to_quantized_type_args(
+            tensor=scaled_b, args=weights, min=q_min, max=q_max
+        )
+        dq_b = q_b.to(eff_scale_b.dtype) * eff_scale_b_3d
+
+        mse_a = ((w_grouped - dq_a) ** 2).mean(dim=-1)
+        mse_b = ((w_grouped - dq_b) ** 2).mean(dim=-1)
+
+        use_4 = mse_b < mse_a
+        adjusted_scale = scale.clone()
+        adjusted_scale[use_4] = adjusted_scale[use_4] * 1.5
+        return adjusted_scale.to(scale.dtype)
+
     @classmethod
     def compress(
         cls, state_dict: TensorStateDict, scheme: QuantizationScheme
@@ -82,6 +136,11 @@ def compress(
         zero_point = state_dict.get("weight_zero_point", None)
         weights = scheme.weights
 
+        if getattr(weights, "four_over_six", False):
+            scale = cls._adjust_scale_for_four_over_six(
+                weight, scale, global_scale, weights
+            )
+
         quantized_weight = quantize(
             x=weight,
             scale=scale,

src/compressed_tensors/quantization/lifecycle/forward_helpers.py

@@ -6,6 +6,7 @@
 import torch
 from compressed_tensors.quantization.quant_args import (
     QuantizationArgs,
+    QuantizationType,
     round_to_quantized_type_args,
 )
 from compressed_tensors.quantization.utils import maybe_pad_tensor_for_block_quant
@@ -187,6 +188,15 @@ def _quantize_dequantize(
     - Double scale/global_scale division
     - Intermediate quantized dtype allocation
     """
+    if (
+        getattr(args, "four_over_six", False)
+        and args.num_bits == 4
+        and args.type == QuantizationType.FLOAT
+    ):
+        return _four_over_six_quantize_dequantize(
+            x, scale, zero_point, q_min, q_max, args, global_scale
+        )
+
     # compute effective scale once
     if global_scale is not None:
         scale = scale / global_scale
@@ -210,6 +220,56 @@ def _quantize_dequantize(
     return dequant * scale
 
 
+@torch.no_grad()
+def _four_over_six_quantize_dequantize(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    zero_point: torch.Tensor | None,
+    q_min: torch.Tensor,
+    q_max: torch.Tensor,
+    args: QuantizationArgs,
+    global_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    Four Over Six adaptive block scaling: for each group, try quantizing
+    with the standard scale (maps max to 6) and an alternative scale
+    (maps max to 4, i.e. scale * 1.5). Pick whichever yields lower MSE.
+    """
+    if global_scale is not None:
+        eff_scale = scale / global_scale
+    else:
+        eff_scale = scale
+
+    # --- Path A: standard (scale to 6) ---
+    scaled_a = x / eff_scale
+    if zero_point is not None:
+        scaled_a = scaled_a + zero_point.to(x.dtype)
+    q_a = round_to_quantized_type_args(tensor=scaled_a, args=args, min=q_min, max=q_max)
+    dq_a = q_a.to(eff_scale.dtype)
+    if zero_point is not None:
+        dq_a = dq_a - zero_point.to(eff_scale.dtype)
+    dq_a = dq_a * eff_scale
+
+    # --- Path B: scale to 4 (scale * 1.5) ---
+    eff_scale_b = eff_scale * 1.5
+    scaled_b = x / eff_scale_b
+    if zero_point is not None:
+        scaled_b = scaled_b + zero_point.to(x.dtype)
+    q_b = round_to_quantized_type_args(tensor=scaled_b, args=args, min=q_min, max=q_max)
+    dq_b = q_b.to(eff_scale_b.dtype)
+    if zero_point is not None:
+        dq_b = dq_b - zero_point.to(eff_scale_b.dtype)
+    dq_b = dq_b * eff_scale_b
+
+    # --- Per-group MSE comparison ---
+    group_dims = tuple(range(scale.ndim, x.ndim))
+    mse_a = ((x - dq_a) ** 2).mean(dim=group_dims, keepdim=True)
+    mse_b = ((x - dq_b) ** 2).mean(dim=group_dims, keepdim=True)
+
+    use_b = mse_b < mse_a
+    return torch.where(use_b, dq_b, dq_a)
+
+
 @torch.no_grad()
 def _quantize(
     x: torch.Tensor,

src/compressed_tensors/quantization/quant_args.py

@@ -224,6 +224,17 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
             "Observers constructor excluding quantization range or symmetry"
         ),
     )
+    four_over_six: bool = Field(
+        default=False,
+        exclude=True,
+        description=(
+            "Enable Four Over Six (4/6) adaptive block scaling for NVFP4 "
+            "quantization. For each group of values, tries scaling to both 4 "
+            "and 6, selecting the scale with lower MSE. Reduces quantization "
+            "error for near-maximal values in FP4. Only applies to FP4 "
+            "quantization (num_bits=4, type=float)."
+        ),
+    )
 
     @field_serializer("zp_dtype")
     def serialize_dtype(self, dtype: torch.dtype):

src/compressed_tensors/quantization/quant_scheme.py

@@ -180,6 +180,47 @@ def is_preset_scheme(name: str) -> bool:
     ),
 )
 
+NVFP4A16_46 = dict(
+    weights=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TENSOR_GROUP,
+        symmetric=True,
+        dynamic=False,
+        group_size=16,
+        scale_dtype=FP8_E4M3_DATA.dtype,
+        zp_dtype=FP8_E4M3_DATA.dtype,
+        four_over_six=True,
+    )
+)
+
+
+NVFP4_46 = dict(
+    weights=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TENSOR_GROUP,
+        symmetric=True,
+        dynamic=False,
+        group_size=16,
+        scale_dtype=FP8_E4M3_DATA.dtype,
+        zp_dtype=FP8_E4M3_DATA.dtype,
+        four_over_six=True,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TENSOR_GROUP,
+        symmetric=True,
+        dynamic=DynamicType.LOCAL,
+        group_size=16,
+        observer="static_minmax",
+        scale_dtype=FP8_E4M3_DATA.dtype,
+        zp_dtype=FP8_E4M3_DATA.dtype,
+        four_over_six=True,
+    ),
+)
+
 MXFP4A16 = dict(
     weights=QuantizationArgs(
         num_bits=4,
@@ -422,7 +463,9 @@ def is_preset_scheme(name: str) -> bool:
     "FP8_DYNAMIC": FP8_DYNAMIC,
     "FP8_BLOCK": FP8_BLOCK,
     "NVFP4A16": NVFP4A16,
+    "NVFP4A16_46": NVFP4A16_46,
     "NVFP4": NVFP4,
+    "NVFP4_46": NVFP4_46,
     "MXFP4A16": MXFP4A16,
     "MXFP4": MXFP4,
     "MXFP8A16": MXFP8A16,

src/compressed_tensors/quantization/utils/helpers.py

@@ -312,6 +312,7 @@ def generate_gparam(
     scale_data: FloatArgs | None = FP8_E4M3_DATA,
     quant_data: FloatArgs | None = FP4_E2M1_DATA,
     dtype: torch.dtype | None = torch.float32,
+    four_over_six: bool = False,
 ):
     """
     Generate a global scale for an entire tensor (input_tensor).
@@ -321,12 +322,19 @@ def generate_gparam(
     E.g. for NVFP4, group (local) scales are in dtype FP8. The global_scale
     attempts to use the entire FP8 dtype range while mapping a per-group max
     to the FP4 max.
+
+    When four_over_six=True, uses 256 instead of 448 for the FP8 scale max.
+    This allows blocks containing the tensor's largest values to have their
+    FP8 scale multiplied by 1.5 (for scale-to-4) without overflowing
+    (256 * 1.5 = 384 < 448).
     """
     min_vals = torch.min(updated_min_val, torch.zeros_like(updated_min_val))
     max_vals = torch.max(updated_max_val, torch.zeros_like(updated_max_val))
     max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals))
     max_val_pos = torch.clamp(max_val_pos, min=torch.finfo(max_val_pos.dtype).tiny)
-    global_scale = scale_data.max * quant_data.max / max_val_pos
+
+    scale_max = 256.0 if four_over_six else scale_data.max
+    global_scale = scale_max * quant_data.max / max_val_pos
 
     # Replace any NaN or Inf with 1.0. NaN arises when max_val_pos was NaN
     # (clamp does not propagate NaN, so it passes through to the division).