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, 18:30 UTC
[Misc] Replace os environ to monkeypatch in test suite (#14516)
Signed-off-by: sibi <85477603+t-sibiraj@users.noreply.github.com> Signed-off-by: Aaron Pham <contact@aarnphm.xyz> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Aaron Pham <contact@aarnphm.xyz>
Sibi committed
Mar 17, 2025 at 11:35 UTC
a73e183e36a818ea95f442ae1751bc66cf4f135d
43 files changed
+1997
-1755
.buildkite/test-pipeline.yaml
+1
-1
@@ -522,7 +522,7 @@ steps:
522
# TODO: investigate and fix
523
# - pytest -v -s spec_decode/e2e/test_integration_dist_tp2.py
524
- VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
525
- - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/disagg_test.py
525
+ - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/test_disagg.py
526
527
- label: Plugin Tests (2 GPUs) # 40min
528
working_dir: "/vllm-workspace/tests"
tests/basic_correctness/test_basic_correctness.py
+63
-52
@@ -47,6 +47,7 @@ def test_vllm_gc_ed():
47
@pytest.mark.parametrize("max_tokens", [5])
48
@pytest.mark.parametrize("enforce_eager", [False])
49
def test_models(
50
+ monkeypatch: pytest.MonkeyPatch,
51
hf_runner,
52
model: str,
53
backend: str,
@@ -63,31 +64,33 @@ def test_models(
64
pytest.skip(
65
f"{backend} does not support gemma2 with full context length.")
66
66
- os.environ["VLLM_ATTENTION_BACKEND"] = backend
67
+ with monkeypatch.context() as m:
68
+ m.setenv("VLLM_ATTENTION_BACKEND", backend)
69
68
- # 5042 tokens for gemma2
69
- # gemma2 has alternating sliding window size of 4096
70
- # we need a prompt with more than 4096 tokens to test the sliding window
71
- prompt = "The following numbers of the sequence " + ", ".join(
72
- str(i) for i in range(1024)) + " are:"
73
- example_prompts = [prompt]
70
+ # 5042 tokens for gemma2
71
+ # gemma2 has alternating sliding window size of 4096
72
+ # we need a prompt with more than 4096 tokens to test the sliding window
73
+ prompt = "The following numbers of the sequence " + ", ".join(
74
+ str(i) for i in range(1024)) + " are:"
75
+ example_prompts = [prompt]
76
75
- with hf_runner(model, dtype=dtype) as hf_model:
76
- hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
77
+ with hf_runner(model, dtype=dtype) as hf_model:
78
+ hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
79
78
- with VllmRunner(model,
79
- max_model_len=8192,
80
- dtype=dtype,
81
- enforce_eager=enforce_eager,
82
- gpu_memory_utilization=0.7) as vllm_model:
83
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
80
+ with VllmRunner(model,
81
+ max_model_len=8192,
82
+ dtype=dtype,
83
+ enforce_eager=enforce_eager,
84
+ gpu_memory_utilization=0.7) as vllm_model:
85
+ vllm_outputs = vllm_model.generate_greedy(example_prompts,
86
+ max_tokens)
87
85
- check_outputs_equal(
86
- outputs_0_lst=hf_outputs,
87
- outputs_1_lst=vllm_outputs,
88
- name_0="hf",
89
- name_1="vllm",
90
- )
88
+ check_outputs_equal(
89
+ outputs_0_lst=hf_outputs,
90
+ outputs_1_lst=vllm_outputs,
91
+ name_0="hf",
92
+ name_1="vllm",
93
+ )
94
95
96
@multi_gpu_test(num_gpus=2)
@@ -104,6 +107,7 @@ def test_models(
107
("meta-llama/Meta-Llama-3-8B", "ray", "FLASHINFER", "A100"),
108
])
109
def test_models_distributed(
110
+ monkeypatch: pytest.MonkeyPatch,
111
hf_runner,
112
vllm_runner,
113
example_prompts,
@@ -116,34 +120,41 @@ def test_models_distributed(
120
if test_suite != TARGET_TEST_SUITE:
121
pytest.skip(f"Skip test for {test_suite}")
122
119
- if model == "meta-llama/Llama-3.2-1B-Instruct" and distributed_executor_backend == "ray" and attention_backend == "" and test_suite == "L4": # noqa
120
- # test Ray Compiled Graph
121
- os.environ['VLLM_USE_RAY_SPMD_WORKER'] = "1"
122
- os.environ['VLLM_USE_RAY_COMPILED_DAG'] = "1"
123
-
124
- if attention_backend:
125
- os.environ["VLLM_ATTENTION_BACKEND"] = attention_backend
126
-
127
- dtype = "half"
128
- max_tokens = 5
129
-
130
- # NOTE: take care of the order. run vLLM first, and then run HF.
131
- # vLLM needs a fresh new process without cuda initialization.
132
- # if we run HF first, the cuda initialization will be done and it
133
- # will hurt multiprocessing backend with fork method (the default method).
134
- with vllm_runner(model,
135
- dtype=dtype,
136
- tensor_parallel_size=2,
137
- distributed_executor_backend=distributed_executor_backend
138
- ) as vllm_model:
139
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
140
-
141
- with hf_runner(model, dtype=dtype) as hf_model:
142
- hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
143
-
144
- check_outputs_equal(
145
- outputs_0_lst=hf_outputs,
146
- outputs_1_lst=vllm_outputs,
147
- name_0="hf",
148
- name_1="vllm",
149
- )
123
+ with monkeypatch.context() as monkeypatch_context:
124
+ if model == "meta-llama/Llama-3.2-1B-Instruct" and distributed_executor_backend == "ray" and attention_backend == "" and test_suite == "L4": # noqa
125
+ # test Ray Compiled Graph
126
+ monkeypatch_context.setenv("VLLM_USE_RAY_SPMD_WORKER", "1")
127
+ monkeypatch_context.setenv("VLLM_USE_RAY_COMPILED_DAG", "1")
128
+
129
+ if attention_backend:
130
+ monkeypatch_context.setenv(
131
+ "VLLM_ATTENTION_BACKEND",
132
+ attention_backend,
133
+ )
134
+
135
+ dtype = "half"
136
+ max_tokens = 5
137
+
138
+ # NOTE: take care of the order. run vLLM first, and then run HF.
139
+ # vLLM needs a fresh new process without cuda initialization.
140
+ # if we run HF first, the cuda initialization will be done and it
141
+ # will hurt multiprocessing backend with fork method
142
+ # (the default method).
143
+ with vllm_runner(
144
+ model,
145
+ dtype=dtype,
146
+ tensor_parallel_size=2,
147
+ distributed_executor_backend=distributed_executor_backend,
148
+ ) as vllm_model:
149
+ vllm_outputs = vllm_model.generate_greedy(example_prompts,
150
+ max_tokens)
151
+
152
+ with hf_runner(model, dtype=dtype) as hf_model:
153
+ hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
154
+
155
+ check_outputs_equal(
156
+ outputs_0_lst=hf_outputs,
157
+ outputs_1_lst=vllm_outputs,
158
+ name_0="hf",
159
+ name_1="vllm",
160
+ )
tests/basic_correctness/test_chunked_prefill.py
+95
-79
@@ -7,16 +7,22 @@ prefill requests are chunked.
7
8
Run `pytest tests/models/test_chunked_prefill.py`.
9
"""
10
-import os
10
+
11
+from __future__ import annotations
12
+
13
+from typing import TYPE_CHECKING
14
15
import pytest
16
14
-from tests.kernels.utils import override_backend_env_variable
17
from vllm.platforms import current_platform
18
+from vllm.utils import STR_BACKEND_ENV_VAR
19
20
from ..models.utils import check_logprobs_close, check_outputs_equal
21
from ..utils import multi_gpu_test
22
23
+if TYPE_CHECKING:
24
+ from .conftest import HfRunner, VllmRunner
25
+
26
MODELS = [
27
"facebook/opt-125m",
28
"meta-llama/Llama-3.2-1B-Instruct",
@@ -24,12 +30,14 @@ MODELS = [
30
31
32
@pytest.fixture(scope="function", autouse=True)
27
-def use_v0_only(monkeypatch):
33
+def use_v0_only(monkeypatch: pytest.MonkeyPatch):
34
"""
35
Since this module is V0 only, set VLLM_USE_V1=0 for
36
all tests in the file.
37
"""
32
- monkeypatch.setenv('VLLM_USE_V1', '0')
38
+ with monkeypatch.context() as m:
39
+ m.setenv('VLLM_USE_V1', '0')
40
+ yield
41
42
43
@pytest.mark.parametrize("model", MODELS)
@@ -42,8 +50,8 @@ def use_v0_only(monkeypatch):
50
@pytest.mark.parametrize("tensor_parallel_size", [1])
51
@pytest.mark.parametrize("attention_backend", ["FLASHINFER", "FLASH_ATTN"])
52
def test_models(
45
- hf_runner,
46
- vllm_runner,
53
+ hf_runner: HfRunner,
54
+ vllm_runner: VllmRunner,
55
example_prompts,
56
model: str,
57
dtype: str,
@@ -52,37 +60,39 @@ def test_models(
60
enforce_eager: bool,
61
tensor_parallel_size: int,
62
attention_backend: str,
55
- monkeypatch,
63
+ monkeypatch: pytest.MonkeyPatch,
64
) -> None:
65
"""
66
Checks exact match decode between huggingface model and vllm runner with
67
chunked prefill.
68
"""
61
- override_backend_env_variable(monkeypatch, attention_backend)
69
+ with monkeypatch.context() as m:
70
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
71
63
- max_num_seqs = chunked_prefill_token_size
64
- max_num_batched_tokens = chunked_prefill_token_size
72
+ max_num_seqs = chunked_prefill_token_size
73
+ max_num_batched_tokens = chunked_prefill_token_size
74
66
- with hf_runner(model, dtype=dtype) as hf_model:
67
- hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
75
+ with hf_runner(model, dtype=dtype) as hf_model:
76
+ hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
77
69
- with vllm_runner(
70
- model,
71
- dtype=dtype,
72
- max_num_batched_tokens=max_num_batched_tokens,
73
- enable_chunked_prefill=True,
74
- tensor_parallel_size=tensor_parallel_size,
75
- enforce_eager=enforce_eager,
76
- max_num_seqs=max_num_seqs,
77
- ) as vllm_model:
78
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
78
+ with vllm_runner(
79
+ model,
80
+ dtype=dtype,
81
+ max_num_batched_tokens=max_num_batched_tokens,
82
+ enable_chunked_prefill=True,
83
+ tensor_parallel_size=tensor_parallel_size,
84
+ enforce_eager=enforce_eager,
85
+ max_num_seqs=max_num_seqs,
86
+ ) as vllm_model:
87
+ vllm_outputs = vllm_model.generate_greedy(example_prompts,
88
+ max_tokens)
89
80
- check_outputs_equal(
81
- outputs_0_lst=hf_outputs,
82
- outputs_1_lst=vllm_outputs,
83
- name_0="hf",
84
- name_1="vllm",
85
- )
90
+ check_outputs_equal(
91
+ outputs_0_lst=hf_outputs,
92
+ outputs_1_lst=vllm_outputs,
93
+ name_0="hf",
94
+ name_1="vllm",
95
+ )
96
97
98
@multi_gpu_test(num_gpus=2)
@@ -90,57 +100,61 @@ def test_models(
100
@pytest.mark.parametrize("model", MODELS)
101
@pytest.mark.parametrize("attention_backend", ["FLASHINFER", "FLASH_ATTN"])
102
def test_models_distributed(
93
- hf_runner,
94
- vllm_runner,
103
+ hf_runner: HfRunner,
104
+ vllm_runner: VllmRunner,
105
example_prompts,
106
model: str,
107
distributed_executor_backend: str,
108
attention_backend: str,
99
- monkeypatch,
109
+ monkeypatch: pytest.MonkeyPatch,
110
) -> None:
101
- override_backend_env_variable(monkeypatch, attention_backend)
102
-
103
- if (model == "meta-llama/Llama-3.2-1B-Instruct"
104
- and distributed_executor_backend == "ray"):
105
- # test Ray Compiled Graph
106
- os.environ['VLLM_USE_RAY_SPMD_WORKER'] = "1"
107
- os.environ['VLLM_USE_RAY_COMPILED_DAG'] = "1"
108
-
109
- dtype = "half"
110
- max_tokens = 5
111
- chunked_prefill_token_size = 16
112
-
113
- # Add a chunked prefill config.
114
- max_num_seqs = min(chunked_prefill_token_size, 256)
115
- assert chunked_prefill_token_size != -1
116
- enable_chunked_prefill = True
117
- max_num_batched_tokens = chunked_prefill_token_size
111
+ with monkeypatch.context() as m:
112
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
113
+ if (model == "meta-llama/Llama-3.2-1B-Instruct"
114
+ and distributed_executor_backend == "ray"):
115
+ # test Ray Compiled Graph
116
+ m.setenv("VLLM_USE_RAY_SPMD_WORKER", "1")
117
+ m.setenv("VLLM_USE_RAY_COMPILED_DAG", "1")
118
+
119
+ dtype = "half"
120
+ max_tokens = 5
121
+ chunked_prefill_token_size = 16
122
+
123
+ # Add a chunked prefill config.
124
+ max_num_seqs = min(chunked_prefill_token_size, 256)
125
+ assert chunked_prefill_token_size != -1
126
+ enable_chunked_prefill = True
127
+ max_num_batched_tokens = chunked_prefill_token_size
128
+
129
+ # NOTE: take care of the order. run vLLM first, and then run HF.
130
+ # vLLM needs a fresh new process without cuda initialization.
131
+ # if we run HF first, the cuda initialization will be done and it
132
+ # will hurt multiprocessing backend with
133
+ # fork method (the default method).
134
119
- # NOTE: take care of the order. run vLLM first, and then run HF.
120
- # vLLM needs a fresh new process without cuda initialization.
121
- # if we run HF first, the cuda initialization will be done and it
122
- # will hurt multiprocessing backend with fork method (the default method).
123
-
124
- with vllm_runner(
125
- model,
126
- dtype=dtype,
127
- tensor_parallel_size=2,
128
- max_num_seqs=max_num_seqs,
129
- enable_chunked_prefill=enable_chunked_prefill,
130
- max_num_batched_tokens=max_num_batched_tokens,
131
- distributed_executor_backend=distributed_executor_backend,
132
- ) as vllm_model:
133
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
135
+ with vllm_runner(
136
+ model,
137
+ dtype=dtype,
138
+ tensor_parallel_size=2,
139
+ max_num_seqs=max_num_seqs,
140
+ enable_chunked_prefill=enable_chunked_prefill,
141
+ max_num_batched_tokens=max_num_batched_tokens,
142
+ distributed_executor_backend=distributed_executor_backend,
143
+ ) as vllm_model:
144
+ vllm_outputs = vllm_model.generate_greedy(
145
+ example_prompts,
146
+ max_tokens,
147
+ )
148
135
- with hf_runner(model, dtype=dtype) as hf_model:
136
- hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
149
+ with hf_runner(model, dtype=dtype) as hf_model:
150
+ hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
151
138
- check_outputs_equal(
139
- outputs_0_lst=hf_outputs,
140
- outputs_1_lst=vllm_outputs,
141
- name_0="hf",
142
- name_1="vllm",
143
- )
152
+ check_outputs_equal(
153
+ outputs_0_lst=hf_outputs,
154
+ outputs_1_lst=vllm_outputs,
155
+ name_0="hf",
156
+ name_1="vllm",
157
+ )
158
159
160
@pytest.mark.parametrize(
@@ -158,7 +172,7 @@ def test_models_distributed(
172
# the async postprocessor
173
@pytest.mark.parametrize("disable_async_output_proc", [True])
174
def test_models_with_fp8_kv_cache(
161
- vllm_runner,
175
+ vllm_runner: VllmRunner,
176
example_prompts,
177
kv_cache_dtype: str,
178
model: str,
@@ -218,7 +232,7 @@ def test_models_with_fp8_kv_cache(
232
@pytest.mark.parametrize("tensor_parallel_size", [1])
233
@pytest.mark.parametrize("dtype", ["half"])
234
def test_with_prefix_caching(
221
- vllm_runner,
235
+ vllm_runner: VllmRunner,
236
max_tokens: int,
237
enforce_eager: bool,
238
chunk_size: int,
@@ -254,8 +268,10 @@ def test_with_prefix_caching(
268
) as vllm_model:
269
outputs[enable] = []
270
for prompt in full_prompts:
257
- outputs[enable] += vllm_model.generate_greedy([prompt],
258
- max_tokens)
271
+ outputs[enable] += vllm_model.generate_greedy(
272
+ [prompt],
273
+ max_tokens,
274
+ )
275
276
check_outputs_equal(
277
outputs_0_lst=outputs[False],
@@ -274,8 +290,8 @@ def test_with_prefix_caching(
290
@pytest.mark.cpu_model
291
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
292
def test_models_cpu(
277
- hf_runner,
278
- vllm_runner,
293
+ hf_runner: HfRunner,
294
+ vllm_runner: VllmRunner,
295
example_prompts,
296
model: str,
297
dtype: str,
@@ -283,7 +299,7 @@ def test_models_cpu(
299
chunked_prefill_token_size: int,
300
enforce_eager: bool,
301
attention_backend: str,
286
- monkeypatch,
302
+ monkeypatch: pytest.MonkeyPatch,
303
) -> None:
304
test_models(
305
hf_runner,
@@ -307,7 +323,7 @@ def test_models_cpu(
323
@pytest.mark.cpu_model
324
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
325
def test_with_prefix_caching_cpu(
310
- vllm_runner,
326
+ vllm_runner: VllmRunner,
327
max_tokens: int,
328
enforce_eager: bool,
329
chunk_size: int,
tests/basic_correctness/test_cumem.py
+35
-37
@@ -123,40 +123,38 @@ def test_cumem_with_cudagraph():
123
# sleep mode with pytorch checkpoint
124
("facebook/opt-125m", False),
125
])
126
-def test_end_to_end(model: str, use_v1: bool):
127
- import os
128
- os.environ["VLLM_USE_V1"] = "1" if use_v1 else "0"
129
- free, total = torch.cuda.mem_get_info()
130
- used_bytes_baseline = total - free # in case other process is running
131
- llm = LLM(model, enable_sleep_mode=True)
132
- prompt = "How are you?"
133
- sampling_params = SamplingParams(temperature=0, max_tokens=10)
134
- output = llm.generate(prompt, sampling_params)
135
-
136
- # the benefit of `llm.sleep(level=2)` is mainly CPU memory usage,
137
- # which is difficult to measure in the test. therefore, we only
138
- # test sleep level 1 here.
139
- llm.sleep(level=1)
140
-
141
- free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
142
- used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
143
- # now the memory usage is mostly cudagraph memory pool,
144
- # and it should be less than the model weights (1B model, 2GiB weights)
145
-
146
- # NOTE: In V1, the memory buffer for logits (max_num_reqs x vocab_size)
147
- # is captured but cannot be releasesd from PyTorch due to a known bug,
148
- # therefore high memory usage after `llm.sleep` is called is expected.
149
- # FIXME(youkaichao & ywang96): Fix memory buffer issue with sleep mode
150
- # in V1.
151
- if use_v1:
152
- assert used_bytes < 7 * GiB_bytes
153
- else:
154
- assert used_bytes < 2 * GiB_bytes
155
-
156
- llm.wake_up()
157
- output2 = llm.generate(prompt, sampling_params)
158
-
159
- # cmp output
160
- assert output[0].outputs[0].text == output2[0].outputs[0].text
161
-
162
- del os.environ["VLLM_USE_V1"]
126
+def test_end_to_end(monkeypatch: pytest.MonkeyPatch, model: str, use_v1: bool):
127
+ with monkeypatch.context() as m:
128
+ m.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
129
+ free, total = torch.cuda.mem_get_info()
130
+ used_bytes_baseline = total - free # in case other process is running
131
+ llm = LLM(model, enable_sleep_mode=True)
132
+ prompt = "How are you?"
133
+ sampling_params = SamplingParams(temperature=0, max_tokens=10)
134
+ output = llm.generate(prompt, sampling_params)
135
+
136
+ # the benefit of `llm.sleep(level=2)` is mainly CPU memory usage,
137
+ # which is difficult to measure in the test. therefore, we only
138
+ # test sleep level 1 here.
139
+ llm.sleep(level=1)
140
+
141
+ free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
142
+ used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
143
+ # now the memory usage is mostly cudagraph memory pool,
144
+ # and it should be less than the model weights (1B model, 2GiB weights)
145
+
146
+ # NOTE: In V1, the memory buffer for logits (max_num_reqs x vocab_size)
147
+ # is captured but cannot be releasesd from PyTorch due to a known bug,
148
+ # therefore high memory usage after `llm.sleep` is called is expected.
149
+ # FIXME(youkaichao & ywang96): Fix memory buffer issue with sleep mode
150
+ # in V1.
151
+ if use_v1:
152
+ assert used_bytes < 7 * GiB_bytes
153
+ else:
154
+ assert used_bytes < 2 * GiB_bytes
155
+
156
+ llm.wake_up()
157
+ output2 = llm.generate(prompt, sampling_params)
158
+
159
+ # cmp output
160
+ assert output[0].outputs[0].text == output2[0].outputs[0].text
tests/compile/test_basic_correctness.py
+106
-101
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from __future__ import annotations
3
4
import dataclasses
4
-from typing import Optional
5
6
import pytest
7
@@ -22,75 +22,76 @@ class TestSetting:
22
fullgraph: bool
23
24
25
-# representative settings for testing
26
-test_settings = [
27
- # basic llama model
28
- TestSetting(
29
- model="meta-llama/Llama-3.2-1B-Instruct",
30
- model_args=[],
31
- pp_size=2,
32
- tp_size=2,
33
- attn_backend="FLASHINFER",
34
- method="generate",
35
- fullgraph=True,
36
- ),
37
- # llama model with quantization
38
- TestSetting(
39
- model="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ",
40
- model_args=["--quantization", "gptq"],
41
- pp_size=1,
42
- tp_size=1,
43
- attn_backend="FLASH_ATTN",
44
- method="generate",
45
- fullgraph=True,
46
- ),
47
- # MoE model
48
- TestSetting(
49
- model="ibm/PowerMoE-3b",
50
- model_args=[],
51
- pp_size=1,
52
- tp_size=2,
53
- attn_backend="FLASH_ATTN",
54
- method="generate",
55
- fullgraph=True,
56
- ),
57
- # embedding model
58
- TestSetting(
59
- model="BAAI/bge-multilingual-gemma2",
60
- model_args=["--task", "embed"],
61
- pp_size=1,
62
- tp_size=1,
63
- attn_backend="FLASH_ATTN",
64
- method="encode",
65
- fullgraph=True,
66
- ),
67
- # encoder-based embedding model (BERT)
68
- TestSetting(
69
- model="BAAI/bge-base-en-v1.5",
70
- model_args=["--task", "embed"],
71
- pp_size=1,
72
- tp_size=1,
73
- attn_backend="XFORMERS",
74
- method="encode",
75
- fullgraph=True,
76
- ),
77
- # vision language model
78
- TestSetting(
79
- model="microsoft/Phi-3.5-vision-instruct",
80
- model_args=["--trust-remote-code", "--max-model-len", "2048"],
81
- pp_size=2,
82
- tp_size=1,
83
- attn_backend="FLASH_ATTN",
84
- method="generate_with_image",
85
- fullgraph=False,
86
- ),
87
-]
88
-
89
-
25
# we cannot afford testing the full Catesian product
26
# of all models and all levels
92
-@pytest.mark.parametrize("test_setting", test_settings)
93
-def test_compile_correctness(test_setting: TestSetting):
27
+@pytest.mark.parametrize(
28
+ "test_setting",
29
+ [
30
+ # basic llama model
31
+ TestSetting(
32
+ model="meta-llama/Llama-3.2-1B-Instruct",
33
+ model_args=[],
34
+ pp_size=2,
35
+ tp_size=2,
36
+ attn_backend="FLASHINFER",
37
+ method="generate",
38
+ fullgraph=True,
39
+ ),
40
+ # llama model with quantization
41
+ TestSetting(
42
+ model="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ",
43
+ model_args=["--quantization", "gptq"],
44
+ pp_size=1,
45
+ tp_size=1,
46
+ attn_backend="FLASH_ATTN",
47
+ method="generate",
48
+ fullgraph=True,
49
+ ),
50
+ # MoE model
51
+ TestSetting(
52
+ model="ibm/PowerMoE-3b",
53
+ model_args=[],
54
+ pp_size=1,
55
+ tp_size=2,
56
+ attn_backend="FLASH_ATTN",
57
+ method="generate",
58
+ fullgraph=True,
59
+ ),
60
+ # embedding model
61
+ TestSetting(
62
+ model="BAAI/bge-multilingual-gemma2",
63
+ model_args=["--task", "embed"],
64
+ pp_size=1,
65
+ tp_size=1,
66
+ attn_backend="FLASH_ATTN",
67
+ method="encode",
68
+ fullgraph=True,
69
+ ),
70
+ # encoder-based embedding model (BERT)
71
+ TestSetting(
72
+ model="BAAI/bge-base-en-v1.5",
73
+ model_args=["--task", "embed"],
74
+ pp_size=1,
75
+ tp_size=1,
76
+ attn_backend="XFORMERS",
77
+ method="encode",
78
+ fullgraph=True,
79
+ ),
80
+ # vision language model
81
+ TestSetting(
82
+ model="microsoft/Phi-3.5-vision-instruct",
83
+ model_args=["--trust-remote-code", "--max-model-len", "2048"],
84
+ pp_size=2,
85
+ tp_size=1,
86
+ attn_backend="FLASH_ATTN",
87
+ method="generate_with_image",
88
+ fullgraph=False,
89
+ ),
90
+ ])
91
+def test_compile_correctness(
92
+ monkeypatch: pytest.MonkeyPatch,
93
+ test_setting: TestSetting,
94
+):
95
# this test is run under multiple suits, with different GPUs.
96
# make sure we only run the test with correct CUDA devices.
97
# don't use "<", as it will duplicate the tests.
@@ -103,41 +104,45 @@ def test_compile_correctness(test_setting: TestSetting):
104
fullgraph = test_setting.fullgraph
105
if cuda_device_count_stateless() != pp_size * tp_size:
106
pytest.skip("Not correct CUDA devices for the test.")
106
- import os
107
- os.environ["VLLM_ATTENTION_BACKEND"] = attn_backend
108
- final_args = ["--enforce-eager"] + model_args + ["-pp", str(pp_size)] + \
109
- ["-tp", str(tp_size)]
107
111
- all_args: list[list[str]] = []
112
- all_envs: list[Optional[dict[str, str]]] = []
108
+ with monkeypatch.context() as m:
109
+ m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
110
+ final_args = [
111
+ "--enforce-eager", *model_args, "-pp",
112
+ str(pp_size), "-tp",
113
+ str(tp_size)
114
+ ]
115
+
116
+ all_args: list[list[str]] = []
117
+ all_envs: list[dict[str, str] | None] = []
118
114
- for level in [
115
- CompilationLevel.NO_COMPILATION,
116
- CompilationLevel.PIECEWISE,
117
- ]:
118
- all_args.append(final_args + [f"-O{level}"])
119
- all_envs.append({})
119
+ for level in [
120
+ CompilationLevel.NO_COMPILATION,
121
+ CompilationLevel.PIECEWISE,
122
+ ]:
123
+ all_args.append(final_args + [f"-O{level}"])
124
+ all_envs.append({})
125
121
- # inductor will change the output, so we only compare if the output
122
- # is close, not exactly the same.
123
- compare_all_settings(
124
- model,
125
- all_args,
126
- all_envs,
127
- method=method if method != "generate" else "generate_close")
128
- all_envs.clear()
129
- all_args.clear()
126
+ # inductor will change the output, so we only compare if the output
127
+ # is close, not exactly the same.
128
+ compare_all_settings(
129
+ model,
130
+ all_args,
131
+ all_envs,
132
+ method=method if method != "generate" else "generate_close")
133
+ all_envs.clear()
134
+ all_args.clear()
135
131
- for level in [
132
- CompilationLevel.NO_COMPILATION,
133
- CompilationLevel.DYNAMO_AS_IS,
134
- CompilationLevel.DYNAMO_ONCE,
135
- ]:
136
- all_args.append(final_args + [f"-O{level}"])
137
- all_envs.append({})
138
- if level != CompilationLevel.DYNAMO_ONCE and not fullgraph:
139
- # "DYNAMO_ONCE" will always use fullgraph
140
- all_envs[-1][
141
- "VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE"] = "0" # type: ignore
136
+ for level in [
137
+ CompilationLevel.NO_COMPILATION,
138
+ CompilationLevel.DYNAMO_AS_IS,
139
+ CompilationLevel.DYNAMO_ONCE,
140
+ ]:
141
+ all_args.append(final_args + [f"-O{level}"])
142
+ all_envs.append({})
143
+ if level != CompilationLevel.DYNAMO_ONCE and not fullgraph:
144
+ # "DYNAMO_ONCE" will always use fullgraph
145
+ all_envs[-1][
146
+ "VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE"] = "0" # type: ignore
147
143
- compare_all_settings(model, all_args * 3, all_envs, method=method)
148
+ compare_all_settings(model, all_args * 3, all_envs, method=method)
tests/compile/test_full_graph.py
+103
-10
@@ -1,22 +1,115 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+from __future__ import annotations
4
+
5
+from typing import Any
6
+
7
import pytest
8
+import torch
9
10
+from tests.quantization.utils import is_quant_method_supported
11
+from vllm import LLM, SamplingParams
12
from vllm.config import CompilationLevel
13
+from vllm.platforms import current_platform
14
15
from ..utils import fork_new_process_for_each_test
8
-from .utils import TEST_MODELS, check_full_graph_support
16
17
11
-@pytest.mark.parametrize("model_info", TEST_MODELS)
18
+@pytest.fixture(params=None, name="model_info")
19
+def models_list_fixture(request):
20
+ TEST_MODELS: list[tuple[str, dict[str, Any]]] = [
21
+ ("facebook/opt-125m", {}),
22
+ ("nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", {
23
+ "dtype": torch.float16,
24
+ "quantization": "compressed-tensors"
25
+ }),
26
+ ("neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic", {
27
+ "dtype": torch.float16,
28
+ "quantization": "compressed-tensors"
29
+ }),
30
+ ("neuralmagic/Llama-3.2-1B-Instruct-quantized.w8a8", {
31
+ "quantization": "compressed-tensors"
32
+ }),
33
+ ("meta-llama/Llama-3.2-1B-Instruct", {}),
34
+ ]
35
+
36
+ if is_quant_method_supported("aqlm"):
37
+ TEST_MODELS.append(("ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf", {
38
+ "quantization": "aqlm"
39
+ }))
40
+
41
+ # TODO: figure out why this fails.
42
+ if False and is_quant_method_supported("gguf"): # noqa: SIM223
43
+ TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", {
44
+ "quantization": "gguf"
45
+ }))
46
+
47
+ if is_quant_method_supported("gptq"):
48
+ TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", {
49
+ "quantization": "gptq"
50
+ }))
51
+
52
+ if is_quant_method_supported("gptq_marlin"):
53
+ TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", {
54
+ "quantization": "gptq_marlin"
55
+ }))
56
+
57
+ if is_quant_method_supported("gptq_marlin_24"):
58
+ TEST_MODELS.append(("alexm-nm/tinyllama-24-marlin24-4bit-g128", {
59
+ "quantization": "gptq_marlin_24"
60
+ }))
61
+
62
+ if is_quant_method_supported("marlin"):
63
+ TEST_MODELS.append(
64
+ ("robertgshaw2/TinyLlama-1.1B-Chat-v1.0-g128-marlin", {
65
+ "quantization": "marlin"
66
+ }))
67
+
68
+ if not current_platform.is_rocm() and is_quant_method_supported("awq"):
69
+ TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {
70
+ "quantization": "AWQ"
71
+ }))
72
+
73
+ return TEST_MODELS
74
+
75
+
76
@pytest.mark.parametrize(
77
"optimization_level",
14
- [CompilationLevel.DYNAMO_ONCE, CompilationLevel.PIECEWISE])
78
+ [CompilationLevel.DYNAMO_ONCE, CompilationLevel.PIECEWISE],
79
+)
80
+@pytest.mark.parametrize("model_info", "", indirect=True)
81
@fork_new_process_for_each_test
16
-def test_full_graph(model_info, optimization_level):
17
- model = model_info[0]
18
- model_kwargs = model_info[1]
19
- check_full_graph_support(model,
20
- model_kwargs,
21
- optimization_level,
22
- tp_size=1)
82
+def test_full_graph(
83
+ monkeypatch: pytest.MonkeyPatch,
84
+ model_info: tuple[str, dict[str, Any]],
85
+ optimization_level: int,
86
+):
87
+ model, model_kwargs = model_info
88
+
89
+ with monkeypatch.context() as m:
90
+ # make sure these models can be captured in full graph mode
91
+ m.setenv("VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1")
92
+ print(f"MODEL={model}")
93
+
94
+ prompts = [
95
+ "Hello, my name is",
96
+ "The president of the United States is",
97
+ "The capital of France is",
98
+ "The future of AI is",
99
+ ]
100
+ sampling_params = SamplingParams(temperature=0)
101
+ llm = LLM(
102
+ model=model,
103
+ enforce_eager=True,
104
+ tensor_parallel_size=1,
105
+ disable_custom_all_reduce=True,
106
+ compilation_config=optimization_level,
107
+ **model_kwargs,
108
+ )
109
+ outputs = llm.generate(prompts, sampling_params)
110
+
111
+ # Print the outputs.
112
+ for output in outputs:
113
+ prompt = output.prompt
114
+ generated_text = output.outputs[0].text
115
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
tests/compile/utils.py
deleted
-93
@@ -1,93 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-import os
4
-
5
-import torch
6
-
7
-from tests.quantization.utils import is_quant_method_supported
8
-from vllm import LLM, SamplingParams
9
-from vllm.platforms import current_platform
10
-
11
-TEST_MODELS = [
12
- ("facebook/opt-125m", {}),
13
- ("nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", {
14
- "dtype": torch.float16,
15
- "quantization": "compressed-tensors"
16
- }),
17
- ("neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic", {
18
- "dtype": torch.float16,
19
- "quantization": "compressed-tensors"
20
- }),
21
- ("neuralmagic/Llama-3.2-1B-Instruct-quantized.w8a8", {
22
- "quantization": "compressed-tensors"
23
- }),
24
- ("meta-llama/Llama-3.2-1B-Instruct", {}),
25
-]
26
-
27
-if is_quant_method_supported("aqlm"):
28
- TEST_MODELS.append(("ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf", {
29
- "quantization": "aqlm"
30
- }))
31
-
32
-# TODO: figure out why this fails.
33
-if False and is_quant_method_supported("gguf"): # noqa: SIM223
34
- TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", {
35
- "quantization": "gguf"
36
- }))
37
-
38
-if is_quant_method_supported("gptq"):
39
- TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", {
40
- "quantization": "gptq"
41
- }))
42
-
43
-if is_quant_method_supported("gptq_marlin"):
44
- TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", {
45
- "quantization": "gptq_marlin"
46
- }))
47
-
48
-if is_quant_method_supported("gptq_marlin_24"):
49
- TEST_MODELS.append(("alexm-nm/tinyllama-24-marlin24-4bit-g128", {
50
- "quantization": "gptq_marlin_24"
51
- }))
52
-
53
-if is_quant_method_supported("marlin"):
54
- TEST_MODELS.append(("robertgshaw2/TinyLlama-1.1B-Chat-v1.0-g128-marlin", {
55
- "quantization": "marlin"
56
- }))
57
-
58
-if not current_platform.is_rocm() and is_quant_method_supported("awq"):
59
- TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {
60
- "quantization": "AWQ"
61
- }))
62
-
63
-
64
-def check_full_graph_support(model,
65
- model_kwargs,
66
- optimization_level,
67
- tp_size=1):
68
- # make sure these models can be captured in full graph mode
69
- os.environ["VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE"] = "1"
70
-
71
- print(f"MODEL={model}")
72
-
73
- prompts = [
74
- "Hello, my name is",
75
- "The president of the United States is",
76
- "The capital of France is",
77
- "The future of AI is",
78
- ]
79
- sampling_params = SamplingParams(temperature=0)
80
- llm = LLM(model=model,
81
- enforce_eager=True,
82
- tensor_parallel_size=tp_size,
83
- disable_custom_all_reduce=True,
84
- compilation_config=optimization_level,
85
- **model_kwargs)
86
-
87
- outputs = llm.generate(prompts, sampling_params)
88
-
89
- # Print the outputs.
90
- for output in outputs:
91
- prompt = output.prompt
92
- generated_text = output.outputs[0].text
93
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
tests/conftest.py
+1
-1
@@ -1110,4 +1110,4 @@ def pytest_collection_modifyitems(config, items):
1110
skip_optional = pytest.mark.skip(reason="need --optional option to run")
1111
for item in items:
1112
if "optional" in item.keywords:
1113
- item.add_marker(skip_optional)
1113
+ item.add_marker(skip_optional)
\ No newline at end of file
tests/distributed/test_comm_ops.py
+63
-22
@@ -3,7 +3,10 @@
3
4
Run `pytest tests/distributed/test_comm_ops.py`.
5
"""
6
-import os
6
+
7
+from __future__ import annotations
8
+
9
+from typing import Any, Callable
10
11
import pytest
12
import ray
@@ -17,12 +20,18 @@ from ..utils import init_test_distributed_environment, multi_process_parallel
20
21
22
@ray.remote(num_gpus=1, max_calls=1)
20
-def all_reduce_test_worker(tp_size: int, pp_size: int, rank: int,
21
- distributed_init_port: str):
23
+def all_reduce_test_worker(
24
+ monkeypatch: pytest.MonkeyPatch,
25
+ tp_size: int,
26
+ pp_size: int,
27
+ rank: int,
28
+ distributed_init_port: str,
29
+):
30
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
31
# so that each worker can see all the GPUs
32
# they will be able to set the device to the correct GPU
25
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
33
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
34
+
35
device = torch.device(f"cuda:{rank}")
36
torch.cuda.set_device(device)
37
init_test_distributed_environment(tp_size, pp_size, rank,
@@ -39,12 +48,17 @@ def all_reduce_test_worker(tp_size: int, pp_size: int, rank: int,
48
49
50
@ray.remote(num_gpus=1, max_calls=1)
42
-def all_gather_test_worker(tp_size: int, pp_size: int, rank: int,
43
- distributed_init_port: str):
51
+def all_gather_test_worker(
52
+ monkeypatch: pytest.MonkeyPatch,
53
+ tp_size: int,
54
+ pp_size: int,
55
+ rank: int,
56
+ distributed_init_port: str,
57
+):
58
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
59
# so that each worker can see all the GPUs
60
# they will be able to set the device to the correct GPU
47
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
61
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
62
device = torch.device(f"cuda:{rank}")
63
torch.cuda.set_device(device)
64
init_test_distributed_environment(tp_size, pp_size, rank,
@@ -67,12 +81,17 @@ def all_gather_test_worker(tp_size: int, pp_size: int, rank: int,
81
82
83
@ray.remote(num_gpus=1, max_calls=1)
70
-def broadcast_tensor_dict_test_worker(tp_size: int, pp_size: int, rank: int,
71
- distributed_init_port: str):
84
+def broadcast_tensor_dict_test_worker(
85
+ monkeypatch: pytest.MonkeyPatch,
86
+ tp_size: int,
87
+ pp_size: int,
88
+ rank: int,
89
+ distributed_init_port: str,
90
+):
91
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
92
# so that each worker can see all the GPUs
93
# they will be able to set the device to the correct GPU
75
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
94
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
95
device = torch.device(f"cuda:{rank}")
96
torch.cuda.set_device(device)
97
init_test_distributed_environment(tp_size, pp_size, rank,
@@ -106,9 +125,14 @@ def broadcast_tensor_dict_test_worker(tp_size: int, pp_size: int, rank: int,
125
126
127
@ray.remote(num_gpus=1, max_calls=1)
109
-def send_recv_tensor_dict_test_worker(tp_size: int, pp_size: int, rank: int,
110
- distributed_init_port: str):
111
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
128
+def send_recv_tensor_dict_test_worker(
129
+ monkeypatch: pytest.MonkeyPatch,
130
+ tp_size: int,
131
+ pp_size: int,
132
+ rank: int,
133
+ distributed_init_port: str,
134
+):
135
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
136
device = torch.device(f"cuda:{rank}")
137
torch.cuda.set_device(device)
138
init_test_distributed_environment(tp_size, pp_size, rank,
@@ -146,9 +170,14 @@ def send_recv_tensor_dict_test_worker(tp_size: int, pp_size: int, rank: int,
170
171
172
@ray.remote(num_gpus=1, max_calls=1)
149
-def send_recv_test_worker(tp_size: int, pp_size: int, rank: int,
150
- distributed_init_port: str):
151
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
173
+def send_recv_test_worker(
174
+ monkeypatch: pytest.MonkeyPatch,
175
+ tp_size: int,
176
+ pp_size: int,
177
+ rank: int,
178
+ distributed_init_port: str,
179
+):
180
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
181
device = torch.device(f"cuda:{rank}")
182
torch.cuda.set_device(device)
183
init_test_distributed_environment(tp_size, pp_size, rank,
@@ -174,8 +203,12 @@ def send_recv_test_worker(tp_size: int, pp_size: int, rank: int,
203
all_reduce_test_worker, all_gather_test_worker,
204
broadcast_tensor_dict_test_worker
205
])
177
-def test_multi_process_tensor_parallel(tp_size, test_target):
178
- multi_process_parallel(tp_size, 1, test_target)
206
+def test_multi_process_tensor_parallel(
207
+ monkeypatch: pytest.MonkeyPatch,
208
+ tp_size: int,
209
+ test_target: Callable[..., Any],
210
+):
211
+ multi_process_parallel(monkeypatch, tp_size, 1, test_target)
212
213
214
@pytest.mark.skipif(torch.cuda.device_count() < 2,
@@ -183,8 +216,12 @@ def test_multi_process_tensor_parallel(tp_size, test_target):
216
@pytest.mark.parametrize("pp_size", [2])
217
@pytest.mark.parametrize(
218
"test_target", [send_recv_test_worker, send_recv_tensor_dict_test_worker])
186
-def test_multi_process_pipeline_parallel(pp_size, test_target):
187
- multi_process_parallel(1, pp_size, test_target)
219
+def test_multi_process_pipeline_parallel(
220
+ monkeypatch: pytest.MonkeyPatch,
221
+ pp_size: int,
222
+ test_target: Callable[..., Any],
223
+):
224
+ multi_process_parallel(monkeypatch, 1, pp_size, test_target)
225
226
227
@pytest.mark.skipif(torch.cuda.device_count() < 4,
@@ -197,5 +234,9 @@ def test_multi_process_pipeline_parallel(pp_size, test_target):
234
broadcast_tensor_dict_test_worker
235
])
236
def test_multi_process_tensor_parallel_pipeline_parallel(
200
- tp_size, pp_size, test_target):
201
- multi_process_parallel(tp_size, pp_size, test_target)
237
+ tp_size: int,
238
+ pp_size: int,
239
+ test_target: Callable[..., Any],
240
+ monkeypatch: pytest.MonkeyPatch,
241
+):
242
+ multi_process_parallel(monkeypatch, tp_size, pp_size, test_target)
tests/distributed/test_custom_all_reduce.py
+101
-82
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import os
3
import random
4
5
import pytest
@@ -23,95 +22,115 @@ for i, v in enumerate(test_sizes):
22
23
24
@ray.remote(num_gpus=1, max_calls=1)
26
-def graph_allreduce(tp_size, pp_size, rank, distributed_init_port):
27
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
28
- device = torch.device(f"cuda:{rank}")
29
- torch.cuda.set_device(device)
30
- init_test_distributed_environment(tp_size, pp_size, rank,
31
- distributed_init_port)
32
- ensure_model_parallel_initialized(tp_size, pp_size)
33
- group = get_tensor_model_parallel_group().device_group
34
-
35
- # A small all_reduce for warmup.
36
- # this is needed because device communicators might be created lazily
37
- # (e.g. NCCL). This will ensure that the communicator is initialized
38
- # before any communication happens, so that this group can be used for
39
- # graph capture immediately.
40
- data = torch.zeros(1)
41
- data = data.to(device=device)
42
- torch.distributed.all_reduce(data, group=group)
43
- torch.cuda.synchronize()
44
- del data
45
-
46
- # we use the first group to communicate once
47
- # and the second group to communicate twice
48
- # and so on
49
- # this is used to demonstrate that each group can
50
- # communicate independently
51
- num_communication = rank // tp_size + 1
52
-
53
- for sz in test_sizes:
54
- for dtype in [torch.float32, torch.float16, torch.bfloat16]:
55
- with graph_capture(device=device) as graph_capture_context:
56
- # use integers so result matches NCCL exactly
57
- inp1 = torch.randint(1,
58
- 16, (sz, ),
59
- dtype=dtype,
60
- device=torch.cuda.current_device())
61
- inp2 = torch.randint(1,
62
- 16, (sz, ),
63
- dtype=dtype,
64
- device=torch.cuda.current_device())
65
- torch.cuda.synchronize()
66
- graph = torch.cuda.CUDAGraph()
67
- with torch.cuda.graph(graph,
68
- stream=graph_capture_context.stream):
69
- for i in range(num_communication):
70
- out1 = tensor_model_parallel_all_reduce(inp1)
71
- # the input buffer is immediately modified to test
72
- # synchronization
73
- dist.all_reduce(inp1, group=group)
74
- out2 = tensor_model_parallel_all_reduce(inp2)
75
- dist.all_reduce(inp2, group=group)
76
- graph.replay()
77
- torch.testing.assert_close(out1, inp1)
78
- torch.testing.assert_close(out2, inp2)
25
+def graph_allreduce(
26
+ monkeypatch: pytest.MonkeyPatch,
27
+ tp_size,
28
+ pp_size,
29
+ rank,
30
+ distributed_init_port,
31
+):
32
+ with monkeypatch.context() as m:
33
+ m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
34
+ device = torch.device(f"cuda:{rank}")
35
+ torch.cuda.set_device(device)
36
+ init_test_distributed_environment(tp_size, pp_size, rank,
37
+ distributed_init_port)
38
+ ensure_model_parallel_initialized(tp_size, pp_size)
39
+ group = get_tensor_model_parallel_group().device_group
40
+
41
+ # A small all_reduce for warmup.
42
+ # this is needed because device communicators might be created lazily
43
+ # (e.g. NCCL). This will ensure that the communicator is initialized
44
+ # before any communication happens, so that this group can be used for
45
+ # graph capture immediately.
46
+ data = torch.zeros(1)
47
+ data = data.to(device=device)
48
+ torch.distributed.all_reduce(data, group=group)
49
+ torch.cuda.synchronize()
50
+ del data
51
+
52
+ # we use the first group to communicate once
53
+ # and the second group to communicate twice
54
+ # and so on
55
+ # this is used to demonstrate that each group can
56
+ # communicate independently
57
+ num_communication = rank // tp_size + 1
58
+
59
+ for sz in test_sizes:
60
+ for dtype in [torch.float32, torch.float16, torch.bfloat16]:
61
+ with graph_capture(device=device) as graph_capture_context:
62
+ # use integers so result matches NCCL exactly
63
+ inp1 = torch.randint(1,
64
+ 16, (sz, ),
65
+ dtype=dtype,
66
+ device=torch.cuda.current_device())
67
+ inp2 = torch.randint(1,
68
+ 16, (sz, ),
69
+ dtype=dtype,
70
+ device=torch.cuda.current_device())
71
+ torch.cuda.synchronize()
72
+ graph = torch.cuda.CUDAGraph()
73
+ with torch.cuda.graph(graph,
74
+ stream=graph_capture_context.stream):
75
+ for i in range(num_communication):
76
+ out1 = tensor_model_parallel_all_reduce(inp1)
77
+ # the input buffer is immediately modified to test
78
+ # synchronization
79
+ dist.all_reduce(inp1, group=group)
80
+ out2 = tensor_model_parallel_all_reduce(inp2)
81
+ dist.all_reduce(inp2, group=group)
82
+ graph.replay()
83
+ torch.testing.assert_close(out1, inp1)
84
+ torch.testing.assert_close(out2, inp2)
85
86
87
@ray.remote(num_gpus=1, max_calls=1)
82
-def eager_allreduce(tp_size, pp_size, rank, distributed_init_port):
83
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
84
- device = torch.device(f"cuda:{rank}")
85
- torch.cuda.set_device(device)
86
- init_test_distributed_environment(tp_size, pp_size, rank,
87
- distributed_init_port)
88
-
89
- # we use the first group to communicate once
90
- # and the second group to communicate twice
91
- # and so on
92
- # this is used to demonstrate that each group can
93
- # communicate independently
94
- num_communication = rank // tp_size + 1
95
- sz = 1024
96
- fa = get_tp_group().ca_comm
97
- inp = torch.ones(sz, dtype=torch.float32, device=device)
98
- out = inp
99
- for _ in range(num_communication):
100
- out = fa.all_reduce(out, registered=False)
101
- torch.testing.assert_close(out, inp * (tp_size**num_communication))
102
-
103
- inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device)
104
- out = inp
105
- for _ in range(num_communication):
106
- out = fa.all_reduce(out, registered=False)
107
- torch.testing.assert_close(out, inp * (tp_size**num_communication))
88
+def eager_allreduce(
89
+ monkeypatch: pytest.MonkeyPatch,
90
+ tp_size,
91
+ pp_size,
92
+ rank,
93
+ distributed_init_port,
94
+):
95
+ with monkeypatch.context() as m:
96
+ m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
97
+ device = torch.device(f"cuda:{rank}")
98
+ torch.cuda.set_device(device)
99
+ init_test_distributed_environment(tp_size, pp_size, rank,
100
+ distributed_init_port)
101
+
102
+ # we use the first group to communicate once
103
+ # and the second group to communicate twice
104
+ # and so on
105
+ # this is used to demonstrate that each group can
106
+ # communicate independently
107
+ num_communication = rank // tp_size + 1
108
+ sz = 1024
109
+ fa = get_tp_group().ca_comm
110
+ inp = torch.ones(sz, dtype=torch.float32, device=device)
111
+ out = inp
112
+ for _ in range(num_communication):
113
+ out = fa.all_reduce(out, registered=False)
114
+ torch.testing.assert_close(out, inp * (tp_size**num_communication))
115
+
116
+ inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device)
117
+ out = inp
118
+ for _ in range(num_communication):
119
+ out = fa.all_reduce(out, registered=False)
120
+ torch.testing.assert_close(out, inp * (tp_size**num_communication))
121
122
123
@pytest.mark.parametrize("tp_size", [2])
124
@pytest.mark.parametrize("pipeline_parallel_size", [1, 2])
125
@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
113
-def test_custom_allreduce(tp_size, pipeline_parallel_size, test_target):
126
+def test_custom_allreduce(
127
+ monkeypatch: pytest.MonkeyPatch,
128
+ tp_size,
129
+ pipeline_parallel_size,
130
+ test_target,
131
+):
132
world_size = tp_size * pipeline_parallel_size
133
if world_size > torch.cuda.device_count():
134
pytest.skip("Not enough GPUs to run the test.")
117
- multi_process_parallel(tp_size, pipeline_parallel_size, test_target)
135
+ multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size,
136
+ test_target)
tests/distributed/test_pipeline_partition.py
+35
-29
@@ -7,33 +7,35 @@ import pytest
7
from vllm.distributed.utils import get_pp_indices
8
9
10
-def test_custom_layer_partition():
11
-
12
- def _verify(partition_str, num_layers, pp_size, goldens):
13
- bak = os.environ.get("VLLM_PP_LAYER_PARTITION", None)
14
- os.environ["VLLM_PP_LAYER_PARTITION"] = partition_str
15
- for pp_rank, golden in enumerate(goldens):
16
- assert get_pp_indices(num_layers, pp_rank, pp_size) == golden
17
- if bak is not None:
18
- os.environ["VLLM_PP_LAYER_PARTITION"] = bak
19
-
20
- # Even partition
21
- _verify("5,5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
22
- # Balanced partition
23
- _verify("4,6,6,4", 20, 4, [(0, 4), (4, 10), (10, 16), (16, 20)])
24
- # Put reminder somewhere
25
- _verify("5,6,5,6", 22, 4, [(0, 5), (5, 11), (11, 16), (16, 22)])
26
- # Invalid partition strings
27
- with pytest.raises(ValueError):
28
- _verify("5,5,5,5,", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
29
- with pytest.raises(ValueError):
30
- _verify("5,5,5,a", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
31
- # Wrong number of partitions
32
- with pytest.raises(ValueError):
33
- _verify("5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
34
- # Wrong number of layers
35
- with pytest.raises(ValueError):
36
- _verify("5,5,5,5", 21, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
10
+def test_custom_layer_partition(monkeypatch: pytest.MonkeyPatch):
11
+
12
+ with monkeypatch.context() as m:
13
+
14
+ def _verify(partition_str, num_layers, pp_size, goldens):
15
+ bak = os.environ.get("VLLM_PP_LAYER_PARTITION", None)
16
+ m.setenv("VLLM_PP_LAYER_PARTITION", partition_str)
17
+ for pp_rank, golden in enumerate(goldens):
18
+ assert get_pp_indices(num_layers, pp_rank, pp_size) == golden
19
+ if bak is not None:
20
+ m.setenv("VLLM_PP_LAYER_PARTITION", bak)
21
+
22
+ # Even partition
23
+ _verify("5,5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
24
+ # Balanced partition
25
+ _verify("4,6,6,4", 20, 4, [(0, 4), (4, 10), (10, 16), (16, 20)])
26
+ # Put reminder somewhere
27
+ _verify("5,6,5,6", 22, 4, [(0, 5), (5, 11), (11, 16), (16, 22)])
28
+ # Invalid partition strings
29
+ with pytest.raises(ValueError):
30
+ _verify("5,5,5,5,", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
31
+ with pytest.raises(ValueError):
32
+ _verify("5,5,5,a", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
33
+ # Wrong number of partitions
34
+ with pytest.raises(ValueError):
35
+ _verify("5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
36
+ # Wrong number of layers
37
+ with pytest.raises(ValueError):
38
+ _verify("5,5,5,5", 21, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
39
40
41
@pytest.mark.parametrize(
@@ -55,6 +57,10 @@ def test_custom_layer_partition():
57
(5, 3, 1, (2, 4)),
58
(5, 3, 2, (4, 5)),
59
])
58
-def test_uneven_auto_partition(num_hidden_layers: int, pp_size: int,
59
- pp_rank: int, indices: tuple[int, int]):
60
+def test_uneven_auto_partition(
61
+ num_hidden_layers: int,
62
+ pp_size: int,
63
+ pp_rank: int,
64
+ indices: tuple[int, int],
65
+):
66
assert indices == get_pp_indices(num_hidden_layers, pp_rank, pp_size)
tests/distributed/test_pp_cudagraph.py
+26
-16
@@ -1,11 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from __future__ import annotations
3
3
-import os
4
+from typing import TYPE_CHECKING
5
6
import pytest
7
8
from ..utils import compare_two_settings, fork_new_process_for_each_test
9
10
+if TYPE_CHECKING:
11
+ from typing_extensions import LiteralString
12
+
13
14
@pytest.mark.parametrize("PP_SIZE, MODEL_NAME", [
15
(2, "JackFram/llama-160m"),
@@ -15,18 +19,24 @@ from ..utils import compare_two_settings, fork_new_process_for_each_test
19
"FLASHINFER",
20
])
21
@fork_new_process_for_each_test
18
-def test_pp_cudagraph(PP_SIZE, MODEL_NAME, ATTN_BACKEND):
19
- cudagraph_args = [
20
- # use half precision for speed and memory savings in CI environment
21
- "--dtype",
22
- "float16",
23
- "--pipeline-parallel-size",
24
- str(PP_SIZE),
25
- "--distributed-executor-backend",
26
- "mp",
27
- ]
28
- os.environ["VLLM_ATTENTION_BACKEND"] = ATTN_BACKEND
29
-
30
- eager_args = cudagraph_args + ["--enforce-eager"]
31
-
32
- compare_two_settings(MODEL_NAME, eager_args, cudagraph_args)
22
+def test_pp_cudagraph(
23
+ monkeypatch: pytest.MonkeyPatch,
24
+ PP_SIZE: int,
25
+ MODEL_NAME: str,
26
+ ATTN_BACKEND: LiteralString,
27
+):
28
+ with monkeypatch.context() as m:
29
+ cudagraph_args = [
30
+ # use half precision for speed and memory savings in CI environment
31
+ "--dtype",
32
+ "float16",
33
+ "--pipeline-parallel-size",
34
+ str(PP_SIZE),
35
+ "--distributed-executor-backend",
36
+ "mp",
37
+ ]
38
+ m.setenv("VLLM_ATTENTION_BACKEND", ATTN_BACKEND)
39
+
40
+ eager_args = cudagraph_args + ["--enforce-eager"]
41
+
42
+ compare_two_settings(MODEL_NAME, eager_args, cudagraph_args)
tests/entrypoints/llm/test_accuracy.py
+2
-2
@@ -49,7 +49,7 @@ TPU_TP_TEST_STR = "" #"tensor_parallel_size=4"
49
@pytest.mark.skipif(not current_platform.is_cuda()
50
and not current_platform.is_tpu(),
51
reason="V1 is currently only supported on CUDA and TPU")
52
-def test_lm_eval_accuracy_v1_engine(monkeypatch):
52
+def test_lm_eval_accuracy_v1_engine(monkeypatch: pytest.MonkeyPatch):
53
"""Run with the V1 Engine."""
54
55
with monkeypatch.context() as m:
@@ -67,7 +67,7 @@ def test_lm_eval_accuracy_v1_engine(monkeypatch):
67
run_test(more_args)
68
69
70
-def test_lm_eval_accuracy_v0_engine(monkeypatch):
70
+def test_lm_eval_accuracy_v0_engine(monkeypatch: pytest.MonkeyPatch):
71
"""Run with the V0 Engine."""
72
73
with monkeypatch.context() as m:
tests/entrypoints/offline_mode/test_offline_mode.py
+27
-22
@@ -53,32 +53,37 @@ def cache_models():
53
54
@pytest.mark.skip_global_cleanup
55
@pytest.mark.usefixtures("cache_models")
56
-def test_offline_mode(monkeypatch):
56
+def test_offline_mode(monkeypatch: pytest.MonkeyPatch):
57
# Set HF to offline mode and ensure we can still construct an LLM
58
- try:
59
- monkeypatch.setenv("HF_HUB_OFFLINE", "1")
60
- monkeypatch.setenv("VLLM_NO_USAGE_STATS", "1")
58
+ with monkeypatch.context() as m:
59
+ try:
60
+ m.setenv("HF_HUB_OFFLINE", "1")
61
+ m.setenv("VLLM_NO_USAGE_STATS", "1")
62
62
- def disable_connect(*args, **kwargs):
63
- raise RuntimeError("No http calls allowed")
63
+ def disable_connect(*args, **kwargs):
64
+ raise RuntimeError("No http calls allowed")
65
65
- monkeypatch.setattr(urllib3.connection.HTTPConnection, "connect",
66
- disable_connect)
67
- monkeypatch.setattr(urllib3.connection.HTTPSConnection, "connect",
68
- disable_connect)
66
+ m.setattr(
67
+ urllib3.connection.HTTPConnection,
68
+ "connect",
69
+ disable_connect,
70
+ )
71
+ m.setattr(
72
+ urllib3.connection.HTTPSConnection,
73
+ "connect",
74
+ disable_connect,
75
+ )
76
70
- # Need to re-import huggingface_hub and friends to setup offline mode
71
- _re_import_modules()
72
- # Cached model files should be used in offline mode
73
- for model_config in MODEL_CONFIGS:
74
- LLM(**model_config)
75
- finally:
76
- # Reset the environment after the test
77
- # NB: Assuming tests are run in online mode
78
- monkeypatch.delenv("HF_HUB_OFFLINE")
79
- monkeypatch.delenv("VLLM_NO_USAGE_STATS")
80
- _re_import_modules()
81
- pass
77
+ # Need to re-import huggingface_hub
78
+ # and friends to setup offline mode
79
+ _re_import_modules()
80
+ # Cached model files should be used in offline mode
81
+ for model_config in MODEL_CONFIGS:
82
+ LLM(**model_config)
83
+ finally:
84
+ # Reset the environment after the test
85
+ # NB: Assuming tests are run in online mode
86
+ _re_import_modules()
87
88
89
def _re_import_modules():
tests/entrypoints/openai/correctness/test_lmeval.py
+3
-2
@@ -70,7 +70,7 @@ def run_test(more_args):
70
@pytest.mark.skipif(not current_platform.is_cuda()
71
and not current_platform.is_tpu(),
72
reason="V1 currently only supported on CUDA and TPU")
73
-def test_lm_eval_accuracy_v1_engine(monkeypatch):
73
+def test_lm_eval_accuracy_v1_engine(monkeypatch: pytest.MonkeyPatch):
74
"""Run with the V1 Engine."""
75
76
with monkeypatch.context() as m:
@@ -85,7 +85,8 @@ def test_lm_eval_accuracy_v1_engine(monkeypatch):
85
86
87
@pytest.mark.parametrize("more_args", MORE_ARGS_LIST)
88
-def test_lm_eval_accuracy_v0_engine(monkeypatch, more_args):
88
+def test_lm_eval_accuracy_v0_engine(monkeypatch: pytest.MonkeyPatch,
89
+ more_args):
90
"""Run with the V0 Engine."""
91
92
with monkeypatch.context() as m:
tests/kernels/test_attention_selector.py
+77
-54
@@ -5,13 +5,12 @@ from unittest.mock import Mock, patch
5
import pytest
6
import torch
7
8
-from tests.kernels.utils import override_backend_env_variable
8
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
9
from vllm.platforms.cpu import CpuPlatform
10
from vllm.platforms.cuda import CudaPlatform
11
from vllm.platforms.openvino import OpenVinoPlatform
12
from vllm.platforms.rocm import RocmPlatform
14
-from vllm.utils import STR_FLASH_ATTN_VAL, STR_INVALID_VAL
13
+from vllm.utils import STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL, STR_INVALID_VAL
14
15
16
@pytest.fixture(autouse=True)
@@ -25,87 +24,111 @@ def clear_cache():
24
"name", ["TORCH_SDPA", "ROCM_FLASH", "XFORMERS", "FLASHINFER", "OPENVINO"])
25
@pytest.mark.parametrize("use_v1", [True, False])
26
@pytest.mark.parametrize("device", ["cpu", "openvino", "hip", "cuda"])
28
-def test_env(name: str, use_v1: bool, device: str, monkeypatch):
27
+def test_env(
28
+ name: str,
29
+ use_v1: bool,
30
+ device: str,
31
+ monkeypatch: pytest.MonkeyPatch,
32
+):
33
"""Test that the attention selector can be set via environment variable.
34
Note that we do not test FlashAttn because it is the default backend.
35
"""
36
33
- monkeypatch.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
34
- override_backend_env_variable(monkeypatch, name)
35
-
36
- if device == "cpu":
37
- with patch("vllm.attention.selector.current_platform", CpuPlatform()):
38
- backend = get_attn_backend(16, torch.float16, torch.float16, 16,
39
- False)
40
- assert backend.get_name() == "TORCH_SDPA"
41
- elif device == "hip":
42
- with patch("vllm.attention.selector.current_platform", RocmPlatform()):
43
- backend = get_attn_backend(16, torch.float16, torch.float16, 16,
44
- False)
45
- EXPECTED = "ROCM_ATTN_VLLM_V1" if use_v1 else "ROCM_FLASH"
46
- assert backend.get_name() == EXPECTED
47
- elif device == "openvino":
48
- with patch("vllm.attention.selector.current_platform",
49
- OpenVinoPlatform()), patch.dict('sys.modules',
50
- {'openvino': Mock()}):
51
- backend = get_attn_backend(16, torch.float16, torch.float16, 16,
52
- False)
53
- assert backend.get_name() == "OPENVINO"
54
- else:
55
- if name in ["XFORMERS", "FLASHINFER"]:
37
+ with monkeypatch.context() as m:
38
+ m.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
39
+ m.setenv(STR_BACKEND_ENV_VAR, name)
40
+
41
+ if device == "cpu":
42
+ with patch("vllm.attention.selector.current_platform",
43
+ CpuPlatform()):
44
+ backend = get_attn_backend(16, torch.float16, torch.float16,
45
+ 16, False)
46
+ assert backend.get_name() == "TORCH_SDPA"
47
+ elif device == "hip":
48
with patch("vllm.attention.selector.current_platform",
57
- CudaPlatform()):
49
+ RocmPlatform()):
50
backend = get_attn_backend(16, torch.float16, torch.float16,
51
16, False)
60
- EXPECTED = "FLASH_ATTN_VLLM_V1" if use_v1 else name
52
+ EXPECTED = "ROCM_ATTN_VLLM_V1" if use_v1 else "ROCM_FLASH"
53
assert backend.get_name() == EXPECTED
54
+ elif device == "openvino":
55
+ with patch("vllm.attention.selector.current_platform",
56
+ OpenVinoPlatform()), patch.dict('sys.modules',
57
+ {'openvino': Mock()}):
58
+ backend = get_attn_backend(16, torch.float16, torch.float16,
59
+ 16, False)
60
+ assert backend.get_name() == "OPENVINO"
61
+ else:
62
+ if name in ["XFORMERS", "FLASHINFER"]:
63
+ with patch("vllm.attention.selector.current_platform",
64
+ CudaPlatform()):
65
+ backend = get_attn_backend(16, torch.float16,
66
+ torch.float16, 16, False)
67
+ EXPECTED = "FLASH_ATTN_VLLM_V1" if use_v1 else name
68
+ assert backend.get_name() == EXPECTED
69
70
64
-def test_flash_attn(monkeypatch):
71
+def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
72
"""Test FlashAttn validation."""
73
# TODO: When testing for v1, pipe in `use_v1` as an argument to
74
# get_attn_backend
75
69
- override_backend_env_variable(monkeypatch, STR_FLASH_ATTN_VAL)
76
+ with monkeypatch.context() as m:
77
+ m.setenv(STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL)
78
71
- # Unsupported CUDA arch
72
- with patch("torch.cuda.get_device_capability", return_value=(7, 5)):
79
+ # Unsupported CUDA arch
80
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda:
81
+ (7, 5))
82
backend = get_attn_backend(16, torch.float16, None, 16, False)
83
assert backend.get_name() != STR_FLASH_ATTN_VAL
84
76
- # Unsupported data type
77
- backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16, False)
78
- assert backend.get_name() != STR_FLASH_ATTN_VAL
85
+ # Reset the monkeypatch for subsequent tests
86
+ monkeypatch.undo()
87
80
- # Unsupported kv cache data type
81
- backend = get_attn_backend(16, torch.float16, "fp8", 16, False)
82
- assert backend.get_name() != STR_FLASH_ATTN_VAL
88
+ # Unsupported data type
89
+ backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16, False)
90
+ assert backend.get_name() != STR_FLASH_ATTN_VAL
91
84
- # Unsupported block size
85
- backend = get_attn_backend(16, torch.float16, None, 8, False)
86
- assert backend.get_name() != STR_FLASH_ATTN_VAL
92
+ # Unsupported kv cache data type
93
+ backend = get_attn_backend(16, torch.float16, "fp8", 16, False)
94
+ assert backend.get_name() != STR_FLASH_ATTN_VAL
95
88
- # flash-attn is not installed
89
- with patch.dict('sys.modules', {'vllm_flash_attn': None}):
96
+ # Unsupported block size
97
+ backend = get_attn_backend(16, torch.float16, None, 8, False)
98
+ assert backend.get_name() != STR_FLASH_ATTN_VAL
99
+
100
+ # flash-attn is not installed
101
+ import sys
102
+ original_module = sys.modules.get('vllm_flash_attn')
103
+ monkeypatch.setitem(sys.modules, 'vllm_flash_attn', None)
104
backend = get_attn_backend(16, torch.float16, None, 16, False)
105
assert backend.get_name() != STR_FLASH_ATTN_VAL
106
93
- # Unsupported head size
94
- backend = get_attn_backend(17, torch.float16, None, 16, False)
95
- assert backend.get_name() != STR_FLASH_ATTN_VAL
107
+ # Restore the original module if it existed
108
+ if original_module is not None:
109
+ monkeypatch.setitem(sys.modules, 'vllm_flash_attn',
110
+ original_module)
111
+ else:
112
+ monkeypatch.delitem(sys.modules, 'vllm_flash_attn', raising=False)
113
97
- # Attention-free models should bypass env and use PlaceholderAttention
98
- backend = get_attn_backend(16, torch.float16, torch.float16, 16, True)
99
- assert backend.get_name() != STR_FLASH_ATTN_VAL
114
+ # Unsupported head size
115
+ backend = get_attn_backend(17, torch.float16, None, 16, False)
116
+ assert backend.get_name() != STR_FLASH_ATTN_VAL
117
+
118
+ # Attention-free models should bypass env and use PlaceholderAttention
119
+ backend = get_attn_backend(16, torch.float16, torch.float16, 16, True)
120
+ assert backend.get_name() != STR_FLASH_ATTN_VAL
121
122
123
@pytest.mark.parametrize("use_v1", [True, False])
103
-def test_invalid_env(use_v1: bool, monkeypatch):
104
- """Ignore the invalid env variable if it is set."""
105
- monkeypatch.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
106
- override_backend_env_variable(monkeypatch, STR_INVALID_VAL)
124
+def test_invalid_env(use_v1: bool, monkeypatch: pytest.MonkeyPatch):
125
+
126
+ with monkeypatch.context() as m, patch(
127
+ "vllm.attention.selector.current_platform", CudaPlatform()):
128
+ m.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
129
+ m.setenv(STR_BACKEND_ENV_VAR, STR_INVALID_VAL)
130
108
- with patch("vllm.attention.selector.current_platform", CudaPlatform()):
131
+ # Test with head size 32
132
backend = get_attn_backend(32, torch.float16, None, 16, False)
133
EXPECTED = "FLASH_ATTN_VLLM_V1" if use_v1 else "FLASH_ATTN"
134
assert backend.get_name() == EXPECTED
tests/kernels/test_awq.py
+30
-30
@@ -1,7 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import os
4
-
3
import pytest
4
import torch
5
@@ -11,36 +9,38 @@ from vllm import _custom_ops as ops # noqa: F401
9
10
@pytest.mark.skipif(not hasattr(torch.ops._C, "awq_dequantize"),
11
reason="AWQ is not supported on this GPU type.")
14
-def test_awq_dequantize_opcheck():
15
- os.environ["VLLM_USE_TRITON_AWQ"] = "0"
16
- qweight = torch.randint(-2000000000,
17
- 2000000000, (8192, 256),
18
- device='cuda',
19
- dtype=torch.int32)
20
- scales = torch.rand((64, 2048), device='cuda', dtype=torch.float16)
21
- zeros = torch.empty((64, 256), device='cuda', dtype=torch.int32)
22
- split_k_iters = 0
23
- thx = 0
24
- thy = 0
25
- opcheck(torch.ops._C.awq_dequantize,
26
- (qweight, scales, zeros, split_k_iters, thx, thy))
12
+def test_awq_dequantize_opcheck(monkeypatch: pytest.MonkeyPatch):
13
+ with monkeypatch.context() as m:
14
+ m.setenv("VLLM_USE_TRITON_AWQ", "0")
15
+ qweight = torch.randint(-2000000000,
16
+ 2000000000, (8192, 256),
17
+ device='cuda',
18
+ dtype=torch.int32)
19
+ scales = torch.rand((64, 2048), device='cuda', dtype=torch.float16)
20
+ zeros = torch.empty((64, 256), device='cuda', dtype=torch.int32)
21
+ split_k_iters = 0
22
+ thx = 0
23
+ thy = 0
24
+ opcheck(torch.ops._C.awq_dequantize,
25
+ (qweight, scales, zeros, split_k_iters, thx, thy))
26
27
28
@pytest.mark.skip(reason="Not working; needs investigation.")
29
@pytest.mark.skipif(not hasattr(torch.ops._C, "awq_gemm"),
30
reason="AWQ is not supported on this GPU type.")
32
-def test_awq_gemm_opcheck():
33
- os.environ["VLLM_USE_TRITON_AWQ"] = "0"
34
- input = torch.rand((2, 8192), device='cuda', dtype=torch.float16)
35
- qweight = torch.randint(-2000000000,
36
- 2000000000, (8192, 256),
37
- device='cuda',
38
- dtype=torch.int32)
39
- scales = torch.randint(-2000000000,
40
- 2000000000, (64, 256),
41
- device='cuda',
42
- dtype=torch.int32)
43
- qzeros = torch.empty((64, 2048), device='cuda', dtype=torch.float16)
44
- split_k_iters = 8
45
- opcheck(torch.ops._C.awq_gemm,
46
- (input, qweight, qzeros, scales, split_k_iters))
31
+def test_awq_gemm_opcheck(monkeypatch: pytest.MonkeyPatch):
32
+ with monkeypatch.context() as m:
33
+ m.setenv("VLLM_USE_TRITON_AWQ", "0")
34
+ input = torch.rand((2, 8192), device='cuda', dtype=torch.float16)
35
+ qweight = torch.randint(-2000000000,
36
+ 2000000000, (8192, 256),
37
+ device='cuda',
38
+ dtype=torch.int32)
39
+ scales = torch.randint(-2000000000,
40
+ 2000000000, (64, 256),
41
+ device='cuda',
42
+ dtype=torch.int32)
43
+ qzeros = torch.empty((64, 2048), device='cuda', dtype=torch.float16)
44
+ split_k_iters = 8
45
+ opcheck(torch.ops._C.awq_gemm,
46
+ (input, qweight, qzeros, scales, split_k_iters))
tests/kernels/test_rocm_attention_selector.py
+10
-8
@@ -1,13 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from unittest.mock import patch
4
-
3
import pytest
4
import torch
5
8
-from tests.kernels.utils import override_backend_env_variable
6
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
7
from vllm.platforms.rocm import RocmPlatform
8
+from vllm.utils import STR_BACKEND_ENV_VAR
9
10
11
@pytest.fixture(autouse=True)
@@ -17,15 +15,19 @@ def clear_cache():
15
_cached_get_attn_backend.cache_clear()
16
17
20
-def test_selector(monkeypatch):
21
- """Test that the attention selector for ROCm.
22
- """
23
- override_backend_env_variable(monkeypatch, "ROCM_FLASH")
18
+def test_selector(monkeypatch: pytest.MonkeyPatch):
19
+ with monkeypatch.context() as m:
20
+ m.setenv(STR_BACKEND_ENV_VAR, "ROCM_FLASH")
21
25
- with patch("vllm.attention.selector.current_platform", RocmPlatform()):
22
+ # Set the current platform to ROCm using monkeypatch
23
+ monkeypatch.setattr("vllm.attention.selector.current_platform",
24
+ RocmPlatform())
25
+
26
+ # Test standard ROCm attention
27
backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
28
assert (backend.get_name() == "ROCM_FLASH"
29
or backend.get_name() == "ROCM_ATTN_VLLM_V1")
30
+
31
# mla test for deepseek related
32
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
33
False, True)
tests/kernels/utils.py
+32
-32
@@ -36,12 +36,12 @@ ALL_OPCHECK_TEST_UTILS: tuple[str, ...] = (
36
37
class QKVInputs(NamedTuple):
38
'''
39
- Data structure for representing unpacked attention inputs,
39
+ Data structure for representing unpacked attention inputs,
40
query/key/values and their sequence lengths.
41
42
Attributes:
43
44
- * {query,key,value}: unpacked (batch_size x padded_seq_len x
44
+ * {query,key,value}: unpacked (batch_size x padded_seq_len x
45
num_heads x head_size) attention inputs
46
* q_seq_lens: query sequence lengths list
47
* kv_seq_lens: shared key/value sequence lengths list
@@ -56,14 +56,14 @@ class QKVInputs(NamedTuple):
56
57
class QKVO(NamedTuple):
58
'''
59
- Data structure for representing unpacked attention inputs,
59
+ Data structure for representing unpacked attention inputs,
60
alongside unpacked known-correct attention output
61
62
Attributes:
63
64
- * qkv: unpacked (batch_size x padded_seq_len x
64
+ * qkv: unpacked (batch_size x padded_seq_len x
65
num_heads x head_size) attention inputs
66
- * ideal_output: unpacked (batch_size x padded_seq_len x
66
+ * ideal_output: unpacked (batch_size x padded_seq_len x
67
num_heads x head_size) known-correct attention output
68
'''
69
@@ -77,7 +77,7 @@ class PackedQKVInputs(NamedTuple):
77
78
Attributes:
79
80
- * {query,key,value}: packed (number_of_tokens x num_heads
80
+ * {query,key,value}: packed (number_of_tokens x num_heads
81
x head_size) attention inputs
82
* q_start_loc_list: list of query start locations within packed tensor
83
* kv_start_loc_list: shared list of key/value start locations within
@@ -97,14 +97,14 @@ class PackedQKVInputs(NamedTuple):
97
98
class PackedQKVO(NamedTuple):
99
'''
100
- Data structure for representing packed attention inputs,
100
+ Data structure for representing packed attention inputs,
101
alongside packed known-correct attention output
102
103
Attributes:
104
105
- * packed_qkv: packed (number_of_tokens x num_heads
105
+ * packed_qkv: packed (number_of_tokens x num_heads
106
x head_size) attention inputs
107
- * ideal_output: packed (number_of_tokens x num_heads
107
+ * ideal_output: packed (number_of_tokens x num_heads
108
x head_size) known-correct attention output
109
'''
110
@@ -134,7 +134,7 @@ class PhaseTestParameters(NamedTuple):
134
135
Attributes:
136
137
- * packed_qkvo: packed (number_of_tokens x num_heads
137
+ * packed_qkvo: packed (number_of_tokens x num_heads
138
x head_size) attention inputs & known-correct
139
output
140
* kv_mmap: KV cache memory mapping, specific to this test phase &
@@ -195,7 +195,7 @@ def make_causal_mask(
195
Create a q_max_seq_len x kv_max_seq_len causal mask
196
197
Arguments:
198
-
198
+
199
* q_max_seq_len: query max seq len
200
* kv_max_seq_len: key/value max seq len
201
@@ -320,9 +320,9 @@ def make_qkv(
320
* max_kv_seq_len: max key/value seq len
321
* num_heads
322
* head_size
323
- * is_encoder_decoder_attn: if True, query seqlen may differ from
324
- key/value seqlen (as is often the case for cross-attention);
325
- o/w, query/key/value seqlens match at each batch index
323
+ * is_encoder_decoder_attn: if True, query seqlen may differ from
324
+ key/value seqlen (as is often the case for cross-attention);
325
+ o/w, query/key/value seqlens match at each batch index
326
(max_kv_seq_len is unused)
327
* force_kv_seq_lens: if not None, overrides kv sequence lengths
328
* attn_type: encoder, decoder self, or enc/dec cross attention
@@ -469,7 +469,7 @@ def pack_qkv(qkv: QKVInputs, device: Union[torch.device,
469
Individually pack each of Q, K and V, each with dimensions batch_size x
470
padded_seq_len x num_heads x head_size, into respective number_of_tokens x
471
num_heads x head_size tensors.
472
-
472
+
473
For Q, number_of_tokens = sum(q_seq_lens).
474
475
For K and V, number_of_tokens = sum(kv_seq_lens)
@@ -619,9 +619,9 @@ def make_kv_cache(num_blocks: int,
619
Returns:
620
621
* kv_cache: 2 x num_blocks x (block_size * num_heads * head_size)
622
- * for backend 'XFORMERS'
622
+ * for backend 'XFORMERS'
623
* kv_cache: 2 x num_blocks x block_size x num_heads x head_size
624
- * for backend 'FLASH_ATTN'
624
+ * for backend 'FLASH_ATTN'
625
'''
626
if backend == 'XFORMERS':
627
kv_cache = torch.rand(
@@ -662,20 +662,20 @@ def split_slot_mapping(slot_mapping_list: torch.Tensor, seq_lens: list[int],
662
Context:
663
* Your goal is to test (1) prefill of N prompts, with prompt-lengths
664
{K_i \\forall i \\in [0,N)}, followed by (2) decoding of a single token
665
- for all N prompts (N tokens total); the resultant sequence lengths
665
+ for all N prompts (N tokens total); the resultant sequence lengths
666
after decode would be {K_i + 1 for i \\in [0,N)}
667
- * The test you want to do requires (1) having the prefill slot mapping
668
- for all tokens present during prefill, the number of which is
669
- M = \\sum_i{K_i}, and (2) having the decode slot mapping for all N
667
+ * The test you want to do requires (1) having the prefill slot mapping
668
+ for all tokens present during prefill, the number of which is
669
+ M = \\sum_i{K_i}, and (2) having the decode slot mapping for all N
670
decoded tokens
671
-
672
- This function consumes a single 1D slot mapping, which is the
671
+
672
+ This function consumes a single 1D slot mapping, which is the
673
concatenation of N slot mappings each of length K_i + 1 (corresponding
674
to the sequence lengths after decode), with a total length of
675
P = \\sum_i{K_i + 1} = M + N
676
677
The prefill-phase slot mapping results from excising the (K_i + 1)-th entry
678
- from each of the N subsequences in the slot mapping (i.e. omitting the
678
+ from each of the N subsequences in the slot mapping (i.e. omitting the
679
decoded token's mapping.)
680
681
The N excised entries are appended to obtain the decode-phase slot mapping
@@ -684,15 +684,15 @@ def split_slot_mapping(slot_mapping_list: torch.Tensor, seq_lens: list[int],
684
685
* slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N
686
post-decode sequences
687
- * seq_lens: list of N post-decode sequence lengths (K_i + 1 in the
687
+ * seq_lens: list of N post-decode sequence lengths (K_i + 1 in the
688
description above)
689
* device: cuda, cpu, etc.
690
691
Returns:
692
693
- * prefill_slot_mapping: Length-M 1D slot mapping (as Tensor)
693
+ * prefill_slot_mapping: Length-M 1D slot mapping (as Tensor)
694
reflecting all N prefill prompts
695
- * decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting
695
+ * decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting
696
all N decoded tokens
697
'''
698
@@ -725,7 +725,7 @@ def make_block_tables_slot_mapping(
725
726
Then the minimum KV cache size in blocks is
727
728
- total_cache_blocks = sum(num_blocks for all seqs)
728
+ total_cache_blocks = sum(num_blocks for all seqs)
729
730
Then, the blocktable mapping counts downward from
731
@@ -734,7 +734,7 @@ def make_block_tables_slot_mapping(
734
to
735
736
block_base_addr
737
-
737
+
738
739
The constructed block-tables and slot-mapping are sized to the
740
lengths of the sequences in their entirety (as reflected by seq_lens),
@@ -749,7 +749,7 @@ def make_block_tables_slot_mapping(
749
750
Return:
751
752
- * block_tables_tensor: block table for sequence
752
+ * block_tables_tensor: block table for sequence
753
* slot_mapping_list: slot mapping for sequence
754
* max_block_idx: the highest block address within this block table
755
'''
@@ -807,7 +807,7 @@ def make_test_metadata(
807
encoder_test_params and cross_test_params arguments allow encoder
808
attention and enc/dec cross-attention (respectively) to use distinct
809
metadata values from decoder self-attention (decoder_test_params.)
810
-
810
+
811
if encoder_test_params and cross_test_params are None, the attention
812
metadata will support decoder-only scenario.
813
@@ -820,7 +820,7 @@ def make_test_metadata(
820
* attn_backend_name: Backend for sourcing attention kernels
821
* is_prompt: prefill if True, o/w decode
822
* seq_lens: list of token counts for each sequence
823
- * decoder_test_params: decoder self-attention test params;
823
+ * decoder_test_params: decoder self-attention test params;
824
this function requires
825
kv_mmap (memory mapping) field
826
* device: CPU or CUDA device
tests/kv_transfer/test_disagg.py
renamed
tests/kv_transfer/test_module.py
renamed
tests/models/decoder_only/language/test_fp8.py
+70
-66
@@ -12,11 +12,10 @@ import pytest
12
from tests.kernels.utils import override_backend_env_variable
13
from tests.quantization.utils import is_quant_method_supported
14
from vllm.platforms import current_platform
15
+from vllm.utils import STR_BACKEND_ENV_VAR
16
17
from ...utils import check_logprobs_close
18
18
-os.environ["TOKENIZERS_PARALLELISM"] = "true"
19
-
19
20
@pytest.mark.quant_model
21
@pytest.mark.skipif(not is_quant_method_supported("fp8"),
@@ -55,45 +54,47 @@ def test_models(
54
backend: str,
55
tensor_parallel_size: int,
56
disable_async_output_proc: bool,
58
- monkeypatch,
57
+ monkeypatch: pytest.MonkeyPatch,
58
) -> None:
59
"""
60
Only checks log probs match to cover the discrepancy in
61
numerical sensitive kernels.
62
"""
64
- override_backend_env_variable(monkeypatch, backend)
65
-
66
- MAX_MODEL_LEN = 1024
67
- NUM_LOG_PROBS = 8
68
-
69
- with vllm_runner(
70
- base_model,
71
- max_model_len=MAX_MODEL_LEN,
72
- tensor_parallel_size=tensor_parallel_size,
73
- enforce_eager=enforce_eager,
74
- kv_cache_dtype="auto",
75
- disable_async_output_proc=disable_async_output_proc,
76
- ) as vllm_model:
77
- baseline_outputs = vllm_model.generate_greedy_logprobs(
78
- example_prompts, max_tokens, NUM_LOG_PROBS)
79
-
80
- with vllm_runner(
81
- test_model,
82
- max_model_len=MAX_MODEL_LEN,
83
- tensor_parallel_size=tensor_parallel_size,
84
- enforce_eager=enforce_eager,
85
- kv_cache_dtype=kv_cache_dtype,
86
- disable_async_output_proc=disable_async_output_proc,
87
- ) as vllm_model:
88
- test_outputs = vllm_model.generate_greedy_logprobs(
89
- example_prompts, max_tokens, NUM_LOG_PROBS)
90
-
91
- check_logprobs_close(
92
- outputs_0_lst=baseline_outputs,
93
- outputs_1_lst=test_outputs,
94
- name_0="fp16_kv_cache",
95
- name_1="fp8_kv_cache",
96
- )
63
+ with monkeypatch.context() as m:
64
+ m.setenv("TOKENIZERS_PARALLELISM", 'true')
65
+ m.setenv(STR_BACKEND_ENV_VAR, backend)
66
+
67
+ MAX_MODEL_LEN = 1024
68
+ NUM_LOG_PROBS = 8
69
+
70
+ with vllm_runner(
71
+ base_model,
72
+ max_model_len=MAX_MODEL_LEN,
73
+ tensor_parallel_size=tensor_parallel_size,
74
+ enforce_eager=enforce_eager,
75
+ kv_cache_dtype="auto",
76
+ disable_async_output_proc=disable_async_output_proc,
77
+ ) as vllm_model:
78
+ baseline_outputs = vllm_model.generate_greedy_logprobs(
79
+ example_prompts, max_tokens, NUM_LOG_PROBS)
80
+
81
+ with vllm_runner(
82
+ test_model,
83
+ max_model_len=MAX_MODEL_LEN,
84
+ tensor_parallel_size=tensor_parallel_size,
85
+ enforce_eager=enforce_eager,
86
+ kv_cache_dtype=kv_cache_dtype,
87
+ disable_async_output_proc=disable_async_output_proc,
88
+ ) as vllm_model:
89
+ test_outputs = vllm_model.generate_greedy_logprobs(
90
+ example_prompts, max_tokens, NUM_LOG_PROBS)
91
+
92
+ check_logprobs_close(
93
+ outputs_0_lst=baseline_outputs,
94
+ outputs_1_lst=test_outputs,
95
+ name_0="fp16_kv_cache",
96
+ name_1="fp8_kv_cache",
97
+ )
98
99
100
@pytest.mark.cpu_model
@@ -119,38 +120,41 @@ def test_cpu_models(
120
test_model: str,
121
max_tokens: int,
122
disable_async_output_proc: bool,
123
+ monkeypatch: pytest.MonkeyPatch,
124
) -> None:
125
"""
126
Only checks log probs match to cover the discrepancy in
127
numerical sensitive kernels.
128
"""
127
-
128
- MAX_MODEL_LEN = 1024
129
- NUM_LOG_PROBS = 8
130
-
131
- with vllm_runner(
132
- base_model,
133
- max_model_len=MAX_MODEL_LEN,
134
- dtype="bfloat16",
135
- kv_cache_dtype="auto",
136
- disable_async_output_proc=disable_async_output_proc,
137
- ) as vllm_model:
138
- baseline_outputs = vllm_model.generate_greedy_logprobs(
139
- example_prompts, max_tokens, NUM_LOG_PROBS)
140
-
141
- with vllm_runner(
142
- test_model,
143
- max_model_len=MAX_MODEL_LEN,
144
- dtype="bfloat16",
145
- kv_cache_dtype=kv_cache_dtype,
146
- disable_async_output_proc=disable_async_output_proc,
147
- ) as vllm_model:
148
- test_outputs = vllm_model.generate_greedy_logprobs(
149
- example_prompts, max_tokens, NUM_LOG_PROBS)
150
-
151
- check_logprobs_close(
152
- outputs_0_lst=baseline_outputs,
153
- outputs_1_lst=test_outputs,
154
- name_0="bf16_kv_cache",
155
- name_1="fp8_kv_cache",
156
- )
129
+ with monkeypatch.context() as m:
130
+ m.setenv("TOKENIZERS_PARALLELISM", 'true')
131
+
132
+ MAX_MODEL_LEN = 1024
133
+ NUM_LOG_PROBS = 8
134
+
135
+ with vllm_runner(
136
+ base_model,
137
+ max_model_len=MAX_MODEL_LEN,
138
+ dtype="bfloat16",
139
+ kv_cache_dtype="auto",
140
+ disable_async_output_proc=disable_async_output_proc,
141
+ ) as vllm_model:
142
+ baseline_outputs = vllm_model.generate_greedy_logprobs(
143
+ example_prompts, max_tokens, NUM_LOG_PROBS)
144
+
145
+ with vllm_runner(
146
+ test_model,
147
+ max_model_len=MAX_MODEL_LEN,
148
+ dtype="bfloat16",
149
+ kv_cache_dtype=kv_cache_dtype,
150
+ disable_async_output_proc=disable_async_output_proc,
151
+ ) as vllm_model:
152
+ test_outputs = vllm_model.generate_greedy_logprobs(
153
+ example_prompts, max_tokens, NUM_LOG_PROBS)
154
+
155
+ check_logprobs_close(
156
+ outputs_0_lst=baseline_outputs,
157
+ outputs_1_lst=test_outputs,
158
+ name_0="bf16_kv_cache",
159
+ name_1="fp8_kv_cache",
160
+ )
tests/models/embedding/language/test_gritlm.py
+53
-43
@@ -1,4 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from __future__ import annotations
3
4
import importlib.util
5
import math
@@ -11,6 +12,7 @@ from scipy.spatial.distance import cosine
12
13
import vllm
14
import vllm.config
15
+from vllm.utils import STR_BACKEND_ENV_VAR
16
17
from ....utils import RemoteOpenAIServer
18
@@ -29,36 +31,34 @@ def _arr(arr):
31
return array("i", arr)
32
33
32
-def test_find_array(monkeypatch):
34
+def test_find_array(monkeypatch: pytest.MonkeyPatch):
35
# GritLM embedding implementation is only supported by XFormers backend.
34
- monkeypatch.setenv("VLLM_ATTENTION_BACKEND", "XFORMERS")
36
+ with monkeypatch.context() as m:
37
+ m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
38
36
- from vllm.model_executor.models.gritlm import GritLMPooler
39
+ from vllm.model_executor.models.gritlm import GritLMPooler
40
38
- # Create an LLM object to get the model config.
39
- llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
40
- pooler = GritLMPooler(model_config=llm.llm_engine.model_config)
41
+ # Create an LLM object to get the model config.
42
+ llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
43
+ pooler = GritLMPooler(model_config=llm.llm_engine.model_config)
44
42
- arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
45
+ arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
46
44
- assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
45
- assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
46
- assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
47
- assert pooler._find_array(arr, _arr([3, 5]), start_idx=0) == -1
47
+ assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
48
+ assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
49
+ assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
50
+ assert pooler._find_array(arr, _arr([3, 5]), start_idx=0) == -1
51
49
- with pytest.raises(ValueError):
50
- pooler._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
52
+ with pytest.raises(ValueError):
53
+ pooler._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
54
55
56
@pytest.fixture(scope="module")
57
def server_embedding():
58
# GritLM embedding implementation is only supported by XFormers backend.
56
- with pytest.MonkeyPatch.context() as mp:
57
- mp.setenv("VLLM_ATTENTION_BACKEND", "XFORMERS")
58
-
59
- args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
60
- with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
61
- yield remote_server
59
+ args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
60
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
61
+ yield remote_server
62
63
64
@pytest.fixture(scope="module")
@@ -69,9 +69,12 @@ def server_generate():
69
70
71
@pytest_asyncio.fixture
72
-async def client_embedding(server_embedding: RemoteOpenAIServer):
73
- async with server_embedding.get_async_client() as async_client:
74
- yield async_client
72
+async def client_embedding(monkeypatch: pytest.MonkeyPatch,
73
+ server_embedding: RemoteOpenAIServer):
74
+ with monkeypatch.context() as m:
75
+ m.setenv("VLLM_ATTENTION_BACKEND", "XFORMERS")
76
+ async with server_embedding.get_async_client() as async_client:
77
+ yield async_client
78
79
80
@pytest_asyncio.fixture
@@ -80,14 +83,20 @@ async def client_generate(server_generate: RemoteOpenAIServer):
83
yield async_client
84
85
83
-def run_llm_encode(llm: vllm.LLM, queries: list[str],
84
- instruction: str) -> list[float]:
86
+def run_llm_encode(
87
+ llm: vllm.LLM,
88
+ queries: list[str],
89
+ instruction: str,
90
+) -> list[float]:
91
outputs = llm.encode([instruction + q for q in queries], )
92
return [output.outputs.embedding for output in outputs]
93
94
89
-async def run_client_embeddings(client: vllm.LLM, queries: list[str],
90
- instruction: str) -> list[float]:
95
+async def run_client_embeddings(
96
+ client: vllm.LLM,
97
+ queries: list[str],
98
+ instruction: str,
99
+) -> list[float]:
100
outputs = await client.embeddings.create(
101
model=MODEL_NAME,
102
input=[instruction + q for q in queries],
@@ -106,7 +115,7 @@ def get_test_data():
115
README.md in https://github.com/ContextualAI/gritlm
116
"""
117
q_instruction = gritlm_instruction(
109
- "Given a scientific paper title, retrieve the paper's abstract")
118
+ "Given a scientific paper title, retrieve the paper's abstract", )
119
queries = [
120
"Bitcoin: A Peer-to-Peer Electronic Cash System",
121
"Generative Representational Instruction Tuning",
@@ -136,31 +145,32 @@ def validate_embed_output(q_rep: list[float], d_rep: list[float]):
145
assert math.isclose(cosine_sim_q1_d1, 0.532, abs_tol=0.001)
146
147
139
-def test_gritlm_offline_embedding(monkeypatch):
148
+def test_gritlm_offline_embedding(monkeypatch: pytest.MonkeyPatch):
149
# GritLM embedding implementation is only supported by XFormers backend.
141
- monkeypatch.setenv("VLLM_ATTENTION_BACKEND", "XFORMERS")
150
+ with monkeypatch.context() as m:
151
+ m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
152
143
- queries, q_instruction, documents, d_instruction = get_test_data()
153
+ queries, q_instruction, documents, d_instruction = get_test_data()
154
145
- llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
155
+ llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
156
147
- d_rep = run_llm_encode(
148
- llm,
149
- documents,
150
- d_instruction,
151
- )
152
- q_rep = run_llm_encode(
153
- llm,
154
- queries,
155
- q_instruction,
156
- )
157
+ d_rep = run_llm_encode(
158
+ llm,
159
+ documents,
160
+ d_instruction,
161
+ )
162
+ q_rep = run_llm_encode(
163
+ llm,
164
+ queries,
165
+ q_instruction,
166
+ )
167
158
- validate_embed_output(q_rep, d_rep)
168
+ validate_embed_output(q_rep, d_rep)
169
170
171
@pytest.mark.asyncio
172
async def test_gritlm_api_server_embedding(
163
- client_embedding: openai.AsyncOpenAI):
173
+ client_embedding: openai.AsyncOpenAI, ):
174
queries, q_instruction, documents, d_instruction = get_test_data()
175
176
d_rep = await run_client_embeddings(
tests/models/test_oot_registration.py
+75
-61
@@ -1,7 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import os
4
-
3
import pytest
4
5
from vllm import LLM, SamplingParams
@@ -11,76 +9,92 @@ from ..utils import fork_new_process_for_each_test
9
10
11
@fork_new_process_for_each_test
14
-def test_plugin(dummy_opt_path, monkeypatch):
12
+def test_plugin(
13
+ monkeypatch: pytest.MonkeyPatch,
14
+ dummy_opt_path: str,
15
+):
16
# V1 shuts down rather than raising an error here.
16
- monkeypatch.setenv("VLLM_USE_V1", "0")
17
- os.environ["VLLM_PLUGINS"] = ""
18
- with pytest.raises(Exception) as excinfo:
19
- LLM(model=dummy_opt_path, load_format="dummy")
20
- error_msg = "has no vLLM implementation and " \
21
- "the Transformers implementation is not compatible with vLLM"
22
- assert (error_msg in str(excinfo.value))
17
+ with monkeypatch.context() as m:
18
+ m.setenv("VLLM_USE_V1", "0")
19
+ m.setenv("VLLM_PLUGINS", "")
20
21
+ with pytest.raises(Exception) as excinfo:
22
+ LLM(model=dummy_opt_path, load_format="dummy")
23
+ error_msg = "has no vLLM implementation and the Transformers implementation is not compatible with vLLM" # noqa: E501
24
+ assert (error_msg in str(excinfo.value))
25
25
-@fork_new_process_for_each_test
26
-def test_oot_registration_text_generation(dummy_opt_path):
27
- os.environ["VLLM_PLUGINS"] = "register_dummy_model"
28
- prompts = ["Hello, my name is", "The text does not matter"]
29
- sampling_params = SamplingParams(temperature=0)
30
- llm = LLM(model=dummy_opt_path, load_format="dummy")
31
- first_token = llm.get_tokenizer().decode(0)
32
- outputs = llm.generate(prompts, sampling_params)
26
34
- for output in outputs:
35
- generated_text = output.outputs[0].text
36
- # make sure only the first token is generated
37
- rest = generated_text.replace(first_token, "")
38
- assert rest == ""
27
+@fork_new_process_for_each_test
28
+def test_oot_registration_text_generation(
29
+ monkeypatch: pytest.MonkeyPatch,
30
+ dummy_opt_path: str,
31
+):
32
+ with monkeypatch.context() as m:
33
+ m.setenv("VLLM_PLUGINS", "register_dummy_model")
34
+ prompts = ["Hello, my name is", "The text does not matter"]
35
+ sampling_params = SamplingParams(temperature=0)
36
+ llm = LLM(model=dummy_opt_path, load_format="dummy")
37
+ first_token = llm.get_tokenizer().decode(0)
38
+ outputs = llm.generate(prompts, sampling_params)
39
+
40
+ for output in outputs:
41
+ generated_text = output.outputs[0].text
42
+ # make sure only the first token is generated
43
+ rest = generated_text.replace(first_token, "")
44
+ assert rest == ""
45
46
47
@fork_new_process_for_each_test
42
-def test_oot_registration_embedding(dummy_gemma2_embedding_path):
43
- os.environ["VLLM_PLUGINS"] = "register_dummy_model"
44
- prompts = ["Hello, my name is", "The text does not matter"]
45
- llm = LLM(model=dummy_gemma2_embedding_path, load_format="dummy")
46
- outputs = llm.embed(prompts)
48
+def test_oot_registration_embedding(
49
+ monkeypatch: pytest.MonkeyPatch,
50
+ dummy_gemma2_embedding_path: str,
51
+):
52
+ with monkeypatch.context() as m:
53
+ m.setenv("VLLM_PLUGINS", "register_dummy_model")
54
+ prompts = ["Hello, my name is", "The text does not matter"]
55
+ llm = LLM(model=dummy_gemma2_embedding_path, load_format="dummy")
56
+ outputs = llm.embed(prompts)
57
48
- for output in outputs:
49
- assert all(v == 0 for v in output.outputs.embedding)
58
+ for output in outputs:
59
+ assert all(v == 0 for v in output.outputs.embedding)
60
61
62
image = ImageAsset("cherry_blossom").pil_image.convert("RGB")
63
64
65
@fork_new_process_for_each_test
56
-def test_oot_registration_multimodal(dummy_llava_path, monkeypatch):
57
- os.environ["VLLM_PLUGINS"] = "register_dummy_model"
58
- prompts = [{
59
- "prompt": "What's in the image?<image>",
60
- "multi_modal_data": {
61
- "image": image
62
- },
63
- }, {
64
- "prompt": "Describe the image<image>",
65
- "multi_modal_data": {
66
- "image": image
67
- },
68
- }]
69
-
70
- sampling_params = SamplingParams(temperature=0)
71
- llm = LLM(model=dummy_llava_path,
72
- load_format="dummy",
73
- max_num_seqs=1,
74
- trust_remote_code=True,
75
- gpu_memory_utilization=0.98,
76
- max_model_len=4096,
77
- enforce_eager=True,
78
- limit_mm_per_prompt={"image": 1})
79
- first_token = llm.get_tokenizer().decode(0)
80
- outputs = llm.generate(prompts, sampling_params)
81
-
82
- for output in outputs:
83
- generated_text = output.outputs[0].text
84
- # make sure only the first token is generated
85
- rest = generated_text.replace(first_token, "")
86
- assert rest == ""
66
+def test_oot_registration_multimodal(
67
+ monkeypatch: pytest.MonkeyPatch,
68
+ dummy_llava_path: str,
69
+):
70
+ with monkeypatch.context() as m:
71
+ m.setenv("VLLM_PLUGINS", "register_dummy_model")
72
+ prompts = [{
73
+ "prompt": "What's in the image?<image>",
74
+ "multi_modal_data": {
75
+ "image": image
76
+ },
77
+ }, {
78
+ "prompt": "Describe the image<image>",
79
+ "multi_modal_data": {
80
+ "image": image
81
+ },
82
+ }]
83
+
84
+ sampling_params = SamplingParams(temperature=0)
85
+ llm = LLM(model=dummy_llava_path,
86
+ load_format="dummy",
87
+ max_num_seqs=1,
88
+ trust_remote_code=True,
89
+ gpu_memory_utilization=0.98,
90
+ max_model_len=4096,
91
+ enforce_eager=True,
92
+ limit_mm_per_prompt={"image": 1})
93
+ first_token = llm.get_tokenizer().decode(0)
94
+ outputs = llm.generate(prompts, sampling_params)
95
+
96
+ for output in outputs:
97
+ generated_text = output.outputs[0].text
98
+ # make sure only the first token is generated
99
+ rest = generated_text.replace(first_token, "")
100
+ assert rest == ""
tests/mq_llm_engine/test_error_handling.py
+17
-14
@@ -235,25 +235,28 @@ async def test_bad_request(tmp_socket):
235
236
237
@pytest.mark.asyncio
238
-async def test_mp_crash_detection(monkeypatch):
238
+async def test_mp_crash_detection(monkeypatch: pytest.MonkeyPatch):
239
+ with monkeypatch.context() as m:
240
240
- parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
241
- parser = make_arg_parser(parser)
242
- args = parser.parse_args([])
241
+ parser = FlexibleArgumentParser(
242
+ description="vLLM's remote OpenAI server.")
243
+ parser = make_arg_parser(parser)
244
+ args = parser.parse_args([])
245
244
- # When LLMEngine is loaded, it will crash.
245
- def mock_init():
246
- raise ValueError
246
+ # When LLMEngine is loaded, it will crash.
247
+ def mock_init():
248
+ raise ValueError
249
248
- monkeypatch.setattr(LLMEngine, "__init__", mock_init)
250
+ m.setattr(LLMEngine, "__init__", mock_init)
251
250
- start = time.perf_counter()
251
- async with build_async_engine_client(args):
252
- pass
253
- end = time.perf_counter()
252
+ start = time.perf_counter()
253
+ async with build_async_engine_client(args):
254
+ pass
255
+ end = time.perf_counter()
256
255
- assert end - start < 60, ("Expected vLLM to gracefully shutdown in <60s "
256
- "if there is an error in the startup.")
257
+ assert end - start < 60, (
258
+ "Expected vLLM to gracefully shutdown in <60s "
259
+ "if there is an error in the startup.")
260
261
262
@pytest.mark.asyncio
tests/multi_step/test_correctness_async_llm.py
+119
-115
@@ -5,7 +5,7 @@ from typing import Optional
5
6
import pytest
7
8
-from tests.kernels.utils import override_backend_env_variable
8
+from vllm.utils import STR_BACKEND_ENV_VAR
9
10
from ..models.utils import check_logprobs_close
11
from ..utils import (completions_with_server_args, get_client_text_generations,
@@ -52,7 +52,7 @@ async def test_multi_step(
52
num_logprobs: Optional[int],
53
attention_backend: str,
54
enable_chunked_prefill: bool,
55
- monkeypatch,
55
+ monkeypatch: pytest.MonkeyPatch,
56
) -> None:
57
"""Test vLLM engine with multi-step scheduling in an OpenAI-protocol
58
client/server environment.
@@ -82,67 +82,70 @@ async def test_multi_step(
82
pytest.skip("Multi-step with Chunked-Prefill only supports"
83
"PP=1 and FLASH_ATTN backend")
84
85
- override_backend_env_variable(monkeypatch, attention_backend)
86
-
87
- prompts = example_prompts
88
- if len(prompts) < num_prompts:
89
- prompts = prompts * ((num_prompts // len(prompts)) + 1)
90
- prompts = prompts[:num_prompts]
91
- assert len(prompts) == num_prompts
92
-
93
- server_args = DEFAULT_SERVER_ARGS + ["--enforce-eager"]
94
- ms_server_args = DEFAULT_SERVER_ARGS + \
95
- ["--num-scheduler-steps", f"{num_scheduler_steps}"]
96
-
97
- if not is_async:
98
- ms_server_args += ["--disable-async-output-proc"]
99
-
100
- if eager_mode:
101
- ms_server_args.append("--enforce-eager")
102
-
103
- if enable_chunked_prefill:
104
- ms_server_args.append("--enable-chunked-prefill")
105
-
106
- distributed_args = [
107
- "--tensor-parallel-size",
108
- str(tp_size),
109
- "--pipeline-parallel-size",
110
- str(pp_size),
111
- ]
112
-
113
- # Spin up client/server & issue completion API requests.
114
- # Default `max_wait_seconds` is 240 but was empirically
115
- # was raised 5x to 1200 *just for this test* due to
116
- # observed timeouts in GHA CI
117
- ref_completions = await completions_with_server_args(
118
- prompts,
119
- model,
120
- server_args + distributed_args,
121
- num_logprobs,
122
- max_wait_seconds=5 * 240)
123
- test_completions = await completions_with_server_args(
124
- prompts,
125
- model,
126
- ms_server_args + distributed_args,
127
- num_logprobs,
128
- max_wait_seconds=5 * 240)
129
-
130
- # Assert multi-step scheduling produces identical tokens
131
- # to single-step scheduling.
132
- ref_generations = get_client_text_generations(ref_completions)
133
- test_generations = get_client_text_generations(test_completions)
134
- assert ref_generations == test_generations
135
-
136
- # Assert multi-step scheduling produces nearly-identical logprobs
137
- # to single-step scheduling.
138
- ref_text_logprobs = get_client_text_logprob_generations(ref_completions)
139
- test_text_logprobs = get_client_text_logprob_generations(test_completions)
140
- check_logprobs_close(
141
- outputs_0_lst=ref_text_logprobs,
142
- outputs_1_lst=test_text_logprobs,
143
- name_0="hf",
144
- name_1="vllm",
145
- )
85
+ with monkeypatch.context() as m:
86
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
87
+
88
+ prompts = example_prompts
89
+ if len(prompts) < num_prompts:
90
+ prompts = prompts * ((num_prompts // len(prompts)) + 1)
91
+ prompts = prompts[:num_prompts]
92
+ assert len(prompts) == num_prompts
93
+
94
+ server_args = DEFAULT_SERVER_ARGS + ["--enforce-eager"]
95
+ ms_server_args = DEFAULT_SERVER_ARGS + \
96
+ ["--num-scheduler-steps", f"{num_scheduler_steps}"]
97
+
98
+ if not is_async:
99
+ ms_server_args += ["--disable-async-output-proc"]
100
+
101
+ if eager_mode:
102
+ ms_server_args.append("--enforce-eager")
103
+
104
+ if enable_chunked_prefill:
105
+ ms_server_args.append("--enable-chunked-prefill")
106
+
107
+ distributed_args = [
108
+ "--tensor-parallel-size",
109
+ str(tp_size),
110
+ "--pipeline-parallel-size",
111
+ str(pp_size),
112
+ ]
113
+
114
+ # Spin up client/server & issue completion API requests.
115
+ # Default `max_wait_seconds` is 240 but was empirically
116
+ # was raised 5x to 1200 *just for this test* due to
117
+ # observed timeouts in GHA CI
118
+ ref_completions = await completions_with_server_args(
119
+ prompts,
120
+ model,
121
+ server_args + distributed_args,
122
+ num_logprobs,
123
+ max_wait_seconds=5 * 240)
124
+ test_completions = await completions_with_server_args(
125
+ prompts,
126
+ model,
127
+ ms_server_args + distributed_args,
128
+ num_logprobs,
129
+ max_wait_seconds=5 * 240)
130
+
131
+ # Assert multi-step scheduling produces identical tokens
132
+ # to single-step scheduling.
133
+ ref_generations = get_client_text_generations(ref_completions)
134
+ test_generations = get_client_text_generations(test_completions)
135
+ assert ref_generations == test_generations
136
+
137
+ # Assert multi-step scheduling produces nearly-identical logprobs
138
+ # to single-step scheduling.
139
+ ref_text_logprobs = get_client_text_logprob_generations(
140
+ ref_completions)
141
+ test_text_logprobs = get_client_text_logprob_generations(
142
+ test_completions)
143
+ check_logprobs_close(
144
+ outputs_0_lst=ref_text_logprobs,
145
+ outputs_1_lst=test_text_logprobs,
146
+ name_0="hf",
147
+ name_1="vllm",
148
+ )
149
150
151
@pytest.mark.parametrize(("tp_size, pp_size"), [
@@ -152,7 +155,7 @@ async def test_multi_step(
155
async def test_multi_step_pp_smoke(
156
tp_size: int,
157
pp_size: int,
155
- monkeypatch,
158
+ monkeypatch: pytest.MonkeyPatch,
159
) -> None:
160
"""
161
Smoke test for the vLLM engine with multi-step scheduling in an
@@ -174,54 +177,55 @@ async def test_multi_step_pp_smoke(
177
attention_backend = "FLASH_ATTN"
178
max_num_seqs = 3
179
177
- override_backend_env_variable(monkeypatch, attention_backend)
178
-
179
- # Prompt from the ShareGPT dataset
180
- prompts = [
181
- "in the jtbd context whats a push?", # codespell:ignore
182
- "in the jtbd context whats a push?", # codespell:ignore
183
- "in the jtbd context whats a push?", # codespell:ignore
184
- "in the jtbd context whats a push?", # codespell:ignore
185
- ]
186
- # Use varying max_tokens to introduce scheduling randomness.
187
- max_tokens = [10 * i for i in range(1, len(prompts) + 1)]
188
- assert len(prompts) == len(max_tokens)
189
-
190
- test_args = [
191
- "--tensor-parallel-size",
192
- str(tp_size), "--pipeline-parallel-size",
193
- str(pp_size), "--max-num-seqs",
194
- str(max_num_seqs)
195
- ]
196
-
197
- server_args = DEFAULT_SERVER_ARGS + test_args
198
- ms_server_args = DEFAULT_SERVER_ARGS + \
199
- ["--num-scheduler-steps", f"{num_scheduler_steps}"] + \
200
- test_args
201
-
202
- # Spin up client/server & issue completion API requests.
203
- # Default `max_wait_seconds` is 240 but was empirically
204
- # was raised 3x to 720 *just for this test* due to
205
- # observed timeouts in GHA CI
206
- ref_completions = await completions_with_server_args(
207
- prompts=prompts,
208
- model_name=model,
209
- server_cli_args=server_args,
210
- num_logprobs=None,
211
- max_wait_seconds=5 * 240,
212
- max_tokens=max_tokens)
213
-
214
- test_completions = await completions_with_server_args(
215
- prompts=prompts,
216
- model_name=model,
217
- server_cli_args=ms_server_args,
218
- num_logprobs=None,
219
- max_wait_seconds=5 * 240,
220
- max_tokens=max_tokens)
221
-
222
- # Assert multi-step scheduling produces identical tokens
223
- # to single-step scheduling.
224
- ref_generations = get_client_text_generations(ref_completions)
225
- test_generations = get_client_text_generations(test_completions)
226
-
227
- assert ref_generations == test_generations
180
+ with monkeypatch.context() as m:
181
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
182
+
183
+ # Prompt from the ShareGPT dataset
184
+ prompts = [
185
+ "in the jtbd context whats a push?", # codespell:ignore
186
+ "in the jtbd context whats a push?", # codespell:ignore
187
+ "in the jtbd context whats a push?", # codespell:ignore
188
+ "in the jtbd context whats a push?", # codespell:ignore
189
+ ]
190
+ # Use varying max_tokens to introduce scheduling randomness.
191
+ max_tokens = [10 * i for i in range(1, len(prompts) + 1)]
192
+ assert len(prompts) == len(max_tokens)
193
+
194
+ test_args = [
195
+ "--tensor-parallel-size",
196
+ str(tp_size), "--pipeline-parallel-size",
197
+ str(pp_size), "--max-num-seqs",
198
+ str(max_num_seqs)
199
+ ]
200
+
201
+ server_args = DEFAULT_SERVER_ARGS + test_args
202
+ ms_server_args = DEFAULT_SERVER_ARGS + \
203
+ ["--num-scheduler-steps", f"{num_scheduler_steps}"] + \
204
+ test_args
205
+
206
+ # Spin up client/server & issue completion API requests.
207
+ # Default `max_wait_seconds` is 240 but was empirically
208
+ # was raised 3x to 720 *just for this test* due to
209
+ # observed timeouts in GHA CI
210
+ ref_completions = await completions_with_server_args(
211
+ prompts=prompts,
212
+ model_name=model,
213
+ server_cli_args=server_args,
214
+ num_logprobs=None,
215
+ max_wait_seconds=5 * 240,
216
+ max_tokens=max_tokens)
217
+
218
+ test_completions = await completions_with_server_args(
219
+ prompts=prompts,
220
+ model_name=model,
221
+ server_cli_args=ms_server_args,
222
+ num_logprobs=None,
223
+ max_wait_seconds=5 * 240,
224
+ max_tokens=max_tokens)
225
+
226
+ # Assert multi-step scheduling produces identical tokens
227
+ # to single-step scheduling.
228
+ ref_generations = get_client_text_generations(ref_completions)
229
+ test_generations = get_client_text_generations(test_completions)
230
+
231
+ assert ref_generations == test_generations
tests/multi_step/test_correctness_llm.py
+164
-161
@@ -7,7 +7,7 @@ from typing import Optional
7
8
import pytest
9
10
-from tests.kernels.utils import override_backend_env_variable
10
+from vllm.utils import STR_BACKEND_ENV_VAR
11
12
from ..models.utils import check_logprobs_close, check_outputs_equal
13
@@ -42,7 +42,7 @@ def test_multi_step_llm(
42
num_prompts: int,
43
num_logprobs: Optional[int],
44
attention_backend: str,
45
- monkeypatch,
45
+ monkeypatch: pytest.MonkeyPatch,
46
) -> None:
47
"""Test vLLM engine with multi-step scheduling via sync LLM Engine.
48
@@ -70,48 +70,49 @@ def test_multi_step_llm(
70
num_logprobs: corresponds to the `logprobs` argument to the OpenAI
71
completions endpoint; `None` -> 1 logprob returned.
72
"""
73
- override_backend_env_variable(monkeypatch, attention_backend)
74
-
75
- prompts = example_prompts
76
- if len(prompts) < num_prompts:
77
- prompts = prompts * ((num_prompts // len(prompts)) + 1)
78
- prompts = prompts[:num_prompts]
79
- assert len(prompts) == num_prompts
80
-
81
- with vllm_runner(
82
- model,
83
- dtype=dtype,
84
- enforce_eager=enforce_eager,
85
- gpu_memory_utilization=0.7,
86
- tensor_parallel_size=tp_size,
87
- enable_chunked_prefill=enable_chunked_prefill,
88
- num_scheduler_steps=num_scheduler_steps,
89
- ) as vllm_model:
90
- vllm_outputs = (vllm_model.generate_greedy(prompts, max_tokens)
91
- if num_logprobs is None else
92
- vllm_model.generate_greedy_logprobs(
93
- prompts, max_tokens, num_logprobs))
94
-
95
- with hf_runner(model, dtype=dtype) as hf_model:
96
- hf_outputs = (hf_model.generate_greedy(prompts, max_tokens)
97
- if num_logprobs is None else
98
- hf_model.generate_greedy_logprobs_limit(
99
- prompts, max_tokens, num_logprobs))
100
-
101
- if num_logprobs is None:
102
- check_outputs_equal(
103
- outputs_0_lst=hf_outputs,
104
- outputs_1_lst=vllm_outputs,
105
- name_0="hf",
106
- name_1="vllm",
107
- )
108
- else:
109
- check_logprobs_close(
110
- outputs_0_lst=hf_outputs,
111
- outputs_1_lst=vllm_outputs,
112
- name_0="hf",
113
- name_1="vllm",
114
- )
73
+ with monkeypatch.context() as m:
74
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
75
+
76
+ prompts = example_prompts
77
+ if len(prompts) < num_prompts:
78
+ prompts = prompts * ((num_prompts // len(prompts)) + 1)
79
+ prompts = prompts[:num_prompts]
80
+ assert len(prompts) == num_prompts
81
+
82
+ with vllm_runner(
83
+ model,
84
+ dtype=dtype,
85
+ enforce_eager=enforce_eager,
86
+ gpu_memory_utilization=0.7,
87
+ tensor_parallel_size=tp_size,
88
+ enable_chunked_prefill=enable_chunked_prefill,
89
+ num_scheduler_steps=num_scheduler_steps,
90
+ ) as vllm_model:
91
+ vllm_outputs = (vllm_model.generate_greedy(prompts, max_tokens)
92
+ if num_logprobs is None else
93
+ vllm_model.generate_greedy_logprobs(
94
+ prompts, max_tokens, num_logprobs))
95
+
96
+ with hf_runner(model, dtype=dtype) as hf_model:
97
+ hf_outputs = (hf_model.generate_greedy(prompts, max_tokens)
98
+ if num_logprobs is None else
99
+ hf_model.generate_greedy_logprobs_limit(
100
+ prompts, max_tokens, num_logprobs))
101
+
102
+ if num_logprobs is None:
103
+ check_outputs_equal(
104
+ outputs_0_lst=hf_outputs,
105
+ outputs_1_lst=vllm_outputs,
106
+ name_0="hf",
107
+ name_1="vllm",
108
+ )
109
+ else:
110
+ check_logprobs_close(
111
+ outputs_0_lst=hf_outputs,
112
+ outputs_1_lst=vllm_outputs,
113
+ name_0="hf",
114
+ name_1="vllm",
115
+ )
116
117
118
@pytest.mark.parametrize("model", MODELS)
@@ -136,7 +137,7 @@ def test_multi_step_llm_w_prompt_logprobs(
137
num_logprobs: Optional[int],
138
num_prompt_logprobs: Optional[int],
139
attention_backend: str,
139
- monkeypatch,
140
+ monkeypatch: pytest.MonkeyPatch,
141
) -> None:
142
"""Test prompt logprobs with multi-step scheduling via sync LLM Engine.
143
@@ -166,47 +167,48 @@ def test_multi_step_llm_w_prompt_logprobs(
167
note that this argument is not supported by the
168
OpenAI completions endpoint.
169
"""
169
- override_backend_env_variable(monkeypatch, attention_backend)
170
-
171
- prompts = example_prompts
172
- if len(prompts) < num_prompts:
173
- prompts = prompts * ((num_prompts // len(prompts)) + 1)
174
- prompts = prompts[:num_prompts]
175
- assert len(prompts) == num_prompts
176
-
177
- with vllm_runner(
178
- model,
179
- dtype=dtype,
180
- enforce_eager=enforce_eager,
181
- gpu_memory_utilization=0.7,
182
- tensor_parallel_size=tp_size,
183
- num_scheduler_steps=num_scheduler_steps,
184
- ) as vllm_model:
185
- vllm_outputs = vllm_model.generate_greedy_logprobs(
186
- prompts,
187
- max_tokens,
188
- num_logprobs,
189
- num_prompt_logprobs=num_prompt_logprobs)
190
-
191
- with vllm_runner(
192
- model,
193
- dtype=dtype,
194
- enforce_eager=enforce_eager,
195
- gpu_memory_utilization=0.7,
196
- tensor_parallel_size=tp_size,
197
- ) as vllm_model:
198
- single_step_vllm_outputs = vllm_model.generate_greedy_logprobs(
199
- prompts,
200
- max_tokens,
201
- num_logprobs,
202
- num_prompt_logprobs=num_prompt_logprobs)
203
-
204
- check_logprobs_close(
205
- outputs_0_lst=single_step_vllm_outputs,
206
- outputs_1_lst=vllm_outputs,
207
- name_0="hf",
208
- name_1="vllm",
209
- )
170
+ with monkeypatch.context() as m:
171
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
172
+
173
+ prompts = example_prompts
174
+ if len(prompts) < num_prompts:
175
+ prompts = prompts * ((num_prompts // len(prompts)) + 1)
176
+ prompts = prompts[:num_prompts]
177
+ assert len(prompts) == num_prompts
178
+
179
+ with vllm_runner(
180
+ model,
181
+ dtype=dtype,
182
+ enforce_eager=enforce_eager,
183
+ gpu_memory_utilization=0.7,
184
+ tensor_parallel_size=tp_size,
185
+ num_scheduler_steps=num_scheduler_steps,
186
+ ) as vllm_model:
187
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
188
+ prompts,
189
+ max_tokens,
190
+ num_logprobs,
191
+ num_prompt_logprobs=num_prompt_logprobs)
192
+
193
+ with vllm_runner(
194
+ model,
195
+ dtype=dtype,
196
+ enforce_eager=enforce_eager,
197
+ gpu_memory_utilization=0.7,
198
+ tensor_parallel_size=tp_size,
199
+ ) as vllm_model:
200
+ single_step_vllm_outputs = vllm_model.generate_greedy_logprobs(
201
+ prompts,
202
+ max_tokens,
203
+ num_logprobs,
204
+ num_prompt_logprobs=num_prompt_logprobs)
205
+
206
+ check_logprobs_close(
207
+ outputs_0_lst=single_step_vllm_outputs,
208
+ outputs_1_lst=vllm_outputs,
209
+ name_0="hf",
210
+ name_1="vllm",
211
+ )
212
213
214
@pytest.mark.parametrize("model", MODELS)
@@ -230,7 +232,7 @@ def test_multi_step_llm_chunked_prefill_prefix_cache(
232
num_prompts: int,
233
num_logprobs: Optional[int],
234
attention_backend: str,
233
- monkeypatch,
235
+ monkeypatch: pytest.MonkeyPatch,
236
) -> None:
237
"""Test vLLM engine with multi-step+"single-step chunked prefill"+APC.
238
@@ -293,77 +295,78 @@ def test_multi_step_llm_chunked_prefill_prefix_cache(
295
#
296
# The Incorrect scheduling behavior - if it occurs - will cause an exception
297
# in the model runner resulting from `do_sample=False`.
296
- override_backend_env_variable(monkeypatch, attention_backend)
297
-
298
- assert len(example_prompts) >= 2
299
- challenge_prompts = copy.deepcopy(example_prompts)
300
- challenge_prompts[0] = ('vLLM is a high-throughput and memory-efficient '
301
- 'inference and serving engine for LLMs.\n'
302
- ) # 24 tok
303
- challenge_prompts[1] = (
304
- 'Briefly describe the major milestones in the '
305
- 'development of artificial intelligence from 1950 to 2020.\n'
306
- ) # 30 tok
307
-
308
- # If necessary, adjust the length of `challenge_prompts` to match
309
- # `num_prompts`
310
- if len(challenge_prompts) < num_prompts:
311
- challenge_prompts = (challenge_prompts *
312
- ((num_prompts // len(challenge_prompts)) + 1))
313
- challenge_prompts = challenge_prompts[:num_prompts]
314
- assert len(challenge_prompts) == num_prompts
315
-
316
- # Single-step scheduler baseline
317
- with vllm_runner(
318
- model,
319
- dtype=dtype,
320
- enforce_eager=enforce_eager,
321
- gpu_memory_utilization=0.7,
322
- tensor_parallel_size=tp_size,
323
- num_scheduler_steps=num_scheduler_steps,
324
- max_model_len=48,
325
- max_num_batched_tokens=48,
326
- max_num_seqs=4,
327
- block_size=16,
328
- ) as vllm_model:
329
- outputs_baseline = (vllm_model.generate_greedy(
330
- challenge_prompts, max_tokens) if num_logprobs is None else
331
- vllm_model.generate_greedy_logprobs(
332
- challenge_prompts, max_tokens, num_logprobs))
333
-
334
- # multi-step+"single-step chunked prefill"+APC
335
- with vllm_runner(
336
- model,
337
- dtype=dtype,
338
- enforce_eager=enforce_eager,
339
- gpu_memory_utilization=0.7,
340
- tensor_parallel_size=tp_size,
341
- enable_chunked_prefill=True,
342
- enable_prefix_caching=True,
343
- num_scheduler_steps=num_scheduler_steps,
344
- max_model_len=48,
345
- max_num_batched_tokens=48,
346
- max_num_seqs=4,
347
- block_size=16,
348
- ) as vllm_model:
349
- outputs_w_features = (vllm_model.generate_greedy(
350
- challenge_prompts, max_tokens) if num_logprobs is None else
351
- vllm_model.generate_greedy_logprobs(
352
- challenge_prompts, max_tokens, num_logprobs))
353
-
354
- if num_logprobs is None:
355
- # No-logprobs test
356
- check_outputs_equal(
357
- outputs_0_lst=outputs_baseline,
358
- outputs_1_lst=outputs_w_features,
359
- name_0="multi-step",
360
- name_1="multi-step+features",
361
- )
362
- else:
363
- # Yes-logprobs test
364
- check_logprobs_close(
365
- outputs_0_lst=outputs_baseline,
366
- outputs_1_lst=outputs_w_features,
367
- name_0="multi-step",
368
- name_1="multi-step+features",
369
- )
298
+ with monkeypatch.context() as m:
299
+ m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
300
+
301
+ assert len(example_prompts) >= 2
302
+ challenge_prompts = copy.deepcopy(example_prompts)
303
+ challenge_prompts[0] = (
304
+ 'vLLM is a high-throughput and memory-efficient '
305
+ 'inference and serving engine for LLMs.\n') # 24 tok
306
+ challenge_prompts[1] = (
307
+ 'Briefly describe the major milestones in the '
308
+ 'development of artificial intelligence from 1950 to 2020.\n'
309
+ ) # 30 tok
310
+
311
+ # If necessary, adjust the length of `challenge_prompts` to match
312
+ # `num_prompts`
313
+ if len(challenge_prompts) < num_prompts:
314
+ challenge_prompts = (challenge_prompts *
315
+ ((num_prompts // len(challenge_prompts)) + 1))
316
+ challenge_prompts = challenge_prompts[:num_prompts]
317
+ assert len(challenge_prompts) == num_prompts
318
+
319
+ # Single-step scheduler baseline
320
+ with vllm_runner(
321
+ model,
322
+ dtype=dtype,
323
+ enforce_eager=enforce_eager,
324
+ gpu_memory_utilization=0.7,
325
+ tensor_parallel_size=tp_size,
326
+ num_scheduler_steps=num_scheduler_steps,
327
+ max_model_len=48,
328
+ max_num_batched_tokens=48,
329
+ max_num_seqs=4,
330
+ block_size=16,
331
+ ) as vllm_model:
332
+ outputs_baseline = (
333
+ vllm_model.generate_greedy(challenge_prompts, max_tokens) if
334
+ num_logprobs is None else vllm_model.generate_greedy_logprobs(
335
+ challenge_prompts, max_tokens, num_logprobs))
336
+
337
+ # multi-step+"single-step chunked prefill"+APC
338
+ with vllm_runner(
339
+ model,
340
+ dtype=dtype,
341
+ enforce_eager=enforce_eager,
342
+ gpu_memory_utilization=0.7,
343
+ tensor_parallel_size=tp_size,
344
+ enable_chunked_prefill=True,
345
+ enable_prefix_caching=True,
346
+ num_scheduler_steps=num_scheduler_steps,
347
+ max_model_len=48,
348
+ max_num_batched_tokens=48,
349
+ max_num_seqs=4,
350
+ block_size=16,
351
+ ) as vllm_model:
352
+ outputs_w_features = (
353
+ vllm_model.generate_greedy(challenge_prompts, max_tokens) if
354
+ num_logprobs is None else vllm_model.generate_greedy_logprobs(
355
+ challenge_prompts, max_tokens, num_logprobs))
356
+
357
+ if num_logprobs is None:
358
+ # No-logprobs test
359
+ check_outputs_equal(
360
+ outputs_0_lst=outputs_baseline,
361
+ outputs_1_lst=outputs_w_features,
362
+ name_0="multi-step",
363
+ name_1="multi-step+features",
364
+ )
365
+ else:
366
+ # Yes-logprobs test
367
+ check_logprobs_close(
368
+ outputs_0_lst=outputs_baseline,
369
+ outputs_1_lst=outputs_w_features,
370
+ name_0="multi-step",
371
+ name_1="multi-step+features",
372
+ )
tests/neuron/1_core/test_block_table.py
+42
-42
@@ -1,5 +1,4 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-import os
2
3
import neuronxcc.nki.language as nl
4
import pytest
@@ -99,6 +98,7 @@ def ref_block_tables_transform(
98
)
99
@torch.inference_mode()
100
def test_load_and_transform_block_tables(
101
+ monkeypatch: pytest.MonkeyPatch,
102
num_tiles,
103
num_blocks_per_tile,
104
q_head_per_kv_head,
@@ -108,46 +108,46 @@ def test_load_and_transform_block_tables(
108
109
device = xm.xla_device()
110
111
- compiler_flags = [
111
+ compiler_flags_str = " ".join([
112
"-O1",
113
"--retry_failed_compilation",
114
- ]
115
- compiler_flags_str = " ".join(compiler_flags)
116
- os.environ["NEURON_CC_FLAGS"] = compiler_flags_str
117
-
118
- torch.manual_seed(10000)
119
- torch.set_printoptions(sci_mode=False)
120
-
121
- # On Neuron, we need B_P_SIZE = 128 blocks to make DMA efficient
122
- B_P_SIZE = 128
123
- if num_blocks_per_tile < B_P_SIZE:
124
- assert B_P_SIZE % num_blocks_per_tile == 0
125
- block_size_tiling_factor = B_P_SIZE // num_blocks_per_tile
126
- else:
127
- block_size_tiling_factor = 1
128
- max_num_blocks = 100000
129
- block_tables = torch.randint(
130
- 0,
131
- max_num_blocks,
132
- (num_tiles * num_blocks_per_tile, ),
133
- dtype=torch.int32,
134
- )
135
- nki_out = nki.jit(nki_load_and_transform_block_tables)[1, 1](
136
- block_tables.to(device=device),
137
- num_tiles,
138
- num_blocks_per_tile,
139
- q_head_per_kv_head,
140
- head_id,
141
- block_size_tiling_factor,
142
- ).cpu()
143
- ref_out = ref_block_tables_transform(
144
- block_tables,
145
- num_tiles,
146
- num_blocks_per_tile,
147
- q_head_per_kv_head,
148
- head_id,
149
- block_size_tiling_factor,
150
- )
151
- assert (nki_out.shape == ref_out.shape
152
- ), f"{nki_out.shape=} != {ref_out.shape=}"
153
- assert torch.all(nki_out == ref_out)
114
+ ])
115
+ with monkeypatch.context() as m:
116
+ m.setenv("NEURON_CC_FLAGS", compiler_flags_str)
117
+
118
+ torch.manual_seed(10000)
119
+ torch.set_printoptions(sci_mode=False)
120
+
121
+ # On Neuron, we need B_P_SIZE = 128 blocks to make DMA efficient
122
+ B_P_SIZE = 128
123
+ if num_blocks_per_tile < B_P_SIZE:
124
+ assert B_P_SIZE % num_blocks_per_tile == 0
125
+ block_size_tiling_factor = B_P_SIZE // num_blocks_per_tile
126
+ else:
127
+ block_size_tiling_factor = 1
128
+ max_num_blocks = 100000
129
+ block_tables = torch.randint(
130
+ 0,
131
+ max_num_blocks,
132
+ (num_tiles * num_blocks_per_tile, ),
133
+ dtype=torch.int32,
134
+ )
135
+ nki_out = nki.jit(nki_load_and_transform_block_tables)[1, 1](
136
+ block_tables.to(device=device),
137
+ num_tiles,
138
+ num_blocks_per_tile,
139
+ q_head_per_kv_head,
140
+ head_id,
141
+ block_size_tiling_factor,
142
+ ).cpu()
143
+ ref_out = ref_block_tables_transform(
144
+ block_tables,
145
+ num_tiles,
146
+ num_blocks_per_tile,
147
+ q_head_per_kv_head,
148
+ head_id,
149
+ block_size_tiling_factor,
150
+ )
151
+ assert (nki_out.shape == ref_out.shape
152
+ ), f"{nki_out.shape=} != {ref_out.shape=}"
153
+ assert torch.all(nki_out == ref_out)
tests/neuron/1_core/test_prefix_prefill.py
+163
-159
@@ -320,6 +320,7 @@ def get_active_block_tables(block_tables, query_lens, seq_lens, block_size,
320
])
321
@torch.inference_mode()
322
def test_contexted_kv_attention(
323
+ monkeypatch: pytest.MonkeyPatch,
324
prefill_batch_size: int,
325
decode_batch_size: int,
326
num_heads: int,
@@ -329,7 +330,6 @@ def test_contexted_kv_attention(
330
large_tile_size,
331
mixed_precision: bool,
332
) -> None:
332
- import os
333
334
import torch_xla.core.xla_model as xm
335
@@ -340,174 +340,178 @@ def test_contexted_kv_attention(
340
341
device = xm.xla_device()
342
343
- compiler_flags = [
343
+ compiler_flags_str = " ".join([
344
"-O1",
345
"--retry_failed_compilation",
346
- ]
347
- compiler_flags_str = " ".join(compiler_flags)
348
- os.environ["NEURON_CC_FLAGS"] = compiler_flags_str
349
-
350
- torch.manual_seed(0)
351
- torch.set_printoptions(sci_mode=False)
352
- torch.set_default_device("cpu")
353
- dtype = torch.float32
354
-
355
- min_ctx_len = 32
356
- max_ctx_len = 1024
357
- min_query_len = 16
358
- max_query_len = 512
359
- num_kv_heads = num_heads // num_queries_per_kv
360
- (
361
- query,
362
- k_active,
363
- v_active,
364
- k_cache,
365
- v_cache,
366
- block_table,
367
- key,
368
- value,
369
- query_lens,
370
- seq_lens,
371
- ) = sample_inputs(
372
- prefill_batch_size=prefill_batch_size,
373
- decode_batch_size=decode_batch_size,
374
- min_query_len=min_query_len,
375
- max_query_len=max_query_len,
376
- min_ctx_len=min_ctx_len,
377
- max_ctx_len=max_ctx_len,
378
- block_size=block_size,
379
- num_heads=num_heads,
380
- num_kv_heads=num_kv_heads,
381
- head_size=head_size,
382
- dtype=dtype,
383
- )
346
+ ])
347
+ with monkeypatch.context() as m:
348
+ m.setenv("NEURON_CC_FLAGS", compiler_flags_str)
349
+
350
+ torch.manual_seed(0)
351
+ torch.set_printoptions(sci_mode=False)
352
+ torch.set_default_device("cpu")
353
+ dtype = torch.float32
354
+
355
+ min_ctx_len = 32
356
+ max_ctx_len = 1024
357
+ min_query_len = 16
358
+ max_query_len = 512
359
+ num_kv_heads = num_heads // num_queries_per_kv
360
+ (
361
+ query,
362
+ k_active,
363
+ v_active,
364
+ k_cache,
365
+ v_cache,
366
+ block_table,
367
+ key,
368
+ value,
369
+ query_lens,
370
+ seq_lens,
371
+ ) = sample_inputs(
372
+ prefill_batch_size=prefill_batch_size,
373
+ decode_batch_size=decode_batch_size,
374
+ min_query_len=min_query_len,
375
+ max_query_len=max_query_len,
376
+ min_ctx_len=min_ctx_len,
377
+ max_ctx_len=max_ctx_len,
378
+ block_size=block_size,
379
+ num_heads=num_heads,
380
+ num_kv_heads=num_kv_heads,
381
+ head_size=head_size,
382
+ dtype=dtype,
383
+ )
384
385
- output_ref = ref_context_attention(
386
- query,
387
- key,
388
- value,
389
- query_lens,
390
- seq_lens,
391
- head_size,
392
- num_queries_per_kv,
393
- return_max_reduce=False,
394
- )
385
+ output_ref = ref_context_attention(
386
+ query,
387
+ key,
388
+ value,
389
+ query_lens,
390
+ seq_lens,
391
+ head_size,
392
+ num_queries_per_kv,
393
+ return_max_reduce=False,
394
+ )
395
396
- # build neuron program
397
- B_P_SIZE = 128
398
- assert (large_tile_size >= B_P_SIZE
399
- ), f"Expect {large_tile_size=} to be larger than {B_P_SIZE=}"
400
-
401
- def ceil_div(a, b):
402
- return (a + b - 1) // b
403
-
404
- def pad_to_multiple(a, b):
405
- return ceil_div(a, b) * b
406
-
407
- def pad_to_next_power_of_2(a):
408
- assert a > 0
409
- return 2**int(a - 1).bit_length()
410
-
411
- # calculate input shapes
412
- max_num_queries = pad_to_next_power_of_2(sum(query_lens))
413
- context_lens = torch.tensor(seq_lens) - torch.tensor(query_lens)
414
- num_active_blocks = ceil_div(context_lens, block_size).sum().item()
415
- num_active_blocks = pad_to_multiple(num_active_blocks,
416
- large_tile_size // block_size)
417
- context_kv_len = num_active_blocks * block_size
418
- assert (context_kv_len %
419
- large_tile_size == 0), f"invalid context_kv_len={context_kv_len}"
396
+ # build neuron program
397
+ B_P_SIZE = 128
398
+ assert (large_tile_size >= B_P_SIZE
399
+ ), f"Expect {large_tile_size=} to be larger than {B_P_SIZE=}"
400
421
- # pad QKV tensors
422
- pad_dims = (
423
- 0,
424
- 0,
425
- 0,
426
- 0,
427
- 0,
428
- max_num_queries - query.shape[0],
429
- )
430
- query = F.pad(query, pad_dims, "constant", 0)
431
- k = F.pad(k_active, pad_dims, "constant", 0)
432
- v = F.pad(v_active, pad_dims, "constant", 0)
433
-
434
- # permute QKV tensors
435
- # query: (1, n_heads, d, seq_q)
436
- # key: (1, n_kv_heads, d, seq_k)
437
- # value: (1, n_kv_heads, seq_v, d)
438
- query = query.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
439
- k = k.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
440
- v = v.unsqueeze(0).permute(0, 2, 1, 3).contiguous()
441
- k_cache = k_cache.permute(0, 2, 1, 3).contiguous()
442
- v_cache = v_cache.permute(0, 2, 1, 3).contiguous()
443
-
444
- # transform block table
445
- active_block_table = get_active_block_tables(
446
- block_table.cpu(),
447
- torch.tensor(query_lens).cpu(),
448
- torch.tensor(seq_lens).cpu(),
449
- block_size,
450
- num_active_blocks,
451
- )
401
+ def ceil_div(a, b):
402
+ return (a + b - 1) // b
403
453
- # Build attention masks
454
- prior_mask, active_mask = (
455
- BlockDiagonalCausalFromBottomRightMask.from_seqlens(
456
- query_lens, seq_lens, block_size=block_size))
457
- prior_mask_padded = F.pad(
458
- prior_mask,
459
- (
404
+ def pad_to_multiple(a, b):
405
+ return ceil_div(a, b) * b
406
+
407
+ def pad_to_next_power_of_2(a):
408
+ assert a > 0
409
+ return 2**int(a - 1).bit_length()
410
+
411
+ # calculate input shapes
412
+ max_num_queries = pad_to_next_power_of_2(sum(query_lens))
413
+ context_lens = torch.tensor(seq_lens) - torch.tensor(query_lens)
414
+ num_active_blocks = ceil_div(context_lens, block_size).sum().item()
415
+ num_active_blocks = pad_to_multiple(num_active_blocks,
416
+ large_tile_size // block_size)
417
+ context_kv_len = num_active_blocks * block_size
418
+ assert (
419
+ context_kv_len %
420
+ large_tile_size == 0), f"invalid context_kv_len={context_kv_len}"
421
+
422
+ # pad QKV tensors
423
+ pad_dims = (
424
0,
461
- context_kv_len - prior_mask.shape[1],
425
0,
463
- max_num_queries - prior_mask.shape[0],
464
- ),
465
- "constant",
466
- 0,
467
- ).bool()
468
- active_mask_padded = F.pad(
469
- active_mask,
470
- (
426
0,
472
- max_num_queries - active_mask.shape[1],
427
0,
474
- max_num_queries - active_mask.shape[0],
475
- ),
476
- "constant",
477
- 0,
478
- ).bool()
479
- attn_mask = torch.concat([prior_mask_padded, active_mask_padded], dim=1)
480
-
481
- attn_mask = reorder_context_mask(attn_mask, large_tile_size, block_size)
482
-
483
- input_args = (
484
- query.to(device=device),
485
- k.to(device=device),
486
- v.to(device=device),
487
- k_cache.to(device=device),
488
- v_cache.to(device=device),
489
- active_block_table.to(device=device),
490
- attn_mask.to(device=device),
491
- )
492
- input_kwargs = dict(
493
- n_kv_head=num_kv_heads,
494
- head_size=head_size,
495
- mixed_precision=mixed_precision,
496
- LARGE_TILE_SZ=large_tile_size,
497
- )
428
+ 0,
429
+ max_num_queries - query.shape[0],
430
+ )
431
+ query = F.pad(query, pad_dims, "constant", 0)
432
+ k = F.pad(k_active, pad_dims, "constant", 0)
433
+ v = F.pad(v_active, pad_dims, "constant", 0)
434
+
435
+ # permute QKV tensors
436
+ # query: (1, n_heads, d, seq_q)
437
+ # key: (1, n_kv_heads, d, seq_k)
438
+ # value: (1, n_kv_heads, seq_v, d)
439
+ query = query.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
440
+ k = k.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
441
+ v = v.unsqueeze(0).permute(0, 2, 1, 3).contiguous()
442
+ k_cache = k_cache.permute(0, 2, 1, 3).contiguous()
443
+ v_cache = v_cache.permute(0, 2, 1, 3).contiguous()
444
+
445
+ # transform block table
446
+ active_block_table = get_active_block_tables(
447
+ block_table.cpu(),
448
+ torch.tensor(query_lens).cpu(),
449
+ torch.tensor(seq_lens).cpu(),
450
+ block_size,
451
+ num_active_blocks,
452
+ )
453
499
- output_nki = flash_attn_varlen_nkifunc(*input_args, **input_kwargs)
454
+ # Build attention masks
455
+ prior_mask, active_mask = (
456
+ BlockDiagonalCausalFromBottomRightMask.from_seqlens(
457
+ query_lens, seq_lens, block_size=block_size))
458
+ prior_mask_padded = F.pad(
459
+ prior_mask,
460
+ (
461
+ 0,
462
+ context_kv_len - prior_mask.shape[1],
463
+ 0,
464
+ max_num_queries - prior_mask.shape[0],
465
+ ),
466
+ "constant",
467
+ 0,
468
+ ).bool()
469
+ active_mask_padded = F.pad(
470
+ active_mask,
471
+ (
472
+ 0,
473
+ max_num_queries - active_mask.shape[1],
474
+ 0,
475
+ max_num_queries - active_mask.shape[0],
476
+ ),
477
+ "constant",
478
+ 0,
479
+ ).bool()
480
+ attn_mask = torch.concat([prior_mask_padded, active_mask_padded],
481
+ dim=1)
482
+
483
+ attn_mask = reorder_context_mask(attn_mask, large_tile_size,
484
+ block_size)
485
+
486
+ input_args = (
487
+ query.to(device=device),
488
+ k.to(device=device),
489
+ v.to(device=device),
490
+ k_cache.to(device=device),
491
+ v_cache.to(device=device),
492
+ active_block_table.to(device=device),
493
+ attn_mask.to(device=device),
494
+ )
495
+ input_kwargs = dict(
496
+ n_kv_head=num_kv_heads,
497
+ head_size=head_size,
498
+ mixed_precision=mixed_precision,
499
+ LARGE_TILE_SZ=large_tile_size,
500
+ )
501
501
- num_actual_tokens = sum(query_lens)
502
- # - o: shape (bs, n_heads, seq_q, d) -> (bs, seq_q, n_heads, d)
503
- output_nki = output_nki.cpu().permute(0, 2, 1, 3)
504
- output_nki = output_nki[0, :num_actual_tokens, :, :]
505
- output_ref_padded = F.pad(
506
- output_ref,
507
- (0, 0, 0, 0, 0, 0, 0, max_num_queries - output_ref.shape[0]),
508
- "constant",
509
- 0,
510
- )
511
- output_ref = output_ref_padded.transpose(0, 1)[0, :num_actual_tokens, :, :]
502
+ output_nki = flash_attn_varlen_nkifunc(*input_args, **input_kwargs)
503
+
504
+ num_actual_tokens = sum(query_lens)
505
+ # - o: shape (bs, n_heads, seq_q, d) -> (bs, seq_q, n_heads, d)
506
+ output_nki = output_nki.cpu().permute(0, 2, 1, 3)
507
+ output_nki = output_nki[0, :num_actual_tokens, :, :]
508
+ output_ref_padded = F.pad(
509
+ output_ref,
510
+ (0, 0, 0, 0, 0, 0, 0, max_num_queries - output_ref.shape[0]),
511
+ "constant",
512
+ 0,
513
+ )
514
+ output_ref = output_ref_padded.transpose(
515
+ 0, 1)[0, :num_actual_tokens, :, :]
516
513
- torch.testing.assert_close(output_nki, output_ref, atol=1e-2, rtol=0)
517
+ torch.testing.assert_close(output_nki, output_ref, atol=1e-2, rtol=0)
tests/plugins_tests/test_platform_plugins.py
+7
-6
@@ -1,10 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import pytest
4
import torch
5
5
-from tests.kernels.utils import override_backend_env_variable
6
from vllm.attention.selector import get_attn_backend
7
-from vllm.utils import STR_INVALID_VAL
7
+from vllm.utils import STR_BACKEND_ENV_VAR, STR_INVALID_VAL
8
9
10
def test_platform_plugins():
@@ -25,8 +25,9 @@ def test_platform_plugins():
25
f" is loaded. The first import:\n{_init_trace}")
26
27
28
-def test_oot_attention_backend(monkeypatch):
28
+def test_oot_attention_backend(monkeypatch: pytest.MonkeyPatch):
29
# ignore the backend env variable if it is set
30
- override_backend_env_variable(monkeypatch, STR_INVALID_VAL)
31
- backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
32
- assert backend.get_name() == "Dummy_Backend"
30
+ with monkeypatch.context() as m:
31
+ m.setenv(STR_BACKEND_ENV_VAR, STR_INVALID_VAL)
32
+ backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
33
+ assert backend.get_name() == "Dummy_Backend"
tests/plugins_tests/test_scheduler_plugins.py
+33
-29
@@ -22,43 +22,47 @@ class DummyV1Scheduler(V1Scheduler):
22
raise Exception("Exception raised by DummyV1Scheduler")
23
24
25
-def test_scheduler_plugins_v0(monkeypatch):
26
- monkeypatch.setenv("VLLM_USE_V1", "0")
27
- with pytest.raises(Exception) as exception_info:
25
+def test_scheduler_plugins_v0(monkeypatch: pytest.MonkeyPatch):
26
+ with monkeypatch.context() as m:
27
+ m.setenv("VLLM_USE_V1", "0")
28
+ with pytest.raises(Exception) as exception_info:
29
29
- engine_args = EngineArgs(
30
- model="facebook/opt-125m",
31
- enforce_eager=True, # reduce test time
32
- scheduler_cls=DummyV0Scheduler,
33
- )
30
+ engine_args = EngineArgs(
31
+ model="facebook/opt-125m",
32
+ enforce_eager=True, # reduce test time
33
+ scheduler_cls=DummyV0Scheduler,
34
+ )
35
35
- engine = LLMEngine.from_engine_args(engine_args=engine_args)
36
+ engine = LLMEngine.from_engine_args(engine_args=engine_args)
37
37
- sampling_params = SamplingParams(max_tokens=1)
38
- engine.add_request("0", "foo", sampling_params)
39
- engine.step()
38
+ sampling_params = SamplingParams(max_tokens=1)
39
+ engine.add_request("0", "foo", sampling_params)
40
+ engine.step()
41
41
- assert str(exception_info.value) == "Exception raised by DummyV0Scheduler"
42
+ assert str(
43
+ exception_info.value) == "Exception raised by DummyV0Scheduler"
44
45
44
-def test_scheduler_plugins_v1(monkeypatch):
45
- monkeypatch.setenv("VLLM_USE_V1", "1")
46
- # Explicitly turn off engine multiprocessing so that the scheduler runs in
47
- # this process
48
- monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
46
+def test_scheduler_plugins_v1(monkeypatch: pytest.MonkeyPatch):
47
+ with monkeypatch.context() as m:
48
+ m.setenv("VLLM_USE_V1", "1")
49
+ # Explicitly turn off engine multiprocessing so
50
+ # that the scheduler runs in this process
51
+ m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
52
50
- with pytest.raises(Exception) as exception_info:
53
+ with pytest.raises(Exception) as exception_info:
54
52
- engine_args = EngineArgs(
53
- model="facebook/opt-125m",
54
- enforce_eager=True, # reduce test time
55
- scheduler_cls=DummyV1Scheduler,
56
- )
55
+ engine_args = EngineArgs(
56
+ model="facebook/opt-125m",
57
+ enforce_eager=True, # reduce test time
58
+ scheduler_cls=DummyV1Scheduler,
59
+ )
60
58
- engine = V1LLMEngine.from_engine_args(engine_args=engine_args)
61
+ engine = V1LLMEngine.from_engine_args(engine_args=engine_args)
62
60
- sampling_params = SamplingParams(max_tokens=1)
61
- engine.add_request("0", "foo", sampling_params)
62
- engine.step()
63
+ sampling_params = SamplingParams(max_tokens=1)
64
+ engine.add_request("0", "foo", sampling_params)
65
+ engine.step()
66
64
- assert str(exception_info.value) == "Exception raised by DummyV1Scheduler"
67
+ assert str(
68
+ exception_info.value) == "Exception raised by DummyV1Scheduler"
tests/prefix_caching/test_prefix_caching.py
+67
-58
@@ -4,25 +4,29 @@
4
Run `pytest tests/prefix_caching/test_prefix_caching.py`.
5
"""
6
7
+from __future__ import annotations
8
+
9
import pytest
10
11
from tests.conftest import VllmRunner
12
from tests.core.utils import SchedulerProxy, create_dummy_prompt
11
-from tests.kernels.utils import override_backend_env_variable
13
from vllm import SamplingParams, TokensPrompt
14
from vllm.core.scheduler import Scheduler
15
from vllm.engine.llm_engine import LLMEngine
16
from vllm.platforms import current_platform
17
+from vllm.utils import STR_BACKEND_ENV_VAR
18
19
from ..models.utils import check_outputs_equal
20
21
22
@pytest.fixture(scope="function", autouse=True)
21
-def use_v0_only(monkeypatch):
23
+def use_v0_only(monkeypatch: pytest.MonkeyPatch):
24
"""
25
This module relies on V0 internals, so set VLLM_USE_V1=0.
26
"""
25
- monkeypatch.setenv('VLLM_USE_V1', '0')
27
+ with monkeypatch.context() as m:
28
+ m.setenv('VLLM_USE_V1', '0')
29
+ yield
30
31
32
MODELS = [
@@ -56,7 +60,7 @@ def test_mixed_requests(
60
cached_position: int,
61
enable_chunked_prefill: bool,
62
block_size: int,
59
- monkeypatch,
63
+ monkeypatch: pytest.MonkeyPatch,
64
) -> None:
65
"""
66
Test the case when some sequences have the prefix cache hit
@@ -67,72 +71,77 @@ def test_mixed_requests(
71
pytest.skip("Flashinfer does not support ROCm/HIP.")
72
if backend == "XFORMERS" and current_platform.is_rocm():
73
pytest.skip("Xformers does not support ROCm/HIP.")
70
- override_backend_env_variable(monkeypatch, backend)
71
-
72
- with hf_runner(model, dtype=dtype) as hf_model:
73
- hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
74
-
75
- cached_prompt = example_prompts[cached_position]
76
- with vllm_runner(
77
- model,
78
- dtype=dtype,
79
- enable_prefix_caching=True,
80
- enable_chunked_prefill=enable_chunked_prefill,
81
- block_size=block_size,
82
- ) as vllm_model:
83
- # Run the first prompt so the cache is populated
84
- vllm_outputs = vllm_model.generate_greedy([cached_prompt], max_tokens)
85
-
86
- # Run all the promopts
87
- greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
88
- req_outputs = vllm_model.model.generate(example_prompts, greedy_params)
89
-
90
- # Verify number of cached tokens
91
- for i in range(len(req_outputs)):
92
- if i == cached_position:
93
- expected_num_cached_tokens = (
94
- len(req_outputs[i].prompt_token_ids) //
95
- block_size) * block_size
96
- else:
97
- expected_num_cached_tokens = 0
98
- assert (
99
- req_outputs[i].num_cached_tokens == expected_num_cached_tokens)
100
-
101
- vllm_outputs = [(
102
- output.prompt_token_ids + list(output.outputs[0].token_ids),
103
- output.prompt + output.outputs[0].text,
104
- ) for output in req_outputs]
105
-
106
- check_outputs_equal(
107
- outputs_0_lst=hf_outputs,
108
- outputs_1_lst=vllm_outputs,
109
- name_0="hf",
110
- name_1="vllm",
111
- )
74
+ with monkeypatch.context() as m:
75
+ m.setenv(STR_BACKEND_ENV_VAR, backend)
76
+
77
+ with hf_runner(model, dtype=dtype) as hf_model:
78
+ hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
79
+
80
+ cached_prompt = example_prompts[cached_position]
81
+ with vllm_runner(
82
+ model,
83
+ dtype=dtype,
84
+ enable_prefix_caching=True,
85
+ enable_chunked_prefill=enable_chunked_prefill,
86
+ block_size=block_size,
87
+ ) as vllm_model:
88
+ # Run the first prompt so the cache is populated
89
+ vllm_outputs = vllm_model.generate_greedy([cached_prompt],
90
+ max_tokens)
91
+
92
+ # Run all the promopts
93
+ greedy_params = SamplingParams(temperature=0.0,
94
+ max_tokens=max_tokens)
95
+ req_outputs = vllm_model.model.generate(example_prompts,
96
+ greedy_params)
97
+
98
+ # Verify number of cached tokens
99
+ for i in range(len(req_outputs)):
100
+ if i == cached_position:
101
+ expected_num_cached_tokens = (
102
+ len(req_outputs[i].prompt_token_ids) //
103
+ block_size) * block_size
104
+ else:
105
+ expected_num_cached_tokens = 0
106
+ assert (req_outputs[i].num_cached_tokens ==
107
+ expected_num_cached_tokens)
108
+
109
+ vllm_outputs = [(
110
+ output.prompt_token_ids + list(output.outputs[0].token_ids),
111
+ output.prompt + output.outputs[0].text,
112
+ ) for output in req_outputs]
113
+
114
+ check_outputs_equal(
115
+ outputs_0_lst=hf_outputs,
116
+ outputs_1_lst=vllm_outputs,
117
+ name_0="hf",
118
+ name_1="vllm",
119
+ )
120
121
122
@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])
123
def test_unstable_prompt_sequence(
124
vllm_runner,
125
backend: str,
118
- monkeypatch,
126
+ monkeypatch: pytest.MonkeyPatch,
127
) -> None:
128
129
if backend == "FLASHINFER" and current_platform.is_rocm():
130
pytest.skip("Flashinfer does not support ROCm/HIP.")
131
if backend == "XFORMERS" and current_platform.is_rocm():
132
pytest.skip("Xformers does not support ROCm/HIP.")
125
- override_backend_env_variable(monkeypatch, backend)
126
-
127
- with vllm_runner(
128
- "Qwen/Qwen2.5-0.5B-Instruct",
129
- enable_chunked_prefill=True,
130
- enable_prefix_caching=True,
131
- max_model_len=4096,
132
- ) as vllm_model:
133
- for prompt in UNSTABLE_PROMPT_SEQUENCE:
134
- vllm_model.generate(TokensPrompt(prompt_token_ids=prompt),
135
- SamplingParams(max_tokens=1))
133
+ with monkeypatch.context() as m:
134
+ m.setenv(STR_BACKEND_ENV_VAR, backend)
135
+
136
+ with vllm_runner(
137
+ "Qwen/Qwen2.5-0.5B-Instruct",
138
+ enable_chunked_prefill=True,
139
+ enable_prefix_caching=True,
140
+ max_model_len=4096,
141
+ ) as vllm_model:
142
+ for prompt in UNSTABLE_PROMPT_SEQUENCE:
143
+ vllm_model.generate(TokensPrompt(prompt_token_ids=prompt),
144
+ SamplingParams(max_tokens=1))
145
146
147
@pytest.mark.parametrize("model", MODELS)
tests/test_regression.py
+4
-12
@@ -56,12 +56,11 @@ def test_gc():
56
assert allocated < 50 * 1024 * 1024
57
58
59
-def test_model_from_modelscope(monkeypatch):
59
+def test_model_from_modelscope(monkeypatch: pytest.MonkeyPatch):
60
# model: https://modelscope.cn/models/qwen/Qwen1.5-0.5B-Chat/summary
61
- MODELSCOPE_MODEL_NAME = "qwen/Qwen1.5-0.5B-Chat"
62
- monkeypatch.setenv("VLLM_USE_MODELSCOPE", "True")
63
- try:
64
- llm = LLM(model=MODELSCOPE_MODEL_NAME)
61
+ with monkeypatch.context() as m:
62
+ m.setenv("VLLM_USE_MODELSCOPE", "True")
63
+ llm = LLM(model="qwen/Qwen1.5-0.5B-Chat")
64
65
prompts = [
66
"Hello, my name is",
@@ -73,10 +72,3 @@ def test_model_from_modelscope(monkeypatch):
72
73
outputs = llm.generate(prompts, sampling_params)
74
assert len(outputs) == 4
76
- finally:
77
- monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
78
-
79
-
80
-if __name__ == "__main__":
81
- import pytest
82
- pytest.main([__file__])
tests/test_utils.py
+32
-31
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa
3
4
import asyncio
4
-import os
5
import socket
6
from collections.abc import AsyncIterator
7
from unittest.mock import patch
@@ -112,16 +112,16 @@ def test_deprecate_kwargs_additional_message():
112
dummy(old_arg=1)
113
114
115
-def test_get_open_port():
116
- os.environ["VLLM_PORT"] = "5678"
117
- # make sure we can get multiple ports, even if the env var is set
118
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s1:
119
- s1.bind(("localhost", get_open_port()))
120
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2:
121
- s2.bind(("localhost", get_open_port()))
122
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
123
- s3.bind(("localhost", get_open_port()))
124
- os.environ.pop("VLLM_PORT")
115
+def test_get_open_port(monkeypatch: pytest.MonkeyPatch):
116
+ with monkeypatch.context() as m:
117
+ m.setenv("VLLM_PORT", "5678")
118
+ # make sure we can get multiple ports, even if the env var is set
119
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s1:
120
+ s1.bind(("localhost", get_open_port()))
121
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2:
122
+ s2.bind(("localhost", get_open_port()))
123
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
124
+ s3.bind(("localhost", get_open_port()))
125
126
127
# Tests for FlexibleArgumentParser
@@ -366,31 +366,32 @@ def test_bind_kv_cache_non_attention():
366
assert ctx['model.layers.28.attn'].kv_cache[0] is kv_cache[1]
367
368
369
-def test_bind_kv_cache_encoder_decoder(monkeypatch):
369
+def test_bind_kv_cache_encoder_decoder(monkeypatch: pytest.MonkeyPatch):
370
# V1 TESTS: ENCODER_DECODER is not supported on V1 yet.
371
- monkeypatch.setenv("VLLM_USE_V1", "0")
371
+ with monkeypatch.context() as m:
372
+ m.setenv("VLLM_USE_V1", "0")
373
373
- from vllm.attention import Attention, AttentionType
374
+ from vllm.attention import Attention, AttentionType
375
375
- # example from bart
376
- ctx = {
377
- 'encoder.layers.0.self_attn.attn':
378
- Attention(32, 128, 0.1, attn_type=AttentionType.ENCODER),
379
- 'decoder.layers.0.encoder_attn.attn':
380
- Attention(32, 128, 0.1, attn_type=AttentionType.ENCODER_DECODER),
381
- 'decoder.layers.0.self_attn.attn':
382
- Attention(32, 128, 0.1, attn_type=AttentionType.DECODER),
383
- }
376
+ # example from bart
377
+ ctx = {
378
+ 'encoder.layers.0.self_attn.attn':
379
+ Attention(32, 128, 0.1, attn_type=AttentionType.ENCODER),
380
+ 'decoder.layers.0.encoder_attn.attn':
381
+ Attention(32, 128, 0.1, attn_type=AttentionType.ENCODER_DECODER),
382
+ 'decoder.layers.0.self_attn.attn':
383
+ Attention(32, 128, 0.1, attn_type=AttentionType.DECODER),
384
+ }
385
385
- kv_cache = [
386
- torch.zeros((1, )),
387
- ]
388
- encoder_kv_cache = ctx['encoder.layers.0.self_attn.attn'].kv_cache
386
+ kv_cache = [
387
+ torch.zeros((1, )),
388
+ ]
389
+ encoder_kv_cache = ctx['encoder.layers.0.self_attn.attn'].kv_cache
390
390
- bind_kv_cache(ctx, [kv_cache])
391
- assert ctx['encoder.layers.0.self_attn.attn'].kv_cache is encoder_kv_cache
392
- assert ctx['decoder.layers.0.encoder_attn.attn'].kv_cache[0] is kv_cache[0]
393
- assert ctx['decoder.layers.0.self_attn.attn'].kv_cache[0] is kv_cache[0]
391
+ bind_kv_cache(ctx, [kv_cache])
392
+ assert ctx['encoder.layers.0.self_attn.attn'].kv_cache is encoder_kv_cache
393
+ assert ctx['decoder.layers.0.encoder_attn.attn'].kv_cache[0] is kv_cache[0]
394
+ assert ctx['decoder.layers.0.self_attn.attn'].kv_cache[0] is kv_cache[0]
395
396
397
def test_bind_kv_cache_pp():
tests/tpu/test_custom_dispatcher.py
+13
-12
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import os
3
+import pytest
4
5
from vllm.config import CompilationLevel
6
@@ -9,16 +9,17 @@ from ..utils import compare_two_settings
9
# --enforce-eager on TPU causes graph compilation
10
# this times out default Health Check in the MQLLMEngine,
11
# so we set the timeout here to 30s
12
-os.environ["VLLM_RPC_TIMEOUT"] = "30000"
12
13
15
-def test_custom_dispatcher():
16
- compare_two_settings(
17
- "google/gemma-2b",
18
- arg1=[
19
- "--enforce-eager",
20
- f"-O{CompilationLevel.DYNAMO_ONCE}",
21
- ],
22
- arg2=["--enforce-eager", f"-O{CompilationLevel.DYNAMO_AS_IS}"],
23
- env1={},
24
- env2={})
14
+def test_custom_dispatcher(monkeypatch: pytest.MonkeyPatch):
15
+ with monkeypatch.context() as m:
16
+ m.setenv("VLLM_RPC_TIMEOUT", "30000")
17
+ compare_two_settings(
18
+ "google/gemma-2b",
19
+ arg1=[
20
+ "--enforce-eager",
21
+ f"-O{CompilationLevel.DYNAMO_ONCE}",
22
+ ],
23
+ arg2=["--enforce-eager", f"-O{CompilationLevel.DYNAMO_AS_IS}"],
24
+ env1={},
25
+ env2={})
tests/tracing/test_tracing.py
+159
-138
@@ -1,10 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa
3
+# type: ignore
4
+from __future__ import annotations
5
3
-import os
6
import threading
7
from collections.abc import Iterable
8
from concurrent import futures
7
-from typing import Callable, Literal
9
+from typing import Callable, Generator, Literal
10
11
import grpc
12
import pytest
@@ -21,12 +23,14 @@ from vllm.tracing import SpanAttributes
23
24
25
@pytest.fixture(scope="function", autouse=True)
24
-def use_v0_only(monkeypatch):
26
+def use_v0_only(monkeypatch: pytest.MonkeyPatch):
27
"""
28
Since this module is V0 only, set VLLM_USE_V1=0 for
29
all tests in the module.
30
"""
29
- monkeypatch.setenv('VLLM_USE_V1', '0')
31
+ with monkeypatch.context() as m:
32
+ m.setenv('VLLM_USE_V1', '0')
33
+ yield
34
35
36
FAKE_TRACE_SERVER_ADDRESS = "localhost:4317"
@@ -67,7 +71,7 @@ class FakeTraceService(TraceServiceServicer):
71
72
73
@pytest.fixture
70
-def trace_service():
74
+def trace_service() -> Generator[FakeTraceService, None, None]:
75
"""Fixture to set up a fake gRPC trace service"""
76
server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
77
service = FakeTraceService()
@@ -80,136 +84,153 @@ def trace_service():
84
server.stop(None)
85
86
83
-def test_traces(trace_service):
84
- os.environ[OTEL_EXPORTER_OTLP_TRACES_INSECURE] = "true"
85
-
86
- sampling_params = SamplingParams(temperature=0.01,
87
- top_p=0.1,
88
- max_tokens=256)
89
- model = "facebook/opt-125m"
90
- llm = LLM(
91
- model=model,
92
- otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS,
93
- )
94
- prompts = ["This is a short prompt"]
95
- outputs = llm.generate(prompts, sampling_params=sampling_params)
96
-
97
- timeout = 5
98
- if not trace_service.evt.wait(timeout):
99
- raise TimeoutError(
100
- f"The fake trace service didn't receive a trace within "
101
- f"the {timeout} seconds timeout")
102
-
103
- request = trace_service.request
104
- assert len(request.resource_spans) == 1, (
105
- f"Expected 1 resource span, "
106
- f"but got {len(request.resource_spans)}")
107
- assert len(request.resource_spans[0].scope_spans) == 1, (
108
- f"Expected 1 scope span, "
109
- f"but got {len(request.resource_spans[0].scope_spans)}")
110
- assert len(request.resource_spans[0].scope_spans[0].spans) == 1, (
111
- f"Expected 1 span, "
112
- f"but got {len(request.resource_spans[0].scope_spans[0].spans)}")
113
-
114
- attributes = decode_attributes(
115
- request.resource_spans[0].scope_spans[0].spans[0].attributes)
116
- assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model
117
- assert attributes.get(
118
- SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id
119
- assert attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE
120
- ) == sampling_params.temperature
121
- assert attributes.get(
122
- SpanAttributes.GEN_AI_REQUEST_TOP_P) == sampling_params.top_p
123
- assert attributes.get(
124
- SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS) == sampling_params.max_tokens
125
- assert attributes.get(SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n
126
- assert attributes.get(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len(
127
- outputs[0].prompt_token_ids)
128
- completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs)
129
- assert attributes.get(
130
- SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) == completion_tokens
131
- metrics = outputs[0].metrics
132
- assert attributes.get(
133
- SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE) == metrics.time_in_queue
134
- ttft = metrics.first_token_time - metrics.arrival_time
135
- assert attributes.get(
136
- SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) == ttft
137
- e2e_time = metrics.finished_time - metrics.arrival_time
138
- assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) == e2e_time
139
- assert metrics.scheduler_time > 0
140
- assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_SCHEDULER
141
- ) == metrics.scheduler_time
142
- # Model forward and model execute should be none, since detailed traces is
143
- # not enabled.
144
- assert metrics.model_forward_time is None
145
- assert metrics.model_execute_time is None
146
-
147
-
148
-def test_traces_with_detailed_steps(trace_service):
149
- os.environ[OTEL_EXPORTER_OTLP_TRACES_INSECURE] = "true"
150
-
151
- sampling_params = SamplingParams(temperature=0.01,
152
- top_p=0.1,
153
- max_tokens=256)
154
- model = "facebook/opt-125m"
155
- llm = LLM(
156
- model=model,
157
- otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS,
158
- collect_detailed_traces="all",
159
- )
160
- prompts = ["This is a short prompt"]
161
- outputs = llm.generate(prompts, sampling_params=sampling_params)
162
-
163
- timeout = 5
164
- if not trace_service.evt.wait(timeout):
165
- raise TimeoutError(
166
- f"The fake trace service didn't receive a trace within "
167
- f"the {timeout} seconds timeout")
168
-
169
- request = trace_service.request
170
- assert len(request.resource_spans) == 1, (
171
- f"Expected 1 resource span, "
172
- f"but got {len(request.resource_spans)}")
173
- assert len(request.resource_spans[0].scope_spans) == 1, (
174
- f"Expected 1 scope span, "
175
- f"but got {len(request.resource_spans[0].scope_spans)}")
176
- assert len(request.resource_spans[0].scope_spans[0].spans) == 1, (
177
- f"Expected 1 span, "
178
- f"but got {len(request.resource_spans[0].scope_spans[0].spans)}")
179
-
180
- attributes = decode_attributes(
181
- request.resource_spans[0].scope_spans[0].spans[0].attributes)
182
- assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model
183
- assert attributes.get(
184
- SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id
185
- assert attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE
186
- ) == sampling_params.temperature
187
- assert attributes.get(
188
- SpanAttributes.GEN_AI_REQUEST_TOP_P) == sampling_params.top_p
189
- assert attributes.get(
190
- SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS) == sampling_params.max_tokens
191
- assert attributes.get(SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n
192
- assert attributes.get(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len(
193
- outputs[0].prompt_token_ids)
194
- completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs)
195
- assert attributes.get(
196
- SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) == completion_tokens
197
- metrics = outputs[0].metrics
198
- assert attributes.get(
199
- SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE) == metrics.time_in_queue
200
- ttft = metrics.first_token_time - metrics.arrival_time
201
- assert attributes.get(
202
- SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) == ttft
203
- e2e_time = metrics.finished_time - metrics.arrival_time
204
- assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) == e2e_time
205
- assert metrics.scheduler_time > 0
206
- assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_SCHEDULER
207
- ) == metrics.scheduler_time
208
- assert metrics.model_forward_time > 0
209
- assert attributes.get(
210
- SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_FORWARD) == pytest.approx(
211
- metrics.model_forward_time / 1000)
212
- assert metrics.model_execute_time > 0
213
- assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_EXECUTE
214
- ) == metrics.model_execute_time
215
- assert metrics.model_forward_time < 1000 * metrics.model_execute_time
87
+def test_traces(
88
+ monkeypatch: pytest.MonkeyPatch,
89
+ trace_service: FakeTraceService,
90
+):
91
+ with monkeypatch.context() as m:
92
+ m.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true")
93
+
94
+ sampling_params = SamplingParams(
95
+ temperature=0.01,
96
+ top_p=0.1,
97
+ max_tokens=256,
98
+ )
99
+ model = "facebook/opt-125m"
100
+ llm = LLM(
101
+ model=model,
102
+ otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS,
103
+ )
104
+ prompts = ["This is a short prompt"]
105
+ outputs = llm.generate(prompts, sampling_params=sampling_params)
106
+
107
+ timeout = 5
108
+ if not trace_service.evt.wait(timeout):
109
+ raise TimeoutError(
110
+ f"The fake trace service didn't receive a trace within "
111
+ f"the {timeout} seconds timeout")
112
+
113
+ request = trace_service.request
114
+ assert len(request.resource_spans) == 1, (
115
+ f"Expected 1 resource span, "
116
+ f"but got {len(request.resource_spans)}")
117
+ assert len(request.resource_spans[0].scope_spans) == 1, (
118
+ f"Expected 1 scope span, "
119
+ f"but got {len(request.resource_spans[0].scope_spans)}")
120
+ assert len(request.resource_spans[0].scope_spans[0].spans) == 1, (
121
+ f"Expected 1 span, "
122
+ f"but got {len(request.resource_spans[0].scope_spans[0].spans)}")
123
+
124
+ attributes = decode_attributes(
125
+ request.resource_spans[0].scope_spans[0].spans[0].attributes)
126
+ assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model
127
+ assert attributes.get(
128
+ SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id
129
+ assert attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE
130
+ ) == sampling_params.temperature
131
+ assert attributes.get(
132
+ SpanAttributes.GEN_AI_REQUEST_TOP_P) == sampling_params.top_p
133
+ assert attributes.get(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS
134
+ ) == sampling_params.max_tokens
135
+ assert attributes.get(
136
+ SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n
137
+ assert attributes.get(
138
+ SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len(
139
+ outputs[0].prompt_token_ids)
140
+ completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs)
141
+ assert attributes.get(
142
+ SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) == completion_tokens
143
+ metrics = outputs[0].metrics
144
+ assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE
145
+ ) == metrics.time_in_queue
146
+ ttft = metrics.first_token_time - metrics.arrival_time
147
+ assert attributes.get(
148
+ SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) == ttft
149
+ e2e_time = metrics.finished_time - metrics.arrival_time
150
+ assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) == e2e_time
151
+ assert metrics.scheduler_time > 0
152
+ assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_SCHEDULER
153
+ ) == metrics.scheduler_time
154
+ # Model forward and model execute should be none, since detailed traces is
155
+ # not enabled.
156
+ assert metrics.model_forward_time is None
157
+ assert metrics.model_execute_time is None
158
+
159
+
160
+def test_traces_with_detailed_steps(
161
+ monkeypatch: pytest.MonkeyPatch,
162
+ trace_service: FakeTraceService,
163
+):
164
+ with monkeypatch.context() as m:
165
+ m.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true")
166
+
167
+ sampling_params = SamplingParams(
168
+ temperature=0.01,
169
+ top_p=0.1,
170
+ max_tokens=256,
171
+ )
172
+ model = "facebook/opt-125m"
173
+ llm = LLM(
174
+ model=model,
175
+ otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS,
176
+ collect_detailed_traces="all",
177
+ )
178
+ prompts = ["This is a short prompt"]
179
+ outputs = llm.generate(prompts, sampling_params=sampling_params)
180
+
181
+ timeout = 5
182
+ if not trace_service.evt.wait(timeout):
183
+ raise TimeoutError(
184
+ f"The fake trace service didn't receive a trace within "
185
+ f"the {timeout} seconds timeout")
186
+
187
+ request = trace_service.request
188
+ assert len(request.resource_spans) == 1, (
189
+ f"Expected 1 resource span, "
190
+ f"but got {len(request.resource_spans)}")
191
+ assert len(request.resource_spans[0].scope_spans) == 1, (
192
+ f"Expected 1 scope span, "
193
+ f"but got {len(request.resource_spans[0].scope_spans)}")
194
+ assert len(request.resource_spans[0].scope_spans[0].spans) == 1, (
195
+ f"Expected 1 span, "
196
+ f"but got {len(request.resource_spans[0].scope_spans[0].spans)}")
197
+
198
+ attributes = decode_attributes(
199
+ request.resource_spans[0].scope_spans[0].spans[0].attributes)
200
+ assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model
201
+ assert attributes.get(
202
+ SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id
203
+ assert attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE
204
+ ) == sampling_params.temperature
205
+ assert attributes.get(
206
+ SpanAttributes.GEN_AI_REQUEST_TOP_P) == sampling_params.top_p
207
+ assert attributes.get(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS
208
+ ) == sampling_params.max_tokens
209
+ assert attributes.get(
210
+ SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n
211
+ assert attributes.get(
212
+ SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len(
213
+ outputs[0].prompt_token_ids)
214
+ completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs)
215
+ assert attributes.get(
216
+ SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) == completion_tokens
217
+ metrics = outputs[0].metrics
218
+ assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE
219
+ ) == metrics.time_in_queue
220
+ ttft = metrics.first_token_time - metrics.arrival_time
221
+ assert attributes.get(
222
+ SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) == ttft
223
+ e2e_time = metrics.finished_time - metrics.arrival_time
224
+ assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) == e2e_time
225
+ assert metrics.scheduler_time > 0
226
+ assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_SCHEDULER
227
+ ) == metrics.scheduler_time
228
+ assert metrics.model_forward_time > 0
229
+ assert attributes.get(
230
+ SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_FORWARD
231
+ ) == pytest.approx(metrics.model_forward_time / 1000)
232
+ assert metrics.model_execute_time > 0
233
+ assert attributes.get(
234
+ SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_EXECUTE
235
+ ) == metrics.model_execute_time
236
+ assert metrics.model_forward_time < 1000 * metrics.model_execute_time
tests/utils.py
+9
-2
@@ -566,6 +566,7 @@ def init_test_distributed_environment(
566
567
568
def multi_process_parallel(
569
+ monkeypatch: pytest.MonkeyPatch,
570
tp_size: int,
571
pp_size: int,
572
test_target: Any,
@@ -582,7 +583,13 @@ def multi_process_parallel(
583
refs = []
584
for rank in range(tp_size * pp_size):
585
refs.append(
585
- test_target.remote(tp_size, pp_size, rank, distributed_init_port))
586
+ test_target.remote(
587
+ monkeypatch,
588
+ tp_size,
589
+ pp_size,
590
+ rank,
591
+ distributed_init_port,
592
+ ), )
593
ray.get(refs)
594
595
ray.shutdown()
@@ -700,7 +707,7 @@ def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
707
"""
708
Get a pytest mark, which skips the test if the GPU doesn't meet
709
a minimum memory requirement in GB.
703
-
710
+
711
This can be leveraged via `@large_gpu_test` to skip tests in environments
712
without enough resources, or called when filtering tests to run directly.
713
"""
tests/v1/e2e/test_ngram_spec_decode.py
+9
-2
@@ -1,5 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from __future__ import annotations
3
+
4
import random
5
+from typing import Any
6
7
import pytest
8
@@ -50,8 +53,12 @@ def model_name():
53
return "meta-llama/Meta-Llama-3-8B-Instruct"
54
55
53
-def test_ngram_correctness(monkeypatch, test_prompts, sampling_config,
54
- model_name):
56
+def test_ngram_correctness(
57
+ monkeypatch: pytest.MonkeyPatch,
58
+ test_prompts: list[list[dict[str, Any]]],
59
+ sampling_config: SamplingParams,
60
+ model_name: str,
61
+):
62
'''
63
Compare the outputs of a original LLM and a speculative LLM
64
should be the same when using ngram speculative decoding.
tests/v1/engine/test_async_llm.py
+7
-4
@@ -80,9 +80,11 @@ async def generate(engine: AsyncLLM,
80
[(TEXT_ENGINE_ARGS, TEXT_PROMPT),
81
(VISION_ENGINE_ARGS, VISION_PROMPT)])
82
@pytest.mark.asyncio
83
-async def test_load(monkeypatch, output_kind: RequestOutputKind,
84
- engine_args_and_prompt: tuple[AsyncEngineArgs,
85
- PromptType]):
83
+async def test_load(
84
+ monkeypatch: pytest.MonkeyPatch,
85
+ output_kind: RequestOutputKind,
86
+ engine_args_and_prompt: tuple[AsyncEngineArgs, PromptType],
87
+):
88
# TODO(rickyx): Remove monkeypatch once we have a better way to test V1
89
# so that in the future when we switch, we don't have to change all the
90
# tests.
@@ -126,7 +128,8 @@ async def test_load(monkeypatch, output_kind: RequestOutputKind,
128
[(TEXT_ENGINE_ARGS, TEXT_PROMPT),
129
(VISION_ENGINE_ARGS, VISION_PROMPT)])
130
@pytest.mark.asyncio
129
-async def test_abort(monkeypatch, output_kind: RequestOutputKind,
131
+async def test_abort(monkeypatch: pytest.MonkeyPatch,
132
+ output_kind: RequestOutputKind,
133
engine_args_and_prompt: tuple[AsyncEngineArgs,
134
PromptType]):
135
tests/v1/engine/test_engine_core.py
+5
-5
@@ -45,7 +45,7 @@ def make_request() -> EngineCoreRequest:
45
46
47
@fork_new_process_for_each_test
48
-def test_engine_core(monkeypatch):
48
+def test_engine_core(monkeypatch: pytest.MonkeyPatch):
49
50
with monkeypatch.context() as m:
51
m.setenv("VLLM_USE_V1", "1")
@@ -159,10 +159,10 @@ def test_engine_core(monkeypatch):
159
160
161
@fork_new_process_for_each_test
162
-def test_engine_core_advanced_sampling(monkeypatch):
162
+def test_engine_core_advanced_sampling(monkeypatch: pytest.MonkeyPatch):
163
"""
164
- A basic end-to-end test to verify that the engine functions correctly
165
- when additional sampling parameters, such as top_p, min_tokens, and
164
+ A basic end-to-end test to verify that the engine functions correctly
165
+ when additional sampling parameters, such as top_p, min_tokens, and
166
presence_penalty, are set.
167
"""
168
with monkeypatch.context() as m:
@@ -209,7 +209,7 @@ def test_engine_core_advanced_sampling(monkeypatch):
209
210
211
@fork_new_process_for_each_test
212
-def test_engine_core_concurrent_batches(monkeypatch):
212
+def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):
213
"""
214
Test that the engine can handle multiple concurrent batches.
215
"""
tests/v1/engine/test_engine_core_client.py
+3
-2
@@ -90,7 +90,8 @@ def echo(self, msg: str, err_msg: Optional[str] = None) -> str:
90
91
@fork_new_process_for_each_test
92
@pytest.mark.parametrize("multiprocessing_mode", [True, False])
93
-def test_engine_core_client(monkeypatch, multiprocessing_mode: bool):
93
+def test_engine_core_client(monkeypatch: pytest.MonkeyPatch,
94
+ multiprocessing_mode: bool):
95
96
with monkeypatch.context() as m:
97
m.setenv("VLLM_USE_V1", "1")
@@ -175,7 +176,7 @@ def test_engine_core_client(monkeypatch, multiprocessing_mode: bool):
176
177
178
@pytest.mark.asyncio(loop_scope="function")
178
-async def test_engine_core_client_asyncio(monkeypatch):
179
+async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
180
181
with monkeypatch.context() as m:
182
m.setenv("VLLM_USE_V1", "1")
tests/v1/sample/test_logprobs.py
+125
-115
@@ -57,7 +57,7 @@ def _repeat_logprob_config(
57
logprob_prompt_logprob_list: BatchLogprobsSpecType,
58
) -> BatchLogprobsSpecType:
59
"""Ensure each test prompt has a logprob config.
60
-
60
+
61
A logprob config specifies the optional (i.e.
62
may-be-`None`) number of sample logprobs and
63
the optional number of prompt logprobs.
@@ -80,7 +80,7 @@ def _repeat_logprob_config(
80
(optional num sample logprob,
81
optional num prompt logprob)
82
tuples
83
-
83
+
84
Returns:
85
list of
86
(optional num sample logprob,optional num prompt logprob)
@@ -255,14 +255,12 @@ def _run_and_validate(
255
[NONE, SAMPLE, PROMPT, SAMPLE_PROMPT])
256
@pytest.mark.parametrize("temperature", [0.0, 2.0])
257
def test_get_logprobs_and_prompt_logprobs(
258
- hf_model,
259
- vllm_model,
260
- batch_logprobs_composition: BatchLogprobsComposition,
261
- temperature: float,
262
- example_prompts,
263
-) -> None:
258
+ hf_model, vllm_model,
259
+ batch_logprobs_composition: BatchLogprobsComposition,
260
+ temperature: float, example_prompts: list[str],
261
+ monkeypatch: pytest.MonkeyPatch) -> None:
262
"""Test V1 Engine logprobs & prompt logprobs
265
-
263
+
264
Exercise a variety of combinations of `logprobs` and `prompt_logprobs`
265
settings and validate that
266
* The generated logprobs and prompt logprobs are consistent with the
@@ -279,7 +277,7 @@ def test_get_logprobs_and_prompt_logprobs(
277
278
To save time, only test one APC-enabled scenario
279
(sample & prompt logprobs enabled, temperature>0.0).
282
-
280
+
281
Args:
282
hf_model: HuggingFace reference model fixture
283
vllm_model: vLLM model fixture
@@ -287,128 +285,140 @@ def test_get_logprobs_and_prompt_logprobs(
285
temperature: "temperature" sampling parameter
286
example_prompts: example prompt fixture
287
"""
290
- do_apc = vllm_model.model.llm_engine.cache_config.enable_prefix_caching
291
- if do_apc and (temperature < 2.0
292
- or batch_logprobs_composition != SAMPLE_PROMPT):
293
- # Skip some test-cases to save time.
294
- pytest.skip()
295
- test_prompts = example_prompts
296
-
297
- max_tokens = 5
298
- hf_outputs = hf_model.generate_greedy(
299
- test_prompts,
300
- max_tokens=max_tokens,
301
- )
302
- hf_logprobs = hf_model.generate_greedy_logprobs(
303
- test_prompts,
304
- max_tokens=max_tokens,
305
- )
306
-
307
- # Batch has mixed sample params
308
- # (different logprobs/prompt logprobs combos)
309
- logprob_prompt_logprob_list = get_test_batch(batch_logprobs_composition)
310
-
311
- # Ensure that each test prompt has a logprob config for testing
312
- logprob_prompt_logprob_list = _repeat_logprob_config(
313
- test_prompts, logprob_prompt_logprob_list)
314
- # Generate SamplingParams
315
- vllm_sampling_params = [
316
- SamplingParams(max_tokens=max_tokens,
317
- logprobs=num_lp,
318
- prompt_logprobs=num_plp,
319
- temperature=temperature,
320
- seed=1984)
321
- for num_lp, num_plp in logprob_prompt_logprob_list
322
- ]
323
- for _ in range(2 if do_apc else 1):
324
- _run_and_validate(
325
- vllm_model=vllm_model,
326
- test_prompts=test_prompts,
327
- vllm_sampling_params=vllm_sampling_params,
328
- hf_logprobs=hf_logprobs,
329
- hf_outputs=hf_outputs,
330
- logprob_prompt_logprob_list=logprob_prompt_logprob_list,
331
- temperature=temperature,
288
+ with monkeypatch.context() as m:
289
+ m.setenv("VLLM_USE_V1", "1")
290
+ do_apc = vllm_model.model.llm_engine.cache_config.enable_prefix_caching
291
+ if do_apc and (temperature < 2.0
292
+ or batch_logprobs_composition != SAMPLE_PROMPT):
293
+ # Skip some test-cases to save time.
294
+ pytest.skip()
295
+ test_prompts = example_prompts
296
+
297
+ max_tokens = 5
298
+ hf_outputs = hf_model.generate_greedy(
299
+ test_prompts,
300
max_tokens=max_tokens,
333
- do_apc=do_apc)
334
-
335
-
336
-def test_max_logprobs():
301
+ )
302
+ hf_logprobs = hf_model.generate_greedy_logprobs(
303
+ test_prompts,
304
+ max_tokens=max_tokens,
305
+ )
306
+
307
+ # Batch has mixed sample params
308
+ # (different logprobs/prompt logprobs combos)
309
+ logprob_prompt_logprob_list = get_test_batch(
310
+ batch_logprobs_composition)
311
+
312
+ # Ensure that each test prompt has a logprob config for testing
313
+ logprob_prompt_logprob_list = _repeat_logprob_config(
314
+ test_prompts, logprob_prompt_logprob_list)
315
+ # Generate SamplingParams
316
+ vllm_sampling_params = [
317
+ SamplingParams(max_tokens=max_tokens,
318
+ logprobs=num_lp,
319
+ prompt_logprobs=num_plp,
320
+ temperature=temperature,
321
+ seed=1984)
322
+ for num_lp, num_plp in logprob_prompt_logprob_list
323
+ ]
324
+ for _ in range(2 if do_apc else 1):
325
+ _run_and_validate(
326
+ vllm_model=vllm_model,
327
+ test_prompts=test_prompts,
328
+ vllm_sampling_params=vllm_sampling_params,
329
+ hf_logprobs=hf_logprobs,
330
+ hf_outputs=hf_outputs,
331
+ logprob_prompt_logprob_list=logprob_prompt_logprob_list,
332
+ temperature=temperature,
333
+ max_tokens=max_tokens,
334
+ do_apc=do_apc)
335
+
336
+
337
+def test_max_logprobs(monkeypatch: pytest.MonkeyPatch):
338
"""vLLM v1 engine should fail a request with `logprobs > max_logprobs`
338
-
339
Should also fail for `prompt_logprobs > max_logprobs`
340
-
340
APC should not matter as this test checks basic request validation.
342
-
343
- Args:
344
- monkeypatch
341
"""
342
+ with monkeypatch.context() as m:
343
+ m.setenv("VLLM_USE_V1", "1")
344
347
- runner = VllmRunner("facebook/opt-125m",
348
- max_logprobs=1,
349
- enable_prefix_caching=False,
350
- max_model_len=256)
351
- vllm_sampling_params = SamplingParams(logprobs=1)
352
- # should pass
353
- runner.generate(["Hello world"], sampling_params=vllm_sampling_params)
345
+ runner = VllmRunner("facebook/opt-125m",
346
+ max_logprobs=1,
347
+ enable_prefix_caching=False,
348
+ max_model_len=256)
349
+ vllm_sampling_params = SamplingParams(logprobs=1)
350
+ # should pass
351
+ runner.generate(["Hello world"], sampling_params=vllm_sampling_params)
352
355
- bad_sampling_params = SamplingParams(logprobs=2)
356
- with pytest.raises(ValueError):
357
- runner.generate(["Hello world"], sampling_params=bad_sampling_params)
353
+ bad_sampling_params = SamplingParams(logprobs=2)
354
+ with pytest.raises(ValueError):
355
+ runner.generate(["Hello world"],
356
+ sampling_params=bad_sampling_params)
357
358
360
-def test_none_logprobs(vllm_model, example_prompts):
359
+def test_none_logprobs(vllm_model, example_prompts,
360
+ monkeypatch: pytest.MonkeyPatch):
361
"""Engine should return `logprobs` and `prompt_logprobs` as `None`
362
-
362
+
363
Args:
364
vllm_model: vLLM model fixture
365
example_prompts: list of example prompts (test fixture)
366
"""
367
- max_tokens = 5
368
-
369
- sampling_params_logprobs_none = SamplingParams(max_tokens=max_tokens,
370
- logprobs=None,
371
- prompt_logprobs=None,
372
- temperature=0.0)
373
- results_logprobs_none = vllm_model.model.generate(
374
- example_prompts, sampling_params=sampling_params_logprobs_none)
367
+ with monkeypatch.context() as m:
368
+ m.setenv("VLLM_USE_V1", "1")
369
+ max_tokens = 5
370
376
- for i in range(len(results_logprobs_none)):
377
- # Check sample logprobs are None
378
- assert results_logprobs_none[i].outputs[0].logprobs is None
379
- assert results_logprobs_none[i].outputs[0].cumulative_logprob is None
380
- # Check prompt logprobs are None
381
- assert results_logprobs_none[i].prompt_logprobs is None
382
-
383
-
384
-def test_zero_logprobs(vllm_model, example_prompts):
371
+ sampling_params_logprobs_none = SamplingParams(
372
+ max_tokens=max_tokens,
373
+ logprobs=None,
374
+ prompt_logprobs=None,
375
+ temperature=0.0,
376
+ )
377
+ results_logprobs_none = vllm_model.model.generate(
378
+ example_prompts,
379
+ sampling_params=sampling_params_logprobs_none,
380
+ )
381
+
382
+ for i in range(len(results_logprobs_none)):
383
+ # Check sample logprobs are None
384
+ assert results_logprobs_none[i].outputs[0].logprobs is None
385
+ assert results_logprobs_none[i].outputs[
386
+ 0].cumulative_logprob is None
387
+ # Check prompt logprobs are None
388
+ assert results_logprobs_none[i].prompt_logprobs is None
389
+
390
+
391
+def test_zero_logprobs(vllm_model, example_prompts,
392
+ monkeypatch: pytest.MonkeyPatch):
393
"""Engine should return sampled token and prompt token logprobs
386
-
394
+
395
Args:
396
vllm_model: vLLM model fixture
397
example_prompts: list of example prompts (test fixture)
398
"""
391
- max_tokens = 5
392
-
393
- sampling_params_logprobs_zero = SamplingParams(max_tokens=max_tokens,
394
- logprobs=0,
395
- prompt_logprobs=0,
396
- temperature=0.0)
397
- results_logprobs_zero = vllm_model.model.generate(
398
- example_prompts, sampling_params=sampling_params_logprobs_zero)
399
-
400
- for i in range(len(results_logprobs_zero)):
401
- # Check that there is one sample logprob dict for each
402
- # sample token
403
- logprobs = results_logprobs_zero[i].outputs[0].logprobs
404
- prompt_logprobs = results_logprobs_zero[i].prompt_logprobs
405
- sampled_token_ids = results_logprobs_zero[i].outputs[0].token_ids
406
- prompt_token_ids = results_logprobs_zero[i].prompt_token_ids
407
- assert logprobs is not None
408
- assert len(sampled_token_ids) == len(logprobs)
409
- assert results_logprobs_zero[i].outputs[
410
- 0].cumulative_logprob is not None
411
- # Check that there is one prompt logprob dict for each
412
- # prompt token
413
- assert prompt_logprobs is not None
414
- assert len(prompt_token_ids) == len(prompt_logprobs)
399
+ with monkeypatch.context() as m:
400
+ m.setenv("VLLM_USE_V1", "1")
401
+ max_tokens = 5
402
+
403
+ sampling_params_logprobs_zero = SamplingParams(max_tokens=max_tokens,
404
+ logprobs=0,
405
+ prompt_logprobs=0,
406
+ temperature=0.0)
407
+ results_logprobs_zero = vllm_model.model.generate(
408
+ example_prompts, sampling_params=sampling_params_logprobs_zero)
409
+
410
+ for i in range(len(results_logprobs_zero)):
411
+ # Check that there is one sample logprob dict for each
412
+ # sample token
413
+ logprobs = results_logprobs_zero[i].outputs[0].logprobs
414
+ prompt_logprobs = results_logprobs_zero[i].prompt_logprobs
415
+ sampled_token_ids = results_logprobs_zero[i].outputs[0].token_ids
416
+ prompt_token_ids = results_logprobs_zero[i].prompt_token_ids
417
+ assert logprobs is not None
418
+ assert len(sampled_token_ids) == len(logprobs)
419
+ assert results_logprobs_zero[i].outputs[
420
+ 0].cumulative_logprob is not None
421
+ # Check that there is one prompt logprob dict for each
422
+ # prompt token
423
+ assert prompt_logprobs is not None
424
+ assert len(prompt_token_ids) == len(prompt_logprobs)
tests/v1/tpu/test_basic.py
+11
-5
@@ -3,11 +3,16 @@
3
4
Run `pytest tests/v1/tpu/test_basic.py`.
5
"""
6
+from __future__ import annotations
7
+
8
+from typing import TYPE_CHECKING
9
+
10
import pytest
11
12
from vllm.platforms import current_platform
13
10
-from ...conftest import VllmRunner
14
+if TYPE_CHECKING:
15
+ from tests.conftest import VllmRunner
16
17
MODELS = [
18
# "Qwen/Qwen2-7B-Instruct",
@@ -28,7 +33,8 @@ TENSOR_PARALLEL_SIZES = [1]
33
@pytest.mark.parametrize("enforce_eager", [True])
34
@pytest.mark.parametrize("tensor_parallel_size", TENSOR_PARALLEL_SIZES)
35
def test_models(
31
- monkeypatch,
36
+ vllm_runner: type[VllmRunner],
37
+ monkeypatch: pytest.MonkeyPatch,
38
model: str,
39
max_tokens: int,
40
enforce_eager: bool,
@@ -41,7 +47,7 @@ def test_models(
47
with monkeypatch.context() as m:
48
m.setenv("VLLM_USE_V1", "1")
49
44
- with VllmRunner(
50
+ with vllm_runner(
51
model,
52
max_model_len=8192,
53
enforce_eager=enforce_eager,
@@ -50,5 +56,5 @@ def test_models(
56
tensor_parallel_size=tensor_parallel_size) as vllm_model:
57
vllm_outputs = vllm_model.generate_greedy(example_prompts,
58
max_tokens)
53
- output = vllm_outputs[0][1]
54
- assert "1024" in output
59
+ output = vllm_outputs[0][1]
60
+ assert "1024" in output