tpu-inference 0.11.1__py3-none-any.whl

tests/test_tpu_info.py

@@ -0,0 +1,120 @@
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+import requests
+
+from tpu_inference.tpu_info import (get_node_name, get_node_worker_id,
+                                    get_num_chips, get_num_cores_per_chip,
+                                    get_tpu_metadata, get_tpu_type)
+
+
+# Mock requests.get for get_tpu_metadata tests
+@patch("tpu_inference.tpu_info.requests.get")
+def test_get_tpu_metadata_success(mock_get):
+    """Test get_tpu_metadata when the request is successful."""
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.text = "test_metadata_value"
+    mock_get.return_value = mock_response
+    assert get_tpu_metadata(key="test-key") == "test_metadata_value"
+
+
+@patch("tpu_inference.tpu_info.requests.get")
+def test_get_tpu_metadata_request_error(mock_get):
+    """Test get_tpu_metadata when a RequestException is raised."""
+    mock_get.side_effect = requests.RequestException("Test RequestException")
+    assert get_tpu_metadata(key="test-key") is None
+
+
+# Test get_tpu_type
+@patch("tpu_inference.tpu_info.get_tpu_metadata")
+@patch.dict(os.environ, {"TPU_ACCELERATOR_TYPE": "env_tpu_type"})
+def test_get_tpu_type_from_env(mock_get_tpu_metadata):
+    """Test get_tpu_type when TPU_ACCELERATOR_TYPE is set in environment."""
+    # The function should return the env var value and not call get_tpu_metadata
+    assert get_tpu_type() == "env_tpu_type"
+    mock_get_tpu_metadata.assert_not_called()
+
+
+@patch.dict(os.environ, {}, clear=True)
+@patch("tpu_inference.tpu_info.get_tpu_metadata",
+       return_value="metadata_tpu_type")
+def test_get_tpu_type_from_metadata(mock_get_tpu_metadata):
+    """Test get_tpu_type when environment variable is not set."""
+    assert get_tpu_type() == "metadata_tpu_type"
+    mock_get_tpu_metadata.assert_called_once_with(key="accelerator-type")
+
+
+# Test get_node_name
+@patch("tpu_inference.tpu_info.get_tpu_metadata")
+@patch.dict(os.environ, {"TPU_NAME": "env_tpu_name"})
+def test_get_node_name_from_env(mock_get_tpu_metadata):
+    """Test get_node_name when TPU_NAME is set in environment."""
+    assert get_node_name() == "env_tpu_name"
+    mock_get_tpu_metadata.assert_not_called()
+
+
+@patch.dict(os.environ, {}, clear=True)
+@patch("tpu_inference.tpu_info.get_tpu_metadata",
+       return_value="metadata_tpu_name")
+def test_get_node_name_from_metadata(mock_get_tpu_metadata):
+    """Test get_node_name when environment variable is not set."""
+    assert get_node_name() == "metadata_tpu_name"
+    mock_get_tpu_metadata.assert_called_once_with(key="instance-id")
+
+
+# Test get_node_worker_id
+@patch("tpu_inference.tpu_info.get_tpu_metadata")
+@patch.dict(os.environ, {"TPU_WORKER_ID": "5"})
+def test_get_node_worker_id_from_env(mock_get_tpu_metadata):
+    """Test get_node_worker_id when TPU_WORKER_ID is set in environment."""
+    assert get_node_worker_id() == 5
+    mock_get_tpu_metadata.assert_not_called()
+
+
+@patch.dict(os.environ, {}, clear=True)
+@patch("tpu_inference.tpu_info.get_tpu_metadata", return_value="10")
+def test_get_node_worker_id_from_metadata(mock_get_tpu_metadata):
+    """Test get_node_worker_id when environment variable is not set."""
+    assert get_node_worker_id() == 10
+    mock_get_tpu_metadata.assert_called_once_with(key="agent-worker-number")
+
+
+# Test get_num_cores_per_chip
+@pytest.mark.parametrize(
+    "tpu_type, expected",
+    [
+        ("v5litepod-4", 1),
+        ("v6e-8", 1),
+        ("v4-8", 2),
+        ("v5p-16", 2),
+        ("unknown-type", 2)  # Default case
+    ])
+@patch("tpu_inference.tpu_info.get_tpu_type")
+def test_get_num_cores_per_chip(mock_get_tpu_type, tpu_type, expected):
+    """Test get_num_cores_per_chip with different TPU types."""
+    mock_get_tpu_type.return_value = tpu_type
+    assert get_num_cores_per_chip() == expected
+
+
+# Test get_num_chips
+@patch("tpu_inference.tpu_info.glob.glob",
+       return_value=["/dev/accel0", "/dev/accel1"])
+def test_get_num_chips_from_accel(mock_glob):
+    """Test get_num_chips when /dev/accel* files exist."""
+    assert get_num_chips() == 2
+
+
+@patch("tpu_inference.tpu_info.glob.glob", return_value=[])
+@patch("tpu_inference.tpu_info.os.listdir", return_value=["0", "1", "2"])
+def test_get_num_chips_from_vfio(mock_listdir, mock_glob):
+    """Test get_num_chips when /dev/accel* files don't exist but /dev/vfio entries do."""
+    assert get_num_chips() == 3
+
+
+@patch("tpu_inference.tpu_info.glob.glob", return_value=[])
+@patch("tpu_inference.tpu_info.os.listdir", side_effect=FileNotFoundError)
+def test_get_num_chips_not_found(mock_listdir, mock_glob, caplog):
+    """Test get_num_chips when neither files nor directory are found."""
+    assert get_num_chips() == 0

tests/test_utils.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+from unittest.mock import MagicMock, patch
+
+import jax.numpy as jnp
+import pytest
+
+# Import the functions to be tested
+from tpu_inference.utils import (GBYTES, enable_megacore,
+                                 get_jax_dtype_from_str_dtype, get_megacore,
+                                 get_padded_head_dim, hbm_usage_bytes,
+                                 hbm_usage_gb, quantize_kv)
+
+
+def test_enable_and_get_megacore():
+    """Tests the enable_megacore and get_megacore functions."""
+    assert not get_megacore()
+    enable_megacore()
+    assert get_megacore()
+
+
+@patch.dict(os.environ, {"TPU_MULTIHOST_BACKEND": "ray"})
+def test_hbm_usage_bytes_ray_backend():
+    """Tests hbm_usage_bytes when TPU_MULTIHOST_BACKEND is ray."""
+    mock_device1 = MagicMock()
+    mock_device1.memory_stats.return_value = {
+        "bytes_in_use": 100 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+    mock_device2 = MagicMock()
+    mock_device2.memory_stats.side_effect = Exception("Memory stats failed")
+
+    devices = [mock_device1, mock_device2]
+    usage = hbm_usage_bytes(devices)
+
+    expected_usage = [(100 * GBYTES, 128 * GBYTES),
+                      (100 * GBYTES, 128 * GBYTES)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", False)
+def test_hbm_usage_bytes_pathways_disabled():
+    """Tests hbm_usage_bytes when VLLM_TPU_USING_PATHWAYS is False."""
+    mock_device1 = MagicMock()
+    mock_device1.memory_stats.return_value = {
+        "bytes_in_use": 100 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+    mock_device2 = MagicMock()
+    mock_device2.memory_stats.return_value = {
+        "bytes_in_use": 50 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+
+    devices = [mock_device1, mock_device2]
+    usage = hbm_usage_bytes(devices)
+
+    expected_usage = [(100 * GBYTES, 128 * GBYTES),
+                      (50 * GBYTES, 128 * GBYTES)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", True)
+@patch("jax.live_arrays")
+@patch("jax.devices")
+def test_hbm_usage_bytes_pathways_enabled(mock_devices, mock_live_arrays):
+    """Tests hbm_usage_bytes when VLLM_TPU_USING_PATHWAYS is True."""
+    # Mock TPU v5p devices
+    mock_jax_device = MagicMock()
+    mock_jax_device.device_kind = "TPU v5p"
+    mock_devices.return_value = [mock_jax_device]
+
+    # Create mock devices
+    mock_device1 = MagicMock()
+    mock_device2 = MagicMock()
+    devices = [mock_device1, mock_device2]
+
+    # Create mock arrays with sharding
+    mock_array1 = MagicMock()
+    mock_array1.dtype.itemsize = 4  # float32
+    mock_array1.size = 1000  # 1000 elements
+    mock_array1.sharding.device_set = {mock_device1, mock_device2
+                                       }  # Sharded across 2 devices
+
+    mock_array2 = MagicMock()
+    mock_array2.dtype.itemsize = 2  # float16
+    mock_array2.size = 500  # 500 elements
+    mock_array2.sharding.device_set = {mock_device1}  # Only on device1
+
+    mock_live_arrays.return_value = [mock_array1, mock_array2]
+
+    usage = hbm_usage_bytes(devices)
+
+    # Expected calculations:
+    # Array1: 4 bytes * 1000 elements / 2 devices = 2000 bytes per device
+    # Array2: 2 bytes * 500 elements / 1 device = 1000 bytes on device1 only
+    # Device1: 2000 + 1000 = 3000 bytes
+    # Device2: 2000 + 0 = 2000 bytes
+    # hbm_limit = 33550237184 (hardcoded in the function)
+    expected_usage = [(3000, 95 * GBYTES), (2000, 95 * GBYTES)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", False)
+def test_hbm_usage_gb_pathways_disabled():
+    """Tests hbm_usage_gb when VLLM_TPU_USING_PATHWAYS is False."""
+    mock_device1 = MagicMock()
+    mock_device1.memory_stats.return_value = {
+        "bytes_in_use": 100 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+    mock_device2 = MagicMock()
+    mock_device2.memory_stats.return_value = {
+        "bytes_in_use": 50.5 * GBYTES,
+        "bytes_limit": 128.0 * GBYTES
+    }
+
+    devices = [mock_device1, mock_device2]
+    usage = hbm_usage_gb(devices)
+
+    expected_usage = [(100.0, 128.0), (50.5, 128.0)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", True)
+@patch("jax.live_arrays")
+@patch("jax.devices")
+def test_hbm_usage_bytes_pathways_no_arrays(mock_devices, mock_live_arrays):
+    """Tests hbm_usage_bytes when VLLM_TPU_USING_PATHWAYS is True but no live arrays."""
+    # Mock TPU v5e devices
+    mock_jax_device = MagicMock()
+    mock_jax_device.device_kind = "TPU v6e"
+    mock_devices.return_value = [mock_jax_device]
+
+    mock_device1 = MagicMock()
+    mock_device2 = MagicMock()
+    devices = [mock_device1, mock_device2]
+
+    # No live arrays
+    mock_live_arrays.return_value = []
+
+    usage = hbm_usage_bytes(devices)
+
+    # No arrays means no memory usage
+    expected_usage = [(0, 32 * GBYTES), (0, 32 * GBYTES)]
+    assert usage == expected_usage
+
+
+@pytest.mark.parametrize(
+    "head_dim, expected_padded_head_dim",
+    [
+        (1, 128),
+        (64, 128),
+        (127, 128),
+        (128, 128),
+        (129, 256),
+        (255, 256),
+        (256, 256),
+        (0, 0),  # Although head_dim is usually positive, testing boundary
+    ],
+)
+def test_get_padded_head_dim(head_dim, expected_padded_head_dim):
+    """Tests the get_padded_head_dim function."""
+    assert get_padded_head_dim(head_dim) == expected_padded_head_dim
+
+
+def test_quantize_kv_float8_e4m3fn():
+    """Tests the quantize_kv function with float8_e4m3fn dtype."""
+    key = jnp.array([-1.0, 0.5, 1.0, 1.5])
+    value = jnp.array([2.0, 0.0, -2.0, -3.0])
+    kv_cache_quantized_dtype = jnp.float8_e4m3fn
+    k_scale = 0.1
+    v_scale = 0.2
+
+    quantized_key, quantized_value = quantize_kv(key, value,
+                                                 kv_cache_quantized_dtype,
+                                                 k_scale, v_scale)
+
+    # Expected key: key / k_scale -> clip -> astype
+    # [-10., 5., 10., 15.] are within float8_e4m3fn range
+    expected_key = jnp.array([-10.0, 5.0, 10.0, 15.0], dtype=jnp.float8_e4m3fn)
+
+    # Expected value: value / v_scale -> clip -> astype
+    # [10., 0., -10., -15.] are within float8_e4m3fn range
+    expected_value = jnp.array([10.0, 0.0, -10.0, -15.0],
+                               dtype=jnp.float8_e4m3fn)
+
+    assert jnp.array_equal(quantized_key, expected_key)
+    assert jnp.array_equal(quantized_value, expected_value)
+
+    # Test clipping
+    dtype_info = jnp.finfo(kv_cache_quantized_dtype)
+    minval, maxval = float(dtype_info.min), float(dtype_info.max)
+
+    # Values that will be outside the range after scaling
+    key_clip = jnp.array([minval * k_scale * 2, maxval * k_scale * 2])
+    value_clip = jnp.array([maxval * v_scale * 2, minval * v_scale * 2])
+    quantized_key_clip, quantized_value_clip = quantize_kv(
+        key_clip, value_clip, kv_cache_quantized_dtype, k_scale, v_scale)
+
+    # Values should be clipped to the min/max of the float8 dtype
+    expected_key_clip = jnp.array([minval, maxval], dtype=jnp.float8_e4m3fn)
+    expected_value_clip = jnp.array([maxval, minval], dtype=jnp.float8_e4m3fn)
+
+    assert jnp.array_equal(quantized_key_clip, expected_key_clip)
+    assert jnp.array_equal(quantized_value_clip, expected_value_clip)
+
+
+def test_get_jax_dtype_from_str_dtype():
+    """
+    Test the get_jax_dtype_from_str_dtype function
+    """
+    assert get_jax_dtype_from_str_dtype("int8") == jnp.int8
+    assert get_jax_dtype_from_str_dtype("bfloat16") == jnp.bfloat16
+    assert get_jax_dtype_from_str_dtype("fp8") == jnp.float8_e4m3fn
+    assert get_jax_dtype_from_str_dtype("fp8_e4m3") == jnp.float8_e4m3
+    assert get_jax_dtype_from_str_dtype("fp8_e5m2") == jnp.float8_e5m2
+    assert get_jax_dtype_from_str_dtype("auto") is None

tests/tpu_backend_test.py

@@ -0,0 +1,59 @@
+import unittest
+from unittest.mock import Mock, patch
+
+from tpu_inference.backend import TPUBackend
+
+
+class TPUBackendTest(unittest.TestCase):
+
+    @patch('tpu_inference.backend.TPUWorker')
+    def test_tpu_backend_initialization(self, mock_tpu_worker_class):
+        """Test that TPUBackend initializes the worker correctly."""
+        mock_host_interface = Mock()
+        mock_worker_kwargs = {'worker_arg': 'test_value'}
+
+        backend = TPUBackend(host_interface=mock_host_interface,
+                             **mock_worker_kwargs)
+
+        # Assert that the TPUWorker was instantiated with the correct arguments
+        mock_tpu_worker_class.assert_called_once_with(
+            host_interface=mock_host_interface, **mock_worker_kwargs)
+
+        # Assert that the worker attribute is an instance of the mock class
+        self.assertEqual(backend.worker, mock_tpu_worker_class.return_value)
+
+    @patch('tpu_inference.backend.VllmSchedulerOutputAdapter')
+    @patch('tpu_inference.backend.TPUWorker')
+    def test_launch_tpu_batch(self, mock_tpu_worker_class, mock_adapter_class):
+        """Test that launch_tpu_batch delegates to the worker correctly."""
+        mock_worker_instance = mock_tpu_worker_class.return_value
+
+        backend = TPUBackend()
+        mock_batch = Mock()
+
+        backend.launch_tpu_batch(mock_batch)
+
+        # Assert that the adapter was created with the correct input
+        mock_adapter_class.assert_called_once_with(mock_batch)
+
+        # Assert that the worker's execute_model method was called with the mock adapter's return value
+        mock_worker_instance.execute_model.assert_called_once_with(
+            mock_adapter_class.return_value)
+
+    @patch('tpu_inference.backend.VllmLoRARequestAdapter')
+    @patch('tpu_inference.backend.TPUWorker')
+    def test_add_lora(self, mock_tpu_worker_class, mock_adapter_class):
+        """Test that add_lora delegates to the worker correctly."""
+        mock_worker_instance = mock_tpu_worker_class.return_value
+
+        backend = TPUBackend()
+        mock_lora_request = Mock()
+
+        backend.add_lora(mock_lora_request)
+
+        # Assert that the adapter was created with the correct input
+        mock_adapter_class.assert_called_once_with(mock_lora_request)
+
+        # Assert that the worker's add_lora method was called with the mock adapter's return value
+        mock_worker_instance.add_lora.assert_called_once_with(
+            mock_adapter_class.return_value)

tpu_inference/__init__.py

@@ -0,0 +1,30 @@
+import os
+
+from tpu_inference import tpu_info as ti
+from tpu_inference.logger import init_logger
+
+logger = init_logger(__name__)
+
+if "proxy" in os.environ.get('JAX_PLATFORMS', '').lower():
+    logger.info("Running vLLM on TPU via Pathways proxy.")
+    # Must run pathwaysutils.initialize() before any JAX operations
+    try:
+        import pathwaysutils
+        pathwaysutils.initialize()
+        logger.info("Module pathwaysutils is imported.")
+    except Exception as e:
+        logger.error(
+            f"Error occurred while importing pathwaysutils or logging TPU info: {e}"
+        )
+else:
+    # Either running on TPU or CPU
+    try:
+        logger.info(f"TPU info: node_name={ti.get_node_name()} | "
+                    f"tpu_type={ti.get_tpu_type()} | "
+                    f"worker_id={ti.get_node_worker_id()} | "
+                    f"num_chips={ti.get_num_chips()} | "
+                    f"num_cores_per_chip={ti.get_num_cores_per_chip()}")
+    except Exception as e:
+        logger.error(
+            f"Error occurred while logging TPU info: {e}. Are you running on CPU?"
+        )

tpu_inference/adapters/vllm_adapters.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from vllm.lora.request import LoRARequest
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+from vllm.v1.outputs import ModelRunnerOutput
+
+from tpu_inference.di.abstracts import (AbstractKVCacheConfig,
+                                        AbstractKVCacheSpec,
+                                        AbstractLoRARequest,
+                                        AbstractModelRunnerOutput,
+                                        AbstractSchedulerOutput)
+
+
+class VllmModelRunnerOutputAdapter(AbstractModelRunnerOutput):
+
+    def __init__(self, vllm_output: ModelRunnerOutput):
+        self.vllm_output = vllm_output
+
+
+class VllmSchedulerOutputAdapter(AbstractSchedulerOutput):
+
+    def __init__(self, vllm_scheduler_output: SchedulerOutput):
+        self.vllm_scheduler_output = vllm_scheduler_output
+
+
+class VllmLoRARequestAdapter(AbstractLoRARequest):
+
+    def __init__(self, vllm_lora_request: LoRARequest):
+        self.vllm_lora_request = vllm_lora_request
+
+
+class VllmKVCacheConfigAdapter(AbstractKVCacheConfig):
+
+    def __init__(self, vllm_kv_cache_config: KVCacheConfig):
+        self.vllm_kv_cache_config = vllm_kv_cache_config
+
+
+class VllmKVCacheSpecAdapter(AbstractKVCacheSpec):
+
+    def __init__(self, vllm_kv_cache_spec: KVCacheSpec):
+        self.vllm_kv_cache_spec = vllm_kv_cache_spec

tpu_inference/adapters/vllm_config_adapters.py

@@ -0,0 +1,134 @@
+"""
+Adapters for wrapping concrete vLLM config objects in tpu_inference interfaces.
+"""
+from typing import Any, Optional
+
+import torch
+
+from tpu_inference.interfaces.config_parts import (ICacheConfig,
+                                                   ICompilationConfig,
+                                                   IModelConfig,
+                                                   IParallelConfig,
+                                                   ISchedulerConfig)
+
+
+class VllmCacheConfigAdapter(ICacheConfig):
+
+    def __init__(self, vllm_cache_config: Any):
+        self._vllm_cache_config = vllm_cache_config
+
+    @property
+    def block_size(self) -> Optional[int]:
+        return self._vllm_cache_config.block_size
+
+    @block_size.setter
+    def block_size(self, value: Optional[int]) -> None:
+        self._vllm_cache_config.block_size = value
+
+
+class VllmCompilationConfigAdapter(ICompilationConfig):
+
+    def __init__(self, vllm_compilation_config: Any):
+        self._vllm_compilation_config = vllm_compilation_config
+
+    @property
+    def level(self) -> Any:
+        return self._vllm_compilation_config.level
+
+    @level.setter
+    def level(self, value: Any) -> None:
+        self._vllm_compilation_config.level = value
+
+    @property
+    def backend(self) -> str:
+        return self._vllm_compilation_config.backend
+
+    @backend.setter
+    def backend(self, value: str) -> None:
+        self._vllm_compilation_config.backend = value
+
+
+class VllmModelConfigAdapter(IModelConfig):
+
+    def __init__(self, vllm_model_config: Any):
+        self._vllm_model_config = vllm_model_config
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self._vllm_model_config.dtype
+
+    @dtype.setter
+    def dtype(self, value: torch.dtype) -> None:
+        self._vllm_model_config.dtype = value
+
+    @property
+    def use_mla(self) -> bool:
+        return self._vllm_model_config.use_mla
+
+
+class VllmParallelConfigAdapter(IParallelConfig):
+
+    def __init__(self, vllm_parallel_config: Any):
+        self._vllm_parallel_config = vllm_parallel_config
+
+    @property
+    def worker_cls(self) -> str:
+        return self._vllm_parallel_config.worker_cls
+
+    @worker_cls.setter
+    def worker_cls(self, value: str) -> None:
+        self._vllm_parallel_config.worker_cls = value
+
+
+class VllmSchedulerConfigAdapter(ISchedulerConfig):
+
+    def __init__(self, vllm_scheduler_config: Any):
+        self._vllm_scheduler_config = vllm_scheduler_config
+
+    @property
+    def max_num_seqs(self) -> int:
+        return self._vllm_scheduler_config.max_num_seqs
+
+    @property
+    def is_multi_step(self) -> bool:
+        return self._vllm_scheduler_config.is_multi_step
+
+    @property
+    def is_multimodal_model(self) -> bool:
+        return self._vllm_scheduler_config.is_multimodal_model
+
+    @property
+    def disable_chunked_mm_input(self) -> bool:
+        return self._vllm_scheduler_config.disable_chunked_mm_input
+
+    @disable_chunked_mm_input.setter
+    def disable_chunked_mm_input(self, value: bool) -> None:
+        self._vllm_scheduler_config.disable_chunked_mm_input = value
+
+    @property
+    def enable_chunked_prefill(self) -> bool:
+        return self._vllm_scheduler_config.enable_chunked_prefill
+
+    @enable_chunked_prefill.setter
+    def enable_chunked_prefill(self, value: bool) -> None:
+        self._vllm_scheduler_config.enable_chunked_prefill = value
+
+    @property
+    def chunked_prefill_enabled(self) -> bool:
+        return self._vllm_scheduler_config.chunked_prefill_enabled
+
+    @chunked_prefill_enabled.setter
+    def chunked_prefill_enabled(self, value: bool) -> None:
+        self._vllm_scheduler_config.chunked_prefill_enabled = value
+
+    @property
+    def max_model_len(self) -> int:
+        return self._vllm_scheduler_config.max_model_len
+
+    @property
+    def max_num_batched_tokens(self) -> int:
+        return self._vllm_scheduler_config.max_num_batched_tokens
+
+    @max_num_batched_tokens.setter
+    def max_num_batched_tokens(self, value: int) -> None:
+        self._vllm_scheduler_config.max_num_batched_tokens = value

tpu_inference/backend.py

@@ -0,0 +1,69 @@
+"""
+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
+
+from tpu_inference.adapters.vllm_adapters import (VllmLoRARequestAdapter,
+                                                  VllmSchedulerOutputAdapter)
+from tpu_inference.di.interfaces import BackendInterface, HostInterface
+from tpu_inference.worker.base import AbstractTpuWorker
+from tpu_inference.worker.tpu_worker_jax import TPUWorker
+
+
+class TPUBackend(BackendInterface):
+    """
+  The main entry point for the host system to interact with the TPU backend.
+
+  This class implements the BackendInterface. It is responsible for creating
+  and managing the concrete TPU worker instance and delegating calls to it.
+  """
+
+    def __init__(self,
+                 host_interface: Optional[HostInterface] = None,
+                 **worker_kwargs):
+        """
+        Initializes the TPUBackend.
+
+        Args:
+            host_interface: An optional object that implements the HostInterface,
+                            providing a way for the backend to communicate with the host.
+            **worker_kwargs: Additional keyword arguments to be passed to the
+                            worker's constructor.
+        """
+        self.worker: AbstractTpuWorker = TPUWorker(
+            host_interface=host_interface, **worker_kwargs)
+
+    def launch_tpu_batch(self, batch_to_launch):
+        """
+        Launches a batch of requests on the TPU worker and returns the result.
+
+        Args:
+            batch_to_launch: The batch of requests to be processed.
+
+        Returns:
+            The result of the model execution.
+        """
+        adapted_batch = VllmSchedulerOutputAdapter(batch_to_launch)
+        return self.worker.execute_model(adapted_batch)
+
+    def add_lora(self, lora_request):
+        """
+        Adds a LoRA adapter to the worker.
+
+        Args:
+            lora_request: The LoRA request to be processed.
+        """
+        adapted_lora_request = VllmLoRARequestAdapter(lora_request)
+        return self.worker.add_lora(adapted_lora_request)