tpu-inference 0.11.1__py3-none-any.whl

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -0,0 +1,272 @@
+import copy
+import functools
+import os
+from collections.abc import Sequence
+from contextlib import nullcontext
+from typing import Any, List, Optional, Tuple
+from unittest.mock import patch
+
+import jax
+import torch
+import torch.nn
+import torchax
+from flax.typing import PRNGKey
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import TORCH_DTYPE_TO_JAX
+from vllm.config import VllmConfig
+from vllm.forward_context import set_forward_context
+from vllm.lora.layers import BaseLayerWithLoRA
+from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
+from vllm.model_executor.model_loader import get_model as vllm_get_model
+from vllm.model_executor.models import supports_lora, supports_multimodal
+from vllm.sequence import IntermediateTensors
+
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.vllm.quantization import get_tpu_quantization_config
+from tpu_inference.layers.vllm.sharding import shard_model_to_tpu
+from tpu_inference.logger import init_logger
+from tpu_inference.models.vllm.vllm_model_wrapper_context import (
+    get_vllm_model_wrapper_context, set_vllm_model_wrapper_context)
+from tpu_inference.runner.lora_utils import replace_lora_metadata
+
+logger = init_logger(__name__)
+
+
+class _VllmRunner(torch.nn.Module):
+
+    def __init__(self, vllm_model: torch.nn.Module):
+        super().__init__()
+        self.vllm_model = vllm_model
+
+    def forward(self, **kwargs) -> torch.Tensor:
+        # We don't support multimodal input in Gemma3, but we need patch it to
+        # None to workaround vLLM Gemma3 model bug that
+        # `get_multimodal_embeddings` returns empty list but it's caller checks
+        # for None.
+        with patch(
+                "vllm.model_executor.models.gemma3_mm."
+                "Gemma3ForConditionalGeneration."
+                "get_multimodal_embeddings",
+                return_value=None):
+            if "hidden_state" in kwargs:
+                return self.compute_logits(kwargs["hidden_state"])
+            else:
+                return self.compute_hidden_state(
+                    kwargs["input_ids"],
+                    kwargs["positions"],
+                    kwargs["intermediate_tensors"],
+                    kwargs["inputs_embeds"],
+                )
+
+    def compute_hidden_state(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors],
+        inputs_embeds: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        hidden_state = self.vllm_model(input_ids, positions,
+                                       intermediate_tensors, inputs_embeds)
+        return hidden_state
+
+    def compute_logits(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        return self.vllm_model.compute_logits(hidden_state)
+
+
+class VllmModelWrapper:
+    """ Wraps a vLLM Pytorch model and let it run on the JAX engine. """
+
+    rng: PRNGKey
+    mesh: Mesh
+    model: _VllmRunner
+
+    def __init__(self, vllm_config: VllmConfig, rng: PRNGKey, mesh: Mesh):
+        self.vllm_config = vllm_config
+        self.rng = rng
+        self.mesh = mesh
+
+        self.vllm_config.quant_config = get_tpu_quantization_config(
+            self.vllm_config, self.mesh)
+
+    def load_weights(self):
+        # Set up to load the model into CPU first.
+        vllm_config_for_load = copy.deepcopy(self.vllm_config)
+        assert self.vllm_config.model_config.dtype in TORCH_DTYPE_TO_JAX, "The model_config.dtype must be a PyTorch dtype."
+        vllm_config_for_load.device_config.device = "cpu"
+
+        if os.getenv("JAX_RANDOM_WEIGHTS", False):
+            vllm_config_for_load.load_config.load_format = "dummy"
+            use_random_weights = True
+        else:
+            use_random_weights = (
+                vllm_config_for_load.load_config.load_format == "dummy")
+        if use_random_weights:
+            logger.info(
+                "Initializing vLLM model with random weights, weight loading skipped."
+            )
+        # The DummyModelLoader in vLLM calls torch._sync for torch_xla path when
+        # it detects the tpu platform, but we don't need it and it causes crash
+        # without proper setup.
+        load_context = patch(
+            "torch._sync",
+            return_value=None) if use_random_weights else nullcontext()
+
+        # Load the vLLM model and wrap it into a new model whose forward
+        # function can calculate the hidden_state and logits.
+        with load_context, jax.default_device(jax.devices('cpu')[0]):
+            vllm_model = vllm_get_model(vllm_config=vllm_config_for_load)
+        lora_manager = None
+        if vllm_config_for_load.lora_config is not None:
+            # Replace layers in the model with LoRA layers.
+            with torchax.default_env():
+                # Argument "device" in load_lora_model is used to set the device
+                # used in punica wrapper.
+                lora_manager, vllm_model = load_lora_model(
+                    vllm_model, vllm_config_for_load, device="jax")
+            replace_set_lora(vllm_model)
+
+        static_forward_context = vllm_config_for_load.compilation_config.static_forward_context
+        self.vllm_config.compilation_config.static_forward_context = static_forward_context
+
+        self.model = _VllmRunner(vllm_model)
+        params_and_buffers = shard_model_to_tpu(self.model, self.mesh)
+
+        # Returning to the jax land, so we need to wrap it into a JaxValue.
+        return jax_view(params_and_buffers), lora_manager
+
+    def jit_step_func(self):
+
+        @functools.partial(
+            jax.jit,
+            donate_argnames=("kv_caches", ),
+            compiler_options={
+                "xla_tpu_all_gather_collective_matmul_mode":
+                "post_spmd_conservative",
+                "xla_tpu_reduce_scatter_collective_matmul_mode":
+                "post_spmd_conservative"
+            },
+            static_argnames=("layer_name_to_kvcache_index", ),
+        )
+        def step_fun(
+            params_and_buffers,  # This has been wrapped into torchax TorchValue
+            kv_caches: List[jax.Array],
+            input_ids: jax.Array,
+            attn_metadata: AttentionMetadata,
+            input_embeds: jax.Array,
+            layer_name_to_kvcache_index: Sequence[Tuple[str, int]],
+            lora_metadata,
+            *args,
+        ) -> Tuple[List[jax.Array], jax.Array]:
+            layer_name_to_kvcache_index = dict(layer_name_to_kvcache_index)
+            lora_metadata = torch_view(lora_metadata)
+            with torchax.default_env(), set_vllm_model_wrapper_context(
+                    kv_caches=kv_caches,
+                    mesh=self.mesh,
+                    layer_name_to_kvcache_index=layer_name_to_kvcache_index
+            ), set_forward_context(attn_metadata=attn_metadata,
+                                   vllm_config=self.vllm_config):
+                # We need to wrap args from jax land into TorchValue with
+                # torch_view in order to call the Torch function.
+                original_lora_metadata = replace_lora_metadata(
+                    self.model, lora_metadata, self.vllm_config.lora_config)
+                hidden_states = torch.func.functional_call(
+                    self.model,
+                    torch_view(params_and_buffers),
+                    kwargs={
+                        "input_ids": torch_view(input_ids),
+                        "positions": torch_view(attn_metadata.input_positions),
+                        "intermediate_tensors": None,
+                        "inputs_embeds": None,
+                    },
+                    tie_weights=False,
+                )
+                replace_lora_metadata(self.model, original_lora_metadata,
+                                      self.vllm_config.lora_config)
+                vllm_model_wrapper_context = get_vllm_model_wrapper_context()
+                new_kv_caches = vllm_model_wrapper_context.kv_caches
+            # Wrap the hidden_states from torch land into a JaxValue for the jax
+            # code to consume.
+            hidden_states = jax_view(hidden_states)
+
+            return new_kv_caches, hidden_states, []
+
+        return step_fun
+
+    def jit_compute_logits_func(self):
+
+        @functools.partial(
+            jax.jit,
+            out_shardings=(NamedSharding(self.mesh,
+                                         PartitionSpec(None, "model"))),
+        )
+        def compute_logits_func(
+            params_and_buffers: Any,
+            hidden_states: jax.Array,
+            lora_metadata,
+        ) -> jax.Array:
+            lora_metadata = torch_view(lora_metadata)
+            with torchax.default_env(), set_vllm_model_wrapper_context(
+                    kv_caches=None, mesh=self.mesh):
+                original_lora_metadata = replace_lora_metadata(
+                    self.model, lora_metadata, self.vllm_config.lora_config)
+                logits = torch.func.functional_call(
+                    self.model,
+                    torch_view(params_and_buffers),
+                    kwargs={
+                        "hidden_state": torch_view(hidden_states),
+                    },
+                    tie_weights=False,
+                )
+                replace_lora_metadata(self.model, original_lora_metadata,
+                                      self.vllm_config.lora_config)
+            return jax_view(logits)
+
+        return compute_logits_func
+
+
+def load_lora_model(model: torch.nn.Module, vllm_config: VllmConfig,
+                    device: str) -> torch.nn.Module:
+    if not supports_lora(model):
+        raise ValueError(
+            f"{model.__class__.__name__} does not support LoRA yet.")
+
+    if supports_multimodal(model):
+        logger.warning("Regarding multimodal models, vLLM currently "
+                       "only supports adding LoRA to language model.")
+
+    # Add LoRA Manager to the Model Runner
+    lora_manager = LRUCacheWorkerLoRAManager(
+        vllm_config,
+        device,
+        model.embedding_modules,
+        model.embedding_padding_modules,
+    )
+    return lora_manager, lora_manager.create_lora_manager(model)
+
+
+# The reason why replace the method is that the set_lora and reset_lora need to
+# run under torchax env.
+def replace_set_lora(model):
+
+    def _tpu_set_lora(
+        self,
+        index: int,
+        lora_a: torch.Tensor,
+        lora_b: torch.Tensor,
+        embeddings_tensor: Optional[torch.Tensor],
+    ):
+        with torchax.default_env():
+            self._original_set_lora(index, lora_a, lora_b, embeddings_tensor)
+
+    def _tpu_reset_lora(self, index: int):
+        with torchax.default_env():
+            self._original_reset_lora(index)
+
+    for _, module in model.named_modules():
+        if isinstance(module, BaseLayerWithLoRA):
+            module._original_set_lora = module.set_lora
+            module._original_reset_lora = module.reset_lora
+            module.set_lora = _tpu_set_lora.__get__(module, module.__class__)
+            module.reset_lora = _tpu_reset_lora.__get__(
+                module, module.__class__)

tpu_inference/models/vllm/vllm_model_wrapper_context.py

@@ -0,0 +1,45 @@
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import jax
+from jax.sharding import Mesh
+
+
+@dataclass
+class VllmModelWrapperContext:
+    kv_caches: List[jax.Array]
+    mesh: Mesh
+    layer_name_to_kvcache_index: Dict[str, int]
+
+
+_vllm_model_wrapper_context: Optional[VllmModelWrapperContext] = None
+
+
+def get_vllm_model_wrapper_context() -> VllmModelWrapperContext:
+    assert _vllm_model_wrapper_context is not None, (
+        "VllmModelWrapperContext is not set. "
+        "Please use `set_vllm_model_wrapper_context` to set the VllmModelWrapperContext."
+    )
+    return _vllm_model_wrapper_context
+
+
+@contextmanager
+def set_vllm_model_wrapper_context(
+    *,
+    kv_caches: List[jax.Array],
+    mesh: Mesh,
+    layer_name_to_kvcache_index: Dict[str, int] = None,
+):
+    global _vllm_model_wrapper_context
+    prev_context = _vllm_model_wrapper_context
+    _vllm_model_wrapper_context = VllmModelWrapperContext(
+        kv_caches=kv_caches,
+        mesh=mesh,
+        layer_name_to_kvcache_index=layer_name_to_kvcache_index,
+    )
+
+    try:
+        yield
+    finally:
+        _vllm_model_wrapper_context = prev_context

tpu_inference/platforms/__init__.py

@@ -0,0 +1,2 @@
+# ruff: noqa
+from tpu_inference.platforms.tpu_jax import TpuPlatform

tpu_inference/platforms/tpu_jax.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import TYPE_CHECKING, Optional, Tuple, Union, cast
+
+import jax.numpy as jnp
+import vllm.envs as envs
+from torchax.ops.mappings import j2t_dtype
+from tpu_info import device
+from vllm.inputs import ProcessorInputs, PromptType
+from vllm.platforms.interface import Platform, PlatformEnum
+from vllm.sampling_params import SamplingParams, SamplingType
+
+from tpu_inference.logger import init_logger
+from tpu_inference.models.jax.utils.quantization.quantization_utils import \
+    update_vllm_config_for_qwix_quantization
+
+if TYPE_CHECKING:
+    from vllm.attention.backends.registry import _Backend
+    from vllm.config import BlockSize, ModelConfig, VllmConfig
+    from vllm.pooling_params import PoolingParams
+else:
+    BlockSize = None
+    ModelConfig = None
+    VllmConfig = None
+    PoolingParams = None
+    _Backend = None
+
+logger = init_logger(__name__)
+
+_DTYPE: dict[str, jnp.dtype] = {
+    "bfloat16": jnp.bfloat16,
+    "float": jnp.float32,
+    "float32": jnp.float32,
+}
+
+
+class TpuPlatform(Platform):
+    _enum = PlatformEnum.TPU
+    device_name: str = "tpu"
+    device_type: str = "tpu"
+    dispatch_key: str = "XLA"
+    ray_device_key: str = "TPU"
+    device_control_env_var: str = "TPU_VISIBLE_CHIPS"
+    simple_compile_backend: str = "openxla"
+
+    supported_quantization: list[str] = [
+        "tpu_int8", "compressed-tensors", "awq", "fp8"
+    ]
+
+    additional_env_vars: list[str] = [
+        "JAX_RANDOM_WEIGHTS", "PHASED_PROFILING_DIR",
+        "TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS",
+        "TPU_MULTIHOST_BACKEND", "VLLM_MLA_DISABLE", "NEW_MODEL_DESIGN",
+        "TPU_BACKEND_TYPE"
+    ]
+
+    @classmethod
+    def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
+                             dtype: jnp.dtype, kv_cache_dtype: Optional[str],
+                             block_size: int, use_v1: bool, use_mla: bool,
+                             has_sink: bool, use_sparse: bool) -> str:
+        from vllm.attention.backends.registry import _Backend
+        if selected_backend != _Backend.PALLAS:
+            logger.info("Cannot use %s backend on TPU.", selected_backend)
+
+        if use_v1:
+            logger.info("Using Pallas V1 backend.")
+            return "tpu_inference.layers.vllm.attention.PallasAttentionBackend"
+        else:
+            logger.info("Using Pallas backend.")
+            return "vllm.attention.backends.pallas.PallasAttentionBackend"
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        try:
+            if envs.VLLM_TPU_USING_PATHWAYS:
+                # Causes mutliprocess accessing IFRT when calling jax.devices()
+                return "TPU v6 lite"
+            else:
+                chip_type, _ = device.get_local_chips()
+                return f"TPU {chip_type.name}"
+        except Exception as e:
+            logger.warning(f"Error getting device name: {e}")
+            return 'TPU'
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
+        return not envs.VLLM_USE_V1
+
+    @classmethod
+    def get_punica_wrapper(cls) -> str:
+        return "tpu_inference.lora.torch_punica_tpu.PunicaWrapperTPU"
+
+    @classmethod
+    def get_infinity_values(cls, dtype: jnp.dtype) -> Tuple[float, float]:
+        return jnp.finfo(dtype).min, jnp.finfo(dtype).max
+
+    @classmethod
+    def can_update_inplace(cls):
+        return False
+
+    @classmethod
+    def get_lora_vocab_padding_size(cls) -> int:
+        return 1
+
+    @classmethod
+    def inference_mode(cls):
+        return True
+
+    @classmethod
+    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        if not envs.VLLM_USE_V1:
+            raise RuntimeError("VLLM_USE_V1=1 must be set for JAX backend.")
+
+        if envs.VLLM_TPU_USING_PATHWAYS:
+            assert not envs.VLLM_ENABLE_V1_MULTIPROCESSING, (
+                "VLLM_ENABLE_V1_MULTIPROCESSING must be 0 when using Pathways(JAX_PLATFORMS=proxy)"
+            )
+
+        from vllm.config import CompilationLevel
+
+        cache_config = vllm_config.cache_config
+        # For v0, the default block size is 16.
+        if cache_config and cache_config.block_size is None:
+            cache_config.block_size = cast(BlockSize, 16)
+        compilation_config = vllm_config.compilation_config
+
+        # TPU only supports DYNAMO_ONCE compilation level
+        # NOTE(xiang): the compilation_config is not used by jax.
+        if compilation_config.level != CompilationLevel.DYNAMO_ONCE:
+            compilation_config.level = CompilationLevel.DYNAMO_ONCE
+
+        if compilation_config.backend == "":
+            compilation_config.backend = "openxla"
+
+        # If we use vLLM's model implementation in PyTorch, we should set it with torch version of the dtype.
+        impl = os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower()
+
+        # NOTE(xiang): convert dtype to jnp.dtype
+        # NOTE(wenlong): skip this logic for mm model preprocessing
+        # For mm model preprocessors, it may need the output dtype to be torch.
+        # In order to avoid a PR to vLLM, we postpone the dtype checking during tpu_worker initialization
+        if not vllm_config.scheduler_config.is_multimodal_model or impl == "vllm":
+            if not isinstance(vllm_config.model_config.dtype, str):
+                logger.warning(
+                    "The model dtype is not properly set for JAX backend. "
+                    "Overwriting it to jnp.bfloat16")
+                vllm_config.model_config.dtype = jnp.bfloat16
+            else:
+                vllm_config.model_config.dtype = _DTYPE.get(
+                    vllm_config.model_config.dtype, jnp.bfloat16)
+
+        if impl == "vllm":
+            vllm_config.model_config.dtype = j2t_dtype(
+                vllm_config.model_config.dtype.dtype)
+
+        if envs.VLLM_USE_V1:
+            # TODO(cuiq): remove this dependency.
+            from vllm.v1.attention.backends.pallas import \
+                PallasAttentionBackend
+            cache_config.block_size = PallasAttentionBackend.get_page_size(
+                vllm_config)  # type: ignore[assignment]
+            min_page_size = PallasAttentionBackend.get_min_page_size(
+                vllm_config)
+            if min_page_size > cache_config.block_size:
+                logger.warning(
+                    "Increase the page size from %s to %s to make sure there's"
+                    "no SMEM OOM",
+                    cache_config.block_size,
+                    min_page_size,
+                )
+                cache_config.block_size = min_page_size  # type: ignore[assignment]
+
+        parallel_config = vllm_config.parallel_config
+        scheduler_config = vllm_config.scheduler_config
+        parallel_config.worker_cls = \
+                        "tpu_inference.worker.tpu_worker_jax.TPUWorker"
+
+        multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
+        if not multihost_backend:  # Single host
+            logger.info("Force using UniProcExecutor for JAX on single host.")
+            parallel_config.distributed_executor_backend = "uni"
+        elif multihost_backend == "ray":
+            from tpu_inference.executors.ray_distributed_executor import \
+                RayDistributedExecutor
+            parallel_config.distributed_executor_backend = RayDistributedExecutor
+            logger.info(
+                "Force using RayDistributedExecutor for JAX on single host.")
+        else:
+            logger.warning(
+                f"Unknown TPU multihost backend: {multihost_backend}. "
+                "Using uniproc_executor.")
+            parallel_config.distributed_executor_backend = "uni"
+
+        if scheduler_config.is_multimodal_model and not \
+            scheduler_config.disable_chunked_mm_input:
+            logger.warning("TPU does not support running Multimodal models"\
+            " without setting `--disable_chunked_mm_input`. " \
+            "Forcing --disable_chunked_mm_input.")
+            scheduler_config.disable_chunked_mm_input = True
+
+        kv_transfer_config = vllm_config.kv_transfer_config
+        if kv_transfer_config is not None:
+            assert kv_transfer_config.kv_connector == "TPUConnector"
+
+        update_vllm_config_for_qwix_quantization(vllm_config)
+
+    @classmethod
+    def is_pin_memory_available(cls):
+        logger.warning("Pin memory is not supported on TPU.")
+        return False
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "vllm.distributed.device_communicators.tpu_communicator.TpuCommunicator"  # noqa
+
+    @classmethod
+    def use_all_gather(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_v1(cls, model_config: ModelConfig) -> bool:
+        # V1 support on TPU is experimental
+        return True
+
+    @classmethod
+    def validate_request(
+        cls,
+        prompt: PromptType,
+        params: Union[SamplingParams, PoolingParams],
+        processed_inputs: ProcessorInputs,
+    ) -> None:
+        """Raises if this request is unsupported on this platform"""
+
+        if isinstance(params, SamplingParams):
+            if params.structured_outputs is not None and not envs.VLLM_USE_V1:
+                raise ValueError("Structured output is not supported on "
+                                 f"{cls.device_name} V0.")
+            if params.sampling_type == SamplingType.RANDOM_SEED:
+                raise ValueError("JAX does not support per-request seed.")
+
+    @classmethod
+    def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,
+                                    model_config: ModelConfig) -> bool:
+        return True
+
+    @classmethod
+    def use_sync_weight_loader(cls) -> bool:
+        """
+        Returns if the current platform needs to sync weight loader.
+        """
+        return True