Update GemLite to support vLLM V1 (#2199)

* update to forward_functional()

* add 8-bit symmetric case

* ruff

* fix test

Author: mobicham

test/quantization/test_config_serialization.py

@@ -63,8 +63,6 @@
     GemliteUIntXWeightOnlyConfig(
         group_size=128,  # Optional, has default of 64
         bit_width=8,  # Optional, has default of 4
-        packing_bitwidth=8,  # Optional, has default of 32
-        contiguous=True,  # Optional, has default of None
     ),
     FPXWeightOnlyConfig(ebits=4, mbits=8),
     # Sparsity configs

torchao/dtypes/uintx/gemlite_layout.py

@@ -22,7 +22,6 @@
 
 try:
     import gemlite
-    from gemlite.core import GemLiteLinearTriton
 except:
     gemlite = None
 
@@ -51,18 +50,6 @@ def _same_metadata(
     )
 
 
-def scale_activations_no_scaling(x):
-    return x, None
-
-
-def scale_activations_int8(x):
-    x_shape = x.shape
-    out_x = x.view(-1, x.shape[-1])
-    scaled_x = torch.abs(out_x).amax(axis=1, keepdim=True) / 127
-    out_x = torch.round(out_x / scaled_x).to(dtype=torch.int8)
-    return out_x.view(x_shape), scaled_x
-
-
 def get_gemlite_quant_kwargs(bit_width, group_size, dtype):
     from torchao.quantization.quant_primitives import MappingType, ZeroPointDomain
 
@@ -93,8 +80,6 @@ def get_gemlite_aqt_kwargs(
     weight,
     group_size=64,
     bit_width=4,
-    packing_bitwidth=32,
-    contiguous=None,
     use_hqq=True,
 ):
     if gemlite is None:
@@ -106,12 +91,7 @@ def get_gemlite_aqt_kwargs(
         4,
         8,
     ], f"gemlite only works with bit_width 4,8 but got {bit_width}"
-    assert packing_bitwidth in [
-        8,
-        16,
-        32,
-        None,
-    ], f"gemlite needs packing_bitwidth in [8, 16, 32] but got {packing_bitwidth}"
+
     assert weight.dtype in [torch.float16, torch.bfloat16], (
         f"gemlite only works with dtype torch.float16 or torch.bfloat16 but got {weight.dtype}"
     )
@@ -127,8 +107,6 @@ def get_gemlite_aqt_kwargs(
     aqt_kwargs["_layout"] = GemlitePackedLayout(
         group_size=group_size,
         bit_width=bit_width,
-        packing_bitwidth=packing_bitwidth,
-        contiguous=contiguous,
     )
     aqt_kwargs["use_hqq"] = use_hqq
     return aqt_kwargs
@@ -138,8 +116,6 @@ def get_gemlite_aqt_kwargs(
 class GemlitePackedLayout(Layout):
     group_size: Optional[int] = 64
     bit_width: int = 4
-    packing_bitwidth: int = None
-    contiguous: bool = None
 
 
 @register_layout(GemlitePackedLayout)
@@ -216,13 +192,18 @@ def from_plain(
         group_size, bit_width = _layout.group_size, _layout.bit_width
         out_features, in_features = int_data.shape
 
-        gemlite_linear = gemlite.helper.A16Wn(device=int_data.device).from_weights(
-            int_data, scale, zero_point, bit_width, group_size, bias=None
-        )
+        if bit_width == 8 and group_size == in_features:
+            gemlite_linear = gemlite.helper.A16W8(device=int_data.device).from_weights(
+                int_data, scales=scale, bias=None
+            )
+        else:
+            gemlite_linear = gemlite.helper.A16Wn(device=int_data.device).from_weights(
+                int_data, scale, zero_point, bit_width, group_size, bias=None
+            )
 
         gemlite_kwargs = {
+            "in_features": in_features,
             "out_features": out_features,
-            "scaled_activations": gemlite_linear.scaled_activations,
             "meta_args": gemlite_linear.get_meta_args(),
         }
 
@@ -253,20 +234,17 @@ def _apply_fn_to_data(self, fn):
 
     def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         device = self.packed_weight.device
-        elements_per_sample = self._layout.packing_bitwidth // self._layout.bit_width
-        in_features = (
-            self.packed_weight.numel() * elements_per_sample
-        ) // self.gemlite_kwargs["out_features"]
         int_data = (
             gemlite.bitpack.unpack_over_rows(
                 self.packed_weight.cuda(),
                 W_nbits=self._layout.bit_width,
-                num_output_rows=in_features,
+                num_output_rows=self.gemlite_kwargs["out_features"],
                 dtype=torch.uint8,
             )
             .t()
             .contiguous()
         ).to(device)
+
         scale = self.scale.t().contiguous()
         zero_point = self.zero_point.t().contiguous()
 
@@ -353,42 +331,21 @@ def block_size(self):
         return (1, self._layout.group_size)
 
 
-# logic taken from gemlite's core.py
-def _matmul_type_fn(batch_size: int, bit_width: int) -> str:
-    if batch_size > 64:
-        return "GEMM"
-    elif batch_size > 1:
-        return "GEMM_SPLITK"
-    else:
-        return gemlite.core.get_default_gemv(bit_width)
-
-
 def _linear_fp_act_int4_weight_gemlite_impl(input_tensor, weight_tensor, bias=None):
     if hasattr(weight_tensor, "tensor_impl"):
         weight_impl = weight_tensor.tensor_impl
     else:
         weight_impl = weight_tensor
 
-    batch_size = input_tensor.view(-1, input_tensor.shape[-1]).shape[0]
-    matmul_type = _matmul_type_fn(batch_size, weight_impl._layout.bit_width)
-
-    if weight_impl.gemlite_kwargs["scaled_activations"]:
-        scale_activations = scale_activations_int8
-    else:
-        scale_activations = scale_activations_no_scaling
-
-    return GemLiteLinearTriton.forward_functional(
+    return gemlite.core.forward_functional(
         x=input_tensor,
         bias=bias,
-        matmul_type=matmul_type,
-        out_features=weight_impl.gemlite_kwargs["out_features"],
-        scale_activations=scale_activations,
-        meta_args=weight_impl.gemlite_kwargs["meta_args"],
         tensor_args=(
             weight_impl.packed_weight,
             weight_impl.scale,
             weight_impl.zero_point,
         ),
+        meta_args=weight_impl.gemlite_kwargs["meta_args"],
     )
 
 

torchao/quantization/quant_api.py

@@ -979,8 +979,6 @@ class GemliteUIntXWeightOnlyConfig(AOBaseConfig):
 
     group_size: Optional[int] = 64
     bit_width: int = 4
-    packing_bitwidth: int = 32
-    contiguous: Optional[bool] = None
     set_inductor_config: bool = True
 
 
@@ -994,8 +992,6 @@ def _gemlite_uintx_weight_only_transform(
 ):
     group_size = config.group_size
     bit_width = config.bit_width
-    packing_bitwidth = config.packing_bitwidth
-    contiguous = config.contiguous
     if config.set_inductor_config:
         torchao.quantization.utils.recommended_inductor_config_setter()
 
@@ -1006,9 +1002,7 @@ def _gemlite_uintx_weight_only_transform(
     use_hqq = True if bit_width == 4 else False
     new_weight = to_affine_quantized_intx(
         weight,
-        **get_gemlite_aqt_kwargs(
-            weight, group_size, bit_width, packing_bitwidth, contiguous, use_hqq
-        ),
+        **get_gemlite_aqt_kwargs(weight, group_size, bit_width, use_hqq),
     )
     module.weight = torch.nn.Parameter(new_weight, requires_grad=False)
     module.extra_repr = types.MethodType(_linear_extra_repr, module)