A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[Bugfix][TPU][V1] Fix recompilation (#15553)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Mar 27, 2025 at 20:15 UTC
4098b72210dc10761bb348b373bbd0fc9b23b0e4
4 files changed
+15
-74
.buildkite/run-tpu-v1-test.sh
+3
-1
@@ -32,7 +32,9 @@ docker run --privileged --net host --shm-size=16G -it \
32
&& echo TEST_5 \
33
&& python3 /workspace/vllm/examples/offline_inference/tpu.py \
34
&& echo TEST_6 \
35
- && pytest -s -v /workspace/vllm/tests/tpu/worker/test_tpu_model_runner.py" \
35
+ && pytest -s -v /workspace/vllm/tests/tpu/worker/test_tpu_model_runner.py \
36
+ && echo TEST_7 \
37
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py" \
38
39
40
# TODO: This test fails because it uses RANDOM_SEED sampling
tests/v1/tpu/test_sampler.py
+5
-64
@@ -1,7 +1,4 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-import tempfile
3
-from time import time
4
-
2
import pytest
3
4
from vllm import LLM, envs
@@ -15,60 +12,6 @@ if not envs.VLLM_USE_V1:
12
)
13
14
18
-@pytest.mark.parametrize("model_name", ["D4nt3/Qwen2.5-two-layers"])
19
-@pytest.mark.skipif(not current_platform.is_tpu(),
20
- reason="This test needs a TPU")
21
-def test_sampler_compilation(model_name: str, monkeypatch):
22
- """
23
- Check that no recompilation happens despite changing sampling parameters.
24
- We can't read XLA metrics from the engine process, hence we measure time.
25
- """
26
- with tempfile.TemporaryDirectory() as temp_dir:
27
- monkeypatch.setenv("VLLM_XLA_CACHE_PATH", temp_dir)
28
- # Compiling model init may still take some time, enforce_eager to skip.
29
- llm = LLM(model_name,
30
- enforce_eager=True,
31
- max_num_seqs=16,
32
- max_model_len=1024,
33
- gpu_memory_utilization=0.5)
34
- prompts = [
35
- "A robot may not injure a human being",
36
- "It is only with the heart that one can see rightly;",
37
- ]
38
- # First inference should be slow
39
- sampling_params = SamplingParams(
40
- temperature=0.7,
41
- # top_p=0.6, # TODO too slow!
42
- top_k=10,
43
- min_p=0.2,
44
- max_tokens=16)
45
- s = time()
46
- _ = llm.generate(prompts, sampling_params)
47
- run1 = time() - s
48
-
49
- # Second request with different params, but for which we
50
- # compiled for in previous eager iteration.
51
- sampling_params = SamplingParams(temperature=0.1,
52
- top_k=12,
53
- min_p=0.8,
54
- max_tokens=24)
55
- s = time()
56
- _ = llm.generate(prompts, sampling_params)
57
- run2 = time() - s
58
- # Much faster after compiling
59
- assert run1 * 0.1 > run2
60
- print("TIMES", run1, run2)
61
-
62
- # Third request with min_p set to "None". It will not trigger
63
- # recompilation as a default 0 value will be used.
64
- sampling_params = SamplingParams(max_tokens=24, temperature=0.0)
65
- s = time()
66
- _ = llm.generate(prompts, sampling_params)
67
- run3 = time() - s
68
- assert run1 * 0.1 > run3
69
- print("TIMES", run1, run3)
70
-
71
-
15
@pytest.mark.parametrize("model_name", ["Qwen/Qwen2.5-1.5B-Instruct"])
16
@pytest.mark.skipif(not current_platform.is_tpu(),
17
reason="This test needs a TPU")
@@ -77,13 +20,11 @@ def test_sampler_different(model_name: str):
20
Test significantly different sampling params to assert the model produces
21
different results.
22
"""
80
- llm = LLM(
81
- model_name,
82
- enforce_eager=True,
83
- max_num_seqs=1,
84
- max_model_len=64,
85
- # TODO: setting to 0.5 or it will go OOM
86
- gpu_memory_utilization=0.5)
23
+ llm = LLM(model_name,
24
+ enforce_eager=False,
25
+ max_num_seqs=1,
26
+ max_model_len=512,
27
+ max_num_batched_tokens=512)
28
prompts = [
29
"Write a short story about a robot that dreams for the first time."
30
]
vllm/v1/sample/tpu/metadata.py
+1
-7
@@ -88,6 +88,7 @@ class TPUSupportedSamplingMetadata:
88
# Copy slice from CPU to corresponding TPU pre-allocated tensor.
89
# Pad value is the default one.
90
cpu_tensor[num_reqs:padded_num_reqs] = fill_val
91
+ # Subtle compilation: len(tpu_tensor) must be >= `padded_num_reqs`
92
tpu_tensor[:padded_num_reqs] = cpu_tensor[:padded_num_reqs]
93
94
# NOTE NickLucche The sync CPU-TPU graph we produce here must be
@@ -101,13 +102,6 @@ class TPUSupportedSamplingMetadata:
102
copy_slice(input_batch.min_p_cpu_tensor, input_batch.min_p,
103
DEFAULT_SAMPLING_PARAMS["min_p"])
104
104
- # copy_slice(input_batch.frequency_penalties_cpu_tensor,
105
- # input_batch.frequency_penalties)
106
- # copy_slice(input_batch.presence_penalties_cpu_tensor,
107
- # input_batch.presence_penalties)
108
- # copy_slice(input_batch.repetition_penalties_cpu_tensor,
109
- # input_batch.repetition_penalties)
110
-
105
xm.mark_step()
106
xm.wait_device_ops()
107
vllm/v1/worker/tpu_model_runner.py
+6
-2
@@ -88,6 +88,8 @@ class TPUModelRunner:
88
self.max_model_len = model_config.max_model_len
89
self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)
90
self.max_num_tokens = scheduler_config.max_num_batched_tokens
91
+ # InputBatch needs to work with sampling tensors greater than padding
92
+ # to avoid dynamic shapes. Also, avoid suboptimal alignment.
93
self.max_num_reqs = max(scheduler_config.max_num_seqs, MIN_NUM_SEQS)
94
95
# Model-related.
@@ -788,6 +790,7 @@ class TPUModelRunner:
790
dummy_hidden = torch.randn((num_tokens, hsize),
791
device=device,
792
dtype=torch.bfloat16)
793
+ # Compile for [8, 16, .., 128,.., `self.max_num_reqs`]
794
while True:
795
indices = torch.zeros(
796
num_reqs_to_sample,
@@ -804,7 +807,9 @@ class TPUModelRunner:
807
out = out.cpu()
808
if num_reqs_to_sample >= self.max_num_reqs:
809
break
807
- num_reqs_to_sample *= 2
810
+ # Make sure to compile the `max_num_reqs` upper-limit case
811
+ num_reqs_to_sample = _get_padded_num_reqs_with_upper_limit(
812
+ num_reqs_to_sample + 1, self.max_num_reqs)
813
xm.wait_device_ops()
814
end = time.perf_counter()
815
logger.info("Compilation finished in in %.2f [secs].", end - start)
@@ -897,7 +902,6 @@ class ModelWrapperV1(nn.Module):
902
903
return hidden_states
904
900
- # @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
905
def sample_from_hidden(
906
self,
907
hidden_states: torch.Tensor,