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

tpu_inference/layers/vllm/quantization/unquantized.py

@@ -0,0 +1,416 @@
+# 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 Any, Optional, Union
+
+import jax
+import jax.numpy as jnp
+import torch
+from jax.experimental.layout import Format, Layout
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torch.nn.parameter import Parameter
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import t2j
+from vllm.attention.layer import Attention
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import (
+    FusedMoE, FusedMoEConfig, UnquantizedFusedMoEMethod)
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize)
+from vllm.model_executor.layers.linear import (LinearBase,
+                                               UnquantizedLinearMethod)
+from vllm.model_executor.layers.quantization import \
+    register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase)
+
+from tpu_inference import envs
+from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe
+from tpu_inference.layers.common.quant_methods import (UNQUANTIZED,
+                                                       get_tpu_quant_method)
+from tpu_inference.layers.vllm.fused_moe import fused_moe_func
+from tpu_inference.layers.vllm.linear_common import (
+    reorder_concatenated_tensor_for_sharding,
+    slice_sharded_tensor_for_concatenation, torch_to_jax_param)
+from tpu_inference.layers.vllm.quantization.common import (
+    JaxCommonConfig, JaxCommonLinearConfig)
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+def align_to(a, b):
+    return (a + b - 1) // b * b
+
+
+@register_quantization_config(get_tpu_quant_method(UNQUANTIZED))
+class VllmUnquantizedConfig(QuantizationConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls) -> str:
+        return UNQUANTIZED
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.float32, torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0  # Always supported
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []  # No extra configs required.
+
+    @classmethod
+    def from_config(cls, _: dict[str, Any]) -> "VllmUnquantizedConfig":
+        return cls()
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional[QuantizeMethodBase]:
+        if isinstance(layer, LinearBase):
+            linear_config = self.get_linear_config(layer)
+            return VllmUnquantizedLinearMethod(linear_config)
+        if isinstance(layer, FusedMoE):
+            moe_config = self.get_moe_config(layer)
+            return VllmUnquantizedFusedMoEMethod(moe_config, self.mesh)
+        if isinstance(layer, Attention):
+            return None
+        return None
+
+
+class VllmUnquantizedLinearMethod(UnquantizedLinearMethod):
+
+    def __init__(self, jax_config: JaxCommonLinearConfig):
+        self.jax_config = jax_config
+
+    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,
+        )
+        delattr(layer, "weight")
+        layer.weight = weight
+
+        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:
+        assert isinstance(layer, LinearBase)
+
+        with jax.named_scope(layer._get_name()):
+            if in_sharding := self.jax_config.get_input_sharding(x):
+                x.shard_(NamedSharding(self.jax_config.mesh, in_sharding))
+
+            if self.jax_config.fuse_matmuls:
+                out = self._apply_fused(layer, x, bias)
+            else:
+                out = self._apply_split(layer, x, bias)
+
+            if out_sharding := self.jax_config.get_output_sharding(out):
+                out.shard_(NamedSharding(self.jax_config.mesh, out_sharding))
+
+        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)
+        weight_jax = jax_view(layer.weight)
+
+        outs = jnp.einsum("mn,pn->mp", x_jax, weight_jax)
+        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.weight, torch.nn.ParameterList)
+
+        x_jax = x.jax()
+        outs = []
+        for i, weight in enumerate(layer.weight):
+            weight_jax = jax_view(weight)
+
+            out = jnp.einsum("mn,pn->mp", x_jax, weight_jax)
+            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)
+
+
+class VllmUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod):
+
+    def __init__(self,
+                 moe: FusedMoEConfig,
+                 mesh: Mesh,
+                 ep_axis_name: str = 'model'):
+        super().__init__(moe)
+        self.mesh = mesh
+        self.use_kernel = envs.USE_MOE_EP_KERNEL and moe.use_ep
+        self.ep_axis_name = ep_axis_name
+        # TODO: Use autotune table once we have it.
+        self.block_size = {
+            "bt": 64,
+            "bf": 1024,
+            "bd1": 1536,
+            "bd2": 1536,
+            "btc": 64,
+            "bfc": 1024,
+            "bd1c": 1536,
+            "bd2c": 1536,
+        }
+
+    def select_gemm_impl(
+        self,
+        prepare_finalize: FusedMoEPrepareAndFinalize,
+        moe: FusedMoEConfig,
+        layer: torch.nn.Module,
+    ) -> FusedMoEPermuteExpertsUnpermute:
+        raise NotImplementedError(
+            "Selecting gemm implementation is currently not supported.")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        assert isinstance(layer, FusedMoE)
+        w13_weight = t2j(layer.w13_weight, use_dlpack=False)
+        w2_weight = t2j(layer.w2_weight, use_dlpack=False)
+
+        num_experts, hidden_size, intermediate_size = w2_weight.shape
+
+        if self.moe.has_bias:
+            w13_bias = t2j(layer.w13_bias, use_dlpack=False)
+            w2_bias = t2j(layer.w2_bias, use_dlpack=False)
+
+        if layer.activation == "swigluoai":
+            # When using swigluoai, vLLM splits gmm output in a interleaved way.
+            # However, interleaved split is not performant on TPU. Therefore,
+            # we preprocess the weight so that splitting gmm output by middle
+            # can still get the same result.
+            w1_weight = w13_weight[:, ::2, :]
+            w3_weight = w13_weight[:, 1::2, :]
+            w13_weight = jnp.concat([w1_weight, w3_weight], axis=1)
+
+            if self.moe.has_bias:
+                w1_bias = w13_bias[:, ::2]
+                w3_bias = w13_bias[:, 1::2]
+                w13_bias = jnp.concat([w1_bias, w3_bias], axis=1)
+
+        if self.use_kernel:
+            # Kernel expects:
+            # w13: (num_experts, 2, hidden_size, intermediate_size)
+            # w2: (num_experts, intermediate_size, hidden_size)
+            # Current format:
+            # w13_weight: (num_experts, 2*intermediate_size, hidden_size)
+            # w2_weight: (num_experts, hidden_size, intermediate_size)
+            num_experts = w13_weight.shape[0]
+            intermediate_size = w13_weight.shape[1] // 2
+            hidden_size = w13_weight.shape[2]
+
+            padded_intermediate_size = align_to(intermediate_size, 256)
+            padded_hidden_size = align_to(hidden_size, 256)
+
+            # Transpose w2_weight to (num_experts, intermediate_size, hidden_size)
+            w13_weight = w13_weight.reshape(num_experts, 2, intermediate_size,
+                                            hidden_size)
+            w13_weight = jnp.swapaxes(w13_weight, 3, 2)
+
+            w2_weight = jnp.swapaxes(w2_weight, 2, 1)
+
+            w13_weight = jnp.pad(
+                w13_weight,
+                ((0, 0), (0, 0), (0, padded_hidden_size - hidden_size),
+                 (0, padded_intermediate_size - intermediate_size)),
+                constant_values=0)
+
+            w2_weight = jnp.pad(
+                w2_weight,
+                ((0, 0), (0, padded_intermediate_size - intermediate_size),
+                 (0, padded_hidden_size - hidden_size)),
+                constant_values=0)
+
+            # Apply EP sharding
+            ep_sharding = NamedSharding(self.mesh, P("model"))
+
+            w13_weight = jax.device_put(
+                w13_weight,
+                Format(Layout((0, 1, 2, 3)),
+                       NamedSharding(self.mesh, P("model", None, None, None))))
+            w2_weight = jax.device_put(
+                w2_weight,
+                Format(Layout((0, 1, 2)),
+                       NamedSharding(self.mesh, P("model", None, None))))
+
+            if self.moe.has_bias:
+                w13_bias = w13_bias.astype(jnp.float32).reshape(
+                    num_experts, 2, 1, intermediate_size)
+                w2_bias = w2_bias.astype(jnp.float32).reshape(
+                    num_experts, 1, hidden_size)
+
+                w13_bias = jnp.pad(
+                    w13_bias,
+                    ((0, 0), (0, 0), (0, 0),
+                     (0, padded_intermediate_size - intermediate_size)),
+                    constant_values=0)
+
+                w2_bias = jnp.pad(w2_bias,
+                                  ((0, 0), (0, 0),
+                                   (0, padded_hidden_size - hidden_size)),
+                                  constant_values=0)
+
+                # Apply EP sharding
+                w13_bias = jax.device_put(
+                    w13_bias, Format(Layout((0, 1, 2, 3)), ep_sharding))
+                w2_bias = jax.device_put(
+                    w2_bias, Format(Layout((0, 1, 2)), ep_sharding))
+        else:
+            if self.moe.has_bias:
+                w13_bias = jnp.expand_dims(w13_bias, 1)
+                w2_bias = jnp.expand_dims(w2_bias, 1)
+
+            if layer.use_ep:
+                ep_sharding = NamedSharding(self.mesh, P("model"))
+                w13_weight = jax.device_put(
+                    w13_weight, Format(Layout((0, 1, 2)), ep_sharding))
+                w2_weight = jax.device_put(
+                    w2_weight, Format(Layout((0, 1, 2)), ep_sharding))
+
+                if self.moe.has_bias:
+                    w13_bias = jax.device_put(
+                        w13_bias, Format(Layout((0, 1, 2)), ep_sharding))
+                    w2_bias = jax.device_put(
+                        w2_bias, Format(Layout((0, 1, 2)), ep_sharding))
+
+            else:
+                output_sizes = [intermediate_size, intermediate_size]
+                n_shards = self.mesh.shape["model"]
+                assert intermediate_size % n_shards == 0
+
+                w13_weight = reorder_concatenated_tensor_for_sharding(
+                    w13_weight, output_sizes, n_shards, dim=1)
+                w13_weight = jax.device_put(
+                    w13_weight,
+                    Format(Layout((0, 1, 2)),
+                           NamedSharding(self.mesh, P(None, "model", None))))
+                w2_weight = jax.device_put(
+                    w2_weight,
+                    Format(Layout((0, 1, 2)),
+                           NamedSharding(self.mesh, P(None, None, "model"))))
+
+                if self.moe.has_bias:
+                    w13_bias = reorder_concatenated_tensor_for_sharding(
+                        w13_bias, output_sizes, n_shards, dim=2)
+
+                    w13_bias = jax.device_put(
+                        w13_bias,
+                        Format(
+                            Layout((0, 1, 2)),
+                            NamedSharding(self.mesh, P(None, None, "model"))))
+                    w2_bias = jax.device_put(
+                        w2_bias,
+                        Format(Layout((0, 1, 2)),
+                               NamedSharding(self.mesh, P(None, None, None))))
+
+        layer.w13_weight = Parameter(torch_view(w13_weight),
+                                     requires_grad=False)
+        layer.w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
+
+        if self.moe.has_bias:
+            layer.w13_bias = Parameter(torch_view(w13_bias),
+                                       requires_grad=False)
+            layer.w2_bias = Parameter(torch_view(w2_bias), requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        assert isinstance(layer, FusedMoE)
+        if layer.scoring_func != "softmax":
+            raise NotImplementedError(
+                "Only softmax is supported for scoring_func")
+
+        x = jax_view(x)
+        w13_weight = jax_view(layer.w13_weight)
+        w2_weight = jax_view(layer.w2_weight)
+        w13_bias = w2_bias = None
+        if self.moe.has_bias:
+            w13_bias = jax_view(layer.w13_bias)
+            w2_bias = jax_view(layer.w2_bias)
+        gating_output = jax_view(router_logits)
+
+        if self.use_kernel:
+            actual_hidden_size = x.shape[-1]
+            padding_size = w13_weight.shape[-2] - actual_hidden_size
+            x = jnp.pad(x, ((0, 0), (0, padding_size)))
+            output = fused_ep_moe(
+                mesh=self.mesh,
+                tokens=x,
+                w1=w13_weight,
+                w2=w2_weight,
+                b1=w13_bias,
+                b2=w2_bias,
+                gating_output=gating_output,
+                top_k=layer.top_k,
+                ep_axis_name=self.ep_axis_name,
+                renormalize_topk_logits=layer.renormalize,
+                act_fn=layer.activation,
+                **self.block_size,
+            )[:, :actual_hidden_size]
+        else:
+            output = fused_moe_func(
+                hidden_states=x,
+                w1=w13_weight,
+                w2=w2_weight,
+                w1_scale=None,
+                w2_scale=None,
+                w1_bias=w13_bias,
+                w2_bias=w2_bias,
+                gating_output=gating_output,
+                topk=layer.top_k,
+                renormalize=layer.renormalize,
+                mesh=self.mesh,
+                use_ep=layer.use_ep,
+                activation=layer.activation,
+            )
+
+        return torch_view(output)

tpu_inference/layers/vllm/sharding.py

@@ -0,0 +1,244 @@
+# 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 os
+
+import jax
+import jax.numpy as jnp
+import torch
+import torchax
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torch.nn import Parameter
+from torch.utils import _pytree as pytree
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import t2j
+from vllm.lora.layers import (ColumnParallelLinearWithLoRA,
+                              MergedColumnParallelLinearWithLoRA,
+                              MergedQKVParallelLinearWithLoRA,
+                              QKVParallelLinearWithLoRA,
+                              ReplicatedLinearWithLoRA,
+                              RowParallelLinearWithLoRA)
+from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead, VocabParallelEmbedding)
+
+from tpu_inference import envs
+from tpu_inference.logger import init_logger
+
+P = PartitionSpec
+
+logger = init_logger(__name__)
+
+TORCH_TO_JAX_DTYPE_MAP = {
+    torch.float32: jnp.float32,
+    torch.float16: jnp.float16,
+    torch.bfloat16: jnp.bfloat16,
+}
+
+
+def shard_model_to_tpu(model: torch.nn.Module,
+                       mesh: Mesh) -> dict[str, torchax.torch.Tensor]:
+    """
+    Shard the model weights and move them to TPU.
+    At the same time, also turn the weight tensors into torchax tensors so that
+    jax code can interop with it and the overall program can be traced and
+    compiled in XLA.
+    Args:
+        model: A PyTorch model whose weights are on CPU main memory.
+        mesh: JAX mesh object for sharding.
+    Returns:
+        Dictionary of parameters and buffers that will be used as arguments of
+        torch.func.functional_call
+    """
+
+    with jax.default_device(jax.devices("cpu")[0]):
+        _shard_module_to_tpu(model, mesh)
+
+        params, buffers = _extract_all_params_buffers(model)
+
+        # For other weight tensors, repliate them on all the TPU chips.
+        params, buffers = pytree.tree_map_only(
+            _tensor_is_in_cpu,
+            lambda tensor: _shard_tensor_to_tpu_replicated(tensor, mesh),
+            (params, buffers))
+
+        return {**params, **buffers}
+
+
+def update_lora(model: torch.nn.Module,
+                initial_params_buffers) -> dict[str, torchax.torch.Tensor]:
+    params, buffers = _extract_all_params_buffers(model)
+    params_buffers = {**params, **buffers}
+    for k, v in params_buffers.items():
+        if 'lora_a_stacked' in k or 'lora_b_stacked' in k:
+            assert k in initial_params_buffers, f"{k} not in initial_params_buffers"
+            initial_params_buffers[k] = v
+
+    return initial_params_buffers
+
+
+def _extract_all_params_buffers(model: torch.nn.Module):
+    return dict(model.named_parameters()), dict(model.named_buffers())
+
+
+def _tensor_is_in_cpu(tensor: torch.tensor) -> bool:
+    # Check if a tensor haven't been converted to torchax tensor.
+    if not isinstance(tensor, torchax.tensor.Tensor):
+        return True
+    # Check if torchax tensor is still in CPU.
+    return tensor.jax_device == jax.devices('cpu')[0]
+
+
+def _convert_to_torchax_and_shard(tensor: torch.Tensor,
+                                  sharding: NamedSharding) -> torch.Tensor:
+    if os.getenv("VLLM_TPU_USING_PATHWAYS", False) and isinstance(
+            tensor, torch.Tensor):
+        np_tensor = tensor.detach().cpu().to(torch.float32).numpy()
+        dtype = TORCH_TO_JAX_DTYPE_MAP.get(tensor.dtype, jnp.float32)
+        return torch_view(jax.device_put(np_tensor, sharding).astype(dtype))
+    else:
+        if isinstance(tensor, torchax.tensor.Tensor):
+            tensor = jax_view(tensor)
+        else:
+            tensor = t2j(tensor)
+        return torch_view(_sharded_device_put(tensor, sharding))
+
+
+def _shard_tensor_to_tpu_replicated(tensor: torch.Tensor,
+                                    mesh: Mesh) -> torchax.tensor.Tensor:
+    return _convert_to_torchax_and_shard(tensor, NamedSharding(mesh, P()))
+
+
+def _shard_vocab_parallel_embedding(layer: VocabParallelEmbedding,
+                                    mesh: Mesh) -> None:
+    weight = _convert_to_torchax_and_shard(
+        layer.weight, NamedSharding(mesh, P('model', None)))
+    layer.weight = Parameter(weight, requires_grad=False)
+
+
+def _shard_lm_head(layer: ParallelLMHead, mesh: Mesh):
+    # TODO(qihqi): currently this is not handling case of tie_word_weights=True.
+    # if that config is set, then we should not create new weights but reuse the
+    # weight from VocabParallelEmbedding
+    weight = _convert_to_torchax_and_shard(
+        layer.weight, NamedSharding(mesh, P('model', None)))
+    layer.weight = Parameter(weight, requires_grad=False)
+    if layer.bias is not None:
+        bias = _convert_to_torchax_and_shard(layer.bias,
+                                             NamedSharding(mesh, P('model')))
+        layer.bias = Parameter(bias, requires_grad=False)
+
+
+def _shard_base_linear_lora_replicated(layer: BaseLinearLayerWithLoRA,
+                                       mesh: Mesh) -> None:
+    # NOTE: lora_a_stacked[i] has shape [max_loras, 1, num_out, num_in]
+    sharded_lora_a_tpu = torch.nn.ParameterList()
+    sharded_lora_b_tpu = torch.nn.ParameterList()
+
+    for i in range(layer.n_slices):
+        sharded_lora_a_tpu.append(
+            _shard_tensor_to_tpu_replicated(layer.lora_a_stacked[i], mesh))
+        sharded_lora_b_tpu.append(
+            _shard_tensor_to_tpu_replicated(layer.lora_b_stacked[i], mesh))
+
+    layer.lora_a_stacked = sharded_lora_a_tpu
+    layer.lora_b_stacked = sharded_lora_b_tpu
+
+
+def _shard_column_linear_lora(layer: ColumnParallelLinearWithLoRA,
+                              mesh: Mesh) -> None:
+    assert layer.n_slices > 0, "layer.n_slices should be greater than 0"
+    # lora_a_stacked[i] has shape [max_loras, 1, max_lora_rank, in_features]
+    sharded_lora_a_tpu = torch.nn.ParameterList()
+    sharded_lora_b_tpu = torch.nn.ParameterList()
+
+    # lora_b_stacked[i] has shape [max_loras, 1, out_features, max_lora_rank]
+    lora_b_partition_spec = P(None, None, 'model', None)
+    lora_b_sharding = NamedSharding(mesh, lora_b_partition_spec)
+    for i in range(layer.n_slices):
+        sharded_lora_a_tpu.append(
+            _shard_tensor_to_tpu_replicated(layer.lora_a_stacked[i], mesh))
+
+        sharded_lora_b_tpu.append(
+            _convert_to_torchax_and_shard(layer.lora_b_stacked[i],
+                                          lora_b_sharding))
+
+    layer.lora_a_stacked = sharded_lora_a_tpu
+    layer.lora_b_stacked = sharded_lora_b_tpu
+
+
+def _shard_qkv_linear_lora(layer: ColumnParallelLinearWithLoRA,
+                           mesh: Mesh) -> None:
+    _shard_column_linear_lora(layer, mesh)
+
+
+def _shard_merged_column_parallel_linear_lora(
+        layer: MergedColumnParallelLinearWithLoRA, mesh: Mesh) -> None:
+    _shard_column_linear_lora(layer, mesh)
+
+
+def _shard_merged_qkv_parallel_linear_lora(
+        layer: MergedQKVParallelLinearWithLoRA, mesh: Mesh) -> None:
+    _shard_column_linear_lora(layer, mesh)
+
+
+def _shard_row_parallel_linear_lora(layer: RowParallelLinearWithLoRA,
+                                    mesh: Mesh) -> None:
+    _shard_base_linear_lora_replicated(layer, mesh)
+
+
+# NOTE: Ordering is important as it calls first matched type of a given module
+MODULE_TYPE_TO_SHARDING_FUNC = [
+    # Shard embedding layers
+    (ParallelLMHead, _shard_lm_head),
+    (VocabParallelEmbedding, _shard_vocab_parallel_embedding),
+    # Shard LoRA layers
+    (ColumnParallelLinearWithLoRA, _shard_column_linear_lora),
+    (QKVParallelLinearWithLoRA, _shard_qkv_linear_lora),
+    (MergedColumnParallelLinearWithLoRA,
+     _shard_merged_column_parallel_linear_lora),
+    (MergedQKVParallelLinearWithLoRA, _shard_merged_qkv_parallel_linear_lora),
+    (RowParallelLinearWithLoRA, _shard_row_parallel_linear_lora),
+    (ReplicatedLinearWithLoRA, _shard_base_linear_lora_replicated),
+]
+
+
+def _shard_module_to_tpu(model: torch.nn.Module, mesh: Mesh) -> None:
+    for path, module in model.named_modules():
+        for module_type, sharding_func in MODULE_TYPE_TO_SHARDING_FUNC:
+            if type(module) is module_type:
+                logger.debug("shard %s with %s", path, sharding_func)
+                sharding_func(module, mesh)
+                break
+
+
+def _sharded_device_put(tensor: jax.Array, sharding) -> jax.Array:
+    if isinstance(tensor, tuple):
+        return tuple(_sharded_device_put(t, sharding) for t in tensor)
+    multihost_backend = envs.TPU_MULTIHOST_BACKEND
+    if multihost_backend != "ray":
+        return jax.device_put(tensor, sharding)
+
+    # NOTE: at here, num_global_devices != num_local_devices
+    # meaning we are in multi-host setup. Each host will run the same process
+    # and each process only need to handle the devices accessible to this host.
+    shape = tensor.shape
+    x_split = [
+        jax.device_put(tensor[i], device) for device, i in
+        sharding.addressable_devices_indices_map(shape).items()
+    ]
+    return jax.make_array_from_single_device_arrays(shape,
+                                                    sharding,
+                                                    x_split,
+                                                    dtype=tensor.dtype)

tpu_inference/logger.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from vllm.logger import _VllmLogger
+from vllm.logger import init_logger as init_vllm_logger
+
+
+def init_logger(name: str) -> _VllmLogger:
+    # Prepend the root "vllm" to the module path to use vllm's configured logger.
+    patched_name = "vllm." + name
+    return init_vllm_logger(patched_name)

tpu_inference/lora/__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.