@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e60f550b38
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 15, 09:45 UTC

[v1] Support multiple KV cache groups in GPU model runner (#17945)

Signed-off-by: Chen Zhang <zhangch99@outlook.com>

Chen Zhang committed May 15, 2025 at 09:54 UTC e60f550b3825cbce2d3c7e882b029e2c1d914d8d
16 files changed +482 -215
tests/v1/core/test_kv_cache_utils.py
+68 -3
@@ -19,7 +19,8 @@ from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType,
19 hash_request_tokens,
20 unify_kv_cache_configs)
21 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
22 - KVCacheGroupSpec, KVCacheTensor)
22 + KVCacheGroupSpec, KVCacheTensor,
23 + SlidingWindowSpec)
24 from vllm.v1.metrics.stats import PrefixCacheStats
25 from vllm.v1.request import Request
26
@@ -54,12 +55,14 @@ def new_kv_cache_spec(block_size=16,
55 num_kv_heads=2,
56 head_size=64,
57 dtype=torch.float32,
57 - use_mla=False):
58 + use_mla=False,
59 + sliding_window=None):
60 return FullAttentionSpec(block_size=block_size,
61 num_kv_heads=num_kv_heads,
62 head_size=head_size,
63 dtype=dtype,
62 - use_mla=use_mla)
64 + use_mla=use_mla,
65 + sliding_window=sliding_window)
66
67
68 def test_none_hash():
@@ -471,6 +474,68 @@ def test_unify_kv_cache_configs():
474 unify_kv_cache_configs(diff_kv_cache_config)
475
476
477 +def test_merge_kv_cache_spec():
478 + same_layer_specs = [
479 + new_kv_cache_spec(num_kv_heads=32),
480 + new_kv_cache_spec(num_kv_heads=32),
481 + ]
482 + merged_layer_spec = same_layer_specs[0].merge(same_layer_specs)
483 + assert merged_layer_spec.block_size == 16
484 + assert merged_layer_spec.num_kv_heads == 32
485 + assert merged_layer_spec.head_size == 64
486 + assert merged_layer_spec.dtype == torch.float32
487 + assert merged_layer_spec.sliding_window is None
488 +
489 + different_layer_specs = [
490 + new_kv_cache_spec(num_kv_heads=32),
491 + new_kv_cache_spec(num_kv_heads=16),
492 + ]
493 + with pytest.raises(AssertionError):
494 + different_layer_specs[0].merge(different_layer_specs)
495 +
496 + full_spec = new_kv_cache_spec(num_kv_heads=32)
497 + different_type_layer_specs = [
498 + full_spec,
499 + SlidingWindowSpec(
500 + block_size=full_spec.block_size,
501 + num_kv_heads=full_spec.num_kv_heads,
502 + head_size=full_spec.head_size,
503 + dtype=full_spec.dtype,
504 + use_mla=full_spec.use_mla,
505 + sliding_window=1,
506 + ),
507 + ]
508 + with pytest.raises(AssertionError):
509 + different_type_layer_specs[0].merge(different_type_layer_specs)
510 + with pytest.raises(AssertionError):
511 + different_type_layer_specs[1].merge(different_type_layer_specs)
512 +
513 + different_sliding_window_layer_specs = [
514 + new_kv_cache_spec(num_kv_heads=32),
515 + new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
516 + new_kv_cache_spec(num_kv_heads=32, sliding_window=2),
517 + ]
518 + with pytest.raises(ValueError):
519 + different_sliding_window_layer_specs[0].merge(
520 + different_sliding_window_layer_specs)
521 +
522 + same_sliding_window_layer_specs = [
523 + new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
524 + new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
525 + ]
526 + merged_layer_spec = same_sliding_window_layer_specs[0].merge(
527 + same_sliding_window_layer_specs)
528 + assert merged_layer_spec.sliding_window == 1
529 +
530 + same_sliding_window_layer_spec_with_none = [
531 + new_kv_cache_spec(num_kv_heads=32, sliding_window=1),
532 + new_kv_cache_spec(num_kv_heads=32, sliding_window=None),
533 + ]
534 + merged_layer_spec = same_sliding_window_layer_spec_with_none[0].merge(
535 + same_sliding_window_layer_spec_with_none)
536 + assert merged_layer_spec.sliding_window == 1
537 +
538 +
539 @pytest.mark.parametrize(
540 ("model_id", "max_model_len", "want_estimated_max_len"), [
541 ("Qwen/Qwen1.5-7B", 16385, 16384),
tests/v1/core/test_prefix_caching.py
+18 -18
@@ -84,7 +84,7 @@ def test_prefill(hash_algo):
84 blocks = manager.allocate_slots(req0, 55,
85 len(computed_blocks.blocks) * 16,
86 computed_blocks)
87 - assert blocks.get_block_ids() == [1, 2, 3, 4]
87 + assert blocks.get_block_ids() == [[1, 2, 3, 4]]
88
89 # Check full block metadata
90 parent_block_hash = None
@@ -107,13 +107,13 @@ def test_prefill(hash_algo):
107 req1 = make_request("1", common_token_ids + unique_token_ids)
108 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
109 assert len(manager.req_to_block_hashes[req1.request_id]) == 3
110 - assert computed_blocks.get_block_ids() == [1, 2, 3]
110 + assert computed_blocks.get_block_ids() == [[1, 2, 3]]
111 assert num_computed_tokens == 3 * 16
112 num_new_tokens = 53 - 3 * 16
113 blocks = manager.allocate_slots(req1, num_new_tokens,
114 len(computed_blocks.blocks) * 16,
115 computed_blocks)
116 - assert blocks.get_block_ids() == [5]
116 + assert blocks.get_block_ids() == [[5]]
117 for block in computed_blocks.blocks:
118 assert block.ref_cnt == 2
119
@@ -141,13 +141,13 @@ def test_prefill(hash_algo):
141 req2 = make_request("2", common_token_ids + unique_token_ids)
142 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
143 assert len(manager.req_to_block_hashes[req2.request_id]) == 3
144 - assert computed_blocks.get_block_ids() == [1, 2, 3]
144 + assert computed_blocks.get_block_ids() == [[1, 2, 3]]
145 assert num_computed_tokens == 3 * 16
146 num_new_tokens = 53 - 3 * 16
147 blocks = manager.allocate_slots(req2, num_new_tokens,
148 len(computed_blocks.blocks) * 16,
149 computed_blocks)
150 - assert blocks.get_block_ids() == [6]
150 + assert blocks.get_block_ids() == [[6]]
151
152 # Although we only have 6 free blocks, we have 8 blocks in
153 # the free block queue due to lazy removal.
@@ -171,7 +171,7 @@ def test_prefill(hash_algo):
171 len(computed_blocks.blocks) * 16,
172 computed_blocks)
173 # This block ID order also checks the eviction order.
174 - assert blocks.get_block_ids() == [7, 8, 9, 10, 4, 5, 6, 3, 2, 1]
174 + assert blocks.get_block_ids() == [[7, 8, 9, 10, 4, 5, 6, 3, 2, 1]]
175 assert manager.block_pool.free_block_queue.num_free_blocks == 0
176 assert manager.block_pool.free_block_queue.free_list_head is None
177 assert manager.block_pool.free_block_queue.free_list_tail is None
@@ -208,7 +208,7 @@ def test_prefill_plp():
208 blocks = manager.allocate_slots(req0, 55,
209 len(computed_blocks.blocks) * 16,
210 computed_blocks)
211 - assert blocks.get_block_ids() == [1, 2, 3, 4]
211 + assert blocks.get_block_ids() == [[1, 2, 3, 4]]
212 req0_block_hashes = [b.block_hash for b in blocks.blocks]
213
214 # Check full block metadata
@@ -233,13 +233,13 @@ def test_prefill_plp():
233 req1 = make_request("1", common_token_ids + unique_token_ids)
234 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
235 assert len(manager.req_to_block_hashes[req1.request_id]) == 3
236 - assert computed_blocks.get_block_ids() == [1, 2, 3]
236 + assert computed_blocks.get_block_ids() == [[1, 2, 3]]
237 assert num_computed_tokens == 3 * 16
238 num_new_tokens = 53 - 3 * 16
239 blocks = manager.allocate_slots(req1, num_new_tokens,
240 len(computed_blocks.blocks) * 16,
241 computed_blocks)
242 - assert blocks.get_block_ids() == [5]
242 + assert blocks.get_block_ids() == [[5]]
243 for block in computed_blocks.blocks:
244 assert block.ref_cnt == 2
245
@@ -277,11 +277,11 @@ def test_prefill_plp():
277 block_ids = blocks.get_block_ids()
278 # Duplicate cached blocks have different ids but same hashes vs request #0
279 assert [b.block_hash for b in blocks.blocks] == req0_block_hashes
280 - assert block_ids != [1, 2, 3, 4]
280 + assert block_ids != [[1, 2, 3, 4]]
281
282 # Request #2 block hashes are valid since request #0 hashes are.
283 # Check block reference counts.
284 - for block_id in block_ids:
284 + for block_id in block_ids[0]:
285 assert manager.block_pool.blocks[block_id].ref_cnt == 1
286
287 manager.free(req2)
@@ -307,7 +307,7 @@ def test_decode():
307 blocks = manager.allocate_slots(req0, 55,
308 len(computed_blocks.blocks) * 16,
309 computed_blocks)
310 - assert blocks.get_block_ids() == [1, 2, 3, 4]
310 + assert blocks.get_block_ids() == [[1, 2, 3, 4]]
311
312 # Append slots without allocating a new block.
313 req0.num_computed_tokens = 55
@@ -379,12 +379,12 @@ def test_evict():
379 # Touch the first 2 blocks.
380 req2 = make_request("2", list(range(2 * 16 + 3)))
381 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
382 - assert computed_blocks.get_block_ids() == [1, 2]
382 + assert computed_blocks.get_block_ids() == [[1, 2]]
383 assert num_computed_tokens == 2 * 16
384 blocks = manager.allocate_slots(req2, 3,
385 len(computed_blocks.blocks) * 16,
386 computed_blocks)
387 - assert blocks.get_block_ids() == [10]
387 + assert blocks.get_block_ids() == [[10]]
388 assert manager.block_pool.free_block_queue.num_free_blocks == 7
389
390
@@ -625,7 +625,7 @@ def test_mm_prefix_caching():
625 blocks = manager.allocate_slots(req0, 59,
626 len(computed_blocks.blocks) * 16,
627 computed_blocks)
628 - assert blocks.get_block_ids() == [1, 2, 3, 4]
628 + assert blocks.get_block_ids() == [[1, 2, 3, 4]]
629 req0.num_computed_tokens = 59
630
631 # Append slots without allocating a new block.
@@ -686,7 +686,7 @@ def test_cache_key_salting():
686 blocks = manager.allocate_slots(req0, 59,
687 len(computed_blocks.blocks) * 16,
688 computed_blocks)
689 - assert blocks.get_block_ids() == [1, 2, 3, 4]
689 + assert blocks.get_block_ids() == [[1, 2, 3, 4]]
690 req0.num_computed_tokens = 59
691
692 # Append slots without allocating a new block.
@@ -797,7 +797,7 @@ def test_reset_prefix_cache():
797 all_token_ids = full_block_token_ids + unique_token_ids
798 req0 = make_request("0", all_token_ids)
799 blocks = manager.allocate_slots(req0, 55)
800 - assert blocks.get_block_ids() == [1, 2, 3, 4]
800 + assert blocks.get_block_ids() == [[1, 2, 3, 4]]
801
802 unique_token_ids = [4] * 7
803 all_token_ids = full_block_token_ids + unique_token_ids
@@ -808,7 +808,7 @@ def test_reset_prefix_cache():
808 blocks = manager.allocate_slots(req1, 7,
809 len(computed_blocks.blocks) * 16,
810 computed_blocks)
811 - assert blocks.get_block_ids() == [5]
811 + assert blocks.get_block_ids() == [[5]]
812
813 # Failed to reset prefix cache because some blocks are not freed yet.
814 assert not manager.reset_prefix_cache()
tests/v1/worker/test_gpu_input_batch.py
+33 -6
@@ -9,9 +9,11 @@ import torch
9
10 from vllm.sampling_params import SamplingParams
11 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
12 +from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
13 + KVCacheGroupSpec, KVCacheTensor)
14 from vllm.v1.sample.metadata import SamplingMetadata
13 -from vllm.v1.worker.gpu_input_batch import (BlockTable, CachedRequestState,
14 - InputBatch)
15 +from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable
16 +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
17
18 VOCAB_SIZE = 1024
19 NUM_OUTPUT_TOKENS = 20
@@ -22,6 +24,27 @@ CUDA_DEVICES = [
24 MAX_NUM_PROMPT_TOKENS = 64
25
26
27 +def get_kv_cache_config() -> KVCacheConfig:
28 + return KVCacheConfig(
29 + num_blocks=10,
30 + tensors={
31 + "layer.0": KVCacheTensor(size=1024),
32 + },
33 + kv_cache_groups=[
34 + KVCacheGroupSpec(
35 + layer_names=["layer.0"],
36 + kv_cache_spec=FullAttentionSpec(
37 + block_size=1,
38 + num_kv_heads=1,
39 + head_size=16,
40 + dtype=torch.float16,
41 + use_mla=False,
42 + ),
43 + ),
44 + ],
45 + )
46 +
47 +
48 def _compare_objs(obj1, obj2):
49 attrs = inspect.getmembers(obj1, lambda a: not (inspect.isroutine(a)))
50 attr_names = set([
@@ -41,6 +64,10 @@ def _compare_objs(obj1, obj2):
64 elif isinstance(a, np.ndarray):
65 if np.allclose(a, b):
66 is_same = True
67 + elif isinstance(a, MultiGroupBlockTable):
68 + for a_i, b_i in zip(a.block_tables, b.block_tables):
69 + _compare_objs(a_i, b_i)
70 + is_same = True
71 elif isinstance(a, (BlockTable, SamplingMetadata)):
72 _compare_objs(a, b)
73 is_same = True # if we make it here must be same
@@ -198,7 +225,7 @@ def _construct_cached_request_state(req_id_suffix: int):
225 sampling_params=_create_sampling_params(),
226 mm_inputs=[],
227 mm_positions=[],
201 - block_ids=[],
228 + block_ids=[[]],
229 generator=None,
230 num_computed_tokens=len(output_token_ids),
231 output_token_ids=output_token_ids,
@@ -220,11 +247,11 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
247 input_batch: InputBatch = InputBatch(
248 max_num_reqs=batch_size,
249 max_model_len=1024,
223 - max_num_blocks_per_req=10,
250 max_num_batched_tokens=1024,
251 device=torch.device(device),
252 pin_memory=is_pin_memory_available(),
253 vocab_size=1024,
254 + kv_cache_config=get_kv_cache_config(),
255 )
256 reqs: list[CachedRequestState] = []
257 req_id_reqs = {}
@@ -310,20 +337,20 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,
337 input_batch: InputBatch = InputBatch(
338 max_num_reqs=batch_size,
339 max_model_len=1024,
313 - max_num_blocks_per_req=10,
340 max_num_batched_tokens=1024,
341 device=torch.device(device),
342 pin_memory=is_pin_memory_available(),
343 vocab_size=1024,
344 + kv_cache_config=get_kv_cache_config(),
345 )
346 ref_input_batch: InputBatch = InputBatch(
347 max_num_reqs=batch_size,
348 max_model_len=1024,
322 - max_num_blocks_per_req=10,
349 max_num_batched_tokens=1024,
350 device=torch.device(device),
351 pin_memory=is_pin_memory_available(),
352 vocab_size=1024,
353 + kv_cache_config=get_kv_cache_config(),
354 )
355
356 reqs: list[CachedRequestState] = []
tests/v1/worker/test_gpu_model_runner.py
+41 -16
@@ -1,15 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -import weakref
2
3 import pytest
5 -import torch
4
7 -from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
5 +from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
6 + SchedulerConfig, VllmConfig)
7 from vllm.sampling_params import SamplingParams
8 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
9 SchedulerOutput)
11 -from vllm.v1.kv_cache_interface import FullAttentionSpec
10 +from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
11 + KVCacheGroupSpec, KVCacheTensor)
12 from vllm.v1.sample.metadata import SamplingMetadata
13 +from vllm.v1.worker.gpu_input_batch import InputBatch
14 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
15
16
@@ -17,13 +18,34 @@ def initialize_kv_cache(runner: GPUModelRunner):
18 """
19 Only perform necessary steps in GPUModelRunner.initialize_kv_cache()
20 """
20 - kv_cache_spec = FullAttentionSpec(block_size=16,
21 - num_kv_heads=1,
22 - head_size=64,
23 - dtype=torch.float16,
24 - use_mla=False)
25 - runner.attn_metadata_builder = runner.attn_backend.get_builder_cls()(
26 - weakref.proxy(runner), kv_cache_spec, runner.input_batch.block_table)
21 + kv_cache_config = KVCacheConfig(
22 + num_blocks=10,
23 + tensors={
24 + "layer.0": KVCacheTensor(size=1024),
25 + },
26 + kv_cache_groups=[
27 + KVCacheGroupSpec(
28 + layer_names=["layer.0"],
29 + kv_cache_spec=FullAttentionSpec(
30 + block_size=16,
31 + num_kv_heads=runner.model_config.get_num_kv_heads(
32 + runner.parallel_config),
33 + head_size=runner.model_config.get_head_size(),
34 + dtype=runner.kv_cache_dtype,
35 + use_mla=False,
36 + ))
37 + ])
38 + runner.kv_cache_config = kv_cache_config
39 + runner.input_batch = InputBatch(
40 + max_num_reqs=runner.max_num_reqs,
41 + max_model_len=runner.max_model_len,
42 + max_num_batched_tokens=runner.max_num_tokens,
43 + device=runner.device,
44 + pin_memory=runner.pin_memory,
45 + vocab_size=runner.model_config.get_vocab_size(),
46 + kv_cache_config=kv_cache_config,
47 + )
48 + runner.initialize_attn_backend(kv_cache_config)
49
50
51 @pytest.fixture
@@ -48,10 +70,12 @@ def model_runner():
70 swap_space=0,
71 cache_dtype="auto",
72 )
73 + parallel_config = ParallelConfig()
74 vllm_config = VllmConfig(
75 model_config=model_config,
76 cache_config=cache_config,
77 scheduler_config=scheduler_config,
78 + parallel_config=parallel_config,
79 )
80
81 device = "cuda"
@@ -73,7 +97,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
97 mm_hashes=[],
98 mm_positions=[],
99 sampling_params=SamplingParams(),
76 - block_ids=[0],
100 + block_ids=[[0]],
101 num_computed_tokens=0,
102 lora_request=None,
103 ))
@@ -111,13 +135,14 @@ def _is_sampling_metadata_changed(model_runner,
135
136 def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
137 req_index = model_runner.input_batch.req_id_to_index[req_id]
114 - block_table = model_runner.input_batch.block_table
138 + block_table = model_runner.input_batch.block_table[0]
139 req_state = model_runner.requests[req_id]
116 - if block_table.num_blocks_per_row[req_index] != len(req_state.block_ids):
140 + if block_table.num_blocks_per_row[req_index] != len(
141 + req_state.block_ids[0]):
142 return False
143 num_blocks = block_table.num_blocks_per_row[req_index]
144 return (block_table.block_table_np[req_index, :num_blocks] ==
120 - req_state.block_ids).all()
145 + req_state.block_ids[0]).all()
146
147
148 def test_update_states_new_request(model_runner):
@@ -200,7 +225,7 @@ def test_update_states_request_resumed(model_runner):
225 req_id=req_id,
226 resumed_from_preemption=False,
227 new_token_ids=[],
203 - new_block_ids=[],
228 + new_block_ids=[[]],
229 num_computed_tokens=0,
230 )
231
tests/weight_loading/models.txt
+1 -1
@@ -2,7 +2,7 @@ gptq_marlin, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
2 gptq_marlin, TheBloke/Llama-2-7B-GPTQ, main
3 gptq_marlin, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, main
4 gptq_marlin, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True
5 -gptq_marlin, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit-32g-actorder_True
5 +#gptq_marlin, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit-32g-actorder_True
6 gptq_marlin, TechxGenus/gemma-1.1-2b-it-GPTQ, main
7 gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
8 gptq, TheBloke/Llama-2-7B-GPTQ, main
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py
+3 -3
@@ -288,7 +288,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
288 for new_req in scheduler_output.scheduled_new_reqs:
289 if new_req.req_id in self._requests_need_load:
290 meta.add_request(token_ids=new_req.prompt_token_ids,
291 - block_ids=new_req.block_ids,
291 + block_ids=new_req.block_ids[0],
292 block_size=self._block_size,
293 is_store=False)
294 total_need_load += 1
@@ -299,7 +299,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
299 # the original prompt tokens.
300 if not self._found_match_for_request(new_req):
301 meta.add_request(token_ids=new_req.prompt_token_ids,
302 - block_ids=new_req.block_ids,
302 + block_ids=new_req.block_ids[0],
303 block_size=self._block_size,
304 is_store=True)
305
@@ -319,7 +319,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
319
320 # NOTE(rob): For resumed req, new_block_ids is all
321 # of the block_ids for the request.
322 - block_ids = cached_req.new_block_ids
322 + block_ids = cached_req.new_block_ids[0]
323
324 meta.add_request(token_ids=token_ids,
325 block_ids=block_ids,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+2 -2
@@ -67,13 +67,13 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
67 max_model_len = self.runner.model_config.max_model_len
68 assert max_model_len == 32768,\
69 "AITER MLA requires max_model_len=32768"
70 - assert self.runner.block_size == 1, "AITER MLA" \
70 + assert self.kv_cache_spec.block_size == 1, "AITER MLA" \
71 "only supports block size 1."
72
73 def _get_paged_kv_tensors(
74 self, block_table: torch.Tensor,
75 seq_lens: torch.Tensor) -> tuple[torch.Tensor, ...]:
76 - page_size = self.runner.block_size
76 + page_size = self.kv_cache_spec.block_size
77 block_table_bounds = (seq_lens + page_size - 1) // page_size
78
79 mask = (torch.arange(block_table.size(1),
vllm/v1/core/kv_cache_manager.py
+21 -13
@@ -32,9 +32,16 @@ class KVCacheBlocks:
32 """Creates a new KVCacheBlocks instance with no blocks."""
33 return cls([])
34
35 - def get_block_ids(self) -> list[int]:
36 - """Converts the KVCacheBlocks instance to a list of block IDs."""
37 - return [block.block_id for block in self.blocks]
35 + def get_block_ids(self) -> list[list[int]]:
36 + """
37 + Converts the KVCacheBlocks instance to block_ids.
38 +
39 + Returns:
40 + list[list[int]]: A two-level list where
41 + * the outer list corresponds to KV cache groups (only 1 group now)
42 + * each inner list contains the block_ids of the blocks in that group
43 + """
44 + return [[block.block_id for block in self.blocks]]
45
46 def get_unhashed_block_ids(self) -> list[int]:
47 """Get block_ids of unhashed blocks from KVCacheBlocks instance."""
@@ -300,9 +307,9 @@ class KVCacheManager:
307 self,
308 request: Request,
309 num_running_requests: int,
303 - ) -> int:
310 + ) -> list[int]:
311 """Calculate the number of common prefix blocks shared by all requests
305 - in the RUNNING state.
312 + in the RUNNING state for each kv cache group.
313
314 The function determines this by selecting any request and iterating
315 through its blocks. A block is considered a common prefix block if its
@@ -332,11 +339,14 @@ class KVCacheManager:
339 requests in the current step.
340
341 Returns:
335 - int: The number of common prefix blocks.
342 + list[int]: The number of common prefix blocks for each kv cache
343 + group.
344 """
345 assert request.status == RequestStatus.RUNNING
338 - return self.single_type_manager.get_num_common_prefix_blocks(
339 - request.request_id, num_running_requests)
346 + return [
347 + self.single_type_manager.get_num_common_prefix_blocks(
348 + request.request_id, num_running_requests)
349 + ]
350
351 def free_block_hashes(self, request: Request) -> None:
352 """Discard the block hashes for the request.
@@ -354,10 +364,8 @@ class KVCacheManager:
364 """
365 return self.block_pool.take_events()
366
357 - def get_block_ids(self, request_id: str) -> list[int]:
367 + def get_block_ids(self, request_id: str) -> list[list[int]]:
368 """Get the block ids of a request."""
369 assert request_id in self.single_type_manager.req_to_blocks
360 - return [
361 - block.block_id
362 - for block in self.single_type_manager.req_to_blocks[request_id]
363 - ]
370 + return KVCacheBlocks(self.single_type_manager.req_to_blocks[request_id]
371 + ).get_block_ids()
vllm/v1/core/kv_cache_utils.py
+6 -7
@@ -577,14 +577,12 @@ def create_kv_cache_group_specs(
577 """
578 kv_cache_groups = []
579 for layer_names_one_group in grouped_layer_names:
580 - layer_spec = kv_cache_spec[layer_names_one_group[0]]
581 - assert all(
582 - kv_cache_spec[layer_name] == layer_spec
583 - for layer_name in layer_names_one_group[1:]), (
584 - "All layers in the same KV cache group must share the same "
585 - "KVCacheSpec.")
580 + layer_specs = [
581 + kv_cache_spec[layer_name] for layer_name in layer_names_one_group
582 + ]
583 + merged_layer_spec = layer_specs[0].merge(layer_specs)
584 kv_cache_groups.append(
587 - KVCacheGroupSpec(layer_names_one_group, layer_spec))
585 + KVCacheGroupSpec(layer_names_one_group, merged_layer_spec))
586 return kv_cache_groups
587
588
@@ -683,6 +681,7 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
681 head_size=spec.head_size,
682 dtype=spec.dtype,
683 use_mla=spec.use_mla,
684 + sliding_window=spec.sliding_window,
685 )
686
687
vllm/v1/core/sched/output.py
+6 -6
@@ -26,7 +26,7 @@ class NewRequestData:
26 mm_hashes: list[str]
27 mm_positions: list[PlaceholderRange]
28 sampling_params: SamplingParams
29 - block_ids: list[int]
29 + block_ids: list[list[int]]
30 num_computed_tokens: int
31 lora_request: Optional[LoRARequest]
32
@@ -34,7 +34,7 @@ class NewRequestData:
34 def from_request(
35 cls,
36 request: Request,
37 - block_ids: list[int],
37 + block_ids: list[list[int]],
38 ) -> NewRequestData:
39 return cls(
40 req_id=request.request_id,
@@ -85,7 +85,7 @@ class CachedRequestData:
85 # request's block IDs instead of appending to the existing block IDs.
86 resumed_from_preemption: bool
87 new_token_ids: list[int]
88 - new_block_ids: list[int]
88 + new_block_ids: list[list[int]]
89 num_computed_tokens: int
90
91 @classmethod
@@ -94,7 +94,7 @@ class CachedRequestData:
94 request: Request,
95 resumed_from_preemption: bool,
96 new_token_ids: list[int],
97 - new_block_ids: list[int],
97 + new_block_ids: list[list[int]],
98 ) -> CachedRequestData:
99 return cls(
100 req_id=request.request_id,
@@ -131,9 +131,9 @@ class SchedulerOutput:
131 # E.g., if a request has [0, 1], it could mean the vision encoder needs
132 # to process that the request's 0-th and 1-th images in the current step.
133 scheduled_encoder_inputs: dict[str, list[int]]
134 - # Number of common prefix blocks for all requests.
134 + # Number of common prefix blocks for all requests in each KV cache group.
135 # This can be used for cascade attention.
136 - num_common_prefix_blocks: int
136 + num_common_prefix_blocks: list[int]
137
138 # Request IDs that are finished in between the previous and the current
139 # steps. This is used to notify the workers about the finished requests
vllm/v1/core/sched/scheduler.py
+10 -6
@@ -173,7 +173,7 @@ class Scheduler(SchedulerInterface):
173 # uses structured decoding.
174 structured_output_request_ids: dict[str, int] = {}
175
176 - req_to_new_block_ids: dict[str, list[int]] = {}
176 + req_to_new_block_ids: dict[str, list[list[int]]] = {}
177 num_scheduled_tokens: dict[str, int] = {}
178 token_budget = self.max_num_scheduled_tokens
179 # Encoder-related.
@@ -477,7 +477,8 @@ class Scheduler(SchedulerInterface):
477
478 # Get the longest common prefix among all requests in the running queue.
479 # This can be potentially used for cascade attention.
480 - num_common_prefix_blocks = 0
480 + num_common_prefix_blocks = [0] * len(
481 + self.kv_cache_config.kv_cache_groups)
482 if self.running:
483 any_request = self.running[0]
484 num_common_prefix_blocks = (
@@ -564,7 +565,7 @@ class Scheduler(SchedulerInterface):
565 request: Request,
566 num_scheduled_tokens: int,
567 num_scheduled_spec_tokens: int,
567 - new_block_ids: list[int],
568 + new_block_ids: list[list[int]],
569 resumed_from_preemption: bool,
570 ) -> CachedRequestData:
571 # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating
@@ -939,7 +940,9 @@ class Scheduler(SchedulerInterface):
940 """
941 if self.connector is None:
942 return False, None
942 - block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
943 + assert len(self.kv_cache_config.kv_cache_groups
944 + ) == 1, "KV connector only supports one KV cache group now"
945 + block_ids = self.kv_cache_manager.get_block_ids(request.request_id)[0]
946 return self.connector.request_finished(request, block_ids)
947
948 def _update_waiting_for_remote_kv(self, request: Request) -> bool:
@@ -956,9 +959,10 @@ class Scheduler(SchedulerInterface):
959 """
960 if request.request_id not in self.finished_recving_kv_req_ids:
961 return False
959 -
962 + assert len(self.kv_cache_config.kv_cache_groups
963 + ) == 1, "KV connector only supports one KV cache group now"
964 # Now that the blocks are ready, actually cache them.
961 - block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
965 + block_ids = self.kv_cache_manager.get_block_ids(request.request_id)[0]
966 num_computed_tokens = len(block_ids) * self.block_size
967 if num_computed_tokens == request.num_tokens:
968 num_computed_tokens -= 1
vllm/v1/kv_cache_interface.py
+42
@@ -1,8 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +import copy
4 from dataclasses import dataclass
5 +from typing import Optional
6
7 import torch
8 +from typing_extensions import Self
9
10 from vllm.config import VllmConfig
11 from vllm.logger import init_logger
@@ -53,6 +56,16 @@ class KVCacheSpec:
56 """
57 raise NotImplementedError
58
59 + @classmethod
60 + def merge(cls, specs: list[Self]) -> Self:
61 + """
62 + Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
63 + """
64 + assert all(spec.type_id == specs[0].type_id for spec in specs[1:]), (
65 + "All layers in the same KV cache group must share the same "
66 + "type_id.")
67 + return copy.deepcopy(specs[0])
68 +
69
70 @dataclass
71 class AttentionSpec(KVCacheSpec):
@@ -71,6 +84,16 @@ class AttentionSpec(KVCacheSpec):
84
85 @dataclass
86 class FullAttentionSpec(AttentionSpec):
87 + sliding_window: Optional[int] = None
88 + """
89 + When hybrid allocator is disabled and the model contains both full
90 + attention layers and sliding window attention layers, sliding
91 + window attention are regarded as full attention in KV cache manager
92 + (blocks are allocated for all tokens), while computed as sliding window
93 + attention in model runner.
94 + In this case, we use FullAttentionSpec and record the sliding window size.
95 + Default to None for not using sliding window attention.
96 + """
97
98 @property
99 def type_id(self) -> str:
@@ -80,6 +103,25 @@ class FullAttentionSpec(AttentionSpec):
103 max_model_len = vllm_config.model_config.max_model_len
104 return cdiv(max_model_len, self.block_size) * self.page_size_bytes
105
106 + @classmethod
107 + def merge(cls, specs: list[Self]) -> Self:
108 + """
109 + Merge a list of FullAttentionSpec objects into a single
110 + FullAttentionSpec object.
111 + """
112 + merged_spec = super().merge(specs)
113 + sliding_window = set(spec.sliding_window for spec in specs
114 + if spec.sliding_window is not None)
115 + if len(sliding_window) == 0:
116 + merged_spec.sliding_window = None
117 + elif len(sliding_window) == 1:
118 + merged_spec.sliding_window = sliding_window.pop()
119 + else:
120 + raise ValueError(
121 + "All sliding window layers in the same KV cache group "
122 + "must have the same window size.")
123 + return merged_spec
124 +
125
126 @dataclass
127 class SlidingWindowSpec(AttentionSpec):
vllm/v1/worker/block_table.py
+47
@@ -4,6 +4,8 @@ import numpy as np
4 import torch
5
6 from vllm.logger import init_logger
7 +from vllm.utils import cdiv
8 +from vllm.v1.kv_cache_interface import KVCacheConfig
9
10 logger = init_logger(__name__)
11
@@ -96,3 +98,48 @@ class BlockTable:
98 def get_numpy_array(self) -> np.ndarray:
99 """Returns the numpy array of the block table."""
100 return self.block_table_np
101 +
102 +
103 +class MultiGroupBlockTable:
104 + """The BlockTables for each KV cache group."""
105 +
106 + def __init__(self, max_num_reqs: int, max_model_len: int,
107 + max_num_batched_tokens: int, pin_memory: bool,
108 + device: torch.device, kv_cache_config: KVCacheConfig) -> None:
109 + max_num_blocks_per_req = [
110 + cdiv(max_model_len, g.kv_cache_spec.block_size)
111 + for g in kv_cache_config.kv_cache_groups
112 + ]
113 + self.block_tables = [
114 + BlockTable(max_num_reqs, max_num_blocks_per_req[i],
115 + max_num_batched_tokens, pin_memory, device)
116 + for i in range(len(kv_cache_config.kv_cache_groups))
117 + ]
118 +
119 + def append_row(self, block_ids: list[list[int]], row_idx: int) -> None:
120 + for i, block_table in enumerate(self.block_tables):
121 + block_table.append_row(block_ids[i], row_idx)
122 +
123 + def add_row(self, block_ids: list[list[int]], row_idx: int) -> None:
124 + for i, block_table in enumerate(self.block_tables):
125 + block_table.add_row(block_ids[i], row_idx)
126 +
127 + def move_row(self, src: int, tgt: int) -> None:
128 + for block_table in self.block_tables:
129 + block_table.move_row(src, tgt)
130 +
131 + def swap_row(self, src: int, tgt: int) -> None:
132 + for block_table in self.block_tables:
133 + block_table.swap_row(src, tgt)
134 +
135 + def commit(self, num_reqs: int) -> None:
136 + for block_table in self.block_tables:
137 + block_table.commit(num_reqs)
138 +
139 + def clear(self) -> None:
140 + for block_table in self.block_tables:
141 + block_table.clear()
142 +
143 + def __getitem__(self, idx: int) -> "BlockTable":
144 + """Returns the BlockTable for the i-th KV cache group."""
145 + return self.block_tables[idx]
vllm/v1/worker/gpu_input_batch.py
+7 -6
@@ -11,10 +11,11 @@ from vllm.lora.request import LoRARequest
11 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
12 from vllm.sampling_params import SamplingParams, SamplingType
13 from vllm.utils import swap_dict_values
14 +from vllm.v1.kv_cache_interface import KVCacheConfig
15 from vllm.v1.outputs import LogprobsTensors
16 from vllm.v1.sample.metadata import SamplingMetadata
17 from vllm.v1.utils import copy_slice
17 -from vllm.v1.worker.block_table import BlockTable
18 +from vllm.v1.worker.block_table import MultiGroupBlockTable
19
20 _SAMPLING_EPS = 1e-5
21
@@ -29,7 +30,7 @@ class CachedRequestState:
30 sampling_params: SamplingParams
31 generator: Optional[torch.Generator]
32
32 - block_ids: list[int]
33 + block_ids: list[list[int]]
34 num_computed_tokens: int
35 output_token_ids: list[int]
36
@@ -58,15 +59,14 @@ class InputBatch:
59 self,
60 max_num_reqs: int,
61 max_model_len: int,
61 - max_num_blocks_per_req: int,
62 max_num_batched_tokens: int,
63 device: torch.device,
64 pin_memory: bool,
65 vocab_size: int,
66 + kv_cache_config: KVCacheConfig,
67 ):
68 self.max_num_reqs = max_num_reqs
69 self.max_model_len = max_model_len
69 - self.max_num_blocks_per_req = max_num_blocks_per_req
70 self.max_num_batched_tokens = max_num_batched_tokens
71 self.device = device
72 self.pin_memory = pin_memory
@@ -99,12 +99,13 @@ class InputBatch:
99 self.num_computed_tokens_cpu_tensor.numpy()
100
101 # Block table.
102 - self.block_table = BlockTable(
102 + self.block_table = MultiGroupBlockTable(
103 max_num_reqs=max_num_reqs,
104 - max_num_blocks_per_req=max_num_blocks_per_req,
104 + max_model_len=max_model_len,
105 max_num_batched_tokens=max_num_batched_tokens,
106 pin_memory=pin_memory,
107 device=device,
108 + kv_cache_config=kv_cache_config,
109 )
110
111 # Sampling-related.
vllm/v1/worker/gpu_model_runner.py
+158 -112
@@ -12,6 +12,8 @@ import torch.distributed
12 import torch.nn as nn
13
14 from vllm.attention import AttentionType, get_attn_backend
15 +from vllm.attention.backends.abstract import (AttentionBackend,
16 + AttentionMetadataBuilder)
17 from vllm.attention.layer import Attention
18 from vllm.attention.utils.fa_utils import get_flash_attn_version
19 from vllm.config import (CompilationLevel, VllmConfig,
@@ -31,8 +33,8 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
33 from vllm.sampling_params import SamplingType
34 from vllm.sequence import IntermediateTensors
35 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
34 - GiB_bytes, LayerBlockType, LazyLoader, cdiv,
35 - check_use_alibi, is_pin_memory_available)
36 + GiB_bytes, LazyLoader, cdiv, check_use_alibi,
37 + is_pin_memory_available)
38 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
39 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
40 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -49,6 +51,7 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
51 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
52 from vllm.v1.spec_decode.utils import is_spec_decode_supported
53 from vllm.v1.utils import bind_kv_cache
54 +from vllm.v1.worker.block_table import BlockTable
55 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
56 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
57
@@ -100,59 +103,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
103 self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
104 cache_config.cache_dtype]
105
103 - # NOTE(woosuk): sliding_window is None for models with interleaved
104 - # attention. Use interleaved_sliding_window instead.
105 - self.sliding_window = model_config.get_sliding_window()
106 - self.interleaved_sliding_window = getattr(
107 - model_config.hf_text_config, "interleaved_sliding_window", None)
108 - self.window_size = (self.sliding_window
109 - or self.interleaved_sliding_window)
110 -
106 self.is_multimodal_model = model_config.is_multimodal_model
112 - self.block_size = cache_config.block_size
107 self.max_model_len = model_config.max_model_len
114 - self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)
108 self.max_num_tokens = scheduler_config.max_num_batched_tokens
109 self.max_num_reqs = scheduler_config.max_num_seqs
110
111 # Model-related.
119 - self.num_attn_layers = model_config.get_num_layers_by_block_type(
120 - parallel_config, LayerBlockType.attention)
112 self.num_query_heads = model_config.get_num_attention_heads(
113 parallel_config)
123 - self.num_kv_heads = model_config.get_num_kv_heads(parallel_config)
124 - self.head_size = model_config.get_head_size()
114 self.hidden_size = model_config.get_hidden_size()
115 self.attention_chunk_size = model_config.attention_chunk_size
116
128 - self.attn_backend = get_attn_backend(
129 - self.head_size,
130 - self.dtype,
131 - self.kv_cache_dtype,
132 - self.block_size,
133 - self.model_config.is_attention_free,
134 - use_mla=self.model_config.use_mla,
135 - )
136 - if self.attn_backend is None:
137 - error_msg = (
138 - f"Error with get_att_backend: {self.head_size=}, "
139 - f"{self.dtype=}, {self.kv_cache_dtype=}, {self.block_size=}, "
140 - f"{self.model_config.is_attention_free=}, "
141 - f"{self.model_config.use_mla=}")
142 - logger.error(error_msg)
143 - raise NotImplementedError(
144 - "Non-Attention backend is not supported by V1 GPUModelRunner.")
145 -
146 - if self.vllm_config.compilation_config.full_cuda_graph:
147 - attn_backend_name = self.attn_backend.__name__
148 - flash_attn_version = get_flash_attn_version()
149 - if attn_backend_name != "FlashAttentionBackend" or \
150 - flash_attn_version != 3:
151 - raise ValueError(
152 - f"full_cuda_graph is only supported with "
153 - f"FA3. Current attention backend is {attn_backend_name}, "
154 - f"FlashAttention version is {flash_attn_version}.")
155 -
117 self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
118
119 # Multi-modal data support
@@ -174,8 +135,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
135 # self.model: nn.Module # Set after load_model
136 # Initialize in initialize_kv_cache
137 self.kv_caches: list[torch.Tensor] = []
138 + self.attn_metadata_builders: list[AttentionMetadataBuilder] = []
139 + self.attn_backends: list[type[AttentionBackend]] = []
140 # self.kv_cache_config: KVCacheConfig
178 - # self.attn_metadata_builder: type[AttentionMetadataBuilder]
141 + # self.input_batch: InputBatch # Persistent batch.
142
143 # req_id -> (input_id -> encoder_output)
144 self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
@@ -200,16 +163,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
163
164 # Request states.
165 self.requests: dict[str, CachedRequestState] = {}
203 - # Persistent batch.
204 - self.input_batch = InputBatch(
205 - max_num_reqs=self.max_num_reqs,
206 - max_model_len=self.max_model_len,
207 - max_num_blocks_per_req=self.max_num_blocks_per_req,
208 - max_num_batched_tokens=self.max_num_tokens,
209 - device=self.device,
210 - pin_memory=self.pin_memory,
211 - vocab_size=model_config.get_vocab_size(),
212 - )
166
167 self.use_cuda_graph = (self.vllm_config.compilation_config.level
168 == CompilationLevel.PIECEWISE
@@ -304,6 +257,31 @@ class GPUModelRunner(LoRAModelRunnerMixin):
257 pin_memory=self.pin_memory)
258 self.seq_lens_np = self.seq_lens_cpu.numpy()
259
260 + def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> bool:
261 + """
262 + Update the order of requests in the batch based on the attention
263 + backend's needs. For example, some attention backends (namely MLA) may
264 + want to separate requests based on if the attention computation will be
265 + compute-bound or memory-bound.
266 +
267 + Args:
268 + scheduler_output: The scheduler output.
269 +
270 + Returns:
271 + True if the batch was reordered, False otherwise.
272 + """
273 + batch_reordered = self.attn_metadata_builders[0].reorder_batch(
274 + self.input_batch, scheduler_output)
275 +
276 + # For models with multiple KV cache groups, the groups should agree on
277 + # the same order of requests. We ensure this by only allowing the first
278 + # group to reorder the batch and asserting that all other groups do not
279 + # reorder the batch.
280 + for i in range(1, len(self.kv_cache_config.kv_cache_groups)):
281 + assert not self.attn_metadata_builders[i].reorder_batch(
282 + self.input_batch, scheduler_output)
283 + return batch_reordered
284 +
285 def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
286 """Update the cached states and the persistent batch with the scheduler
287 output.
@@ -440,7 +418,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
418 # Update the block IDs.
419 if not req_data.resumed_from_preemption:
420 # Append the new blocks to the existing block IDs.
443 - req_state.block_ids.extend(req_data.new_block_ids)
421 + for i in range(len(self.kv_cache_config.kv_cache_groups)):
422 + req_state.block_ids[i].extend(req_data.new_block_ids[i])
423 else:
424 # The request is resumed from preemption.
425 # Replace the existing block IDs with the new ones.
@@ -498,11 +477,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
477 if removed_req_indices:
478 self.input_batch.condense(removed_req_indices)
479
501 - # Some attention backends (namely MLA) may want to separate requests
502 - # based on if the attention computation will be compute-bound or
503 - # memory-bound. This gives them a hook to do that.
504 - batch_reordered = self.attn_metadata_builder.reorder_batch(
505 - self.input_batch, scheduler_output)
480 + batch_reordered = self._may_reorder_batch(scheduler_output)
481
482 if batch_changed or batch_reordered:
483 self.input_batch.refresh_sampling_metadata()
@@ -570,21 +545,29 @@ class GPUModelRunner(LoRAModelRunnerMixin):
545 torch.from_numpy(token_indices),
546 out=self.input_ids_cpu[:total_num_scheduled_tokens])
547
573 - # Calculate the slot mapping.
574 - # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
575 - # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
576 - # where K is the max_num_blocks_per_req and the block size is 2.
577 - # NOTE(woosuk): We can't simply use `token_indices // block_size` here
578 - # because M (max_model_len) is not necessarily divisible by block_size.
579 - block_table_indices = (req_indices * self.max_num_blocks_per_req +
580 - positions_np // self.block_size)
581 - block_table_cpu = self.input_batch.block_table.get_cpu_tensor()
582 - block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()
583 - block_offsets = positions_np % self.block_size
584 - np.add(block_numbers * self.block_size,
585 - block_offsets,
586 - out=self.input_batch.block_table.
587 - slot_mapping_np[:total_num_scheduled_tokens])
548 + # Calculate the slot mapping for each KV cache group.
549 + for kv_cache_group_id, kv_cache_group_spec in enumerate(
550 + self.kv_cache_config.kv_cache_groups):
551 + block_size = kv_cache_group_spec.kv_cache_spec.block_size
552 + block_table: BlockTable = self.input_batch.block_table[
553 + kv_cache_group_id]
554 + # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
555 + # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
556 + # where K is the max_num_blocks_per_req and the block size is 2.
557 + # NOTE(woosuk): We can't simply use `token_indices // block_size`
558 + # here because M (max_model_len) is not necessarily divisible by
559 + # block_size.
560 + block_table_indices = (
561 + req_indices * block_table.max_num_blocks_per_req +
562 + positions_np // block_size)
563 + block_table_cpu = block_table.get_cpu_tensor()
564 + block_numbers = block_table_cpu.flatten(
565 + )[block_table_indices].numpy()
566 + block_offsets = positions_np % block_size
567 + np.add(
568 + block_numbers * block_size,
569 + block_offsets,
570 + out=block_table.slot_mapping_np[:total_num_scheduled_tokens])
571
572 # Prepare the attention metadata.
573 self.query_start_loc_np[0] = 0
@@ -626,10 +609,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
609 attn_metadata: dict[str, FlashAttentionMetadata] = {}
610 # Prepare the attention metadata for each KV cache group and make layers
611 # in the same group share the same metadata.
629 - # NOTE(Chen): there is exactly one KV cache group that contains all
630 - # attetnion layers in the model for now, so the current logic for
631 - # getting attn_metadata is not related to kv_cache_group information.
632 - # Will extend this part to support multiple KV cache groups later.
612 for kv_cache_group_id, kv_cache_group_spec in enumerate(
613 self.kv_cache_config.kv_cache_groups):
614
@@ -638,15 +617,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
617 if self.cascade_attn_enabled:
618 common_prefix_len = self._compute_cascade_attn_prefix_len(
619 num_scheduled_tokens,
641 - scheduler_output.num_common_prefix_blocks,
620 + scheduler_output.
621 + num_common_prefix_blocks[kv_cache_group_id],
622 + kv_cache_group_spec.kv_cache_spec,
623 + self.attn_metadata_builders[kv_cache_group_id],
624 )
625
644 - attn_metadata_i = self.attn_metadata_builder.build(
645 - num_reqs=num_reqs,
646 - num_actual_tokens=total_num_scheduled_tokens,
647 - max_query_len=max_num_scheduled_tokens,
648 - common_prefix_len=common_prefix_len,
649 - common_attn_metadata=common_attn_metadata)
626 + attn_metadata_i = (
627 + self.attn_metadata_builders[kv_cache_group_id].build(
628 + num_reqs=num_reqs,
629 + num_actual_tokens=total_num_scheduled_tokens,
630 + max_query_len=max_num_scheduled_tokens,
631 + common_prefix_len=common_prefix_len,
632 + common_attn_metadata=common_attn_metadata))
633 for layer_name in kv_cache_group_spec.layer_names:
634 attn_metadata[layer_name] = attn_metadata_i
635
@@ -684,6 +667,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
667 self,
668 num_scheduled_tokens: np.ndarray,
669 num_common_prefix_blocks: int,
670 + kv_cache_spec: KVCacheSpec,
671 + attn_metadata_builder: AttentionMetadataBuilder,
672 ) -> int:
673 """Compute the length of the common prefix for cascade attention.
674
@@ -702,7 +687,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
687 Returns:
688 int: Length of common prefix in tokens.
689 """
705 - common_prefix_len = num_common_prefix_blocks * self.block_size
690 + common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size
691 if common_prefix_len == 0:
692 # Common case.
693 return 0
@@ -751,15 +736,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
736 common_prefix_len,
737 self.input_batch.num_computed_tokens_cpu[:num_reqs].min())
738 # common_prefix_len should be a multiple of the block size.
754 - common_prefix_len = (common_prefix_len // self.block_size *
755 - self.block_size)
756 - use_cascade = self.attn_metadata_builder.use_cascade_attention(
739 + common_prefix_len = (common_prefix_len // kv_cache_spec.block_size *
740 + kv_cache_spec.block_size)
741 + use_sliding_window = (isinstance(kv_cache_spec, SlidingWindowSpec) or
742 + (isinstance(kv_cache_spec, FullAttentionSpec)
743 + and kv_cache_spec.sliding_window is not None))
744 + assert isinstance(kv_cache_spec, AttentionSpec)
745 + use_cascade = attn_metadata_builder.use_cascade_attention(
746 common_prefix_len=common_prefix_len,
747 query_lens=num_scheduled_tokens,
748 num_query_heads=self.num_query_heads,
760 - num_kv_heads=self.num_kv_heads,
749 + num_kv_heads=kv_cache_spec.num_kv_heads,
750 use_alibi=self.use_alibi,
762 - use_sliding_window=self.window_size is not None,
751 + use_sliding_window=use_sliding_window,
752 num_sms=self.num_sms,
753 )
754 return common_prefix_len if use_cascade else 0
@@ -1577,7 +1566,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1566 dtype=np.int32)
1567
1568 if skip_attn:
1580 - attn_metadata = None
1569 + attn_metadata: Optional[dict[str, FlashAttentionMetadata]] = None
1570 else:
1571 query_start_loc = self.query_start_loc[:num_reqs + 1]
1572 seq_lens = self.seq_lens[:num_reqs]
@@ -1585,13 +1574,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1574 common_attn_metadata = CommonAttentionMetadata(
1575 query_start_loc=query_start_loc, seq_lens=seq_lens)
1576
1588 - attn_metadata = self.attn_metadata_builder.build(
1589 - num_reqs=num_tokens,
1590 - num_actual_tokens=num_tokens,
1591 - max_query_len=num_tokens,
1592 - common_prefix_len=0,
1593 - common_attn_metadata=common_attn_metadata,
1594 - )
1577 + attn_metadata = {}
1578 + for kv_cache_group_id, kv_cache_group_spec in enumerate(
1579 + self.kv_cache_config.kv_cache_groups):
1580 + attn_metadata_i = (
1581 + self.attn_metadata_builders[kv_cache_group_id].build(
1582 + num_reqs=num_tokens,
1583 + num_actual_tokens=num_tokens,
1584 + max_query_len=num_tokens,
1585 + common_prefix_len=0,
1586 + common_attn_metadata=common_attn_metadata,
1587 + ))
1588 + for layer_name in kv_cache_group_spec.layer_names:
1589 + attn_metadata[layer_name] = attn_metadata_i
1590
1591 with self.maybe_dummy_run_with_lora(self.lora_config,
1592 num_scheduled_tokens):
@@ -1822,6 +1817,56 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1817 logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
1818 elapsed_time, cuda_graph_size / (1 << 30))
1819
1820 + def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
1821 + """
1822 + Initialize the attention backends and attention metadata builders.
1823 + """
1824 + assert len(self.attn_backends) == 0 and len(
1825 + self.attn_metadata_builders
1826 + ) == 0, "Attention backends are already initialized"
1827 + for i, kv_cache_group_spec in enumerate(
1828 + kv_cache_config.kv_cache_groups):
1829 + kv_cache_spec = kv_cache_group_spec.kv_cache_spec
1830 + if not isinstance(kv_cache_spec, AttentionSpec):
1831 + raise NotImplementedError(
1832 + "Only AttentionSpec is supported for now.")
1833 + attn_backend_i = get_attn_backend(
1834 + kv_cache_spec.head_size,
1835 + self.dtype,
1836 + kv_cache_spec.dtype,
1837 + kv_cache_spec.block_size,
1838 + self.model_config.is_attention_free,
1839 + use_mla=kv_cache_spec.use_mla,
1840 + )
1841 + if attn_backend_i is None:
1842 + error_msg = (
1843 + f"Error with get_attn_backend: {kv_cache_spec.head_size=}, "
1844 + f"{self.dtype=}, {kv_cache_spec.dtype=}, "
1845 + f"{kv_cache_spec.block_size=}, "
1846 + f"{self.model_config.is_attention_free=}, "
1847 + f"{kv_cache_spec.use_mla=}")
1848 + logger.error(error_msg)
1849 + raise NotImplementedError(
1850 + "Non-Attention backend is not supported by V1 "
1851 + "GPUModelRunner.")
1852 +
1853 + if self.vllm_config.compilation_config.full_cuda_graph:
1854 + attn_backend_name = attn_backend_i.__name__
1855 + flash_attn_version = get_flash_attn_version()
1856 + if attn_backend_name != "FlashAttentionBackend" or \
1857 + flash_attn_version != 3:
1858 + raise ValueError(
1859 + f"full_cuda_graph is only supported with "
1860 + f"FA3. Current attention backend is "
1861 + f"{attn_backend_name}, FlashAttention version is "
1862 + f"{flash_attn_version}.")
1863 +
1864 + block_table_i = self.input_batch.block_table[i]
1865 + attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
1866 + weakref.proxy(self), kv_cache_spec, block_table_i)
1867 + self.attn_backends.append(attn_backend_i)
1868 + self.attn_metadata_builders.append(attn_metadata_builder_i)
1869 +
1870 def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
1871 """
1872 Initialize KV cache based on `kv_cache_config`.
@@ -1829,15 +1874,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1874 kv_cache_config: Configuration for the KV cache, including the KV
1875 cache size of each layer
1876 """
1832 - if len(kv_cache_config.kv_cache_groups) > 1:
1833 - raise NotImplementedError(
1834 - "Hybrid models with more than one KV cache type are not "
1835 - "supported yet.")
1877 self.kv_cache_config = kv_cache_config
1878 + self.input_batch = InputBatch(
1879 + max_num_reqs=self.max_num_reqs,
1880 + max_model_len=self.max_model_len,
1881 + max_num_batched_tokens=self.max_num_tokens,
1882 + device=self.device,
1883 + pin_memory=self.pin_memory,
1884 + vocab_size=self.model_config.get_vocab_size(),
1885 + kv_cache_config=kv_cache_config,
1886 + )
1887 + self.initialize_attn_backend(kv_cache_config)
1888
1889 kv_caches: dict[str, torch.Tensor] = {}
1890
1840 - for kv_cache_group in kv_cache_config.kv_cache_groups:
1891 + for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
1892 kv_cache_spec = kv_cache_group.kv_cache_spec
1893 for layer_name in kv_cache_group.layer_names:
1894 tensor_config = kv_cache_config.tensors[layer_name]
@@ -1852,7 +1903,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1903 # the min of all `num_blocks`. Verify it here.
1904 assert num_blocks >= kv_cache_config.num_blocks
1905 if isinstance(kv_cache_spec, AttentionSpec):
1855 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(
1906 + kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(
1907 num_blocks, kv_cache_spec.block_size,
1908 kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
1909 dtype = kv_cache_spec.dtype
@@ -1872,11 +1923,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1923 if has_kv_transfer_group():
1924 get_kv_transfer_group().register_kv_caches(kv_caches)
1925
1875 - self.attn_metadata_builder = self.attn_backend.get_builder_cls()(
1876 - weakref.proxy(self),
1877 - kv_cache_config.kv_cache_groups[0].kv_cache_spec,
1878 - self.input_batch.block_table)
1879 -
1926 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
1927 """
1928 Generates the KVCacheSpec by parsing the kv cache format from each
vllm/v1/worker/tpu_model_runner.py
+19 -16
@@ -171,19 +171,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
171 self.kv_caches: list[torch.Tensor] = []
172 # req_id -> (input_id -> encoder_output)
173 self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
174 + # self.input_batch: InputBatch # Persistent batch.
175
176 # Request states.
177 self.requests: dict[str, CachedRequestState] = {}
177 - # Persistent batch.
178 - self.input_batch = InputBatch(
179 - max_num_reqs=self.max_num_reqs,
180 - max_model_len=self.max_model_len,
181 - max_num_blocks_per_req=self.max_num_blocks_per_req,
182 - max_num_batched_tokens=self.max_num_tokens,
183 - device=self.device,
184 - pin_memory=self.pin_memory,
185 - vocab_size=self.vocab_size,
186 - )
178
179 # Cached torch/numpy tensor
180 # The pytorch tensor and numpy array share the same buffer.
@@ -199,7 +190,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
190
191 self.block_table_cpu = torch.zeros(
192 (self.max_num_reqs, self.max_num_blocks_per_req),
202 - dtype=self.input_batch.block_table.get_cpu_tensor().dtype,
193 + dtype=torch.int32,
194 device="cpu")
195
196 self.query_start_loc_cpu = torch.zeros(self.max_num_tokens + 1,
@@ -524,12 +515,12 @@ class TPUModelRunner(LoRAModelRunnerMixin):
515 # NOTE(woosuk): We use torch.index_select instead of np.take here
516 # because torch.index_select is much faster than np.take for large
517 # tensors.
527 - block_table_cpu = self.input_batch.block_table.get_cpu_tensor()
518 + block_table_cpu = self.input_batch.block_table[0].get_cpu_tensor()
519 block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()
520 block_offsets = positions_np % self.block_size
521 np.add(block_numbers * self.block_size,
522 block_offsets,
532 - out=self.input_batch.block_table.
523 + out=self.input_batch.block_table[0].
524 slot_mapping_np[:total_num_scheduled_tokens])
525
526 # Prepare the attention metadata.
@@ -554,15 +545,15 @@ class TPUModelRunner(LoRAModelRunnerMixin):
545 self.position_ids = self.positions_cpu[:
546 padded_total_num_scheduled_tokens].to(
547 self.device)
557 - self.input_batch.block_table.slot_mapping_cpu[
548 + self.input_batch.block_table[0].slot_mapping_cpu[
549 total_num_scheduled_tokens:] = _PAD_SLOT_ID
550 slot_mapping = (
560 - self.input_batch.block_table.
551 + self.input_batch.block_table[0].
552 slot_mapping_cpu[:padded_total_num_scheduled_tokens].to(
553 self.device))
554 block_tables = self.block_table_cpu[:self.max_num_reqs]
555 block_tables[:num_reqs, :self.max_num_blocks_per_req] = (
565 - self.input_batch.block_table.get_cpu_tensor()[:num_reqs])
556 + self.input_batch.block_table[0].get_cpu_tensor()[:num_reqs])
557 block_tables = block_tables.to(self.device)
558 query_start_loc = self.query_start_loc_cpu[:self.max_num_reqs + 1].to(
559 self.device)
@@ -1263,6 +1254,18 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1254 "Hybrid models with more than one KV cache type are not "
1255 "supported yet.")
1256
1257 + self.input_batch = InputBatch(
1258 + max_num_reqs=self.max_num_reqs,
1259 + max_model_len=self.max_model_len,
1260 + max_num_batched_tokens=self.max_num_tokens,
1261 + device=self.device,
1262 + pin_memory=self.pin_memory,
1263 + vocab_size=self.model_config.get_vocab_size(),
1264 + kv_cache_config=kv_cache_config,
1265 + )
1266 + assert self.block_table_cpu.dtype == self.input_batch.block_table[
1267 + 0].get_cpu_tensor().dtype
1268 +
1269 kv_caches: dict[str, torch.Tensor] = {}
1270
1271 for kv_cache_group in kv_cache_config.kv_cache_groups: