Skip to content

Commit 59237ea

Browse files
MengqingCaoMengqingCao
authored
[CI/UT] Add test for chunk prefill and prefix cache on v1/AscendScheduler (#1505)
### What this PR does / why we need it? Add test for chunked prefill and prefix cache on v1/AscendScheduler Covered scenarios: - `Qwen/Qwen3-0.6B-Base` and `deepseek-ai/DeepSeek-V2-Lite-Chat` --- multicard CI time increased by 19 min - `V1 + default scheduler` vs `V1 + default scheduler + enable prefix cache` - `V1 + Ascend scheduler` vs `V1 + Ascend scheduler + enable prefix cache` vs `V1 + Ascend scheduler + enable prefix cache + enable chunked prefill` - `Qwen/Qwen3-0.6B-Base` --- singlecard CI time increased by 8 min - `V1 + Ascend scheduler` vs `V1 + Ascend scheduler + enable chunked prefill` should rebase after #1498 and #1446 ### Does this PR introduce _any_ user-facing change? N/A ### How was this patch tested? CI passed with new added test. Signed-off-by: MengqingCao <cmq0113@163.com>
1 parent 6b80c5a commit 59237ea

File tree

6 files changed

+241
-0
lines changed

6 files changed

+241
-0
lines changed

‎tests/conftest.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import contextlib
2121
import gc
22+
import os
2223
from typing import Any, List, Optional, Tuple, TypeVar, Union
2324

2425
import numpy as np
@@ -59,6 +60,9 @@
5960
PromptAudioInput = _PromptMultiModalInput[Tuple[np.ndarray, int]]
6061
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
6162

63+
_TEST_DIR = os.path.dirname(__file__)
64+
_TEST_PROMPTS = [os.path.join(_TEST_DIR, "e2e", "prompts", "example.txt")]
65+
6266

6367
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
6468
destroy_model_parallel()
@@ -367,6 +371,20 @@ def prompt_template(request):
367371
return PROMPT_TEMPLATES[request.param]
368372

369373

374+
def _read_prompts(filename: str) -> list[str]:
375+
with open(filename) as f:
376+
prompts = f.readlines()
377+
return prompts
378+
379+
380+
@pytest.fixture
381+
def example_prompts() -> list[str]:
382+
prompts = []
383+
for filename in _TEST_PROMPTS:
384+
prompts += _read_prompts(filename)
385+
return prompts
386+
387+
370388
@pytest.fixture(scope="session")
371389
def ilama_lora_files():
372390
return snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider")

‎tests/e2e/multicard/test_prefix_caching.py

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3+
"""Compare the with and without prefix caching on V1 scheduler or AscendScheduler."""
4+
5+
import os
6+
7+
import pytest
8+
9+
from tests.conftest import VllmRunner
10+
from tests.model_utils import check_outputs_equal
11+
12+
MODELS = [
13+
# for MHA
14+
"Qwen/Qwen3-8B-Base",
15+
# for MLA
16+
"deepseek-ai/DeepSeek-V2-Lite-Chat"
17+
]
18+
19+
# A prompt containing a large markdown table. The table is randomly generated by GPT-4.
20+
LONG_PROMPT = "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as follows.\n# Table\n" + """
21+
| ID | Name | Age | Occupation | Country | Email | Phone Number | Address |
22+
|-----|---------------|-----|---------------|---------------|------------------------|----------------|------------------------------|
23+
| 1 | John Doe | 29 | Engineer | USA | john.doe@example.com | 555-1234 | 123 Elm St, Springfield, IL |
24+
| 2 | Jane Smith | 34 | Doctor | Canada | jane.smith@example.com | 555-5678 | 456 Oak St, Toronto, ON |
25+
| 3 | Alice Johnson | 27 | Teacher | UK | alice.j@example.com | 555-8765 | 789 Pine St, London, UK |
26+
| 4 | Bob Brown | 45 | Artist | Australia | bob.b@example.com | 555-4321 | 321 Maple St, Sydney, NSW |
27+
| 5 | Carol White | 31 | Scientist | New Zealand | carol.w@example.com | 555-6789 | 654 Birch St, Wellington, NZ |
28+
| 6 | Dave Green | 28 | Lawyer | Ireland | dave.g@example.com | 555-3456 | 987 Cedar St, Dublin, IE |
29+
| 7 | Emma Black | 40 | Musician | USA | emma.b@example.com | 555-1111 | 246 Ash St, New York, NY |
30+
| 8 | Frank Blue | 37 | Chef | Canada | frank.b@example.com | 555-2222 | 135 Spruce St, Vancouver, BC |
31+
| 9 | Grace Yellow | 50 | Engineer | UK | grace.y@example.com | 555-3333 | 864 Fir St, Manchester, UK |
32+
| 10 | Henry Violet | 32 | Artist | Australia | henry.v@example.com | 555-4444 | 753 Willow St, Melbourne, VIC|
33+
| 11 | Irene Orange | 26 | Scientist | New Zealand | irene.o@example.com | 555-5555 | 912 Poplar St, Auckland, NZ |
34+
| 12 | Jack Indigo | 38 | Teacher | Ireland | jack.i@example.com | 555-6666 | 159 Elm St, Cork, IE |
35+
| 13 | Karen Red | 41 | Lawyer | USA | karen.r@example.com | 555-7777 | 357 Cedar St, Boston, MA |
36+
| 14 | Leo Brown | 30 | Chef | Canada | leo.b@example.com | 555-8888 | 246 Oak St, Calgary, AB |
37+
| 15 | Mia Green | 33 | Musician | UK | mia.g@example.com | 555-9999 | 975 Pine St, Edinburgh, UK |
38+
| 16 | Noah Yellow | 29 | Doctor | Australia | noah.y@example.com | 555-0000 | 864 Birch St, Brisbane, QLD |
39+
| 17 | Olivia Blue | 35 | Engineer | New Zealand | olivia.b@example.com | 555-1212 | 753 Maple St, Hamilton, NZ |
40+
| 18 | Peter Black | 42 | Artist | Ireland | peter.b@example.com | 555-3434 | 912 Fir St, Limerick, IE |
41+
| 19 | Quinn White | 28 | Scientist | USA | quinn.w@example.com | 555-5656 | 159 Willow St, Seattle, WA |
42+
| 20 | Rachel Red | 31 | Teacher | Canada | rachel.r@example.com | 555-7878 | 357 Poplar St, Ottawa, ON |
43+
| 21 | Steve Green | 44 | Lawyer | UK | steve.g@example.com | 555-9090 | 753 Elm St, Birmingham, UK |
44+
| 22 | Tina Blue | 36 | Musician | Australia | tina.b@example.com | 555-1213 | 864 Cedar St, Perth, WA |
45+
| 23 | Umar Black | 39 | Chef | New Zealand | umar.b@example.com | 555-3435 | 975 Spruce St, Christchurch, NZ|
46+
| 24 | Victor Yellow | 43 | Engineer | Ireland | victor.y@example.com | 555-5657 | 246 Willow St, Galway, IE |
47+
| 25 | Wendy Orange | 27 | Artist | USA | wendy.o@example.com | 555-7879 | 135 Elm St, Denver, CO |
48+
| 26 | Xavier Green | 34 | Scientist | Canada | xavier.g@example.com | 555-9091 | 357 Oak St, Montreal, QC |
49+
| 27 | Yara Red | 41 | Teacher | UK | yara.r@example.com | 555-1214 | 975 Pine St, Leeds, UK |
50+
| 28 | Zack Blue | 30 | Lawyer | Australia | zack.b@example.com | 555-3436 | 135 Birch St, Adelaide, SA |
51+
| 29 | Amy White | 33 | Musician | New Zealand | amy.w@example.com | 555-5658 | 159 Maple St, Wellington, NZ |
52+
| 30 | Ben Black | 38 | Chef | Ireland | ben.b@example.com | 555-7870 | 246 Fir St, Waterford, IE |
53+
"""
54+
55+
INPUT_PROMPTS = [
56+
LONG_PROMPT +
57+
"Question: what is the age of John Doe? Your answer: The age of John Doe is ",
58+
LONG_PROMPT +
59+
"Question: what is the age of Zack Blue? Your answer: The age of Zack Blue is "
60+
]
61+
62+
63+
@pytest.mark.skipif(os.getenv("VLLM_USE_V1") == "0",
64+
reason="mtp is not supported on v1")
65+
@pytest.mark.parametrize("model", MODELS)
66+
@pytest.mark.parametrize("max_tokens", [50])
67+
def test_prefix_cache_with_v1_scheduler(model: str, max_tokens: int) -> None:
68+
with VllmRunner(model,
69+
enforce_eager=True,
70+
max_model_len=2048,
71+
tensor_parallel_size=2,
72+
gpu_memory_utilization=0.7) as vllm_model:
73+
prefix_cache_output = vllm_model.generate_greedy(
74+
INPUT_PROMPTS, max_tokens)
75+
76+
with VllmRunner(model,
77+
enable_prefix_caching=False,
78+
enforce_eager=True,
79+
max_model_len=2048,
80+
tensor_parallel_size=2,
81+
gpu_memory_utilization=0.7) as vllm_model:
82+
vllm_output = vllm_model.generate_greedy(INPUT_PROMPTS, max_tokens)
83+
84+
check_outputs_equal(
85+
outputs_0_lst=vllm_output,
86+
outputs_1_lst=prefix_cache_output,
87+
name_0="vllm_output",
88+
name_1="prefix_cache_output",
89+
)
90+
91+
92+
@pytest.mark.skipif(os.getenv("VLLM_USE_V1") == "0",
93+
reason="mtp is not supported on v1")
94+
@pytest.mark.parametrize("model", MODELS)
95+
@pytest.mark.parametrize("max_tokens", [50])
96+
def test_prefix_cache_with_ascend_scheduler(model: str,
97+
max_tokens: int) -> None:
98+
99+
with VllmRunner(model,
100+
additional_config={
101+
'ascend_scheduler_config': {
102+
'enabled': True,
103+
},
104+
},
105+
enforce_eager=True,
106+
max_model_len=2048,
107+
tensor_parallel_size=2,
108+
gpu_memory_utilization=0.7) as vllm_model:
109+
vllm_output = vllm_model.generate_greedy(INPUT_PROMPTS, max_tokens)
110+
111+
with VllmRunner(model,
112+
additional_config={
113+
'ascend_scheduler_config': {
114+
'enabled': True,
115+
'enable_prefix_caching': True,
116+
},
117+
},
118+
enforce_eager=True,
119+
max_model_len=2048,
120+
tensor_parallel_size=2,
121+
gpu_memory_utilization=0.7) as vllm_model:
122+
prefix_cache_output = vllm_model.generate_greedy(
123+
INPUT_PROMPTS, max_tokens)
124+
125+
with VllmRunner(model,
126+
additional_config={
127+
'ascend_scheduler_config': {
128+
'enabled': True,
129+
'enable_prefix_caching': True,
130+
"enable_chunked_prefill": True,
131+
},
132+
},
133+
enforce_eager=True,
134+
max_model_len=2048,
135+
tensor_parallel_size=2,
136+
gpu_memory_utilization=0.7) as vllm_model:
137+
chunk_prefill_prefix_cache_output = vllm_model.generate_greedy(
138+
INPUT_PROMPTS, max_tokens)
139+
140+
check_outputs_equal(
141+
outputs_0_lst=vllm_output,
142+
outputs_1_lst=prefix_cache_output,
143+
name_0="vllm_output",
144+
name_1="prefix_cache_output",
145+
)
146+
147+
check_outputs_equal(
148+
outputs_0_lst=chunk_prefill_prefix_cache_output,
149+
outputs_1_lst=prefix_cache_output,
150+
name_0="chunk_prefill_prefix_cache_output",
151+
name_1="prefix_cache_output",
152+
)

‎tests/e2e/prompts/example.txt

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.
2+
Briefly describe the major milestones in the development of artificial intelligence from 1950 to 2020.
3+
Compare and contrast artificial intelligence with human intelligence in terms of processing information.
4+
Describe the basic components of a neural network and how it can be trained.
5+
Write a short story about a robot that dreams for the first time.
6+
Analyze the impact of the COVID-19 pandemic on global economic structures and future business models.
7+
Explain the cultural significance of the Mona Lisa painting, and how its perception might vary in Western versus Eastern societies.
8+
Translate the following English sentence into Japanese, French, and Swahili: 'The early bird catches the worm.'

‎tests/e2e/singlecard/core/test_ascend_scheduler.py renamed to ‎tests/e2e/singlecard/core/ascend_scheduler/test_ascend_scheduler.py

File renamed without changes.

‎tests/e2e/singlecard/core/test_ascend_scheduler_e2e.py renamed to ‎tests/e2e/singlecard/core/ascend_scheduler/test_ascend_scheduler_e2e.py

File renamed without changes.

‎tests/e2e/singlecard/core/ascend_scheduler/test_chunk_prefill.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
"""Compare the with and without chunked prefill on AscendScheduler
3+
4+
It tests chunked prefill. Chunked prefill can be enabled by
5+
`additional_config={'ascend_scheduler_config': {'enabled': True, 'enable_chunked_prefill': True,},}`.
6+
If prefill size exceeds max_num_batched_tokens, prefill requests are chunked.
7+
8+
Run `pytest tests/e2e/singlecard/core/ascend_scheduler/test_chunk_prefill.py`.
9+
"""
10+
import os
11+
12+
import pytest
13+
14+
from tests.conftest import VllmRunner
15+
from tests.model_utils import check_outputs_equal
16+
17+
MODELS = [
18+
"Qwen/Qwen3-0.6B-Base",
19+
]
20+
21+
22+
@pytest.mark.skipif(os.getenv("VLLM_USE_V1") == "0", reason="only test on v1")
23+
@pytest.mark.parametrize("model", MODELS)
24+
@pytest.mark.parametrize("max_tokens",
25+
[4]) # cannot align results when max_tokens > 4
26+
@pytest.mark.parametrize("chunked_prefill_token_size", [16])
27+
def test_chunked_prefill_with_ascend_scheduler(
28+
example_prompts, model: str, max_tokens: int,
29+
chunked_prefill_token_size: int) -> None:
30+
max_num_seqs = chunked_prefill_token_size
31+
max_num_batched_tokens = chunked_prefill_token_size
32+
with VllmRunner(model,
33+
additional_config={
34+
'ascend_scheduler_config': {
35+
'enabled': True,
36+
'enable_chunked_prefill': True,
37+
},
38+
},
39+
max_num_seqs=max_num_seqs,
40+
max_num_batched_tokens=max_num_batched_tokens,
41+
enforce_eager=True,
42+
max_model_len=2048,
43+
gpu_memory_utilization=0.7) as vllm_model:
44+
chunked_prefill_output = vllm_model.generate_greedy(
45+
example_prompts, max_tokens)
46+
47+
with VllmRunner(model,
48+
additional_config={
49+
'ascend_scheduler_config': {
50+
'enabled': True,
51+
},
52+
},
53+
enforce_eager=True,
54+
max_model_len=2048,
55+
gpu_memory_utilization=0.7) as vllm_model:
56+
vllm_output = vllm_model.generate_greedy(example_prompts, max_tokens)
57+
58+
check_outputs_equal(
59+
outputs_0_lst=vllm_output,
60+
outputs_1_lst=chunked_prefill_output,
61+
name_0="vllm_output",
62+
name_1="chunked_prefill_output",
63+
)

0 commit comments

Comments
 (0)