tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tpu_inference/layers/vllm/linear_common.py

@@ -0,0 +1,221 @@
+# 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 Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
+from torchax.interop import torch_view
+from torchax.ops.mappings import t2j
+
+from tpu_inference import envs
+from tpu_inference.kernels.quantized_matmul.kernel import (
+    quantized_matmul_kernel, xla_quantized_matmul)
+
+
+def sharded_quantized_matmul(x: jax.Array, w_q: jax.Array, w_s: jax.Array,
+                             mesh: Mesh, weight_sharding: P) -> jax.Array:
+    """
+    Wrapper around the quantized matmul kernel.
+
+    Args:
+        x:  Activation.
+        w_q: Weight quantized array. [n_output_features, n_input_features]
+        w_s: Weight quantization scale. [n_output_features]
+        mesh: Mesh to shard on.
+        weight_sharding: PartitionSpec for the weight tensor.
+
+    Returns:
+        Output of the quantized matmul.
+    """
+
+    # NOTE (jacobplatin/kyuyeunk) there have been numeric issues (concerning) NaNs
+    # with the kernel and thus we disable it for now.
+    if envs.ENABLE_QUANTIZED_MATMUL_KERNEL:
+        out_axis, in_axis = weight_sharding
+        x_sharding = P(None, in_axis)
+        scale_sharding = P(out_axis, )
+        out_sharding = P(None, out_axis)
+
+        x = jax.lax.with_sharding_constraint(x,
+                                             NamedSharding(mesh, x_sharding))
+
+        def wrapper(x, w_q, w_s):
+            output = quantized_matmul_kernel(x, w_q, w_s, x_q_dtype=w_q.dtype)
+            if in_axis:
+                output = jax.lax.psum(output, axis_name=in_axis)
+            return output
+
+        return jax.shard_map(wrapper,
+                             mesh=mesh,
+                             in_specs=(x_sharding, weight_sharding,
+                                       scale_sharding),
+                             out_specs=(out_sharding),
+                             check_vma=False)(x, w_q, w_s)
+    else:
+        return xla_quantized_matmul(x, w_q, w_s)
+
+
+def reorder_concatenated_tensor_for_sharding(concatenated_tensor: jax.Array,
+                                             split_sizes: list[int],
+                                             n_shards: int, dim: int):
+    """
+    Reorder a replicated concatenated tensor such that when sharded on multiple chips, each shard is a concatenation of the shards of the individual tensors.
+    For example, let the concatenated_tensor be:
+        AAAAAAAAAAAABBBBBBBBCCCC
+            12 As     8 Bs  4 Cs
+    and let the split_sizes = [12, 8, 4] and n_shards = 4.
+    The output is:
+        AAABBCAAABBCAAABBCAAABBC
+    In other words, it reorders the input tensor into 4 segements, with each segment corresponding to a shard and being AAABBC.
+    Args:
+        concatenated_tensor: the tensor, concatenated on the dimension specified by `dim`.
+        split_sizes: each individual tensor's size on the dimension specified by `dim`.
+        n_shards: num of shards.
+        dim: the dimension on which the concatenated_tensor is concatenated.
+    """
+    # Split the concatenated tensor into individual tensors.
+    split_tensors = []
+    start_offset = 0
+    old_shape = concatenated_tensor.shape
+    # New shape ensures each split_tensor[i] maps to a tensor in ith shards
+    new_shape = old_shape[:dim] + (n_shards, -1) + old_shape[dim + 1:]
+    for split_size in split_sizes:
+        split_tensor = jax.lax.slice_in_dim(concatenated_tensor,
+                                            start_offset,
+                                            start_offset + split_size,
+                                            axis=dim)
+        split_tensors.append(split_tensor.reshape(new_shape))
+        start_offset += split_size
+    # While maintaining 0th dim as a shard dim, we concatenate along 1th dim to
+    # to create concatenated tnensor where 0th dim maps to shard dim.
+    reordered_tensor = jnp.concatenate(split_tensors, axis=dim + 1)
+    return reordered_tensor.reshape(old_shape)
+
+
+def slice_sharded_tensor_for_concatenation(sharded_tensor: jax.Array,
+                                           split_sizes: list[int],
+                                           n_shards: int):
+    """
+    Slice the input tensor which is sharded on multiple chips (on the last dim) into individual tensors with the same sharding.
+    For example, let the sharded_tensor be:
+        AAABBC | AAABBC | AAABBC | AAABBC
+        Shard0   Shard1   Shard2   Shard3
+    and let the split_sizes = [12, 8, 4] and n_shards = 4.
+    The output is a list of 3 tensors:
+         AAA   |  AAA   |  AAA   |  AAA
+          BB   |   BB   |   BB   |   BB
+           C   |    C   |    C   |    C
+        Shard0   Shard1   Shard2   Shard3
+    In other words, each individual tensor is a slice of the input tensor with the same sharding.
+    Args:
+        sharded_tensor: the input tensor, sharded on the last dim.
+        split_sizes: each individual tensor's size on the last dim.
+        n_shards: num of shards.
+    """
+    new_shape = sharded_tensor.shape[:-1] + (n_shards, -1)
+    # New shape ensures each sharded_tensor[:, i] maps to a tensor in ith shards
+    sharded_tensor = sharded_tensor.reshape(new_shape)
+
+    split_tensors = []
+    start_offset = 0
+    for split_size in split_sizes:
+        assert split_size % n_shards == 0
+        sz = split_size // n_shards  # size of this split tensor per shard
+        end_offset = start_offset + sz
+        # Because we are slicing over last dim, sharding dim remains intact.
+        # Therefore, splitting happens locally.
+        split_tensor = sharded_tensor[..., start_offset:end_offset]
+        split_tensors.append(split_tensor.reshape(new_shape[:-2] + (-1, )))
+        start_offset = end_offset
+
+    return split_tensors
+
+
+def torch_to_jax_param(
+    tensor: torch.Tensor,
+    sharding: NamedSharding,
+    output_sizes: Optional[int],
+    n_shards: int,
+    fused: bool,
+    dim: int = 0,
+    jax_dtype: Optional[jnp.dtype] = None,
+) -> Union[torch.nn.Parameter, torch.nn.ParameterList]:
+    if output_sizes is None:
+        output_sizes = [tensor.shape[0]]
+
+    tensor = t2j(tensor, use_dlpack=False)
+    if jax_dtype:
+        tensor = tensor.astype(jax_dtype)
+
+    if fused:
+        tensor = reorder_concatenated_tensor_for_sharding(
+            tensor, output_sizes, n_shards, dim)
+        tensor = jax.device_put(tensor, sharding)
+        param = torch.nn.Parameter(torch_view(tensor), requires_grad=False)
+    else:
+        tensors = []
+        start_offset = 0
+        for size in output_sizes:
+            end_offset = start_offset + size
+
+            tensor_split = jax.lax.slice_in_dim(tensor,
+                                                start_offset,
+                                                end_offset,
+                                                axis=dim)
+            tensor_split = jax.device_put(tensor_split, sharding)
+            tensor_split = torch.nn.Parameter(torch_view(tensor_split),
+                                              requires_grad=False)
+            tensors.append(tensor_split)
+
+            start_offset = end_offset
+        param = torch.nn.ParameterList(tensors)
+    return param
+
+
+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 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)

tpu_inference/layers/vllm/quantization/__init__.py

@@ -0,0 +1,55 @@
+# 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
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizationConfig
+
+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
+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
+
+
+def get_tpu_quantization_config(vllm_config: VllmConfig,
+                                mesh: Mesh) -> QuantizationConfig:
+    model_config = copy.deepcopy(vllm_config.model_config)
+    # TODO(kyuyeunk): Add support for "tpu_int8".
+    method_to_config: dict[str, str] = {
+        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:
+        raise NotImplementedError(
+            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)
+    quant_config.set_configs(vllm_config, mesh)
+
+    model_config.quantization = quant_methods.get_tpu_quant_method(
+        quant_config.get_name())
+    return VllmConfig.get_quantization_config(model_config,
+                                              vllm_config.load_config)

tpu_inference/layers/vllm/quantization/awq.py

@@ -0,0 +1,221 @@
+# 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 torchax.interop import jax_view, torch_view
+from vllm.logger import init_logger
+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.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 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.vllm.quantization.unquantized import \
+    VllmUnquantizedLinearMethod
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+@register_quantization_config(get_tpu_quant_method(AWQ))
+class VllmAWQConfig(AWQConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls):
+        return AWQ
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        # NOTE: AWQ checkpoint was quantized with float16. But on TPUs, using
+        # bfloat16 is significantly preferred over float16. This might lead to
+        # some numeric output change.
+        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.modules_to_not_convert):
+                return VllmUnquantizedLinearMethod(linear_config)
+            return VllmAWQLinearMethod(self, linear_config)
+        elif isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "AWQ FusedMoE is currently not supported in torchax-jax")
+        return None
+
+
+class VllmAWQLinearMethod(AWQLinearMethod):
+
+    def __init__(self, quant_config: VllmAWQConfig,
+                 jax_config: JaxCommonLinearConfig):
+        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)
+
+    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)
+        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)
+        delattr(layer, "scales")
+        layer.scales = scales
+
+        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,
+            )
+            delattr(layer, "bias")
+            layer.bias = bias
+
+    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:
+        x_jax = jax_view(x)
+
+        qweight = jax_view(layer.qweight)
+        qzeros = jnp.expand_dims(jax_view(layer.qzeros), 1)
+        scales = jnp.expand_dims(jax_view(layer.scales), 1)
+
+        qweight = qweight.astype(jnp.int8)
+        qzeros = qzeros.astype(jnp.int8)
+
+        weight = (qweight - qzeros) * scales
+        weight = weight.reshape((-1, weight.shape[-1]))
+        outs = jnp.einsum("bd,df->bf", x_jax, weight)
+
+        if bias is not None and not layer.skip_bias_add:
+            outs += bias.jax()
+
+        outs = slice_sharded_tensor_for_concatenation(
+            outs, self.jax_config.output_sizes, self.jax_config.n_shards)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)
+
+    def _apply_split(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        assert isinstance(layer.qweight, torch.nn.ParameterList)
+
+        x_jax = jax_view(x)
+        params = zip(layer.qweight, layer.qzeros, layer.scales)
+        outs = []
+        for i, (qweight, qzeros, scales) in enumerate(params):
+            qweight = jax_view(qweight)
+            scales = jnp.expand_dims(jax_view(scales), 1)
+            qzeros = jnp.expand_dims(jax_view(qzeros), 1)
+
+            qweight = qweight.astype(jnp.int8)
+            qzeros = qzeros.astype(jnp.int8)
+
+            weight = (qweight - qzeros) * scales
+            weight = weight.reshape((-1, weight.shape[-1]))
+            out = jnp.einsum("bd,df->bf", x_jax, weight)
+
+            if bias is not None and not layer.skip_bias_add:
+                out += jax_view(bias[i])
+
+            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/common.py

@@ -0,0 +1,124 @@
+# 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
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEConfig
+# yapf: disable
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+                                               LinearBase,
+                                               MergedColumnParallelLinear,
+                                               QKVParallelLinear,
+                                               ReplicatedLinear,
+                                               RowParallelLinear)
+
+from tpu_inference.layers.vllm.linear_common import \
+    get_model_matmul_fusion_assignment
+from tpu_inference.utils import TPU_SECOND_LAST_MINOR
+
+# yapf: enable
+
+P = PartitionSpec
+
+logger = init_logger(__name__)
+
+
+class JaxCommonLinearConfig:
+
+    def __init__(self, vllm_config: VllmConfig, mesh: Mesh, layer: LinearBase):
+        assert isinstance(layer, LinearBase)
+
+        self.mesh = mesh
+        self.output_sizes = [layer.output_size]
+        self.weight_sharding = P(None, None)
+        self.fuse_matmuls = True
+        self.enable_sp = vllm_config.compilation_config.pass_config.enable_sp
+        self.input_sharding = None
+        self.output_sharding = None
+
+        if isinstance(layer, RowParallelLinear):
+            self.weight_sharding = P(None, "model")
+            if self.enable_sp:
+                self.output_sharding = P("model", None)
+        elif isinstance(layer, ColumnParallelLinear):
+            self.weight_sharding = P("model", None)
+            if self.enable_sp:
+                self.input_sharding = P("model", None)
+
+            if isinstance(layer, MergedColumnParallelLinear) or isinstance(
+                    layer, QKVParallelLinear):
+                self.output_sizes = layer.output_sizes
+
+            self.fuse_matmuls = get_model_matmul_fusion_assignment(
+                vllm_config.model_config.model,
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                vllm_config.parallel_config.tensor_parallel_size,
+                layer._get_name())
+        elif isinstance(layer, ReplicatedLinear):
+            self.weight_sharding = P(None, None)
+        else:
+            logger.warning(
+                "Unsupported linear layer type of %s. Can potentially yield "
+                " 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)
+
+    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
+        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:
+                return None
+        return self.output_sharding
+
+
+class JaxCommonConfig:
+    vllm_config: VllmConfig
+    mesh: Mesh
+
+    @classmethod
+    def set_configs(cls, vllm_config: VllmConfig, mesh: Mesh):
+        cls.vllm_config = vllm_config
+        cls.mesh = mesh
+
+    def get_linear_config(self, layer: LinearBase) -> JaxCommonLinearConfig:
+        assert isinstance(layer, LinearBase)
+        return JaxCommonLinearConfig(self.vllm_config, self.mesh, layer)
+
+    def get_moe_config(self, layer: FusedMoE) -> FusedMoEConfig:
+        assert isinstance(layer, FusedMoE)
+        moe_config = layer.moe_config
+        use_ep = self.vllm_config.parallel_config.enable_expert_parallel
+        moe_config.moe_parallel_config.use_ep = use_ep
+        return moe_config

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.