@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 646d62f636

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, 06:00 UTC

[Core] Use tuple for kv cache group block ids (#19175)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Jun 9, 2025 at 22:01 UTC 646d62f636197d5d809d21090752700e094b8b86

12 files changed +140 -142

tests/v1/core/test_prefix_caching.py

+22 -22

@@ -117,7 +117,7 @@ def test_prefill(hash_algo):

117 blocks = manager.allocate_slots(req0, 55,

118 len(computed_blocks.blocks[0]) * 16,

119 computed_blocks)

120 - assert blocks.get_block_ids() == [[1, 2, 3, 4]]

120 + assert blocks.get_block_ids() == ([1, 2, 3, 4], )

121

122 # Check full block metadata

123 parent_block_hash = None

@@ -141,13 +141,13 @@ def test_prefill(hash_algo):

141 req1 = make_request("1", common_token_ids + unique_token_ids)

142 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

143 assert len(manager.req_to_block_hashes[req1.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(req1, num_new_tokens,

148 len(computed_blocks.blocks[0]) * 16,

149 computed_blocks)

150 - assert blocks.get_block_ids() == [[5]]

150 + assert blocks.get_block_ids() == ([5], )

151 for block in computed_blocks.blocks[0]:

152 assert block.ref_cnt == 2

153

@@ -175,13 +175,13 @@ def test_prefill(hash_algo):

175 req2 = make_request("2", common_token_ids + unique_token_ids)

176 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

177 assert len(manager.req_to_block_hashes[req2.request_id]) == 3

178 - assert computed_blocks.get_block_ids() == [[1, 2, 3]]

178 + assert computed_blocks.get_block_ids() == ([1, 2, 3], )

179 assert num_computed_tokens == 3 * 16

180 num_new_tokens = 53 - 3 * 16

181 blocks = manager.allocate_slots(req2, num_new_tokens,

182 len(computed_blocks.blocks[0]) * 16,

183 computed_blocks)

184 - assert blocks.get_block_ids() == [[6]]

184 + assert blocks.get_block_ids() == ([6], )

185

186 # Although we only have 6 free blocks, we have 8 blocks in

187 # the free block queue due to lazy removal.

@@ -205,7 +205,7 @@ def test_prefill(hash_algo):

205 len(computed_blocks.blocks[0]) * 16,

206 computed_blocks)

207 # This block ID order also checks the eviction order.

208 - assert blocks.get_block_ids() == [[7, 8, 9, 10, 4, 5, 6, 3, 2, 1]]

208 + assert blocks.get_block_ids() == ([7, 8, 9, 10, 4, 5, 6, 3, 2, 1], )

209 assert manager.block_pool.free_block_queue.num_free_blocks == 0

210 assert manager.block_pool.free_block_queue.free_list_head is None

211 assert manager.block_pool.free_block_queue.free_list_tail is None

@@ -236,8 +236,8 @@ def test_prefill_hybrid_model():

236 blocks = manager.allocate_slots(req0, 55,

237 len(computed_blocks.blocks[0]) * 16,

238 computed_blocks)

239 - assert blocks.get_block_ids() == [[1, 2, 3, 4], [5, 6, 7, 8],

240 - [9, 10, 11, 12]]

239 + assert blocks.get_block_ids() == ([1, 2, 3, 4], [5, 6, 7,

240 + 8], [9, 10, 11, 12])

241

242 # Check full block metadata

243 parent_block_hash = None

@@ -263,14 +263,14 @@ def test_prefill_hybrid_model():

263 req1 = make_request("1", common_token_ids + unique_token_ids)

264 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

265 assert len(manager.req_to_block_hashes[req1.request_id]) == 3

266 - assert computed_blocks.get_block_ids() == [[1, 2, 3], [0, 6, 7],

267 - [0, 10, 11]]

266 + assert computed_blocks.get_block_ids() == ([1, 2, 3], [0, 6,

267 + 7], [0, 10, 11])

268 assert num_computed_tokens == 3 * 16

269 num_new_tokens = 53 - 3 * 16

270 blocks = manager.allocate_slots(req1, num_new_tokens,

271 len(computed_blocks.blocks[0]) * 16,

272 computed_blocks)

273 - assert blocks.get_block_ids() == [[13], [14], [15]]

273 + assert blocks.get_block_ids() == ([13], [14], [15])

274 for block_per_group in computed_blocks.blocks:

275 for block in block_per_group:

276 if block != manager.block_pool.null_block:

@@ -374,7 +374,7 @@ def test_prefill_plp():

374 blocks = manager.allocate_slots(req0, 55,

375 len(computed_blocks.blocks[0]) * 16,

376 computed_blocks)

377 - assert blocks.get_block_ids() == [[1, 2, 3, 4]]

377 + assert blocks.get_block_ids() == ([1, 2, 3, 4], )

378 req0_block_hashes = [b.block_hash for b in blocks.blocks[0]]

379

380 # Check full block metadata

@@ -400,13 +400,13 @@ def test_prefill_plp():

400 req1 = make_request("1", common_token_ids + unique_token_ids)

401 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

402 assert len(manager.req_to_block_hashes[req1.request_id]) == 3

403 - assert computed_blocks.get_block_ids() == [[1, 2, 3]]

403 + assert computed_blocks.get_block_ids() == ([1, 2, 3], )

404 assert num_computed_tokens == 3 * 16

405 num_new_tokens = 53 - 3 * 16

406 blocks = manager.allocate_slots(req1, num_new_tokens,

407 len(computed_blocks.blocks[0]) * 16,

408 computed_blocks)

409 - assert blocks.get_block_ids() == [[5]]

409 + assert blocks.get_block_ids() == ([5], )

410 for block in computed_blocks.blocks[0]:

411 assert block.ref_cnt == 2

412

@@ -444,7 +444,7 @@ def test_prefill_plp():

444 block_ids = blocks.get_block_ids()

445 # Duplicate cached blocks have different ids but same hashes vs request #0

446 assert [b.block_hash for b in blocks.blocks[0]] == req0_block_hashes

447 - assert block_ids != [[1, 2, 3, 4]]

447 + assert block_ids != ([1, 2, 3, 4], )

448

449 # Request #2 block hashes are valid since request #0 hashes are.

450 # Check block reference counts.

@@ -474,7 +474,7 @@ def test_decode():

474 blocks = manager.allocate_slots(req0, 55,

475 len(computed_blocks.blocks[0]) * 16,

476 computed_blocks)

477 - assert blocks.get_block_ids() == [[1, 2, 3, 4]]

477 + assert blocks.get_block_ids() == ([1, 2, 3, 4], )

478

479 # Append slots without allocating a new block.

480 req0.num_computed_tokens = 55

@@ -546,12 +546,12 @@ def test_evict():

546 # Touch the first 2 blocks.

547 req2 = make_request("2", list(range(2 * 16 + 3)))

548 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

549 - assert computed_blocks.get_block_ids() == [[1, 2]]

549 + assert computed_blocks.get_block_ids() == ([1, 2], )

550 assert num_computed_tokens == 2 * 16

551 blocks = manager.allocate_slots(req2, 3,

552 len(computed_blocks.blocks[0]) * 16,

553 computed_blocks)

554 - assert blocks.get_block_ids() == [[10]]

554 + assert blocks.get_block_ids() == ([10], )

555 assert manager.block_pool.free_block_queue.num_free_blocks == 7

556

557

@@ -865,7 +865,7 @@ def test_mm_prefix_caching():

865 blocks = manager.allocate_slots(req0, 59,

866 len(computed_blocks.blocks[0]) * 16,

867 computed_blocks)

868 - assert blocks.get_block_ids() == [[1, 2, 3, 4]]

868 + assert blocks.get_block_ids() == ([1, 2, 3, 4], )

869 req0.num_computed_tokens = 59

870

871 # Append slots without allocating a new block.

@@ -926,7 +926,7 @@ def test_cache_key_salting():

926 blocks = manager.allocate_slots(req0, 59,

927 len(computed_blocks.blocks[0]) * 16,

928 computed_blocks)

929 - assert blocks.get_block_ids() == [[1, 2, 3, 4]]

929 + assert blocks.get_block_ids() == ([1, 2, 3, 4], )

930 req0.num_computed_tokens = 59

931

932 # Append slots without allocating a new block.

@@ -1042,7 +1042,7 @@ def test_reset_prefix_cache():

1042 all_token_ids = full_block_token_ids + unique_token_ids

1043 req0 = make_request("0", all_token_ids)

1044 blocks = manager.allocate_slots(req0, 55)

1045 - assert blocks.get_block_ids() == [[1, 2, 3, 4]]

1045 + assert blocks.get_block_ids() == ([1, 2, 3, 4], )

1046

1047 unique_token_ids = [4] * 7

1048 all_token_ids = full_block_token_ids + unique_token_ids

@@ -1053,7 +1053,7 @@ def test_reset_prefix_cache():

1053 blocks = manager.allocate_slots(req1, 7,

1054 len(computed_blocks.blocks[0]) * 16,

1055 computed_blocks)

1056 - assert blocks.get_block_ids() == [[5]]

1056 + assert blocks.get_block_ids() == ([5], )

1057

1058 # Failed to reset prefix cache because some blocks are not freed yet.

1059 assert not manager.reset_prefix_cache()

tests/v1/tpu/worker/test_tpu_model_runner.py

+4 -4

@@ -71,7 +71,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

71 mm_hashes=[],

72 mm_positions=[],

73 sampling_params=SamplingParams(),

74 - block_ids=[[0]], # block_ids should be list[list[int]]

74 + block_ids=([0], ), # block_ids should be tuple[list[int]]

75 num_computed_tokens=0,

76 lora_request=None,

77 ))

@@ -116,10 +116,10 @@ def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:

116 # This is safe since we currently only use single KV cache groups

117 block_table = multi_group_block_table[0]

118

119 - # req_state.block_ids is now list[list[int]] for MultiGroupBlockTable

119 + # req_state.block_ids is now tuple[list[int], ...] for MultiGroupBlockTable

120 # Extract the first group's block IDs

121 if isinstance(req_state.block_ids[0], list):

122 - # New format: list[list[int]] - extract first group

122 + # New format: tuple[list[int], ...] - extract first group

123 req_block_ids = req_state.block_ids[0]

124 else:

125 # Legacy format: list[int] - use directly

@@ -210,7 +210,7 @@ def test_update_states_request_resumed(model_runner):

210 req_id=req_id,

211 resumed_from_preemption=False,

212 new_token_ids=[],

213 - new_block_ids=[[]],

213 + new_block_ids=([], ),

214 num_computed_tokens=0,

215 )

216

tests/v1/worker/test_gpu_input_batch.py

+1 -1

@@ -203,7 +203,7 @@ def _construct_cached_request_state(req_id_suffix: int):

203 sampling_params=_create_sampling_params(),

204 mm_inputs=[],

205 mm_positions=[],

206 - block_ids=[[]],

206 + block_ids=([], ),

207 generator=None,

208 num_computed_tokens=len(output_token_ids),

209 output_token_ids=output_token_ids,

tests/v1/worker/test_gpu_model_runner.py

+2 -2

@@ -123,7 +123,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

123 mm_hashes=[],

124 mm_positions=[],

125 sampling_params=SamplingParams(),

126 - block_ids=[[0]],

126 + block_ids=([0], ),

127 num_computed_tokens=0,

128 lora_request=None,

129 ))

@@ -251,7 +251,7 @@ def test_update_states_request_resumed(model_runner):

251 req_id=req_id,

252 resumed_from_preemption=False,

253 new_token_ids=[],

254 - new_block_ids=[[]],

254 + new_block_ids=([], ),

255 num_computed_tokens=0,

256 )

257

vllm/v1/core/block_pool.py

+4 -4

@@ -89,8 +89,8 @@ class BlockPool:

89 BlockHashWithGroupId(block_hash, group_id))

90 if not cached_blocks_one_group:

91 return None

92 - first_block_id = next(iter(cached_blocks_one_group))

93 - cached_blocks.append(cached_blocks_one_group[first_block_id])

92 + first_block = next(iter(cached_blocks_one_group.values()))

93 + cached_blocks.append(first_block)

94 return cached_blocks

96 def cache_full_blocks(

@@ -260,7 +260,7 @@ class BlockPool:

260 return True

261 return False

262

263 - def touch(self, blocks: list[list[KVCacheBlock]]) -> None:

263 + def touch(self, blocks: tuple[list[KVCacheBlock], ...]) -> None:

264 """Touch a block increases its reference count by 1, and may remove

265 the block from the free queue. This is used when a block is hit by

266 another request with the same prefix.

@@ -299,7 +299,7 @@ class BlockPool:

299 bool: True if the prefix cache is successfully reset,

300 False otherwise.

301 """

302 - num_used_blocks = (self.num_gpu_blocks - self.get_num_free_blocks())

302 + num_used_blocks = self.num_gpu_blocks - self.get_num_free_blocks()

303 if num_used_blocks != 1: # The null block is always marked as used

304 logger.warning(

305 "Failed to reset prefix cache because some "

vllm/v1/core/kv_cache_coordinator.py

+53 -48

@@ -5,8 +5,7 @@ from typing import Callable, Optional

5 from vllm.v1.core.block_pool import BlockPool

6 from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

7 from vllm.v1.core.single_type_kv_cache_manager import (

8 - FullAttentionManager, SingleTypeKVCacheManager,

9 - get_manager_for_kv_cache_spec)

8 + FullAttentionManager, get_manager_for_kv_cache_spec)

9 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig

10 from vllm.v1.request import Request

@@ -30,25 +29,21 @@ class KVCacheCoordinator(ABC):

30 self.block_pool = BlockPool(kv_cache_config.num_blocks, enable_caching,

31 enable_kv_cache_events)

33 - self.single_type_managers: list[SingleTypeKVCacheManager] = []

33 # Needs special handling for find_longest_cache_hit if eagle is enabled

34 self.use_eagle = use_eagle

37 -

38 - for i in range(len(self.kv_cache_config.kv_cache_groups)):

39 - kv_cache_spec = self.kv_cache_config.kv_cache_groups[

40 - i].kv_cache_spec

41 - self.single_type_managers.append(

42 - get_manager_for_kv_cache_spec(

43 - kv_cache_spec=kv_cache_spec,

44 - block_pool=self.block_pool,

45 - kv_cache_group_id=i,

46 - caching_hash_fn=caching_hash_fn,

47 - ))

35 + self.single_type_managers = tuple(

36 + get_manager_for_kv_cache_spec(

37 + kv_cache_spec=kv_cache_group.kv_cache_spec,

38 + block_pool=self.block_pool,

39 + kv_cache_group_id=i,

40 + caching_hash_fn=caching_hash_fn,

41 + ) for i, kv_cache_group in enumerate(

42 + self.kv_cache_config.kv_cache_groups))

44 def get_num_blocks_to_allocate(

45 self, request_id: str, num_tokens: int,

51 - new_computed_blocks: list[list[KVCacheBlock]]) -> int:

46 + new_computed_blocks: tuple[list[KVCacheBlock], ...]) -> int:

47 """

48 Get the number of blocks needed to be allocated for the request.

@@ -70,7 +65,7 @@ class KVCacheCoordinator(ABC):

66 def save_new_computed_blocks(

67 self, request_id: str,

73 - new_computed_blocks: list[list[KVCacheBlock]]) -> None:

68 + new_computed_blocks: tuple[list[KVCacheBlock], ...]) -> None:

69 """

70 Add the new computed blocks to the request.

@@ -84,7 +79,7 @@ class KVCacheCoordinator(ABC):

79 new_computed_blocks[i])

81 def allocate_new_blocks(self, request_id: str,

87 - num_tokens: int) -> list[list[KVCacheBlock]]:

82 + num_tokens: int) -> tuple[list[KVCacheBlock], ...]:

83 """

84 Allocate new blocks for the request to give it at least `num_tokens`

85 token slots.

@@ -97,11 +92,9 @@ class KVCacheCoordinator(ABC):

92 Returns:

93 The new allocated blocks.

94 """

100 - new_blocks = []

101 - for manager in self.single_type_managers:

102 - new_blocks.append(

103 - manager.allocate_new_blocks(request_id, num_tokens))

104 - return new_blocks

95 + return tuple(

96 + manager.allocate_new_blocks(request_id, num_tokens)

97 + for manager in self.single_type_managers)

99 def cache_blocks(self, request: Request, block_hashes: list[BlockHash],

100 num_computed_tokens: int) -> None:

@@ -159,19 +152,20 @@ class KVCacheCoordinator(ABC):

152 for manager in self.single_type_managers:

153 manager.remove_skipped_blocks(request_id, num_computed_tokens)

154

162 - def get_blocks(self, request_id: str) -> list[list[KVCacheBlock]]:

155 + def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:

156 """

157 Get the blocks for the request.

158 """

166 - return [

159 + return tuple(

160 manager.req_to_blocks.get(request_id) or []

168 - for manager in self.single_type_managers

169 - ]

161 + for manager in self.single_type_managers)

162

163 @abstractmethod

164 def find_longest_cache_hit(

173 - self, block_hashes: list[BlockHash],

174 - max_cache_hit_length: int) -> tuple[list[list[KVCacheBlock]], int]:

165 + self,

166 + block_hashes: list[BlockHash],

167 + max_cache_hit_length: int,

168 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:

169 pass

170

171

@@ -195,8 +189,10 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

189 "UnitaryKVCacheCoordinator assumes only one kv cache group")

190

191 def find_longest_cache_hit(

198 - self, block_hashes: list[BlockHash],

199 - max_cache_hit_length: int) -> tuple[list[list[KVCacheBlock]], int]:

192 + self,

193 + block_hashes: list[BlockHash],

194 + max_cache_hit_length: int,

195 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:

196 hit_blocks = self.single_type_managers[0].find_longest_cache_hit(

197 block_hashes=block_hashes,

198 max_length=max_cache_hit_length,

@@ -275,11 +271,24 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

271 "KVCacheCoordinator assumes the block_size of full attention "

272 "layers is divisible by other layers now.")

273

274 + if max(self.full_attention_group_ids) < min(self.other_group_ids):

275 + self.full_attn_first = True

276 + elif max(self.other_group_ids) < min(self.full_attention_group_ids):

277 + self.full_attn_first = False

278 + else:

279 + raise ValueError(

280 + "HybridKVCacheCoordinator assumes the full "

281 + "attention group ids and other attention group ids "

282 + "do not interleave, either full attention group ids "

283 + "are before other attention group ids or vice versa."

284 + "This is for simplifying merging hit_blocks_full_attn and "

285 + "hit_blocks_other_attn to hit_blocks.")

286 +

287 def find_longest_cache_hit(

288 self,

289 block_hashes: list[BlockHash],

290 max_cache_hit_length: int,

282 - ) -> tuple[list[list[KVCacheBlock]], int]:

291 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:

292 """

293 Find the longest cache hit for the request.

294

@@ -318,27 +327,25 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

327 ))

328 hit_length = len(hit_blocks_other_attn[0]) * self.other_block_size

329

321 - # NOTE: the prefix cache hit length must be a multiply of block_size as

330 + # NOTE: the prefix cache hit length must be a multiple of block_size as

331 # we don't support partial block cache hit yet. The cache hit length

323 - # of other attention is ensured to be a multiply of the block size of

332 + # of other attention is ensured to be a multiple of the block size of

333 # full attention layers in current implementation, because hit_length is

325 - # a multiply of other attention's block size, and other attention's

326 - # block size is a multiply of full attention's block size (verified in

334 + # a multiple of other attention's block size, and other attention's

335 + # block size is a multiple of full attention's block size (verified in

336 # `verify_and_split_kv_cache_groups`).

337 assert hit_length % self.full_attention_block_size == 0

338

339 # Truncate the full attention cache hit to the length of the

340 # cache hit of the other attention.

332 - for i in range(len(hit_blocks_full_attn)):

333 - del hit_blocks_full_attn[i][hit_length //

334 - self.full_attention_block_size:]

341 + for group_hit_blocks in hit_blocks_full_attn:

342 + del group_hit_blocks[hit_length // self.full_attention_block_size:]

343

344 # Merge the hit blocks of full attention and other attention.

337 - hit_blocks = hit_blocks_other_attn

338 - for group_id, blocks in enumerate(hit_blocks_full_attn):

339 - # NOTE: there is only one full attention group in most cases. So

340 - # the time complexity of insert is fine.

341 - hit_blocks.insert(group_id, blocks)

345 + if self.full_attn_first:

346 + hit_blocks = hit_blocks_full_attn + hit_blocks_other_attn

347 + else:

348 + hit_blocks = hit_blocks_other_attn + hit_blocks_full_attn

349 return hit_blocks, hit_length

350

351

@@ -351,8 +358,6 @@ def get_kv_cache_coordinator(

358 use_eagle, enable_caching,

359 caching_hash_fn,

360 enable_kv_cache_events)

354 - else:

355 - return HybridKVCacheCoordinator(kv_cache_config, max_model_len,

356 - use_eagle, enable_caching,

357 - caching_hash_fn,

358 - enable_kv_cache_events)

361 + return HybridKVCacheCoordinator(kv_cache_config, max_model_len, use_eagle,

362 + enable_caching, caching_hash_fn,

363 + enable_kv_cache_events)

vllm/v1/core/kv_cache_manager.py

+18 -20

@@ -21,11 +21,11 @@ logger = init_logger(__name__)

21 @dataclass

22 class KVCacheBlocks:

23 """

24 - The allocation result of KVCacheManager, work as the interface between

25 - Scheduler and KVCacheManager, to hide KVCacheManager's internal data

24 + The allocation result of KVCacheManager, work as the interface between

25 + Scheduler and KVCacheManager, to hide KVCacheManager's internal data

26 structure from the Scheduler.

27 """

28 - blocks: list[list[KVCacheBlock]]

28 + blocks: tuple[list[KVCacheBlock], ...]

29 """

30 blocks[i][j] refers to the i-th kv_cache_group and the j-th block of tokens.

31 We don't use block of tokens as the outer dimension because it assumes all

@@ -37,21 +37,19 @@ class KVCacheBlocks:

37 def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":

38 """Adds two KVCacheBlocks instances."""

39 return KVCacheBlocks(

40 - [blk1 + blk2 for blk1, blk2 in zip(self.blocks, other.blocks)])

40 + tuple(blk1 + blk2

41 + for blk1, blk2 in zip(self.blocks, other.blocks)))

42 - def get_block_ids(self) -> list[list[int]]:

43 + def get_block_ids(self) -> tuple[list[int], ...]:

44 """

45 Converts the KVCacheBlocks instance to block_ids.

47 Returns:

47 - list[list[int]]: A two-level list where

48 - * the outer list corresponds to KV cache groups

48 + tuple[list[int], ...]: A tuple of lists where

49 + * the outer tuple corresponds to KV cache groups

50 * each inner list contains the block_ids of the blocks in that group

51 """

51 - block_ids = []

52 - for group in self.blocks:

53 - block_ids.append([blk.block_id for blk in group])

54 - return block_ids

52 + return tuple([blk.block_id for blk in group] for group in self.blocks)

54 def get_unhashed_block_ids(self) -> list[int]:

55 """Get block_ids of unhashed blocks from KVCacheBlocks instance."""

@@ -63,7 +61,7 @@ class KVCacheBlocks:

62 def new_empty(self) -> "KVCacheBlocks":

63 """Creates a new KVCacheBlocks instance with no blocks."""

66 - return KVCacheBlocks([[] for _ in range(len(self.blocks))])

64 + return KVCacheBlocks(tuple([] for _ in range(len(self.blocks))))

67 class KVCacheManager:

@@ -232,9 +230,8 @@ class KVCacheManager:

230 if new_computed_blocks is not None:

231 new_computed_block_list = new_computed_blocks.blocks

232 else:

235 - new_computed_block_list = [

236 - [] for _ in range(len(self.kv_cache_config.kv_cache_groups))

237 - ]

233 + new_computed_block_list = tuple(

234 + [] for _ in range(len(self.kv_cache_config.kv_cache_groups)))

235

236 # Free the blocks that are skipped during the attention computation

237 # (e.g., tokens outside the sliding window).

@@ -267,7 +264,7 @@ class KVCacheManager:

264 if self.enable_caching:

265 self.block_pool.touch(new_computed_block_list)

266 else:

270 - assert all(not blocks for blocks in new_computed_block_list), (

267 + assert not any(new_computed_block_list), (

268 "Computed blocks should be empty when "

269 "prefix caching is disabled")

270

@@ -378,17 +375,18 @@ class KVCacheManager:

375 """

376 return self.block_pool.take_events()

377

381 - def get_block_ids(self, request_id: str) -> list[list[int]]:

378 + def get_block_ids(self, request_id: str) -> tuple[list[int], ...]:

379 """Get the block ids of a request."""

380 return KVCacheBlocks(

381 self.coordinator.get_blocks(request_id)).get_block_ids()

382

386 - def cache_blocks(self, request: Request, block_hashes: list[BlockHash],

387 - num_computed_tokens: int) -> None:

383 + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:

384 """Cache the blocks for the request."""

385 + block_hashes = self.req_to_block_hashes[request.request_id]

386 self.coordinator.cache_blocks(request, block_hashes,

387 num_computed_tokens)

388

389 def create_empty_block_list(self) -> KVCacheBlocks:

390 """Creates a new KVCacheBlocks instance with no blocks."""

394 - return KVCacheBlocks([[] for _ in range(self.num_kv_cache_groups)])

391 + return KVCacheBlocks(tuple([]

392 + for _ in range(self.num_kv_cache_groups)))

vllm/v1/core/sched/output.py

+4 -4

@@ -27,7 +27,7 @@ class NewRequestData:

27 mm_hashes: list[str]

28 mm_positions: list[PlaceholderRange]

29 sampling_params: SamplingParams

30 - block_ids: list[list[int]]

30 + block_ids: tuple[list[int], ...]

31 num_computed_tokens: int

32 lora_request: Optional[LoRARequest]

@@ -35,7 +35,7 @@ class NewRequestData:

35 def from_request(

36 cls,

37 request: Request,

38 - block_ids: list[list[int]],

38 + block_ids: tuple[list[int], ...],

39 ) -> NewRequestData:

40 return cls(

41 req_id=request.request_id,

@@ -86,7 +86,7 @@ class CachedRequestData:

86 # request's block IDs instead of appending to the existing block IDs.

87 resumed_from_preemption: bool

88 new_token_ids: list[int]

89 - new_block_ids: list[list[int]]

89 + new_block_ids: tuple[list[int], ...]

90 num_computed_tokens: int

92 @classmethod

@@ -95,7 +95,7 @@ class CachedRequestData:

95 request: Request,

96 resumed_from_preemption: bool,

97 new_token_ids: list[int],

98 - new_block_ids: list[list[int]],

98 + new_block_ids: tuple[list[int], ...],

99 ) -> CachedRequestData:

100 return cls(

101 req_id=request.request_id,

vllm/v1/core/sched/scheduler.py

+4 -8

@@ -180,7 +180,7 @@ class Scheduler(SchedulerInterface):

180 # uses structured decoding.

181 structured_output_request_ids: dict[str, int] = {}

182

183 - req_to_new_block_ids: dict[str, list[list[int]]] = {}

183 + req_to_new_block_ids: dict[str, tuple[list[int], ...]] = {}

184 num_scheduled_tokens: dict[str, int] = {}

185 token_budget = self.max_num_scheduled_tokens

186 # Encoder-related.

@@ -471,7 +471,7 @@ class Scheduler(SchedulerInterface):

471 token_budget -= num_new_tokens

472 request.status = RequestStatus.RUNNING

473 request.num_computed_tokens = num_computed_tokens

474 - # Count the number of prifix cached tokens.

474 + # Count the number of prefix cached tokens.

475 if request.num_cached_tokens < 0:

476 request.num_cached_tokens = num_computed_tokens

477 # Encoder-related.

@@ -588,7 +588,7 @@ class Scheduler(SchedulerInterface):

588 request: Request,

589 num_scheduled_tokens: int,

590 num_scheduled_spec_tokens: int,

591 - new_block_ids: list[list[int]],

591 + new_block_ids: tuple[list[int], ...],

592 resumed_from_preemption: bool,

593 ) -> CachedRequestData:

594 # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating

@@ -1015,11 +1015,7 @@ class Scheduler(SchedulerInterface):

1015 num_computed_tokens = min(num_computed_tokens, request.num_tokens)

1016 if num_computed_tokens == request.num_tokens:

1017 num_computed_tokens -= 1

1018 - self.kv_cache_manager.cache_blocks(

1019 - request,

1020 - self.kv_cache_manager.req_to_block_hashes[request.request_id],

1021 - num_computed_tokens,

1022 - )

1018 + self.kv_cache_manager.cache_blocks(request, num_computed_tokens)

1019

1020 # Update the request state for scheduling.

1021 request.num_computed_tokens = num_computed_tokens

vllm/v1/core/single_type_kv_cache_manager.py

+24 -26

@@ -197,7 +197,7 @@ class SingleTypeKVCacheManager(ABC):

197 block_pool: BlockPool,

198 kv_cache_spec: KVCacheSpec,

199 use_eagle: bool,

200 - ) -> list[list[KVCacheBlock]]:

200 + ) -> tuple[list[KVCacheBlock], ...]:

201 """

202 Get the longest cache hit prefix of the blocks that is not longer than

203 `max_length`. The prefix should be a common prefix hit for all the

@@ -222,7 +222,7 @@ class SingleTypeKVCacheManager(ABC):

222 element is a list of cached blocks for the i-th kv cache group

223 in `kv_cache_group_ids`.

224 For example, sliding window manager should return a list like

225 - [[NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]] for block size 4

225 + ([NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]) for block size 4

226 and sliding window 8 and len(kv_cache_group_ids) = 1.

227 """

228

@@ -254,27 +254,25 @@ class FullAttentionManager(SingleTypeKVCacheManager):

254 block_pool: BlockPool,

255 kv_cache_spec: KVCacheSpec,

256 use_eagle: bool,

257 - ) -> list[list[KVCacheBlock]]:

257 + ) -> tuple[list[KVCacheBlock], ...]:

258 assert isinstance(kv_cache_spec, FullAttentionSpec), (

259 "FullAttentionManager can only be used for full attention groups")

260 - computed_blocks: list[list[KVCacheBlock]] = [

261 - [] for _ in range(len(kv_cache_group_ids))

262 - ]

260 + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(

261 + [] for _ in range(len(kv_cache_group_ids)))

262 max_num_blocks = max_length // kv_cache_spec.block_size

264 - for i in range(max_num_blocks):

265 - block_hash = block_hashes[i]

263 + for i, block_hash in zip(range(max_num_blocks), block_hashes):

264 # block_hashes is a chain of block hashes. If a block hash is not

265 # in the cached_block_hash_to_id, the following block hashes are

266 # not computed yet for sure.

267 if cached_block := block_pool.get_cached_block(

268 block_hash, kv_cache_group_ids):

271 - for j in range(len(kv_cache_group_ids)):

272 - computed_blocks[j].append(cached_block[j])

269 + for computed, cached in zip(computed_blocks, cached_block):

270 + computed.append(cached)

271 else:

272 break

275 - if use_eagle and len(computed_blocks[0]) > 0:

276 - for j in range(len(kv_cache_group_ids)):

277 - computed_blocks[j].pop()

273 + if use_eagle and computed_blocks[0]:

274 + for computed in computed_blocks:

275 + computed.pop()

276 return computed_blocks

277

278 def remove_skipped_blocks(self, request_id: str,

@@ -311,7 +309,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

309 block_pool: BlockPool,

310 kv_cache_spec: KVCacheSpec,

311 use_eagle: bool,

314 - ) -> list[list[KVCacheBlock]]:

312 + ) -> tuple[list[KVCacheBlock], ...]:

313 assert isinstance(kv_cache_spec, SlidingWindowSpec), (

314 "SlidingWindowManager can only be used for sliding window groups")

315

@@ -332,23 +330,23 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

330 # sliding_window_contiguous_blocks),

331 # which is good for low cache hit rate scenarios.

332 max_num_blocks = max_length // kv_cache_spec.block_size

335 - computed_blocks = [[block_pool.null_block] * max_num_blocks

336 - for _ in range(len(kv_cache_group_ids))]

333 + computed_blocks = tuple([block_pool.null_block] * max_num_blocks

334 + for _ in range(len(kv_cache_group_ids)))

335 num_contiguous_blocks = 0

336 match_found = False

337 # Search from right to left and early stop when a match is found.

338 for i in range(max_num_blocks - 1, -1, -1):

339 if cached_block := block_pool.get_cached_block(

340 block_hashes[i], kv_cache_group_ids):

343 - for j in range(len(kv_cache_group_ids)):

344 - computed_blocks[j][i] = cached_block[j]

341 + for computed, cached in zip(computed_blocks, cached_block):

342 + computed[i] = cached

343 num_contiguous_blocks += 1

346 - if (num_contiguous_blocks >= sliding_window_contiguous_blocks):

344 + if num_contiguous_blocks >= sliding_window_contiguous_blocks:

345 # Trim the trailing blocks.

346 # E.g., [NULL, NULL, 8, 3, NULL, 9] -> [NULL, NULL, 8, 3]

347 # when sliding_window_contiguous_blocks=2.

350 - for j in range(len(kv_cache_group_ids)):

351 - del computed_blocks[j][i + num_contiguous_blocks:]

348 + for computed in computed_blocks:

349 + del computed[i + num_contiguous_blocks:]

350 match_found = True

351 break

352 else:

@@ -356,11 +354,11 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

354 if not match_found:

355 # The first `num_contiguous_blocks` is a cache hit even if

356 # `num_contiguous_blocks < sliding_window_contiguous_blocks`.

359 - for j in range(len(kv_cache_group_ids)):

360 - del computed_blocks[j][num_contiguous_blocks:]

361 - if use_eagle and len(computed_blocks[0]) > 0:

362 - for j in range(len(kv_cache_group_ids)):

363 - computed_blocks[j].pop()

357 + for computed in computed_blocks:

358 + del computed[num_contiguous_blocks:]

359 + if use_eagle and computed_blocks[0]:

360 + for computed in computed_blocks:

361 + computed.pop()

362 return computed_blocks

363

364 def remove_skipped_blocks(self, request_id: str,

vllm/v1/worker/block_table.py

+3 -2

@@ -112,11 +112,12 @@ class MultiGroupBlockTable:

112 for block_size in block_sizes

113 ]

114

115 - def append_row(self, block_ids: list[list[int]], row_idx: int) -> None:

115 + def append_row(self, block_ids: tuple[list[int], ...],

116 + row_idx: int) -> None:

117 for i, block_table in enumerate(self.block_tables):

118 block_table.append_row(block_ids[i], row_idx)

119

119 - def add_row(self, block_ids: list[list[int]], row_idx: int) -> None:

120 + def add_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:

121 for i, block_table in enumerate(self.block_tables):

122 block_table.add_row(block_ids[i], row_idx)

123

vllm/v1/worker/gpu_input_batch.py

+1 -1

@@ -30,7 +30,7 @@ class CachedRequestState:

30 sampling_params: SamplingParams

31 generator: Optional[torch.Generator]

33 - block_ids: list[list[int]]

33 + block_ids: tuple[list[int], ...]

34 num_computed_tokens: int

35 output_token_ids: list[int]