@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6603288736
Fork 0

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 14, 05:30 UTC

[CI][V0 Deprecation] Removed V0 Only Chunked Prefill and Prefix Caching Tests (#22871)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Robert Shaw committed Aug 18, 2025 at 20:39 UTC 660328873619f7a711b33c7a6f1c5669a1ccdea2
6 files changed -595
.buildkite/test-pipeline.yaml
-18
@@ -88,15 +88,6 @@ steps:
88 - pytest -v -s basic_correctness/test_cpu_offload.py
89 - VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT=1 pytest -v -s basic_correctness/test_preemption.py
90
91 -- label: Chunked Prefill Test
92 - mirror_hardwares: [amdexperimental]
93 - source_file_dependencies:
94 - - vllm/
95 - - tests/basic_correctness/test_chunked_prefill
96 - commands:
97 - - VLLM_ATTENTION_BACKEND=XFORMERS pytest -v -s basic_correctness/test_chunked_prefill.py
98 - - VLLM_ATTENTION_BACKEND=FLASH_ATTN pytest -v -s basic_correctness/test_chunked_prefill.py
99 -
91 - label: Core Test # 10min
92 mirror_hardwares: [amdexperimental]
93 fast_check: true
@@ -295,15 +286,6 @@ steps:
286 - python3 offline_inference/basic/score.py
287 - VLLM_USE_V1=0 python3 offline_inference/profiling.py --model facebook/opt-125m run_num_steps --num-steps 2
288
298 -- label: Prefix Caching Test # 9min
299 - mirror_hardwares: [amdexperimental]
300 - source_file_dependencies:
301 - - vllm/
302 - - tests/prefix_caching
303 - commands:
304 - - pytest -v -s prefix_caching
305 -
306 -
289 - label: Platform Tests (CUDA)
290 mirror_hardwares: [amdexperimental]
291 source_file_dependencies:
.github/CODEOWNERS
-1
@@ -31,7 +31,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
31 # Test ownership
32 /.buildkite/lm-eval-harness @mgoin @simon-mo
33 /tests/async_engine @njhill @robertgshaw2-redhat @simon-mo
34 -/tests/basic_correctness/test_chunked_prefill @rkooo567 @comaniac
34 /tests/distributed/test_multi_node_assignment.py @youkaichao
35 /tests/distributed/test_pipeline_parallel.py @youkaichao
36 /tests/distributed/test_same_node.py @youkaichao
tests/basic_correctness/test_chunked_prefill.py deleted
-296
@@ -1,296 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Compare the outputs of HF and vLLM when using greedy sampling.
4 -
5 -It tests chunked prefill. Chunked prefill can be enabled by
6 -enable_chunked_prefill=True. If prefill size exceeds max_num_batched_tokens,
7 -prefill requests are chunked.
8 -
9 -Run `pytest tests/models/test_chunked_prefill.py`.
10 -"""
11 -
12 -from __future__ import annotations
13 -
14 -from typing import TYPE_CHECKING
15 -
16 -import pytest
17 -
18 -from vllm.platforms import current_platform
19 -from vllm.utils import STR_BACKEND_ENV_VAR
20 -
21 -from ..models.utils import check_logprobs_close, check_outputs_equal
22 -from ..utils import multi_gpu_test
23 -
24 -if TYPE_CHECKING:
25 - from .conftest import HfRunner, VllmRunner
26 -
27 -MODELS = [
28 - "facebook/opt-125m",
29 - "meta-llama/Llama-3.2-1B-Instruct",
30 -]
31 -
32 -
33 -@pytest.fixture(scope="function", autouse=True)
34 -def use_v0_only(monkeypatch: pytest.MonkeyPatch):
35 - """
36 - Since this module is V0 only, set VLLM_USE_V1=0 for
37 - all tests in the file.
38 - """
39 - with monkeypatch.context() as m:
40 - m.setenv('VLLM_USE_V1', '0')
41 - yield
42 -
43 -
44 -@pytest.mark.parametrize("model", MODELS)
45 -@pytest.mark.parametrize("dtype", ["half"])
46 -@pytest.mark.parametrize("max_tokens", [32])
47 -@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16])
48 -@pytest.mark.parametrize("enforce_eager", [False, True])
49 -# NOTE: Increasing this in this suite will fail CI because we currently cannot
50 -# reset distributed env properly. Use a value > 1 just when you test.
51 -@pytest.mark.parametrize("tensor_parallel_size", [1])
52 -@pytest.mark.parametrize("attention_backend", [
53 - pytest.param("FLASHINFER",
54 - marks=pytest.mark.skipif(
55 - current_platform.is_rocm(),
56 - reason="FLASHINFER isn't supported on ROCm")),
57 - "FLASH_ATTN"
58 -])
59 -def test_models(
60 - hf_runner: HfRunner,
61 - vllm_runner: VllmRunner,
62 - example_prompts,
63 - model: str,
64 - dtype: str,
65 - max_tokens: int,
66 - chunked_prefill_token_size: int,
67 - enforce_eager: bool,
68 - tensor_parallel_size: int,
69 - attention_backend: str,
70 - monkeypatch: pytest.MonkeyPatch,
71 -) -> None:
72 - """
73 - Checks exact match decode between huggingface model and vllm runner with
74 - chunked prefill.
75 - """
76 - with monkeypatch.context() as m:
77 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
78 -
79 - max_num_seqs = chunked_prefill_token_size
80 - max_num_batched_tokens = chunked_prefill_token_size
81 -
82 - with hf_runner(model, dtype=dtype) as hf_model:
83 - hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
84 -
85 - with vllm_runner(
86 - model,
87 - dtype=dtype,
88 - max_num_batched_tokens=max_num_batched_tokens,
89 - enable_chunked_prefill=True,
90 - tensor_parallel_size=tensor_parallel_size,
91 - enforce_eager=enforce_eager,
92 - max_num_seqs=max_num_seqs,
93 - ) as vllm_model:
94 - vllm_outputs = vllm_model.generate_greedy(example_prompts,
95 - max_tokens)
96 -
97 - check_outputs_equal(
98 - outputs_0_lst=hf_outputs,
99 - outputs_1_lst=vllm_outputs,
100 - name_0="hf",
101 - name_1="vllm",
102 - )
103 -
104 -
105 -@multi_gpu_test(num_gpus=2)
106 -@pytest.mark.parametrize("distributed_executor_backend", ["ray", "mp"])
107 -@pytest.mark.parametrize("model", MODELS)
108 -@pytest.mark.parametrize("attention_backend", [
109 - pytest.param("FLASHINFER",
110 - marks=pytest.mark.skipif(
111 - current_platform.is_rocm(),
112 - reason="FLASHINFER isn't supported on ROCm")),
113 - "FLASH_ATTN"
114 -])
115 -def test_models_distributed(
116 - hf_runner: HfRunner,
117 - vllm_runner: VllmRunner,
118 - example_prompts,
119 - model: str,
120 - distributed_executor_backend: str,
121 - attention_backend: str,
122 - monkeypatch: pytest.MonkeyPatch,
123 -) -> None:
124 - with monkeypatch.context() as m:
125 - m.setenv(STR_BACKEND_ENV_VAR, attention_backend)
126 - if (model == "meta-llama/Llama-3.2-1B-Instruct"
127 - and distributed_executor_backend == "ray"):
128 - # test Ray Compiled Graph
129 - m.setenv("VLLM_USE_RAY_SPMD_WORKER", "1")
130 - m.setenv("VLLM_USE_RAY_COMPILED_DAG", "1")
131 -
132 - dtype = "half"
133 - max_tokens = 5
134 - chunked_prefill_token_size = 16
135 -
136 - # Add a chunked prefill config.
137 - max_num_seqs = min(chunked_prefill_token_size, 256)
138 - assert chunked_prefill_token_size != -1
139 - enable_chunked_prefill = True
140 - max_num_batched_tokens = chunked_prefill_token_size
141 -
142 - # NOTE: take care of the order. run vLLM first, and then run HF.
143 - # vLLM needs a fresh new process without cuda initialization.
144 - # if we run HF first, the cuda initialization will be done and it
145 - # will hurt multiprocessing backend with
146 - # fork method (the default method).
147 -
148 - with vllm_runner(
149 - model,
150 - dtype=dtype,
151 - tensor_parallel_size=2,
152 - max_num_seqs=max_num_seqs,
153 - enable_chunked_prefill=enable_chunked_prefill,
154 - max_num_batched_tokens=max_num_batched_tokens,
155 - distributed_executor_backend=distributed_executor_backend,
156 - ) as vllm_model:
157 - vllm_outputs = vllm_model.generate_greedy(
158 - example_prompts,
159 - max_tokens,
160 - )
161 -
162 - with hf_runner(model, dtype=dtype) as hf_model:
163 - hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
164 -
165 - check_outputs_equal(
166 - outputs_0_lst=hf_outputs,
167 - outputs_1_lst=vllm_outputs,
168 - name_0="hf",
169 - name_1="vllm",
170 - )
171 -
172 -
173 -@pytest.mark.parametrize(
174 - "kv_cache_dtype,model",
175 - [("fp8_e4m3",
176 - "nm-testing/TinyLlama-1.1B-compressed-tensors-kv-cache-scheme")])
177 -# Due to low-precision numerical divergence, we only test logprob of 4 tokens
178 -@pytest.mark.parametrize("max_tokens", [4])
179 -@pytest.mark.parametrize("chunked_prefill_token_size", [4, 16])
180 -@pytest.mark.parametrize("enforce_eager", [False, True])
181 -# NOTE: Increasing this in this suite will fail CI because we currently cannot
182 -# reset distributed env properly. Use a value > 1 just when you test.
183 -@pytest.mark.parametrize("tensor_parallel_size", [1])
184 -# Due to low-precision numerical divergence, this test is too sensitive to
185 -# the async postprocessor
186 -@pytest.mark.parametrize("disable_async_output_proc", [True])
187 -@pytest.mark.skipif(current_platform.is_rocm(),
188 - reason="machete_prepack_B isn't supported on ROCm")
189 -def test_models_with_fp8_kv_cache(
190 - vllm_runner: VllmRunner,
191 - example_prompts,
192 - kv_cache_dtype: str,
193 - model: str,
194 - max_tokens: int,
195 - chunked_prefill_token_size: int,
196 - enforce_eager: bool,
197 - tensor_parallel_size: int,
198 - disable_async_output_proc: bool,
199 -) -> None:
200 - """
201 - Check output logprobs match between no_chunked_prefill and chunked_prefill
202 - with fp8 kv cache. General fp8 kv-cache tests are covered in test_fp8.py,
203 - so here we only check chunked prefill.
204 - """
205 - NUM_LOG_PROBS = 8
206 -
207 - max_num_seqs = chunked_prefill_token_size
208 - max_num_batched_tokens = chunked_prefill_token_size
209 -
210 - with vllm_runner(
211 - model,
212 - tensor_parallel_size=tensor_parallel_size,
213 - enforce_eager=enforce_eager,
214 - max_num_seqs=max_num_seqs,
215 - kv_cache_dtype=kv_cache_dtype,
216 - disable_async_output_proc=disable_async_output_proc,
217 - ) as vllm_model:
218 - no_chunked_prefill_outputs = vllm_model.generate_greedy_logprobs(
219 - example_prompts, max_tokens, NUM_LOG_PROBS)
220 -
221 - with vllm_runner(
222 - model,
223 - max_num_batched_tokens=max_num_batched_tokens,
224 - enable_chunked_prefill=True,
225 - tensor_parallel_size=tensor_parallel_size,
226 - enforce_eager=enforce_eager,
227 - max_num_seqs=max_num_seqs,
228 - kv_cache_dtype=kv_cache_dtype,
229 - disable_async_output_proc=disable_async_output_proc,
230 - ) as vllm_model:
231 - chunked_prefill_outputs = vllm_model.generate_greedy_logprobs(
232 - example_prompts, max_tokens, NUM_LOG_PROBS)
233 -
234 - check_logprobs_close(
235 - outputs_0_lst=no_chunked_prefill_outputs,
236 - outputs_1_lst=chunked_prefill_outputs,
237 - name_0="no_chunked_prefill",
238 - name_1="chunked_prefill",
239 - )
240 -
241 -
242 -@pytest.mark.parametrize("max_tokens", [16])
243 -@pytest.mark.parametrize("enforce_eager", [False])
244 -@pytest.mark.parametrize("chunk_size", [30, 32])
245 -# NOTE: Increasing this in this suite will fail CI because we currently cannot
246 -# reset distributed env properly. Use a value > 1 just when you test.
247 -@pytest.mark.parametrize("tensor_parallel_size", [1])
248 -@pytest.mark.parametrize("dtype", ["half"])
249 -def test_with_prefix_caching(
250 - vllm_runner: VllmRunner,
251 - max_tokens: int,
252 - enforce_eager: bool,
253 - chunk_size: int,
254 - tensor_parallel_size: int,
255 - dtype: str,
256 -) -> None:
257 - """
258 - Checks exact match decode with and without prefix caching
259 - with chunked prefill enabled.
260 - """
261 - model = "meta-llama/Llama-3.2-1B-Instruct"
262 - # The common prompt has 142 tokens with Llama-2 tokenizer.
263 - common_prompt = "You are a helpful AI assistant " * 20
264 - unique_prompts = [
265 - "Question", # Warmup
266 - "Question", # Fully cached
267 - "Another question", # Partial cached
268 - ]
269 - full_prompts = [f"{common_prompt}\n{p}" for p in unique_prompts]
270 -
271 - max_num_batched_tokens = max_num_seqs = chunk_size
272 - outputs = {} # type: ignore
273 - for enable in (True, False):
274 - with vllm_runner(
275 - model,
276 - dtype=dtype,
277 - max_num_batched_tokens=max_num_batched_tokens,
278 - enable_chunked_prefill=True,
279 - enable_prefix_caching=enable,
280 - tensor_parallel_size=tensor_parallel_size,
281 - enforce_eager=enforce_eager,
282 - max_num_seqs=max_num_seqs,
283 - ) as vllm_model:
284 - outputs[enable] = []
285 - for prompt in full_prompts:
286 - outputs[enable] += vllm_model.generate_greedy(
287 - [prompt],
288 - max_tokens,
289 - )
290 -
291 - check_outputs_equal(
292 - outputs_0_lst=outputs[False],
293 - outputs_1_lst=outputs[True],
294 - name_0="w/o prefix caching",
295 - name_1="with prefix caching",
296 - )
tests/prefix_caching/__init__.py
tests/prefix_caching/test_disable_sliding_window.py deleted
-49
@@ -1,49 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Compare the with and without prefix caching.
4 -
5 -Run `pytest tests/prefix_caching/test_prefix_caching.py`.
6 -"""
7 -import pytest
8 -
9 -from vllm import LLM
10 -from vllm.distributed import cleanup_dist_env_and_memory
11 -
12 -MODEL_LEN_LEN = [
13 - # Example models with sliding window.
14 - ("bigcode/starcoder2-3b", 4096, 16384),
15 - # ("mistralai/Mistral-7B-v0.1", 4096, 32768), << OOM in CI
16 -
17 - # Confirm model with sliding window works.
18 - # config has "use_sliding_window": false
19 - ("Qwen/Qwen1.5-0.5B-Chat", 32768, 32768),
20 - # config has no sliding window attribute.
21 - ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", 2048, 2048),
22 -]
23 -
24 -
25 -@pytest.mark.parametrize("model_len_len", MODEL_LEN_LEN)
26 -def test_disable_sliding_window(model_len_len, ):
27 - model, sliding_len, full_len = model_len_len
28 - disabled_llm = LLM(model, disable_sliding_window=True)
29 - disabled_llm.generate("Hi my name is")
30 - model_config = disabled_llm.llm_engine.model_config
31 - assert model_config.max_model_len == sliding_len, (
32 - "Max len expected to equal sliding_len of %s, but got %s", sliding_len,
33 - model_config.max_model_len)
34 -
35 - del disabled_llm
36 - cleanup_dist_env_and_memory()
37 -
38 - enabled_llm = LLM(model,
39 - enforce_eager=True,
40 - disable_sliding_window=False,
41 - enable_prefix_caching=False)
42 - enabled_llm.generate("Hi my name is")
43 - model_config = enabled_llm.llm_engine.model_config
44 - assert model_config.max_model_len == full_len, (
45 - "Max len expected to equal full_len of %s, but got %s", full_len,
46 - model_config.max_model_len)
47 -
48 - del enabled_llm
49 - cleanup_dist_env_and_memory()
tests/prefix_caching/test_prefix_caching.py deleted
-231
@@ -1,231 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Compare the with and without prefix caching.
4 -
5 -Run `pytest tests/prefix_caching/test_prefix_caching.py`.
6 -"""
7 -
8 -from __future__ import annotations
9 -
10 -import pytest
11 -
12 -from tests.conftest import VllmRunner
13 -from tests.core.utils import SchedulerProxy, create_dummy_prompt
14 -from vllm import SamplingParams, TokensPrompt
15 -from vllm.core.scheduler import Scheduler
16 -from vllm.engine.llm_engine import LLMEngine
17 -from vllm.platforms import current_platform
18 -from vllm.utils import STR_BACKEND_ENV_VAR
19 -
20 -from ..models.utils import check_outputs_equal
21 -
22 -
23 -@pytest.fixture(scope="function", autouse=True)
24 -def use_v0_only(monkeypatch: pytest.MonkeyPatch):
25 - """
26 - This module relies on V0 internals, so set VLLM_USE_V1=0.
27 - """
28 - with monkeypatch.context() as m:
29 - m.setenv('VLLM_USE_V1', '0')
30 - yield
31 -
32 -
33 -MODELS = [
34 - "distilbert/distilgpt2",
35 -]
36 -
37 -UNSTABLE_PROMPT_SEQUENCE = [
38 - ([0] * 588) + ([1] * 1332) + ([2] * 30) + ([3] * 1),
39 - ([0] * 588) + ([1] * 1332) + ([4] * 3) + ([5] * 50),
40 - ([0] * 588) + ([1] * 1332) + ([2] * 30) + ([6] * 95),
41 - ([0] * 588) + ([1] * 1332) + ([4] * 3) + ([7] * 174),
42 - ([0] * 588) + ([8] * 1539),
43 -]
44 -
45 -
46 -@pytest.mark.parametrize("model", MODELS)
47 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])
48 -@pytest.mark.parametrize("dtype", ["half"])
49 -@pytest.mark.parametrize("max_tokens", [5])
50 -@pytest.mark.parametrize("cached_position", [0, 1])
51 -@pytest.mark.parametrize("enable_chunked_prefill", [True, False])
52 -@pytest.mark.parametrize("block_size", [16])
53 -def test_mixed_requests(
54 - hf_runner,
55 - vllm_runner,
56 - example_prompts,
57 - model: str,
58 - backend: str,
59 - dtype: str,
60 - max_tokens: int,
61 - cached_position: int,
62 - enable_chunked_prefill: bool,
63 - block_size: int,
64 - monkeypatch: pytest.MonkeyPatch,
65 -) -> None:
66 - """
67 - Test the case when some sequences have the prefix cache hit
68 - and the others don't. The cached position determines where
69 - the sequence is at among the batch of prefills.
70 - """
71 - if backend == "FLASHINFER" and current_platform.is_rocm():
72 - pytest.skip("Flashinfer does not support ROCm/HIP.")
73 - if backend == "XFORMERS" and current_platform.is_rocm():
74 - pytest.skip("Xformers does not support ROCm/HIP.")
75 - with monkeypatch.context() as m:
76 - m.setenv(STR_BACKEND_ENV_VAR, backend)
77 -
78 - with hf_runner(model, dtype=dtype) as hf_model:
79 - hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
80 -
81 - cached_prompt = example_prompts[cached_position]
82 - with vllm_runner(
83 - model,
84 - dtype=dtype,
85 - enable_prefix_caching=True,
86 - enable_chunked_prefill=enable_chunked_prefill,
87 - block_size=block_size,
88 - ) as vllm_model:
89 - # Run the first prompt so the cache is populated
90 - vllm_outputs = vllm_model.generate_greedy([cached_prompt],
91 - max_tokens)
92 -
93 - # Run all the promopts
94 - greedy_params = SamplingParams(temperature=0.0,
95 - max_tokens=max_tokens)
96 - req_outputs = vllm_model.llm.generate(example_prompts,
97 - greedy_params)
98 -
99 - # Verify number of cached tokens
100 - for i in range(len(req_outputs)):
101 - if i == cached_position:
102 - expected_num_cached_tokens = (
103 - len(req_outputs[i].prompt_token_ids) //
104 - block_size) * block_size
105 - else:
106 - expected_num_cached_tokens = 0
107 - assert (req_outputs[i].num_cached_tokens ==
108 - expected_num_cached_tokens)
109 -
110 - vllm_outputs = [(
111 - output.prompt_token_ids + list(output.outputs[0].token_ids),
112 - output.prompt + output.outputs[0].text,
113 - ) for output in req_outputs]
114 -
115 - check_outputs_equal(
116 - outputs_0_lst=hf_outputs,
117 - outputs_1_lst=vllm_outputs,
118 - name_0="hf",
119 - name_1="vllm",
120 - )
121 -
122 -
123 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])
124 -def test_unstable_prompt_sequence(
125 - vllm_runner,
126 - backend: str,
127 - monkeypatch: pytest.MonkeyPatch,
128 -) -> None:
129 -
130 - if backend == "FLASHINFER" and current_platform.is_rocm():
131 - pytest.skip("Flashinfer does not support ROCm/HIP.")
132 - if backend == "XFORMERS" and current_platform.is_rocm():
133 - pytest.skip("Xformers does not support ROCm/HIP.")
134 - with monkeypatch.context() as m:
135 - m.setenv(STR_BACKEND_ENV_VAR, backend)
136 -
137 - with vllm_runner(
138 - "Qwen/Qwen2.5-0.5B-Instruct",
139 - enable_chunked_prefill=True,
140 - enable_prefix_caching=True,
141 - max_model_len=4096,
142 - ) as vllm_model:
143 - for prompt in UNSTABLE_PROMPT_SEQUENCE:
144 - vllm_model.generate(TokensPrompt(prompt_token_ids=prompt),
145 - SamplingParams(max_tokens=1))
146 -
147 -
148 -@pytest.mark.parametrize("model", MODELS)
149 -def test_fully_cached_prefill_needs_uncached_token(model):
150 - block_size = 16
151 - max_num_batched_tokens = 16
152 - num_output_tokens = 5
153 - # Make a vllm engine
154 - runner = VllmRunner(
155 - model_name=model,
156 - gpu_memory_utilization=0.7,
157 - enable_chunked_prefill=True,
158 - enforce_eager=True,
159 - enable_prefix_caching=True,
160 - block_size=block_size,
161 - max_num_batched_tokens=max_num_batched_tokens,
162 - max_num_seqs=max_num_batched_tokens,
163 - )
164 - engine: LLMEngine = runner.llm.llm_engine
165 -
166 - scheduler: Scheduler = SchedulerProxy(engine.scheduler[0]) # type: ignore
167 - engine.scheduler[0] = scheduler
168 -
169 - # SeqA
170 - seqA_tokens = list(range(2 * block_size))
171 - seqA, seq_groupA = create_dummy_prompt(
172 - request_id="0",
173 - prompt_tokens=seqA_tokens,
174 - max_tokens=num_output_tokens,
175 - block_size=block_size,
176 - )
177 -
178 - scheduler.add_seq_group(seq_groupA)
179 -
180 - assert seqA.data.get_num_computed_tokens() == 0
181 -
182 - # Prefill seqA
183 - while not seqA.is_finished():
184 - engine.step()
185 -
186 - # seqB
187 - seqB_tokens = [t + 1 for t in seqA_tokens] # shift by 1
188 - seqB, seq_groupB = create_dummy_prompt(
189 - request_id="1",
190 - prompt_tokens=seqB_tokens,
191 - max_tokens=num_output_tokens,
192 - block_size=block_size,
193 - )
194 -
195 - # seqC is the same as seqA
196 - seqC, seq_groupC = create_dummy_prompt(
197 - request_id="2",
198 - prompt_tokens=seqA_tokens,
199 - max_tokens=num_output_tokens,
200 - block_size=block_size,
201 - )
202 -
203 - scheduler.add_seq_group(seq_groupB)
204 - scheduler.add_seq_group(seq_groupC)
205 -
206 - # Even seqC is fully cached, it should not be prefilled since we
207 - # require at least 1 uncached token.
208 - engine.step()
209 -
210 - sched_metas, sched_out, _ = scheduler.last_schedule_ret()
211 - assert len(sched_out.scheduled_seq_groups) == 1
212 - assert (sched_out.scheduled_seq_groups[0].seq_group.request_id ==
213 - seq_groupB.request_id)
214 - assert (sched_out.scheduled_seq_groups[0].token_chunk_size ==
215 - max_num_batched_tokens)
216 -
217 - # When seqB is finished, seqC could be prefilled.
218 - while not seqB.is_finished():
219 - engine.step()
220 - sched_metas, sched_out, _ = scheduler.last_schedule_ret()
221 - assert len(sched_out.scheduled_seq_groups) == 1
222 - assert (sched_out.scheduled_seq_groups[0].seq_group.request_id ==
223 - seq_groupB.request_id)
224 -
225 - engine.step()
226 - sched_metas, sched_out, _ = scheduler.last_schedule_ret()
227 - assert len(sched_out.scheduled_seq_groups) == 1
228 - assert (sched_out.scheduled_seq_groups[0].seq_group.request_id ==
229 - seq_groupC.request_id)
230 - assert sched_out.scheduled_seq_groups[0].token_chunk_size == len(
231 - seqA_tokens)