tpu-inference 0.11.1.dev202512030818__py3-none-any.whl → 0.13.2rc3__py3-none-any.whl

tests/e2e/test_runai_model_streamer_loader.py

@@ -0,0 +1,104 @@
+# 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 the RunAI Model Streamer loader.
+#
+# The RunAI Model Streamer is a high-performance model loader that serves as an
+# alternative to the default Hugging Face loader. Instead of downloading a model
+# to local disk, it streams the weights from object storage (like GCS) into
+# GPU memory. This streaming process is significantly faster than the
+# traditional disk-based loading method.
+
+# The tests in this file verify that loading model weights using the
+# streamer produces the same results as loading the same model using the
+# standard Hugging Face loader. This ensures the correctness of the streamer
+# integration.
+
+# The tests are performed by:
+# 1. Loading a model from Google Cloud Storage using the `runai_streamer` format.
+# 2. Generating output with this model.
+# 3. Loading the same model from Hugging Face using the default loader.
+# 4. Generating output with this second model.
+# 5. Asserting that the outputs from both models are identical.
+
+from __future__ import annotations
+
+import time
+
+import pytest
+from vllm import LLM, SamplingParams
+
+
+@pytest.fixture
+def sampling_config():
+    return SamplingParams(temperature=0, max_tokens=10, ignore_eos=True)
+
+
+@pytest.fixture
+# TODO(amacaskill): Replace with GKE owned GCS bucket.
+def gcs_model_name():
+    return "gs://vertex-model-garden-public-us/llama3/llama3-8b-hf"
+
+
+@pytest.fixture
+def hf_model_name():
+    return "meta-llama/Meta-Llama-3-8B"
+
+
+@pytest.fixture
+def prompt():
+    return "Hello, my name is"
+
+
+def test_correctness(
+    sampling_config: SamplingParams,
+    gcs_model_name: str,
+    hf_model_name: str,
+    prompt: str,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    '''
+    Compare the outputs of a model loaded from GCS via runai_model_streamer
+    and a model loaded from Hugging Face. The outputs should be the same.
+    These tests attempt to use tensor_parallel_size=1. The model is 16GB,
+    # and v6e has 32GB of HBM, so it will fit.
+    '''
+    # Set ENV variables so that runai_model_streamer uses anonymous GCS access.
+    monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "fake-project")
+    monkeypatch.setenv("RUNAI_STREAMER_GCS_USE_ANONYMOUS_CREDENTIALS", "true")
+    monkeypatch.setenv("CLOUD_STORAGE_EMULATOR_ENDPOINT",
+                       "https://storage.googleapis.com")
+    gcs_llm = LLM(model=gcs_model_name,
+                  load_format="runai_streamer",
+                  max_model_len=128,
+                  max_num_seqs=16,
+                  max_num_batched_tokens=256)
+    gcs_outputs = gcs_llm.generate([prompt], sampling_config)
+    gcs_output_text = gcs_outputs[0].outputs[0].text
+    del gcs_llm
+    time.sleep(10)  # Wait for TPUs to be released
+
+    # Test with Hugging Face model
+    hf_llm = LLM(model=hf_model_name,
+                 max_model_len=128,
+                 max_num_seqs=16,
+                 max_num_batched_tokens=256)
+    hf_outputs = hf_llm.generate([prompt], sampling_config)
+    hf_output_text = hf_outputs[0].outputs[0].text
+    del hf_llm
+    time.sleep(10)  # Wait for TPUs to be released
+
+    assert gcs_output_text == hf_output_text, (
+        f"Outputs do not match! "
+        f"GCS output: {gcs_output_text}, HF output: {hf_output_text}")

tests/e2e/test_sampling_params.py

@@ -0,0 +1,269 @@
+# 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 sampling parameters.
+#
+# Sampling parameters control how the model selects tokens during generation.
+# These tests verify that temperature, top_p, top_k, and logprobs work correctly.
+#
+# The tests in this file verify that:
+# 1. Temperature=0 produces deterministic (greedy) outputs
+# 2. Higher temperature produces more varied outputs
+# 3. top_p (nucleus sampling) correctly constrains token selection
+# 4. top_k correctly limits the number of candidate tokens
+# 5. logprobs returns probability information for generated tokens
+
+from __future__ import annotations
+
+import pytest
+from vllm import LLM, SamplingParams
+
+
+@pytest.fixture(scope="module")
+def llm():
+    """Create a shared LLM instance for all tests in this module."""
+    return LLM(
+        model='meta-llama/Llama-3.2-1B-Instruct',
+        max_model_len=1024,
+        max_num_seqs=4,
+        enable_prefix_caching=False,
+    )
+
+
+class TestTemperature:
+    """Tests for temperature sampling parameter."""
+
+    def test_temperature_zero_is_deterministic(self, llm: LLM):
+        """Temperature=0 should produce identical outputs across multiple runs."""
+        prompt = "What is 2 + 2? Answer with just the number:"
+        sampling_params = SamplingParams(temperature=0, max_tokens=10)
+
+        outputs1 = llm.generate([prompt], sampling_params)
+        outputs2 = llm.generate([prompt], sampling_params)
+
+        assert outputs1[0].outputs[0].text == outputs2[0].outputs[0].text
+
+    def test_high_temperature_produces_variation(self, llm: LLM):
+        """High temperature should produce varied outputs across multiple runs."""
+        prompt = "Write a random word:"
+        sampling_params = SamplingParams(temperature=2,
+                                         max_tokens=10,
+                                         top_k=4096)
+
+        # Run multiple times and collect unique outputs
+        unique_outputs = set()
+        num_runs = 10
+        for _ in range(num_runs):
+            outputs = llm.generate([prompt], sampling_params)
+            unique_outputs.add(outputs[0].outputs[0].text)
+
+        # With high temperature, we expect some variation
+        assert len(unique_outputs) > 1, (
+            "High temperature should produce varied outputs")
+
+
+class TestTopP:
+    """Tests for top_p (nucleus sampling) parameter."""
+
+    def test_top_p_restricts_sampling(self, llm: LLM):
+        """top_p=1.0 vs lower values should affect output diversity."""
+        prompt = "Name a color:"
+
+        # With top_p=1.0 (consider all tokens)
+        sampling_params_full = SamplingParams(temperature=0.8,
+                                              top_p=1.0,
+                                              max_tokens=5)
+
+        # With top_p=0.1 (very restrictive, only top tokens)
+        sampling_params_restricted = SamplingParams(temperature=0.8,
+                                                    top_p=0.1,
+                                                    max_tokens=5)
+
+        # Collect outputs with full nucleus
+        full_outputs = set()
+        for _ in range(10):
+            outputs = llm.generate([prompt], sampling_params_full)
+            full_outputs.add(outputs[0].outputs[0].text)
+
+        # Collect outputs with restricted nucleus
+        restricted_outputs = set()
+        for _ in range(10):
+            outputs = llm.generate([prompt], sampling_params_restricted)
+            restricted_outputs.add(outputs[0].outputs[0].text)
+
+        # Restricted top_p should generally produce less variety
+        # (though this isn't guaranteed, it's a reasonable expectation)
+        assert len(
+            restricted_outputs) >= 1, "Should produce at least one output"
+        assert len(full_outputs) >= 1, "Should produce at least one output"
+
+    def test_top_p_with_temperature_zero(self, llm: LLM):
+        """top_p should have no effect when temperature=0 (greedy)."""
+        prompt = "The capital of France is"
+
+        sampling_params_1 = SamplingParams(temperature=0,
+                                           top_p=0.1,
+                                           max_tokens=10)
+        sampling_params_2 = SamplingParams(temperature=0,
+                                           top_p=0.9,
+                                           max_tokens=10)
+
+        outputs1 = llm.generate([prompt], sampling_params_1)
+        outputs2 = llm.generate([prompt], sampling_params_2)
+
+        # Both should produce identical outputs since temperature=0
+        assert outputs1[0].outputs[0].text == outputs2[0].outputs[0].text
+
+
+class TestTopK:
+    """Tests for top_k sampling parameter."""
+
+    def test_top_k_restricts_sampling(self, llm: LLM):
+        """top_k should limit the candidate tokens for sampling."""
+        prompt = "Pick a number between 1 and 10:"
+
+        # top_k=1 is equivalent to greedy (always pick the most likely)
+        sampling_params_k1 = SamplingParams(temperature=1.0,
+                                            top_k=1,
+                                            max_tokens=5)
+
+        # top_k=-1 considers all tokens
+        sampling_params_all = SamplingParams(temperature=1.0,
+                                             top_k=-1,
+                                             max_tokens=5)
+
+        # With top_k=1, outputs should be deterministic
+        outputs_k1_run1 = llm.generate([prompt], sampling_params_k1)
+        outputs_k1_run2 = llm.generate([prompt], sampling_params_k1)
+        assert outputs_k1_run1[0].outputs[0].text == outputs_k1_run2[
+            0].outputs[0].text
+
+        # With top_k=-1 and temperature=1.0, we may see variation
+        all_outputs = set()
+        for _ in range(10):
+            outputs = llm.generate([prompt], sampling_params_all)
+            all_outputs.add(outputs[0].outputs[0].text)
+
+        # Should produce at least one valid output
+        assert len(all_outputs) >= 1
+
+    def test_top_k_with_temperature_zero(self, llm: LLM):
+        """top_k should have no effect when temperature=0 (greedy)."""
+        prompt = "The largest planet is"
+
+        sampling_params_k5 = SamplingParams(temperature=0,
+                                            top_k=5,
+                                            max_tokens=10)
+        sampling_params_k50 = SamplingParams(temperature=0,
+                                             top_k=50,
+                                             max_tokens=10)
+
+        outputs1 = llm.generate([prompt], sampling_params_k5)
+        outputs2 = llm.generate([prompt], sampling_params_k50)
+
+        # Both should produce identical outputs since temperature=0
+        assert outputs1[0].outputs[0].text == outputs2[0].outputs[0].text
+
+
+class TestLogprobs:
+    """Tests for logprobs parameter."""
+
+    def test_logprobs_returns_probabilities(self, llm: LLM):
+        """logprobs parameter should return log probabilities for tokens."""
+        prompt = "Hello"
+        sampling_params = SamplingParams(temperature=0,
+                                         max_tokens=5,
+                                         logprobs=5)
+
+        outputs = llm.generate([prompt], sampling_params)
+        output = outputs[0].outputs[0]
+
+        # Check that logprobs are returned
+        assert output.logprobs is not None, "logprobs should be returned"
+        assert len(output.logprobs) > 0, "logprobs should contain entries"
+
+        # Each token should have logprob information
+        for token_logprobs in output.logprobs:
+            assert token_logprobs is not None
+            # Should have up to 5 top logprobs as requested
+            assert len(token_logprobs) <= 5
+
+    def test_logprobs_none_returns_no_probabilities(self, llm: LLM):
+        """When logprobs=None, no log probabilities should be returned."""
+        prompt = "Hello"
+        sampling_params = SamplingParams(temperature=0,
+                                         max_tokens=5,
+                                         logprobs=None)
+
+        outputs = llm.generate([prompt], sampling_params)
+        output = outputs[0].outputs[0]
+
+        # logprobs should be None when not requested
+        assert output.logprobs is None, "logprobs should be None when not requested"
+
+    def test_logprobs_values_are_valid(self, llm: LLM):
+        """Log probabilities should be valid (negative or zero)."""
+        prompt = "The sky is"
+        sampling_params = SamplingParams(temperature=0,
+                                         max_tokens=3,
+                                         logprobs=3)
+
+        outputs = llm.generate([prompt], sampling_params)
+        output = outputs[0].outputs[0]
+
+        assert output.logprobs is not None
+        for token_logprobs in output.logprobs:
+            for token_id, logprob_obj in token_logprobs.items():
+                # Log probabilities should be <= 0
+                assert logprob_obj.logprob <= 0, (
+                    f"Log probability should be <= 0, got {logprob_obj.logprob}"
+                )
+
+
+class TestCombinedParameters:
+    """Tests for combinations of sampling parameters."""
+
+    def test_top_p_and_top_k_combined(self, llm: LLM):
+        """top_p and top_k can be used together."""
+        prompt = "List a fruit:"
+        sampling_params = SamplingParams(
+            temperature=0.7,
+            top_p=0.9,
+            top_k=50,
+            max_tokens=10,
+        )
+
+        outputs = llm.generate([prompt], sampling_params)
+        assert len(outputs[0].outputs[0].text) > 0
+
+    def test_all_params_with_logprobs(self, llm: LLM):
+        """All sampling parameters should work together with logprobs."""
+        prompt = "Complete this sentence: The weather today is"
+        sampling_params = SamplingParams(
+            temperature=0.5,
+            top_p=0.95,
+            top_k=40,
+            max_tokens=10,
+            logprobs=3,
+        )
+
+        outputs = llm.generate([prompt], sampling_params)
+        output = outputs[0].outputs[0]
+
+        # Should have generated text
+        assert len(output.text) > 0
+
+        # Should have logprobs
+        assert output.logprobs is not None
+        assert len(output.logprobs) > 0

tests/e2e/test_speculative_decoding.py

@@ -0,0 +1,311 @@
+# 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 os
+import random
+import string
+import time
+
+import pytest
+from vllm import LLM, SamplingParams
+
+
+# TODO (Qiliang Cui): remove this when XLA fixes the recursive jit call issue.
+def _is_v7x():
+    # jax.devices() will hang so use IS_FOR_V7X to indicate the version.
+    return os.environ.get("IS_FOR_V7X", "false") == "true"
+
+
+def _get_tensor_parallel_size():
+    # Work around an XLA issue.
+    if _is_v7x():
+        return 2
+    return 1
+
+
+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,
+                      tensor_parallel_size=_get_tensor_parallel_size())
+        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,
+                       tensor_parallel_size=_get_tensor_parallel_size())
+        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,
+                      tensor_parallel_size=_get_tensor_parallel_size())
+
+        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,
+                       tensor_parallel_size=_get_tensor_parallel_size(),
+                       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,
+        }, 1.2 if _is_v7x() else 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,
+        }, 1.5 if _is_v7x() else 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.2 if _is_v7x() else 1.8)