[ Misc ] Clean Up CompressedTensorsW8A8 (#6113)

Author: robertgshaw2-redhat

tests/quantization/test_compressed_tensors.py

@@ -9,8 +9,7 @@
 from vllm import SamplingParams
 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
     CompressedTensorsLinearMethod, CompressedTensorsW4A16Sparse24,
-    CompressedTensorsW8A8DynamicToken, CompressedTensorsW8A8StaticTensor,
-    CompressedTensorsWNA16)
+    CompressedTensorsW8A8, CompressedTensorsWNA16)
 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
     QuantizationType)
 
@@ -38,9 +37,10 @@ def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
                           CompressedTensorsLinearMethod)
         assert isinstance(down_proj.quant_method,
                           CompressedTensorsLinearMethod)
-        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8StaticTensor)
+        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8)
 
         assert qkv_proj.scheme.strategy == strategy
+        assert qkv_proj.scheme.is_static_input_scheme
         expected_type = (torch.int8 if quant_type == QuantizationType.INT else
                          torch.float8_e4m3fn)
 
@@ -79,7 +79,8 @@ def test_compressed_tensors_w8a8_dynanmic_per_token(vllm_runner, model_args):
         qkv_proj = layer.self_attn.qkv_proj
 
         assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
-        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8DynamicToken)
+        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8)
+        assert not qkv_proj.scheme.is_static_input_scheme
         assert qkv_proj.scheme.strategy == strategy
         assert qkv_proj.weight.dtype is torch.int8
 

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -9,8 +9,7 @@
 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
     W4A16SPARSE24_SUPPORTED_BITS, WNA16_SUPPORTED_BITS,
     CompressedTensorsScheme, CompressedTensorsW4A16Sparse24,
-    CompressedTensorsW8A8DynamicToken, CompressedTensorsW8A8StaticTensor,
-    CompressedTensorsWNA16)
+    CompressedTensorsW8A8, CompressedTensorsWNA16)
 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
     CompressionFormat, QuantizationArgs, QuantizationStrategy,
     find_first_name_or_class_match)
@@ -150,12 +149,12 @@ def _get_schema(self, weight_quant: BaseModel,
 
         if self.quant_format == CompressionFormat.int_quantized.value:
             if self._is_static_tensor_w8a8(weight_quant, input_quant):
-                return CompressedTensorsW8A8StaticTensor(
-                    strategy=weight_quant.strategy)
+                return CompressedTensorsW8A8(strategy=weight_quant.strategy,
+                                             is_static_input_scheme=True)
 
             if self._is_dynamic_token_w8a8(weight_quant, input_quant):
-                return CompressedTensorsW8A8DynamicToken(
-                    strategy=weight_quant.strategy)
+                return CompressedTensorsW8A8(strategy=weight_quant.strategy,
+                                             is_static_input_scheme=False)
 
         raise NotImplementedError(
             "No compressed-tensors compatible scheme was found.")

vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py

@@ -3,9 +3,6 @@
     CompressedTensorsUnquantized)
 from .compressed_tensors_w4a16_24 import (  # noqa: F401
     W4A16SPARSE24_SUPPORTED_BITS, CompressedTensorsW4A16Sparse24)
-from .compressed_tensors_w8a8_dynamictoken import (  # noqa: F401, E501
-    CompressedTensorsW8A8DynamicToken)
-from .compressed_tensors_w8a8_statictensor import (  # noqa: F401, E501
-    CompressedTensorsW8A8StaticTensor)
+from .compressed_tensors_w8a8 import CompressedTensorsW8A8  # noqa: F401
 from .compressed_tensors_wNa16 import WNA16_SUPPORTED_BITS  # noqa: F401
 from .compressed_tensors_wNa16 import CompressedTensorsWNA16  # noqa: F401

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8.py

@@ -3,6 +3,7 @@
 import torch
 from torch.nn import Parameter
 
+from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsScheme)
 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
@@ -12,8 +13,9 @@
 
 class CompressedTensorsW8A8(CompressedTensorsScheme):
 
-    def __init__(self, strategy: str):
+    def __init__(self, strategy: str, is_static_input_scheme: bool):
         self.strategy = strategy
+        self.is_static_input_scheme = is_static_input_scheme
 
     # Cutlass kernels support only per-tensor and per-channel cases.
     # So if we have a fused module (QKV, MLP) with per tensor scales (thus N
@@ -36,6 +38,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             layer.weight_scale = Parameter(weight_scale_channel,
                                            requires_grad=False)
 
+        # transpose weights for cutlass.
+        weight = layer.weight
+        layer.weight = Parameter(weight.t(), requires_grad=False)
+
     def create_weights(self, layer: torch.nn.Module,
                        output_partition_sizes: List[int],
                        input_size_per_partition: int,
@@ -75,3 +81,29 @@ def create_weights(self, layer: torch.nn.Module,
             "output_dim": 0,
             "weight_loader": weight_loader,
         })
+
+        # INPUT SCALE
+        # Static quantization:  load from disk.
+        if self.is_static_input_scheme:
+            input_scale = Parameter(torch.empty(1, dtype=torch.float32),
+                                    requires_grad=False)
+            layer.register_parameter("input_scale", input_scale)
+            set_weight_attrs(input_scale, {
+                "weight_loader": weight_loader,
+                "ignore_warning": True,
+            })
+        # Dynamic quantization: set to None.
+        else:
+            layer.input_scale = None
+
+    def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor):
+        # ops.scaled_int8_quant supports both dynamic and static quant.
+        # * dynamic, layer.input_scale is None and x_scale computed from x.
+        # * static, layer.input_scale is scalar and x_scale is input_scale.
+        x_q, x_scale = ops.scaled_int8_quant(x, layer.input_scale)
+
+        return ops.cutlass_scaled_mm(x_q,
+                                     layer.weight,
+                                     scale_a=x_scale,
+                                     scale_b=layer.weight_scale,
+                                     out_dtype=x.dtype)