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

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

@@ -1,21 +1,41 @@
+# 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 jax
 import jax.numpy as jnp
 import torch
 from compressed_tensors.quantization import QuantizationStrategy
-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.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_int8 import \
     CompressedTensorsW8A8Int8
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import \
     convert_to_channelwise
 
-from tpu_inference.layers.vllm.linear_common import (
-    sharded_quantized_matmul, slice_sharded_tensor_for_concatenation,
-    torch_to_jax_param)
-from tpu_inference.layers.vllm.quantization.common import JaxCommonLinearConfig
+from tpu_inference.layers.common.utils import \
+    slice_sharded_tensor_for_concatenation
+from tpu_inference.layers.vllm.linear import sharded_quantized_matmul
+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 \
+    VllmQuantLinearConfig
+from tpu_inference.logger import init_logger
 
 P = PartitionSpec
 logger = init_logger(__name__)
@@ -24,23 +44,15 @@ logger = init_logger(__name__)
 class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
 
     def __init__(self, strategy: str, is_static_input_scheme: bool,
-                 input_symmetric: bool, jax_config: JaxCommonLinearConfig):
+                 input_symmetric: bool, linear_config: VllmQuantLinearConfig):
         super().__init__(strategy, is_static_input_scheme, input_symmetric)
 
-        self.jax_config = jax_config
-        self.is_channelwise = (self.strategy == QuantizationStrategy.CHANNEL),
+        self.linear_config = linear_config
+        self.is_channelwise = (self.strategy == QuantizationStrategy.CHANNEL)
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        weight = torch_to_jax_param(
-            layer.weight,
-            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,
-        )
+        weight = t2j(layer.weight, use_dlpack=False)
         delattr(layer, "weight")
-        layer.weight = weight
 
         weight_scale = layer.weight_scale
         is_fused_module = len(layer.logical_widths) > 1
@@ -48,31 +60,55 @@ class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
             weight_scale = convert_to_channelwise(weight_scale,
                                                   layer.logical_widths)
         weight_scale = weight_scale.squeeze(-1)
-
-        weight_scale = torch_to_jax_param(
-            weight_scale,
-            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,
-        )
+        weight_scale = t2j(weight_scale, use_dlpack=False)
         delattr(layer, "weight_scale")
-        layer.weight_scale = weight_scale
 
         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_int8_linear_weights(
+            weight: jax.Array,
+            weight_scale: jax.Array,
+            bias: jax.Array | None,
+        ) -> LinearWeights:
+            return process_lienar_weights(
+                LinearWeights(
+                    weight=weight,
+                    weight_scale=weight_scale,
+                    zero_point=None,
+                    bias=bias,
+                ),
+                fused=self.linear_config.fuse_matmuls,
+                output_sizes=self.linear_config.output_sizes,
+                reorder_size=self.linear_config.n_shards,
+            )
+
+        weights = process_int8_linear_weights(weight, weight_scale, 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,
+            ))
+
+        if self.linear_config.fuse_matmuls:
+            layer.weight = Parameter(weights.weight, requires_grad=False)
+            layer.weight_scale = Parameter(weights.weight_scale,
+                                           requires_grad=False)
+            if bias is not None:
+                layer.bias = Parameter(weights.bias, requires_grad=False)
+        else:
+            layer.weight = to_parameter_list(weights.weight)
+            layer.weight_scale = to_parameter_list(weights.weight_scale)
+            if bias is not None:
+                layer.bias = to_parameter_list(weights.bias)
 
         # TODO(kyuyeunk): Support static range input quantization.
         assert getattr(layer, "input_scale", None) is None
@@ -82,7 +118,7 @@ class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
     def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,
                       bias: Optional[torch.Tensor]) -> 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)
@@ -99,14 +135,14 @@ class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
             x_jax,
             weight_jax,
             weight_scale_jax,
-            self.jax_config.mesh,
-            self.jax_config.weight_sharding,
+            self.linear_config.mesh,
+            self.linear_config.weight_sharding,
         )
         if bias is not None and not layer.skip_bias_add:
             outs += jax_view(bias)
 
         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)
 
@@ -125,8 +161,8 @@ class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
                 x_jax,
                 weight_jax,
                 weight_scale_jax,
-                self.jax_config.mesh,
-                self.jax_config.weight_sharding,
+                self.linear_config.mesh,
+                self.linear_config.weight_sharding,
             )
             if bias is not None and not layer.skip_bias_add:
                 out += jax_view(bias[i])

tpu_inference/layers/vllm/quantization/{common.py → configs.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 torchax
 from jax.sharding import Mesh, PartitionSpec
 from vllm.config import VllmConfig
@@ -11,9 +25,10 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                ReplicatedLinear,
                                                RowParallelLinear)
 
-from tpu_inference.layers.vllm.linear_common import \
+from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.vllm.process_weights.linear_weights import \
     get_model_matmul_fusion_assignment
-from tpu_inference.utils import TPU_SECOND_LAST_MINOR
+from tpu_inference.utils import TPU_SECOND_LAST_MINOR, get_mesh_shape_product
 
 # yapf: enable
 
@@ -22,7 +37,7 @@ P = PartitionSpec
 logger = init_logger(__name__)
 
 
-class JaxCommonLinearConfig:
+class VllmQuantLinearConfig:
 
     def __init__(self, vllm_config: VllmConfig, mesh: Mesh, layer: LinearBase):
         assert isinstance(layer, LinearBase)
@@ -35,14 +50,18 @@ class JaxCommonLinearConfig:
         self.input_sharding = None
         self.output_sharding = None
 
+        self.tp_size = get_mesh_shape_product(self.mesh,
+                                              ShardingAxisName.MLP_TENSOR)
+
         if isinstance(layer, RowParallelLinear):
-            self.weight_sharding = P(None, "model")
+            self.weight_sharding = P(None, ShardingAxisName.ATTN_HEAD)
             if self.enable_sp:
-                self.output_sharding = P("model", None)
+                self.output_sharding = P(ShardingAxisName.MLP_TENSOR, None)
         elif isinstance(layer, ColumnParallelLinear):
-            self.weight_sharding = P("model", None)
+            self.weight_sharding = P(ShardingAxisName.ATTN_HEAD, None)
+
             if self.enable_sp:
-                self.input_sharding = P("model", None)
+                self.input_sharding = P(ShardingAxisName.MLP_TENSOR, None)
 
             if isinstance(layer, MergedColumnParallelLinear) or isinstance(
                     layer, QKVParallelLinear):
@@ -61,35 +80,28 @@ class JaxCommonLinearConfig:
                 " bad performance.", type(layer))
 
         self.bias_sharding = P(self.weight_sharding[0])
-        if isinstance(self.weight_sharding[0], tuple):
-            self.n_shards = 1
-            for axis in self.weight_sharding[0]:
-                self.n_shards *= self.mesh.shape.get(axis, 1)
-        else:
-            self.n_shards = self.mesh.shape.get(self.weight_sharding[0], 1)
+        self.n_shards = get_mesh_shape_product(self.mesh,
+                                               self.weight_sharding[0])
 
     def get_input_sharding(self, x: torchax.tensor.Tensor):
-        if self.enable_sp:
-            token_num = x.shape[0]
-            # NOTE(chengjiyao): make sure the sharded token_num is larger than TPU_SECOND_LAST_MINOR
-            if token_num // self.mesh.shape["model"] >= TPU_SECOND_LAST_MINOR:
-                return self.input_sharding
-            else:
-                return None
+        if not self.enable_sp:
+            return None
+        token_num = x.shape[0]
+        # NOTE(chengjiyao): make sure the sharded token_num is larger than TPU_SECOND_LAST_MINOR
+        if token_num // self.tp_size < TPU_SECOND_LAST_MINOR:
+            return None
         return self.input_sharding
 
     def get_output_sharding(self, x: torchax.tensor.Tensor):
         if self.enable_sp:
             token_num = x.shape[0]
             # NOTE(chengjiyao): make sure the sharded token_num is larger than TPU_SECOND_LAST_MINOR
-            if token_num // self.mesh.shape["model"] >= TPU_SECOND_LAST_MINOR:
-                return self.output_sharding
-            else:
+            if token_num // self.tp_size < TPU_SECOND_LAST_MINOR:
                 return None
         return self.output_sharding
 
 
-class JaxCommonConfig:
+class VllmQuantConfig:
     vllm_config: VllmConfig
     mesh: Mesh
 
@@ -98,9 +110,9 @@ class JaxCommonConfig:
         cls.vllm_config = vllm_config
         cls.mesh = mesh
 
-    def get_linear_config(self, layer: LinearBase) -> JaxCommonLinearConfig:
+    def get_linear_config(self, layer: LinearBase) -> VllmQuantLinearConfig:
         assert isinstance(layer, LinearBase)
-        return JaxCommonLinearConfig(self.vllm_config, self.mesh, layer)
+        return VllmQuantLinearConfig(self.vllm_config, self.mesh, layer)
 
     def get_moe_config(self, layer: FusedMoE) -> FusedMoEConfig:
         assert isinstance(layer, FusedMoE)

tpu_inference/layers/vllm/quantization/fp8.py

@@ -0,0 +1,119 @@
+# 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 torch
+from jax.sharding import PartitionSpec
+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 \
+    register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizeMethodBase
+from vllm.model_executor.layers.quantization.fp8 import (Fp8Config,
+                                                         Fp8LinearMethod)
+from vllm.model_executor.layers.quantization.utils.quant_utils import \
+    is_layer_skipped
+
+from tpu_inference.layers.common.quant_methods import FP8, get_tpu_quant_method
+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(FP8))
+class VllmFp8Config(Fp8Config, VllmQuantConfig):
+
+    @classmethod
+    def get_name(cls):
+        return FP8
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.bfloat16]
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[Union["LinearMethodBase", "QuantizeMethodBase"]]:
+        if isinstance(layer, LinearBase):
+            linear_config = self.get_linear_config(layer)
+            if is_layer_skipped(prefix, self.ignored_layers):
+                return VllmUnquantizedLinearMethod(linear_config)
+            return VllmFp8LinearMethod(self, linear_config)
+        elif isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "FP8 FusedMoE is currently not supported in torchax-jax")
+        return None
+
+
+class VllmFp8LinearMethod(Fp8LinearMethod):
+
+    def __init__(self, quant_config: VllmFp8Config,
+                 jax_config: VllmQuantLinearConfig):
+        super().__init__(quant_config)
+        self.jax_config = jax_config
+        self._configure_sharding()
+
+    def _configure_sharding(self) -> None:
+
+        raise NotImplementedError(
+            "Configure PartitionSpec for weight_sharding and scale_sharding "
+            "based on layer type (RowParallel/ColumnParallel)")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        raise NotImplementedError(
+            "Convert layer.weight, layer.weight_scale, and optionally "
+            "layer.input_scale and layer.bias from torch tensors to JAX arrays "
+            "using torch_to_jax_param() with appropriate sharding")
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        with jax.named_scope(layer._get_name()):
+            if self.jax_config.fuse_matmuls:
+                out = self._apply_fused(layer, x, bias)
+            else:
+                out = self._apply_split(layer, x, bias)
+
+        return out
+
+    def _apply_fused(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        raise NotImplementedError(
+            "Implement single matmul for fused outputs: "
+            "quantize input to fp8, perform fp8 matmul with weight and scales, "
+            "dequantize output, and add bias if present")
+
+    def _apply_split(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        raise NotImplementedError(
+            "Implement separate matmuls per output partition: "
+            "split weight/scale by output_sizes, perform fp8 matmul for each, "
+            "concatenate results, and add bias if present")