@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4b946d693e
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

[V0 Deprecation] Remove V0 Core tests (#25082)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Sep 17, 2025 at 09:32 UTC 4b946d693e0af15740e9ca9c0e059d5f333b1083
20 files changed -5706
.buildkite/test-pipeline.yaml
-11
@@ -91,17 +91,6 @@ steps:
91 - pytest -v -s basic_correctness/test_cpu_offload.py
92 - VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT=1 pytest -v -s basic_correctness/test_preemption.py
93
94 -- label: Core Test # 22min
95 - timeout_in_minutes: 35
96 - mirror_hardwares: [amdexperimental]
97 - fast_check: true
98 - source_file_dependencies:
99 - - vllm/core
100 - - vllm/distributed
101 - - tests/core
102 - commands:
103 - - pytest -v -s core
104 -
94 - label: Entrypoints Unit Tests # 5min
95 timeout_in_minutes: 10
96 working_dir: "/vllm-workspace/tests"
tests/core/__init__.py
tests/core/block/__init__.py
tests/core/block/conftest.py deleted
-15
@@ -1,15 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -
7 -@pytest.fixture()
8 -def should_do_global_cleanup_after_test() -> bool:
9 - """Disable the global cleanup fixture for tests in this directory. This
10 - provides a ~10x speedup for unit tests that don't load a model to GPU.
11 -
12 - This requires that tests in this directory clean up after themselves if they
13 - use the GPU.
14 - """
15 - return False
tests/core/block/e2e/__init__.py
tests/core/block/e2e/conftest.py deleted
-71
@@ -1,71 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from collections.abc import Iterable
5 -from typing import Callable, Optional
6 -
7 -import pytest
8 -
9 -from vllm import LLM
10 -from vllm.distributed import cleanup_dist_env_and_memory
11 -from vllm.model_executor.utils import set_random_seed
12 -
13 -
14 -@pytest.fixture
15 -def baseline_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
16 - baseline_llm_kwargs, seed):
17 - return create_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
18 - baseline_llm_kwargs, seed)
19 -
20 -
21 -@pytest.fixture
22 -def test_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
23 - test_llm_kwargs, seed):
24 - return create_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
25 - test_llm_kwargs, seed)
26 -
27 -
28 -def create_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
29 - distinct_llm_kwargs, seed):
30 - kwargs = {
31 - **common_llm_kwargs,
32 - **per_test_common_llm_kwargs,
33 - **distinct_llm_kwargs,
34 - }
35 -
36 - def generator_inner():
37 - llm = LLM(**kwargs)
38 -
39 - set_random_seed(seed)
40 -
41 - yield llm
42 - del llm
43 - cleanup_dist_env_and_memory()
44 -
45 - for llm in generator_inner():
46 - yield llm
47 - del llm
48 -
49 -
50 -def get_text_from_llm_generator(llm_generator: Iterable[LLM],
51 - prompts,
52 - sampling_params,
53 - llm_cb: Optional[Callable[[LLM],
54 - None]] = None):
55 - for llm in llm_generator:
56 - if llm_cb:
57 - llm_cb(llm)
58 - outputs = llm.generate(prompts, sampling_params, use_tqdm=True)
59 - text = [output.outputs[0].text for output in outputs]
60 - del llm
61 -
62 - return text
63 -
64 -
65 -def get_token_ids_from_llm_generator(llm_generator, prompts, sampling_params):
66 - for llm in llm_generator:
67 - outputs = llm.generate(prompts, sampling_params, use_tqdm=True)
68 - token_ids = [output.outputs[0].token_ids for output in outputs]
69 - del llm
70 -
71 - return token_ids
tests/core/block/e2e/test_correctness.py deleted
-479
@@ -1,479 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from itertools import cycle
5 -
6 -import pytest
7 -
8 -from vllm import SamplingParams
9 -
10 -from .conftest import get_token_ids_from_llm_generator
11 -
12 -
13 -@pytest.mark.parametrize(
14 - "common_llm_kwargs",
15 - [{
16 - # Use a small model for a fast test.
17 - "model": "facebook/opt-125m",
18 -
19 - # skip cuda graph creation for fast test.
20 - "enforce_eager": True,
21 -
22 - # Allow only 5 sequences of ~1024 tokens in worst case.
23 - "block_size": 16,
24 - "num_gpu_blocks_override": 5 * (64 + 1),
25 - }])
26 -@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
27 -@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
28 -@pytest.mark.parametrize("test_llm_kwargs", [{
29 - "preemption_mode": "swap"
30 -}, {
31 - "preemption_mode": "recompute"
32 -}])
33 -@pytest.mark.parametrize("batch_size", [10])
34 -@pytest.mark.parametrize("seed", [1])
35 -def test_block_manager_with_preemption(baseline_llm_generator,
36 - test_llm_generator, batch_size):
37 - """Verify block manager produces same outputs even when there is preemption.
38 -
39 - This constructs two LLM, each with limited number of GPU blocks. The limit
40 - is decided such that as the sequences in the batch grow, sequences must be
41 - preempted and removed from cache.
42 -
43 - If the output token ids are equivalent, then we have confidence that the KV
44 - cache is not corrupted.
45 -
46 - NOTE: We want a significant number of generated tokens so that any incorrect
47 - KV mapping has time to build up error.
48 -
49 - NOTE(Kuntai): Though we have removed block manager v1, this test is still
50 - useful as it asserts the behavior of block manager v2 (now it is called
51 - SelfAttnBlockSpaceManager) is the same when swapping / preemption, so we
52 - keep this test.
53 - """
54 - output_len = 1024
55 - temperature = 0.0
56 -
57 - # We want to ensure equality even with preemption.
58 - # We force the total block size to be 1 + cdiv(output_len, block_size)
59 - # so that only one sequence can fit at a time (once the sequences grow).
60 -
61 - prompts = [
62 - "Hello, my name is",
63 - "The president of the United States is",
64 - "The capital of France is",
65 - "The future of AI is",
66 - ]
67 -
68 - prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
69 -
70 - sampling_params = SamplingParams(
71 - max_tokens=output_len,
72 - ignore_eos=True,
73 - temperature=temperature,
74 - )
75 -
76 - baseline_token_ids = get_token_ids_from_llm_generator(
77 - baseline_llm_generator, prompts, sampling_params)
78 -
79 - test_token_ids = get_token_ids_from_llm_generator(test_llm_generator,
80 - prompts, sampling_params)
81 -
82 - for expected_token_ids, actual_token_ids in zip(baseline_token_ids,
83 - test_token_ids):
84 - assert expected_token_ids == actual_token_ids
85 -
86 - assert baseline_token_ids == test_token_ids
87 -
88 -
89 -@pytest.mark.parametrize(
90 - "common_llm_kwargs",
91 - [{
92 - # Use a small model for a fast test.
93 - "model": "facebook/opt-125m",
94 -
95 - # Our prompts will generate 128 tokens; since the prompts themselves are
96 - # small, we don't need much KV space beyond 128.
97 - "max_model_len": 160,
98 -
99 - # skip cuda graph creation for fast test.
100 - "enforce_eager": True,
101 - }])
102 -@pytest.mark.parametrize(
103 - "per_test_common_llm_kwargs",
104 - [
105 - {
106 - "block_size": 16,
107 -
108 - # Allow only 2 sequences of ~128 tokens in worst case.
109 - # Note 8 = 128/block_size
110 - "num_gpu_blocks_override": 2 * (8 + 1),
111 - },
112 - {
113 - "block_size": 8,
114 -
115 - # Allow only 2 sequences of ~128 tokens in worst case.
116 - # Note 16 = 128/block_size
117 - "num_gpu_blocks_override": 2 * (16 + 2),
118 - }
119 - ])
120 -@pytest.mark.parametrize("baseline_llm_kwargs", [{
121 - "num_lookahead_slots": 0,
122 -}])
123 -@pytest.mark.parametrize(
124 - "test_llm_kwargs",
125 - [
126 - {
127 - # We run one test with block_size < lookahead_slots, one test with
128 - # block_size > lookahead_slots
129 - "num_lookahead_slots": 10,
130 - "preemption_mode": "swap",
131 - },
132 - {
133 - "num_lookahead_slots": 10,
134 - "preemption_mode": "recompute",
135 - }
136 - ])
137 -@pytest.mark.parametrize("batch_size", [4])
138 -@pytest.mark.parametrize("seed", [1])
139 -def test_lookahead_greedy_equality_with_preemption(baseline_llm_generator,
140 - test_llm_generator,
141 - batch_size):
142 - """Verify vLLM produces the same output with greedy sampling, when lookahead
143 - scheduling is used vs. not.
144 -
145 - Lookahead scheduling is not expected to modify the output, as it simply
146 - allocates empty slots ahead of the known token ids in a sliding fashion.
147 -
148 - This test constrains the total number of blocks to force preemption. It also
149 - varies the block size so that the lookahead size is less than and greater
150 - than the block size.
151 - """
152 - output_len = 128
153 - temperature = 0.0
154 -
155 - prompts = [
156 - "Hello, my name is",
157 - "The president of the United States is",
158 - "The capital of France is",
159 - "The future of AI is",
160 - ]
161 -
162 - prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
163 -
164 - sampling_params = SamplingParams(
165 - max_tokens=output_len,
166 - ignore_eos=True,
167 - temperature=temperature,
168 - )
169 -
170 - print('Getting token ids without lookahead scheduling')
171 - baseline_token_ids = get_token_ids_from_llm_generator(
172 - baseline_llm_generator, prompts, sampling_params)
173 -
174 - print('Getting token ids with lookahead scheduling')
175 - test_token_ids = get_token_ids_from_llm_generator(test_llm_generator,
176 - prompts, sampling_params)
177 -
178 - for expected_token_ids, actual_token_ids in zip(baseline_token_ids,
179 - test_token_ids):
180 - assert expected_token_ids == actual_token_ids
181 -
182 - assert baseline_token_ids == test_token_ids
183 -
184 -
185 -@pytest.mark.parametrize(
186 - "common_llm_kwargs",
187 - [
188 - {
189 - # Use a small model for a fast test.
190 - "model": "facebook/opt-125m",
191 -
192 - # skip cuda graph creation for fast test.
193 - "enforce_eager": True,
194 - "enable_chunked_prefill": True,
195 - },
196 - ])
197 -@pytest.mark.parametrize("per_test_common_llm_kwargs",
198 - [{
199 - "block_size": 16,
200 - "max_num_batched_tokens": 2,
201 - "max_num_seqs": 2,
202 - }, {
203 - "block_size": 16,
204 - "max_num_batched_tokens": 3,
205 - "max_num_seqs": 2,
206 - }, {
207 - "block_size": 16,
208 - "max_num_batched_tokens": 256,
209 - "max_num_seqs": 10,
210 - }])
211 -@pytest.mark.parametrize("baseline_llm_kwargs", [
212 - {},
213 -])
214 -@pytest.mark.parametrize("test_llm_kwargs", [
215 - {
216 - "num_lookahead_slots": 0,
217 - },
218 - {
219 - "num_lookahead_slots": 5,
220 - },
221 -])
222 -@pytest.mark.parametrize("batch_size", [4])
223 -@pytest.mark.parametrize("seed", [1])
224 -def test_chunked_prefill_block_manager(baseline_llm_generator,
225 - test_llm_generator, batch_size):
226 - """Verify that chunked prefill works with SelfAttnBlockSpaceManager,
227 - with and without lookahead scheduling.
228 - """
229 - output_len = 32
230 - temperature = 0.0
231 -
232 - prompts = [
233 - "Hello, my name is",
234 - "The president of the United States is",
235 - ("1 + " * 50) + " 1 = ", # Longer prompt.
236 - "The capital of France is",
237 - "The future of AI is",
238 - ]
239 -
240 - prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
241 -
242 - sampling_params = SamplingParams(
243 - max_tokens=output_len,
244 - ignore_eos=True,
245 - temperature=temperature,
246 - )
247 -
248 - print('Getting token ids with BlockManager')
249 - baseline_token_ids = get_token_ids_from_llm_generator(
250 - baseline_llm_generator, prompts, sampling_params)
251 -
252 - print('Getting token ids with BlockManager, with lookahead slots.')
253 - test_token_ids = get_token_ids_from_llm_generator(test_llm_generator,
254 - prompts, sampling_params)
255 -
256 - for expected_token_ids, actual_token_ids in zip(baseline_token_ids,
257 - test_token_ids):
258 - assert expected_token_ids == actual_token_ids
259 -
260 - assert baseline_token_ids == test_token_ids
261 -
262 -
263 -@pytest.mark.parametrize(
264 - "common_llm_kwargs",
265 - [{
266 - # Use a small model for a fast test.
267 - "model": "facebook/opt-125m",
268 -
269 - # skip cuda graph creation for fast test.
270 - "enforce_eager": True,
271 -
272 - # Allow only 5 sequences of ~1024 tokens in worst case.
273 - "block_size": 16,
274 - "num_gpu_blocks_override": 5 * (64 + 1),
275 -
276 - # Enable prefill cache
277 - "enable_prefix_caching": True,
278 - }])
279 -@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
280 -@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
281 -@pytest.mark.parametrize("test_llm_kwargs", [{
282 - "preemption_mode": "swap"
283 -}, {
284 - "preemption_mode": "recompute"
285 -}])
286 -@pytest.mark.parametrize("batch_size", [10])
287 -@pytest.mark.parametrize("seed", [1])
288 -def test_block_manager_prefix_caching_enabled_with_preemption(
289 - baseline_llm_generator, test_llm_generator, batch_size):
290 - """Verify block manager produces same outputs even when there is preemption.
291 -
292 - This constructs two LLM, each with limited number of GPU blocks. The limit
293 - is decided such that as the sequences in the batch grow, sequences must be
294 - preempted and removed from cache.
295 -
296 - If the output token ids are equivalent, then we have confidence that the KV
297 - cache is not corrupted.
298 -
299 - NOTE: We want a significant number of generated tokens so that any incorrect
300 - KV mapping has time to build up error.
301 -
302 - NOTE(Kuntai): Though we have removed block manager v1, this test is still
303 - useful as it asserts the behavior of block manager v2 (now it is called
304 - SelfAttnBlockSpaceManager) is the same when swapping / preemption, so we
305 - keep this test.
306 - """
307 - output_len = 1024
308 - temperature = 0.0
309 -
310 - # We want to ensure equality even with preemption.
311 - # We force the total block size to be 1 + cdiv(output_len, block_size)
312 - # so that only one sequence can fit at a time (once the sequences grow).
313 -
314 - prompts = [
315 - "Hello, my name is",
316 - "The president of the United States is",
317 - "The capital of France is",
318 - "The future of AI is",
319 - ]
320 -
321 - prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
322 -
323 - sampling_params = SamplingParams(
324 - max_tokens=output_len,
325 - ignore_eos=True,
326 - temperature=temperature,
327 - )
328 -
329 - print('Getting token ids from block manager')
330 - baseline_token_ids = get_token_ids_from_llm_generator(
331 - baseline_llm_generator, prompts, sampling_params)
332 -
333 - print('Getting token ids from block manager, with preemption')
334 - test_token_ids = get_token_ids_from_llm_generator(test_llm_generator,
335 - prompts, sampling_params)
336 -
337 - for expected_token_ids, actual_token_ids in zip(baseline_token_ids,
338 - test_token_ids):
339 - assert expected_token_ids == actual_token_ids
340 -
341 - assert baseline_token_ids == test_token_ids
342 -
343 -
344 -@pytest.mark.parametrize(
345 - "common_llm_kwargs",
346 - [{
347 - # Use a small model for a fast test.
348 - "model": "facebook/opt-125m",
349 -
350 - # skip cuda graph creation for fast test.
351 - "enforce_eager": True,
352 -
353 - # Allow only 5 sequences of ~1024 tokens in worst case.
354 - "block_size": 16,
355 - "num_gpu_blocks_override": 5 * (64 + 1),
356 - }])
357 -@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
358 -@pytest.mark.parametrize("baseline_llm_kwargs", [{
359 - "enable_prefix_caching": False
360 -}])
361 -@pytest.mark.parametrize("test_llm_kwargs", [{
362 - "enable_prefix_caching": True,
363 - "preemption_mode": "swap"
364 -}, {
365 - "enable_prefix_caching": True,
366 - "preemption_mode": "recompute"
367 -}])
368 -@pytest.mark.parametrize("batch_size", [10])
369 -@pytest.mark.parametrize("seed", [1])
370 -def test_auto_prefix_caching_with_preemption(baseline_llm_generator,
371 - test_llm_generator, batch_size):
372 - """Verify block manager v2 with auto prefix caching enabled produces same
373 - outputs as auto prefix caching disabled, even when there is preemption.
374 -
375 - This constructs two LLM, each with limited number of GPU blocks. The limit
376 - is decided such that as the sequences in the batch grow, sequences must be
377 - preempted and removed from cache.
378 -
379 - If the output token ids are equivalent, then we have confidence that auto
380 - prefix caching itself at least don't cause result error.
381 - """
382 - output_len = 1024
383 - temperature = 0.0
384 -
385 - # We want to ensure equality even with preemption.
386 - # We force the total block size to be 1 + cdiv(output_len, block_size)
387 - # so that only one sequence can fit at a time (once the sequences grow).
388 - prompts = [
389 - "Hello, my name is",
390 - "The president of the United States is",
391 - "The capital of France is",
392 - "The future of AI is",
393 - ]
394 -
395 - prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
396 -
397 - sampling_params = SamplingParams(
398 - max_tokens=output_len,
399 - ignore_eos=True,
400 - temperature=temperature,
401 - )
402 -
403 - print('Getting token ids with APC disabled')
404 - baseline_token_ids = get_token_ids_from_llm_generator(
405 - baseline_llm_generator, prompts, sampling_params)
406 -
407 - print('Getting token ids with APC enabled')
408 - test_token_ids = get_token_ids_from_llm_generator(test_llm_generator,
409 - prompts, sampling_params)
410 -
411 - for expected_token_ids, actual_token_ids in zip(baseline_token_ids,
412 - test_token_ids):
413 - assert expected_token_ids == actual_token_ids
414 -
415 - assert baseline_token_ids == test_token_ids
416 -
417 -
418 -@pytest.mark.parametrize(
419 - "common_llm_kwargs",
420 - [{
421 - # Use a small model for a fast test.
422 - "model": "facebook/opt-125m",
423 -
424 - # skip cuda graph creation for fast test.
425 - "enforce_eager": True,
426 -
427 - # we keep the blocks small, so that hit eviction quickly
428 - "max_model_len": 48,
429 - "block_size": 16,
430 - "num_gpu_blocks_override": 3,
431 - }])
432 -@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
433 -@pytest.mark.parametrize("baseline_llm_kwargs", [{
434 - "enable_prefix_caching": False
435 -}])
436 -@pytest.mark.parametrize("test_llm_kwargs", [{
437 - "enable_prefix_caching": True,
438 -}])
439 -@pytest.mark.parametrize("seed", [1])
440 -def test_auto_prefix_caching_after_eviction_start(baseline_llm_generator,
441 - test_llm_generator):
442 - """Verify block manager v2 with auto prefix caching could work normally
443 - even when eviction started.
444 - With APC enabled, all blocks are held by native block at the beginning.
445 - Then blocks are managed by evictor instead. If cache hit at the evictor's
446 - block, then it could be reused, or we need to recompute its kv cache.
447 - """
448 - output_len = 10
449 - temperature = 0.0
450 -
451 - prompts = [
452 - "You are a helpful assistant. Please answer truthfully and write "
453 - "out your thinking step by step to be sure you get the right answer. "
454 - "If you make a mistake, attempt to correct it. who are you?",
455 - "You are a helpful assistant. Please answer truthfully and write out "
456 - "your thinking step by step to be sure you get the right answer. You "
457 - "are helpful and harmless and you follow ethical guidelines. "
458 - "who are you?"
459 - ]
460 -
461 - sampling_params = SamplingParams(
462 - max_tokens=output_len,
463 - ignore_eos=True,
464 - temperature=temperature,
465 - )
466 -
467 - print('Getting token ids with APC disabled')
468 - baseline_token_ids = get_token_ids_from_llm_generator(
469 - baseline_llm_generator, prompts, sampling_params)
470 -
471 - print('Getting token ids with APC enabled')
472 - test_token_ids = get_token_ids_from_llm_generator(test_llm_generator,
473 - prompts, sampling_params)
474 -
475 - for expected_token_ids, actual_token_ids in zip(baseline_token_ids,
476 - test_token_ids):
477 - assert expected_token_ids == actual_token_ids
478 -
479 - assert baseline_token_ids == test_token_ids
tests/core/block/e2e/test_correctness_sliding_window.py deleted
-185
@@ -1,185 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import random
5 -
6 -import pytest
7 -
8 -from tests.kernels.utils import override_backend_env_variable
9 -from vllm import LLM, SamplingParams
10 -from vllm.platforms import current_platform
11 -
12 -from .conftest import get_text_from_llm_generator
13 -
14 -# relatively small model with 4k sliding window
15 -MODEL = "bigcode/starcoder2-3b"
16 -BLOCK_SIZE = 16
17 -
18 -
19 -@pytest.mark.parametrize(
20 - "common_llm_kwargs",
21 - [{
22 - "model": MODEL,
23 -
24 - # skip cuda graph creation for fast test.
25 - "enforce_eager": True,
26 - "block_size": BLOCK_SIZE,
27 - # needed due to https://github.com/vllm-project/vllm/issues/1908#issuecomment-2101122008
28 - "num_gpu_blocks_override": 100000 // BLOCK_SIZE,
29 - }])
30 -@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
31 -@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
32 -@pytest.mark.parametrize("test_llm_kwargs", [{}])
33 -@pytest.mark.parametrize("batch_size", [5])
34 -@pytest.mark.parametrize("seed", [1])
35 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "XFORMERS"])
36 -def test_sliding_window_retrieval(baseline_llm_generator, test_llm_generator,
37 - batch_size, seed, backend, monkeypatch):
38 - """
39 - The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
40 - asks for value of one of them (which is outside the sliding window).
41 - If we tell it upfront which we are going to be looking for, then
42 - it answers correctly (mostly).
43 -
44 - Additionally, we compare the results of the v1 and v2 managers.
45 - """
46 - if backend == "XFORMERS" and current_platform.is_rocm():
47 - pytest.skip("Xformers does not support ROCm/HIP.")
48 -
49 - override_backend_env_variable(monkeypatch, backend)
50 -
51 - sampling_params = SamplingParams(
52 - max_tokens=1024,
53 - ignore_eos=True,
54 - temperature=0.0,
55 - )
56 -
57 - prompts, answer, indices = prep_prompts(batch_size)
58 -
59 - baseline_texts = get_text_from_llm_generator(baseline_llm_generator,
60 - prompts,
61 - sampling_params,
62 - llm_cb=check_window(prompts))
63 -
64 - check_answers(indices, answer, baseline_texts)
65 -
66 - print('Getting token ids from block manager v2')
67 - test_texts = get_text_from_llm_generator(test_llm_generator, prompts,
68 - sampling_params)
69 - check_answers(indices, answer, test_texts)
70 -
71 - cmp = [
72 - expected_text == actual_text
73 - for expected_text, actual_text in zip(baseline_texts, test_texts)
74 - ]
75 - print(cmp)
76 - # make sure it's mostly OK; this is possibly because https://github.com/vllm-project/vllm/pull/4768
77 - # however, https://github.com/vllm-project/vllm/issues/3385#issuecomment-1995924290
78 - # states that xformers and flash_attn have different ideas about the window
79 - # size anyways
80 - assert sum(cmp) > 0.7 * len(cmp)
81 -
82 -
83 -@pytest.mark.parametrize(
84 - "common_llm_kwargs",
85 - [{
86 - "model": MODEL,
87 -
88 - # skip cuda graph creation for fast test.
89 - "enforce_eager": True,
90 - "block_size": BLOCK_SIZE,
91 - "num_gpu_blocks_override": 100000 // BLOCK_SIZE,
92 - }])
93 -@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
94 -@pytest.mark.parametrize("test_llm_kwargs", [{"enable_chunked_prefill": True}])
95 -@pytest.mark.parametrize("batch_size", [5])
96 -@pytest.mark.parametrize("seed", [1])
97 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "XFORMERS"])
98 -def test_sliding_window_chunked_prefill(test_llm_generator, batch_size, seed,
99 - backend, monkeypatch):
100 - """
101 - This is similar to test_sliding_window_retrieval, however, it doesn't
102 - compare against the v1 block manager since v1 doesn't support
103 - chunked prefill with sliding window.
104 -
105 - The results with and without chunked prefill are not the same due to
106 - numerical instabilities.
107 - """
108 - if backend == "XFORMERS" and current_platform.is_rocm():
109 - pytest.skip("Xformers does not support ROCm/HIP.")
110 - override_backend_env_variable(monkeypatch, backend)
111 -
112 - sampling_params = SamplingParams(
113 - max_tokens=10,
114 - ignore_eos=True,
115 - temperature=0.0,
116 - )
117 -
118 - prompts, answer, indices = prep_prompts(batch_size)
119 -
120 - # We don't compare with the baseline model here, since the results
121 - # slightly different due to different tailing in attention.
122 - test_texts = get_text_from_llm_generator(test_llm_generator,
123 - prompts,
124 - sampling_params,
125 - llm_cb=check_window(prompts))
126 - check_answers(indices, answer, test_texts)
127 -
128 -
129 -def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
130 - """
131 - Generate prompts which a bunch of assignments,
132 - then asking for the value of one of them.
133 - The prompt is just under 10k tokens; sliding window is 4k
134 - so the answer is outside sliding window, but should still be correct.
135 -
136 - Args:
137 - batch_size: number of prompts to generate
138 - ln_range: an argument to control the length of the prompt
139 - """
140 - prompts: list[str] = []
141 - answer: list[int] = []
142 - indices: list[int] = []
143 - random.seed(1)
144 - for _ in range(batch_size):
145 - idx = random.randint(30, 90)
146 - indices.append(idx)
147 - prompt = "```python\n# We set a number of variables, " + \
148 - f"x{idx} will be important later\n"
149 - ln = random.randint(*ln_range)
150 - for k in range(30, ln):
151 - v = random.randint(10, 99)
152 - if k == idx:
153 - answer.append(v)
154 - prompt += f"x{k} = {v}\n"
155 - prompt += f"# Now, we check the value of x{idx}:\n"
156 - prompt += f"assert x{idx} == "
157 - prompts.append(prompt)
158 - return prompts, answer, indices
159 -
160 -
161 -def check_answers(indices: list[int],
162 - answer: list[int],
163 - outputs: list[str],
164 - accept_rate: float = 0.7):
165 - answer2 = [int(text[0:2].strip()) for text in outputs]
166 - print(list(zip(indices, zip(answer, answer2))))
167 - numok = 0
168 - for a1, a2 in zip(answer, answer2):
169 - if a1 == a2:
170 - numok += 1
171 - frac_ok = numok / len(answer)
172 - print(f"Num OK: {numok}/{len(answer)} {frac_ok}")
173 - assert frac_ok >= accept_rate
174 -
175 -
176 -def check_window(prompts: list[str]):
177 -
178 - def inner(llm: LLM):
179 - sliding_window = llm.llm_engine.model_config.get_sliding_window()
180 - assert sliding_window and sliding_window > 0
181 - assert any(
182 - len(llm.get_tokenizer().tokenize(prompt)) > sliding_window
183 - for prompt in prompts)
184 -
185 - return inner
tests/core/block/test_block_manager.py deleted
-341
@@ -1,341 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from vllm.core.block_manager import SelfAttnBlockSpaceManager
7 -from vllm.core.interfaces import AllocStatus
8 -from vllm.sequence import Logprob, SequenceStatus
9 -from vllm.utils import chunk_list
10 -
11 -from ..utils import create_dummy_prompt, create_seq_group
12 -
13 -
14 -@pytest.mark.parametrize("block_size", [16])
15 -@pytest.mark.parametrize("num_gpu_blocks", [8, 40, 80])
16 -@pytest.mark.parametrize("num_seqs_per_group", [1, 4])
17 -@pytest.mark.parametrize("watermark", [0.0, 0.5])
18 -def test_can_allocate_seq_group(block_size: int, num_seqs_per_group: int,
19 - num_gpu_blocks: int, watermark: float):
20 - block_manager = SelfAttnBlockSpaceManager(
21 - block_size=block_size,
22 - num_gpu_blocks=num_gpu_blocks,
23 - num_cpu_blocks=1024,
24 - watermark=watermark,
25 - )
26 - num_watermark_blocks = int(watermark * num_gpu_blocks)
27 -
28 - num_output_blocks_per_seq = 1
29 -
30 - # NOTE: This should be num_output_blocks_per_seq * num_seqs_per_group, but
31 - # the current implementation assumes all seqs are new prompts / don't have
32 - # different output lens.
33 - num_output_blocks = num_output_blocks_per_seq
34 -
35 - for num_prompt_blocks in range(1, num_gpu_blocks - num_output_blocks):
36 - seq_group = create_seq_group(
37 - seq_prompt_len=block_size * num_prompt_blocks,
38 - seq_output_lens=[
39 - block_size * num_output_blocks_per_seq
40 - for _ in range(num_seqs_per_group)
41 - ],
42 - )
43 -
44 - assert num_prompt_blocks + num_output_blocks <= num_gpu_blocks
45 -
46 - can_allocate_result = block_manager.can_allocate(seq_group)
47 -
48 - num_required_blocks = num_prompt_blocks + num_output_blocks
49 -
50 - if num_gpu_blocks - num_required_blocks < num_watermark_blocks:
51 - assert can_allocate_result == AllocStatus.NEVER
52 - elif num_gpu_blocks >= num_required_blocks:
53 - assert can_allocate_result == AllocStatus.OK
54 - else:
55 - assert can_allocate_result == AllocStatus.LATER
56 -
57 -
58 -@pytest.mark.parametrize("block_size", [1, 8])
59 -@pytest.mark.parametrize("prompt_len", [1, 7, 8])
60 -@pytest.mark.parametrize("num_slots_to_append", [1, 8, 129])
61 -@pytest.mark.parametrize("num_lookahead_slots", [0, 10])
62 -def test_append_slots(block_size, prompt_len, num_slots_to_append,
63 - num_lookahead_slots):
64 - """Verify append_slots consumes the correct number of blocks from the block
65 - table.
66 - """
67 -
68 - num_gpu_blocks = 1024
69 - watermark = 0.1
70 - block_manager = SelfAttnBlockSpaceManager(
71 - block_size=block_size,
72 - num_gpu_blocks=num_gpu_blocks,
73 - num_cpu_blocks=0,
74 - watermark=watermark,
75 - )
76 -
77 - seq_group = create_seq_group(
78 - seq_prompt_len=prompt_len,
79 - seq_output_lens=[0],
80 - )
81 -
82 - # Allocate seq
83 - assert block_manager.can_allocate(seq_group)
84 - block_manager.allocate(seq_group)
85 -
86 - # Seq seq to RUNNING
87 - seq = seq_group.get_seqs()[0]
88 - seq.status = SequenceStatus.RUNNING
89 -
90 - # Append tokens to the sequeqnce
91 - for token_id in range(num_slots_to_append):
92 - seq.append_token_id(token_id, {token_id: Logprob(0.0)})
93 -
94 - # Append slots for new tokens and lookahead slots.
95 - free_blocks_before_append = block_manager.get_num_free_gpu_blocks()
96 - block_manager.append_slots(seq, num_lookahead_slots)
97 - num_consumed_blocks = (free_blocks_before_append -
98 - block_manager.get_num_free_gpu_blocks())
99 -
100 - # Expect consumed blocks to be new blocks required to support the new slots.
101 - expected_consumed_blocks = len(
102 - list(
103 - chunk_list(
104 - list(
105 - range(prompt_len + num_slots_to_append +
106 - num_lookahead_slots)),
107 - block_size))) - len(
108 - list(chunk_list(list(range(prompt_len)), block_size)))
109 - assert num_consumed_blocks == expected_consumed_blocks
110 -
111 -
112 -@pytest.mark.parametrize("block_size", [8])
113 -@pytest.mark.parametrize("num_cpu_blocks", [4])
114 -@pytest.mark.parametrize("num_gpu_blocks", [4])
115 -@pytest.mark.parametrize("num_lookahead_slots", [0, 2, 10])
116 -@pytest.mark.parametrize("enable_caching", [False, True])
117 -def test_swap(block_size, num_cpu_blocks, num_gpu_blocks, num_lookahead_slots,
118 - enable_caching):
119 - """Verify blocks number on src/desc device is correct after swapping in/out
120 - sequence group (not missing or extra blocks).
121 - """
122 - block_manager = SelfAttnBlockSpaceManager(block_size,
123 - num_cpu_blocks,
124 - num_gpu_blocks,
125 - watermark=0,
126 - enable_caching=enable_caching)
127 - prompt, seq_group = create_dummy_prompt("1", prompt_length=block_size - 1)
128 - prompt.status = SequenceStatus.WAITING
129 - block_manager.allocate(seq_group)
130 -
131 - # Emulate a forward pass by appending a single token.
132 - # The block manager then knows how many unprocessed
133 - # tokens will be written in the next forward pass.
134 - token_id = 0
135 - prompt.status = SequenceStatus.RUNNING
136 - prompt.append_token_id(token_id, {token_id: Logprob(0.0)})
137 -
138 - # Swap seq group from GPU -> CPU.
139 - gpu_blocks = block_manager.get_block_table(prompt)
140 - assert block_manager.can_swap_out(seq_group)
141 - before_cpu_blocks = block_manager.get_num_free_cpu_blocks()
142 - before_gpu_blocks = block_manager.get_num_free_gpu_blocks()
143 - mapping = block_manager.swap_out(seq_group)
144 - mapping_keys = [key for key, _ in mapping]
145 - assert mapping_keys == gpu_blocks
146 - after_cpu_blocks = block_manager.get_num_free_cpu_blocks()
147 - after_gpu_blocks = block_manager.get_num_free_gpu_blocks()
148 - assert before_cpu_blocks == after_cpu_blocks + len(gpu_blocks)
149 - assert before_gpu_blocks + len(gpu_blocks) == after_gpu_blocks
150 - prompt.status = SequenceStatus.SWAPPED
151 -
152 - # Swap seq group from CPU -> GPU.
153 - assert block_manager.can_swap_in(seq_group, num_lookahead_slots)
154 - before_cpu_blocks = block_manager.get_num_free_cpu_blocks()
155 - before_gpu_blocks = block_manager.get_num_free_gpu_blocks()
156 - mapping = block_manager.swap_in(seq_group)
157 - cpu_blocks = block_manager.get_block_table(prompt)
158 - mapping_keys = [key for key, _ in mapping]
159 - assert mapping_keys == [cpu_blocks[0]]
160 - after_cpu_blocks = block_manager.get_num_free_cpu_blocks()
161 - after_gpu_blocks = block_manager.get_num_free_gpu_blocks()
162 - assert before_gpu_blocks == after_gpu_blocks + len(cpu_blocks)
163 -
164 -
165 -@pytest.mark.parametrize("block_size", [8])
166 -@pytest.mark.parametrize("num_gpu_blocks", [4])
167 -@pytest.mark.parametrize("num_lookahead_slots", [3, 8, 10])
168 -@pytest.mark.parametrize("enable_caching", [True, False])
169 -def test_can_swap(block_size, num_gpu_blocks, num_lookahead_slots,
170 - enable_caching):
171 - """ Verify the block manager can correctly determine if a sequence group
172 - can be swapped in/out.
173 - """
174 - num_cpu_blocks = num_gpu_blocks
175 - block_manager = SelfAttnBlockSpaceManager(block_size,
176 - num_cpu_blocks,
177 - num_gpu_blocks,
178 - watermark=0,
179 - enable_caching=enable_caching)
180 - prompt, seq_group = create_dummy_prompt(
181 - "1", prompt_length=(num_gpu_blocks - 1) * block_size - 1)
182 - prompt.status = SequenceStatus.WAITING
183 - block_manager.allocate(seq_group)
184 - prompt.status = SequenceStatus.RUNNING
185 -
186 - # Swap seq group from GPU -> CPU.
187 - gpu_blocks = block_manager.get_block_table(prompt)
188 - assert block_manager.can_swap_out(seq_group)
189 - before_cpu_blocks = block_manager.get_num_free_cpu_blocks()
190 - before_gpu_blocks = block_manager.get_num_free_gpu_blocks()
191 - mapping = block_manager.swap_out(seq_group)
192 - mapping_keys = [key for key, _ in mapping]
193 - assert mapping_keys == gpu_blocks
194 - after_cpu_blocks = block_manager.get_num_free_cpu_blocks()
195 - after_gpu_blocks = block_manager.get_num_free_gpu_blocks()
196 - assert before_cpu_blocks == after_cpu_blocks + len(gpu_blocks)
197 - assert before_gpu_blocks + len(gpu_blocks) == after_gpu_blocks
198 - prompt.status = SequenceStatus.SWAPPED
199 -
200 - # At this moment, we still have enough free blocks to swap in the seq group.
201 - if num_lookahead_slots <= block_size:
202 - assert block_manager.can_swap_in(seq_group,
203 - num_lookahead_slots) == AllocStatus.OK
204 - else:
205 - assert block_manager.can_swap_in(
206 - seq_group, num_lookahead_slots) == AllocStatus.NEVER
207 -
208 - # During Swapped out, 2 cached blocks were evicted from the GPU,
209 - # so the prompt1 can't be swapped in
210 - prompt2_len = 2 * block_size - 1
211 - prompt2, seq_group2 = create_dummy_prompt(
212 - "2",
213 - prompt_length=prompt2_len,
214 - prompt_tokens=[10000 + i for i in range(prompt2_len)])
215 - prompt2.status = SequenceStatus.WAITING
216 - block_manager.allocate(seq_group2)
217 -
218 - # Swap seq group from CPU -> GPU.
219 - if num_lookahead_slots <= block_size:
220 - assert block_manager.can_swap_in(
221 - seq_group, num_lookahead_slots) == AllocStatus.LATER
222 - else:
223 - assert block_manager.can_swap_in(
224 - seq_group, num_lookahead_slots) == AllocStatus.NEVER
225 -
226 -
227 -@pytest.mark.parametrize("num_lookahead_slots", [0, 2, 10])
228 -@pytest.mark.parametrize("enable_caching", [False, True])
229 -def test_swap_in_infeasible(num_lookahead_slots, enable_caching):
230 - """Verifies that swapping fails if there is not enough free blocks
231 - to account for unseen tokens and lookahead_slots.
232 - """
233 - block_size = 8
234 - num_cpu_blocks = 1
235 - num_gpu_blocks = 1
236 - block_manager = SelfAttnBlockSpaceManager(block_size,
237 - num_cpu_blocks,
238 - num_gpu_blocks,
239 - watermark=0,
240 - enable_caching=enable_caching)
241 - prompt_length = block_size - 3
242 - assert prompt_length > 0
243 - prompt, seq_group = create_dummy_prompt("1", prompt_length=prompt_length)
244 - prompt.status = SequenceStatus.WAITING
245 - block_manager.allocate(seq_group)
246 - # Emulate a forward pass by appending a single token.
247 - # The block manager then knows how many unprocessed
248 - # tokens will be written in the next forward pass.
249 - token_id = 0
250 - prompt.status = SequenceStatus.RUNNING
251 - prompt.append_token_id(token_id, {token_id: Logprob(0.0)})
252 -
253 - # Swap seq group from GPU -> CPU.
254 - assert block_manager.can_swap_out(seq_group)
255 - block_manager.swap_out(seq_group)
256 - prompt.status = SequenceStatus.SWAPPED
257 -
258 - # Swap seq group from CPU -> GPU.
259 - # The number of unseen tokens is 1. If the number of existing
260 - # tokens plus the unseen ones and number of lookahead slots exceeds
261 - # the total number of available GPU blocks then the swap
262 - # should fail.
263 - num_unseen_tokens = 1
264 - if (num_lookahead_slots + num_unseen_tokens +
265 - prompt_length) <= (block_size * num_gpu_blocks):
266 - assert block_manager.can_swap_in(seq_group,
267 - num_lookahead_slots) == AllocStatus.OK
268 - else:
269 - assert block_manager.can_swap_in(
270 - seq_group, num_lookahead_slots) == AllocStatus.NEVER
271 -
272 -
273 -# TODO(cade/kaiyang): add comprehensive tests for swapping at allocator level.
274 -
275 -
276 -@pytest.mark.parametrize("block_size", [8, 16])
277 -@pytest.mark.parametrize("prompt_len", [10, 300, 1000])
278 -@pytest.mark.parametrize("num_slots_to_append", [50])
279 -@pytest.mark.parametrize("sliding_window", [20, 32, 200, 512])
280 -def test_sliding_window(block_size, prompt_len, num_slots_to_append,
281 - sliding_window):
282 - """Verify append_slots consumes the correct number of blocks from the block
283 - table.
284 - """
285 -
286 - num_gpu_blocks = 1024
287 - watermark = 0.1
288 - block_manager = SelfAttnBlockSpaceManager(
289 - block_size=block_size,
290 - num_gpu_blocks=num_gpu_blocks,
291 - num_cpu_blocks=0,
292 - watermark=watermark,
293 - sliding_window=sliding_window,
294 - )
295 -
296 - def check_used(min_n, max_n=None):
297 - if max_n is None:
298 - max_n = min_n
299 - used = num_gpu_blocks - block_manager.get_num_free_gpu_blocks()
300 - assert min_n <= used
301 - assert used <= max_n
302 -
303 - def num_blocks(num_tokens):
304 - return (num_tokens + block_size - 1) // block_size
305 -
306 - check_used(0)
307 -
308 - seq_group = create_seq_group(
309 - seq_prompt_len=prompt_len,
310 - seq_output_lens=[0],
311 - )
312 -
313 - check_used(0)
314 -
315 - # Allocate seq
316 - assert block_manager.can_allocate(seq_group)
317 - block_manager.allocate(seq_group)
318 -
319 - check_used(num_blocks(prompt_len))
320 -
321 - # Seq seq to RUNNING
322 - seq = seq_group.get_seqs()[0]
323 - seq.status = SequenceStatus.RUNNING
324 -
325 - seq.data.update_num_computed_tokens(prompt_len)
326 - check_used(num_blocks(prompt_len))
327 -
328 - # this is how we compute it in SelfAttnBlockSpaceManager.__init__
329 - sliding_blocks = (sliding_window // block_size) + 2
330 - # plus one block for null block
331 - sliding_blocks += 1
332 -
333 - # Append tokens to the sequeqnce
334 - for token_id in range(num_slots_to_append):
335 - seq.append_token_id(token_id, {token_id: Logprob(0.0)})
336 - seq.data.update_num_computed_tokens(1)
337 - block_manager.append_slots(seq, num_lookahead_slots=0)
338 - if prompt_len < sliding_window + 10:
339 - check_used(0, sliding_blocks + 1)
340 - else:
341 - check_used(sliding_blocks, sliding_blocks + 1)
tests/core/block/test_block_table.py deleted
-577
@@ -1,577 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from vllm.core.block.block_table import BlockTable
7 -from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
8 -from vllm.utils import Device, cdiv, chunk_list
9 -
10 -
11 -@pytest.mark.parametrize("block_size", [16])
12 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
13 -def test_allocate_naive(block_size: int, sequence_len: int):
14 - """Test the allocation of blocks using the naive allocator.
15 -
16 - This test creates a CpuGpuBlockAllocator with the specified block size and
17 - number of blocks. It then allocates multiple BlockTables with varying
18 - sequence lengths and verifies that the number of free blocks decreases as
19 - expected after each allocation.
20 - """
21 - assert block_size > 1
22 - num_gpu_blocks = 1024
23 -
24 - allocator = CpuGpuBlockAllocator.create(
25 - allocator_type="naive",
26 - num_gpu_blocks=num_gpu_blocks,
27 - num_cpu_blocks=1024,
28 - block_size=block_size,
29 - )
30 -
31 - token_ids = list(range(sequence_len))
32 - num_blocks_per_alloc = len(list(chunk_list(token_ids, block_size)))
33 -
34 - block_tables: list[BlockTable] = []
35 - for i in range(5):
36 - assert allocator.get_num_free_blocks(
37 - device=Device.GPU) == num_gpu_blocks - i * num_blocks_per_alloc
38 -
39 - block_tables.append(
40 - BlockTable(
41 - block_size=block_size,
42 - block_allocator=allocator,
43 - ))
44 - block_tables[-1].allocate(token_ids=token_ids, device=Device.GPU)
45 -
46 -
47 -@pytest.mark.parametrize("block_size", [16])
48 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
49 -def test_allocate_prefix_caching(block_size: int, sequence_len: int):
50 - """Test the allocation of blocks using the prefix caching allocator.
51 -
52 - This test creates a CpuGpuBlockAllocator with the specified block size and
53 - number of blocks, using the prefix caching allocator. It then allocates
54 - multiple BlockTables with varying sequence lengths and verifies that the
55 - number of free blocks decreases as expected after each allocation.
56 -
57 - The test expects all sequences to share allocations, except for their last
58 - block, which may be mutable. It calculates the expected number of immutable
59 - and mutable blocks per allocation based on the sequence length and block
60 - size.
61 - """
62 - assert block_size > 1
63 - num_gpu_blocks = 1024
64 -
65 - allocator = CpuGpuBlockAllocator.create(
66 - allocator_type="prefix_caching",
67 - num_gpu_blocks=num_gpu_blocks,
68 - num_cpu_blocks=1024,
69 - block_size=block_size,
70 - )
71 -
72 - token_ids = list(range(sequence_len))
73 - chunked_tokens = list(chunk_list(token_ids, block_size))
74 - num_mutable_blocks_per_alloc = 0 if len(
75 - chunked_tokens[-1]) == block_size else 1
76 - num_immutable_blocks_per_alloc = len(
77 - chunked_tokens) - num_mutable_blocks_per_alloc
78 -
79 - block_tables: list[BlockTable] = []
80 - for alloc_i in range(1, 6):
81 -
82 - block_tables.append(
83 - BlockTable(
84 - block_size=block_size,
85 - block_allocator=allocator,
86 - ))
87 - block_tables[-1].allocate(token_ids=token_ids, device=Device.GPU)
88 -
89 - # Expect all sequences to share allocations, except for their last block
90 - # (which may be mutable).
91 - assert allocator.get_num_free_blocks(
92 - device=Device.GPU) == num_gpu_blocks - (
93 - num_immutable_blocks_per_alloc + num_mutable_blocks_per_alloc *
94 - (alloc_i))
95 -
96 -
97 -@pytest.mark.parametrize("block_size", [16])
98 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
99 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
100 -@pytest.mark.parametrize("device", ["cpu", "gpu"])
101 -def test_allocate_free(block_size: int, sequence_len: int, allocator_type: str,
102 - device: str):
103 - """Test the allocation and freeing of blocks using different allocators and
104 - devices.
105 -
106 - This test creates a CpuGpuBlockAllocator with the specified block size,
107 - number of blocks, allocator type, and device. It then allocates a BlockTable
108 - multiple times with the same sequence and verifies that the number of free
109 - blocks remains consistent after each allocation and freeing.
110 - """
111 - device = Device[device.upper()]
112 -
113 - num_device_blocks = 1024
114 - allocator = CpuGpuBlockAllocator.create(
115 - allocator_type=allocator_type,
116 - num_gpu_blocks=num_device_blocks,
117 - num_cpu_blocks=num_device_blocks,
118 - block_size=block_size,
119 - )
120 -
121 - token_ids = list(range(sequence_len))
122 - num_blocks_per_alloc = len(list(chunk_list(token_ids, block_size)))
123 -
124 - block_table = BlockTable(
125 - block_size=block_size,
126 - block_allocator=allocator,
127 - )
128 -
129 - for i in range(5):
130 - block_table.allocate(token_ids=token_ids, device=device)
131 - assert allocator.get_num_free_blocks(
132 - device) == num_device_blocks - num_blocks_per_alloc
133 - assert all(block_id is not None
134 - for block_id in block_table.physical_block_ids)
135 -
136 - block_table.free()
137 - assert allocator.get_num_free_blocks(device) == num_device_blocks
138 -
139 -
140 -@pytest.mark.parametrize("block_size", [1, 8])
141 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
142 -@pytest.mark.parametrize("append_len", [1, 16, 129])
143 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
144 -def test_append_token_ids_allocation(block_size: int, sequence_len: int,
145 - append_len: int, allocator_type: str):
146 - """Test the allocation behavior when appending token IDs to a BlockTable.
147 -
148 - This test creates a CpuGpuBlockAllocator with the specified block size,
149 - number of blocks, and allocator type. It then allocates a BlockTable with an
150 - initial sequence and appends additional token IDs to it. The test verifies
151 - that the number of allocated blocks before and after appending matches the
152 - expected values.
153 - """
154 -
155 - num_gpu_blocks = 1024
156 -
157 - allocator = CpuGpuBlockAllocator.create(
158 - allocator_type=allocator_type,
159 - num_gpu_blocks=num_gpu_blocks,
160 - num_cpu_blocks=1024,
161 - block_size=block_size,
162 - )
163 -
164 - token_ids = list(range(sequence_len))
165 - token_ids_to_append = list(range(append_len))
166 -
167 - block_table = BlockTable(
168 - block_size=block_size,
169 - block_allocator=allocator,
170 - )
171 -
172 - num_expected_blocks_before_append = len(
173 - list(chunk_list(token_ids, block_size)))
174 - num_expected_appended_blocks = len(
175 - list(chunk_list(token_ids + token_ids_to_append,
176 - block_size))) - num_expected_blocks_before_append
177 -
178 - block_table.allocate(token_ids=token_ids, device=Device.GPU)
179 -
180 - assert len(
181 - block_table.physical_block_ids) == num_expected_blocks_before_append
182 - block_table.append_token_ids(token_ids_to_append)
183 - assert len(
184 - block_table.physical_block_ids
185 - ) == num_expected_blocks_before_append + num_expected_appended_blocks
186 -
187 -
188 -@pytest.mark.parametrize("block_size", [1, 8])
189 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
190 -@pytest.mark.parametrize("num_empty_slots", [1, 16, 129])
191 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
192 -def test_ensure_num_empty_slots_allocation(block_size: int, sequence_len: int,
193 - num_empty_slots: int,
194 - allocator_type: str):
195 - """Test the allocation behavior when ensuring a certain number of empty
196 - slots in a BlockTable.
197 -
198 - This test creates a CpuGpuBlockAllocator with the specified block size,
199 - number of blocks, and allocator type. It then allocates a BlockTable with an
200 - initial sequence and ensures a certain number of empty slots. The test
201 - verifies that the number of allocated blocks before and after ensuring empty
202 - slots matches the expected values. It also checks that filling up the empty
203 - slots does not consume additional blocks.
204 - """
205 - num_gpu_blocks = 1024
206 -
207 - allocator = CpuGpuBlockAllocator.create(
208 - allocator_type=allocator_type,
209 - num_gpu_blocks=num_gpu_blocks,
210 - num_cpu_blocks=1024,
211 - block_size=block_size,
212 - )
213 -
214 - token_ids = list(range(sequence_len))
215 -
216 - block_table = BlockTable(
217 - block_size=block_size,
218 - block_allocator=allocator,
219 - )
220 -
221 - num_expected_blocks_before_append = len(
222 - list(chunk_list(token_ids, block_size)))
223 - num_expected_appended_blocks = len(
224 - list(chunk_list(token_ids + [-1] * num_empty_slots,
225 - block_size))) - num_expected_blocks_before_append
226 -
227 - block_table.allocate(token_ids=token_ids, device=Device.GPU)
228 -
229 - # Assert that the empty slots consume the expected number of additional
230 - # blocks.
231 - assert len(
232 - block_table.physical_block_ids) == num_expected_blocks_before_append
233 - block_table.ensure_num_empty_slots(num_empty_slots)
234 - assert len(
235 - block_table.physical_block_ids
236 - ) == num_expected_blocks_before_append + num_expected_appended_blocks
237 -
238 - # Now, ensure no additional blocks consumed as we fill up the empty slots.
239 - num_free_blocks = allocator.get_num_free_blocks(device=Device.GPU)
240 - block_table.append_token_ids(token_ids=list(range(num_empty_slots)))
241 - assert num_free_blocks == allocator.get_num_free_blocks(device=Device.GPU)
242 -
243 -
244 -@pytest.mark.parametrize("block_size", [1, 8])
245 -@pytest.mark.parametrize("sequence_len", [1, 9])
246 -@pytest.mark.parametrize("append_len", [1, 16, 129])
247 -@pytest.mark.parametrize("append_size", [1, 4, 129])
248 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
249 -def test_append_token_ids_correct_content(block_size: int, sequence_len: int,
250 - append_len: int, allocator_type: str,
251 - append_size: int):
252 - """Verify token ids are correctly appended. Appends various amounts of
253 - token ids in various append sizes, and verifies the final sequence is
254 - correct.
255 - """
256 - num_gpu_blocks = 1024
257 -
258 - allocator = CpuGpuBlockAllocator.create(
259 - allocator_type=allocator_type,
260 - num_gpu_blocks=num_gpu_blocks,
261 - num_cpu_blocks=1024,
262 - block_size=block_size,
263 - )
264 -
265 - token_ids = list(range(sequence_len))
266 - token_ids_to_append = list(range(append_len))
267 -
268 - block_table = BlockTable(
269 - block_size=block_size,
270 - block_allocator=allocator,
271 - )
272 - block_table.allocate(token_ids=token_ids, device=Device.GPU)
273 -
274 - appended_so_far: list[int] = []
275 - for append in chunk_list(token_ids_to_append, append_size):
276 - block_table.append_token_ids(append)
277 - appended_so_far.extend(append)
278 -
279 - assert block_table._get_all_token_ids() == token_ids + appended_so_far
280 -
281 - assert block_table._get_all_token_ids() == token_ids + token_ids_to_append
282 -
283 -
284 -@pytest.mark.parametrize("seq_len", [1, 9, 129])
285 -@pytest.mark.parametrize("block_size", [1, 8])
286 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
287 -def test_fork(seq_len: int, block_size: int, allocator_type: str):
288 - """Create a sequence using the specified allocator.
289 - 1. Assert that after forking the sequence, the free block count is the
290 - same.
291 - 2. Assert that the forked sequence has the same physical mappings.
292 - 3. Then free the original sequence; verify that the free block count is
293 - the same.
294 - 4. Finally, free the forked sequence and verify that the free block
295 - count drops to zero.
296 - """
297 - num_gpu_blocks = 1024
298 -
299 - allocator = CpuGpuBlockAllocator.create(
300 - allocator_type=allocator_type,
301 - num_gpu_blocks=num_gpu_blocks,
302 - num_cpu_blocks=0,
303 - block_size=block_size,
304 - )
305 -
306 - token_ids = list(range(seq_len))
307 -
308 - block_table = BlockTable(
309 - block_size=block_size,
310 - block_allocator=allocator,
311 - )
312 -
313 - block_table.allocate(token_ids)
314 -
315 - num_free_blocks_before_fork = allocator.get_num_free_blocks(
316 - device=Device.GPU)
317 -
318 - forked_block_table = block_table.fork()
319 -
320 - # Expect physical_block_ids and token_ids to match.
321 - assert (block_table.physical_block_ids ==
322 - forked_block_table.physical_block_ids)
323 - assert block_table._get_all_token_ids(
324 - ) == forked_block_table._get_all_token_ids()
325 -
326 - # Do not expect any additional allocations.
327 - assert allocator.get_num_free_blocks(
328 - device=Device.GPU) == num_free_blocks_before_fork
329 -
330 - # Free the original blocks. Assert num free blocks does not change, since
331 - # refcount is nonzero.
332 - block_table.free()
333 - assert allocator.get_num_free_blocks(
334 - device=Device.GPU) == num_free_blocks_before_fork
335 -
336 - # Expect the forked block table to be unaffected by the free.
337 - assert all(block_id is not None
338 - for block_id in forked_block_table.physical_block_ids)
339 -
340 - # Free the forked blocks. Assert num free blocks does change, since
341 - # refcount is now zero.
342 - forked_block_table.free()
343 - assert allocator.get_num_free_blocks(device=Device.GPU) == num_gpu_blocks
344 -
345 -
346 -@pytest.mark.parametrize("block_size", [8])
347 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
348 -@pytest.mark.parametrize("append_len", [1, 16, 129])
349 -@pytest.mark.parametrize("appender", ["forked", "original"])
350 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
351 -def test_cow(block_size: int, sequence_len: int, append_len: int,
352 - allocator_type: str, appender: str):
353 - """Fork a sequence; append to the forked sequence; verify there's a CoW.
354 - """
355 - num_gpu_blocks = 1024
356 -
357 - allocator = CpuGpuBlockAllocator.create(
358 - allocator_type=allocator_type,
359 - num_gpu_blocks=num_gpu_blocks,
360 - num_cpu_blocks=0,
361 - block_size=block_size,
362 - )
363 -
364 - token_ids = list(range(sequence_len))
365 - token_ids_to_append = list(range(append_len))
366 -
367 - original_block_table = BlockTable(
368 - block_size=block_size,
369 - block_allocator=allocator,
370 - )
371 -
372 - num_expected_non_cow_blocks = cdiv(sequence_len, block_size)
373 - num_expected_cow_blocks = cdiv(sequence_len + append_len,
374 - block_size) - (sequence_len // block_size)
375 -
376 - original_block_table.allocate(token_ids=token_ids, device=Device.GPU)
377 - original_block_ids = original_block_table.physical_block_ids[:]
378 -
379 - print("original_block_ids = {}".format(original_block_ids))
380 - forked_block_table = original_block_table.fork()
381 -
382 - # Expect no additional allocation (copy on _write_).
383 - assert allocator.get_num_free_blocks(
384 - Device.GPU) == (num_gpu_blocks - num_expected_non_cow_blocks)
385 -
386 - if appender == "forked":
387 - appender_block_table = forked_block_table
388 - static_block_table = original_block_table
389 - elif appender == "original":
390 - appender_block_table = original_block_table
391 - static_block_table = forked_block_table
392 - else:
393 - raise ValueError(f"unknown test config {appender=}")
394 -
395 - # Write tokens.
396 - appender_block_table.append_token_ids(token_ids_to_append)
397 -
398 - # Expect the non-appending block table to have no change.
399 - assert static_block_table.physical_block_ids == original_block_ids
400 - assert appender_block_table.physical_block_ids != original_block_ids
401 -
402 - # Expect the blocks changed during append to have a CoW.
403 - assert allocator.get_num_free_blocks(
404 - Device.GPU) == num_gpu_blocks - (num_expected_non_cow_blocks +
405 - num_expected_cow_blocks)
406 -
407 - cows = allocator.clear_copy_on_writes()
408 - if sequence_len % block_size > 0:
409 - # If the last block in the sequence is not full, then when appending we
410 - # expect a CoW.
411 - assert cows
412 -
413 - cow_block_id = sequence_len // block_size
414 - expected_src = static_block_table.physical_block_ids[cow_block_id]
415 - expected_dst = appender_block_table.physical_block_ids[cow_block_id]
416 -
417 - assert (expected_src, expected_dst) in cows
418 - else:
419 - # Otherwise, there should be no copy-on-write.
420 - assert not cows
421 -
422 - static_block_table.free()
423 - appender_block_table.free()
424 -
425 - # After free, expect all blocks to be freed.
426 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
427 -
428 -
429 -@pytest.mark.parametrize("block_size", [8])
430 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
431 -@pytest.mark.parametrize("append_len", [1, 16, 129])
432 -@pytest.mark.parametrize("lookahead_slots", [1, 16, 129])
433 -@pytest.mark.parametrize("appender", ["forked", "original"])
434 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
435 -def test_cow_lookahead_simple(block_size: int, sequence_len: int,
436 - append_len: int, lookahead_slots: int,
437 - allocator_type: str, appender: str):
438 - """Similar to test_cow, except with lookahead allocation. The assertions are
439 - less rigorous due to the complexity of the property under test.
440 - """
441 - num_gpu_blocks = 1024
442 -
443 - allocator = CpuGpuBlockAllocator.create(
444 - allocator_type=allocator_type,
445 - num_gpu_blocks=num_gpu_blocks,
446 - num_cpu_blocks=0,
447 - block_size=block_size,
448 - )
449 -
450 - token_ids = list(range(sequence_len))
451 - token_ids_to_append = list(range(append_len))
452 -
453 - original_block_table = BlockTable(
454 - block_size=block_size,
455 - block_allocator=allocator,
456 - )
457 -
458 - original_block_table.allocate(token_ids=token_ids, device=Device.GPU)
459 -
460 - # Allocate lookahead slots.
461 - original_block_table.ensure_num_empty_slots(lookahead_slots)
462 - original_block_ids = original_block_table.physical_block_ids[:]
463 -
464 - forked_block_table = original_block_table.fork()
465 -
466 - if appender == "forked":
467 - appender_block_table = forked_block_table
468 - static_block_table = original_block_table
469 - elif appender == "original":
470 - appender_block_table = original_block_table
471 - static_block_table = forked_block_table
472 - else:
473 - raise ValueError(f"unknown test config {appender=}")
474 -
475 - # Write tokens.
476 - appender_block_table.append_token_ids(token_ids_to_append)
477 -
478 - # Expect the non-appending block table to have no change.
479 - assert static_block_table.physical_block_ids == original_block_ids
480 - assert appender_block_table.physical_block_ids != original_block_ids
481 -
482 - cows = allocator.clear_copy_on_writes()
483 -
484 - # Always expect copy-on-write
485 - assert cows
486 -
487 - if sequence_len % block_size > 0:
488 - # If the last block in the sequence is not full, then when appending we
489 - # expect a CoW.
490 - assert cows
491 -
492 - cow_block_id = sequence_len // block_size
493 - expected_src = static_block_table.physical_block_ids[cow_block_id]
494 - expected_dst = appender_block_table.physical_block_ids[cow_block_id]
495 -
496 - assert (expected_src, expected_dst) in cows
497 -
498 - static_block_table.free()
499 - appender_block_table.free()
500 -
501 - # After free, expect all blocks to be freed.
502 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
503 -
504 -
505 -@pytest.mark.parametrize("block_size", [1, 8])
506 -@pytest.mark.parametrize("sequence_len", [1, 16, 129])
507 -@pytest.mark.parametrize("num_new_tokens", [1, 16, 129])
508 -@pytest.mark.parametrize("num_lookahead_slots", [1, 7, 8])
509 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
510 -def test_num_blocks_touched_by_append_slots(block_size: int, sequence_len: int,
511 - num_new_tokens: int,
512 - num_lookahead_slots: int,
513 - allocator_type: str):
514 - """Verify correct calculation of get_num_blocks_touched_by_append_slots.
515 -
516 - This is done by using copy-on-write, which requires any modified block to
517 - be copied before write if the refcount > 1. We set the refcount>1 by forking
518 - a sequence, then measure the free blocks before and after an append. If the
519 - number of consumed blocks equals what `get_num_blocks_touched_by_append_
520 - slots` returns, then the calculation is correct.
521 - """
522 -
523 - num_gpu_blocks = 1024
524 -
525 - allocator = CpuGpuBlockAllocator.create(
526 - allocator_type=allocator_type,
527 - num_gpu_blocks=num_gpu_blocks,
528 - num_cpu_blocks=0,
529 - block_size=block_size,
530 - )
531 -
532 - token_ids = list(range(sequence_len))
533 - token_ids_to_append = list(range(num_new_tokens))
534 -
535 - block_table = BlockTable(
536 - block_size=block_size,
537 - block_allocator=allocator,
538 - )
539 -
540 - block_table.allocate(token_ids=token_ids, device=Device.GPU)
541 -
542 - # Add lookahead before fork so both sequences have the same lookahead
543 - # blocks.
544 - block_table.ensure_num_empty_slots(num_empty_slots=num_lookahead_slots)
545 -
546 - # Fork sequence so that every block has refcount > 1.
547 - _ = block_table.fork()
548 -
549 - # Determine how many blocks should be touched.
550 - expected_num_touched_blocks = (
551 - block_table.get_num_blocks_touched_by_append_slots(
552 - token_ids=token_ids_to_append,
553 - num_lookahead_slots=num_lookahead_slots))
554 -
555 - # Measure how many blocks are touched by measuring num_free_blocks before
556 - # and after the append.
557 - #
558 - # We expect append_token_ids to CoW all mutated blocks that have refcount>1.
559 - num_free_blocks_before_append = allocator.get_num_free_blocks(Device.GPU)
560 - block_table.append_token_ids(token_ids_to_append, num_lookahead_slots)
561 - num_consumed_blocks = (num_free_blocks_before_append -
562 - allocator.get_num_free_blocks(Device.GPU))
563 -
564 - # TODO(cade) ensure equality when num_lookahead_slots > 0.
565 - # The reason we have < is because lookahead blocks are not copied eagerly;
566 - # they are copied on first write. This will cause issues for beam search +
567 - # speculative decoding. This is acceptable for now as it is a large effort
568 - # to combine the two. To fix this, we can ensure single sequence ownership
569 - # of lookahead blocks by appending empty slots to each block, which will
570 - # trigger the CoW.
571 - #
572 - # Until then, we can accept that the consumed tokens are <= the expected
573 - # tokens when appending with lookahead.
574 - if num_lookahead_slots > 0:
575 - assert num_consumed_blocks <= expected_num_touched_blocks
576 - else:
577 - assert num_consumed_blocks == expected_num_touched_blocks
tests/core/block/test_common.py deleted
-45
@@ -1,45 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import random
5 -
6 -import pytest
7 -
8 -from vllm.core.block.common import RefCounter
9 -
10 -
11 -@pytest.mark.parametrize("seed", list(range(20)))
12 -@pytest.mark.parametrize("num_incrs", [1, 100])
13 -@pytest.mark.parametrize("num_blocks", [1024])
14 -def test_incr(seed: int, num_incrs: int, num_blocks: int):
15 - random.seed(seed)
16 -
17 - all_block_indices = list(range(num_blocks))
18 - counter = RefCounter(all_block_indices=all_block_indices)
19 -
20 - block_id = random.randint(0, num_blocks - 1)
21 - for i in range(num_incrs):
22 - value = counter.incr(block_id)
23 - assert value == i + 1
24 -
25 -
26 -@pytest.mark.parametrize("seed", list(range(20)))
27 -@pytest.mark.parametrize("num_incrs", [1, 100])
28 -@pytest.mark.parametrize("num_blocks", [1024])
29 -def test_incr_decr(seed: int, num_incrs: int, num_blocks: int):
30 - random.seed(seed)
31 -
32 - all_block_indices = list(range(num_blocks))
33 - counter = RefCounter(all_block_indices=all_block_indices)
34 -
35 - block_id = random.randint(0, num_blocks - 1)
36 - for i in range(num_incrs):
37 - value = counter.incr(block_id)
38 - assert value == i + 1
39 -
40 - for i in range(num_incrs):
41 - value = counter.decr(block_id)
42 - assert value == num_incrs - (i + 1)
43 -
44 - with pytest.raises(AssertionError):
45 - counter.decr(block_id)
tests/core/block/test_cpu_gpu_block_allocator.py deleted
-96
@@ -1,96 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
7 -from vllm.utils import Device, chunk_list
8 -
9 -
10 -@pytest.mark.parametrize("num_cpu_blocks", [0, 512])
11 -@pytest.mark.parametrize("num_gpu_blocks", [1024])
12 -@pytest.mark.parametrize("block_size", [16])
13 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
14 -def test_allocate_mutable_block(num_cpu_blocks: int, num_gpu_blocks: int,
15 - block_size: int, allocator_type: str):
16 - allocator = CpuGpuBlockAllocator.create(
17 - allocator_type=allocator_type,
18 - num_gpu_blocks=num_gpu_blocks,
19 - num_cpu_blocks=num_cpu_blocks,
20 - block_size=block_size,
21 - )
22 -
23 - assert allocator.get_num_free_blocks(Device.CPU) == num_cpu_blocks
24 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
25 -
26 - cpu_blocks = [
27 - allocator.allocate_mutable_block(prev_block=None, device=Device.CPU)
28 - for _ in range(num_cpu_blocks)
29 - ]
30 - assert allocator.get_num_free_blocks(Device.CPU) == 0
31 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
32 -
33 - gpu_blocks = [
34 - allocator.allocate_mutable_block(prev_block=None, device=Device.GPU)
35 - for _ in range(num_gpu_blocks)
36 - ]
37 - assert allocator.get_num_free_blocks(Device.CPU) == 0
38 - assert allocator.get_num_free_blocks(Device.GPU) == 0
39 -
40 - _ = [allocator.free(block) for block in cpu_blocks]
41 - assert allocator.get_num_free_blocks(Device.CPU) == num_cpu_blocks
42 - assert allocator.get_num_free_blocks(Device.GPU) == 0
43 -
44 - _ = [allocator.free(block) for block in gpu_blocks]
45 - assert allocator.get_num_free_blocks(Device.CPU) == num_cpu_blocks
46 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
47 -
48 -
49 -@pytest.mark.parametrize("num_cpu_blocks", [0, 512])
50 -@pytest.mark.parametrize("num_gpu_blocks", [1024])
51 -@pytest.mark.parametrize("block_size", [2])
52 -@pytest.mark.parametrize("allocator_type", ["naive", "prefix_caching"])
53 -def test_allocate_immutable_block(num_cpu_blocks: int, num_gpu_blocks: int,
54 - block_size: int, allocator_type: str):
55 - allocator = CpuGpuBlockAllocator.create(
56 - allocator_type=allocator_type,
57 - num_gpu_blocks=num_gpu_blocks,
58 - num_cpu_blocks=num_cpu_blocks,
59 - block_size=block_size,
60 - )
61 -
62 - unique_token_ids = list(
63 - range((num_cpu_blocks + num_gpu_blocks) * block_size))
64 - gpu_token_ids = list(
65 - chunk_list(unique_token_ids[:num_gpu_blocks * block_size], block_size))
66 - cpu_token_ids = list(
67 - chunk_list(unique_token_ids[num_gpu_blocks * block_size:], block_size))
68 -
69 - assert allocator.get_num_free_blocks(Device.CPU) == num_cpu_blocks
70 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
71 -
72 - cpu_blocks = [
73 - allocator.allocate_immutable_block(prev_block=None,
74 - token_ids=token_ids,
75 - device=Device.CPU)
76 - for token_ids in cpu_token_ids
77 - ]
78 - assert allocator.get_num_free_blocks(Device.CPU) == 0
79 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
80 -
81 - gpu_blocks = [
82 - allocator.allocate_immutable_block(prev_block=None,
83 - token_ids=token_ids,
84 - device=Device.GPU)
85 - for token_ids in gpu_token_ids
86 - ]
87 - assert allocator.get_num_free_blocks(Device.CPU) == 0
88 - assert allocator.get_num_free_blocks(Device.GPU) == 0
89 -
90 - _ = [allocator.free(block) for block in cpu_blocks]
91 - assert allocator.get_num_free_blocks(Device.CPU) == num_cpu_blocks
92 - assert allocator.get_num_free_blocks(Device.GPU) == 0
93 -
94 - _ = [allocator.free(block) for block in gpu_blocks]
95 - assert allocator.get_num_free_blocks(Device.CPU) == num_cpu_blocks
96 - assert allocator.get_num_free_blocks(Device.GPU) == num_gpu_blocks
tests/core/block/test_naive_block.py deleted
-148
@@ -1,148 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Optional
5 -
6 -import pytest
7 -
8 -from vllm.core.block.interfaces import Block, BlockAllocator
9 -from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator
10 -
11 -
12 -class TestNaiveBlockAllocator:
13 -
14 - @staticmethod
15 - def create_allocate_lambda(allocate_type: str,
16 - allocator: NaiveBlockAllocator,
17 - prev_block: Optional[Block],
18 - token_ids: list[int]):
19 - if allocate_type == "immutable":
20 - allocate_block = lambda: allocator.allocate_immutable_block(
21 - prev_block=prev_block, token_ids=token_ids)
22 - elif allocate_type == "mutable":
23 - allocate_block = lambda: allocator.allocate_mutable_block(
24 - prev_block=prev_block)
25 - else:
26 - raise ValueError()
27 -
28 - return allocate_block
29 -
30 - @staticmethod
31 - @pytest.mark.parametrize("allocate_type", ["immutable", "mutable"])
32 - @pytest.mark.parametrize("num_blocks", [1, 1024])
33 - @pytest.mark.parametrize("block_size", [1, 16])
34 - def test_allocate_ooms(allocate_type: str, num_blocks: int,
35 - block_size: int):
36 - allocator = NaiveBlockAllocator(create_block=NaiveBlock,
37 - num_blocks=num_blocks,
38 - block_size=block_size)
39 - allocate_block = TestNaiveBlockAllocator.create_allocate_lambda(
40 - allocate_type,
41 - allocator,
42 - prev_block=None,
43 - token_ids=list(range(block_size)))
44 -
45 - [allocate_block() for _ in range(num_blocks)]
46 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
47 - allocate_block()
48 -
49 - @staticmethod
50 - @pytest.mark.parametrize("allocate_type", ["immutable", "mutable"])
51 - @pytest.mark.parametrize("num_blocks", [1, 1024])
52 - @pytest.mark.parametrize("block_size", [1, 16])
53 - def test_free_prevents_oom(allocate_type: str, num_blocks: int,
54 - block_size: int):
55 - allocator = NaiveBlockAllocator(create_block=NaiveBlock,
56 - num_blocks=num_blocks,
57 - block_size=block_size)
58 - allocate_block = TestNaiveBlockAllocator.create_allocate_lambda(
59 - allocate_type,
60 - allocator,
61 - prev_block=None,
62 - token_ids=list(range(block_size)))
63 -
64 - blocks = [allocate_block() for _ in range(num_blocks)]
65 -
66 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
67 - allocate_block()
68 -
69 - block_to_free = blocks.pop()
70 -
71 - for _ in range(100):
72 - block_id = block_to_free.block_id
73 - allocator.free(block_to_free)
74 - assert block_to_free.block_id is None
75 -
76 - new_block = allocate_block()
77 - assert new_block.block_id == block_id
78 -
79 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
80 - allocate_block()
81 -
82 - block_to_free = new_block
83 -
84 - @staticmethod
85 - @pytest.mark.parametrize("allocate_type", ["immutable", "mutable"])
86 - @pytest.mark.parametrize("num_blocks", [1024])
87 - @pytest.mark.parametrize("block_size", [16])
88 - def test_get_num_free_blocks(allocate_type: str, num_blocks: int,
89 - block_size: int):
90 - allocator = NaiveBlockAllocator(create_block=NaiveBlock,
91 - num_blocks=num_blocks,
92 - block_size=block_size)
93 - allocate_block = TestNaiveBlockAllocator.create_allocate_lambda(
94 - allocate_type,
95 - allocator,
96 - prev_block=None,
97 - token_ids=list(range(block_size)))
98 -
99 - assert allocator.get_num_free_blocks() == num_blocks
100 -
101 - blocks = [allocate_block() for _ in range(num_blocks)]
102 -
103 - for i, block in enumerate(blocks):
104 - assert allocator.get_num_free_blocks() == i
105 - allocator.free(block)
106 -
107 - @staticmethod
108 - @pytest.mark.parametrize("num_blocks", [4])
109 - @pytest.mark.parametrize("block_size", [8])
110 - def test_naive_block_get_num_full_blocks_touched(num_blocks, block_size):
111 - """ Verify the allocator can correctly return the number of
112 - full blocks touched.
113 - """
114 - allocator_src = NaiveBlockAllocator(create_block=NaiveBlock,
115 - num_blocks=num_blocks,
116 - block_size=block_size)
117 - allocator_dst = NaiveBlockAllocator(create_block=NaiveBlock,
118 - num_blocks=num_blocks,
119 - block_size=block_size)
120 -
121 - # Create a chain of cacheable blocks in the dst
122 - allocate_block = TestNaiveBlockAllocator.create_allocate_lambda(
123 - "immutable",
124 - allocator_src,
125 - prev_block=None,
126 - token_ids=list(range(block_size)))
127 - src_blocks = [allocate_block() for _ in range(num_blocks - 1)]
128 -
129 - # All blocks are cached
130 - assert allocator_dst.get_num_full_blocks_touched(
131 - src_blocks) == num_blocks - 1
132 -
133 - # Insert one non-full block in the src
134 - allocate_non_full_block = \
135 - TestNaiveBlockAllocator.create_allocate_lambda(
136 - "mutable", allocator_src,
137 - prev_block=src_blocks[-1],token_ids=[]
138 - )
139 - src_blocks.append(allocate_non_full_block())
140 - src_blocks[-1].append_token_ids([0])
141 -
142 - assert allocator_dst.get_num_full_blocks_touched(
143 - src_blocks) == num_blocks - 1
144 - # Fill up the last source block and then invoke
145 - # get_num_blocks_touched
146 - src_blocks[-1].append_token_ids([0] * (block_size - 1))
147 - assert allocator_dst.get_num_full_blocks_touched(
148 - src_blocks) == num_blocks
tests/core/block/test_prefix_caching_block.py deleted
-1035
@@ -1,1035 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import math
5 -import random
6 -from typing import Optional
7 -from unittest.mock import MagicMock
8 -
9 -import pytest
10 -
11 -from tests.core.utils import create_dummy_lora_sequence, create_dummy_sequence
12 -from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
13 -from vllm.core.block.interfaces import Block, BlockAllocator
14 -from vllm.core.block.prefix_caching_block import (ComputedBlocksTracker,
15 - PrefixCachingBlock,
16 - PrefixCachingBlockAllocator)
17 -from vllm.sequence import Logprob
18 -from vllm.utils import Device
19 -
20 -
21 -class TestPrefixCachingBlock:
22 -
23 - @staticmethod
24 - @pytest.mark.parametrize("seed", list(range(10)))
25 - @pytest.mark.parametrize("block_size", [1, 16])
26 - @pytest.mark.parametrize("is_curr_block_full", [True, False])
27 - def test_first_block_has_correct_content_hash(seed: int, block_size: int,
28 - is_curr_block_full: bool):
29 - """Verify a block which is first in the sequence has the correct hash.
30 - """
31 - random.seed(seed)
32 - num_to_fill = block_size if is_curr_block_full else random.randint(
33 - 0, block_size - 1)
34 - token_ids = list(range(num_to_fill))
35 - mock_allocator = MagicMock(spec=PrefixCachingBlockAllocator)
36 -
37 - block_with_prev = PrefixCachingBlock(prev_block=None,
38 - token_ids=token_ids,
39 - block_size=block_size,
40 - allocator=mock_allocator)
41 -
42 - if is_curr_block_full:
43 - # Expect hash since block is full.
44 - assert block_with_prev.content_hash == (
45 - PrefixCachingBlock.hash_block_tokens(
46 - is_first_block=True,
47 - prev_block_hash=None,
48 - cur_block_token_ids=token_ids))
49 - else:
50 - # Do not expect hash since block is not full.
51 - assert block_with_prev.content_hash is None
52 -
53 - @staticmethod
54 - @pytest.mark.parametrize("seed", list(range(10)))
55 - @pytest.mark.parametrize("block_size", [1, 16])
56 - @pytest.mark.parametrize("is_curr_block_full", [True, False])
57 - @pytest.mark.parametrize("prev_block_has_hash", [True, False])
58 - def test_nth_block_has_correct_content_hash(seed: int, block_size: int,
59 - is_curr_block_full: bool,
60 - prev_block_has_hash: bool):
61 - """Verify a block which is not first in the sequence has the correct
62 - hash.
63 - """
64 -
65 - random.seed(seed)
66 -
67 - previous_block = MagicMock(spec=PrefixCachingBlock)
68 - prev_block_hash = random.randint(0, 1000)
69 - previous_block.content_hash = (prev_block_hash if prev_block_has_hash
70 - else hash('None'))
71 -
72 - num_to_fill = block_size if is_curr_block_full else random.randint(
73 - 0, block_size - 1)
74 - token_ids = list(range(num_to_fill))
75 - mock_allocator = MagicMock(spec=PrefixCachingBlockAllocator)
76 -
77 - block_with_prev = PrefixCachingBlock(
78 - prev_block=previous_block,
79 - token_ids=token_ids,
80 - block_size=block_size,
81 - allocator=mock_allocator,
82 - )
83 -
84 - if is_curr_block_full and prev_block_has_hash:
85 - # Expect hash since block is full and previous block has hash.
86 - assert (block_with_prev.content_hash ==
87 - PrefixCachingBlock.hash_block_tokens(
88 - is_first_block=False,
89 - prev_block_hash=prev_block_hash,
90 - cur_block_token_ids=token_ids))
91 - else:
92 - # Do not expect hash since block is not full or the previous block
93 - # does not have a hash.
94 - assert block_with_prev.content_hash is None
95 -
96 - @staticmethod
97 - @pytest.mark.parametrize("block_size", [1, 2, 16])
98 - @pytest.mark.parametrize("num_tokens", list(range(3)))
99 - @pytest.mark.parametrize("num_empty_trailing_blocks", [0, 1, 10])
100 - def test_blocks_have_correct_hash_in_chain(block_size: int,
101 - num_tokens: int,
102 - num_empty_trailing_blocks: int):
103 - """Create two chains of logical blocks with the same contents.
104 - Assert the hashes are equal.
105 - """
106 - random.seed(0)
107 -
108 - token_ids = [random.randint(0, 50_000) for _ in range(num_tokens)]
109 -
110 - first_chain, second_chain = (TestPrefixCachingBlock.create_chain(
111 - block_size=block_size,
112 - token_ids=token_ids,
113 - num_empty_trailing_blocks=num_empty_trailing_blocks)
114 - for _ in range(2))
115 -
116 - for first_chain_block, second_chain_block in zip(
117 - first_chain, second_chain):
118 - assert (first_chain_block.content_hash ==
119 - second_chain_block.content_hash)
120 -
121 - if not first_chain or not second_chain:
122 - assert first_chain == second_chain
123 - assert num_tokens == 0
124 -
125 - @staticmethod
126 - def create_chain(block_size: int,
127 - token_ids: list[int],
128 - num_empty_trailing_blocks=0) -> list[PrefixCachingBlock]:
129 - """Helper method which creates a chain of blocks.
130 - """
131 - blocks: list[PrefixCachingBlock] = []
132 - num_blocks = math.ceil(
133 - len(token_ids) / block_size) + num_empty_trailing_blocks
134 -
135 - if num_blocks == 0:
136 - return []
137 -
138 - allocator = MagicMock(spec=PrefixCachingBlockAllocator)
139 -
140 - prev_block = None
141 - for block_number in range(0, num_blocks):
142 - prev_block = PrefixCachingBlock(
143 - prev_block=prev_block,
144 - token_ids=[],
145 - block_size=block_size,
146 - allocator=allocator,
147 - )
148 -
149 - tokens_to_append = token_ids[block_number *
150 - block_size:(block_number + 1) *
151 - block_size]
152 - if tokens_to_append:
153 - prev_block.append_token_ids(tokens_to_append)
154 -
155 - blocks.append(prev_block)
156 -
157 - return blocks
158 -
159 -
160 -class TestPrefixCachingBlockAllocator:
161 -
162 - @staticmethod
163 - def create_allocate_lambda(allocate_type: str, allocator: BlockAllocator,
164 - prev_block: Optional[Block],
165 - token_ids: list[int]):
166 - if allocate_type == "immutable":
167 - allocate_block = lambda: allocator.allocate_immutable_block(
168 - prev_block=prev_block, token_ids=token_ids)
169 - elif allocate_type == "mutable":
170 - allocate_block = lambda: allocator.allocate_mutable_block(
171 - prev_block=prev_block)
172 - else:
173 - raise ValueError()
174 -
175 - return allocate_block
176 -
177 - @staticmethod
178 - @pytest.mark.parametrize("num_blocks", [1, 1024])
179 - @pytest.mark.parametrize("block_size", [1, 16])
180 - def test_allocate_mutable_ooms(num_blocks: int, block_size: int):
181 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
182 - block_size=block_size)
183 - allocate_block = TestPrefixCachingBlockAllocator.create_allocate_lambda(
184 - allocate_type="mutable",
185 - allocator=allocator,
186 - prev_block=None,
187 - token_ids=list(range(block_size)),
188 - )
189 -
190 - [allocate_block() for _ in range(num_blocks)]
191 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
192 - allocate_block()
193 -
194 - @staticmethod
195 - @pytest.mark.parametrize("num_blocks", [1, 1024])
196 - @pytest.mark.parametrize("block_size", [1, 16])
197 - def test_allocate_immutable_does_not_oom_single_hash(
198 - num_blocks: int, block_size: int):
199 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
200 - block_size=block_size)
201 - allocate_block = TestPrefixCachingBlockAllocator.create_allocate_lambda(
202 - allocate_type="immutable",
203 - allocator=allocator,
204 - prev_block=None,
205 - token_ids=list(range(block_size)),
206 - )
207 -
208 - blocks = [allocate_block() for _ in range(num_blocks)]
209 -
210 - # Expect no OOM. If these were mutable blocks, this would OOM.
211 - non_oom_block = allocate_block()
212 -
213 - # Expect all blocks to have same physical block index.
214 - for block in blocks:
215 - assert (block.block_id == non_oom_block.block_id)
216 -
217 - @staticmethod
218 - @pytest.mark.parametrize("num_blocks", [1, 1024])
219 - @pytest.mark.parametrize("block_size", [1, 16])
220 - def test_allocate_immutable_ooms_many_hash(num_blocks: int,
221 - block_size: int):
222 - """Consume all blocks using many different hashes/block content.
223 -
224 - Do this by creating a sequence that is very long.
225 - Expect next block to OOM.
226 - """
227 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
228 - block_size=block_size)
229 -
230 - # Create token ids that will exhaust all blocks.
231 - token_ids = list(range(num_blocks * block_size))
232 -
233 - chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
234 - block_size=block_size,
235 - token_ids=token_ids,
236 - allocator=allocator,
237 - )
238 -
239 - # Expect allocation with unseen hash to fail.
240 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
241 - allocator.allocate_immutable_block(prev_block=chain[-1],
242 - token_ids=list(
243 - range(block_size)))
244 -
245 - # Expect mutable allocation to fail.
246 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
247 - allocator.allocate_mutable_block(prev_block=chain[-1])
248 -
249 - # Expect allocation of exact same chain to pass.
250 - second_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
251 - block_size=block_size,
252 - token_ids=token_ids,
253 - allocator=allocator,
254 - )
255 -
256 - # Expect physical block indices to be the same in both chains.
257 - assert chain and second_chain
258 - for first_chain_block, second_chain_block in zip(chain, second_chain):
259 - assert (first_chain_block.block_id == second_chain_block.block_id)
260 -
261 - @staticmethod
262 - @pytest.mark.parametrize("num_blocks", [1, 1024])
263 - @pytest.mark.parametrize("block_size", [1, 16])
264 - def test_free_prevents_oom(num_blocks: int, block_size: int):
265 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
266 - block_size=block_size)
267 -
268 - # Create token ids that will exhaust all blocks.
269 - token_ids = list(range(num_blocks * block_size))
270 -
271 - chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
272 - block_size=block_size,
273 - token_ids=token_ids,
274 - allocator=allocator,
275 - )
276 -
277 - # Expect mutable allocation to fail.
278 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
279 - allocator.allocate_mutable_block(prev_block=None)
280 -
281 - block_to_free = chain[-1]
282 -
283 - # Expect free/allocate loop to succeed many times.
284 - for i in range(100):
285 - block_id = block_to_free.block_id
286 - allocator.free(block_to_free)
287 - assert block_to_free.block_id is None, i
288 -
289 - new_block = allocator.allocate_mutable_block(prev_block=None)
290 - assert new_block.block_id == block_id, i
291 -
292 - with pytest.raises(BlockAllocator.NoFreeBlocksError):
293 - allocator.allocate_mutable_block(prev_block=None)
294 -
295 - block_to_free = new_block
296 -
297 - @staticmethod
298 - @pytest.mark.parametrize("num_blocks", [1024])
299 - @pytest.mark.parametrize("block_size", [16])
300 - @pytest.mark.parametrize("seed", list(range(20)))
301 - def test_get_num_free_blocks(num_blocks: int, block_size: int, seed: int):
302 - random.seed(seed)
303 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
304 - block_size=block_size)
305 - num_blocks_to_consume = random.randint(1, num_blocks - 1)
306 -
307 - # Create token ids that will exhaust all blocks.
308 - token_ids = list(range(num_blocks_to_consume * block_size))
309 -
310 - chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
311 - block_size=block_size,
312 - token_ids=token_ids,
313 - allocator=allocator,
314 - )
315 -
316 - # Free each block in chain, assert num free blocks includes new free
317 - # block.
318 - for i, block in enumerate(chain):
319 - assert allocator.get_num_free_blocks() == (num_blocks -
320 - num_blocks_to_consume +
321 - i)
322 - allocator.free(block)
323 -
324 - @staticmethod
325 - @pytest.mark.parametrize("num_blocks", [4])
326 - @pytest.mark.parametrize("block_size", [8])
327 - def test_prefix_caching_block_get_num_full_blocks_touched(
328 - num_blocks, block_size):
329 - """ Verify the allocator can correctly return the number of
330 - blocks touched, when there are cached prefixes.
331 - """
332 - allocator_src = PrefixCachingBlockAllocator(num_blocks=num_blocks,
333 - block_size=block_size)
334 - allocator_dst = PrefixCachingBlockAllocator(num_blocks=num_blocks,
335 - block_size=block_size)
336 -
337 - # Create token ids that will exhaust all blocks except the last
338 - token_ids = list(range((num_blocks - 1) * block_size))
339 -
340 - # Create a chain of cacheable blocks in the dst
341 - cached_blocks = TestPrefixCachingBlockAllocator.create_immutable_chain(
342 - block_size=block_size,
343 - token_ids=token_ids,
344 - allocator=allocator_dst,
345 - )
346 -
347 - # Create a chain of the same blocks in the src
348 - blocks_to_swap_in = \
349 - TestPrefixCachingBlockAllocator.create_immutable_chain(
350 - block_size=block_size,
351 - token_ids=token_ids,
352 - allocator=allocator_src,
353 - )
354 - # All blocks are cached
355 - assert allocator_dst.get_num_full_blocks_touched(
356 - blocks_to_swap_in) == 0
357 -
358 - # Free the first block in the dst
359 - allocator_dst.free(cached_blocks[0])
360 -
361 - # Now the first block becomes dangling, the swapped blocks need
362 - # to reclaim the first block in the dst
363 - assert allocator_dst.get_num_full_blocks_touched(
364 - blocks_to_swap_in) == 1
365 -
366 - # Insert one non-full block in the src
367 - non_full_block = allocator_src.allocate_mutable_block(
368 - blocks_to_swap_in[-1])
369 - non_full_block.append_token_ids([0])
370 - blocks_to_swap_in.append(non_full_block)
371 - assert allocator_dst.get_num_full_blocks_touched(
372 - blocks_to_swap_in) == 1
373 - # Fill up the last mutable block and invoke get_num_blocks_touched.
374 - # Note: The last block is not cached so it will be touched.
375 - non_full_block.append_token_ids([0] * (block_size - 1))
376 - assert allocator_dst.get_num_full_blocks_touched(
377 - blocks_to_swap_in) == 2
378 -
379 - @staticmethod
380 - @pytest.mark.parametrize("num_blocks", [1024])
381 - @pytest.mark.parametrize("block_size", [16])
382 - @pytest.mark.parametrize("seed", list(range(20)))
383 - def test_get_num_free_blocks_shared(num_blocks: int, block_size: int,
384 - seed: int):
385 - """Verify sharing occurs by allocating two sequences that share prefixes
386 - and incrementally freeing blocks.
387 - """
388 - random.seed(seed)
389 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
390 - block_size=block_size)
391 - num_blocks_to_consume = random.randint(1, num_blocks - 1)
392 -
393 - # Create token ids that will exhaust all blocks.
394 - token_ids = list(range(num_blocks_to_consume * block_size))
395 -
396 - first_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
397 - block_size=block_size,
398 - token_ids=token_ids,
399 - allocator=allocator,
400 - )
401 - second_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
402 - block_size=block_size,
403 - token_ids=token_ids,
404 - allocator=allocator,
405 - )
406 -
407 - # Free each block in the first chain. Since all blocks are shared, the
408 - # free count should stay constant.
409 - for i, block in enumerate(first_chain):
410 - assert allocator.get_num_free_blocks() == (num_blocks -
411 - num_blocks_to_consume)
412 - allocator.free(block)
413 -
414 - # Free each block in the second chain. Since the refcount is now zero,
415 - # the free count should increment with each free.
416 - for i, block in enumerate(second_chain):
417 - assert allocator.get_num_free_blocks() == (num_blocks -
418 - num_blocks_to_consume +
419 - i)
420 - allocator.free(block)
421 -
422 - @staticmethod
423 - @pytest.mark.parametrize("num_blocks", [1024])
424 - @pytest.mark.parametrize("block_size", [16])
425 - @pytest.mark.parametrize("seed", list(range(20)))
426 - def test_get_common_computed_block_ids(num_blocks: int, block_size: int,
427 - seed: int):
428 - """Verify get_common_computed_block_ids could get correct result
429 - by create two immutable chain sharing prefix at specified pos,
430 - and compare whether we also could get right result
431 - from get_common_computed_block_ids.
432 - """
433 - random.seed(seed)
434 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks * 2,
435 - block_size=block_size)
436 - num_blocks_to_consume = random.randint(1, num_blocks - 1)
437 -
438 - # Create token ids that will exhaust all blocks.
439 - token_ids = list(range(num_blocks_to_consume * block_size))
440 -
441 - first_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
442 - block_size=block_size,
443 - token_ids=token_ids,
444 - allocator=allocator,
445 - )
446 -
447 - # After zero_point, second_chain's token_ids would be set -1, which
448 - # make it different from here comparing with first_chain
449 - zero_point = random.randint(1, len(token_ids) - 1)
450 - zero_point_blocks = zero_point // block_size
451 - token_ids[zero_point:] = [-1] * (len(token_ids) - zero_point)
452 -
453 - second_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
454 - block_size=block_size,
455 - token_ids=token_ids,
456 - allocator=allocator,
457 - )
458 -
459 - first_computed_ids = [
460 - first_chain[i].block_id for i in range(num_blocks_to_consume)
461 - ]
462 - second_computed_ids = [
463 - second_chain[i].block_id for i in range(num_blocks_to_consume)
464 - ]
465 - res = allocator.get_common_computed_block_ids(
466 - [first_computed_ids, second_computed_ids])
467 -
468 - assert (len(res) == zero_point_blocks)
469 -
470 - # Test case that assume those prompted block after first immutable would
471 - # be freed into hashless allocator, while first immutable block get ref
472 - # increased.
473 - @staticmethod
474 - @pytest.mark.parametrize("num_blocks", [3])
475 - @pytest.mark.parametrize("block_size", [16])
476 - @pytest.mark.parametrize("seed", list(range(10)))
477 - def test_alloc_promotion(num_blocks: int, block_size: int, seed: int):
478 - random.seed(seed)
479 -
480 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
481 - block_size=block_size)
482 - token_ids = list(range(block_size))
483 -
484 - block = allocator.allocate_immutable_block(prev_block=None,
485 - token_ids=token_ids)
486 -
487 - assert allocator._refcounter.get(block.block_id) == 1
488 - m = allocator.allocate_mutable_block(prev_block=None)
489 -
490 - block_id = m.block_id
491 - for i in range(block_size):
492 - m.append_token_ids([i])
493 -
494 - # After block get promoted to immutable from mutable, if there is
495 - # already same content hash block, then it shall be released into
496 - # hashless_allocator
497 - # And first immutable block's ref get increased by 1
498 - assert m.block_id == block.block_id
499 - assert block_id in allocator._hashless_allocator._free_block_indices
500 - assert allocator._refcounter.get(block.block_id) == 2
501 -
502 - # Test case when eviction and allocation are mixed,
503 - # make sure they work as expected
504 - @staticmethod
505 - @pytest.mark.parametrize("num_blocks", [3])
506 - @pytest.mark.parametrize("block_size", [16])
507 - @pytest.mark.parametrize("seed", list(range(10)))
508 - def test_eviction_alloc_mixed(num_blocks: int, block_size: int, seed: int):
509 - random.seed(seed)
510 -
511 - all_blocks_list = [i for i in range(num_blocks)]
512 - zero_ref = {i: 0 for i in range(num_blocks)}
513 - one_ref = {i: 1 for i in range(num_blocks)}
514 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
515 - block_size=block_size)
516 - token_ids = list(range(num_blocks * block_size))
517 -
518 - # Verify initial/pre-alloc state
519 -
520 - # Ensure all blocks are free inside hashless allocator
521 - assert list(allocator._hashless_allocator._free_block_indices
522 - ) == all_blocks_list
523 - # Ensure no tracked blocks
524 - assert len(allocator._block_tracker.keys()) == num_blocks
525 - for block_id in range(num_blocks):
526 - assert not allocator._block_tracker[block_id].active
527 - # Ensure no cached blocks
528 - assert len(allocator._cached_blocks.values()) == 0
529 - # Ensure no evicted blocks
530 - assert len(allocator.evictor.free_table.keys()) == 0
531 - # Ensure 0s ref counts for all blocks
532 - assert allocator._refcounter._refcounts == zero_ref
533 -
534 - # Allocate immutable chains with only one block residuled in
535 - new_block = []
536 - for i in range(num_blocks):
537 - block = allocator.allocate_immutable_block(
538 - prev_block=None,
539 - token_ids=token_ids[block_size * i:block_size * (i + 1)])
540 - new_block.append(block)
541 -
542 - # Verify post-alloc state
543 -
544 - # Ensure no blocks are free inside hashless allocator
545 - assert (len(allocator._hashless_allocator._free_block_indices) == 0)
546 - # Ensure all blocks are tracked
547 - assert len(allocator._block_tracker.keys()) == num_blocks
548 - for block_id in range(num_blocks):
549 - assert allocator._block_tracker[block_id].active
550 - # Ensure all blocks are cached (all promoted)
551 - assert len(allocator._cached_blocks.values()) == num_blocks
552 - # Ensure no evicted blocks
553 - assert len(allocator.evictor.free_table.keys()) == 0
554 - # Ensure 1s ref counts for all blocks
555 - assert allocator._refcounter._refcounts == one_ref
556 -
557 - # Free all blocks, and now all blocks shall be in the evictor
558 - # there shall be no tracking data left in _block_tracker
559 - # all blocks shall be tracked in _cached_blocks
560 - # all blocks' ref shall be zero
561 - for block in new_block:
562 - allocator.free(block)
563 -
564 - # Verify post-free state
565 -
566 - # Ensure no tracked blocks
567 - assert len(allocator._block_tracker.keys()) == num_blocks
568 - for block_id in range(num_blocks):
569 - assert not allocator._block_tracker[block_id].active
570 - # Ensure no blocks in hashless allocator (all promoted)
571 - assert len(allocator._hashless_allocator._free_block_indices) == 0
572 - # Ensure all blocks are cached
573 - assert list(allocator._cached_blocks.values()) == all_blocks_list
574 - # Ensure all blocks are inside the evictor
575 - assert list(allocator.evictor.free_table.keys()) == all_blocks_list
576 - # Ensure 0s refcounts
577 - assert allocator._refcounter._refcounts == zero_ref
578 -
579 - # Allocate a mutable block, and the first block shall be evicted
580 - # and set its content hash into None, ref to 1
581 - mutable = allocator.allocate_mutable_block(prev_block=None)
582 -
583 - assert mutable.block_id == 0
584 - assert mutable.content_hash is None
585 - assert allocator._block_tracker[0].active
586 - assert allocator._refcounter.get(0) == 1
587 - assert 0 not in allocator._cached_blocks
588 - assert 0 not in allocator.evictor
589 -
590 - # Since this mutable block has no hash yet, it shall be released into
591 - # hashless allocator
592 - allocator.free(mutable)
593 -
594 - assert not allocator._block_tracker[0].active
595 - assert allocator._refcounter._refcounts == zero_ref
596 - assert 0 not in allocator._cached_blocks
597 - assert 0 not in allocator.evictor
598 - assert 0 in allocator._hashless_allocator._free_block_indices
599 -
600 - # When allocate immutable with first block_size tokens, we
601 - # shall get free block from hashless allocator, thus no block left
602 - # in hashless
603 - block = allocator.allocate_immutable_block(
604 - prev_block=None, token_ids=token_ids[:block_size])
605 -
606 - assert block.block_id == 0
607 - assert len(allocator._hashless_allocator._free_block_indices) == 0
608 - assert allocator._block_tracker[0].active
609 - assert 0 in allocator._cached_blocks.values()
610 - assert allocator._refcounter.get(0) == 1
611 - assert 0 not in allocator.evictor
612 -
613 - # allocate mutable block again, it shall be popped from evictor
614 - mutable = allocator.allocate_mutable_block(prev_block=None)
615 - assert len(allocator._hashless_allocator._free_block_indices) == 0
616 - assert mutable.block_id not in allocator.evictor.free_table
617 - assert allocator._refcounter.get(mutable.block_id) == 1
618 -
619 - # Test case where two last accessed times are equal
620 - @staticmethod
621 - @pytest.mark.parametrize("num_blocks", [1024])
622 - @pytest.mark.parametrize("block_size", [16])
623 - @pytest.mark.parametrize("seed", list(range(20)))
624 - def test_eviction_order(num_blocks: int, block_size: int, seed: int):
625 - """This test case simulate the two chain created and free in order,
626 - and together they would exhaust the initial freed blocks.
627 -
628 - So the next block created after those two chain shall use the block
629 - from the first chain as that block has long access time.
630 - While first chain has two blocks, it shall pick up the last one, as
631 - it has larger token number.
632 - """
633 -
634 - random.seed(seed)
635 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
636 - block_size=block_size)
637 - num_blocks_to_consume = num_blocks + 1
638 -
639 - token_ids = list(range(num_blocks_to_consume * block_size))
640 -
641 - num_blocks_in_first_chain = 2
642 - num_tokens_in_first_chain = block_size * num_blocks_in_first_chain
643 - # First chain takes the first block
644 - first_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
645 - block_size=block_size,
646 - token_ids=token_ids[:num_tokens_in_first_chain],
647 - allocator=allocator,
648 - )
649 - # There should only be one block allocated at this point
650 - assert allocator.get_num_free_blocks() == (num_blocks -
651 - num_blocks_in_first_chain)
652 -
653 - # Set the last accessed time of the first block to 1
654 - blocks_ids = [block.block_id for block in first_chain]
655 - allocator.mark_blocks_as_accessed(blocks_ids, 1)
656 -
657 - # Second chain takes the rest of the blocks
658 - second_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
659 - block_size=block_size,
660 - token_ids=token_ids[num_tokens_in_first_chain:-block_size],
661 - allocator=allocator,
662 - )
663 -
664 - # There shouldn't be any blocks left at this point
665 - assert allocator.get_num_free_blocks() == (0)
666 -
667 - assert len(first_chain) == num_blocks_in_first_chain
668 - last_block_id = first_chain[-1].block_id
669 - # Free each block in the first chain.
670 - for i, block in enumerate(first_chain):
671 - allocator.free(block)
672 -
673 - # Set the last accessed time on all of the blocks in the second chain
674 - # to 2
675 - blocks_ids = [block.block_id for block in second_chain]
676 - allocator.mark_blocks_as_accessed(blocks_ids, 2)
677 -
678 - # Free each block in the second chain.
679 - for i, block in enumerate(second_chain):
680 - allocator.free(block)
681 -
682 - # Allocate a new block and check that it's the least recently used block
683 - # from the first chain.
684 - new_block = TestPrefixCachingBlockAllocator.create_immutable_chain(
685 - block_size=block_size,
686 - token_ids=token_ids[-block_size:],
687 - allocator=allocator,
688 - )
689 -
690 - assert new_block[0].block_id == last_block_id
691 -
692 - # Test case for cache mertics
693 - @staticmethod
694 - def test_metric():
695 - block_size = 16
696 - allocator = PrefixCachingBlockAllocator(num_blocks=4,
697 - block_size=block_size)
698 - # Test when no query (0/0)
699 - assert allocator.get_prefix_cache_hit_rate() == 0.0
700 -
701 - token_ids = list(range(block_size))
702 - allocator.allocate_immutable_block(prev_block=None,
703 - token_ids=token_ids)
704 - # Test 0/1 hit rate
705 - assert allocator.get_prefix_cache_hit_rate() == 0.0
706 -
707 - allocator.allocate_immutable_block(prev_block=None,
708 - token_ids=token_ids)
709 - # Test 1/2 hit rate
710 - assert allocator.get_prefix_cache_hit_rate() == 0.5
711 -
712 - # Test more than one block
713 - for _ in range(2, 1005):
714 - allocator.allocate_immutable_block(prev_block=None,
715 - token_ids=token_ids)
716 - assert allocator.get_prefix_cache_hit_rate() > 0.99
717 -
718 - # Test case for marking cache hit blocks as computed right after
719 - # a batch of prefill sequences are scheduled.
720 - @staticmethod
721 - def test_touch_block():
722 - block_size = 16
723 - common_blocks = 4
724 - allocator = PrefixCachingBlockAllocator(num_blocks=8,
725 - block_size=block_size)
726 -
727 - common_token_ids = list(range(block_size * common_blocks))
728 -
729 - # Mimic the behavior of allocating the same block chain
730 - # (i.e., common prefix) for a batch of 3 different prefill sequences.
731 - for _ in range(3):
732 - blocks = TestPrefixCachingBlockAllocator.create_immutable_chain(
733 - block_size=block_size,
734 - token_ids=common_token_ids,
735 - allocator=allocator,
736 - )
737 - block_hashes = [block.content_hash for block in blocks]
738 - # The allocated blocks should be marked as touched
739 - # but not computed.
740 - computed_block_ids = allocator.find_cached_blocks_prefix(
741 - block_hashes)
742 - assert len(computed_block_ids) == 0
743 -
744 - allocator.mark_blocks_as_computed([])
745 - computed_block_ids = allocator.find_cached_blocks_prefix(
746 - block_hashes=block_hashes)
747 - assert len(computed_block_ids) == common_blocks
748 -
749 - @staticmethod
750 - def test_find_cached_blocks_prefix():
751 - """
752 - This test verifies the behavior of find_cached_blocks_prefix.
753 - """
754 - block_size = 4
755 - num_blocks = 8
756 - total_test_blocks = 12
757 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
758 - block_size=block_size)
759 -
760 - token_ids = list(range(total_test_blocks * block_size))
761 - block_tokens_seq1 = token_ids[:num_blocks * block_size]
762 - blocks_seq1 = TestPrefixCachingBlockAllocator.create_immutable_chain(
763 - block_size=block_size,
764 - token_ids=block_tokens_seq1,
765 - allocator=allocator,
766 - )
767 - block_hashes_seq1 = [block.content_hash for block in blocks_seq1]
768 - allocator.mark_blocks_as_computed([])
769 -
770 - # All blocks should be cached.
771 - cached_blocks_seq1 = allocator.find_cached_blocks_prefix(
772 - block_hashes=block_hashes_seq1)
773 - assert len(cached_blocks_seq1) == num_blocks
774 -
775 - # Free the first sequence.
776 - for block in blocks_seq1:
777 - allocator.free(block)
778 -
779 - # All blocks should be still be cached if not required to be allocated.
780 - cached_blocks = allocator.find_cached_blocks_prefix(
781 - block_hashes=block_hashes_seq1)
782 - assert len(cached_blocks) == num_blocks
783 -
784 - block_tokens_seq2 = token_ids[num_blocks * block_size:]
785 - blocks_seq2 = TestPrefixCachingBlockAllocator.create_immutable_chain(
786 - block_size=block_size,
787 - token_ids=block_tokens_seq2,
788 - allocator=allocator,
789 - )
790 - block_hashes_seq2 = [block.content_hash for block in blocks_seq2]
791 - allocator.mark_blocks_as_computed([])
792 - cached_blocks = allocator.find_cached_blocks_prefix(
793 - block_hashes=block_hashes_seq2)
794 - assert len(cached_blocks) == len(blocks_seq2)
795 -
796 - # Half of the blocks from seq1 should still be cached.
797 - num_evicted_blocks = len(blocks_seq2)
798 - cached_blocks = allocator.find_cached_blocks_prefix(
799 - block_hashes=block_hashes_seq1)
800 - assert len(cached_blocks) == len(blocks_seq1) - num_evicted_blocks
801 -
802 - # Test reset prefix cache
803 - @staticmethod
804 - @pytest.mark.parametrize("num_blocks", [10])
805 - @pytest.mark.parametrize("block_size", [16])
806 - def test_reset_prefix_cache(num_blocks: int, block_size: int):
807 - """This test case simulates the case of resetting the prefix cache."""
808 -
809 - allocator = PrefixCachingBlockAllocator(num_blocks=num_blocks,
810 - block_size=block_size)
811 - token_ids = list(range(3 * block_size))
812 -
813 - first_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
814 - block_size=block_size,
815 - token_ids=token_ids,
816 - allocator=allocator,
817 - )
818 - second_chain = TestPrefixCachingBlockAllocator.create_immutable_chain(
819 - block_size=block_size,
820 - token_ids=token_ids,
821 - allocator=allocator,
822 - )
823 -
824 - # Free each block in the first chain.
825 - for block in first_chain:
826 - allocator.free(block)
827 -
828 - # Failed to reset prefix cache because some blocks are not freed yet.
829 - assert not allocator.reset_prefix_cache()
830 - assert allocator.get_prefix_cache_hit_rate() > 0.0
831 -
832 - # Free each block in the second chain.
833 - for block in second_chain:
834 - allocator.free(block)
835 -
836 - # Reset prefix cache.
837 - assert allocator.reset_prefix_cache()
838 - assert allocator.get_prefix_cache_hit_rate() == 0.0
839 -
840 - @staticmethod
841 - def create_immutable_chain(
842 - block_size: int,
843 - token_ids: list[int],
844 - allocator: PrefixCachingBlockAllocator,
845 - extra_hash: Optional[int] = None,
846 - ) -> list[PrefixCachingBlock]:
847 - """Helper method which creates a chain of blocks.
848 - """
849 - blocks: list[Block] = []
850 - num_blocks = math.ceil(len(token_ids) / block_size)
851 -
852 - if num_blocks == 0:
853 - return []
854 -
855 - prev_block = None
856 - for block_number in range(0, num_blocks):
857 - block_token_ids = token_ids[block_number *
858 - block_size:(block_number + 1) *
859 - block_size]
860 - prev_block = allocator.allocate_immutable_block(
861 - prev_block=prev_block,
862 - token_ids=block_token_ids,
863 - extra_hash=extra_hash)
864 - blocks.append(prev_block)
865 -
866 - return blocks
867 -
868 -
869 -class TestComputedBlocksTracker:
870 -
871 - @staticmethod
872 - def _get_mock_allocator():
873 - return MagicMock(spec=PrefixCachingBlockAllocator)
874 -
875 - @staticmethod
876 - def test_get_num_cached_tokens():
877 - """
878 - Test it correctly computes the number of cached tokens for a given
879 - sequence:
880 -
881 - - The cache token count is derived from the number of cached blocks.
882 - - The cache token count is updated when the allocator is updated.
883 - - When a sequence is removed, the cache token count should be updated
884 - accordingly.
885 -
886 - # TODO(rickyx): This behaviour for prefill sequence is a hack until
887 - we fix the computed blocks tracking.
888 - - The cache token count for prefill sequence doesn't change while
889 - the sequence is in continuous prefill (chunked prefill).
890 - """
891 - block_size = 4
892 - mock_allocator = TestComputedBlocksTracker._get_mock_allocator()
893 - tracker = ComputedBlocksTracker(
894 - allocator=mock_allocator,
895 - block_size=block_size,
896 - enable_caching=True,
897 - )
898 -
899 - # Not yet allocated.
900 - tokens = [0, 1, 2, 3, 4, 5]
901 - seq1 = create_dummy_sequence(request_id=0,
902 - token_ids=tokens,
903 - block_size=block_size)
904 - mock_allocator.find_cached_blocks_prefix.return_value = []
905 - assert tracker.get_num_cached_tokens(seq1) == 0
906 -
907 - mock_allocator.find_cached_blocks_prefix.return_value = [
908 - None
909 - ] # 1 block cached.
910 - # Result is cached for prefill sequence.
911 - assert tracker.get_num_cached_tokens(seq1) == 0
912 -
913 - # Mark the sequence as non-prefill.
914 - seq1.data.update_num_computed_tokens(len(tokens)) # 6 tokens computed.
915 - assert not seq1.is_prefill()
916 -
917 - # Recomputes for decoding sequence.
918 - assert tracker.get_num_cached_tokens(seq1) == 4
919 -
920 - # Append new tokens to the sequence.
921 - num_new_tokens = 3
922 - for i in range(num_new_tokens):
923 - seq1.append_token_id(i, {i: Logprob(logprob=0.0)})
924 -
925 - assert tracker.get_num_cached_tokens(seq1) == 4
926 -
927 - # Update the allocator.
928 - mock_allocator.find_cached_blocks_prefix.return_value = [
929 - None
930 - ] * 2 # 2 blocks cached.
931 - assert tracker.get_num_cached_tokens(seq1) == 8
932 -
933 - # Remove the sequence.
934 - tracker.remove_seq(seq1.seq_id)
935 -
936 - # Re-create the sequence with the same request id to simulate recompute.
937 - seq1 = create_dummy_sequence(request_id=0,
938 - token_ids=tokens,
939 - block_size=block_size)
940 - mock_allocator.find_cached_blocks_prefix.return_value = [
941 - ] # no cached block
942 - assert tracker.get_num_cached_tokens(seq1) == 0
943 -
944 - @staticmethod
945 - def test_correct_block_hash():
946 - """
947 - Test that the block hash is correctly computed for a sequence (should
948 - match the underlying block allocator's block hash). So the number of
949 - cached tokens is correctly retrieved.
950 - """
951 - block_size = 4
952 - allocator = CpuGpuBlockAllocator.create(
953 - allocator_type="prefix_caching",
954 - num_gpu_blocks=16,
955 - num_cpu_blocks=16,
956 - block_size=block_size,
957 - )
958 - gpu_allocator = allocator._allocators[Device.GPU]
959 -
960 - tracker = ComputedBlocksTracker(
961 - allocator=allocator,
962 - block_size=block_size,
963 - enable_caching=True,
964 - )
965 -
966 - tokens = list(range(block_size * 4)) # 4 blocks.
967 - seq = create_dummy_sequence(request_id=0,
968 - token_ids=tokens,
969 - block_size=block_size)
970 - _ = TestPrefixCachingBlockAllocator.create_immutable_chain(
971 - block_size=block_size,
972 - token_ids=tokens,
973 - allocator=gpu_allocator,
974 - )
975 - allocator.mark_blocks_as_computed([])
976 -
977 - assert tracker.get_num_cached_tokens(seq) == len(tokens)
978 -
979 - @staticmethod
980 - def test_correct_extra_hash():
981 - """
982 - Test that the block hash is correctly computed based on the extra hash,
983 - ensuring it matches the allocator's block hash, specifically for the
984 - LoRA case, and that the correct number of cached tokens is retrieved.
985 - """
986 - block_size = 4
987 - allocator = CpuGpuBlockAllocator.create(
988 - allocator_type="prefix_caching",
989 - num_gpu_blocks=16,
990 - num_cpu_blocks=16,
991 - block_size=block_size,
992 - )
993 - gpu_allocator = allocator._allocators[Device.GPU]
994 -
995 - tracker = ComputedBlocksTracker(
996 - allocator=allocator,
997 - block_size=block_size,
998 - enable_caching=True,
999 - )
1000 -
1001 - tokens = list(range(block_size * 4))
1002 -
1003 - # Create a dummy LoRA sequence with a specific LoRA ID.
1004 - lora_seq = create_dummy_lora_sequence(request_id=0,
1005 - token_ids=tokens,
1006 - block_size=block_size,
1007 - lora_int_id=1)
1008 -
1009 - _ = TestPrefixCachingBlockAllocator.create_immutable_chain(
1010 - block_size=block_size,
1011 - token_ids=tokens,
1012 - allocator=gpu_allocator,
1013 - extra_hash=lora_seq.extra_hash(),
1014 - )
1015 -
1016 - allocator.mark_blocks_as_computed([])
1017 -
1018 - # Create different dummy sequences that have the same token IDs
1019 - # but different LoRA IDs.
1020 - seq = create_dummy_sequence(request_id=1,
1021 - token_ids=tokens,
1022 - block_size=block_size)
1023 -
1024 - different_lora_seq = create_dummy_lora_sequence(request_id=2,
1025 - token_ids=tokens,
1026 - block_size=block_size,
1027 - lora_int_id=2)
1028 -
1029 - # Due to the different LoRA IDs, corresponding blocks are not cached.
1030 - assert tracker.get_num_cached_tokens(seq) == 0
1031 - assert tracker.get_num_cached_tokens(different_lora_seq) == 0
1032 -
1033 - # The number of cached tokens matches the length of the tokens
1034 - # for the cached LoRA sequence.
1035 - assert tracker.get_num_cached_tokens(lora_seq) == len(tokens)
tests/core/conftest.py deleted
-12
@@ -1,12 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import pytest
4 -
5 -
6 -@pytest.fixture(scope="function", autouse=True)
7 -def use_v0_only(monkeypatch):
8 - """
9 - Since this module is V0 only, set VLLM_USE_V1=0 for
10 - all tests in the module.
11 - """
12 - monkeypatch.setenv('VLLM_USE_V1', '0')
tests/core/test_chunked_prefill_scheduler.py deleted
-858
@@ -1,858 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from unittest.mock import MagicMock
5 -
6 -import pytest # noqa
7 -
8 -from vllm.config import CacheConfig, SchedulerConfig
9 -from vllm.core.scheduler import Scheduler
10 -from vllm.engine.arg_utils import EngineArgs
11 -from vllm.engine.llm_engine import LLMEngine
12 -from vllm.sampling_params import SamplingParams
13 -from vllm.sequence import Logprob, SequenceGroup
14 -
15 -from .utils import create_dummy_prompt
16 -
17 -
18 -def get_sequence_groups(scheduler_output):
19 - return [s.seq_group for s in scheduler_output.scheduled_seq_groups]
20 -
21 -
22 -def append_new_token(seq_group: SequenceGroup, token_id: int):
23 - for seq in seq_group.get_seqs():
24 - seq.append_token_id(token_id, {token_id: Logprob(token_id)})
25 -
26 -
27 -def schedule_and_update_computed_tokens(scheduler):
28 - metas, out, _ = scheduler.schedule()
29 - for s, meta in zip(out.scheduled_seq_groups, metas):
30 - s.seq_group.update_num_computed_tokens(meta.token_chunk_size)
31 - return metas, out
32 -
33 -
34 -def test_simple():
35 - """Verify basic scheduling works."""
36 - block_size = 4
37 - num_seq_group = 4
38 - max_model_len = 16
39 - max_num_batched_tokens = 64
40 - scheduler_config = SchedulerConfig("generate",
41 - max_num_batched_tokens,
42 - num_seq_group,
43 - max_model_len,
44 - enable_chunked_prefill=True)
45 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
46 - cache_config.num_cpu_blocks = 8
47 - cache_config.num_gpu_blocks = 8
48 - scheduler = Scheduler(scheduler_config, cache_config, None)
49 - running: list[SequenceGroup] = []
50 -
51 - # Add seq groups to scheduler.
52 - for i in range(num_seq_group):
53 - _, seq_group = create_dummy_prompt(str(i),
54 - prompt_length=block_size,
55 - block_size=block_size)
56 - scheduler.add_seq_group(seq_group)
57 - running.append(seq_group)
58 -
59 - # Schedule seq groups prompts.
60 - num_tokens = block_size * num_seq_group
61 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
62 - assert set(get_sequence_groups(out)) == set(running)
63 - assert out.num_batched_tokens == num_tokens
64 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
65 - and not out.blocks_to_swap_out)
66 - assert len(seq_group_meta) == num_seq_group
67 - for s in running:
68 - append_new_token(s, 1)
69 -
70 - # Schedule seq groups generation.
71 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
72 - assert set(get_sequence_groups(out)) == set(running)
73 - assert out.num_batched_tokens == num_seq_group
74 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
75 - and not out.blocks_to_swap_out)
76 - assert len(seq_group_meta) == num_seq_group
77 -
78 -
79 -def test_chunk():
80 - """Verify prefills are chunked properly."""
81 - block_size = 4
82 - max_seqs = 60
83 - max_model_len = 80
84 - max_num_batched_tokens = 64
85 - scheduler_config = SchedulerConfig(
86 - "generate",
87 - max_num_batched_tokens,
88 - max_seqs,
89 - max_model_len,
90 - enable_chunked_prefill=True,
91 - )
92 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
93 - cache_config.num_cpu_blocks = 32
94 - cache_config.num_gpu_blocks = 32
95 - scheduler = Scheduler(scheduler_config, cache_config, None)
96 - running: list[SequenceGroup] = []
97 -
98 - # Add seq groups to scheduler.
99 - for i in range(2):
100 - _, seq_group = create_dummy_prompt(str(i),
101 - prompt_length=60,
102 - block_size=block_size)
103 - scheduler.add_seq_group(seq_group)
104 - running.append(seq_group)
105 -
106 - # Verify the second request is chunked.
107 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
108 - print()
109 - assert set(get_sequence_groups(out)) == set(running)
110 - assert seq_group_meta[0].token_chunk_size == 60
111 - # Verify it is chunked.
112 - assert seq_group_meta[1].token_chunk_size == 4
113 - assert out.num_prefill_groups == 2
114 - assert out.num_batched_tokens == 64
115 - # Only the first seq group has a new token appended.
116 - append_new_token(running[0], 1)
117 -
118 - # One chunked prefill, and one decoding.
119 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
120 - assert set(get_sequence_groups(out)) == set(running)
121 - # The first one is prefill. Scheduler guarantees ordering.
122 - assert seq_group_meta[0].token_chunk_size == 56
123 - # The second one is a chunked prefill.
124 - assert seq_group_meta[1].token_chunk_size == 1
125 - assert out.num_prefill_groups == 1
126 - assert out.num_batched_tokens == 57
127 -
128 -
129 -def test_concurrent_chunking():
130 - """Verify prefills are chunked properly when
131 - --max-num-partial-prefills is > 1"""
132 - block_size = 4
133 - max_seqs = 60
134 - max_model_len = 2000
135 - max_num_batched_tokens = 64
136 - scheduler_config = SchedulerConfig(
137 - "generate",
138 - max_num_batched_tokens,
139 - max_seqs,
140 - max_model_len,
141 - enable_chunked_prefill=True,
142 - max_num_partial_prefills=2, # Up to 2 partial prefills at a time
143 - )
144 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
145 - cache_config.num_cpu_blocks = 32
146 - cache_config.num_gpu_blocks = 32
147 - scheduler = Scheduler(scheduler_config, cache_config, None)
148 - running: list[SequenceGroup] = []
149 -
150 - # Add seq groups to scheduler.
151 - for i in range(2):
152 - _, seq_group = create_dummy_prompt(str(i),
153 - prompt_length=60,
154 - block_size=block_size)
155 - scheduler.add_seq_group(seq_group)
156 - running.append(seq_group)
157 -
158 - # Verify both requests are chunked with half of max_num_batched_tokens each
159 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
160 - assert set(get_sequence_groups(out)) == set(running)
161 - assert seq_group_meta[0].token_chunk_size == 32
162 - assert seq_group_meta[1].token_chunk_size == 32
163 - assert out.num_prefill_groups == 2
164 - assert out.num_batched_tokens == 64
165 -
166 - # After one iteration, both should have 60 - 32 = 28 tokens left to prefill
167 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
168 - assert set(get_sequence_groups(out)) == set(running)
169 - assert seq_group_meta[0].token_chunk_size == 28
170 - assert seq_group_meta[1].token_chunk_size == 28
171 - assert out.num_prefill_groups == 2
172 - assert out.num_batched_tokens == 56
173 -
174 -
175 -def test_concurrent_chunking_large_requests():
176 - """Verify large prefill requests are run one at a time"""
177 - block_size = 4
178 - max_seqs = 60
179 - max_model_len = 2000
180 - max_num_batched_tokens = 64
181 - scheduler_config = SchedulerConfig(
182 - "generate",
183 - max_num_batched_tokens,
184 - max_seqs,
185 - max_model_len,
186 - enable_chunked_prefill=True,
187 - max_num_partial_prefills=2, # Up to 2 partial prefills at a time
188 - )
189 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
190 - cache_config.num_cpu_blocks = 3200 # large KV cache size for large requests
191 - cache_config.num_gpu_blocks = 3200
192 - scheduler = Scheduler(scheduler_config, cache_config, None)
193 -
194 - # Add seq groups to scheduler.
195 - for i in range(2):
196 - _, seq_group = create_dummy_prompt(
197 - str(i),
198 - prompt_length=1200, # Very large prompt
199 - block_size=block_size)
200 - scheduler.add_seq_group(seq_group)
201 -
202 - # Verify only a single request is chunked, and it gets all 64 tokens
203 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
204 - assert len(get_sequence_groups(out)) == 1
205 - assert seq_group_meta[0].token_chunk_size == 64
206 - assert out.num_prefill_groups == 1
207 - assert out.num_batched_tokens == 64
208 -
209 -
210 -def test_short_prompts_jump_long_prompts_in_queue():
211 - """Verify large prefill requests are punted behind smaller ones if
212 - another large prefill request is already running"""
213 - block_size = 4
214 - max_seqs = 60
215 - max_model_len = 2000
216 - max_num_batched_tokens = 64
217 - scheduler_config = SchedulerConfig(
218 - "generate",
219 - max_num_batched_tokens,
220 - max_seqs,
221 - max_model_len,
222 - enable_chunked_prefill=True,
223 - max_num_partial_prefills=2, # Up to 2 partial prefills at a time
224 - )
225 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
226 - cache_config.num_cpu_blocks = 3200 # large KV cache size for large requests
227 - cache_config.num_gpu_blocks = 3200
228 - scheduler = Scheduler(scheduler_config, cache_config, None)
229 - long_seqs: list[SequenceGroup] = []
230 - short_seqs: list[SequenceGroup] = []
231 -
232 - # Add 2 large seq groups to scheduler.
233 - for i in range(2):
234 - _, seq_group = create_dummy_prompt(
235 - str(i),
236 - prompt_length=1200, # Very large prompt
237 - block_size=block_size)
238 - scheduler.add_seq_group(seq_group)
239 - long_seqs.append(seq_group)
240 - assert seq_group.is_prefill()
241 -
242 - # Add 2 small seq groups behind them
243 - for i in range(2):
244 - _, seq_group = create_dummy_prompt(
245 - str(i + 2),
246 - prompt_length=40, # Very small prompt
247 - block_size=block_size)
248 - scheduler.add_seq_group(seq_group)
249 - short_seqs.append(seq_group)
250 - assert seq_group.is_prefill()
251 -
252 - # Verify one large req and 1 small req chunked
253 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
254 - assert seq_group_meta[0].token_chunk_size == 32 # large req gets 32 tokens
255 - assert seq_group_meta[1].token_chunk_size == 32 # small req gets 32 tokens
256 -
257 - # all 4 are prefilling
258 - assert long_seqs[0].is_prefill()
259 - assert long_seqs[1].is_prefill()
260 - assert short_seqs[0].is_prefill()
261 - assert short_seqs[1].is_prefill()
262 - # First short and first long sequences have been scheduled
263 - assert long_seqs[0].first_seq.get_num_computed_tokens() == 32
264 - assert long_seqs[1].first_seq.get_num_computed_tokens() == 0
265 - assert short_seqs[0].first_seq.get_num_computed_tokens() == 32
266 - assert short_seqs[1].first_seq.get_num_computed_tokens() == 0
267 -
268 - assert out.num_prefill_groups == 2
269 - assert out.num_batched_tokens == 64
270 -
271 - # in the second iteration,
272 - # the first small request had only 8 tokens left
273 - # so it went to decode
274 - # The other small req is scheduled
275 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
276 - # the new small req got 64 - (32+8) tokens
277 - assert seq_group_meta[0].token_chunk_size == 24
278 - assert seq_group_meta[1].token_chunk_size == 32 # large req still got 32
279 - # the other small request had only 8 tokens left
280 - assert seq_group_meta[2].token_chunk_size == 8 # 40-32
281 -
282 - # The first small request got to decode now
283 - assert long_seqs[0].is_prefill()
284 - assert long_seqs[1].is_prefill()
285 - assert not short_seqs[0].is_prefill()
286 - assert short_seqs[1].is_prefill()
287 - # Both small requests have started in front of the second long request
288 - assert long_seqs[0].first_seq.get_num_computed_tokens() == 64
289 - assert long_seqs[1].first_seq.get_num_computed_tokens() == 0
290 - assert short_seqs[0].first_seq.get_num_computed_tokens() == 40
291 - assert short_seqs[1].first_seq.get_num_computed_tokens() == 24
292 -
293 - assert out.num_prefill_groups == 3
294 - assert out.num_batched_tokens == 64
295 - # the first small seq group has a new token appended.
296 - append_new_token(short_seqs[0], 1)
297 -
298 - # in the third iteration,
299 - # the first small request is already decoding
300 - # the second small request only has 16 tokens left and will enter decoding
301 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
302 - assert seq_group_meta[0].token_chunk_size == 32 # large still got 32
303 - # small req finished prefilling 40-24=16 tokens
304 - assert seq_group_meta[1].token_chunk_size == 16
305 - assert seq_group_meta[2].token_chunk_size == 1 # decode
306 - assert out.num_prefill_groups == 2
307 - assert out.num_batched_tokens == 49 # (32+16+1 decode)
308 -
309 - # both small requests have now reached decode
310 - assert long_seqs[0].is_prefill()
311 - assert long_seqs[1].is_prefill()
312 - assert not short_seqs[0].is_prefill()
313 - assert not short_seqs[1].is_prefill()
314 - assert long_seqs[0].first_seq.get_num_computed_tokens() == 96
315 - assert long_seqs[1].first_seq.get_num_computed_tokens() == 0
316 - assert short_seqs[0].first_seq.get_num_computed_tokens() == 41
317 - assert short_seqs[1].first_seq.get_num_computed_tokens() == 40
318 -
319 - # both the small seq groups have a new token appended
320 - append_new_token(short_seqs[0], 1)
321 - append_new_token(short_seqs[1], 1)
322 -
323 - # in the fourth iteration, both small requests are decoding
324 - # so large request gets all the budget
325 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
326 -
327 - # large req gets 62 tokens (minus 2 for decode)
328 - assert seq_group_meta[0].token_chunk_size == 62
329 - assert seq_group_meta[1].token_chunk_size == 1 # decode
330 - assert seq_group_meta[2].token_chunk_size == 1 # decode
331 - assert out.num_prefill_groups == 1
332 - assert out.num_batched_tokens == 64
333 -
334 - assert long_seqs[0].first_seq.get_num_computed_tokens() == 158
335 -
336 - # assert long_seqs[0].is_prefill()
337 - # assert long_seqs[1].is_prefill()
338 - # assert not short_seqs[0].is_prefill()
339 - # assert not short_seqs[1].is_prefill()
340 -
341 - # # both the small seq groups have a new token appended
342 - # append_new_token(short_seqs[0], 1)
343 - # append_new_token(short_seqs[1], 1)
344 -
345 - # # in the fifth iteration, large request gets all the budget
346 - # # while both small requests are decoding
347 - # seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
348 - # assert seq_group_meta[0].token_chunk_size == 62
349 - # assert seq_group_meta[1].token_chunk_size == 1 # decode
350 - # assert seq_group_meta[2].token_chunk_size == 1 # decode
351 - # assert out.num_prefill_groups == 1
352 - # assert out.num_batched_tokens == 64
353 -
354 -
355 -def test_complex():
356 - block_size = 4
357 - max_seqs = 60
358 - max_model_len = 80
359 - max_num_batched_tokens = 64
360 - scheduler_config = SchedulerConfig(
361 - "generate",
362 - max_num_batched_tokens,
363 - max_seqs,
364 - max_model_len,
365 - enable_chunked_prefill=True,
366 - )
367 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
368 - cache_config.num_cpu_blocks = 64
369 - cache_config.num_gpu_blocks = 64
370 - scheduler = Scheduler(scheduler_config, cache_config, None)
371 - running: list[SequenceGroup] = []
372 -
373 - # Add seq groups to scheduler.
374 - for i in range(2):
375 - _, seq_group = create_dummy_prompt(str(i),
376 - prompt_length=60,
377 - block_size=block_size)
378 - scheduler.add_seq_group(seq_group)
379 - running.append(seq_group)
380 - assert seq_group.is_prefill()
381 -
382 - # Verify the second request is chunked.
383 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
384 -
385 - assert set(get_sequence_groups(out)) == set(running)
386 - assert seq_group_meta[0].token_chunk_size == 60
387 - # Verify it is chunked.
388 - assert seq_group_meta[1].token_chunk_size == 4
389 - assert not running[0].is_prefill()
390 - assert running[1].is_prefill()
391 - assert out.num_prefill_groups == 2
392 - assert out.num_batched_tokens == 64
393 - # Only the first seq group has a new token appended.
394 - append_new_token(running[0], 1)
395 -
396 - # Add 2 more requests.
397 - for i in range(2, 4):
398 - _, seq_group = create_dummy_prompt(str(i),
399 - prompt_length=60,
400 - block_size=block_size)
401 - scheduler.add_seq_group(seq_group)
402 - running.append(seq_group)
403 -
404 - # Decoding & chunked prefill & first chunk of 3rd request is scheduled.
405 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
406 - assert len(get_sequence_groups(out)) == 3
407 - # The first one is the first chunked prefill.
408 - assert seq_group_meta[0].token_chunk_size == 7
409 - # The second one is the second new chunked prefill.
410 - assert seq_group_meta[1].token_chunk_size == 56
411 - # The last one is decode.
412 - assert seq_group_meta[2].token_chunk_size == 1
413 - # Two of them are in chunked prefill.
414 - assert out.num_prefill_groups == 2
415 - assert out.num_batched_tokens == 64
416 - # The first 2 requests are now in decodine phase.
417 - append_new_token(running[0], 1)
418 - assert not running[0].is_prefill()
419 - append_new_token(running[1], 1)
420 - assert not running[1].is_prefill()
421 - # The third request is still in prefill stage.
422 - assert running[2].is_prefill()
423 -
424 -
425 -def test_maximal_decoding():
426 - """Verify decoding requests are prioritized."""
427 - block_size = 4
428 - max_seqs = 2
429 - max_model_len = 8
430 - max_num_batched_tokens = 2
431 - scheduler_config = SchedulerConfig(
432 - "generate",
433 - max_num_batched_tokens,
434 - max_seqs,
435 - max_model_len,
436 - enable_chunked_prefill=True,
437 - )
438 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
439 - cache_config.num_cpu_blocks = 8
440 - cache_config.num_gpu_blocks = 8
441 - scheduler = Scheduler(scheduler_config, cache_config, None)
442 - running: list[SequenceGroup] = []
443 -
444 - # Add seq groups to scheduler.
445 - for i in range(2):
446 - _, seq_group = create_dummy_prompt(str(i),
447 - prompt_length=2,
448 - block_size=block_size)
449 - scheduler.add_seq_group(seq_group)
450 - running.append(seq_group)
451 - assert seq_group.is_prefill()
452 -
453 - # The first prefill is scheduled.
454 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
455 - assert len(get_sequence_groups(out)) == 1
456 - assert seq_group_meta[0].token_chunk_size == 2
457 - assert not running[0].is_prefill()
458 - assert running[1].is_prefill()
459 - assert out.num_prefill_groups == 1
460 - assert out.num_batched_tokens == 2
461 - # Only the first seq group has a new token appended.
462 - append_new_token(running[0], 1)
463 -
464 - # Create one more seq_group.
465 - _, seq_group = create_dummy_prompt("3",
466 - prompt_length=2,
467 - block_size=block_size)
468 - scheduler.add_seq_group(seq_group)
469 - running.append(seq_group)
470 - assert seq_group.is_prefill()
471 - # The first decoding + second chunk is scheduled.
472 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
473 - assert len(get_sequence_groups(out)) == 2
474 - assert seq_group_meta[0].token_chunk_size == 1
475 - assert seq_group_meta[1].token_chunk_size == 1
476 - assert not running[0].is_prefill()
477 - assert running[1].is_prefill()
478 - assert running[2].is_prefill()
479 - assert out.num_prefill_groups == 1
480 - assert out.num_batched_tokens == 2
481 - append_new_token(running[0], 1)
482 -
483 - # Decoding + running prefill is prioritized.
484 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
485 - assert len(get_sequence_groups(out)) == 2
486 - assert seq_group_meta[0].token_chunk_size == 1
487 - assert seq_group_meta[1].token_chunk_size == 1
488 - assert not running[0].is_prefill()
489 - assert not running[1].is_prefill()
490 - assert out.num_prefill_groups == 1
491 - assert out.num_batched_tokens == 2
492 - append_new_token(running[0], 1)
493 - append_new_token(running[1], 1)
494 -
495 - # Only decoding is prioritized.
496 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
497 - assert len(get_sequence_groups(out)) == 2
498 - assert seq_group_meta[0].token_chunk_size == 1
499 - assert seq_group_meta[1].token_chunk_size == 1
500 - assert not running[0].is_prefill()
501 - assert not running[1].is_prefill()
502 - assert out.num_prefill_groups == 0
503 - assert out.num_batched_tokens == 2
504 - append_new_token(running[0], 1)
505 - append_new_token(running[1], 1)
506 -
507 - # After aborting the decoding request, the fcfs new prefill is prioritized.
508 - scheduler.abort_seq_group(running[0].request_id)
509 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
510 - assert len(get_sequence_groups(out)) == 2
511 - assert seq_group_meta[0].token_chunk_size == 1
512 - assert seq_group_meta[1].token_chunk_size == 1
513 - assert not running[1].is_prefill()
514 - assert running[2].is_prefill()
515 - assert out.num_prefill_groups == 1
516 - assert out.num_batched_tokens == 2
517 -
518 -
519 -def test_prompt_limit():
520 - """Verify max_num_batched_tokens < max_model_len is possible."""
521 - block_size = 4
522 - max_seqs = 32
523 - max_model_len = 64
524 - max_num_batched_tokens = 32
525 - scheduler_config = SchedulerConfig(
526 - "generate",
527 - max_num_batched_tokens,
528 - max_seqs,
529 - max_model_len,
530 - enable_chunked_prefill=True,
531 - )
532 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
533 - cache_config.num_cpu_blocks = 16
534 - cache_config.num_gpu_blocks = 16
535 - scheduler = Scheduler(scheduler_config, cache_config, None)
536 - running: list[SequenceGroup] = []
537 -
538 - _, seq_group = create_dummy_prompt("1",
539 - prompt_length=48,
540 - block_size=block_size)
541 - scheduler.add_seq_group(seq_group)
542 - running.append(seq_group)
543 - assert seq_group.is_prefill()
544 -
545 - # The prompt length > max_num_batched_tokens should be still scheduled.
546 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
547 - assert len(get_sequence_groups(out)) == 1
548 - assert seq_group_meta[0].token_chunk_size == 32
549 - assert running[0].is_prefill()
550 - assert out.num_prefill_groups == 1
551 - assert out.num_batched_tokens == 32
552 -
553 -
554 -def test_prompt_limit_exceed():
555 - block_size = 4
556 - max_seqs = 64
557 - max_model_len = 32
558 - max_num_batched_tokens = 64
559 - scheduler_config = SchedulerConfig("generate",
560 - max_num_batched_tokens,
561 - max_seqs,
562 - max_model_len,
563 - enable_chunked_prefill=True)
564 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
565 - cache_config.num_cpu_blocks = 16
566 - cache_config.num_gpu_blocks = 16
567 - scheduler = Scheduler(scheduler_config, cache_config, None)
568 - running: list[SequenceGroup] = []
569 - _, seq_group = create_dummy_prompt("2",
570 - prompt_length=48,
571 - block_size=block_size)
572 - scheduler.add_seq_group(seq_group)
573 - running.append(seq_group)
574 - assert seq_group.is_prefill()
575 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
576 - assert len(out.ignored_seq_groups) == 1
577 - assert out.ignored_seq_groups[0] == seq_group
578 -
579 -
580 -def test_chunked_prefill_preempt():
581 - """Verify preempt works with chunked prefill requests"""
582 - block_size = 4
583 - max_seqs = 30
584 - max_model_len = 200
585 - max_num_batched_tokens = 30
586 - scheduler_config = SchedulerConfig(
587 - "generate",
588 - max_num_batched_tokens,
589 - max_seqs,
590 - max_model_len,
591 - enable_chunked_prefill=True,
592 - )
593 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
594 - cache_config.num_cpu_blocks = 16
595 - cache_config.num_gpu_blocks = 16
596 - scheduler = Scheduler(scheduler_config, cache_config, None)
597 -
598 - _, seq_group = create_dummy_prompt("1",
599 - prompt_length=60,
600 - block_size=block_size)
601 - scheduler.add_seq_group(seq_group)
602 - _, out = schedule_and_update_computed_tokens(scheduler)
603 - # The request is chunked.
604 - # prefill scheduled now.
605 - assert len(out.scheduled_seq_groups) == 1
606 - assert out.num_prefill_groups == 1
607 - assert seq_group.is_prefill()
608 - assert out.num_batched_tokens == max_num_batched_tokens
609 -
610 - # The request should be preempted.
611 - scheduler.block_manager.can_append_slots = MagicMock()
612 -
613 - def cannot_append_second_group1(seq_group, num_lookahead_slots):
614 - return seq_group.request_id != "1"
615 -
616 - scheduler.block_manager.can_append_slots.side_effect = (
617 - cannot_append_second_group1)
618 -
619 - # The running prefill is now preempted.
620 - _, out = schedule_and_update_computed_tokens(scheduler)
621 - assert len(out.scheduled_seq_groups) == 0
622 - assert out.num_batched_tokens == 0
623 - assert out.blocks_to_swap_out == []
624 - assert out.blocks_to_swap_in == []
625 -
626 - # Make sure we can reschedule preempted request.
627 - _, out = schedule_and_update_computed_tokens(scheduler)
628 - assert len(out.scheduled_seq_groups) == 1
629 - assert out.num_prefill_groups == 1
630 - assert seq_group.is_prefill()
631 - assert out.num_batched_tokens == max_num_batched_tokens
632 - assert seq_group.get_num_uncomputed_tokens() == 30
633 -
634 - # We should be able to run prefill twice as it is chunked.
635 - def cannot_append_second_group2(seq_group, num_lookahead_slots):
636 - return True
637 -
638 - scheduler.block_manager.can_append_slots.side_effect = (
639 - cannot_append_second_group2)
640 - _, out = schedule_and_update_computed_tokens(scheduler)
641 - assert len(out.scheduled_seq_groups) == 1
642 - assert out.num_prefill_groups == 1
643 - assert not seq_group.is_prefill()
644 - assert out.num_batched_tokens == max_num_batched_tokens
645 -
646 -
647 -def test_chunked_prefill_spec_prefill():
648 - """Verify that the num_lookahead_slots is set appropriately for an all"""
649 - """prefill batch."""
650 - block_size = 4
651 - max_seqs = 30
652 - max_model_len = 200
653 - max_num_batched_tokens = 30
654 - num_lookahead_slots = 4
655 - scheduler_config = SchedulerConfig(
656 - "generate",
657 - max_num_batched_tokens,
658 - max_seqs,
659 - max_model_len,
660 - enable_chunked_prefill=True,
661 - num_lookahead_slots=num_lookahead_slots,
662 - )
663 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
664 - cache_config.num_cpu_blocks = 16
665 - cache_config.num_gpu_blocks = 16
666 - scheduler = Scheduler(scheduler_config, cache_config, None)
667 -
668 - _, seq_group = create_dummy_prompt("1",
669 - prompt_length=30,
670 - block_size=block_size)
671 - scheduler.add_seq_group(seq_group)
672 - _, out = schedule_and_update_computed_tokens(scheduler)
673 - # The request is chunked.
674 - # prefill scheduled now.
675 - assert len(out.scheduled_seq_groups) == 1
676 - assert out.num_prefill_groups == 1
677 - assert out.num_batched_tokens == max_num_batched_tokens
678 - print(out.num_lookahead_slots)
679 - assert out.num_lookahead_slots == 0
680 -
681 -
682 -def test_chunked_prefill_max_seqs():
683 - block_size = 4
684 - max_seqs = 2
685 - max_model_len = 80
686 - max_num_batched_tokens = 64
687 - scheduler_config = SchedulerConfig(
688 - "generate",
689 - max_num_batched_tokens,
690 - max_seqs,
691 - max_model_len,
692 - enable_chunked_prefill=True,
693 - )
694 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
695 - cache_config.num_cpu_blocks = 128
696 - cache_config.num_gpu_blocks = 128
697 - scheduler = Scheduler(scheduler_config, cache_config, None)
698 - running: list[SequenceGroup] = []
699 -
700 - _, seq_group = create_dummy_prompt("1",
701 - prompt_length=65,
702 - block_size=block_size)
703 - scheduler.add_seq_group(seq_group)
704 - running.append(seq_group)
705 - # The first prefill is chunked.
706 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
707 - assert seq_group_meta[0].token_chunk_size == max_num_batched_tokens
708 - assert len(get_sequence_groups(out)) == 1
709 -
710 - # Add new requests.
711 - for i in range(4):
712 - _, seq_group = create_dummy_prompt(str(i),
713 - prompt_length=65,
714 - block_size=block_size)
715 - scheduler.add_seq_group(seq_group)
716 - running.append(seq_group)
717 -
718 - # Make sure only 2 requests are scheduled.
719 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
720 - assert out.num_batched_tokens == max_num_batched_tokens
721 - assert len(get_sequence_groups(out)) == 2
722 - assert not running[0].is_prefill()
723 - assert running[1].is_prefill()
724 - append_new_token(running[0], 1)
725 -
726 - # Although we have enough token budget, we can only schedule max_seqs.
727 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
728 - assert seq_group_meta[0].token_chunk_size == 2
729 - assert seq_group_meta[1].token_chunk_size == 1
730 - assert out.num_batched_tokens == 3
731 - assert len(get_sequence_groups(out)) == max_seqs
732 - assert not running[0].is_prefill()
733 - assert not running[1].is_prefill()
734 -
735 -
736 -def test_prefix_caching():
737 - """Verify allocating full blocks when prefix caching is enabled."""
738 - block_size = 4
739 - max_seqs = 10
740 - max_model_len = 80
741 - max_num_batched_tokens = 64
742 - scheduler_config = SchedulerConfig(
743 - "generate",
744 - max_num_batched_tokens,
745 - max_seqs,
746 - max_model_len,
747 - enable_chunked_prefill=True,
748 - )
749 - cache_config = CacheConfig(block_size,
750 - 1.0,
751 - 1,
752 - "auto",
753 - enable_prefix_caching=True)
754 - cache_config.num_cpu_blocks = 0
755 - cache_config.num_gpu_blocks = 32
756 - scheduler = Scheduler(scheduler_config, cache_config, None)
757 - running: list[SequenceGroup] = []
758 -
759 - # Add seq groups to scheduler.
760 - for i in range(2):
761 - _, seq_group = create_dummy_prompt(str(i),
762 - block_size=block_size,
763 - prompt_length=50)
764 - scheduler.add_seq_group(seq_group)
765 - running.append(seq_group)
766 -
767 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
768 - assert set(get_sequence_groups(out)) == set(running)
769 - assert seq_group_meta[0].token_chunk_size == 50
770 - # Verify it is chunked. Note that although the budget is 64-50=14,
771 - # we only allocate full blocks for prefix caching, so only 4*(14//4)=12
772 - # tokens are allocated.
773 - assert seq_group_meta[1].token_chunk_size == 12
774 - assert out.num_prefill_groups == 2
775 - assert out.num_batched_tokens == 62
776 -
777 -
778 -def test_prefix_caching_with_concurrent_partial_prefills():
779 - """Verify allocating full blocks when prefix caching is enabled with
780 - --max-num-partial-prefills > 1."""
781 - block_size = 4
782 - max_seqs = 10
783 - max_model_len = 8000
784 - max_num_batched_tokens = 60 # With two slots, each slot will get 30 tokens
785 - scheduler_config = SchedulerConfig("generate",
786 - max_num_batched_tokens,
787 - max_seqs,
788 - max_model_len,
789 - enable_chunked_prefill=True,
790 - max_num_partial_prefills=2)
791 - cache_config = CacheConfig(block_size,
792 - 1.0,
793 - 1,
794 - "auto",
795 - enable_prefix_caching=True)
796 - cache_config.num_cpu_blocks = 0
797 - cache_config.num_gpu_blocks = 32
798 - scheduler = Scheduler(scheduler_config, cache_config, None)
799 - running: list[SequenceGroup] = []
800 -
801 - # Add seq groups to scheduler.
802 - for i in range(2):
803 - _, seq_group = create_dummy_prompt(str(i),
804 - block_size=block_size,
805 - prompt_length=50)
806 - scheduler.add_seq_group(seq_group)
807 - running.append(seq_group)
808 -
809 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
810 - assert set(get_sequence_groups(out)) == set(running)
811 - # To partially prefill both sequences, both can chunk up to 30 tokens
812 - # But the next lowest multiple of the block size (4) is 28
813 - assert seq_group_meta[0].token_chunk_size == 28
814 - assert seq_group_meta[1].token_chunk_size == 28
815 - assert out.num_prefill_groups == 2
816 - assert out.num_batched_tokens == 56
817 -
818 - # On the next iteration, both sequences should finish prefill
819 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
820 - assert set(get_sequence_groups(out)) == set(running)
821 - # Both sequences have 50 - 28 = 22 tokens left to prefill.
822 - # This is not a multiple of the block size, but we don't care since we don't
823 - # cache the final partial block of prefix sequences
824 - assert seq_group_meta[0].token_chunk_size == 22
825 - assert seq_group_meta[1].token_chunk_size == 22
826 - assert out.num_prefill_groups == 2
827 - assert out.num_batched_tokens == 44
828 -
829 -
830 -@pytest.mark.parametrize("model", ["facebook/opt-125m"])
831 -@pytest.mark.parametrize("max_num_partial_prefills", [2, 4, 8])
832 -def test_chunked_prefill_with_actual_engine(model: str,
833 - max_num_partial_prefills: int):
834 - """Make sure the model can actually sample with concurrent
835 - partial prefills
836 - """
837 -
838 - prompt = "hello" * 40
839 -
840 - engine_args = EngineArgs(
841 - model=model,
842 - max_num_partial_prefills=max_num_partial_prefills,
843 - max_num_batched_tokens=40,
844 - max_num_seqs=8,
845 - enable_chunked_prefill=True,
846 - gpu_memory_utilization=0.8,
847 - )
848 -
849 - engine = LLMEngine.from_engine_args(engine_args)
850 - sampling_params = SamplingParams(temperature=0)
851 -
852 - for req_num in range(max_num_partial_prefills):
853 - engine.add_request(f"{req_num}", prompt, sampling_params)
854 - # first step
855 - request_outputs = engine.step()
856 - # means all are prefilling
857 - assert len(request_outputs) == 0
858 - assert len(engine.scheduler[0].running) == max_num_partial_prefills
tests/core/test_num_computed_tokens_update.py deleted
-67
@@ -1,67 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from tests.conftest import VllmRunner
7 -from tests.core.utils import create_dummy_prompt
8 -from vllm.engine.llm_engine import LLMEngine
9 -from vllm.sequence import SequenceGroup
10 -
11 -MODEL = "JackFram/llama-160m"
12 -
13 -
14 -def add_seq_group_to_engine(engine: LLMEngine, seq_group: SequenceGroup):
15 - scheduler = engine.scheduler[0]
16 - scheduler.add_seq_group(seq_group)
17 -
18 -
19 -@pytest.mark.parametrize("enable_chunked_prefill", [False, True])
20 -@pytest.mark.parametrize("enforce_eager", [False, True])
21 -def test_num_computed_tokens_update(enable_chunked_prefill: bool,
22 - enforce_eager: bool):
23 -
24 - # Make a vllm engine
25 - runner = VllmRunner(model_name=MODEL,
26 - gpu_memory_utilization=0.7,
27 - enable_chunked_prefill=enable_chunked_prefill,
28 - enforce_eager=enforce_eager)
29 - engine: LLMEngine = runner.llm.llm_engine
30 -
31 - num_prompt_steps = 1
32 -
33 - num_output_tokens_list = [4, 8, 12, 15, 16, 17]
34 -
35 - # Create sequence and add to engine
36 - prompt_len = 10
37 -
38 - for req_idx, num_output_tokens in enumerate(num_output_tokens_list):
39 - seq, seq_group = create_dummy_prompt(request_id=str(req_idx),
40 - prompt_length=prompt_len,
41 - min_tokens=num_output_tokens,
42 - max_tokens=num_output_tokens)
43 - add_seq_group_to_engine(engine, seq_group)
44 -
45 - assert seq.data.get_num_computed_tokens() == 0
46 -
47 - for _ in range(num_prompt_steps):
48 - # prompt steps
49 - engine.step()
50 -
51 - if not seq.is_finished():
52 - prompt_num_computed_tokens = seq.data.get_num_computed_tokens()
53 - # Test correctness of num_computed_tokens after the prompt steps
54 - assert prompt_num_computed_tokens == \
55 - prompt_len + num_prompt_steps - 1
56 -
57 - decode_step_counter = 0
58 - while not seq.is_finished():
59 - # Test correctness of num_computed_tokens after the decode steps
60 - assert seq.data.get_num_computed_tokens(
61 - ) == prompt_num_computed_tokens + decode_step_counter
62 - engine.step()
63 - decode_step_counter += 1
64 -
65 - # Test correctness of num_computed_tokens after the sequence finish.
66 - assert seq.data.get_num_computed_tokens(
67 - ) == prompt_len + num_output_tokens - 1
tests/core/test_scheduler.py deleted
-1338
@@ -1,1338 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import time
5 -from collections import deque
6 -from typing import Optional
7 -from unittest.mock import MagicMock
8 -
9 -import pytest # noqa
10 -import torch
11 -from torch import Use # noqa
12 -
13 -from vllm.config import CacheConfig, SchedulerConfig
14 -from vllm.config.lora import LoRAConfig
15 -from vllm.core.interfaces import AllocStatus
16 -from vllm.core.scheduler import Scheduler, SchedulingBudget
17 -from vllm.lora.request import LoRARequest
18 -from vllm.sequence import SequenceGroup, SequenceStatus
19 -
20 -from .utils import (append_new_token, append_new_token_seq,
21 - append_new_token_seq_group, create_dummy_prompt,
22 - get_sequence_groups, schedule_and_update_computed_tokens)
23 -
24 -
25 -def test_scheduler_add_seq_group():
26 - block_size = 4
27 - scheduler_config = SchedulerConfig(
28 - "generate",
29 - max_num_batched_tokens=100,
30 - max_num_seqs=64,
31 - max_model_len=1,
32 - )
33 - cache_config = CacheConfig(block_size, 1.0, 1, cache_dtype="auto")
34 - cache_config.num_cpu_blocks = 4
35 - cache_config.num_gpu_blocks = 4
36 - scheduler = Scheduler(scheduler_config, cache_config, None)
37 -
38 - # Add seq group to scheduler.
39 - num_seq_group = 4
40 - for i in range(num_seq_group):
41 - _, seq_group = create_dummy_prompt(str(i),
42 - block_size,
43 - block_size=block_size)
44 - scheduler.add_seq_group(seq_group)
45 - assert scheduler.get_num_unfinished_seq_groups() == i + 1
46 -
47 -
48 -def test_scheduler_abort_seq_group():
49 - block_size = 4
50 - scheduler_config = SchedulerConfig(
51 - "generate",
52 - max_num_batched_tokens=100,
53 - max_num_seqs=64,
54 - max_model_len=1,
55 - )
56 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
57 - cache_config.num_cpu_blocks = 4
58 - cache_config.num_gpu_blocks = 4
59 - scheduler = Scheduler(scheduler_config, cache_config, None)
60 -
61 - # Add multiple seq groups to scheduler.
62 - num_seq_group = 4
63 - request_ids: set[str] = set()
64 - for i in range(num_seq_group):
65 - _, seq_group = create_dummy_prompt(str(i), block_size)
66 - scheduler.add_seq_group(seq_group)
67 - request_ids.add(str(i))
68 -
69 - # Abort all added seq groups.
70 - assert scheduler.get_num_unfinished_seq_groups() == num_seq_group
71 - scheduler.abort_seq_group(request_ids)
72 - assert scheduler.get_num_unfinished_seq_groups() == 0
73 -
74 -
75 -def test_scheduler_schedule_simple():
76 - block_size = 4
77 - num_seq_group = 4
78 - max_model_len = 16
79 - scheduler_config = SchedulerConfig(
80 - "generate",
81 - max_num_batched_tokens=64,
82 - max_num_seqs=num_seq_group,
83 - max_model_len=max_model_len,
84 - )
85 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
86 - cache_config.num_cpu_blocks = 8
87 - cache_config.num_gpu_blocks = 8
88 - scheduler = Scheduler(scheduler_config, cache_config, None)
89 - running: list[SequenceGroup] = []
90 -
91 - # Add seq groups to scheduler.
92 - for i in range(num_seq_group):
93 - _, seq_group = create_dummy_prompt(str(i),
94 - prompt_length=block_size,
95 - block_size=block_size)
96 - scheduler.add_seq_group(seq_group)
97 - running.append(seq_group)
98 -
99 - # Schedule seq groups prompts.
100 - num_tokens = block_size * num_seq_group
101 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
102 - assert set(get_sequence_groups(out)) == set(running)
103 - assert out.num_batched_tokens == num_tokens
104 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
105 - and not out.blocks_to_swap_out)
106 - assert len(seq_group_meta) == num_seq_group
107 - append_new_token(out, 1)
108 -
109 - # Schedule seq groups generation.
110 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
111 - assert set(get_sequence_groups(out)) == set(running)
112 - assert out.num_batched_tokens == num_seq_group
113 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
114 - and not out.blocks_to_swap_out)
115 - assert len(seq_group_meta) == num_seq_group
116 - append_new_token(out, 1)
117 -
118 -
119 -def test_scheduler_prefill_prioritized():
120 - """Verify running batched tokens are not applied to prefill requests."""
121 - block_size = 4
122 - max_model_len = 30
123 - max_batched_num_tokens = 30
124 - scheduler_config = SchedulerConfig(
125 - "generate",
126 - max_num_batched_tokens=max_batched_num_tokens,
127 - max_num_seqs=2,
128 - max_model_len=max_model_len,
129 - )
130 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
131 - cache_config.num_cpu_blocks = 16
132 - cache_config.num_gpu_blocks = 16
133 - scheduler = Scheduler(scheduler_config, cache_config, None)
134 -
135 - # Add seq groups to scheduler.
136 - _, seq_group_a = create_dummy_prompt("1", 1, block_size=block_size)
137 - scheduler.add_seq_group(seq_group_a)
138 -
139 - # Schedule seq groups prompts.
140 - _, out = schedule_and_update_computed_tokens(scheduler)
141 - assert get_sequence_groups(out) == [seq_group_a]
142 -
143 - # Add a new prefill request B.
144 - _, seq_group_b = create_dummy_prompt("2", 30, block_size=block_size)
145 - scheduler.add_seq_group(seq_group_b)
146 -
147 - # Verify prefill requests are prioritized. Since max_batched_num_tokens
148 - # is 1, new prefill request has to be scheduled first.
149 - _, out = schedule_and_update_computed_tokens(scheduler)
150 - assert get_sequence_groups(out) == [seq_group_b]
151 -
152 -
153 -def test_scheduler_schedule_preempt_abort():
154 - block_size = 4
155 - max_model_len = 16
156 - scheduler_config = SchedulerConfig(
157 - "generate",
158 - max_num_batched_tokens=64,
159 - max_num_seqs=2,
160 - max_model_len=max_model_len,
161 - )
162 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
163 - cache_config.num_cpu_blocks = 2
164 - cache_config.num_gpu_blocks = 2
165 - scheduler = Scheduler(scheduler_config, cache_config, None)
166 -
167 - # Add seq groups to scheduler.
168 - seq_a, seq_group_a = create_dummy_prompt("1",
169 - block_size,
170 - block_size=block_size)
171 - seq_b, seq_group_b = create_dummy_prompt("2",
172 - block_size,
173 - block_size=block_size)
174 - scheduler.add_seq_group(seq_group_a)
175 - scheduler.add_seq_group(seq_group_b)
176 -
177 - # Schedule seq groups prompts.
178 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
179 - assert get_sequence_groups(out) == [seq_group_a, seq_group_b]
180 - assert out.num_batched_tokens == block_size * 2 # seq_a and seq_b
181 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
182 - and not out.blocks_to_swap_out)
183 - assert len(seq_group_meta) == 2
184 - assert scheduler.get_num_unfinished_seq_groups() == 2
185 -
186 - # Append "generated" tokens, allowing the sequence to mark prompt tokens as
187 - # processed.
188 - append_new_token(out, 1)
189 -
190 - # Schedule seq groups generation and preempt seq group b.
191 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
192 - assert get_sequence_groups(out) == [seq_group_a]
193 - assert out.num_batched_tokens == 1
194 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
195 - and not out.blocks_to_swap_out)
196 - assert len(seq_group_meta) == 1
197 - assert scheduler.get_num_unfinished_seq_groups() == 2
198 - assert out.preempted == 1
199 -
200 - # Abort seq group a. Re-schedule seq group b prompt with recomputation.
201 - scheduler.abort_seq_group("1")
202 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
203 - assert get_sequence_groups(out) == [seq_group_b]
204 - assert out.num_batched_tokens == 5 # 4 prompt + 1 generation.
205 - assert (not out.blocks_to_copy and not out.blocks_to_swap_in
206 - and not out.blocks_to_swap_out)
207 - assert len(seq_group_meta) == 1
208 - assert scheduler.get_num_unfinished_seq_groups() == 1
209 -
210 -
211 -def test_scheduler_max_seqs():
212 - block_size = 4
213 - num_seq_group = 4
214 - max_seq_group = 2
215 - max_model_len = 16
216 - scheduler_config = SchedulerConfig(
217 - "generate",
218 - max_num_batched_tokens=64,
219 - max_num_seqs=max_seq_group,
220 - max_model_len=max_model_len,
221 - )
222 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
223 - cache_config.num_cpu_blocks = 8
224 - cache_config.num_gpu_blocks = 8
225 - scheduler = Scheduler(scheduler_config, cache_config, None)
226 -
227 - all_seq_groups: list[SequenceGroup] = []
228 - # Add seq groups to scheduler.
229 - for i in range(num_seq_group):
230 - _, seq_group = create_dummy_prompt(str(i),
231 - prompt_length=block_size,
232 - block_size=block_size)
233 - all_seq_groups.append(seq_group)
234 -
235 - # Append 1 seq group
236 - scheduler.add_seq_group(all_seq_groups[0])
237 -
238 - # Schedule seq groups prompts.
239 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
240 - assert set(get_sequence_groups(out)) == set([all_seq_groups[0]])
241 - append_new_token(out, 1)
242 -
243 - # Schedule seq groups generation.
244 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
245 - assert set(get_sequence_groups(out)) == set([all_seq_groups[0]])
246 - append_new_token(out, 1)
247 -
248 - # Append 2 more seq group
249 - scheduler.add_seq_group(all_seq_groups[1])
250 - scheduler.add_seq_group(all_seq_groups[2])
251 -
252 - # Schedule seq groups prompts.
253 - # Only 1 seq group should be scheduled since max_seq_group is 2
254 - # and one is prompting.
255 - _, out = schedule_and_update_computed_tokens(scheduler)
256 - assert set(get_sequence_groups(out)) == set([all_seq_groups[1]])
257 -
258 -
259 -def test_scheduler_delay_factor():
260 - block_size = 4
261 - scheduler_config = SchedulerConfig(
262 - "generate",
263 - max_num_batched_tokens=100,
264 - max_num_seqs=64,
265 - max_model_len=16,
266 - delay_factor=0.5,
267 - )
268 - cache_config = CacheConfig(block_size, 1.0, 1, "auto")
269 - cache_config.num_cpu_blocks = 8
270 - cache_config.num_gpu_blocks = 8
271 - scheduler = Scheduler(scheduler_config, cache_config, None)
272 -
273 - # schedule first prompt
274 - seq_group_meta, seq_group = create_dummy_prompt("0",
275 - prompt_length=block_size,
276 - block_size=block_size)
277 - scheduler.add_seq_group(seq_group)
278 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
279 - assert out.num_prefill_groups > 0
280 - assert seq_group_meta[0].request_id == '0'
281 - append_new_token(out, 1)
282 -
283 - # wait for a second before scheduling next prompt
284 - time.sleep(1)
285 - seq_group_meta, seq_group = create_dummy_prompt("1",
286 - prompt_length=block_size,
287 - block_size=block_size)
288 - scheduler.add_seq_group(seq_group)
289 -
290 - # second prompt should *not* be scheduled
291 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
292 - assert out.num_prefill_groups == 0
293 - assert seq_group_meta[0].request_id == '0'
294 - append_new_token(out, 1)
295 -
296 - # wait for more than 0.5 second and try again
297 - time.sleep(0.6)
298 - seq_group_meta, out = schedule_and_update_computed_tokens(scheduler)
299 - assert out.num_prefill_groups > 0
300 - assert seq_group_meta[0].request_id == '1'
301 - append_new_token(out, 1)
302 -
303 -
304 -def initialize_scheduler(
305 - *,
306 - max_num_seqs=1000,
307 - max_token_budget=1000,
308 - max_model_len=1000,
309 - lora_config=None,
310 - block_size=4,
311 - num_cpu_blocks=8,
312 - num_gpu_blocks=8,
313 - enable_prefix_caching=False,
314 - enable_chunked_prefill=False,
315 -):
316 - block_size = block_size
317 - scheduler_config = SchedulerConfig(
318 - "generate",
319 - max_num_batched_tokens=max_token_budget,
320 - max_num_seqs=max_num_seqs,
321 - max_model_len=max_model_len,
322 - enable_chunked_prefill=enable_chunked_prefill,
323 - )
324 - cache_config = CacheConfig(
325 - block_size,
326 - 1.0,
327 - 1,
328 - "auto",
329 - enable_prefix_caching=enable_prefix_caching,
330 - )
331 - cache_config.num_cpu_blocks = num_cpu_blocks
332 - cache_config.num_gpu_blocks = num_gpu_blocks
333 - scheduler = Scheduler(scheduler_config, cache_config, lora_config)
334 - return scheduler
335 -
336 -
337 -def create_token_budget(token_budget: int = 10000,
338 - max_num_seqs: int = 10000) -> SchedulingBudget:
339 - return SchedulingBudget(
340 - token_budget=token_budget,
341 - max_num_seqs=max_num_seqs,
342 - )
343 -
344 -
345 -def add_token_budget(budget: SchedulingBudget,
346 - num_batched_tokens: int = 0,
347 - num_curr_seqs: int = 0):
348 - mock_seq_group = create_dummy_prompt('10', prompt_length=60)[1]
349 - budget.add_num_batched_tokens(mock_seq_group.request_id,
350 - num_batched_tokens)
351 - budget.add_num_seqs(mock_seq_group.request_id, num_curr_seqs)
352 -
353 -
354 -def test_prefill_schedule_max_prompt_len():
355 - """
356 - Test prompt longer than max_prompt_len is aborted.
357 - """
358 - block_size = 4
359 - scheduler = initialize_scheduler(max_model_len=30, block_size=block_size)
360 - _, seq_group = create_dummy_prompt("0",
361 - prompt_length=60,
362 - block_size=block_size)
363 - scheduler.add_seq_group(seq_group)
364 - budget = create_token_budget()
365 - output = scheduler._schedule_prefills(budget, None)
366 - remaining_waiting = scheduler.waiting
367 - assert len(output.ignored_seq_groups) == 1
368 - assert len(output.seq_groups) == 0
369 - assert budget.num_batched_tokens == 0
370 - assert budget.num_curr_seqs == 0
371 - assert len(remaining_waiting) == 0
372 -
373 -
374 -def test_prefill_schedule_token_budget():
375 - """
376 - Test token budget respected.
377 - """
378 - block_size = 4
379 - scheduler = initialize_scheduler(block_size=block_size,
380 - num_cpu_blocks=64,
381 - num_gpu_blocks=64)
382 - budget = create_token_budget(token_budget=0)
383 - for i in range(2):
384 - _, seq_group = create_dummy_prompt(str(i),
385 - prompt_length=60,
386 - block_size=block_size)
387 - scheduler.add_seq_group(seq_group)
388 -
389 - # 0 token budget == nothing is scheduled.
390 - output = scheduler._schedule_prefills(budget, None)
391 - remaining_waiting = scheduler.waiting
392 - assert len(output.ignored_seq_groups) == 0
393 - assert len(output.seq_groups) == 0
394 - assert budget.num_batched_tokens == 0
395 - assert budget.num_curr_seqs == 0
396 - assert len(remaining_waiting) == 2
397 -
398 - # 60 token budget == 1 request scheduled.
399 - budget = create_token_budget(token_budget=60)
400 - output = scheduler._schedule_prefills(budget, None)
401 - remaining_waiting = scheduler.waiting
402 - assert len(output.ignored_seq_groups) == 0
403 - assert len(output.seq_groups) == 1
404 - assert budget.num_batched_tokens == 60
405 - assert budget.num_curr_seqs == 1
406 - assert len(remaining_waiting) == 1
407 -
408 - # Test when current_batched_tokens respected.
409 - scheduler = initialize_scheduler(block_size=block_size,
410 - num_cpu_blocks=16,
411 - num_gpu_blocks=16)
412 - budget = create_token_budget(token_budget=60)
413 - add_token_budget(budget, 30, 0)
414 - _, seq_group = create_dummy_prompt(str(i),
415 - prompt_length=60,
416 - block_size=block_size)
417 - # Cannot schedule a prompt that doesn't fit the budget.
418 - scheduler.add_seq_group(seq_group)
419 - output = scheduler._schedule_prefills(budget, None)
420 - remaining_waiting = scheduler.waiting
421 - assert len(output.ignored_seq_groups) == 0
422 - assert len(output.seq_groups) == 0
423 - assert budget.num_batched_tokens == 30
424 - assert budget.num_curr_seqs == 0
425 - assert len(remaining_waiting) == 1
426 - budget = create_token_budget(token_budget=90)
427 - add_token_budget(budget, 30, 0)
428 - output = scheduler._schedule_prefills(budget, None)
429 - remaining_waiting = scheduler.waiting
430 - assert len(output.seq_groups) == 1
431 - assert budget.num_batched_tokens == 90
432 - assert budget.num_curr_seqs == 1
433 - assert len(remaining_waiting) == 0
434 -
435 -
436 -def test_prefill_schedule_max_seqs():
437 - """
438 - Test max seq respected.
439 - """
440 - block_size = 4
441 - scheduler = initialize_scheduler(block_size=block_size,
442 - num_cpu_blocks=64,
443 - num_gpu_blocks=64)
444 - budget = create_token_budget(max_num_seqs=2)
445 - for i in range(3):
446 - _, seq_group = create_dummy_prompt(str(i),
447 - prompt_length=60,
448 - block_size=block_size)
449 - scheduler.add_seq_group(seq_group)
450 - output = scheduler._schedule_prefills(budget, None)
451 - remaining_waiting = scheduler.waiting
452 - assert len(output.ignored_seq_groups) == 0
453 - assert len(output.seq_groups) == 2
454 - assert budget.num_batched_tokens == 120
455 - assert budget.num_curr_seqs == 2
456 - assert len(remaining_waiting) == 1
457 -
458 - # Verify curr_num_seqs respected.
459 - scheduler.waiting = deque()
460 - budget = create_token_budget(max_num_seqs=2)
461 - add_token_budget(budget, 0, 2)
462 - _, seq_group = create_dummy_prompt(str(i),
463 - prompt_length=60,
464 - block_size=block_size)
465 - scheduler.add_seq_group(seq_group)
466 - output = scheduler._schedule_prefills(budget, None)
467 - remaining_waiting = scheduler.waiting
468 - assert len(output.ignored_seq_groups) == 0
469 - assert len(output.seq_groups) == 0
470 - assert budget.num_batched_tokens == 0
471 - assert budget.num_curr_seqs == 2
472 - assert len(remaining_waiting) == 1
473 -
474 -
475 -def test_prefill_schedule_max_lora():
476 - """
477 - Test max lora is respected and prioritized.
478 - """
479 - block_size = 4
480 - lora_config = LoRAConfig(max_lora_rank=8, max_loras=1)
481 - scheduler = initialize_scheduler(lora_config=lora_config,
482 - block_size=block_size,
483 - num_cpu_blocks=64,
484 - num_gpu_blocks=64)
485 - budget = create_token_budget(token_budget=120)
486 - curr_loras: set[int] = set()
487 - for i in range(2):
488 - _, seq_group = create_dummy_prompt(str(i),
489 - prompt_length=60,
490 - block_size=block_size,
491 - lora_request=LoRARequest(
492 - lora_name=str(i),
493 - lora_int_id=i + 1,
494 - lora_path="abc"))
495 - scheduler.add_seq_group(seq_group)
496 - # Add two more requests to verify lora is prioritized.
497 - # 0: LoRA, 1: LoRA, 2: regular, 3: regular
498 - # In the first iteration, index 0, 2 is scheduled.
499 - # If a request is not scheduled because it hits max lora, it is
500 - # prioritized. Verify that.
501 - for i in range(2, 4):
502 - _, seq_group = create_dummy_prompt(str(i),
503 - prompt_length=60,
504 - block_size=block_size)
505 - scheduler.add_seq_group(seq_group)
506 - # Schedule 2 requests (0 and 2)
507 - output = scheduler._schedule_prefills(budget, curr_loras)
508 - remaining_waiting = scheduler.waiting
509 - assert len(output.ignored_seq_groups) == 0
510 - assert len(output.seq_groups) == 2
511 - assert budget.num_batched_tokens == 120
512 - assert budget.num_curr_seqs == 2
513 - assert len(remaining_waiting) == 2
514 - assert len(curr_loras) == 1
515 - # The second lora request is scheduled next as FCFS policy.
516 - # Reset curr_loras so that it can be scheduled.
517 - curr_loras = set()
518 - budget = create_token_budget(token_budget=60)
519 - output = scheduler._schedule_prefills(budget, curr_loras)
520 - remaining_waiting = scheduler.waiting
521 - assert len(output.seq_groups) == 1
522 - assert output.seq_groups[0].seq_group.request_id == "1"
523 - assert len(remaining_waiting) == 1
524 - assert len(curr_loras) == 1
525 - assert budget.num_batched_tokens == 60
526 -
527 -
528 -def test_prefill_schedule_no_block_manager_capacity():
529 - """
530 - Test sequence cannot be scheduled due to block manager has no capacity.
531 - """
532 - block_size = 4
533 - scheduler = initialize_scheduler(block_size=block_size,
534 - num_gpu_blocks=128,
535 - num_cpu_blocks=128)
536 - budget = create_token_budget()
537 - for i in range(3):
538 - _, seq_group = create_dummy_prompt(str(i),
539 - prompt_length=60,
540 - block_size=block_size)
541 - scheduler.add_seq_group(seq_group)
542 - scheduler.block_manager.can_allocate = MagicMock()
543 - scheduler.block_manager.can_allocate.return_value = AllocStatus.LATER
544 - output = scheduler._schedule_prefills(budget, None)
545 - remaining_waiting = scheduler.waiting
546 - assert len(output.ignored_seq_groups) == 0
547 - assert len(output.seq_groups) == 0
548 - assert budget.num_batched_tokens == 0
549 - assert budget.num_curr_seqs == 0
550 - assert len(remaining_waiting) == 3
551 -
552 - scheduler = initialize_scheduler()
553 - budget = create_token_budget()
554 - for i in range(3):
555 - _, seq_group = create_dummy_prompt(str(i),
556 - prompt_length=60,
557 - block_size=block_size)
558 - scheduler.add_seq_group(seq_group)
559 - scheduler.block_manager.can_allocate = MagicMock()
560 - scheduler.block_manager.can_allocate.return_value = AllocStatus.NEVER
561 - output = scheduler._schedule_prefills(budget, None)
562 - remaining_waiting = scheduler.waiting
563 - assert len(output.ignored_seq_groups) == 3
564 - assert len(output.seq_groups) == 0
565 - assert budget.num_batched_tokens == 0
566 - assert budget.num_curr_seqs == 0
567 - assert len(remaining_waiting) == 0
568 -
569 -
570 -def test_decode_schedule_preempted():
571 - """
572 - Test decodes cannot be scheduled and preempted.
573 - """
574 - block_size = 4
575 - scheduler = initialize_scheduler(block_size=block_size,
576 - num_cpu_blocks=64,
577 - num_gpu_blocks=64)
578 - curr_loras = None
579 - for i in range(3):
580 - _, seq_group = create_dummy_prompt(str(i),
581 - prompt_length=60,
582 - block_size=block_size)
583 - scheduler._allocate_and_set_running(seq_group)
584 - append_new_token_seq_group(60, seq_group, 1)
585 - scheduler._add_seq_group_to_running(seq_group)
586 - scheduler.block_manager.can_append_slots = MagicMock()
587 -
588 - def cannot_append_second_group(seq_group, num_lookahead_slots):
589 - return seq_group.request_id != "1"
590 -
591 - scheduler.block_manager.can_append_slots.side_effect = (
592 - cannot_append_second_group)
593 -
594 - # 1 cannot be scheduled, and the lowest priority (request 2)
595 - # should be preempted. 1 will also be preempted.
596 - budget = create_token_budget()
597 - output = scheduler._schedule_running(budget, curr_loras)
598 - remaining_running = scheduler.running
599 - assert len(remaining_running) == 0
600 - assert len(output.decode_seq_groups) == 1
601 - assert len(output.prefill_seq_groups) == 0
602 - assert output.decode_seq_groups[0].seq_group.request_id == "0"
603 - assert len(output.preempted) == 2
604 - # Verify budgets are updated.
605 - assert budget.num_batched_tokens == 1
606 - # NOTE: When enable_chunk is False, num_seqs budget is not updated.
607 - # assert budget.num_curr_seqs == 1
608 - # Both should be preempted, not swapped.
609 - assert output.blocks_to_swap_out == []
610 - # Nothing is copied.
611 - assert output.blocks_to_copy == []
612 -
613 -
614 -def test_schedule_decode_blocks_to_copy_update():
615 - """
616 - Verify blocks_to_copy is updated.
617 - """
618 - block_size = 4
619 - scheduler = initialize_scheduler(block_size=4,
620 - num_cpu_blocks=16,
621 - num_gpu_blocks=16)
622 - _, seq_group = create_dummy_prompt("1",
623 - prompt_length=60,
624 - block_size=block_size)
625 - curr_loras = None
626 - scheduler._allocate_and_set_running(seq_group)
627 - append_new_token_seq_group(60, seq_group, 1)
628 - scheduler._add_seq_group_to_running(seq_group)
629 -
630 - # The last request should be swapped out.
631 - scheduler.block_manager.append_slots = MagicMock()
632 - scheduler.block_manager.append_slots.return_value = [(2, 3)]
633 -
634 - budget = create_token_budget()
635 - output = scheduler._schedule_running(budget, curr_loras)
636 - remaining_running = scheduler.running
637 - assert len(remaining_running) == 0
638 - assert len(output.decode_seq_groups) == 1
639 - assert len(output.prefill_seq_groups) == 0
640 - assert len(output.preempted) == 0
641 - assert len(output.swapped_out) == 0
642 - # Nothing is preempted.
643 - assert output.blocks_to_swap_out == []
644 - # Since append_slot returns the source -> dist mapping, it should
645 - # be applied.
646 - assert output.blocks_to_copy == [(2, 3)]
647 -
648 -
649 -def test_schedule_swapped_max_loras():
650 - block_size = 4
651 - lora_config = LoRAConfig(max_lora_rank=8, max_loras=1)
652 - scheduler = initialize_scheduler(lora_config=lora_config,
653 - block_size=block_size,
654 - num_cpu_blocks=32,
655 - num_gpu_blocks=32)
656 - curr_loras: set[int] = set()
657 - blocks_to_swap_out: list[tuple[int, int]] = []
658 - for i in range(2):
659 - _, seq_group = create_dummy_prompt(str(i),
660 - prompt_length=60,
661 - block_size=block_size,
662 - lora_request=LoRARequest(
663 - lora_name=str(i),
664 - lora_int_id=i + 1,
665 - lora_path="abc"))
666 - scheduler._allocate_and_set_running(seq_group)
667 - append_new_token_seq_group(60, seq_group, 1)
668 - scheduler._swap_out(seq_group, blocks_to_swap_out)
669 - scheduler._add_seq_group_to_swapped(seq_group)
670 -
671 - budget = create_token_budget()
672 - output = scheduler._schedule_swapped(budget, curr_loras)
673 - remaining_swapped = scheduler.swapped
674 - assert len(remaining_swapped) == 1
675 - assert budget.num_batched_tokens == 1
676 - assert budget.num_curr_seqs == 1
677 - assert len(output.decode_seq_groups) == 1
678 - assert len(output.prefill_seq_groups) == 0
679 - assert len(curr_loras) == 1
680 -
681 -
682 -def test_schedule_swapped_cannot_swap_in():
683 - block_size = 4
684 - scheduler = initialize_scheduler(block_size=block_size,
685 - num_cpu_blocks=32,
686 - num_gpu_blocks=32)
687 - curr_loras = None
688 - blocks_to_swap_out: list[tuple[int, int]] = []
689 - for i in range(2):
690 - _, seq_group = create_dummy_prompt(str(i),
691 - prompt_length=60,
692 - block_size=block_size)
693 - scheduler._allocate_and_set_running(seq_group)
694 - append_new_token_seq_group(60, seq_group, 1)
695 - scheduler._swap_out(seq_group, blocks_to_swap_out)
696 - scheduler._add_seq_group_to_swapped(seq_group)
697 -
698 - # The last request should be swapped out.
699 - scheduler.block_manager.can_swap_in = MagicMock()
700 - scheduler.block_manager.can_swap_in.return_value = AllocStatus.LATER
701 - # Since we cannot swap in, none of the requests are swapped in.
702 - budget = create_token_budget()
703 - output = scheduler._schedule_swapped(budget, curr_loras)
704 - remaining_swapped = scheduler.swapped
705 - assert len(remaining_swapped) == 2
706 - assert budget.num_batched_tokens == 0
707 - assert budget.num_curr_seqs == 0
708 - assert len(output.decode_seq_groups) == 0
709 - assert len(output.prefill_seq_groups) == 0
710 -
711 -
712 -def test_infeasible_swap():
713 - block_size = 4
714 - scheduler = initialize_scheduler(block_size=block_size,
715 - num_cpu_blocks=32,
716 - num_gpu_blocks=32)
717 - curr_loras = None
718 - blocks_to_swap_out: list[tuple[int, int]] = []
719 - for i in range(2):
720 - _, seq_group = create_dummy_prompt(str(i),
721 - prompt_length=60,
722 - block_size=block_size)
723 - scheduler._allocate_and_set_running(seq_group)
724 - append_new_token_seq_group(60, seq_group, 1)
725 - scheduler._swap_out(seq_group, blocks_to_swap_out)
726 - scheduler._add_seq_group_to_swapped(seq_group)
727 -
728 - # The last request should be swapped out.
729 - scheduler.block_manager.can_swap_in = MagicMock()
730 - scheduler.block_manager.can_swap_in.return_value = AllocStatus.NEVER
731 - # Since we cannot swap in, none of the requests are swapped in.
732 - budget = create_token_budget()
733 - output = scheduler._schedule_swapped(budget, curr_loras)
734 - remaining_swapped = scheduler.swapped
735 - assert len(remaining_swapped) == 0
736 - assert len(output.infeasible_seq_groups) == 2
737 - assert budget.num_batched_tokens == 0
738 - assert budget.num_curr_seqs == 0
739 - assert len(output.decode_seq_groups) == 0
740 - assert len(output.prefill_seq_groups) == 0
741 -
742 -
743 -def test_schedule_swapped_blocks_to_copy():
744 - block_size = 4
745 - scheduler = initialize_scheduler(block_size=block_size,
746 - num_cpu_blocks=32,
747 - num_gpu_blocks=32)
748 - curr_loras = None
749 - _, seq_group = create_dummy_prompt("1",
750 - prompt_length=60,
751 - block_size=block_size)
752 - scheduler._allocate_and_set_running(seq_group)
753 - append_new_token_seq_group(60, seq_group, 1)
754 - blocks_to_swap_out: list[tuple[int, int]] = []
755 - scheduler._swap_out(seq_group, blocks_to_swap_out)
756 - scheduler._add_seq_group_to_swapped(seq_group)
757 -
758 - # The last request should be swapped out.
759 - scheduler.block_manager.append_slots = MagicMock()
760 - scheduler.block_manager.append_slots.return_value = [(2, 3)]
761 -
762 - budget = create_token_budget()
763 - output = scheduler._schedule_swapped(budget, curr_loras)
764 - remaining_swapped = scheduler.swapped
765 - assert len(remaining_swapped) == 0
766 - assert len(output.decode_seq_groups) == 1
767 - assert len(output.prefill_seq_groups) == 0
768 - assert output.blocks_to_copy == [(2, 3)]
769 -
770 -
771 -def test_scheduling_budget():
772 - TOKEN_BUDGET = 4
773 - MAX_SEQS = 4
774 - budget = SchedulingBudget(token_budget=TOKEN_BUDGET, max_num_seqs=MAX_SEQS)
775 - assert budget.can_schedule(num_new_tokens=1, num_new_seqs=1)
776 - assert budget.can_schedule(num_new_tokens=4, num_new_seqs=4)
777 - assert not budget.can_schedule(num_new_tokens=1, num_new_seqs=5)
778 - assert not budget.can_schedule(num_new_tokens=5, num_new_seqs=1)
779 - assert not budget.can_schedule(num_new_tokens=5, num_new_seqs=5)
780 - assert budget.remaining_token_budget() == TOKEN_BUDGET
781 -
782 - # Verify add/subtract num batched tokens.
783 - _, seq_group = create_dummy_prompt("1", 3)
784 - budget.add_num_batched_tokens(seq_group.request_id, 2)
785 - assert budget.remaining_token_budget() == 2
786 - assert budget.num_batched_tokens == 2
787 - assert budget.can_schedule(num_new_tokens=2, num_new_seqs=1)
788 - assert not budget.can_schedule(num_new_tokens=3, num_new_seqs=1)
789 - # Verify adding another seq group is no-op.
790 - budget.add_num_batched_tokens(seq_group.request_id, 2)
791 - assert budget.remaining_token_budget() == 2
792 - assert budget.num_batched_tokens == 2
793 - budget.subtract_num_batched_tokens(seq_group.request_id, 2)
794 - assert budget.remaining_token_budget() == 4
795 - assert budget.num_batched_tokens == 0
796 - budget.subtract_num_batched_tokens(seq_group.request_id, 2)
797 - assert budget.remaining_token_budget() == 4
798 - assert budget.num_batched_tokens == 0
799 -
800 - # Verify add/subtract max seqs.
801 - _, seq_group = create_dummy_prompt("1", 3)
802 - budget.add_num_seqs(seq_group.request_id, 2)
803 - assert budget.can_schedule(num_new_tokens=1, num_new_seqs=2)
804 - assert not budget.can_schedule(num_new_tokens=1, num_new_seqs=3)
805 - assert budget.num_curr_seqs == 2
806 - # Verify adding another seq group is no-op.
807 - budget.add_num_seqs(seq_group.request_id, 2)
808 - assert budget.num_curr_seqs == 2
809 - budget.subtract_num_seqs(seq_group.request_id, 2)
810 - assert budget.num_curr_seqs == 0
811 - budget.subtract_num_seqs(seq_group.request_id, 2)
812 - assert budget.num_curr_seqs == 0
813 -
814 -
815 -@pytest.mark.parametrize("enable_prefix_caching", [True, False])
816 -def test_prefix_caching_aware_prefills(enable_prefix_caching):
817 - """
818 - Test the below scenario:
819 -
820 - For 3 sequences, seqA, seqB, seqC, share the first block as prefix.
821 -
822 - The test verifies the below scenarios:
823 - 1. SeqA is first scheduled.
824 - 2. SeqB and SeqC can be prefilled together in a single schedule round
825 - even though there are not enough token budgets to prefill both without
826 - considering prefix caching.
827 - """
828 -
829 - block_size = 4
830 - max_num_batched_tokens = 12
831 - max_seq_group = 3
832 - scheduler = initialize_scheduler(
833 - block_size=block_size,
834 - num_cpu_blocks=16,
835 - num_gpu_blocks=16,
836 - max_token_budget=max_num_batched_tokens,
837 - max_num_seqs=max_seq_group,
838 - max_model_len=max_num_batched_tokens,
839 - enable_prefix_caching=enable_prefix_caching,
840 - )
841 -
842 - seqA_tokens = list(range(8))
843 - num_shared_tokens = 4
844 - seqB_tokens = seqA_tokens[:num_shared_tokens] + list(range(
845 - 12, 16)) # Shared prefix first 4.
846 - seqC_tokens = seqA_tokens[:num_shared_tokens] + list(range(
847 - 16, 20)) # Shared prefix first 4.
848 -
849 - seqA, seqA_group = create_dummy_prompt("0",
850 - prompt_tokens=seqA_tokens,
851 - block_size=block_size)
852 - seqB, seqB_group = create_dummy_prompt("1",
853 - prompt_tokens=seqB_tokens,
854 - block_size=block_size)
855 - seqC, seqC_group = create_dummy_prompt("2",
856 - prompt_tokens=seqC_tokens,
857 - block_size=block_size)
858 -
859 - # Schedule seqA prefill.
860 - scheduler.add_seq_group(seqA_group)
861 - metas, out, _ = scheduler.schedule()
862 - assert (len(out.scheduled_seq_groups) == 1
863 - and out.scheduled_seq_groups[0].seq_group == seqA_group)
864 - assert out.scheduled_seq_groups[0].token_chunk_size == len(seqA_tokens)
865 -
866 - # Schedule seqA decode.
867 - append_new_token_seq_group(len(seqA_tokens), seqA_group, 999)
868 - metas, out, _ = scheduler.schedule()
869 -
870 - assert len(out.scheduled_seq_groups) == 1
871 - assert out.scheduled_seq_groups[0].seq_group == seqA_group
872 - assert out.scheduled_seq_groups[0].token_chunk_size == 1
873 -
874 - # Schedule seqB and seqC prefills should work with prefix caching.
875 - scheduler.add_seq_group(seqB_group)
876 - scheduler.add_seq_group(seqC_group)
877 - metas, out, _ = scheduler.schedule()
878 -
879 - if enable_prefix_caching:
880 - assert len(out.scheduled_seq_groups) == 2
881 - assert set([
882 - out.scheduled_seq_groups[0].seq_group,
883 - out.scheduled_seq_groups[1].seq_group,
884 - ]) == set([seqB_group, seqC_group])
885 - assert len(metas) == 2
886 - for meta in metas:
887 - assert meta.token_chunk_size == 8
888 - assert (len(meta.computed_block_nums) == num_shared_tokens //
889 - block_size) # 1 Block for the 8 tokens.
890 - else:
891 - assert len(out.scheduled_seq_groups) == 1
892 - assert len(metas) == 1
893 - assert metas[0].token_chunk_size == 8
894 - assert len(metas[0].computed_block_nums) == 0 # No blocks computed.
895 -
896 -
897 -def test_no_multiple_partial_prefills_with_chunked_prefill_and_prefix_caching(
898 -):
899 - """
900 - This test verifies that we don't schedule new prefills if there's already
901 - a continuous prefill in progress even though the new prefills with shared
902 - prefix can fit in the token budget:
903 -
904 - - SeqA is being chunked prefill.
905 - - SeqB with the same prompt shouldn't be scheduled for prefill even though
906 - there's enough token budget to prefill the cached tokens.
907 - - Neither should seqC be scheduled.
908 -
909 - - When seqA is in decoding phase, seqB and seqC can be scheduled.
910 - - Entire seqB should be prefilled since it's a full prefix cache hit.
911 - - SeqC would be partially prefilled with the prefix shared, and the
912 - remaining unique tokens would be prefilled (rounded down to be
913 - block-size aligned).
914 - """
915 -
916 - block_size = 2
917 - max_num_batched_tokens = 4
918 - max_seq_group = 3
919 - scheduler = initialize_scheduler(
920 - block_size=block_size,
921 - num_cpu_blocks=16,
922 - num_gpu_blocks=16,
923 - max_token_budget=max_num_batched_tokens,
924 - max_num_seqs=max_seq_group,
925 - max_model_len=100,
926 - enable_prefix_caching=True,
927 - enable_chunked_prefill=True,
928 - )
929 -
930 - seqA_tokens = list(range(8))
931 - seqB_tokens = seqA_tokens
932 - seqC_shared_prefix_len = 4
933 - seqC_tokens = seqA_tokens[:seqC_shared_prefix_len] + list(range(12, 20))
934 -
935 - seqA, seqA_group = create_dummy_prompt("0",
936 - prompt_tokens=seqA_tokens,
937 - block_size=block_size)
938 - seqB, seqB_group = create_dummy_prompt("1",
939 - prompt_tokens=seqB_tokens,
940 - block_size=block_size)
941 -
942 - # Chunked prefill seqA.
943 - scheduler.add_seq_group(seqA_group)
944 - metas, out = schedule_and_update_computed_tokens(scheduler)
945 - assert len(out.scheduled_seq_groups) == 1
946 - assert out.scheduled_seq_groups[0].seq_group == seqA_group
947 - assert out.scheduled_seq_groups[0].token_chunk_size == 4
948 -
949 - # seqB should not be scheduled with ongoing prefills.
950 - scheduler.add_seq_group(seqB_group)
951 - metas, out = schedule_and_update_computed_tokens(scheduler)
952 - assert len(out.scheduled_seq_groups) == 1
953 - assert out.scheduled_seq_groups[0].seq_group == seqA_group
954 - assert out.scheduled_seq_groups[0].token_chunk_size == 4
955 -
956 - # both seqB and seqC can now be scheduled with seqA is over.
957 - # seqA is in decoding phase.
958 - append_new_token_seq(seqA, 999)
959 - seqC, seqC_group = create_dummy_prompt("2",
960 - prompt_tokens=seqC_tokens,
961 - block_size=block_size)
962 - scheduler.add_seq_group(seqC_group)
963 - metas, out = schedule_and_update_computed_tokens(scheduler)
964 - assert len(out.scheduled_seq_groups) == 3
965 -
966 - metas = {meta.request_id: meta for meta in metas}
967 - assert metas[seqA_group.request_id].token_chunk_size == 1 # Decode
968 - assert (metas[seqB_group.request_id].token_chunk_size == 8
969 - ) # Fully cached prefill
970 - assert (
971 - metas[seqC_group.request_id].token_chunk_size == 6
972 - ), "A partial prefix of C (4 tokens) should be prefilled, with the "
973 - "remaining tokens fit into 3 token budget (4-1 from the seqA). It will "
974 - "then be rounded down to 2 tokens on block size, thus 6 tokens in total."
975 -
976 -
977 -def test_no_batches_mixed_with_prompt_tokens_and_prompt_embeds():
978 - """
979 - Test that the scheduler does not schedule batches with prompt tokens and
980 - prompt embeddings co-mingled.
981 - """
982 - block_size = 2
983 - max_seq_group = 3
984 - scheduler = initialize_scheduler(
985 - block_size=block_size,
986 - num_cpu_blocks=16,
987 - num_gpu_blocks=16,
988 - max_num_seqs=max_seq_group,
989 - max_model_len=100,
990 - enable_prefix_caching=True,
991 - )
992 -
993 - # the odd indexed inputs should be passed in via embeddings,
994 - # evens via token_ids
995 - seq_length = 7
996 - embedding_size = 5
997 - num_seqs = 11
998 - seq_tokens: list[list[int]] = []
999 - seq_embeds: list[Optional[torch.Tensor]] = []
1000 - for i in range(num_seqs):
1001 - if i % 2:
1002 - seq_tokens.append(list(range(seq_length)))
1003 - seq_embeds.append(None)
1004 - else:
1005 - seq_tokens.append([0] * seq_length)
1006 - seq_embeds.append(torch.rand(embedding_size))
1007 -
1008 - seq_and_seq_groups = [
1009 - create_dummy_prompt(f"{i}",
1010 - prompt_tokens=seq_tokens[i],
1011 - prompt_embeds=seq_embeds[i],
1012 - block_size=block_size)
1013 - for i in range(len(seq_tokens))
1014 - ]
1015 -
1016 - for _, seq_group in seq_and_seq_groups:
1017 - scheduler.add_seq_group(seq_group)
1018 -
1019 - while not all(seq.is_finished() for seq, _ in seq_and_seq_groups):
1020 - unfinished_seq_groups = [
1021 - seq_group for _, seq_group in seq_and_seq_groups
1022 - if not seq_group.is_finished()
1023 - ]
1024 - _, out = schedule_and_update_computed_tokens(scheduler)
1025 - assert len(out.scheduled_seq_groups) > 0
1026 - batch_is_prompt_embeds = out.scheduled_seq_groups[
1027 - 0].seq_group.uses_prompt_embeds()
1028 - expected_scheduled_seq_groups = [
1029 - seq_group for seq_group in unfinished_seq_groups
1030 - if seq_group.uses_prompt_embeds() == batch_is_prompt_embeds
1031 - ]
1032 -
1033 - # We should have as many scheduled groups as possible, without mixing
1034 - assert len(out.scheduled_seq_groups) == min(
1035 - max_seq_group, len(expected_scheduled_seq_groups))
1036 - assert all(scheduled_seq_group.seq_group.uses_prompt_embeds() ==
1037 - batch_is_prompt_embeds
1038 - for scheduled_seq_group in out.scheduled_seq_groups)
1039 -
1040 - # Finish the scheduled groups
1041 - for scheduled_seq_group in out.scheduled_seq_groups:
1042 - for seq in scheduled_seq_group.seq_group.seqs:
1043 - seq.status = SequenceStatus.FINISHED_STOPPED
1044 - scheduler.free_finished_seq_groups()
1045 -
1046 -
1047 -def test_remove_seq_from_computed_blocks_tracker():
1048 - """
1049 - Test that computed_blocks_tracker correctly removes stale sequences
1050 - during scheduling.
1051 -
1052 - The test covers 9 scheduling branches where stale seqs are removed:
1053 - - 1 in _schedule_swapped
1054 - - 1 in _schedule_priority_preemption
1055 - - 7 in _schedule_prefill
1056 -
1057 - Each branch is tested to ensure proper cleanup of
1058 - _seq_id_to_num_tokens_computed.
1059 - """
1060 - # Budget can not schedule in swapped
1061 - block_size = 2
1062 - max_seq_group = 3
1063 - seq_tokens_with_swapped: list[list[int]] = []
1064 - blocks_to_swap_out: list[tuple[int, int]] = []
1065 - curr_loras: set[int] = set()
1066 -
1067 - scheduler = initialize_scheduler(
1068 - block_size=block_size,
1069 - num_cpu_blocks=64,
1070 - num_gpu_blocks=16,
1071 - max_num_seqs=max_seq_group,
1072 - enable_prefix_caching=True,
1073 - )
1074 - budget = create_token_budget(token_budget=15)
1075 -
1076 - seq_length = 16
1077 - num_seqs = 3
1078 - for i in range(num_seqs):
1079 - seq_tokens_with_swapped.append([i] * seq_length)
1080 -
1081 - seq_and_seq_groups = [
1082 - create_dummy_prompt(f"{i}",
1083 - prompt_tokens=seq_tokens_with_swapped[i],
1084 - block_size=block_size)
1085 - for i in range(len(seq_tokens_with_swapped))
1086 - ]
1087 -
1088 - for _, seq_group in seq_and_seq_groups:
1089 - scheduler._allocate_and_set_running(seq_group)
1090 - scheduler._swap_out(seq_group, blocks_to_swap_out)
1091 - scheduler._add_seq_group_to_swapped(seq_group)
1092 -
1093 - scheduler._schedule_swapped(budget, curr_loras)
1094 - seq_id_to_num_tokens_computed = (
1095 - scheduler.block_manager._computed_blocks_tracker.
1096 - _seq_id_to_num_tokens_computed.get(1))
1097 - assert seq_id_to_num_tokens_computed is None
1098 -
1099 - # Prefill schedule don't have a space for another LoRA, so
1100 - # we ignore this request for now.
1101 - block_size = 4
1102 - lora_config = LoRAConfig(max_lora_rank=8, max_loras=1)
1103 - scheduler = initialize_scheduler(lora_config=lora_config,
1104 - block_size=block_size,
1105 - num_cpu_blocks=64,
1106 - num_gpu_blocks=64,
1107 - enable_prefix_caching=True)
1108 - budget = create_token_budget(token_budget=120)
1109 - num_seqs = 2
1110 - for i in range(num_seqs):
1111 - _, seq_group = create_dummy_prompt(str(i),
1112 - prompt_length=seq_length,
1113 - block_size=block_size,
1114 - lora_request=LoRARequest(
1115 - lora_name=str(i),
1116 - lora_int_id=i + 1,
1117 - lora_path="abc"))
1118 - scheduler.add_seq_group(seq_group)
1119 -
1120 - scheduler._schedule_prefills(budget, curr_loras)
1121 - seq_id_to_num_tokens_computed = (
1122 - scheduler.block_manager._computed_blocks_tracker.
1123 - _seq_id_to_num_tokens_computed.get(1))
1124 - assert seq_id_to_num_tokens_computed is None
1125 -
1126 - # Priority preemption schedule
1127 - scheduler._schedule_priority_preemption(budget)
1128 - seq_id_to_num_tokens_computed = (
1129 - scheduler.block_manager._computed_blocks_tracker.
1130 - _seq_id_to_num_tokens_computed.get(1))
1131 - assert seq_id_to_num_tokens_computed is None
1132 -
1133 - # Prefill scheduler does not schedule batches with prompt tokens and
1134 - # prompt embeddings co-mingled.
1135 - block_size = 2
1136 - max_seq_group = 3
1137 - scheduler = initialize_scheduler(
1138 - block_size=block_size,
1139 - num_cpu_blocks=16,
1140 - num_gpu_blocks=16,
1141 - max_num_seqs=max_seq_group,
1142 - max_model_len=100,
1143 - enable_prefix_caching=True,
1144 - )
1145 - seq_length = 7
1146 - embedding_size = 5
1147 - seq_tokens_with_embedding: list[list[int]] = []
1148 - seq_embeds: list[Optional[torch.Tensor]] = []
1149 -
1150 - seq_tokens_with_embedding.append(list(range(seq_length)))
1151 - seq_embeds.append(None)
1152 - seq_tokens_with_embedding.append([0] * seq_length)
1153 - seq_embeds.append(torch.rand(embedding_size))
1154 -
1155 - seq_and_seq_groups = [
1156 - create_dummy_prompt(f"{i}",
1157 - prompt_tokens=seq_tokens_with_embedding[i],
1158 - prompt_embeds=seq_embeds[i],
1159 - block_size=block_size)
1160 - for i in range(len(seq_tokens_with_embedding))
1161 - ]
1162 -
1163 - for _, seq_group in seq_and_seq_groups:
1164 - scheduler.add_seq_group(seq_group)
1165 -
1166 - scheduler._schedule_default()
1167 - seq_id_to_num_tokens_computed = (
1168 - scheduler.block_manager._computed_blocks_tracker.
1169 - _seq_id_to_num_tokens_computed.get(1))
1170 - assert seq_id_to_num_tokens_computed is None
1171 -
1172 - # Prefill scheduler budget num_batched_tokens
1173 - # >= scheduler_config max_num_batched_tokens
1174 - block_size = 2
1175 - max_seq_group = 3
1176 - seq_tokens_prefill_budget: list[list[int]] = []
1177 -
1178 - scheduler = initialize_scheduler(
1179 - block_size=block_size,
1180 - max_token_budget=8,
1181 - num_cpu_blocks=16,
1182 - num_gpu_blocks=16,
1183 - max_num_seqs=max_seq_group,
1184 - max_model_len=5,
1185 - enable_prefix_caching=True,
1186 - )
1187 - seq_length = 4
1188 - num_seqs = 3
1189 - for i in range(num_seqs):
1190 - seq_tokens_prefill_budget.append([i] * seq_length)
1191 -
1192 - seq_and_seq_groups = [
1193 - create_dummy_prompt(f"{i}",
1194 - prompt_tokens=seq_tokens_prefill_budget[i],
1195 - block_size=block_size)
1196 - for i in range(len(seq_tokens_prefill_budget))
1197 - ]
1198 -
1199 - for _, seq_group in seq_and_seq_groups:
1200 - scheduler.add_seq_group(seq_group)
1201 -
1202 - scheduler._schedule_default()
1203 - seq_id_to_num_tokens_computed = (
1204 - scheduler.block_manager._computed_blocks_tracker.
1205 - _seq_id_to_num_tokens_computed.get(2))
1206 - assert seq_id_to_num_tokens_computed is None
1207 -
1208 - # Budget can not schedule in waiting
1209 - block_size = 2
1210 - max_seq_group = 3
1211 -
1212 - scheduler = initialize_scheduler(
1213 - block_size=block_size,
1214 - max_token_budget=30,
1215 - num_cpu_blocks=16,
1216 - num_gpu_blocks=16,
1217 - max_num_seqs=max_seq_group,
1218 - max_model_len=30,
1219 - enable_prefix_caching=True,
1220 - )
1221 - seq_length = 16
1222 - num_seqs = 3
1223 - seq_tokens_prefill_budget_waiting: list[list[int]] = []
1224 -
1225 - for i in range(num_seqs):
1226 - seq_tokens_prefill_budget_waiting.append(list(range(seq_length)))
1227 -
1228 - seq_and_seq_groups = [
1229 - create_dummy_prompt(f"{i}",
1230 - prompt_tokens=seq_tokens_prefill_budget_waiting[i],
1231 - block_size=block_size)
1232 - for i in range(len(seq_tokens_prefill_budget_waiting))
1233 - ]
1234 -
1235 - for _, seq_group in seq_and_seq_groups:
1236 - scheduler.add_seq_group(seq_group)
1237 -
1238 - scheduler._schedule_default()
1239 - seq_id_to_num_tokens_computed = (
1240 - scheduler.block_manager._computed_blocks_tracker.
1241 - _seq_id_to_num_tokens_computed.get(1))
1242 - assert seq_id_to_num_tokens_computed is None
1243 -
1244 - # Sequence num_new_tokens > prompt_limit marked FINISHED_IGNORED
1245 - block_size = 2
1246 - max_seq_group = 3
1247 - scheduler = initialize_scheduler(
1248 - block_size=block_size,
1249 - num_cpu_blocks=16,
1250 - num_gpu_blocks=16,
1251 - max_num_seqs=max_seq_group,
1252 - max_model_len=30,
1253 - enable_prefix_caching=True,
1254 - )
1255 -
1256 - seq_length = 31
1257 - seq_tokens_prompt_limit: list[list[int]] = []
1258 - seq_tokens_prompt_limit.append(list(range(seq_length)))
1259 - seq_and_seq_groups = [
1260 - create_dummy_prompt("0",
1261 - prompt_tokens=seq_tokens_prompt_limit[0],
1262 - block_size=block_size)
1263 - ]
1264 - for _, seq_group in seq_and_seq_groups:
1265 - scheduler.add_seq_group(seq_group)
1266 - scheduler._schedule_default()
1267 - seq_id_to_num_tokens_computed = (
1268 - scheduler.block_manager._computed_blocks_tracker.
1269 - _seq_id_to_num_tokens_computed.get(0))
1270 - assert seq_id_to_num_tokens_computed is None
1271 -
1272 - # Budget can not allocate, AllocStatus is NEVER marked FINISHED_IGNORED
1273 - block_size = 2
1274 - max_seq_group = 3
1275 - scheduler = initialize_scheduler(
1276 - block_size=block_size,
1277 - num_cpu_blocks=160,
1278 - num_gpu_blocks=160,
1279 - max_num_seqs=max_seq_group,
1280 - max_model_len=320,
1281 - enable_prefix_caching=True,
1282 - )
1283 -
1284 - seq_length = 320
1285 - num_seqs = 1
1286 - seq_tokens_never: list[list[int]] = []
1287 - for i in range(num_seqs):
1288 - seq_tokens_never.append(list(range(seq_length)))
1289 -
1290 - seq_and_seq_groups = [
1291 - create_dummy_prompt(f"{i}",
1292 - prompt_tokens=seq_tokens_never[i],
1293 - block_size=block_size)
1294 - for i in range(len(seq_tokens_never))
1295 - ]
1296 -
1297 - for _, seq_group in seq_and_seq_groups:
1298 - scheduler.add_seq_group(seq_group)
1299 -
1300 - scheduler._schedule_default()
1301 - seq_id_to_num_tokens_computed = (
1302 - scheduler.block_manager._computed_blocks_tracker.
1303 - _seq_id_to_num_tokens_computed.get(0))
1304 - assert seq_id_to_num_tokens_computed is None
1305 -
1306 - # Budget can not allocate, AllocStatus is LATER
1307 - block_size = 2
1308 - max_seq_group = 3
1309 - scheduler = initialize_scheduler(
1310 - block_size=block_size,
1311 - num_cpu_blocks=160,
1312 - num_gpu_blocks=160,
1313 - max_num_seqs=max_seq_group,
1314 - max_model_len=320,
1315 - enable_prefix_caching=True,
1316 - )
1317 -
1318 - seq_length = 160
1319 - num_seqs = 2
1320 - seq_tokens_later: list[list[int]] = []
1321 - for i in range(num_seqs):
1322 - seq_tokens_later.append(list(range(seq_length)))
1323 -
1324 - seq_and_seq_groups = [
1325 - create_dummy_prompt(f"{i}",
1326 - prompt_tokens=seq_tokens_later[i],
1327 - block_size=block_size)
1328 - for i in range(len(seq_tokens_later))
1329 - ]
1330 -
1331 - for _, seq_group in seq_and_seq_groups:
1332 - scheduler.add_seq_group(seq_group)
1333 -
1334 - scheduler._schedule_default()
1335 - seq_id_to_num_tokens_computed = (
1336 - scheduler.block_manager._computed_blocks_tracker.
1337 - _seq_id_to_num_tokens_computed.get(1))
1338 - assert seq_id_to_num_tokens_computed is None
tests/core/test_serialization.py deleted
-36
@@ -1,36 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import msgspec
5 -
6 -from vllm.executor.msgspec_utils import decode_hook, encode_hook
7 -from vllm.sequence import ExecuteModelRequest
8 -
9 -from .utils import create_batch
10 -
11 -
12 -def test_msgspec_serialization():
13 - num_lookahead_slots = 4
14 - seq_group_metadata_list, _, _ = create_batch(16, num_lookahead_slots)
15 - execute_model_req = ExecuteModelRequest(
16 - seq_group_metadata_list=seq_group_metadata_list,
17 - num_lookahead_slots=num_lookahead_slots,
18 - running_queue_size=4)
19 -
20 - encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
21 - decoder = msgspec.msgpack.Decoder(ExecuteModelRequest,
22 - dec_hook=decode_hook)
23 - req = decoder.decode(encoder.encode(execute_model_req))
24 - expected = execute_model_req.seq_group_metadata_list
25 - actual = req.seq_group_metadata_list
26 - assert (len(expected) == len(actual))
27 - expected = expected[0]
28 - actual = actual[0]
29 -
30 - assert expected.block_tables == actual.block_tables
31 - assert expected.is_prompt == actual.is_prompt
32 - assert expected.request_id == actual.request_id
33 - assert (expected.seq_data[0].prompt_token_ids ==
34 - actual.seq_data[0].prompt_token_ids)
35 - assert (expected.seq_data[0].output_token_ids ==
36 - actual.seq_data[0].output_token_ids)
tests/core/utils.py deleted
-392
@@ -1,392 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import time
5 -from collections import defaultdict
6 -from collections.abc import Sequence as GenericSequence
7 -from itertools import count
8 -from typing import Any, Optional, Union
9 -
10 -import torch
11 -
12 -from vllm.core.scheduler import Scheduler, SchedulerOutputs
13 -from vllm.inputs import EncoderDecoderInputs, embeds_inputs, token_inputs
14 -from vllm.lora.request import LoRARequest
15 -from vllm.sampling_params import SamplingParams
16 -from vllm.sequence import (Logprob, Sequence, SequenceData, SequenceGroup,
17 - SequenceGroupMetadata)
18 -
19 -
20 -def create_dummy_prompt(
21 - request_id: str,
22 - prompt_length: int = -1,
23 - block_size: Optional[int] = None,
24 - lora_request: Optional[LoRARequest] = None,
25 - prompt_tokens: Optional[list[int]] = None,
26 - prompt_embeds: Optional[torch.Tensor] = None,
27 - min_tokens: int = 0,
28 - max_tokens: int = 16,
29 -) -> tuple[Sequence, SequenceGroup]:
30 - if not block_size:
31 - block_size = prompt_length
32 -
33 - if prompt_tokens is None:
34 - # Create dummy prompt sequence with tokens 0...block_size-1
35 - # and prompt "0 ... block_size".
36 - prompt_tokens = list(range(prompt_length))
37 -
38 - prompt_str = " ".join([str(t) for t in prompt_tokens])
39 - inputs = token_inputs(
40 - prompt_token_ids=prompt_tokens,
41 - prompt=prompt_str) if prompt_embeds is None else embeds_inputs(
42 - prompt_embeds=prompt_embeds)
43 - prompt = Sequence(
44 - int(request_id),
45 - inputs=inputs,
46 - block_size=block_size,
47 - )
48 - seq_group = SequenceGroup(
49 - request_id=request_id,
50 - seqs=[prompt],
51 - arrival_time=time.time(),
52 - sampling_params=SamplingParams(max_tokens=max_tokens,
53 - min_tokens=min_tokens),
54 - lora_request=lora_request,
55 - )
56 -
57 - return prompt, seq_group
58 -
59 -
60 -def create_dummy_lora_sequence(request_id: int, token_ids: list[int],
61 - block_size: int, lora_int_id: int) -> Sequence:
62 - return Sequence(seq_id=request_id,
63 - inputs=token_inputs(token_ids),
64 - block_size=block_size,
65 - lora_request=LoRARequest(lora_name="dummy",
66 - lora_path="/dummy",
67 - lora_int_id=lora_int_id))
68 -
69 -
70 -def create_dummy_sequence(request_id: int, token_ids: list[int],
71 - block_size: int) -> Sequence:
72 - return Sequence(
73 - seq_id=request_id,
74 - inputs=token_inputs(token_ids),
75 - block_size=block_size,
76 - )
77 -
78 -
79 -def create_dummy_prompt_encoder_decoder(
80 - request_id: str,
81 - decoder_prompt_length: int,
82 - encoder_prompt_length: int,
83 - block_size: Optional[int] = None,
84 - lora_request: Optional[LoRARequest] = None,
85 -) -> tuple[Sequence, Sequence, SequenceGroup]:
86 - if not block_size:
87 - block_size = decoder_prompt_length
88 -
89 - # Create dummy prompt sequence with tokens 0...block_size-1
90 - # and prompt "0 ... block_size". Note that the prompt string
91 - # doesn't actually match the tokens
92 - decoder_prompt_tokens = list(range(decoder_prompt_length))
93 - decoder_prompt_str = " ".join([str(t) for t in decoder_prompt_tokens])
94 - encoder_prompt_tokens = list(reversed(list(range(encoder_prompt_length))))
95 - encoder_prompt_str = " ".join([str(t) for t in encoder_prompt_tokens])
96 -
97 - inputs: EncoderDecoderInputs = {
98 - "decoder": token_inputs(decoder_prompt_tokens,
99 - prompt=decoder_prompt_str),
100 - "encoder": token_inputs(encoder_prompt_tokens,
101 - prompt=encoder_prompt_str),
102 - }
103 -
104 - decoder_prompt = Sequence(int(request_id),
105 - inputs=inputs["decoder"],
106 - block_size=block_size)
107 -
108 - encoder_prompt = Sequence(int(request_id),
109 - inputs=inputs["encoder"],
110 - block_size=block_size)
111 -
112 - seq_group = SequenceGroup(request_id=request_id,
113 - seqs=[decoder_prompt],
114 - arrival_time=time.time(),
115 - lora_request=lora_request,
116 - encoder_seq=encoder_prompt)
117 -
118 - return decoder_prompt, encoder_prompt, seq_group
119 -
120 -
121 -def create_seq_group(
122 - seq_prompt_len: int = 1024,
123 - seq_output_lens: GenericSequence[int] = (128, ),
124 - request_id: str = '0',
125 - seq_id_start: int = 0,
126 - sampling_params: Optional[SamplingParams] = None) -> SequenceGroup:
127 -
128 - assert len(seq_output_lens) > 0
129 -
130 - if sampling_params is None:
131 - sampling_params = SamplingParams()
132 -
133 - prompt_token_ids = [0] * seq_prompt_len
134 -
135 - seqs: list[Sequence] = []
136 - for seq_id_offset, output_len in enumerate(seq_output_lens):
137 - seq = Sequence(
138 - seq_id=seq_id_start + seq_id_offset,
139 - inputs=token_inputs(prompt_token_ids),
140 - block_size=16,
141 - )
142 -
143 - for i in range(output_len):
144 - seq.append_token_id(
145 - token_id=i,
146 - logprobs={i: Logprob(0.0)},
147 - )
148 - seqs.append(seq)
149 -
150 - seq_group = SequenceGroup(
151 - request_id=request_id,
152 - seqs=seqs,
153 - sampling_params=sampling_params,
154 - arrival_time=time.time(),
155 - )
156 -
157 - return seq_group
158 -
159 -
160 -def create_seq_group_encoder_decoder(
161 - seq_prompt_len: int = 1024,
162 - seq_output_lens: GenericSequence[int] = (128, ),
163 - request_id: str = '0',
164 - seq_id_start: int = 0,
165 - sampling_params: Optional[SamplingParams] = None) -> SequenceGroup:
166 -
167 - assert len(seq_output_lens) > 0
168 -
169 - if sampling_params is None:
170 - sampling_params = SamplingParams()
171 -
172 - prompt_token_ids = [0] * seq_prompt_len
173 -
174 - inputs: EncoderDecoderInputs = {
175 - "decoder": token_inputs(prompt_token_ids),
176 - "encoder": token_inputs(prompt_token_ids),
177 - }
178 -
179 - seqs = []
180 - for seq_id_offset, output_len in enumerate(seq_output_lens):
181 - # Construct decoder input sequences
182 - seq = Sequence(
183 - seq_id=seq_id_start + seq_id_offset,
184 - inputs=inputs["decoder"],
185 - block_size=16,
186 - )
187 -
188 - for i in range(output_len):
189 - seq.append_token_id(
190 - token_id=i,
191 - logprobs={i: Logprob(0.0)},
192 - )
193 - seqs.append(seq)
194 -
195 - # Encoder input sequence
196 - encoder_seq = Sequence(
197 - seq_id=seq_id_start + len(seq_output_lens),
198 - inputs=inputs["encoder"],
199 - block_size=16,
200 - )
201 -
202 - return SequenceGroup(request_id=request_id,
203 - seqs=seqs,
204 - sampling_params=sampling_params,
205 - arrival_time=time.time(),
206 - encoder_seq=encoder_seq)
207 -
208 -
209 -def round_up_to_next_block(seq_len: int, block_size: int) -> int:
210 - return (seq_len + block_size - 1) // block_size
211 -
212 -
213 -# Helper functions for scheduler tests
214 -
215 -
216 -def get_sequence_groups(scheduler_output):
217 - return [s.seq_group for s in scheduler_output.scheduled_seq_groups]
218 -
219 -
220 -def append_new_token(out, token_id: int):
221 - seq_groups = get_sequence_groups(out)
222 - for seq_group in seq_groups:
223 - for seq in seq_group.get_seqs():
224 - seq.append_token_id(token_id, {token_id: Logprob(token_id)})
225 -
226 -
227 -def schedule_and_update_computed_tokens(scheduler):
228 - metas, out, _ = scheduler.schedule()
229 - for s in out.scheduled_seq_groups:
230 - s.seq_group.update_num_computed_tokens(s.token_chunk_size)
231 - return metas, out
232 -
233 -
234 -def append_new_token_seq(seq: Sequence, token_id: int):
235 - seq.append_token_id(token_id, {token_id: Logprob(token_id)})
236 -
237 -
238 -def append_new_token_seq_group(token_chunk_size, seq_group, token_id: int):
239 - seq_group.update_num_computed_tokens(token_chunk_size)
240 - for seq in seq_group.get_seqs():
241 - seq.append_token_id(token_id, {token_id: Logprob(token_id)})
242 -
243 -
244 -class SchedulerProxy:
245 - """
246 - A proxy class to forward calls to the scheduler.
247 - """
248 -
249 - def __init__(self, scheduler: Scheduler):
250 - self.scheduler_ = scheduler
251 - self.call_history: dict[str, list[Any]] = defaultdict(list)
252 -
253 - def __getattr__(self, name: str) -> Any:
254 -
255 - def wrapper(*args, **kwargs):
256 - result = getattr(self.scheduler_, name)(*args, **kwargs)
257 - self.call_history[name].append((args, kwargs, result))
258 - return result
259 -
260 - return wrapper
261 -
262 - def last_schedule_ret(
263 - self, ) -> tuple[list[SequenceGroupMetadata], SchedulerOutputs, Any]:
264 - _, _, ret = self.call_history["schedule"][-1]
265 - return ret
266 -
267 -
268 -def create_seq_group_metadata_from_prompts(
269 - prompts: list[list[int]],
270 - num_gpu_blocks: int,
271 - block_size: int,
272 - final_prompt_lens: list[int],
273 - continuations: Optional[list[list[int]]] = None,
274 - seq_ids: Optional[list[int]] = None,
275 -) -> list[SequenceGroupMetadata]:
276 -
277 - if continuations is None:
278 - continuations = [[] for _ in prompts]
279 -
280 - if seq_ids is None:
281 - seq_ids = list(i for i, _ in enumerate(prompts))
282 -
283 - free_gpu_blocks = list(range(num_gpu_blocks))
284 -
285 - block_allocations = {
286 - i: [
287 - free_gpu_blocks.pop()
288 - for _ in range(round_up_to_next_block(final_len, block_size))
289 - ]
290 - for i, final_len in enumerate(final_prompt_lens)
291 - }
292 -
293 - seq_grou_metadata_list = []
294 - for i, (prompt_token_ids,
295 - cont_token_ids) in enumerate(zip(prompts, continuations)):
296 - data = SequenceData.from_seqs(prompt_token_ids, cont_token_ids)
297 - data.update_num_computed_tokens(
298 - len(prompt_token_ids) + len(cont_token_ids) - 1)
299 - seq_data = {i: data}
300 - seq_grou_metadata_list.append(
301 - SequenceGroupMetadata(
302 - request_id=str(i),
303 - is_prompt=len(cont_token_ids) == 0,
304 - seq_data=seq_data,
305 - sampling_params=SamplingParams(temperature=0.0),
306 - block_tables={i: block_allocations[i][:]},
307 - ))
308 - return seq_grou_metadata_list
309 -
310 -
311 -def create_chunked_seq_group_metadata_from_prompt(
312 - prompt: list[int],
313 - num_gpu_blocks: int,
314 - chunk_size: int,
315 - block_size: int,
316 - seq_id: Optional[int] = None) -> list[SequenceGroupMetadata]:
317 -
318 - if seq_id is None:
319 - seq_id = 0
320 -
321 - free_gpu_blocks = list(range(num_gpu_blocks))
322 -
323 - block_allocations = [
324 - free_gpu_blocks.pop()
325 - for _ in range(round_up_to_next_block(len(prompt), block_size))
326 - ]
327 -
328 - seq_group_metadata_list = []
329 - for i, idx in enumerate(range(0, len(prompt), chunk_size)):
330 - chunk_ids = prompt[idx:idx + chunk_size]
331 - data = SequenceData.from_seqs(prompt)
332 - data.update_num_computed_tokens(idx)
333 - seq_data = {i: data}
334 - seq_group_metadata_list.append(
335 - SequenceGroupMetadata(
336 - request_id=str(seq_id),
337 - is_prompt=True,
338 - do_sample=idx + chunk_size >= len(prompt), # terminal chunk
339 - seq_data=seq_data,
340 - sampling_params=SamplingParams(temperature=0.0),
341 - block_tables={i: block_allocations},
342 - token_chunk_size=len(chunk_ids)))
343 - return seq_group_metadata_list
344 -
345 -
346 -def create_batch(batch_size,
347 - k,
348 - prompt_len: Union[int, list[int]] = 10,
349 - prev_output_token_len: int = 10,
350 - seq_ids: Optional[list[int]] = None,
351 - num_gpu_blocks: Optional[int] = None,
352 - block_size: Optional[int] = None,
353 - prefill_chunk_size: Optional[int] = None):
354 - if block_size is None:
355 - block_size = 8
356 -
357 - if num_gpu_blocks is None:
358 - num_gpu_blocks = 2048 // block_size
359 -
360 - iterator = count()
361 -
362 - if isinstance(prompt_len, int):
363 - prompt_lens = [prompt_len for _ in range(batch_size)]
364 - else:
365 - prompt_lens = prompt_len
366 -
367 - prompts = [[next(iterator) for _ in range(p_len)] for p_len in prompt_lens]
368 -
369 - if prefill_chunk_size:
370 - # Create a batch of chunked prompts.
371 - if not seq_ids:
372 - seq_ids = list(range(len(prompts)))
373 - seq_group_metadata_list = []
374 - for p, sid in zip(prompts, seq_ids):
375 - seq_group_metadata_list += \
376 - create_chunked_seq_group_metadata_from_prompt(
377 - p, num_gpu_blocks, prefill_chunk_size, block_size, sid)
378 - seq_group_metadata_list = seq_group_metadata_list[:batch_size]
379 - prev_output_tokens = []
380 - else:
381 - prev_output_tokens = [[
382 - next(iterator) for _ in range(prev_output_token_len)
383 - ] for _ in range(batch_size)]
384 - final_prompt_lens = [
385 - len(prompt) + len(prev_output_token) + k + 1
386 - for prompt, prev_output_token in zip(prompts, prev_output_tokens)
387 - ]
388 -
389 - seq_group_metadata_list = create_seq_group_metadata_from_prompts(
390 - prompts, num_gpu_blocks, block_size, final_prompt_lens,
391 - prev_output_tokens, seq_ids)
392 - return seq_group_metadata_list, prompts, prev_output_tokens