tpu-inference 0.11.1__py3-none-any.whl

tpu_inference/utils.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+from collections import defaultdict
+from collections.abc import Sequence
+from typing import Any, Callable, List, Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax._src import dtypes
+from jax._src import mesh as mesh_lib
+from jax._src import xla_bridge as xb
+from jax._src.lib import xla_client as xc
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from vllm import envs, utils
+
+from tpu_inference.logger import init_logger
+
+GBYTES = 1024 * 1024 * 1024
+TPU_HEAD_SIZE_ALIGNMENT = 128
+TPU_SECOND_LAST_MINOR = 8
+
+# This is used to translate from a string name for a dtype
+# to formal jax.numpy DType.  One use case for this is
+# converting the `--kv_cache_dtype` flag to a dtype.
+TPU_STR_DTYPE_TO_JAX_DTYPE = {
+    "bfloat16": jnp.bfloat16,
+    "fp8": jnp.float8_e4m3fn,
+    "fp8_e4m3": jnp.float8_e4m3,
+    "fp8_e5m2": jnp.float8_e5m2,
+    "int8": jnp.int8,
+}
+
+_megacore = False
+logger = init_logger(__name__)
+
+
+def enable_megacore() -> None:
+    global _megacore
+    _megacore = True
+
+
+def get_megacore() -> bool:
+    return _megacore
+
+
+def get_num_kv_heads_by_tp(num_kv_heads: int, tp_size: int) -> int:
+    if tp_size <= num_kv_heads:
+        assert num_kv_heads % tp_size == 0
+        return num_kv_heads
+    else:
+        assert tp_size % num_kv_heads == 0
+        return tp_size
+
+
+def hbm_usage_bytes(devices: Any) -> List[Tuple[int, int]]:
+    usage = []
+    if envs.VLLM_TPU_USING_PATHWAYS:
+        return pathways_hbm_usage_gb(devices)
+
+    multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
+    if multihost_backend == "ray":
+        # MemoryStats is only supported for addressable PjRt devices.
+        # Assume all the devices have similar memory usage for now.
+        # TODO(ranlihao): find a proper way to get the memory usage of each device.
+        for device in devices:
+            try:
+                hbm_used = device.memory_stats()["bytes_in_use"]
+                hbm_limit = device.memory_stats()["bytes_limit"]
+                logger.info(
+                    "Get memory stats for device %s. Assuming all devices have the same usage.",
+                    device)
+                usage.extend([(hbm_used, hbm_limit)] * len(devices))
+                break
+            except Exception as e:
+                logger.warning(
+                    "Failed to get memory stats for device %s: %s. ", device,
+                    e)
+    else:
+        for device in devices:
+            hbm_used = device.memory_stats()["bytes_in_use"]
+            hbm_limit = device.memory_stats()["bytes_limit"]
+            usage.append((hbm_used, hbm_limit))
+
+    return usage
+
+
+def get_device_name(num_devices: int | None = None):
+    kind = jax.devices()[0].device_kind
+    if 'TPU' not in kind:
+        raise RuntimeError('Expected TPU devices')
+    suffix = ''
+    if kind.endswith(' lite'):
+        kind = kind[:-len(' lite')]
+        suffix = 'e'
+    elif kind.endswith('e'):
+        kind = kind[:-1]
+        suffix = 'e'
+    elif kind.endswith('p'):
+        kind = kind[:-1]
+        suffix = 'p'
+    elif kind == 'TPU7x':
+        kind = 'TPU v7'
+    assert kind[:-1] == 'TPU v', kind
+    kind += suffix
+    if num_devices is not None:
+        kind += f'-{num_devices}'
+    return kind
+
+
+def get_device_hbm_limit() -> int:
+
+    device_kind = get_device_name()
+    if device_kind == "TPU v5p" or device_kind == "TPU v5":
+        return 95 * GBYTES
+    elif device_kind == "TPU v5e":
+        return 16 * GBYTES
+    elif device_kind == "TPU v6e" or device_kind == "TPU v4":
+        return 32 * GBYTES
+    elif device_kind == "TPU v7":
+        return 192 * GBYTES
+    else:
+        raise ValueError(f"Unknown device kind: {device_kind}")
+
+
+def pathways_hbm_usage_gb(devices: Any) -> List[Tuple[float, float]]:
+    live_arrays = jax.live_arrays()
+    hbm_used = defaultdict(int)
+    hbm_limit = get_device_hbm_limit()
+    for array in live_arrays:
+        assert hasattr(array, 'sharding') and hasattr(
+            array.sharding, 'device_set'
+        ), "This function must not be called within jax tracer (e.g. jit, vmap, grad)"
+        for device in array.sharding.device_set:
+            hbm_used[device] += array.dtype.itemsize * array.size // len(
+                array.sharding.device_set)
+    return [(hbm_used[device], hbm_limit) for device in devices]
+
+
+def hbm_usage_gb(devices: Any) -> List[Tuple[float, float]]:
+    usage = hbm_usage_bytes(devices)
+    usage = [(round(used / GBYTES, 2), round(limit / GBYTES, 2))
+             for used, limit in usage]
+    return usage
+
+
+def get_padded_head_dim(head_dim: int) -> int:
+    """Pads head_dim up to the nearest multiple of 128 for kernel performance."""
+    return (head_dim + 127) // 128 * 128
+
+
+def get_padded_num_heads(num_heads: int, sharding_size: int) -> int:
+    if num_heads >= sharding_size:
+        assert num_heads % sharding_size == 0
+    else:
+        assert sharding_size % num_heads == 0
+        num_heads = sharding_size
+    return num_heads
+
+
+def get_dtype_packing(dtype):
+    bits = dtypes.bit_width(dtype)
+    return 32 // bits
+
+
+def make_optimized_mesh(axis_shapes: Sequence[int],
+                        axis_names: Sequence[str],
+                        *,
+                        devices: Sequence[xc.Device] | None = None):
+    if devices is None:
+        devices = xb.devices()
+    # Sort the devices in case it's passed in an arbitary order
+    devices = sorted(devices, key=lambda x: x.coords)
+
+    def _is_1D(axis_shapes):
+        return sum(x > 1 for x in axis_shapes) == 1
+
+    if _is_1D(axis_shapes):
+        dev_kind = devices[0].device_kind
+        device_num = len(devices)
+        if dev_kind == "TPU v6 lite":
+            ordered_devices = None
+            # NOTE(chengjiyao):
+            # The coords of v6e-8 are
+            # (0,0,0)
+            # (1,0,0)
+            # (0,1,0)
+            # (1,1,0)
+            # (0,2,0)
+            # (1,2,0)
+            # (0,3,0)
+            # (1,3,0)
+            if device_num == 8:
+                ordered_devices = np.array([
+                    devices[0],
+                    devices[1],
+                    devices[2],
+                    devices[3],
+                    devices[7],
+                    devices[6],
+                    devices[5],
+                    devices[4],
+                ])
+            # NOTE(chengjiyao):
+            # The coords of v6e-4 are
+            # (0,0,0)
+            # (1,0,0)
+            # (0,1,0)
+            # (1,1,0)
+            elif device_num == 4:
+                ordered_devices = np.array([
+                    devices[0],
+                    devices[1],
+                    devices[3],
+                    devices[2],
+                ])
+            if ordered_devices is not None:
+                ordered_devices = np.array(ordered_devices)
+                ordered_devices = ordered_devices.reshape(axis_shapes)
+                mesh = mesh_lib.Mesh(ordered_devices, axis_names)
+                logger.info("Use customized mesh: %s", mesh)
+                return mesh
+
+    return jax.make_mesh(axis_shapes, axis_names, devices=devices)
+
+
+def device_array(mesh: Mesh, *args, sharding=None, **kwargs) -> jax.Array:
+    """
+    Create a device array with the specified mesh and sharding.
+
+    Args:
+        mesh: The JAX mesh to use for device placement
+        *args: Positional arguments to pass to jax.device_put
+        sharding: Optional sharding specification. If None, uses PartitionSpec(None)
+        **kwargs: Keyword arguments to pass to jax.device_put
+
+    Returns:
+        A JAX array placed on the specified devices
+    """
+    if sharding is None:
+        sharding = NamedSharding(mesh, PartitionSpec(None))
+    return jax.device_put(*args, device=sharding, **kwargs)
+
+
+def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
+    """
+    A wrapper function of vllm.utils.get_hash_fn_by_name to support builtin
+    """
+    if hash_fn_name == "builtin":
+        return hash
+    return utils.get_hash_fn_by_name(hash_fn_name)
+
+
+def quantize_kv(key: jax.Array, value: jax.Array,
+                kv_cache_quantized_dtype: jnp.dtype, k_scale: float,
+                v_scale: float) -> Tuple[jax.Array, jax.Array]:
+    """
+        Quantize the key and value tensors.
+
+        Args:
+            key: The key tensor to quantize.
+            value: The value tensor to quantize.
+            kv_cache_quantized_dtype: The dtype to quantize the key and value tensors to.
+            q_scale: The scale to quantize the key and value tensors by.
+            k_scale: The scale to quantize the key tensor by.
+            v_scale: The scale to quantize the value tensor by.
+
+        Returns:
+            Tuple[jax.Array, jax.Array]: The quantized key and value tensors.
+        """
+    dtype_info = jnp.finfo(kv_cache_quantized_dtype)
+    minval, maxval = float(dtype_info.min), float(dtype_info.max)
+    key = key.astype(jnp.float32) / k_scale
+    key = jnp.clip(key, minval, maxval)
+    key = key.astype(kv_cache_quantized_dtype)
+    value = value.astype(jnp.float32) / v_scale
+    value = jnp.clip(value, minval, maxval)
+    value = value.astype(kv_cache_quantized_dtype)
+
+    return key, value
+
+
+def get_jax_dtype_from_str_dtype(str_dtype: str) -> jnp.dtype:
+    """
+    Get the JAX dtype from a string dtype.
+
+    Args:
+        str_dtype: The string dtype to get the JAX dtype from.
+
+    Returns:
+        jnp.dtype: The JAX dtype.
+    """
+    str_dtype = str_dtype.lower().strip()
+    return TPU_STR_DTYPE_TO_JAX_DTYPE.get(str_dtype)

tpu_inference/worker/_temporary_vllm_compat.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+
+#
+# WARNING: This is a temporary compatibility module.
+#
+#
+# THE PROBLEM:
+# The ideal dependency injection pattern dictates that the "producer" of data
+# (in this case, the vLLM engine) should be responsible for adapting its data
+# into the abstract format that the "consumer" (the TPU worker) expects.
+#
+# However, this would require a simultaneous code change in both the `vllm` and
+# `tpu_inference` repositories. Such cross-repository changes are difficult to
+# coordinate, slow to land, and can easily cause breakages if the releases
+# are not perfectly synchronized.
+#
+#
+# THE TEMPORARY SOLUTION:
+# To enable independent development and deployment, we are temporarily violating
+# this pattern. We are making the consumer (`tpu_inference`) responsible for
+# detecting and adapting the producer's raw data.
+#
+# This function checks if it has received a raw `vllm.SchedulerOutput` and,
+# if so, wraps it in the appropriate adapter. This allows `vllm` to continue
+# sending its raw data type without modification, decoupling the release cycles.
+#
+#
+# THE FUTURE (HOW TO REMOVE THIS):
+# This entire file should be deleted once the `vllm` repository has been updated.
+# The required change in `vllm` is small and looks like this:
+#
+# --- SKELETON CODE FOR FUTURE vLLM CHANGE ---
+# In the vLLM engine, where `execute_model` is called:
+#
+# from tpu_inference.adapters.vllm_adapters import VllmSchedulerOutputAdapter
+# from vllm.v1.core.sched.output import SchedulerOutput
+#
+# # ... inside some method ...
+#
+# # OLD CODE:
+# # concrete_work = SchedulerOutput(...)
+# # self.tpu_backend.execute_model(concrete_work)
+#
+# # NEW CODE:
+# concrete_work = SchedulerOutput(...)
+# adapted_work = VllmSchedulerOutputAdapter(concrete_work)  # This line is added
+# self.tpu_backend.execute_model(adapted_work)          # Pass the adapter
+#
+# --- END SKELETON CODE ---
+#
+
+import logging
+from typing import Union
+
+from vllm.lora.request import LoRARequest as VllmLoRARequest
+# Import the concrete vLLM type for the check
+from vllm.v1.core.sched.output import SchedulerOutput as VllmSchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig as VllmKVCacheConfig
+
+from tpu_inference.adapters.vllm_adapters import (VllmKVCacheConfigAdapter,
+                                                  VllmLoRARequestAdapter,
+                                                  VllmSchedulerOutputAdapter)
+from tpu_inference.di.abstracts import (AbstractKVCacheConfig,
+                                        AbstractLoRARequest,
+                                        AbstractSchedulerOutput)
+
+logger = logging.getLogger(__name__)
+
+
+def adapt_scheduler_output_if_needed(
+    scheduler_output: Union[AbstractSchedulerOutput, VllmSchedulerOutput]
+) -> AbstractSchedulerOutput:
+    """
+    Checks if the input is a raw VllmSchedulerOutput and wraps it.
+    If it's already an AbstractSchedulerOutput, it's passed through.
+    """
+    if isinstance(scheduler_output, VllmSchedulerOutput):
+        # logger.warning(
+        #     "Received raw VllmSchedulerOutput. Performing temporary, on-the-fly "
+        #     "adaptation. This is a compatibility feature and should be removed "
+        #     "once the vLLM engine is updated to provide an adapted object.")
+        return VllmSchedulerOutputAdapter(scheduler_output)
+
+    if isinstance(scheduler_output, AbstractSchedulerOutput):
+        return scheduler_output
+
+    raise TypeError(
+        f"Unsupported type for scheduler_output: {type(scheduler_output)}")
+
+
+def adapt_kv_cache_config_if_needed(
+    kv_cache_config: Union[AbstractKVCacheConfig, VllmKVCacheConfig]
+) -> AbstractKVCacheConfig:
+    """
+    Checks if the input is a raw VllmKVCacheConfig and wraps it.
+    If it's already an AbstractKVCacheConfig, it's passed through.
+    """
+    if isinstance(kv_cache_config, VllmKVCacheConfig):
+        # logger.warning(
+        #     "Received raw VllmKVCacheConfig. Performing temporary, on-the-fly "
+        #     "adaptation. This is a compatibility feature and should be removed "
+        #     "once the vLLM engine is updated to provide an adapted object.")
+        return VllmKVCacheConfigAdapter(kv_cache_config)
+
+    if isinstance(kv_cache_config, AbstractKVCacheConfig):
+        return kv_cache_config
+
+    raise TypeError(
+        f"Unsupported type for kv_cache_config: {type(kv_cache_config)}")
+
+
+def adapt_lora_request_if_needed(
+    lora_request: Union[AbstractLoRARequest, VllmLoRARequest]
+) -> AbstractLoRARequest:
+    """
+    Checks if the input is a raw VllmLoRARequest and wraps it.
+    If it's already an AbstractLoRARequest, it's passed through.
+    """
+    if isinstance(lora_request, VllmLoRARequest):
+        # logger.warning(
+        #     "Received raw VllmLoRARequest. Performing temporary, on-the-fly "
+        #     "adaptation. This is a compatibility feature and should be removed "
+        #     "once the vLLM engine is updated to provide an adapted object.")
+        return VllmLoRARequestAdapter(lora_request)
+
+    if isinstance(lora_request, AbstractLoRARequest):
+        return lora_request
+
+    raise TypeError(f"Unsupported type for lora_request: {type(lora_request)}")

tpu_inference/worker/base.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from typing import Optional, Union
+
+import torch.nn as nn
+from vllm.lora.request import LoRARequest
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import ModelRunnerOutput
+
+from tpu_inference.di.abstracts import (AbstractKVCacheConfig,
+                                        AbstractKVCacheSpec,
+                                        AbstractLoRARequest,
+                                        AbstractSchedulerOutput)
+from tpu_inference.di.interfaces import HostInterface
+
+
+class AbstractTpuWorker(ABC):
+    """Base class for TPU workers.
+
+    This class defines a pure, host-agnostic contract for what a TPU worker
+    must be able to do. It is intentionally decoupled from any specific host
+    system like vLLM or SGLang.
+
+    Architectural Note on Dependencies:
+    This abstract class only depends on other abstractions (e.g., HostInterface).
+    It does NOT hold configuration objects from any specific host (e.g.,
+    VllmConfig). Doing so would create a "leaky abstraction," forcing all
+    future implementations to depend on a concrete detail from a single host.
+
+    The responsibility for managing concrete configuration is pushed down to the
+    concrete subclasses (e.g., TPUWorkerJax), which keeps this base class
+    pure and truly reusable across different host systems.
+    """
+
+    def __init__(self, host_interface: Optional[HostInterface] = None):
+        self.host_interface = host_interface
+
+    @abstractmethod
+    def initialize_cache(self, num_gpu_blocks: int,
+                         num_cpu_blocks: int) -> None:
+        """Initialize the cache with the given number of blocks."""
+        pass
+
+    @abstractmethod
+    def init_device(self):
+        """Initialize the TPU device and distributed environment."""
+        pass
+
+    @abstractmethod
+    def determine_available_memory(self) -> int:
+        """Determine available memory for the TPU worker."""
+        pass
+
+    @abstractmethod
+    def execute_model(
+        self,
+        scheduler_output: Union[AbstractSchedulerOutput, SchedulerOutput],
+    ) -> Optional[ModelRunnerOutput]:
+        pass
+
+    @abstractmethod
+    def profile(self, is_start: bool = True):
+        pass
+
+    @abstractmethod
+    def add_lora(
+        self,
+        lora_request: Union[AbstractLoRARequest, LoRARequest],
+    ) -> bool:
+        pass
+
+    @abstractmethod
+    def load_model(self) -> None:
+        pass
+
+    @abstractmethod
+    def compile_or_warm_up_model(self) -> None:
+        pass
+
+    @abstractmethod
+    def get_model(self) -> nn.Module:
+        pass
+
+    @abstractmethod
+    def get_kv_cache_spec(self) -> dict[str, AbstractKVCacheSpec]:
+        pass
+
+    @abstractmethod
+    def initialize_from_config(
+        self,
+        kv_cache_config: Union[AbstractKVCacheConfig, KVCacheConfig],
+    ) -> None:
+        """Allocate  KV cache with the specified kv_cache_config."""
+        pass
+
+    def check_health(self) -> None:
+        # worker will always be healthy as long as it's running.
+        return