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

tests/e2e/test_speculative_decoding.py

@@ -0,0 +1,291 @@
+# 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 __future__ import annotations
+
+import random
+import string
+import time
+
+import pytest
+from vllm import LLM, SamplingParams
+
+
+def get_ngram_test_prompts():
+    num_prompts = 100
+    prompts = []
+
+    for _ in range(num_prompts):
+        w = random.choice(list(string.ascii_lowercase))
+        prompts.append(
+            f"Keep repeating: {w} {w} {w} {w} {w} {w} {w} {w} {w} {w}")
+
+    return prompts
+
+
+def get_eagle3_test_prompts():
+    num_prompts = 100
+    prompts = []
+
+    for _ in range(num_prompts):
+        prompts.append(
+            "Predict the continuation of this sequence: 1 2 3 4 5 6 7 8")
+
+    return prompts
+
+
+def get_test_prompts(speculative_config: dict):
+    if speculative_config['method'] == 'ngram':
+        return get_ngram_test_prompts()
+    elif speculative_config['method'] == 'eagle3':
+        return get_eagle3_test_prompts()
+    else:
+        raise NotImplementedError(
+            f"{speculative_config['method']} is not supported yet.")
+
+
+@pytest.fixture
+def sampling_config():
+    return SamplingParams(temperature=0,
+                          max_tokens=32,
+                          ignore_eos=True,
+                          repetition_penalty=1,
+                          frequency_penalty=0,
+                          presence_penalty=0,
+                          min_p=0,
+                          logprobs=None)
+
+
+@pytest.fixture
+def model_name():
+    return "Qwen/Qwen2.5-0.5B-Instruct"
+
+
+# TODO(pooyam): run vLLM engine with InProcClient (`VLLM_ENABLE_V1_MULTIPROCESSING = 0`) mode to avoid TPU contention among processes.
+def _test_correctness_helper(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+    model_name: str,
+    speculative_config: dict,
+):
+    '''
+    Helper function to test ngram correctness.
+    Compare the outputs of a original LLM and a speculative LLM
+    should be the same when using ngram speculative decoding.
+    '''
+    with monkeypatch.context():
+        test_prompts = get_test_prompts(speculative_config)
+
+        ref_llm = LLM(model=model_name, max_model_len=1024, max_num_seqs=4)
+        ref_outputs = ref_llm.generate(test_prompts, sampling_config)
+
+        del ref_llm
+
+        # Waiting for TPUs to be released.
+        time.sleep(10)
+
+        spec_llm = LLM(model=model_name,
+                       speculative_config=speculative_config,
+                       max_model_len=1024,
+                       max_num_seqs=4)
+        spec_outputs = spec_llm.generate(test_prompts, sampling_config)
+
+        matches = 0
+        misses = 0
+        for ref_output, spec_output in zip(ref_outputs, spec_outputs):
+            if ref_output.outputs[0].text == spec_output.outputs[0].text:
+                matches += 1
+            else:
+                misses += 1
+                print(f"ref_output: {ref_output.outputs[0].text}")
+                print(f"spec_output: {spec_output.outputs[0].text}")
+
+        assert misses == 0
+        del spec_llm
+
+        # Waiting for TPUs to be released.
+        time.sleep(10)
+
+
+def test_ngram_correctness_greedy(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+    model_name: str,
+):
+    '''
+    Compare the outputs of a original LLM and a speculative LLM
+    should be the same when using ngram speculative decoding with greedy sampling.
+    '''
+    _test_correctness_helper(
+        monkeypatch, sampling_config, model_name, {
+            "method": "ngram",
+            "prompt_lookup_max": 5,
+            "prompt_lookup_min": 3,
+            "num_speculative_tokens": 3,
+        })
+
+
+def test_ngram_correctness_random(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+    model_name: str,
+):
+    '''
+    Compare the outputs of a original LLM and a speculative LLM
+    should be the same when using ngram speculative decoding with random sampling.
+    '''
+    # Modify sampling config for random sampling
+    sampling_config.temperature = 0.01
+    sampling_config.top_p = 0.9
+    sampling_config.top_k = 5
+
+    _test_correctness_helper(
+        monkeypatch, sampling_config, model_name, {
+            "method": "ngram",
+            "prompt_lookup_max": 5,
+            "prompt_lookup_min": 3,
+            "num_speculative_tokens": 3,
+        })
+
+
+def _test_performance_helper(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+    speculative_config: dict,
+    min_speedup: float,
+):
+    '''
+    Helper function to test speculative decoding performance.
+    Compares timing between reference LLM and speculative LLM using Llama 3 8B.
+    '''
+    model_name = "meta-llama/Llama-3.1-8B-Instruct"
+
+    with monkeypatch.context():
+        # Use a smaller set of prompts for performance testing
+        test_prompts = get_test_prompts(speculative_config)
+
+        # Test reference LLM timing
+        ref_llm = LLM(model=model_name,
+                      max_model_len=1024,
+                      max_num_seqs=1,
+                      enable_prefix_caching=False)
+
+        start_time = time.time()
+        _ = ref_llm.generate(test_prompts, sampling_config)
+        ref_time = time.time() - start_time
+
+        del ref_llm
+
+        # Waiting for TPUs to be released
+        time.sleep(10)
+
+        # Test speculative LLM timing with max_num_seqs=1
+        spec_llm = LLM(model=model_name,
+                       speculative_config=speculative_config,
+                       max_model_len=1024,
+                       max_num_seqs=1,
+                       enable_prefix_caching=False)
+
+        start_time = time.time()
+        _ = spec_llm.generate(test_prompts, sampling_config)
+        spec_time = time.time() - start_time
+
+        del spec_llm
+        # Waiting for TPUs to be released
+        time.sleep(10)
+
+        speedup = ref_time / spec_time
+        print(f"Reference LLM time: {ref_time:.2f}s")
+        print(f"Speculative LLM time: {spec_time:.2f}s")
+        print(f"Speedup: {speedup:.2f}x")
+
+        # TODO(pooyam): Make this tighter once we have better performance.
+        assert speedup >= min_speedup, f"Expected at least {min_speedup}x speedup for {speculative_config['method']}, got {speedup:.2f}x"
+
+
+def test_ngram_performance_greedy(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+):
+    '''
+    Test that speculative decoding provides significant performance improvement.
+    Compares timing between reference LLM and speculative LLM using Llama 3 8B.
+    Expects spec_llm to be at least 3.x faster than ref_llm.
+    '''
+    _test_performance_helper(
+        monkeypatch, sampling_config, {
+            "method": "ngram",
+            "prompt_lookup_max": 2,
+            "prompt_lookup_min": 2,
+            "num_speculative_tokens": 4,
+        }, 3.0)
+
+
+def test_ngram_performance_random(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+):
+    '''
+    Test that speculative decoding provides significant performance improvement.
+    Compares timing between reference LLM and speculative LLM using Llama 3 8B.
+    Expects spec_llm to be at least 3.x faster than ref_llm.
+    '''
+    sampling_config.temperature = 0.01
+    sampling_config.top_p = 0.9
+    sampling_config.top_k = 5
+
+    _test_performance_helper(
+        monkeypatch, sampling_config, {
+            "method": "ngram",
+            "prompt_lookup_max": 2,
+            "prompt_lookup_min": 2,
+            "num_speculative_tokens": 4,
+        }, 3.0)
+
+
+def test_eagle3_correctness(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+):
+    '''
+    Compare the outputs of a original LLM and a speculative LLM
+    should be the same when using eagle-3 speculative decoding.
+    '''
+    model_name = 'meta-llama/Meta-Llama-3-8B-Instruct'
+
+    _test_correctness_helper(
+        monkeypatch, sampling_config, model_name, {
+            'model': "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B",
+            "num_speculative_tokens": 3,
+            "method": "eagle3",
+            "draft_tensor_parallel_size": 1
+        })
+
+
+def test_eagle3_performance(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+):
+    '''
+    Test that speculative decoding provides significant performance improvement.
+    Compares timing between reference LLM and speculative LLM using Llama 3 8B.
+    Expects spec_llm to be at least 1.8 faster than ref_llm.
+    '''
+    _test_performance_helper(
+        monkeypatch, sampling_config, {
+            "method": "eagle3",
+            "model": "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B",
+            "num_speculative_tokens": 2,
+            "draft_tensor_parallel_size": 1
+        }, 1.8)

tests/e2e/test_structured_decoding.py

@@ -0,0 +1,46 @@
+# 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.
+
+# This file contains end-to-end tests for structured decoding.
+#
+# Structured decoding allows constraining the model's output to follow a
+# specific format, such as choosing from a predefined set of options or
+# following a JSON schema. This is useful for classification tasks,
+# structured data extraction, and ensuring outputs conform to expected formats.
+
+# The tests in this file verify that:
+# 1. Choice-based structured decoding correctly constrains output to valid options
+# 2. The model produces deterministic results when given structured constraints
+
+from __future__ import annotations
+
+from vllm import LLM, SamplingParams
+from vllm.sampling_params import StructuredOutputsParams
+
+
+def test_structured_decoding():
+    llm = LLM(model='meta-llama/Llama-3.2-1B-Instruct',
+              max_model_len=1024,
+              max_num_seqs=1,
+              enable_prefix_caching=False)
+
+    choices = ['Positive', 'Negative']
+    structured_outputs_params = StructuredOutputsParams(choice=choices)
+    sampling_params = SamplingParams(
+        structured_outputs=structured_outputs_params)
+    outputs = llm.generate(
+        prompts="Classify this sentiment: tpu-inference is wonderful!",
+        sampling_params=sampling_params,
+    )
+    assert outputs[0].outputs[0].text in choices

tests/executors/__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.

tests/executors/test_ray_distributed_executor.py

@@ -0,0 +1,199 @@
+# 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 unittest
+from unittest.mock import MagicMock, patch
+
+
+# Mock VllmConfig and its nested configs to avoid dependencies on the actual
+# classes, which can be complex to instantiate for testing.
+class MockVllmConfig:
+
+    def __init__(self):
+        self.parallel_config = MagicMock()
+        self.parallel_config.world_size = 4
+        self.parallel_config.tensor_parallel_size = 2
+        self.parallel_config.pipeline_parallel_size = 1
+        self.parallel_config.ray_workers_use_nsight = False
+        self.parallel_config.placement_group = None
+        self.parallel_config.max_parallel_loading_workers = None
+
+        self.sharding_config = MagicMock()
+        self.sharding_config.total_devices = 2
+
+        self.model_config = MagicMock()
+        self.cache_config = MagicMock()
+        self.lora_config = MagicMock()
+        self.load_config = MagicMock()
+        self.scheduler_config = MagicMock()
+        self.speculative_config = MagicMock()
+        self.prompt_adapter_config = MagicMock()
+        self.observability_config = MagicMock()
+        self.device_config = MagicMock()
+        self.ec_transfer_config = MagicMock()
+
+
+@patch(
+    "vllm.v1.executor.ray_distributed_executor.RayDistributedExecutor.__init__",
+    lambda x, y: None)
+@patch("tpu_inference.executors.ray_distributed_executor.envs")
+@patch("tpu_inference.executors.ray_distributed_executor.ray")
+@patch("tpu_inference.executors.ray_distributed_executor.current_platform")
+@patch("tpu_inference.executors.ray_distributed_executor.get_ip",
+       return_value="127.0.0.1")
+@patch("tpu_inference.executors.ray_distributed_executor.get_open_port",
+       return_value=12345)
+@patch(
+    "tpu_inference.executors.ray_distributed_executor.available_resources_per_node"
+)
+@patch("tpu_inference.executors.ray_distributed_executor._wait_until_pg_ready")
+class TestTpuRayDistributedExecutor(unittest.TestCase):
+
+    def setUp(self):
+        # Import the class under test inside the test method to ensure
+        # patches are applied.
+        from tpu_inference.executors.ray_distributed_executor import \
+            RayDistributedExecutor
+        self.RayDistributedExecutor = RayDistributedExecutor
+
+        self.vllm_config = MockVllmConfig()
+        # Reset placement group for each test as it might be modified.
+        self.vllm_config.parallel_config.placement_group = None
+        self.vllm_config.kv_transfer_config = None
+
+    def test_init_executor_basic_flow(self, mock_wait_until_pg_ready,
+                                      mock_avail_resources, mock_get_port,
+                                      mock_get_ip, mock_platform, mock_ray,
+                                      mock_envs):
+        # --- Setup mocks ---
+        mock_envs.VLLM_USE_RAY_COMPILED_DAG = True
+        mock_envs.VLLM_USE_RAY_SPMD_WORKER = True
+        mock_envs.VLLM_RAY_BUNDLE_INDICES = ""
+
+        mock_platform.ray_device_key = "TPU"
+        mock_platform.device_name = "tpu"
+        mock_platform.device_control_env_var = "TPU_VISIBLE_CHIPS"
+        mock_platform.additional_env_vars = []
+
+        mock_ray.is_initialized.return_value = False
+        mock_ray.nodes.return_value = [{"Resources": {"TPU": 4}}]
+        mock_ray.get_runtime_context.return_value.get_node_id.return_value = "node_1"
+        mock_avail_resources.return_value = {"node_1": {"TPU": 4}}
+
+        mock_wait_until_pg_ready.return_value = None
+
+        mock_placement_group = MagicMock()
+        mock_placement_group.bundle_specs = [{"TPU": 1}] * 4
+        mock_ray.util.placement_group.return_value = mock_placement_group
+
+        mock_worker = MagicMock()
+        mock_worker.get_node_and_gpu_ids.remote.return_value = [("node_1",
+                                                                 [0, 1, 2, 3])]
+        mock_ray.remote.return_value.remote.return_value = mock_worker
+
+        # Simulate remote calls on the worker
+        mock_ray.get.side_effect = [
+            ["127.0.0.1"] * 4,  # worker_ips
+            *[("node_1", [i]) for i in range(4)]  # worker_node_and_tpu_ids
+        ]
+
+        executor = self.RayDistributedExecutor(self.vllm_config)
+        # Members of the parent class
+        executor.uses_ray = True
+        executor.vllm_config = self.vllm_config
+        executor.parallel_config = self.vllm_config.parallel_config
+        executor.collective_rpc = MagicMock()
+        executor.collective_rpc.return_value = None
+
+        # --- Initialization ---
+        executor._init_executor()
+
+        # --- Assertions ---
+        mock_ray.init.assert_called_once()
+        self.assertIsNotNone(executor.parallel_config.placement_group)
+        self.assertEqual(len(executor.workers), 4)
+
+    def test_initialize_ray_cluster_no_tpu_on_driver_raises_error(
+            self, mock_wait_until_pg_ready, mock_avail_resources,
+            mock_get_port, mock_get_ip, mock_platform, mock_ray, mock_envs):
+        # --- Setup Mocks ---
+        mock_platform.ray_device_key = "TPU"
+        mock_platform.device_name = "tpu"
+
+        mock_ray.is_initialized.return_value = False
+        mock_ray.nodes.return_value = [{"Resources": {"TPU": 4}}]
+        mock_ray.get_runtime_context.return_value.get_node_id.return_value = "driver_node"
+        # Simulate no TPUs on the driver node
+        mock_avail_resources.return_value = {
+            "driver_node": {
+                "CPU": 8
+            },
+            "worker_node": {
+                "TPU": 4
+            }
+        }
+
+        executor = self.RayDistributedExecutor(self.vllm_config)
+        executor.vllm_config = self.vllm_config
+        executor.parallel_config = self.vllm_config.parallel_config
+
+        # --- Test and Assert ---
+        with self.assertRaisesRegex(ValueError,
+                                    "Current node has no TPU available"):
+            executor._initialize_ray_cluster()
+
+    def test_init_workers_ray_sorts_correctly(self, mock_wait_until_pg_ready,
+                                              mock_avail_resources,
+                                              mock_get_port, mock_get_ip,
+                                              mock_platform, mock_ray,
+                                              mock_envs):
+        # --- Setup Mocks ---
+        mock_envs.VLLM_RAY_BUNDLE_INDICES = ""
+        mock_platform.ray_device_key = "TPU"
+        mock_get_ip.return_value = "10.0.0.1"  # Driver IP
+
+        mock_pg = MagicMock()
+        mock_pg.bundle_specs = [{"TPU": 1}] * 4
+
+        mock_workers = [MagicMock() for _ in range(4)]
+        mock_ray.remote.return_value.return_value.remote.side_effect = mock_workers
+
+        # Simulate IPs for workers created with ranks 0, 1, 2, 3
+        worker_ips = ["10.0.0.2", "10.0.0.3", "10.0.0.1", "10.0.0.4"]
+        mock_ray.get.side_effect = [
+            worker_ips,  # worker_ips
+            *[('node_1', ['0', '1', '2', '3']),
+              ('node_2', ['4', '5', '6', '7']),
+              ('node_3', ['8', '9', '10', '11']),
+              ('node_4', ['12', '13', '14', '15'])]  # worker_node_and_tpu_ids
+        ]
+
+        executor = self.RayDistributedExecutor(self.vllm_config)
+        executor.use_ray_spmd_worker = True
+        executor.parallel_config = self.vllm_config.parallel_config
+        executor.vllm_config = self.vllm_config
+        executor.parallel_config.ray_workers_use_nsight = False
+        executor.collective_rpc = MagicMock()
+        executor.collective_rpc.return_value = None
+
+        # --- Call method under test ---
+        executor._init_workers_ray(mock_pg)
+
+        # --- Assertions ---
+        # Expected sorted order of workers: driver, then by IP
+        # Original workers: 0 (10.0.0.2), 1 (10.0.0.3), 2 (10.0.0.1), 3 (10.0.0.2)
+        # Sorted workers: 2 (driver), 0, 3 (same IP), 1
+        self.assertEqual(executor.workers, [
+            mock_workers[2], mock_workers[0], mock_workers[1], mock_workers[3]
+        ])

tests/experimental/__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.