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
[V1] TPU CI - Fix test_compilation.py (#15570)
Signed-off-by: Alexander Matveev <amatveev@redhat.com>
Alexander Matveev committed
Mar 26, 2025 at 17:51 UTC
9d119a86ae9a1655a972afc7b1f701b7c7191876
2 files changed
+18
-41
.buildkite/run-tpu-v1-test.sh
+1
-1
@@ -22,7 +22,7 @@ docker run --privileged --net host --shm-size=16G -it \
22
&& export VLLM_USE_V1=1 \
23
&& export VLLM_XLA_CHECK_RECOMPILATION=1 \
24
&& echo TEST_1 \
25
- && pytest /workspace/vllm/tests/tpu/test_compilation.py \
25
+ && pytest -v -s /workspace/vllm/tests/tpu/test_compilation.py \
26
&& echo TEST_2 \
27
&& pytest -v -s /workspace/vllm/tests/v1/tpu/test_basic.py \
28
&& echo TEST_3 \
tests/tpu/test_compilation.py
+17
-40
@@ -5,12 +5,8 @@ import os
5
import tempfile
6
7
import depyf
8
-import pytest
8
10
-from vllm.config import CompilationLevel
9
12
-
13
-@pytest.mark.skip(reason="Not working; needs investigation.")
10
def test_tpu_compilation():
11
temp_dir = tempfile.mkdtemp()
12
with depyf.prepare_debug(temp_dir):
@@ -22,27 +18,24 @@ def test_tpu_compilation():
18
"The greatest glory in living lies not in never falling,",
19
]
20
answers = [
25
- " or, through inaction, allow a human being to come to harm.",
26
- " what is essential is invisible to the eye.",
27
- " but in rising every time we fall.",
21
+ " or, through inaction",
22
+ " what is essential ",
23
+ " but in rising ",
24
]
29
- N = 1
25
+
26
# Currently, top-p sampling is disabled. `top_p` should be 1.0.
27
+ N = 1
28
sampling_params = SamplingParams(temperature=0.7,
29
top_p=1.0,
30
n=N,
31
max_tokens=16)
32
36
- # Set `enforce_eager=True` to avoid ahead-of-time compilation.
37
- # In real workloads, `enforace_eager` should be `False`.
38
-
39
- # disable custom dispatcher, let Dynamo takes over
40
- # all the control
33
llm = LLM(model="Qwen/Qwen2.5-1.5B-Instruct",
42
- max_model_len=512,
43
- max_num_seqs=64,
44
- enforce_eager=True,
45
- compilation_config={"level": CompilationLevel.DYNAMO_AS_IS})
34
+ max_num_batched_tokens=256,
35
+ max_model_len=256,
36
+ max_num_seqs=32,
37
+ enforce_eager=False)
38
+
39
outputs = llm.generate(prompts, sampling_params)
40
for output, answer in zip(outputs, answers):
41
prompt = output.prompt
@@ -56,16 +49,11 @@ def test_tpu_compilation():
49
for i, compiled_code in enumerate(compiled_codes):
50
print("{} file: {}".format(i + 1, compiled_code))
51
59
- # We should only trigger Dynamo compilation 4 times:
60
- # 1. forward pass (symbolic)
61
- # 2. compute_logits (symbolic)
62
- # 3. forward pass (shape 16)
63
- # 4. forward pass (shape 32)
64
- # and later calls should not trigger Dynamo compilation again.
65
- # NOTE: It might still trigger XLA compilation.
66
-
52
+ # We should only trigger Dynamo compilation 2 times:
53
+ # 1. Forward pass without kv_caches
54
+ # 2. Forward pass with kv_caches
55
# Check we have 4 compiled codes
68
- assert len(compiled_codes) == 4
56
+ assert len(compiled_codes) == 2
57
58
kv_cache_prefix = "kv_cache"
59
attn_prefix = "ragged_paged_attention"
@@ -77,24 +65,13 @@ def test_tpu_compilation():
65
for i, compiled_fn in enumerate(compiled_fns):
66
print("{} file: {}".format(i + 1, compiled_fn))
67
80
- # The first compilation is symbolic, so it should not have any kv_caches
68
+ # The first compilation should not have any kv_caches
69
with open(compiled_fns[0]) as f:
70
content = f.read()
71
assert kv_cache_prefix not in content
72
85
- # The second compilation is symbolic, so it should not have any kv_caches
86
- with open(compiled_fns[1]) as f:
87
- content = f.read()
88
- assert kv_cache_prefix not in content
89
-
90
- # The third compilation is shape 16, so it should have kv_caches and the
91
- # ragged_paged_attention
92
- with open(compiled_fns[2]) as f:
93
- content = f.read()
94
- assert (kv_cache_prefix in content and attn_prefix in content)
95
-
96
- # The forth compilation is shape 32, so it should have kv_caches and the
73
+ # The second compilation should have kv_caches and the
74
# ragged_paged_attention
98
- with open(compiled_fns[3]) as f:
75
+ with open(compiled_fns[1]) as f:
76
content = f.read()
77
assert (kv_cache_prefix in content and attn_prefix in content)