remove unneccessary helpers

Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>

Author: shanjiaz

src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -22,7 +22,7 @@
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
 from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
-from compressed_tensors.quantization.utils import calculate_qparam_shape, can_quantize
+from compressed_tensors.quantization.utils import can_quantize
 from torch import Tensor
 
 
@@ -64,7 +64,6 @@ def compression_param_info(
         """
         pack_factor = 32 // quantization_args.num_bits
         packed_size = math.ceil(weight_shape[1] / pack_factor)
-        packed_size_zp = math.ceil(weight_shape[0] / pack_factor)
         output = {
             "weight_packed": (torch.Size((weight_shape[0], packed_size)), torch.int32),
             "weight_shape": (torch.Size((2,)), torch.int32),
@@ -75,17 +74,20 @@ def compression_param_info(
             QuantizationStrategy.GROUP.value,
             QuantizationStrategy.CHANNEL.value,
         ]:
-            # Use centralized calculation for consistency and correctness
-            num_groups, scale_shape = calculate_qparam_shape(
-                weight_shape, quantization_args
+            scale_cols = (
+                1
+                if quantization_args.strategy == QuantizationStrategy.CHANNEL.value
+                else math.ceil(weight_shape[1] / quantization_args.group_size)
+            )
+            output["weight_scale"] = (
+                torch.Size((weight_shape[0], scale_cols)),
+                quantization_args.scale_dtype,
             )
-            output["weight_scale"] = (scale_shape, quantization_args.scale_dtype)
 
             # Add weight_zero_point for asymmetric quantization
-            # Zero point has same num_groups as scale, but with packed rows
             if not quantization_args.symmetric:
                 output["weight_zero_point"] = (
-                    torch.Size((packed_size_zp, num_groups)),
+                    torch.Size((math.ceil(weight_shape[0] / pack_factor), scale_cols)),
                     torch.int32,
                 )
 
@@ -201,9 +203,7 @@ def compress_zp(
             QuantizationStrategy.GROUP.value,
             QuantizationStrategy.CHANNEL.value,
         ]:
-            return pack_to_int32(
-                zero_point, quantization_args.num_bits, packed_dim=0
-            ).contiguous()
+            return pack_to_int32(zero_point, quantization_args.num_bits, packed_dim=0)
         return zero_point
 
 

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -35,7 +35,7 @@
 from compressed_tensors.quantization.lifecycle.forward import (
     wrap_module_forward_quantized,
 )
-from compressed_tensors.quantization.utils import calculate_qparam_shape, strategy_cdiv
+from compressed_tensors.quantization.utils import strategy_cdiv
 from compressed_tensors.utils import (
     disable_hf_hook,
     get_execution_device,
@@ -198,25 +198,26 @@ def initialize_qparams(
         return
 
     # 1. Infer expected scale/zp shape
-    if strategy == QuantizationStrategy.TOKEN:
+    if strategy == QuantizationStrategy.TENSOR:
+        expected_shape = (1,)
+
+    elif strategy == QuantizationStrategy.TOKEN:
         raise ValueError("Cannot perform static token quantization")
 
-    elif strategy in (
-        QuantizationStrategy.TENSOR,
-        QuantizationStrategy.CHANNEL,
-        QuantizationStrategy.GROUP,
-        QuantizationStrategy.TENSOR_GROUP,
-    ):
-        # Validate shape requirements
-        if strategy == QuantizationStrategy.CHANNEL and len(observed_shape) < 2:
+    elif strategy == QuantizationStrategy.CHANNEL:
+        if len(observed_shape) < 2:
             raise ValueError("Channel quant requires at least 2 observed dimensions")
-        if strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP):
-            assert quantization_args.group_size is not None
-            if len(observed_shape) < 1:
-                raise ValueError("Group quant requires at least 1 observed dimension")
 
-        # Use unified helper to calculate expected shape
-        _, expected_shape = calculate_qparam_shape(observed_shape, quantization_args)
+        expected_shape = (observed_shape[-2], 1)
+
+    elif strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP):
+        assert quantization_args.group_size is not None
+        if len(observed_shape) < 1:
+            raise ValueError("Group quant requires at least 1 observed dimension")
+
+        group_size = quantization_args.group_size
+        num_groups = strategy_cdiv(observed_shape[-1], group_size, strategy)
+        expected_shape = (*observed_shape[:-1], num_groups)
 
         # initialize activation ordering if applicable
         if actorder == ActivationOrdering.GROUP:

src/compressed_tensors/quantization/utils/helpers.py

@@ -53,7 +53,6 @@
     "calculate_qparams",
     "generate_gparam",
     "strategy_cdiv",
-    "calculate_qparam_shape",
 ]
 
 # target the self_attn layer
@@ -460,50 +459,6 @@ def strategy_cdiv(
     return dividend
 
 
-def calculate_qparam_shape(
-    weight_shape: torch.Size,
-    quantization_args: QuantizationArgs,
-) -> Tuple[int, torch.Size]:
-    """
-    Calculate the number of groups and scale/zero_point shape for quantization.
-
-    This centralizes the logic for determining quantization parameter shapes,
-    ensuring consistency with initialize_qparams and avoiding floor division bugs.
-
-    :param weight_shape: shape of the weight tensor to be quantized
-    :param quantization_args: quantization configuration
-    :return: tuple of (num_groups, expected_shape) where:
-        - num_groups: number of quantization groups
-        - expected_shape: shape for scale/zero_point tensors
-          (weight_shape[0], num_groups)
-    """
-    strategy = quantization_args.strategy
-
-    if strategy == QuantizationStrategy.TENSOR:
-        num_groups = 1
-        expected_shape = (1,)
-
-    elif strategy == QuantizationStrategy.CHANNEL:
-        num_groups = 1
-        expected_shape = (weight_shape[0], 1)
-
-    elif strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP):
-        group_size = quantization_args.group_size
-        if group_size is None:
-            raise ValueError(f"{strategy} quantization requires group_size to be set")
-
-        num_groups = strategy_cdiv(weight_shape[-1], group_size, strategy)
-        expected_shape = (weight_shape[0], num_groups)
-
-    else:
-        raise ValueError(
-            f"Unsupported quantization strategy: {strategy}. "
-            f"Supported strategies: TENSOR, CHANNEL, GROUP, TENSOR_GROUP"
-        )
-
-    return num_groups, expected_shape
-
-
 def _get_dtype_eps(dtype: torch.dtype) -> float:
     if dtype == FP8_E4M3_DATA.dtype:
         return 0.125