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

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

[V1] Remove pre-allocation for KV cache (#16941)

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

Woosuk Kwon committed Apr 22, 2025 at 00:52 UTC c4ab9f3e71f2b3309133886d1cc31e1dc8f614b4

5 files changed +59 -139

tests/v1/core/test_kv_cache_utils.py

+4 -11

@@ -496,8 +496,7 @@ def test_allocate_with_lookahead():

496

497 # Test case 1: Requires additional lookahead tokens

498 kv_cache_manager = KVCacheManager(kv_cache_config=config,

499 - max_model_len=100,

500 - num_preallocate_tokens=0)

499 + max_model_len=100)

500 blocks = kv_cache_manager.allocate_slots(

501 request,

502 num_tokens=3,

@@ -507,25 +506,19 @@ def test_allocate_with_lookahead():

506

507 # Test case 2: With precomputed blocks

508 kv_cache_manager = KVCacheManager(kv_cache_config=config,

510 - max_model_len=100,

511 - num_preallocate_tokens=4)

512 - # num_preallocate_blocks = 4 // 4 - 2 // 4 = 1

509 + max_model_len=100)

510 # required_blocks = ceil((3 + 2) /4) = 2

514 - # total_blocks = 1 + 2 = 3

511 blocks = kv_cache_manager.allocate_slots(

512 request,

513 num_tokens=3,

514 num_lookahead_tokens=2,

515 )

520 - assert len(blocks) == 3

516 + assert len(blocks) == 2

517

518 # Test case 3: With precomputed blocks

523 - # num_preallocate_blocks = 4 // 4 - 4 // 4 = 0

519 # required_blocks = ceil((3 + 4) / 4) = 2

525 - # total_blocks = 0 + 2 = 2

520 kv_cache_manager = KVCacheManager(kv_cache_config=config,

527 - max_model_len=100,

528 - num_preallocate_tokens=4)

521 + max_model_len=100)

522 blocks = kv_cache_manager.allocate_slots(

523 request,

524 num_tokens=3,

tests/v1/core/test_prefix_caching.py

+44 -97

@@ -8,7 +8,7 @@ import torch

9 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

10 from vllm.sampling_params import SamplingParams

11 -from vllm.utils import cdiv, sha256

11 +from vllm.utils import sha256

12 from vllm.v1.core.block_pool import BlockPool

13 from vllm.v1.core.kv_cache_manager import KVCacheManager, Request

14 from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,

@@ -61,7 +61,6 @@ def test_prefill(hash_algo):

61 max_model_len=8192,

62 enable_caching=True,

63 caching_hash_algo=hash_algo,

64 - num_preallocate_tokens=16,

64 )

66 # choose the hash function according to the parameter

@@ -80,7 +79,7 @@ def test_prefill(hash_algo):

79 assert not computed_blocks

80 assert num_computed_tokens == 0

81 blocks = manager.allocate_slots(req0, 55, computed_blocks)

83 - assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]

82 + assert [b.block_id for b in blocks] == [1, 2, 3, 4]

84 # Check full block metadata

85 parent_block_hash = None

@@ -92,8 +91,8 @@ def test_prefill(hash_algo):

91 assert manager.block_pool.blocks[block_id].ref_cnt == 1

92 parent_block_hash = block_hash.hash_value

95 - # Check partial/preallocated block metadata

96 - for block_id in (4, 5):

94 + # Check partial block metadata

95 + for block_id in (4, ):

96 assert manager.block_pool.blocks[block_id].block_hash is None

97 assert manager.block_pool.blocks[block_id].ref_cnt == 1

@@ -107,12 +106,12 @@ def test_prefill(hash_algo):

106 assert num_computed_tokens == 3 * 16

107 num_new_tokens = 53 - 3 * 16

108 blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)

110 - assert [b.block_id for b in blocks] == [6, 7]

109 + assert [b.block_id for b in blocks] == [5]

110 for block in computed_blocks:

111 assert block.ref_cnt == 2

112

114 - # At this point, we should have 3 free blocks left.

115 - assert manager.block_pool.free_block_queue.num_free_blocks == 3

113 + # At this point, we should have 5 free blocks left.

114 + assert manager.block_pool.free_block_queue.num_free_blocks == 5

115

116 manager.free(req0)

117 manager.free(req1)

@@ -120,14 +119,14 @@ def test_prefill(hash_algo):

119 # All blocks should be available.

120 assert manager.block_pool.free_block_queue.num_free_blocks == 10

121 # The order should be

123 - # [unallocated (8, 9, 10)]

124 - # [unique_req0 (5, 4)]

125 - # [unique_req1 (7, 6)]

122 + # [unallocated (6, 7, 8, 9, 10)]

123 + # [unique_req0 (4)]

124 + # [unique_req1 (5)]

125 # [common (3, 2, 1)]

126 assert [

127 b.block_id

128 for b in manager.block_pool.free_block_queue.get_all_free_blocks()

130 - ] == [8, 9, 10, 5, 4, 7, 6, 3, 2, 1]

129 + ] == [6, 7, 8, 9, 10, 4, 5, 3, 2, 1]

130

131 # Cache hit in the common prefix when the original block is already free.

132 # Incomplete 1 block (6 tokens)

@@ -139,29 +138,29 @@ def test_prefill(hash_algo):

138 assert num_computed_tokens == 3 * 16

139 num_new_tokens = 53 - 3 * 16

140 blocks = manager.allocate_slots(req2, num_new_tokens, computed_blocks)

142 - assert [b.block_id for b in blocks] == [8, 9]

141 + assert [b.block_id for b in blocks] == [6]

142

144 - # Although we only have 5 free blocks, we have 8 blocks in

143 + # Although we only have 6 free blocks, we have 8 blocks in

144 # the free block queue due to lazy removal.

146 - assert manager.block_pool.free_block_queue.num_free_blocks == 5

145 + assert manager.block_pool.free_block_queue.num_free_blocks == 6

146 assert all([

147 b.ref_cnt == 0

148 for b in manager.block_pool.free_block_queue.get_all_free_blocks()

149 ])

150 assert len([

151 b for b in manager.block_pool.free_block_queue.get_all_free_blocks()

153 - ]) == 5

152 + ]) == 6

153

154 manager.free(req2)

155

156 # Cache miss and eviction.

158 - req3 = make_request("3", [99] * (16 * 9))

157 + req3 = make_request("3", [99] * (16 * 10))

158 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)

159 assert not computed_blocks

160 assert num_computed_tokens == 0

162 - blocks = manager.allocate_slots(req3, 16 * 9, computed_blocks)

161 + blocks = manager.allocate_slots(req3, 16 * 10, computed_blocks)

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

164 - assert [b.block_id for b in blocks] == [10, 5, 4, 7, 6, 9, 8, 3, 2, 1]

163 + assert [b.block_id for b in blocks] == [7, 8, 9, 10, 4, 5, 6, 3, 2, 1]

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

165 assert manager.block_pool.free_block_queue.free_list_head is None

166 assert manager.block_pool.free_block_queue.free_list_tail is None

@@ -178,7 +177,6 @@ def test_prefill_plp():

177 make_kv_cache_config(16, 11),

178 max_model_len=8192,

179 enable_caching=True,

181 - num_preallocate_tokens=16,

180 )

181 # the default hash function is hash

182 hash_fn = hash

@@ -197,7 +195,7 @@ def test_prefill_plp():

195 assert not computed_blocks

196 assert num_computed_tokens == 0

197 blocks = manager.allocate_slots(req0, 55, computed_blocks)

200 - assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]

198 + assert [b.block_id for b in blocks] == [1, 2, 3, 4]

199 req0_block_hashes = [b.block_hash for b in blocks]

200

201 # Check full block metadata

@@ -210,8 +208,8 @@ def test_prefill_plp():

208 assert manager.block_pool.blocks[block_id].ref_cnt == 1

209 parent_block_hash = block_hash.hash_value

210

213 - # Check partial/preallocated block metadata

214 - for block_id in (4, 5):

211 + # Check partial block metadata

212 + for block_id in (4, ):

213 assert manager.block_pool.blocks[block_id].block_hash is None

214 assert manager.block_pool.blocks[block_id].ref_cnt == 1

215

@@ -226,12 +224,12 @@ def test_prefill_plp():

224 assert num_computed_tokens == 3 * 16

225 num_new_tokens = 53 - 3 * 16

226 blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)

229 - assert [b.block_id for b in blocks] == [6, 7]

227 + assert [b.block_id for b in blocks] == [5]

228 for block in computed_blocks:

229 assert block.ref_cnt == 2

230

233 - # At this point, we should have 3 free blocks left.

234 - assert manager.block_pool.free_block_queue.num_free_blocks == 3

231 + # At this point, we should have 5 free blocks left.

232 + assert manager.block_pool.free_block_queue.num_free_blocks == 5

233

234 manager.free(req0)

235 manager.free(req1)

@@ -239,14 +237,14 @@ def test_prefill_plp():

237 # All blocks should be available.

238 assert manager.block_pool.free_block_queue.num_free_blocks == 10

239 # The order should be

242 - # [unallocated (8, 9, 10)]

243 - # [unique_req0 (5, 4)]

244 - # [unique_req1 (7, 6)]

240 + # [unallocated (6, 7, 8, 9, 10)]

241 + # [unique_req0 (4)]

242 + # [unique_req1 (5)]

243 # [common (3, 2, 1)]

244 assert [

245 b.block_id

246 for b in manager.block_pool.free_block_queue.get_all_free_blocks()

249 - ] == [8, 9, 10, 5, 4, 7, 6, 3, 2, 1]

247 + ] == [6, 7, 8, 9, 10, 4, 5, 3, 2, 1]

248

249 # Request #2 is a prompt-logprobs request:

250 # NO cache hit in the common prefix; duplicates request #0 cached blocks

@@ -262,7 +260,7 @@ def test_prefill_plp():

260 block_ids = [b.block_id for b in blocks]

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

262 assert [b.block_hash for b in blocks] == req0_block_hashes

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

263 + assert block_ids != [1, 2, 3, 4]

264

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

266 # Check block reference counts.

@@ -277,7 +275,6 @@ def test_decode():

275 make_kv_cache_config(16, 11),

276 max_model_len=8192,

277 enable_caching=True,

280 - num_preallocate_tokens=16,

278 )

279

280 # Complete 3 blocks (48 tokens)

@@ -291,7 +288,7 @@ def test_decode():

288 assert not computed_blocks

289 assert num_computed_tokens == 0

290 blocks = manager.allocate_slots(req0, 55, computed_blocks)

294 - assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]

291 + assert [b.block_id for b in blocks] == [1, 2, 3, 4]

292

293 # Append slots without allocating a new block.

294 req0.num_computed_tokens = 55

@@ -299,28 +296,18 @@ def test_decode():

296 req0.append_output_token_ids(8)

297 new_blocks = manager.allocate_slots(req0, 4)

298 assert new_blocks is not None and len(new_blocks) == 0

302 - assert manager.req_to_blocks[req0.request_id][-2].block_hash is None

299 + assert manager.req_to_blocks[req0.request_id][-1].block_hash is None

300

304 - # Append slots without allocating a new block, but start using the

305 - # preallocated block.

301 + # Append slots with allocating a new block.

302 req0.num_computed_tokens = 59

307 - # 6 tokens to fill the previous block, and 10 tokens to fill

303 + # 9 tokens to fill the previous block, and 10 tokens to fill

304 # the preallocated block.

309 - for _ in range(5 + 10):

305 + for _ in range(9 + 10):

306 req0.append_output_token_ids(7)

311 - new_blocks = manager.allocate_slots(req0, 15)

312 - assert new_blocks is not None and len(new_blocks) == 0

307 + new_blocks = manager.allocate_slots(req0, 19)

308 + assert new_blocks is not None and len(new_blocks) == 1

309 assert manager.req_to_blocks[req0.request_id][-2].block_hash is not None

314 -

315 - # Append slots with allocating a new block.

316 - req0.num_computed_tokens = 74

317 - # 6 tokens to fill the previous block, and 10 tokens to fill

318 - # the preallocated block.

319 - for _ in range(6 + 11):

320 - req0.append_output_token_ids(12)

321 - new_blocks = manager.allocate_slots(req0, 17)

322 - # Plus one preallocated block.

323 - assert new_blocks is not None and len(new_blocks) == 2

310 + assert manager.req_to_blocks[req0.request_id][-1].block_hash is None

311

312

313 def test_evict():

@@ -328,7 +315,6 @@ def test_evict():

315 make_kv_cache_config(16, 11),

316 max_model_len=8192,

317 enable_caching=True,

331 - num_preallocate_tokens=16,

318 )

319

320 last_token_id = 5 * 16 + 7

@@ -337,7 +323,7 @@ def test_evict():

323 assert not computed_blocks

324 assert num_computed_tokens == 0

325 blocks = manager.allocate_slots(req0, 5 * 16 + 7, computed_blocks)

340 - assert len(blocks) == 7 # 5 full + 1 partial + 1 preallocated

326 + assert len(blocks) == 6 # 5 full + 1 partial

327

328 # 3 blocks.

329 req1 = make_request("1", list(range(last_token_id,

@@ -349,7 +335,8 @@ def test_evict():

335 assert len(blocks) == 3 # 3 full blocks

336 last_token_id += 3 * 16

337

352 - assert manager.block_pool.free_block_queue.num_free_blocks == 0

338 + # 10 - (6 + 3) == 1

339 + assert manager.block_pool.free_block_queue.num_free_blocks == 1

340

341 manager.free(req0)

342 manager.free(req1)

@@ -357,7 +344,7 @@ def test_evict():

344 assert [

345 b.block_id

346 for b in manager.block_pool.free_block_queue.get_all_free_blocks()

360 - ] == [7, 6, 5, 4, 3, 2, 1, 10, 9, 8]

347 + ] == [10, 6, 5, 4, 3, 2, 1, 9, 8, 7]

348

349 # Touch the first 2 blocks.

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

@@ -365,8 +352,8 @@ def test_evict():

352 assert [b.block_id for b in computed_blocks] == [1, 2]

353 assert num_computed_tokens == 2 * 16

354 blocks = manager.allocate_slots(req2, 3, computed_blocks)

368 - assert [b.block_id for b in blocks] == [7, 6]

369 - assert manager.block_pool.free_block_queue.num_free_blocks == 6

355 + assert [b.block_id for b in blocks] == [10]

356 + assert manager.block_pool.free_block_queue.num_free_blocks == 7

357

358

359 def test_hash_block_correct_reuse():

@@ -379,7 +366,6 @@ def test_hash_block_correct_reuse():

366 make_kv_cache_config(16, 2),

367 max_model_len=8192,

368 enable_caching=True,

382 - num_preallocate_tokens=0,

369 )

370

371 # Allocate 1 block and cache it.

@@ -416,7 +402,6 @@ def test_computed_blocks_not_evicted():

402 make_kv_cache_config(block_size, 3),

403 max_model_len=8192,

404 enable_caching=True,

419 - num_preallocate_tokens=0,

405 )

406

407 # Allocate a block and cache it.

@@ -465,7 +450,6 @@ def test_basic_prefix_caching_disabled():

450 make_kv_cache_config(block_size, 5),

451 max_model_len=8192,

452 enable_caching=False,

468 - num_preallocate_tokens=0,

453 )

454

455 req1 = make_request("1", list(range(10))) # 2 blocks and some more

@@ -496,40 +480,6 @@ def test_basic_prefix_caching_disabled():

480 assert not blocks

481

482

499 -@pytest.mark.parametrize("num_preallocate_tokens", list(range(0, 8)))

500 -@pytest.mark.parametrize("block_size", [4])

501 -def test_preallocate_blocks(num_preallocate_tokens: int, block_size: int):

502 - """

503 - This tests that the preallocated blocks are correctly added.

504 - """

505 - manager = KVCacheManager(

506 - make_kv_cache_config(block_size, 11),

507 - max_model_len=8192,

508 - enable_caching=True,

509 - num_preallocate_tokens=num_preallocate_tokens,

510 - )

511 - num_preallocated_blocks = cdiv(num_preallocate_tokens, block_size)

512 -

513 - req = make_request("0", list(range(block_size * 30)))

514 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)

515 - assert not computed_blocks

516 - assert num_computed_tokens == 0

517 - # Just ask for 1 block.

518 - blocks = manager.allocate_slots(req, block_size, computed_blocks)

519 - req.num_computed_tokens = block_size

520 - assert len(blocks) == 1 + num_preallocated_blocks

521 -

522 - # Assume all computed, only when num_preallocate_tokens > 0, we need to

523 - # consume the previously preallocated blocks.

524 - if num_preallocated_blocks > 0:

525 - manager.allocate_slots(req, block_size * (len(blocks) - 1))

526 - req.num_computed_tokens = block_size * len(blocks)

527 -

528 - # Append 1 block.

529 - blocks = manager.allocate_slots(req, block_size)

530 - assert len(blocks) == 1 + num_preallocated_blocks

531 -

532 -

483 @pytest.mark.parametrize("hash_fn", [sha256, hash])

484 def test_cache_blocks(hash_fn):

485 """

@@ -588,7 +538,6 @@ def test_mm_prefix_caching():

538 make_kv_cache_config(16, 11),

539 max_model_len=8192,

540 enable_caching=True,

591 - num_preallocate_tokens=16,

541 )

542

543 # Common prompt tokens (T is text tokens and P is image placeholder tokens)

@@ -626,7 +575,7 @@ def test_mm_prefix_caching():

575 assert block_hashes[2].extra_keys == ("bbb", )

576

577 blocks = manager.allocate_slots(req0, 59, computed_blocks)

629 - assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]

578 + assert [b.block_id for b in blocks] == [1, 2, 3, 4]

579 req0.num_computed_tokens = 59

580

581 # Append slots without allocating a new block.

@@ -667,7 +616,6 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():

616 make_kv_cache_config(block_size, 11),

617 max_model_len=8192,

618 enable_caching=True,

670 - num_preallocate_tokens=0,

619 )

620 # Complete 3 blocks (48 tokens)

@@ -721,7 +669,6 @@ def test_reset_prefix_cache():

669 make_kv_cache_config(16, 11),

670 max_model_len=8192,

671 enable_caching=True,

724 - num_preallocate_tokens=0,

672 )

673

674 full_block_token_ids = [i for i in range(3) for _ in range(16)]

tests/v1/core/test_scheduler.py

+4 -8

@@ -804,20 +804,17 @@ def _assert_right_kv_cache_manager(

804 """Check whether KVCacheManager is correct after allocate."""

805

806 # Make sure the request stats are right.

807 - EXPECTED_ACTUAL_BLOCKS = num_tokens // block_size

808 - EXPECTED_TOTAL_BLOCKS = (EXPECTED_ACTUAL_BLOCKS +

809 - scheduler.kv_cache_manager.num_preallocate_blocks)

807 + EXPECTED_TOTAL_BLOCKS = num_tokens // block_size

808 for req_id in req_ids:

809 blocks = scheduler.kv_cache_manager.req_to_blocks[req_id]

810 hashes = scheduler.kv_cache_manager.req_to_block_hashes[req_id]

811 assert (scheduler.kv_cache_manager.num_cached_block[req_id] ==

814 - EXPECTED_ACTUAL_BLOCKS)

812 + EXPECTED_TOTAL_BLOCKS)

813 assert len(blocks) == EXPECTED_TOTAL_BLOCKS

816 - assert len(hashes) == EXPECTED_ACTUAL_BLOCKS

814 + assert len(hashes) == EXPECTED_TOTAL_BLOCKS

815

816 # Make sure we actually touched all the blocks.

819 - BLOCKS_PER_REQ = (num_tokens / block_size +

820 - scheduler.kv_cache_manager.num_preallocate_blocks)

817 + BLOCKS_PER_REQ = num_tokens / block_size

818 assert (scheduler.kv_cache_manager.block_pool.get_num_free_blocks() ==

819 num_total_blocks - num_requests * BLOCKS_PER_REQ)

820

@@ -1052,7 +1049,6 @@ def test_kv_connector_handles_preemption():

1049 block_size=BLOCK_SIZE,

1050 num_blocks=NUM_BLOCKS,

1051 )

1055 - scheduler.kv_cache_manager.num_preallocate_blocks = 0

1052

1053 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE

1054 scheduler.connector.get_num_new_matched_tokens = Mock(name="method")

vllm/v1/core/kv_cache_manager.py

+2 -21

@@ -25,7 +25,6 @@ class KVCacheManager:

25 max_model_len: int,

26 enable_caching: bool = True,

27 caching_hash_algo: str = "builtin",

28 - num_preallocate_tokens: int = 64,

28 log_stats: bool = False,

29 ) -> None:

30 assert len(kv_cache_config.kv_cache_groups) == 1, (

@@ -42,22 +41,8 @@ class KVCacheManager:

41 self.log_stats = log_stats

42 # FIXME: make prefix cache stats conditional on log_stats

43 self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

45 - # NOTE(woosuk): To avoid frequent block allocation, we preallocate some

46 - # blocks for each request. For example, when a request reaches the end

47 - # of its block table, we preallocate N blocks in advance. This way, we

48 - # reduce the overhead of updating free_block_ids and ref_cnts for each

49 - # request every step (at the cost of some memory waste).

50 - # NOTE(woosuk): This is different from the "lookahead" slots since this

51 - # does not guarantee that the request always has N empty blocks. After

52 - # the request gets N empty blocks, it starts to use the blocks without

53 - # further allocation. When it uses up all the N empty blocks, it gets

54 - # N new empty blocks.

55 - self.num_preallocate_tokens = num_preallocate_tokens

56 - self.num_preallocate_blocks = cdiv(num_preallocate_tokens,

57 - self.block_size)

45 self.block_pool = BlockPool(self.num_gpu_blocks, enable_caching)

60 -

46 self.specialized_manager = get_specialized_manager(

47 kv_cache_spec=kv_cache_spec,

48 block_pool=self.block_pool,

@@ -256,13 +241,9 @@ class KVCacheManager:

241 # No new block is needed.

242 new_blocks = []

243 else:

259 - # Get new blocks from the free block pool considering

260 - # preallocated blocks.

261 - num_preallocate_blocks = max(

262 - 0, self.num_preallocate_blocks -

263 - num_lookahead_tokens // self.block_size)

244 + # Get new blocks from the free block pool.

245 num_new_blocks = min(

265 - num_new_blocks + num_preallocate_blocks,

246 + num_new_blocks,

247 self.block_pool.get_num_free_blocks(),

248 # Should not exceed the maximum number of blocks per request.

249 # This is especially because the block table has the shape

vllm/v1/core/sched/scheduler.py

+5 -2

@@ -358,8 +358,11 @@ class Scheduler(SchedulerInterface):

358 new_encoder_budget = encoder_budget

359

360 new_blocks = self.kv_cache_manager.allocate_slots(

361 - request, num_new_tokens + num_external_tokens,

362 - computed_blocks)

361 + request,

362 + num_new_tokens + num_external_tokens,

363 + computed_blocks,

364 + num_lookahead_tokens=self.num_lookahead_tokens,

365 + )

366 if new_blocks is None:

367 # The request cannot be scheduled.

368 break