compressed-tensors 0.9.5a20250530__tar.gz → 0.9.5a20250602__tar.gz

{compressed_tensors-0.9.5a20250530/src/compressed_tensors.egg-info → compressed_tensors-0.9.5a20250602}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.9.5a20250530
+Version: 0.9.5a20250602
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/src/compressed_tensors/quantization/lifecycle/forward.py

@@ -227,31 +227,42 @@ def _process_quantization(
             perm = torch.argsort(g_idx)
             x = safe_permute(x, perm, dim=1)
 
-        # TODO: experiment with vectorizing for loop for performance
-        end = 0
-        for index, group_count in enumerate(group_sizes):
-            sc = scale[:, index].view(-1, 1)
-            zp = zero_point[:, index].view(-1, 1) if zero_point is not None else None
-
-            start = end
-            end = start + group_count
-            if do_quantize:
-                output[:, start:end] = _quantize(
-                    x=x[:, start:end],
-                    scale=sc,
-                    zero_point=zp,
-                    q_min=q_min,
-                    q_max=q_max,
-                    args=args,
-                    dtype=dtype,
-                    global_scale=global_scale,
-                )
+        x = torch.reshape(
+            x,
+            (
+                x.shape[0],
+                ceil(x.shape[1] / group_size),
+                group_size,
+            ),
+        )
 
-            if do_dequantize:
-                input = output[:, start:end] if do_quantize else x[:, start:end]
-                output[:, start:end] = _dequantize(
-                    x_q=input, scale=sc, zero_point=zp, global_scale=global_scale
-                )
+        if do_quantize:
+            output = _quantize(
+                x=x,
+                scale=scale.unsqueeze(-1),
+                zero_point=zero_point.unsqueeze(-1) if zero_point is not None else None,
+                dtype=dtype,
+                global_scale=global_scale,
+                q_min=q_min,
+                q_max=q_max,
+                args=args,
+            )
+
+        if do_dequantize:
+            input = output if do_quantize else x
+            output = _dequantize(
+                x_q=input,
+                scale=scale.unsqueeze(-1),
+                zero_point=zero_point.unsqueeze(-1) if zero_point is not None else None,
+                global_scale=global_scale,
+            )
+
+        output = torch.reshape(
+            output,
+            (output.shape[0], output.shape[1] * output.shape[2]),
+        )
+
+        output = output.to(output_dtype)
 
         if not is_column_order:
             output = safe_permute(output, torch.argsort(perm), dim=1)

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/src/compressed_tensors/quantization/quant_args.py

@@ -53,6 +53,7 @@ class FP4_E2M1_DATA(FloatArgs):
     min = -6.0
 
     @staticmethod
+    @torch.compile
     def cast_to_fp4(x):
         sign = torch.sign(x)
         x = torch.abs(x)

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/src/compressed_tensors/quantization/utils/helpers.py

@@ -81,7 +81,7 @@ def calculate_qparams(
         currently only applied/supported for Fp4
 
     :return: tuple of the calculated scale(s) and zero point(s). For FP4, the calculated
-        scale if of dtype FP8
+        scale is of dtype FP8
     """
     # based on the implementations for consuming quantized values,
     # 0.0 must always be representable within the quantized range
@@ -490,7 +490,6 @@ def generate_global_scale(
     attempts to use the entire FP8 dtype range while mapping a per-group max
     to the FP4 max.
     """
-    scale_dtype = scale_data.dtype
     tensor_amax = torch.abs(input_tensor.data).max().to(dtype)
     global_scale = scale_data.max * quant_data.max / tensor_amax
     return global_scale.to(dtype)

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.9.5.a20250530'
+__version__ = version = '0.9.5.a20250602'
 __version_tuple__ = version_tuple = (0, 9, 5)

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602/src/compressed_tensors.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.9.5a20250530
+Version: 0.9.5a20250602
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/tests/test_compressors/quantized_compressors/test_fp8_quant.py

@@ -61,8 +61,8 @@ def make_dummy_g_idx(columns: int, group_size: int) -> torch.Tensor:
         [
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1), dtype=torch.int8),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8), dtype=torch.int8),
         ],
         [
             QuantizationStrategy.CHANNEL,
@@ -79,7 +79,7 @@ def test_quant_format(strategy, group_size, sc, zp):
         "dummy.weight_zero_point": torch.tensor(zp, dtype=torch.float32),
     }
     if group_size is not None:
-        dense_state_dict["dummy.weight_g_idx"] = make_dummy_g_idx(512, group_size)
+        dense_state_dict["dummy.weight_g_idx"] = make_dummy_g_idx(1024, group_size)
 
     quant_config = get_dummy_quant_config(strategy=strategy, group_size=group_size)
 

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/tests/test_compressors/quantized_compressors/test_int_quant.py

@@ -53,8 +53,8 @@ def get_dummy_quant_config(strategy, group_size=None, symmetric=True):
             QuantizationStrategy.GROUP,
             True,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1), dtype=torch.int8),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8), dtype=torch.int8),
         ],
         [
             QuantizationStrategy.CHANNEL,

{compressed_tensors-0.9.5a20250530 → compressed_tensors-0.9.5a20250602}/tests/test_quantization/lifecycle/test_forward.py

@@ -108,8 +108,8 @@ def test_forward_quantize(
             "int",
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1)),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8)),
             None,
         ),
         (
@@ -117,8 +117,8 @@ def test_forward_quantize(
             "int",
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1)),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8)),
             make_dummy_g_idx(1024, 128),
         ),
         (
@@ -135,8 +135,8 @@ def test_forward_quantize(
             "float",
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1)),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8)),
             None,
         ),
         (
@@ -144,8 +144,8 @@ def test_forward_quantize(
             "float",
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1)),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8)),
             make_dummy_g_idx(1024, 128),
         ),
     ],
@@ -174,8 +174,8 @@ def test_quantize(num_bits, type, strategy, group_size, scale, zero_point, g_idx
             "int",
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1)),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8)),
             None,
         ),
         (
@@ -183,8 +183,8 @@ def test_quantize(num_bits, type, strategy, group_size, scale, zero_point, g_idx
             "int",
             QuantizationStrategy.GROUP,
             128,
-            torch.rand((512, 8, 1)) * 0.01,
-            torch.zeros((512, 8, 1)),
+            torch.rand((512, 8)) * 0.01,
+            torch.zeros((512, 8)),
             make_dummy_g_idx(1024, 128),
         ),
     ],