tpu-inference 0.11.1.dev202511220812__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tpu_inference/layers/vllm/quantization/awq.py

@@ -1,11 +1,26 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Optional, Union
 
 import jax
 import jax.numpy as jnp
 import torch
-from jax.sharding import NamedSharding, PartitionSpec
+from jax.sharding import PartitionSpec
+from torch.nn.parameter import Parameter
 from torchax.interop import jax_view, torch_view
-from vllm.logger import init_logger
+from torchax.ops.mappings import t2j
 from vllm.model_executor.layers.fused_moe.layer import FusedMoE
 from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
 from vllm.model_executor.layers.quantization import \
@@ -14,24 +29,29 @@ from vllm.model_executor.layers.quantization.awq import (AWQConfig,
                                                          AWQLinearMethod)
 from vllm.model_executor.layers.quantization.base_config import \
     QuantizeMethodBase
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    is_layer_skipped, unpack_quantized_values_into_int32)
-from vllm.scalar_type import scalar_types
+from vllm.model_executor.layers.quantization.utils.quant_utils import \
+    is_layer_skipped
 
 from tpu_inference.layers.common.quant_methods import AWQ, get_tpu_quant_method
-from tpu_inference.layers.vllm.linear_common import (
-    slice_sharded_tensor_for_concatenation, torch_to_jax_param)
-from tpu_inference.layers.vllm.quantization.common import (
-    JaxCommonConfig, JaxCommonLinearConfig)
+from tpu_inference.layers.common.quantization import awq_u32_unpack_u4
+from tpu_inference.layers.common.utils import \
+    slice_sharded_tensor_for_concatenation
+from tpu_inference.layers.vllm.process_weights.linear_weights import (
+    LinearWeights, process_lienar_weights, shard_linear_weights,
+    to_parameter_list)
+from tpu_inference.layers.vllm.quantization.configs import (
+    VllmQuantConfig, VllmQuantLinearConfig)
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedLinearMethod
+from tpu_inference.logger import init_logger
 
 P = PartitionSpec
+
 logger = init_logger(__name__)
 
 
 @register_quantization_config(get_tpu_quant_method(AWQ))
-class VllmAWQConfig(AWQConfig, JaxCommonConfig):
+class VllmAWQConfig(AWQConfig, VllmQuantConfig):
 
     @classmethod
     def get_name(cls):
@@ -39,7 +59,7 @@ class VllmAWQConfig(AWQConfig, JaxCommonConfig):
 
     def get_supported_act_dtypes(self) -> list[torch.dtype]:
         # NOTE: AWQ checkpoint was quantized with float16. But on TPUs, using
-        # bfloat16 is signifcantly preferred over foat16. This might lead to
+        # bfloat16 is significantly preferred over float16. This might lead to
         # some numeric output change.
         return [torch.bfloat16]
 
@@ -60,72 +80,79 @@ class VllmAWQConfig(AWQConfig, JaxCommonConfig):
 class VllmAWQLinearMethod(AWQLinearMethod):
 
     def __init__(self, quant_config: VllmAWQConfig,
-                 jax_config: JaxCommonLinearConfig):
+                 linear_config: VllmQuantLinearConfig):
         super().__init__(quant_config)
-        self.jax_config = jax_config
-
-        out_sharding, in_sharding = self.jax_config.weight_sharding[:]
-        self.jax_config.weight_sharding = P(in_sharding, None, out_sharding)
-        self.jax_config.scale_sharding = P(in_sharding, out_sharding)
+        self.linear_config = linear_config
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        qweight = layer.qweight
-        qweight = unpack_awq_weight(qweight, qweight.packed_dim)
-
-        group_size = self.quant_config.group_size
-        # Reshape so that each qweight[i] were quantized with same scales[i].
-        qweight = qweight.reshape((-1, group_size, layer.output_size))
-        qweight = torch_to_jax_param(qweight,
-                                     NamedSharding(
-                                         self.jax_config.mesh,
-                                         self.jax_config.weight_sharding),
-                                     self.jax_config.output_sizes,
-                                     self.jax_config.n_shards,
-                                     self.jax_config.fuse_matmuls,
-                                     dim=2,
-                                     jax_dtype=jnp.uint4)
+        assert layer.qweight.packed_dim == layer.qweight.ndim - 1
+        weight = t2j(layer.qweight, use_dlpack=False)
         delattr(layer, "qweight")
-        layer.qweight = qweight
-
-        qzeros = layer.qzeros
-        qzeros = unpack_awq_weight(qzeros, qzeros.packed_dim)
-        qzeros = torch_to_jax_param(qzeros,
-                                    NamedSharding(
-                                        self.jax_config.mesh,
-                                        self.jax_config.scale_sharding),
-                                    self.jax_config.output_sizes,
-                                    self.jax_config.n_shards,
-                                    self.jax_config.fuse_matmuls,
-                                    dim=1,
-                                    jax_dtype=jnp.uint4)
-        delattr(layer, "qzeros")
-        layer.qzeros = qzeros
-
-        scales = torch_to_jax_param(layer.scales,
-                                    NamedSharding(
-                                        self.jax_config.mesh,
-                                        self.jax_config.scale_sharding),
-                                    self.jax_config.output_sizes,
-                                    self.jax_config.n_shards,
-                                    self.jax_config.fuse_matmuls,
-                                    dim=1)
+
+        weight_scale = t2j(layer.scales, use_dlpack=False)
         delattr(layer, "scales")
-        layer.scales = scales
+
+        assert layer.qzeros.packed_dim == layer.qzeros.ndim - 1
+        zero_point = t2j(layer.qzeros, use_dlpack=False)
+        delattr(layer, "qzeros")
 
         if layer.bias is not None and not layer.skip_bias_add:
             if layer.return_bias:
                 logger.warning_once("Bias might return incorrect value.")
-
-            bias = torch_to_jax_param(
-                layer.bias,
-                NamedSharding(self.jax_config.mesh,
-                              self.jax_config.bias_sharding),
-                self.jax_config.output_sizes,
-                self.jax_config.n_shards,
-                self.jax_config.fuse_matmuls,
-            )
+            bias = t2j(layer.bias, use_dlpack=False)
             delattr(layer, "bias")
-            layer.bias = bias
+        else:
+            bias = None
+
+        @jax.jit
+        def process_awq_linear_weights(
+            weight: jax.Array,
+            weight_scale: jax.Array,
+            zero_point: jax.Array,
+            bias: jax.Array | None,
+        ) -> LinearWeights:
+            weight = awq_u32_unpack_u4(weight)
+            group_size = self.quant_config.group_size
+            weight = weight.reshape((-1, group_size, weight.shape[-1]))
+
+            zero_point = awq_u32_unpack_u4(zero_point)
+
+            return process_lienar_weights(
+                LinearWeights(
+                    weight=weight,
+                    weight_scale=weight_scale,
+                    zero_point=zero_point,
+                    bias=bias,
+                ),
+                fused=self.linear_config.fuse_matmuls,
+                output_sizes=self.linear_config.output_sizes,
+                reorder_size=self.linear_config.n_shards,
+                transposed=False,
+            )
+
+        weights = process_awq_linear_weights(weight, weight_scale, zero_point,
+                                             bias)
+        weights = torch_view(
+            shard_linear_weights(
+                weights,
+                mesh=self.linear_config.mesh,
+                weight_p_spec=self.linear_config.weight_sharding,
+                bias_p_spec=self.linear_config.bias_sharding,
+                transposed=False,
+            ))
+
+        if self.linear_config.fuse_matmuls:
+            layer.qweight = Parameter(weights.weight, requires_grad=False)
+            layer.scales = Parameter(weights.weight_scale, requires_grad=False)
+            layer.qzeros = Parameter(weights.zero_point, requires_grad=False)
+            if bias is not None:
+                layer.bias = Parameter(weights.bias, requires_grad=False)
+        else:
+            layer.qweight = to_parameter_list(weights.weight)
+            layer.scales = to_parameter_list(weights.weight_scale)
+            layer.qzeros = to_parameter_list(weights.zero_point)
+            if bias is not None:
+                layer.bias = to_parameter_list(weights.bias)
 
     def apply(self,
               layer: torch.nn.Module,
@@ -133,7 +160,7 @@ class VllmAWQLinearMethod(AWQLinearMethod):
               bias: Optional[torch.Tensor] = None) -> torch.Tensor:
 
         with jax.named_scope(layer._get_name()):
-            if self.jax_config.fuse_matmuls:
+            if self.linear_config.fuse_matmuls:
                 out = self._apply_fused(layer, x, bias)
             else:
                 out = self._apply_split(layer, x, bias)
@@ -161,7 +188,7 @@ class VllmAWQLinearMethod(AWQLinearMethod):
             outs += bias.jax()
 
         outs = slice_sharded_tensor_for_concatenation(
-            outs, self.jax_config.output_sizes, self.jax_config.n_shards)
+            outs, self.linear_config.output_sizes, self.linear_config.n_shards)
         out = jnp.concatenate(outs, axis=-1)
         return torch_view(out)
 
@@ -192,16 +219,3 @@ class VllmAWQLinearMethod(AWQLinearMethod):
             outs.append(out)
         out = jnp.concatenate(outs, axis=-1)
         return torch_view(out)
-
-
-def unpack_awq_weight(weight: torch.Tensor, packed_dim: int):
-    weight = unpack_quantized_values_into_int32(weight, scalar_types.uint4,
-                                                packed_dim)
-
-    # AWQ packs 8 uint4 into 32-bits in this order: (0, 2, 4, 6, 1, 3, 5, 7).
-    # Following list maps the order used by AWQ into an ascending order.
-    reverse_awq_order = (0, 4, 1, 5, 2, 6, 3, 7)
-
-    orig_shape = weight.shape
-    weight = weight.reshape(orig_shape[:-1] + (-1, 8))
-    return weight[..., reverse_awq_order].reshape(orig_shape)

tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py

@@ -1,9 +1,22 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Optional
 
 import torch
 from jax.sharding import PartitionSpec
 from vllm.attention.layer import Attention
-from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.layer import FusedMoE
 from vllm.model_executor.layers.linear import LinearBase
 from vllm.model_executor.layers.quantization import \
@@ -18,22 +31,23 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
 
 from tpu_inference.layers.common.quant_methods import (COMPRESSED_TENSORS,
                                                        get_tpu_quant_method)
-from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
 from tpu_inference.layers.vllm.quantization.compressed_tensors.compressed_tensors_moe import \
-    VllmCompressedTensorsW8A8Fp8MoEMethod
+    VllmCompressedTensorsMoEMethod
 from tpu_inference.layers.vllm.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import \
     VllmCompressedTensorsW8A8Fp8
 from tpu_inference.layers.vllm.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_int8 import \
     VllmCompressedTensorsW8A8Int8
+from tpu_inference.layers.vllm.quantization.configs import VllmQuantConfig
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedConfig
+from tpu_inference.logger import init_logger
 
 P = PartitionSpec
 logger = init_logger(__name__)
 
 
 @register_quantization_config(get_tpu_quant_method(COMPRESSED_TENSORS))
-class VllmCompressedTensorsConfig(CompressedTensorsConfig, JaxCommonConfig):
+class VllmCompressedTensorsConfig(CompressedTensorsConfig, VllmQuantConfig):
 
     @classmethod
     def get_name(cls) -> str:
@@ -84,14 +98,14 @@ class VllmCompressedTensorsConfig(CompressedTensorsConfig, JaxCommonConfig):
             return VllmCompressedTensorsW8A8Fp8(
                 weight_quant=weight_quant,
                 is_static_input_scheme=is_static_input_scheme,
-                jax_config=linear_config,
+                linear_config=linear_config,
             )
         if self._is_dynamic_token_w8a8(weight_quant, input_quant):
             return VllmCompressedTensorsW8A8Int8(
                 strategy=weight_quant.strategy,
                 is_static_input_scheme=False,
                 input_symmetric=input_quant.symmetric,
-                jax_config=linear_config,
+                linear_config=linear_config,
             )
         raise NotImplementedError(
             "No compressed-tensors compatible scheme was found.")
@@ -113,8 +127,9 @@ class VllmCompressedTensorsConfig(CompressedTensorsConfig, JaxCommonConfig):
             layer.scheme = scheme
             return CompressedTensorsLinearMethod(self)
         if isinstance(layer, FusedMoE):
-            return VllmCompressedTensorsW8A8Fp8MoEMethod(
-                self, layer.quant_config, self.mesh)
+            layer.moe_config = self.get_moe_config(layer)
+            return VllmCompressedTensorsMoEMethod.get_moe_method(
+                self, layer, layer_name=prefix)
         if isinstance(layer, Attention):
             return CompressedTensorsKVCacheMethod(self)
         return None