tpu-inference 0.12.0.dev20251213__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tpu_inference/layers/vllm/process_weights/linear_weights.py

@@ -0,0 +1,174 @@
+# 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 dataclasses import dataclass, fields
+
+import jax
+import torch
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torch.nn import ParameterList
+from torch.nn.parameter import Parameter
+from torchax.tensor import Tensor
+
+from tpu_inference.layers.common.utils import \
+    reorder_concatenated_tensor_for_sharding
+from tpu_inference.logger import init_logger
+
+P = PartitionSpec
+
+logger = init_logger(__name__)
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class LinearWeights:
+    weight: jax.Array | Tensor | list[jax.Array | Tensor]
+    weight_scale: jax.Array | Tensor | list[jax.Array | Tensor] | None
+    zero_point: jax.Array | Tensor | list[jax.Array | Tensor] | None
+    bias: jax.Array | Tensor | list[jax.Array | Tensor] | None
+
+
+MODEL_MATMUL_FUSION_TRUTH_TABLE = {
+    ("Qwen/Qwen2.5-7B-Instruct", 1024, 1, "QKVParallelLinear"):
+    True,
+    ("Qwen/Qwen2.5-7B-Instruct", 1024, 1, "MergedColumnParallelLinear"):
+    False,
+    ("Qwen/Qwen2.5-7B-Instruct", 2048, 1, "QKVParallelLinear"):
+    False,
+    ("Qwen/Qwen2.5-7B-Instruct", 2048, 1, "MergedColumnParallelLinear"):
+    False,
+    ("meta-llama/Llama-3.1-8B-Instruct", 1024, 1, "QKVParallelLinear"):
+    False,
+    ("meta-llama/Llama-3.1-8B-Instruct", 1024, 1, "MergedColumnParallelLinear"):
+    False,
+    ("meta-llama/Llama-3.1-8B-Instruct", 2048, 1, "QKVParallelLinear"):
+    False,
+    ("meta-llama/Llama-3.1-8B-Instruct", 2048, 1, "MergedColumnParallelLinear"):
+    False,
+    ("RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8", 1024, 1, "QKVParallelLinear"):
+    False,
+    ("RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8", 1024, 1, "MergedColumnParallelLinear"):
+    False,
+    ("RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8", 2048, 1, "QKVParallelLinear"):
+    False,
+    ("RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8", 2048, 1, "MergedColumnParallelLinear"):
+    False,
+}
+
+
+def to_parameter_list(tensor: list[torch.Tensor]):
+    tensor = [Parameter(t, requires_grad=False) for t in tensor]
+    return ParameterList(tensor)
+
+
+def get_model_matmul_fusion_assignment(model_name: str, batch_size: int,
+                                       tp_size: int, layer_name: str):
+    key = (model_name, batch_size, tp_size, layer_name)
+    return MODEL_MATMUL_FUSION_TRUTH_TABLE.get(key, True)
+
+
+def process_lienar_weights(
+    weights: LinearWeights,
+    fused: bool = False,
+    output_sizes: list[int] | None = None,
+    reorder_size: int | None = None,
+    transposed: bool = True,
+    per_tensor: bool = False,
+) -> LinearWeights:
+    weight = weights.weight
+    weight_scale = weights.weight_scale
+    zero_point = weights.zero_point
+    bias = weights.bias
+
+    dim = 0 if transposed else -1
+    if output_sizes is None:
+        output_sizes = [weight.shape[dim]]
+
+    if fused:
+        assert reorder_size is not None
+        weight = reorder_concatenated_tensor_for_sharding(
+            weight, output_sizes, reorder_size, dim)
+
+        if weight_scale is not None and not per_tensor:
+            weight_scale = reorder_concatenated_tensor_for_sharding(
+                weight_scale, output_sizes, reorder_size, dim)
+        if zero_point is not None:
+            zero_point = reorder_concatenated_tensor_for_sharding(
+                zero_point, output_sizes, reorder_size, dim)
+        if bias is not None:
+            bias = reorder_concatenated_tensor_for_sharding(
+                bias, output_sizes, reorder_size, dim)
+    else:
+
+        def slice_tensor(tensor):
+            tensors = []
+            start = 0
+            for size in output_sizes:
+                end = start + size
+                tensor_split = jax.lax.slice_in_dim(tensor,
+                                                    start,
+                                                    end,
+                                                    axis=dim)
+                tensors.append(tensor_split)
+                start = end
+            return tensors
+
+        weight = slice_tensor(weight)
+        if weight_scale is not None and not per_tensor:
+            weight_scale = slice_tensor(weight_scale)
+        if zero_point is not None:
+            zero_point = slice_tensor(zero_point)
+        if bias is not None:
+            bias = slice_tensor(bias)
+
+    return LinearWeights(
+        weight=weight,
+        weight_scale=weight_scale,
+        zero_point=zero_point,
+        bias=bias,
+    )
+
+
+def shard_linear_weights(
+    weights: LinearWeights,
+    mesh: Mesh,
+    weight_p_spec: PartitionSpec,
+    bias_p_spec: PartitionSpec,
+    transposed: bool = True,
+    per_tensor: bool = False,
+) -> LinearWeights:
+
+    if not transposed:
+        # By defualt, we use transposed weights. If it is not transposed,
+        # we need to transpose the sharding as well.
+        weight_p_spec = PartitionSpec(*weight_p_spec[::-1])
+        bias_p_spec = PartitionSpec(weight_p_spec[0])
+
+    weight_sharding = NamedSharding(mesh, weight_p_spec)
+    bias_sharding = NamedSharding(mesh, bias_p_spec)
+
+    weight_shardings = LinearWeights(
+        weight=weight_sharding,
+        weight_scale=NamedSharding(mesh, P()) if per_tensor else bias_sharding,
+        zero_point=bias_sharding,
+        bias=bias_sharding,
+    )
+
+    for field in fields(LinearWeights):
+        key = field.name
+        if (weight := getattr(weights, key, None)) is not None:
+            sharding = getattr(weight_shardings, key)
+            weight = jax.device_put(weight, sharding)
+            setattr(weights, key, weight)
+    return weights

tpu_inference/layers/vllm/quantization/__init__.py

@@ -1,3 +1,17 @@
+# 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.
+
 import copy
 
 from jax.sharding import Mesh
@@ -7,9 +21,10 @@ from vllm.model_executor.layers.quantization.base_config import \
 
 from tpu_inference.layers.common import quant_methods
 from tpu_inference.layers.vllm.quantization.awq import VllmAWQConfig
-from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
 from tpu_inference.layers.vllm.quantization.compressed_tensors.compressed_tensors import \
-    VllmCompressedTensorsConfig  # noqa: E501
+    VllmCompressedTensorsConfig
+from tpu_inference.layers.vllm.quantization.configs import VllmQuantConfig
+from tpu_inference.layers.vllm.quantization.fp8 import VllmFp8Config
 from tpu_inference.layers.vllm.quantization.mxfp4 import VllmMxfp4Config
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedConfig
@@ -23,6 +38,7 @@ def get_tpu_quantization_config(vllm_config: VllmConfig,
         None: VllmUnquantizedConfig,
         quant_methods.COMPRESSED_TENSORS: VllmCompressedTensorsConfig,
         quant_methods.AWQ: VllmAWQConfig,
+        quant_methods.FP8: VllmFp8Config,
         quant_methods.MXFP4: VllmMxfp4Config,
     }
     if model_config.quantization not in method_to_config:
@@ -30,7 +46,7 @@ def get_tpu_quantization_config(vllm_config: VllmConfig,
             f"{model_config.quantization} quantization method not supported."
             f" Supported methods are {method_to_config.keys()}")
     quant_config = method_to_config[model_config.quantization]
-    assert issubclass(quant_config, JaxCommonConfig)
+    assert issubclass(quant_config, VllmQuantConfig)
     quant_config.set_configs(vllm_config, mesh)
 
     model_config.quantization = quant_methods.get_tpu_quant_method(

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 \
     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.")