@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8ac8375270

[Core] Preserve Marconi caching with selective hybrid cache retention (#47782)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Jul 13, 2026 at 21:24 UTC 8ac8375270c1c86b0a8eaba7e738908a9a7303e1

9 files changed +502 -236

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+16 -16

@@ -110,7 +110,7 @@ def test_hybrid_mamba_align_partial_hash_hit():

110 )

111

112 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

113 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

113 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

114 assert num_computed == 0

115 blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks)

116 assert blocks is not None

@@ -125,7 +125,7 @@ def test_hybrid_mamba_align_partial_hash_hit():

125 assert partial_mamba_block[0].block_hash_num_tokens == 6

126

127 req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)

128 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

128 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

129 assert num_computed == 6

130 assert [len(group) for group in computed_blocks.blocks] == [3, 2]

131

@@ -185,7 +185,7 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():

185 )

186

187 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

188 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

188 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

189 assert num_computed == 0

190 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

191

@@ -258,7 +258,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():

258 )

259

260 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

261 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

261 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

262 assert num_computed == 0

263 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

264

@@ -279,7 +279,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():

279 manager.new_step_starts()

280

281 req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256)

282 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

282 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

283 assert num_computed == 6

284 # The later request hits the moved (private-copy) entry, not the source.

285 assert computed_blocks.get_block_ids()[1][1] == moved_block_id

@@ -338,7 +338,7 @@ def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit():

338 )

339

340 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

341 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

341 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

342 assert num_computed == 0

343 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

344 manager.new_step_starts()

@@ -351,13 +351,13 @@ def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit():

351

352 # A request hitting the moved entry in the SAME step must be deferred.

353 req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256)

354 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

354 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

355 assert num_computed == 6

356 assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is None

357

358 # Next step the moved entry is consumable.

359 manager.new_step_starts()

360 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

360 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

361 assert num_computed == 6

362 assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is not None

363

@@ -397,7 +397,7 @@ def test_hybrid_full_attention_partial_hash_hit_uses_cow():

397 )

398

399 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

400 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

400 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

401 assert num_computed == 0

402 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

403 manager.free(req0)

@@ -410,7 +410,7 @@ def test_hybrid_full_attention_partial_hash_hit_uses_cow():

410 assert partial_full_block is not None

411

412 req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)

413 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

413 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

414 assert num_computed == 6

415 assert [len(group) for group in computed_blocks.blocks] == [2, 2]

416

@@ -471,7 +471,7 @@ def test_hybrid_partial_hit_cow_target_starts_uncached():

471 )

472

473 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

474 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

474 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

475 assert num_computed == 0

476 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

477 manager.free(req0)

@@ -488,7 +488,7 @@ def test_hybrid_partial_hit_cow_target_starts_uncached():

488 assert partial_mamba_block is not None

489

490 req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)

491 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

491 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

492 assert num_computed == 6

493

494 new_blocks = manager.allocate_slots(

@@ -586,7 +586,7 @@ def test_hybrid_partial_hash_truncates_full_attention_hit_length():

586 block_size=block_size,

587 )

588

589 - computed_blocks, num_computed = manager.get_computed_blocks(req)

589 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req)

590 assert num_computed == 6

591 assert [len(group) for group in computed_blocks.blocks] == [2, 2]

592

@@ -627,7 +627,7 @@ def test_cow_retained_blocks_returned_for_release():

627 hash_block_size=hash_block_size,

628 )

629 req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

630 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

630 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

631 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

632

633 # The owner's move queues a copy and retains both endpoints.

@@ -797,7 +797,7 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks():

797 # The owner prefills in scheduler-split style: stop at the block boundary

798 # (4), then at the prompt's last hash boundary (6, partial entries).

799 req0 = make_request("0", [7] * 6, hash_block_size, sha256)

800 - computed_blocks, num_computed = manager.get_computed_blocks(req0)

800 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

801 assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None

802 req0.num_computed_tokens = 4

803 manager.new_step_starts()

@@ -808,7 +808,7 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks():

808 # A longer request with eagle: full attention drops the partial tail, so

809 # the joint hit must fall back to the block boundary the FA blocks cover.

810 req1 = make_request("1", [7] * 6 + [9] * 2, hash_block_size, sha256)

811 - computed_blocks, num_computed = manager.get_computed_blocks(req1)

811 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)

812 assert num_computed == 4

813 assert all(

814 len(group) * block_size >= num_computed for group in computed_blocks.blocks

tests/v1/core/test_prefix_caching.py

+330 -86

@@ -244,7 +244,7 @@ def test_prefill(hash_fn):

244 unique_token_ids = [3] * 7

245 all_token_ids = common_token_ids + unique_token_ids

246 req0 = make_request("0", all_token_ids, block_size, hash_fn)

247 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

247 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

248 assert len(req0.block_hashes) == 3

249 assert not computed_blocks.blocks[0]

250 assert num_computed_tokens == 0

@@ -274,7 +274,7 @@ def test_prefill(hash_fn):

274 # Incomplete 1 block (5 tokens)

275 unique_token_ids = [3] * 5

276 req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)

277 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

277 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

278 assert len(req1.block_hashes) == 3

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

280 assert num_computed_tokens == 3 * 16

@@ -307,7 +307,7 @@ def test_prefill(hash_fn):

307 # Incomplete 1 block (6 tokens)

308 unique_token_ids = [3] * 6

309 req2 = make_request("2", common_token_ids + unique_token_ids, block_size, hash_fn)

310 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

310 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

311 assert len(req2.block_hashes) == 3

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

313 assert num_computed_tokens == 3 * 16

@@ -327,7 +327,7 @@ def test_prefill(hash_fn):

327

328 # Cache miss and eviction.

329 req3 = make_request("3", [99] * (16 * 10), block_size, hash_fn)

330 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)

330 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req3)

331 assert not computed_blocks.blocks[0]

332 assert num_computed_tokens == 0

333 blocks = manager.allocate_slots(

@@ -369,7 +369,7 @@ def test_prefill_hybrid_model():

369 unique_token_ids = [3] * 7

370 all_token_ids = common_token_ids + unique_token_ids

371 req0 = make_request("0", all_token_ids, block_size, hash_fn)

372 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

372 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

373 assert len(req0.block_hashes) == 3

374 assert not computed_blocks.blocks[0]

375 assert num_computed_tokens == 0

@@ -405,7 +405,7 @@ def test_prefill_hybrid_model():

405 unique_token_ids = [3] * 5

406 all_token_ids = common_token_ids + unique_token_ids

407 req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)

408 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

408 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

409 assert len(req1.block_hashes) == 3

410 assert computed_blocks.get_block_ids() == ([1, 2, 3], [0, 6, 7], [0, 10, 11])

411 assert num_computed_tokens == 3 * 16

@@ -539,7 +539,7 @@ def test_prefill_hybrid_model_eagle():

539 unique_token_ids = [6] * 7

540 all_token_ids = common_token_ids + unique_token_ids

541 req0 = make_request("0", all_token_ids, block_size, hash_fn)

542 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

542 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

543 assert len(req0.block_hashes) == len(all_token_ids) // block_size

544 assert not computed_blocks.blocks[0]

545 assert num_computed_tokens == 0

@@ -576,7 +576,7 @@ def test_prefill_hybrid_model_eagle():

576 unique_token_ids = [6] * 5

577 all_token_ids = common_token_ids + unique_token_ids

578 req1 = make_request("1", all_token_ids, block_size, hash_fn)

579 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

579 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

580 assert len(req1.block_hashes) == num_full_blocks

581 assert computed_blocks.get_block_ids() == (

582 [1, 2, 3, 4, 5],

@@ -712,7 +712,7 @@ def _test_partial_request_hit(

712 req = make_request(request_id, prompt_token_ids, block_size, sha256)

713 for hash_with_group_id in hash_to_evict:

714 manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)

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

715 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)

716 assert len(req.block_hashes) == num_full_blocks

717 assert num_computed_tokens == expect_hit_length * block_size

718 for block_per_group in computed_blocks.blocks:

@@ -873,7 +873,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):

873

874 # First request: no cache hit initially

875 req0 = make_request("0", all_token_ids, block_size, hash_fn)

876 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

876 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

877

878 assert len(req0.block_hashes) == 3

879 assert not computed_blocks.blocks[0] # No cache hit initially

@@ -890,7 +890,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):

890

891 # Second request: should hit cached blocks for common prefix

892 req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn)

893 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

893 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

894

895 # Should hit cached blocks for all groups

896 assert num_computed_tokens == 3 * block_size

@@ -950,7 +950,7 @@ def test_prefill_hybrid_model_combinations_eagle(

950

951 # First request: no cache hit initially

952 req0 = make_request("0", all_token_ids, block_size, hash_fn)

953 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

953 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

954

955 assert len(req0.block_hashes) == num_full_blocks

956 assert not computed_blocks.blocks[0] # No cache hit initially

@@ -966,7 +966,7 @@ def test_prefill_hybrid_model_combinations_eagle(

966 # Second request: should hit cached blocks for common prefix

967 all_token_ids = common_token_ids + [6] * 5

968 req1 = make_request("1", all_token_ids, block_size, hash_fn)

969 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

969 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

970

971 # Should hit cached blocks for all groups

972 assert num_computed_tokens == expect_hit_length * block_size

@@ -1019,7 +1019,7 @@ def test_prefill_hybrid_model_mamba_align():

1019 # First request: allocate_slots should not crash with the assertion error

1020 # in MambaManager.cache_blocks() when null blocks are present.

1021 req0 = make_request("0", all_token_ids, block_size, hash_fn)

1022 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1022 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1023 assert num_computed_tokens == 0

1024

1025 blocks = manager.allocate_slots(req0, 55, num_computed_tokens, computed_blocks)

@@ -1048,28 +1048,33 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():

1048 # Request: 3 blocks

1049 prefix = [i for i in range(3) for _ in range(block_size)]

1050 req_0 = make_request("0", prefix, block_size, hash_fn)

1051 - computed_blocks, num_computed = manager.get_computed_blocks(req_0)

1052 - num_uncached_common = manager.coordinator.num_uncached_common_prefix_tokens

1051 + computed_blocks, num_computed, shared_prefix_boundary = manager.get_computed_blocks(

1052 + req_0

1053 + )

1054 assert num_computed == 0 # nothing cached yet

1054 - assert num_uncached_common == 0

1055 + assert shared_prefix_boundary == 0

1056 manager.allocate_slots(req_0, 3 * block_size, 0, computed_blocks)

1057

1058 # Request: 3 blocks (shared with above) + 7 different tokens

1059 req_1 = make_request("1", prefix + [100] * 7, block_size, hash_fn)

1059 - computed_blocks, num_computed = manager.get_computed_blocks(req_1)

1060 - num_uncached_common = manager.coordinator.num_uncached_common_prefix_tokens

1060 + computed_blocks, num_computed, shared_prefix_boundary = manager.get_computed_blocks(

1061 + req_1

1062 + )

1063 assert num_computed == 3 * block_size # we should observe a 3-block cache hit

1062 - assert num_uncached_common == 0

1064 + assert shared_prefix_boundary == 0

1065 manager.allocate_slots(req_1, 7, 3 * block_size, computed_blocks)

1066

1067 # Request: 3 blocks, but only 2 blocks shared (replace the last token in 3rd block):

1068 req_2 = make_request("2", prefix[:-1] + [101], block_size, hash_fn)

1067 - computed_blocks, num_computed = manager.get_computed_blocks(req_2)

1068 - num_uncached_common = manager.coordinator.num_uncached_common_prefix_tokens

1069 + computed_blocks, num_computed, shared_prefix_boundary = manager.get_computed_blocks(

1070 + req_2

1071 + )

1072 assert num_computed == 0 # mamba_align doesn't cache intermediate blocks

1070 - assert num_uncached_common == 2 * block_size # heuristic detects a shared prefix

1073 + # boundary == junction: 2 shared blocks (num_computed 0 + 2 uncached blocks).

1074 + assert shared_prefix_boundary == 2 * block_size

1075

1072 - # Next, validate scheduler logic for num_uncached_common_prefix_tokens > 0

1076 + # Next, validate the align-mode Marconi split. It takes the shared-prefix

1077 + # junction (absolute) and stops the chunk there (req_2 has num_computed 0).

1078 # Create minimal mock with just the needed attributes

1079 mock = SimpleNamespace(

1080 cache_config=SimpleNamespace(block_size=block_size),

@@ -1077,11 +1082,11 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():

1082 hash_block_size=block_size,

1083 mamba_partial_cache_hit=False,

1084 )

1085 + req_2.shared_prefix_boundary = shared_prefix_boundary

1086 num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(

1087 self=mock,

1088 request=req_2,

1089 num_new_tokens=3 * block_size,

1084 - num_uncached_common_prefix_tokens=num_uncached_common,

1090 )

1091 assert num_new_tokens_adjusted == 2 * block_size # adjust to the common prefix

1092

@@ -1118,7 +1123,7 @@ def test_hybrid_model_mamba_align_with_dynamic_draft_tokens():

1123

1124 all_token_ids = [i for i in range(3) for _ in range(block_size)] + [3] * 7

1125 req0 = make_request("0", all_token_ids, block_size, hash_fn)

1121 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1126 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1127 assert num_computed_tokens == 0

1128 blocks = manager.allocate_slots(

1129 req0, len(all_token_ids), num_computed_tokens, computed_blocks

@@ -1170,7 +1175,7 @@ def test_prefill_plp():

1175 unique_token_ids = [3] * 7

1176 all_token_ids = common_token_ids + unique_token_ids

1177 req0 = make_request("0", all_token_ids, block_size, hash_fn, prompt_logprobs=5)

1173 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1178 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1179 assert len(req0.block_hashes) == 3

1180 assert not computed_blocks.blocks[0]

1181 assert num_computed_tokens == 0

@@ -1202,7 +1207,7 @@ def test_prefill_plp():

1207 # Incomplete 1 block (5 tokens)

1208 unique_token_ids = [3] * 5

1209 req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)

1205 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1210 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1211 assert len(req1.block_hashes) == 3

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

1213 assert num_computed_tokens == 3 * 16

@@ -1236,7 +1241,7 @@ def test_prefill_plp():

1241 req2 = make_request(

1242 "2", common_token_ids + unique_token_ids, block_size, hash_fn, prompt_logprobs=5

1243 )

1239 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1244 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

1245 assert len(req2.block_hashes) == 3

1246 assert not computed_blocks.blocks[0]

1247 assert num_computed_tokens == 0

@@ -1273,7 +1278,7 @@ def test_decode():

1278 # Incomplete 1 block (7 tokens)

1279 unique_token_ids = [3] * 7

1280 req0 = make_request("0", common_token_ids + unique_token_ids, block_size, sha256)

1276 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1281 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1282 assert not computed_blocks.blocks[0]

1283 assert num_computed_tokens == 0

1284 blocks = manager.allocate_slots(

@@ -1331,7 +1336,7 @@ def test_evict():

1336

1337 last_token_id = 5 * 16 + 7

1338 req0 = make_request("0", list(range(last_token_id)), block_size, sha256)

1334 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1339 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1340 assert not computed_blocks.blocks[0]

1341 assert num_computed_tokens == 0

1342 blocks = manager.allocate_slots(

@@ -1344,7 +1349,7 @@ def test_evict():

1349 req1 = make_request(

1350 "1", list(range(last_token_id, last_token_id + 3 * 16)), block_size, sha256

1351 )

1347 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1352 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1353 assert not computed_blocks.blocks[0]

1354 assert num_computed_tokens == 0

1355 blocks = manager.allocate_slots(

@@ -1369,7 +1374,7 @@ def test_evict():

1374

1375 # Touch the first 2 blocks.

1376 req2 = make_request("2", list(range(2 * 16 + 3)), block_size, sha256)

1372 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1377 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

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

1379 assert num_computed_tokens == 2 * 16

1380 blocks = manager.allocate_slots(

@@ -1395,7 +1400,7 @@ def test_hash_block_correct_reuse():

1400 # Allocate 1 block and cache it.

1401 num_tokens = block_size * 1

1402 req = make_request("0", list(range(num_tokens)), block_size, sha256)

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

1403 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)

1404 assert not computed_blocks.blocks[0]

1405 assert num_computed_tokens == 0

1406 blocks = manager.allocate_slots(

@@ -1409,7 +1414,7 @@ def test_hash_block_correct_reuse():

1414 # Allocate a new block that's not full, make sure hash info on the

1415 # block is cleared.

1416 req = make_request("1", list(range(num_tokens - 1)), block_size, sha256)

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

1417 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)

1418 assert not computed_blocks.blocks[0]

1419 assert num_computed_tokens == 0

1420 blocks = manager.allocate_slots(

@@ -1436,7 +1441,7 @@ def test_computed_blocks_not_evicted():

1441 # Allocate a block and cache it.

1442 num_tokens = block_size * 1

1443 req0 = make_request("0", list(range(num_tokens)), block_size, sha256)

1439 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1444 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1445 assert not computed_blocks.blocks[0]

1446 assert num_computed_tokens == 0

1447 blocks = manager.allocate_slots(

@@ -1449,7 +1454,7 @@ def test_computed_blocks_not_evicted():

1454 req1 = make_request(

1455 "1", list(range(num_tokens, num_tokens * 2)), block_size, sha256

1456 )

1452 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1457 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1458 assert not computed_blocks.blocks[0]

1459 assert num_computed_tokens == 0

1460 blocks = manager.allocate_slots(

@@ -1465,7 +1470,7 @@ def test_computed_blocks_not_evicted():

1470 # Now if we have a cache hit on the first block, we should evict the second

1471 # cached block rather than the first one.

1472 req2 = make_request("2", list(range(num_tokens * 2)), block_size, sha256)

1468 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1473 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

1474 assert len(computed_blocks.blocks[0]) == 1

1475 assert computed_blocks.blocks[0][0].block_id == 1

1476 assert num_computed_tokens == block_size

@@ -1496,7 +1501,7 @@ def test_basic_prefix_caching_disabled():

1501 "1", list(range(10)), block_size, sha256

1502 ) # 2 blocks and some more

1503

1499 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1504 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1505 assert not computed_blocks.blocks[0]

1506 assert num_computed_tokens == 0

1507 blocks = manager.allocate_slots(

@@ -1509,7 +1514,7 @@ def test_basic_prefix_caching_disabled():

1514

1515 # No caching.

1516 req2 = make_request("2", list(range(16)), block_size, sha256) # shared prefix

1512 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1517 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

1518 assert not computed_blocks.blocks[0]

1519 assert num_computed_tokens == 0

1520 blocks = manager.allocate_slots(

@@ -1519,7 +1524,7 @@ def test_basic_prefix_caching_disabled():

1524

1525 # New requests should not have any blocks.

1526 req3 = make_request("3", list(range(4)), block_size, sha256)

1522 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)

1527 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req3)

1528 assert not computed_blocks.blocks[0]

1529 assert num_computed_tokens == 0

1530 blocks = manager.allocate_slots(

@@ -1701,7 +1706,7 @@ def test_mm_prefix_caching():

1706 mm_positions=mm_positions,

1707 mm_hashes=mm_hashes,

1708 )

1704 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1709 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1710

1711 # Completed block should have hashes

1712 assert not computed_blocks.blocks[0]

@@ -1766,7 +1771,7 @@ def test_mm_prefix_caching():

1771 mm_positions=mm_positions,

1772 mm_hashes=mm_hashes,

1773 )

1769 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1774 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1775 assert len(computed_blocks.blocks[0]) == 3

1776 assert num_computed_tokens == 3 * 16

1777

@@ -1788,7 +1793,7 @@ def test_cache_key_salting():

1793 common_token_ids = [i for i in range(3) for _ in range(block_size)]

1794 token_ids = common_token_ids + [3] * 11

1795 req0 = make_request("0", token_ids, block_size, sha256, cache_salt="salt1")

1791 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1796 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1797

1798 # Completed block should have hashes

1799 assert not computed_blocks.blocks[0]

@@ -1827,7 +1832,7 @@ def test_cache_key_salting():

1832 # Test cache hit with a new request that has the same salt.

1833 token_ids = common_token_ids + [4] * 11

1834 req1 = make_request("1", token_ids, block_size, sha256, cache_salt="salt1")

1830 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1835 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1836 # Should match only a prefix of 3 blocks.

1837 assert len(computed_blocks.blocks[0]) == 3

1838 assert num_computed_tokens == 3 * block_size

@@ -1835,7 +1840,7 @@ def test_cache_key_salting():

1840 # Test cache miss with same content but different salt.

1841 token_ids = common_token_ids + [4] * 11

1842 req2 = make_request("2", token_ids, block_size, sha256, cache_salt="salt2")

1838 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1843 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

1844 assert len(computed_blocks.blocks[0]) == 0

1845 assert num_computed_tokens == 0

1846 block_hashes = req2.block_hashes

@@ -1869,7 +1874,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():

1875 common_token_ids = [i for i in range(3) for _ in range(16)]

1876 req0 = make_request("0", common_token_ids, block_size, sha256)

1872 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1877 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

1878 assert not computed_blocks.blocks[0]

1879 assert num_computed_tokens == 0

1880 manager.allocate_slots(

@@ -1881,7 +1886,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():

1886

1887 # | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |

1888 req1 = make_request("1", common_token_ids * 2, block_size, sha256)

1884 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1889 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

1890 assert computed_blocks.blocks[0] == block_part0

1891 assert num_computed_tokens == 3 * 16

1892 manager.allocate_slots(

@@ -1899,7 +1904,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():

1905 # | Req1-5(F)| Req2-0 | Req2-1 | ... |

1906 req2 = make_request("2", [7] * block_size * 2, block_size, sha256)

1902 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1907 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

1908 assert not computed_blocks.blocks[0]

1909 assert num_computed_tokens == 0

1910 manager.allocate_slots(

@@ -1914,7 +1919,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():

1919 # In this case, the ref_cnt of the computed blocks should not be changed.

1920 assert manager.block_pool.free_block_queue.num_free_blocks == 5

1921 req3 = make_request("3", common_token_ids * 3, block_size, sha256)

1917 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)

1922 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req3)

1923 assert computed_blocks.blocks[0] == block_part1

1924 assert num_computed_tokens == 6 * 16

1925 # Req3 cannot be allocated.

@@ -1949,7 +1954,7 @@ def test_reset_prefix_cache():

1954 unique_token_ids = [4] * 7

1955 all_token_ids = full_block_token_ids + unique_token_ids

1956 req1 = make_request("1", all_token_ids, block_size, sha256)

1952 - computed_blocks, _ = manager.get_computed_blocks(req1)

1957 + computed_blocks, _, _ = manager.get_computed_blocks(req1)

1958 assert len(req1.block_hashes) == 3

1959 assert len(computed_blocks.blocks[0]) == 3

1960 blocks = manager.allocate_slots(

@@ -1984,7 +1989,7 @@ def test_prefix_cache_stats_disabled():

1989

1990 # Call all functions that check whether log_stats is disabled.

1991 req = make_request("0", list(range(16)), block_size, sha256)

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

1992 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)

1993 assert not computed_blocks.blocks[0]

1994 assert num_computed_tokens == 0

1995 manager.allocate_slots(

@@ -2591,7 +2596,7 @@ def test_eagle_enabled_removes_last_block():

2596 req = make_request("divisible_request", token_ids, block_size, sha256)

2597

2598 # Prime the cache

2594 - computed_blocks, _ = manager.get_computed_blocks(req)

2599 + computed_blocks, _, _ = manager.get_computed_blocks(req)

2600 manager.allocate_slots(

2601 req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks

2602 )

@@ -2599,7 +2604,7 @@ def test_eagle_enabled_removes_last_block():

2604

2605 # New request with same tokens + Eagle enabled

2606 req_eagle = make_request("eagle_divisible", token_ids, block_size, sha256)

2602 - computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)

2607 + computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)

2608

2609 # Should retain 1 block:

2610 # 1. Original 3 blocks → pop last hash → 2 matched blocks

@@ -2623,7 +2628,7 @@ def test_eagle_with_partial_blocks():

2628 req = make_request("partial_block_test", token_ids, block_size, sha256)

2629

2630 # Prime the cache

2626 - computed_blocks, _ = manager.get_computed_blocks(req)

2631 + computed_blocks, _, _ = manager.get_computed_blocks(req)

2632 manager.allocate_slots(

2633 req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks

2634 )

@@ -2631,7 +2636,7 @@ def test_eagle_with_partial_blocks():

2636

2637 # New request with Eagle enabled

2638 req_eagle = make_request("partial_eagle", token_ids, block_size, sha256)

2634 - computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)

2639 + computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)

2640 # Original match: 2 full blocks → Eagle removes 1 → 1 remaining

2641 assert len(computed_blocks.blocks[0]) == 1

2642 assert num_tokens == 1 * block_size

@@ -2664,7 +2669,7 @@ def test_eagle_with_sliding_window():

2669 req = make_request("partial_block_test", token_ids, block_size, sha256)

2670

2671 # Prime the cache

2667 - computed_blocks, _ = manager.get_computed_blocks(req)

2672 + computed_blocks, _, _ = manager.get_computed_blocks(req)

2673 manager.allocate_slots(

2674 req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks

2675 )

@@ -2675,7 +2680,7 @@ def test_eagle_with_sliding_window():

2680

2681 # New request with Eagle enabled

2682 req_eagle = make_request("partial_eagle", token_ids, block_size, sha256)

2678 - computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)

2683 + computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)

2684 # Original match: 2 full blocks → Eagle removes 1 → 1 remaining

2685 assert len(computed_blocks.blocks[0]) == 1

2686 assert num_tokens == 1 * block_size

@@ -2695,7 +2700,7 @@ def test_eagle_with_sliding_window():

2700 req_after_evict = make_request(

2701 "partial_eagle_after_evict", token_ids, block_size, sha256

2702 )

2698 - computed_blocks, num_tokens = manager.get_computed_blocks(req_after_evict)

2703 + computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_after_evict)

2704 # Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is

2705 # not considered. But after dropping the last matched block due to eagle,

2706 # there will be no matched prefix.

@@ -2754,7 +2759,7 @@ def test_eagle_swa_alignment_caches_extra_block():

2759 # Prime the cache with a long prompt (16 swa blocks = 4 aligned segments).

2760 token_ids = [i for i in range(16) for _ in range(block_size)]

2761 req0 = make_request("0", token_ids, block_size, sha256)

2757 - computed_blocks, _ = manager.get_computed_blocks(req0)

2762 + computed_blocks, _, _ = manager.get_computed_blocks(req0)

2763 blocks = manager.allocate_slots(

2764 req0,

2765 len(token_ids),

@@ -2768,7 +2773,7 @@ def test_eagle_swa_alignment_caches_extra_block():

2773 # Without the fix, ``num_computed_tokens`` is 0; with the fix, it lands at

2774 # an alignment boundary (multiple of 32 tokens, minus the EAGLE drop).

2775 req1 = make_request("1", token_ids, block_size, sha256)

2771 - _, num_computed_tokens = manager.get_computed_blocks(req1)

2776 + _, num_computed_tokens, _ = manager.get_computed_blocks(req1)

2777 assert num_computed_tokens > 0, (

2778 "EAGLE + SWA with sliding_window <= alignment failed to find any "

2779 "cache hit; the +1 block past each segment boundary must be cached."

@@ -2823,7 +2828,7 @@ def test_eagle_swa_boundary_caches_post_boundary_block():

2828

2829 token_ids = [i for i in range(5) for _ in range(block_size)]

2830 req0 = make_request("0", token_ids, block_size, sha256)

2826 - computed_blocks, _ = manager.get_computed_blocks(req0)

2831 + computed_blocks, _, _ = manager.get_computed_blocks(req0)

2832 blocks = manager.allocate_slots(

2833 req0,

2834 len(token_ids),

@@ -2838,7 +2843,7 @@ def test_eagle_swa_boundary_caches_post_boundary_block():

2843 manager.free(req0)

2844

2845 req1 = make_request("1", token_ids + [999], block_size, sha256)

2841 - _, num_computed_tokens = manager.get_computed_blocks(req1)

2846 + _, num_computed_tokens, _ = manager.get_computed_blocks(req1)

2847 assert num_computed_tokens == 4 * block_size

2848

2849

@@ -2879,7 +2884,7 @@ def test_eagle_grouped_swa_siblings_use_same_cache_mask():

2884

2885 token_ids = [i for i in range(9) for _ in range(block_size)]

2886 req0 = make_request("0", token_ids, block_size, sha256)

2882 - computed_blocks, _ = manager.get_computed_blocks(req0)

2887 + computed_blocks, _, _ = manager.get_computed_blocks(req0)

2888 blocks = manager.allocate_slots(

2889 req0,

2890 len(token_ids),

@@ -2894,7 +2899,7 @@ def test_eagle_grouped_swa_siblings_use_same_cache_mask():

2899 manager.free(req0)

2900

2901 req1 = make_request("1", token_ids + [999], block_size, sha256)

2897 - _, num_computed_tokens = manager.get_computed_blocks(req1)

2902 + _, num_computed_tokens, _ = manager.get_computed_blocks(req1)

2903 assert num_computed_tokens == 8 * block_size

2904

2905

@@ -2938,7 +2943,7 @@ def test_different_block_size():

2943 common_token_ids = [i for i in range(10) for _ in range(block_size)]

2944

2945 req0 = make_request("0", common_token_ids, block_size, sha256)

2941 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

2946 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

2947 assert not computed_blocks.blocks[0]

2948 assert not computed_blocks.blocks[1]

2949 assert num_computed_tokens == 0

@@ -2947,13 +2952,13 @@ def test_different_block_size():

2952 )

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

2954 req1 = make_request("1", common_token_ids[: 7 * block_size + 1], block_size, sha256)

2950 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

2955 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

2956 assert len(computed_blocks.blocks[0]) == 3

2957 assert len(computed_blocks.blocks[1]) == 6

2958 assert num_computed_tokens == 6 * 16

2959

2960 req2 = make_request("2", common_token_ids[: 6 * block_size + 1], block_size, sha256)

2956 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

2961 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)

2962 assert len(computed_blocks.blocks[0]) == 3

2963 assert len(computed_blocks.blocks[1]) == 6

2964 assert num_computed_tokens == 6 * 16

@@ -2967,7 +2972,7 @@ def test_different_block_size():

2972 manager.block_pool.cached_block_hash_to_block.pop(

2973 make_block_hash_with_group_id(req1.block_hashes[5], 1), 10

2974 )

2970 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

2975 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

2976 assert len(computed_blocks.blocks[0]) == 2

2977 assert len(computed_blocks.blocks[1]) == 4

2978 assert num_computed_tokens == 4 * 16

@@ -3020,7 +3025,7 @@ def test_hybrid_cache_blocks_swa_tail_window_only():

3025 # 8 hash-blocks of 8 tokens (64 tokens, two lcm-aligned segments).

3026 token_ids = [i for i in range(8) for _ in range(block_size)]

3027 req = make_request("0", token_ids, block_size, sha256)

3023 - computed_blocks, _ = manager.get_computed_blocks(req)

3028 + computed_blocks, _, _ = manager.get_computed_blocks(req)

3029 blocks = manager.allocate_slots(

3030 req,

3031 8 * block_size,

@@ -3089,7 +3094,7 @@ def test_hybrid_cache_blocks_clamped_to_lcm():

3094 # to 96 tokens — SWA caches 6 hashes, full-attn caches 3.

3095 token_ids = [i for i in range(7) for _ in range(block_size)]

3096 req = make_request("0", token_ids, block_size, sha256)

3092 - computed_blocks, _ = manager.get_computed_blocks(req)

3097 + computed_blocks, _, _ = manager.get_computed_blocks(req)

3098 blocks = manager.allocate_slots(

3099 req,

3100 7 * block_size,

@@ -3154,7 +3159,7 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch):

3159 # 96-token replay boundary, and the 128-token interval boundary.

3160 token_ids = [i for i in range(16) for _ in range(block_size)]

3161 req = make_request("0", token_ids, block_size, sha256)

3157 - computed_blocks, _ = manager.get_computed_blocks(req)

3162 + computed_blocks, _, _ = manager.get_computed_blocks(req)

3163 blocks = manager.allocate_slots(

3164 req,

3165 len(token_ids),

@@ -3296,13 +3301,13 @@ def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch):

3301

3302 token_ids = fill_request("fill_0", 0)

3303 replay_req = make_request("replay", token_ids[:1800], hash_block_size, sha256)

3299 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(replay_req)

3304 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(replay_req)

3305 assert num_computed_tokens == 1024

3306 assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256]

3307

3308 fill_request("fill_1", 100_000)

3309 replay_req = make_request("replay_again", token_ids[:1800], hash_block_size, sha256)

3305 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(replay_req)

3310 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(replay_req)

3311 assert num_computed_tokens == 1024

3312 assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256]

3313

@@ -3345,7 +3350,7 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatc

3350

3351 token_ids = [i for i in range(16) for _ in range(block_size)]

3352 req0 = make_request("0", token_ids, block_size, sha256)

3348 - computed_blocks, _ = manager.get_computed_blocks(req0)

3353 + computed_blocks, _, _ = manager.get_computed_blocks(req0)

3354 blocks = manager.allocate_slots(

3355 req0,

3356 len(token_ids),

@@ -3369,14 +3374,14 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatc

3374 assert retained_swa_block[0].ref_cnt == 0

3375

3376 req1 = make_request("1", token_ids, block_size, sha256)

3372 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

3377 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

3378 # Full prompt hits intentionally recompute the final block for logits, so

3379 # the longest usable hit is the previous LCM boundary: 96 tokens.

3380 assert num_computed_tokens == 12 * block_size

3381 assert len(computed_blocks.blocks[1]) == 12

3382

3383 shorter_req = make_request("2", token_ids[: 12 * block_size], block_size, sha256)

3379 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(shorter_req)

3384 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(shorter_req)

3385 assert num_computed_tokens == 0

3386 assert len(computed_blocks.blocks[1]) == 0

3387

@@ -3429,7 +3434,7 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch):

3434 # 104 tokens, and that tail is two 8-token blocks wide: hashes 11 and 12.

3435 token_ids = [i for i in range(15) for _ in range(block_size)] + [15] * 7

3436 req0 = make_request("0", token_ids, block_size, sha256)

3432 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

3437 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

3438 assert num_computed_tokens == 0

3439 blocks = manager.allocate_slots(

3440 req0,

@@ -3451,7 +3456,7 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch):

3456 manager.free(req0)

3457

3458 req1 = make_request("1", token_ids, block_size, sha256)

3454 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

3459 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

3460 assert num_computed_tokens == 12 * block_size

3461 assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 12]

3462

@@ -3858,7 +3863,7 @@ def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(monkeypatch):

3863

3864 token_ids = [i for i in range(16) for _ in range(block_size)]

3865 req = make_request("0", token_ids, block_size, sha256)

3861 - computed_blocks, _ = manager.get_computed_blocks(req)

3866 + computed_blocks, _, _ = manager.get_computed_blocks(req)

3867 blocks = manager.allocate_slots(

3868 req,

3869 len(token_ids),

@@ -3941,7 +3946,7 @@ def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch):

3946

3947 token_ids = [i for i in range(16) for _ in range(block_size)]

3948 req = make_request("0", token_ids, block_size, sha256)

3944 - computed_blocks, _ = manager.get_computed_blocks(req)

3949 + computed_blocks, _, _ = manager.get_computed_blocks(req)

3950 blocks = manager.allocate_slots(

3951 req,

3952 len(token_ids),

@@ -3965,7 +3970,7 @@ def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch):

3970

3971 # A replay of the same prompt hits the latest replayable boundary (240).

3972 replay = make_request("1", token_ids, block_size, sha256)

3968 - _, num_computed = manager.get_computed_blocks(replay)

3973 + _, num_computed, _ = manager.get_computed_blocks(replay)

3974 assert num_computed == 240

3975

3976

@@ -3977,7 +3982,7 @@ def test_pure_swa_retention_latest_only(monkeypatch):

3982

3983 token_ids = [i for i in range(16) for _ in range(block_size)]

3984 req = make_request("0", token_ids, block_size, sha256)

3980 - computed_blocks, _ = manager.get_computed_blocks(req)

3985 + computed_blocks, _, _ = manager.get_computed_blocks(req)

3986 blocks = manager.allocate_slots(

3987 req,

3988 len(token_ids),

@@ -3997,7 +4002,7 @@ def test_pure_swa_retention_latest_only(monkeypatch):

4002 assert cached == {14}

4003

4004 replay = make_request("1", token_ids, block_size, sha256)

4000 - _, num_computed = manager.get_computed_blocks(replay)

4005 + _, num_computed, _ = manager.get_computed_blocks(replay)

4006 assert num_computed == 240

4007

4008

@@ -4010,7 +4015,7 @@ def test_pure_swa_retention_dense_default_caches_all(monkeypatch):

4015

4016 token_ids = [i for i in range(16) for _ in range(block_size)]

4017 req = make_request("0", token_ids, block_size, sha256)

4013 - computed_blocks, _ = manager.get_computed_blocks(req)

4018 + computed_blocks, _, _ = manager.get_computed_blocks(req)

4019 blocks = manager.allocate_slots(

4020 req,

4021 len(token_ids),

@@ -4052,7 +4057,7 @@ def test_mamba_reachable_block_mask_sparsifies_retention():

4057 kv_cache_spec=spec,

4058 use_eagle=False,

4059 retention_interval=retention_interval,

4055 - num_prompt_tokens=num_prompt_tokens,

4060 + reachable_boundaries=(num_prompt_tokens - 1,),

4061 )

4062 return None if m is None else {i for i, v in enumerate(m) if v}

4063

@@ -4065,3 +4070,242 @@ def test_mamba_reachable_block_mask_sparsifies_retention():

4070 assert retained(64) == {3, 7, 11, 14, 15}

4071 # interval 0 -> only the latest replay boundary (block 14).

4072 assert retained(0) == {14}

4073 +

4074 +

4075 +def test_mamba_reachable_block_mask_pins_shared_prefix():

4076 + """A Marconi-detected shared prefix (``shared_prefix_boundary``) lands before

4077 + ``num_prompt`` so the replay-boundary rule alone would drop it. The mask must

4078 + pin the single state block ending on that boundary so sparse retention does

4079 + not defeat cross-request shared-prefix reuse."""

4080 + from vllm.v1.core.single_type_kv_cache_manager import MambaManager

4081 +

4082 + block_size = 16

4083 + spec = MambaSpec(

4084 + block_size=block_size,

4085 + shapes=(1, 1),

4086 + dtypes=(torch.float32,),

4087 + mamba_cache_mode="align",

4088 + )

4089 +

4090 + def retained(retention_interval, shared_prefix_boundary, end_block=16):

4091 + boundaries = [255] # replay boundary (num_prompt 256 - 1)

4092 + if shared_prefix_boundary:

4093 + boundaries.append(shared_prefix_boundary)

4094 + m = MambaManager.reachable_block_mask(

4095 + start_block=0,

4096 + end_block=end_block,

4097 + alignment_tokens=block_size,

4098 + kv_cache_spec=spec,

4099 + use_eagle=False,

4100 + retention_interval=retention_interval,

4101 + reachable_boundaries=boundaries,

4102 + )

4103 + return None if m is None else {i for i, v in enumerate(m) if v}

4104 +

4105 + # interval 0 keeps only replay boundary 14; the shared prefix at token 96

4106 + # (state block 96//16 - 1 = 5) is now pinned too.

4107 + assert retained(0, 96) == {5, 14}

4108 + # A non-aligned boundary floors to the enclosing aligned boundary.

4109 + assert retained(0, 100) == {5, 14}

4110 + # Coexists with segment tails (interval 64 -> {3,7,11,15} + replay 14).

4111 + assert retained(64, 96) == {3, 5, 7, 11, 14, 15}

4112 + # Dense default ignores the hint (nothing to sparsify).

4113 + assert retained(None, 96) is None

4114 + # Out-of-range boundary is a no-op (only replay 14 remains).

4115 + assert retained(0, 16 * block_size * 2) == {14}

4116 + # No boundary given -> unchanged replay-only behavior.

4117 + assert retained(0, 0) == {14}

4118 + assert retained(0, None) == {14}

4119 +

4120 +

4121 +def test_mamba_shared_prefix_survives_zero_retention(monkeypatch):

4122 + """Manager-level check of the full wiring: a pinned shared-prefix boundary

4123 + (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style

4124 + detection) keeps its Mamba state block cached under

4125 + ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``, which otherwise retains only the

4126 + end-of-prompt replay boundary. Without this, a shared prefix (junction

4127 + before ``num_prompt``) would be recomputed by every sharing request."""

4128 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0")

4129 + block_size = 16

4130 +

4131 + # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14.

4132 + token_ids = [i for i in range(16) for _ in range(block_size)]

4133 +

4134 + def cached_mamba_blocks(shared_prefix_boundary):

4135 + # Fresh manager per scenario so cached blocks don't leak between runs.

4136 + manager = make_kv_cache_manager(

4137 + _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"]),

4138 + max_model_len=8192,

4139 + enable_caching=True,

4140 + hash_block_size=block_size,

4141 + )

4142 + req = make_request("r", token_ids, block_size, sha256)

4143 + req.shared_prefix_boundary = shared_prefix_boundary

4144 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req)

4145 + blocks = manager.allocate_slots(

4146 + req, len(token_ids), num_computed, computed_blocks

4147 + )

4148 + assert blocks is not None

4149 + pool = manager.block_pool

4150 + return {

4151 + i

4152 + for i in range(16)

4153 + if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])

4154 + is not None

4155 + }

4156 +

4157 + # Without a pinned boundary, retention=0 keeps only the replay boundary (14).

4158 + assert cached_mamba_blocks(0) == {14}

4159 + # Pinning the shared prefix at token 96 (state block 5) retains it too, so a

4160 + # later request sharing that prefix can hit the Mamba state.

4161 + assert cached_mamba_blocks(96) == {5, 14}

4162 +

4163 +

4164 +def test_mamba_shared_prefix_reuse_under_zero_retention(monkeypatch):

4165 + """Full cross-request Marconi flow: a partial shared prefix cached by the

4166 + detecting request must stay reusable by a later request under

4167 + ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``. Without the pin the junction is

4168 + masked out and the later request misses; with it (and under dense) the reuse

4169 + is preserved."""

4170 + block_size = 16

4171 +

4172 + def last_req_hit(retention, pin):

4173 + if retention is None:

4174 + monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False)

4175 + else:

4176 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", str(retention))

4177 + manager = make_kv_cache_manager(

4178 + _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]),

4179 + max_model_len=8192,

4180 + enable_caching=True,

4181 + hash_block_size=block_size,

4182 + )

4183 + shared = [7 for _ in range(2 * block_size)] # 2-block shared prefix

4184 +

4185 + def distinct(v):

4186 + return [v for _ in range(2 * block_size)]

4187 +

4188 + # req0 primes the shared prefix's (dense) attention cache; align-mode

4189 + # Mamba keeps only its own tail, not the shared-prefix junction.

4190 + req0 = make_request("0", shared + distinct(50), block_size, sha256)

4191 + cb, nc, _ = manager.get_computed_blocks(req0)

4192 + manager.allocate_slots(req0, len(req0.all_token_ids), nc, cb)

4193 +

4194 + # req1 detects the shared prefix (attn hit, Mamba lag) and, via the

4195 + # Marconi chunk, caches the junction's Mamba state.

4196 + req1 = make_request("1", shared + distinct(60), block_size, sha256)

4197 + cb, nc, boundary = manager.get_computed_blocks(req1)

4198 + assert boundary == 2 * block_size # junction detected at 2 shared blocks

4199 + if pin:

4200 + req1.shared_prefix_boundary = boundary

4201 + # Simulate the Marconi chunk: schedule up to the junction (boundary - nc).

4202 + manager.allocate_slots(req1, boundary - nc, nc, cb)

4203 +

4204 + # req2 shares the same prefix: does it reuse the cached junction?

4205 + req2 = make_request("2", shared + distinct(70), block_size, sha256)

4206 + _, nc2, _ = manager.get_computed_blocks(req2)

4207 + return nc2

4208 +

4209 + # Dense retains the junction -> reuse works (baseline ceiling).

4210 + assert last_req_hit(retention=None, pin=False) == 2 * block_size

4211 + # retention=0 without the pin masks the junction out -> reuse lost (the bug).

4212 + assert last_req_hit(retention=0, pin=False) == 0

4213 + # retention=0 with the pin keeps the junction -> reuse restored.

4214 + assert last_req_hit(retention=0, pin=True) == 2 * block_size

4215 +

4216 +

4217 +def test_swa_reachable_block_mask_pins_shared_prefix():

4218 + """SWA analog of the Mamba pin: the shared-prefix junction must keep the

4219 + ``need``-block sliding-window tail ending on that boundary (not a single

4220 + block), so a windowed hit can land there under sparse retention."""

4221 + from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager

4222 +

4223 + block_size = 16

4224 +

4225 + def retained(retention, boundary, window, end_block=16):

4226 + spec = SlidingWindowSpec(

4227 + block_size=block_size,

4228 + num_kv_heads=1,

4229 + head_size=1,

4230 + dtype=torch.float32,

4231 + sliding_window=window,

4232 + )

4233 + boundaries = [255] # replay boundary (num_prompt 256 - 1)

4234 + if boundary:

4235 + boundaries.append(boundary)

4236 + m = SlidingWindowManager.reachable_block_mask(

4237 + start_block=0,

4238 + end_block=end_block,

4239 + alignment_tokens=block_size,

4240 + kv_cache_spec=spec,

4241 + use_eagle=False,

4242 + retention_interval=retention,

4243 + reachable_boundaries=boundaries,

4244 + )

4245 + return None if m is None else {i for i, v in enumerate(m) if v}

4246 +

4247 + # window == block_size -> need = cdiv(15, 16) = 1: single-block tail (like

4248 + # Mamba). Junction at token 96 -> block 5; replay boundary -> block 14.

4249 + assert retained(0, 96, block_size) == {5, 14}

4250 + # window == 3 * block_size -> need = cdiv(47, 16) = 3: the junction keeps a

4251 + # 3-block WINDOW {3,4,5} (the SWA distinction), plus replay window {12,13,14}.

4252 + assert retained(0, 96, 3 * block_size) == {3, 4, 5, 12, 13, 14}

4253 + # Coexists with segment tails (interval 64, need=1): {3,7,11,15} + replay 14

4254 + # + junction 5.

4255 + assert retained(64, 96, block_size) == {3, 5, 7, 11, 14, 15}

4256 + # Dense (all blocks reachable) ignores the hint.

4257 + assert retained(None, 96, block_size) is None

4258 + # No boundary -> unchanged replay-only behavior.

4259 + assert retained(0, 0, block_size) == {14}

4260 +

4261 +

4262 +def test_swa_shared_prefix_reuse_under_zero_retention(monkeypatch):

4263 + """SWA cross-request analog: a partial shared prefix's sliding-window tail

4264 + must stay reusable under ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``. Without

4265 + the pin the junction window is masked out and a later request misses; with

4266 + it (and under dense) reuse is preserved."""

4267 + block_size = 16

4268 +

4269 + def last_req_hit(retention, pin):

4270 + if retention is None:

4271 + monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False)

4272 + else:

4273 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", str(retention))

4274 + manager = make_kv_cache_manager(

4275 + _make_hybrid_kv_cache_config(block_size, 200, ["full", "sliding_window"]),

4276 + max_model_len=8192,

4277 + enable_caching=True,

4278 + hash_block_size=block_size,

4279 + )

4280 + shared = [7 for _ in range(4 * block_size)] # 4-block shared prefix

4281 +

4282 + def distinct(v):

4283 + return [v for _ in range(2 * block_size)]

4284 +

4285 + # req0 primes the (dense) full-attention cache; SWA under RET=0 keeps

4286 + # only its own replay window, not the shared-prefix boundary.

4287 + req0 = make_request("0", shared + distinct(50), block_size, sha256)

4288 + cb, nc, _ = manager.get_computed_blocks(req0)

4289 + manager.allocate_slots(req0, len(req0.all_token_ids), nc, cb)

4290 +

4291 + # req1 detects the shared prefix (full-attn hit, SWA lag). No chunk for

4292 + # SWA -- the pin fires during normal prefill.

4293 + req1 = make_request("1", shared + distinct(60), block_size, sha256)

4294 + cb, nc, boundary = manager.get_computed_blocks(req1)

4295 + if retention == 0:

4296 + assert boundary == 4 * block_size # junction detected when SWA lags

4297 + if pin and boundary:

4298 + req1.shared_prefix_boundary = boundary

4299 + manager.allocate_slots(req1, len(req1.all_token_ids), nc, cb)

4300 +

4301 + # req2 shares the same prefix: does its SWA window reuse the junction?

4302 + req2 = make_request("2", shared + distinct(70), block_size, sha256)

4303 + _, nc2, _ = manager.get_computed_blocks(req2)

4304 + return nc2

4305 +

4306 + # Dense keeps every SWA tail -> reuse works (ceiling).

4307 + assert last_req_hit(retention=None, pin=False) == 4 * block_size

4308 + # retention=0 without the pin drops the junction window -> reuse lost (bug).

4309 + assert last_req_hit(retention=0, pin=False) == 0

4310 + # retention=0 with the pin keeps the junction window -> reuse restored.

4311 + assert last_req_hit(retention=0, pin=True) == 4 * block_size

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

+4 -1

@@ -235,6 +235,9 @@ class MooncakeStoreCoordinator:

235 manager_cls = KVCacheSpecRegistry.get_manager_class(spec)

236 assert manager_cls is not None

237 use_eagle = g_idx in self.eagle_group_ids

238 + reachable_boundaries = (

239 + () if num_prompt_tokens is None else (num_prompt_tokens - 1,)

240 + )

241 mask = manager_cls.reachable_block_mask(

242 start_block=start_chunk,

243 end_block=end_chunk,

@@ -242,7 +245,7 @@ class MooncakeStoreCoordinator:

245 kv_cache_spec=spec,

246 use_eagle=use_eagle,

247 retention_interval=retention_interval,

245 - num_prompt_tokens=num_prompt_tokens,

248 + reachable_boundaries=reachable_boundaries,

249 )

250 if mask is not None:

251 assert len(mask) == end_chunk - start_chunk

vllm/v1/core/kv_cache_coordinator.py

+20 -12

@@ -370,7 +370,10 @@ class KVCacheCoordinator(ABC):

370 self,

371 block_hashes: list[BlockHash],

372 max_cache_hit_length: int,

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

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

374 + """Returns the per-group hit blocks, the hit length, and the number of

375 + ``num_uncached_common_prefix_tokens`` (a shared prefix that a

376 + sparse-retention group has not cached yet; 0 unless hybrid)."""

377 pass

378

379 def new_step_starts(self) -> None:

@@ -422,11 +425,11 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

425 self,

426 block_hashes: list[BlockHash],

427 max_cache_hit_length: int,

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

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

429 blocks: tuple[list[KVCacheBlock], ...] = tuple(

430 [] for _ in range(self.num_single_type_manager)

431 )

429 - return blocks, 0

432 + return blocks, 0, 0

433

434

435 class UnitaryKVCacheCoordinator(KVCacheCoordinator):

@@ -486,7 +489,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

489 self,

490 block_hashes: list[BlockHash],

491 max_cache_hit_length: int,

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

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

493 hit_blocks, hit_length = self.single_type_managers[0].find_longest_cache_hit(

494 block_hashes=block_hashes,

495 max_length=max_cache_hit_length,

@@ -498,7 +501,8 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

501 dcp_world_size=self.dcp_world_size,

502 pcp_world_size=self.pcp_world_size,

503 )

501 - return hit_blocks, hit_length

504 + # Single group: nothing "uncached common" -- no other group to lag it.

505 + return hit_blocks, hit_length, 0

506

507

508 class SpecGroup(NamedTuple):

@@ -675,7 +679,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

679 self,

680 block_hashes: list[BlockHash],

681 max_cache_hit_length: int,

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

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

683 """

684 Find the longest cache hit using an iterative fixed-point algorithm.

685

@@ -691,7 +695,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

695 Returns:

696 A tuple containing:

697 - A tuple of the cache hit blocks for each single type manager.

694 - - The number of tokens of the longest cache hit.

698 + - The number of tokens of the reconciled (combined) cache hit.

699 + - ``num_uncached_common_prefix_tokens``: a shared prefix that a

700 + sparse-retention group has not cached yet (0 unless hybrid).

701 """

702

703 num_groups = len(self.kv_cache_config.kv_cache_groups)

@@ -794,12 +800,14 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

800 del blks[num_blocks:]

801 hit_length_by_group[group_id] = hit_length

802

797 - # Uncached shared prefix detection: If any attn. group cached a longer prefix

798 - # than the current prefix, it is an uncached common prefix across requests:

799 - self.num_uncached_common_prefix_tokens = longest_hit_length - hit_length

800 - return tuple(

803 + # Uncached shared prefix detection: if any attn. group cached a longer

804 + # prefix than the reconciled hit, it is an uncached common prefix across

805 + # requests that a sparse-retention group hasn't cached yet.

806 + num_uncached_common_prefix_tokens = longest_hit_length - hit_length

807 + cache_hit_blocks = tuple(

808 blocks if blocks is not None else [] for blocks in hit_blocks_by_group

802 - ), hit_length

809 + )

810 + return cache_hit_blocks, hit_length, num_uncached_common_prefix_tokens

811

812 def find_longest_cache_hit_per_group(

813 self,

vllm/v1/core/kv_cache_manager.py

+18 -4

@@ -204,7 +204,7 @@ class KVCacheManager:

204 self.prefix_cache_stats = PrefixCacheStats()

205 return stats

206

207 - def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:

207 + def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int]:

208 """Get the computed (cached) blocks for the request.

209 Note that the computed blocks must be full.

210

@@ -215,13 +215,18 @@ class KVCacheManager:

215 A tuple containing:

216 - A list of blocks that are computed for the request.

217 - The number of computed tokens.

218 + - ``shared_prefix_boundary``: the block-aligned token position of

219 + a shared prefix that a sparse-retention group (Mamba / sliding

220 + window) has not cached yet (Marconi-style APC), or 0 if none.

221 + Pinned so ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL`` does not drop

222 + the junction and defeat cross-request reuse.

223 """

224 # We skip finding the prefix cache hit when prefix caching is

225 # disabled or the request is marked as skipping kv cache read

226 # (which happens when the request requires prompt logprobs

227 # or calls a pooling model with all pooling).

228 if not self.enable_caching or request.skip_reading_prefix_cache:

224 - return self.empty_kv_cache_blocks, 0

229 + return self.empty_kv_cache_blocks, 0, 0

230

231 # NOTE: When all tokens hit the cache, we must recompute the last token

232 # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.

@@ -230,7 +235,7 @@ class KVCacheManager:

235 # num_computed_tokens to be block-size aligned. Removing this limitation

236 # could slightly improve performance in the future.

237 max_cache_hit_length = request.num_tokens - 1

233 - computed_blocks, num_new_computed_tokens = (

238 + computed_blocks, num_new_computed_tokens, num_uncached = (

239 self.coordinator.find_longest_cache_hit(

240 request.block_hashes, max_cache_hit_length

241 )

@@ -256,6 +261,14 @@ class KVCacheManager:

261 group_idx,

262 )

263

264 + # The junction to pin is where the lagging sparse-retention group stops

265 + # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e.

266 + # the longest single-group hit. Sub-block gaps are left to the mask,

267 + # which floors to the alignment boundary (a no-op there).

268 + shared_prefix_boundary = (

269 + num_new_computed_tokens + num_uncached if num_uncached else 0

270 + )

271 +

272 if self.log_stats:

273 assert self.prefix_cache_stats is not None

274 self.prefix_cache_stats.record(

@@ -264,7 +277,8 @@ class KVCacheManager:

277 preempted=request.num_preemptions > 0,

278 )

279

267 - return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens

280 + blocks = self.create_kv_cache_blocks(computed_blocks)

281 + return blocks, num_new_computed_tokens, shared_prefix_boundary

282

283 def allocate_slots(

284 self,

vllm/v1/core/sched/scheduler.py

+71 -89

@@ -350,85 +350,69 @@ class Scheduler(SchedulerInterface):

350 num_new_tokens: int,

351 num_new_local_computed_tokens: int = 0,

352 num_external_computed_tokens: int = 0,

353 - num_uncached_common_prefix_tokens: int = 0,

353 ) -> int:

355 - num_computed_tokens = (

354 + """Clip a prefill chunk so it ends where Mamba state must be cached.

355 +

356 + In "align" cache mode the SSM state is only materialized at chunk

357 + ends, so chunk ends are steered onto cacheable positions: block

358 + boundaries by default, plus mandatory early stops (the prompt's

359 + partial-tail hash boundary, a detected shared-prefix junction).

360 + """

361 + start = (

362 request.num_computed_tokens

363 + num_new_local_computed_tokens

364 + num_external_computed_tokens

365 )

360 - # Perform block-aligned splitting at prefill phase, including:

361 - # * non-resumed requests: num_computed_tokens < num_prompt_tokens + 0

362 - # * resumed requests: num_computed_tokens < (

363 - # num_prompt_tokens + num_output_tokens

364 - # )

365 - # NOTE: Use `request.num_tokens - 1` to bypass normal decoding.

366 - if num_computed_tokens < max(request.num_prompt_tokens, request.num_tokens - 1):

367 - # To enable block-aligned caching of the Mamba state, `num_new_tokens`

368 - # must be a multiple of `block_size`.

369 - # As an exception, if `num_new_tokens` is less than `block_size`, the

370 - # state is simply not cached, requiring no special handling.

371 - # Additionally, when Eagle mode is enabled, FullAttn prunes the last

372 - # matching block. To prevent this from causing a Mamba cache miss, the

373 - # last chunk must be not smaller than `block_size`.

374 - block_size = self.cache_config.block_size

375 - last_cache_position = request.num_tokens - request.num_tokens % block_size

376 - # eagle prune

377 - if self.use_eagle:

378 - last_cache_position = max(last_cache_position - block_size, 0)

379 - num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens

380 - next_boundary = (num_computed_tokens // block_size + 1) * block_size

381 - if (

382 - num_computed_tokens % block_size != 0

383 - and next_boundary <= last_cache_position

384 - and num_computed_tokens_after_sched > next_boundary

385 - ):

386 - # Resumed mid-block (partial hash hit): stop the first chunk

387 - # at the next block boundary so chunk ends re-align to blocks

388 - # and the boundary state is materialized before running past.

389 - num_new_tokens = next_boundary - num_computed_tokens

390 - elif num_computed_tokens_after_sched < last_cache_position:

391 - # Align the chunk END (not its length) to block_size;

392 - # identical to flooring the length when the start is

393 - # block-aligned. May yield 0 (insufficient budget to reach

394 - # the next boundary); the caller then skips the request.

395 - aligned_end = num_computed_tokens_after_sched // block_size * block_size

396 - num_new_tokens = max(aligned_end - num_computed_tokens, 0)

397 - elif (

398 - num_computed_tokens

399 - < last_cache_position

400 - < num_computed_tokens_after_sched

401 - ):

402 - # force to cache the last chunk

403 - num_new_tokens = last_cache_position - num_computed_tokens

404 - elif self.mamba_partial_cache_hit:

405 - # Prefill of the final partial block: stop once at the

406 - # prompt's last hash boundary so the mamba partial tail entry

407 - # can be registered.

408 - tail_boundary = (

409 - request.num_prompt_tokens

410 - // self.hash_block_size

411 - * self.hash_block_size

412 - )

413 - if (

414 - num_computed_tokens

415 - < tail_boundary

416 - < num_computed_tokens_after_sched

417 - and tail_boundary < request.num_prompt_tokens

418 - and tail_boundary > last_cache_position

419 - ):

420 - num_new_tokens = tail_boundary - num_computed_tokens

421 -

422 - # Marconi cache admission optimization:

423 - # cache common prefixes by scheduling num_new_tokens = common prefix length

424 - if (

425 - num_uncached_common_prefix_tokens >= block_size

426 - and num_new_tokens > num_uncached_common_prefix_tokens

427 - ):

428 - num_new_tokens = num_uncached_common_prefix_tokens

429 - # keep alignment to block_size

430 - num_new_tokens = num_new_tokens // block_size * block_size

431 - return num_new_tokens

366 + # Split only during prefill: `request.num_tokens - 1` extends this to

367 + # resumed requests replaying their output tokens.

368 + if start >= max(request.num_prompt_tokens, request.num_tokens - 1):

369 + return num_new_tokens

370 +

371 + block_size = self.cache_config.block_size

372 + # The last block-aligned position whose state can be cached. With

373 + # Eagle, FullAttn prunes the last matching block, so back off one

374 + # block to avoid a Mamba cache miss.

375 + last_cache_position = request.num_tokens - request.num_tokens % block_size

376 + if self.use_eagle:

377 + last_cache_position = max(last_cache_position - block_size, 0)

378 +

379 + end = start + num_new_tokens

380 + # Until `last_cache_position`, chunk ends must land on block

381 + # boundaries. May yield an empty chunk (budget cannot reach the next

382 + # boundary); the caller then skips the request.

383 + if end < last_cache_position:

384 + end = end // block_size * block_size

385 +

386 + next_block_boundary = (start // block_size + 1) * block_size

387 + tail_boundary = (

388 + request.num_prompt_tokens // self.hash_block_size * self.hash_block_size

389 + if self.mamba_partial_cache_hit

390 + else 0

391 + )

392 + stops = (

393 + # Resumed mid-block (fine-grained partial hash hit): re-align to

394 + # the block grid before running on, so the crossed boundary's

395 + # state is materialized (unless it is past the cacheable range).

396 + next_block_boundary

397 + if start % block_size != 0 and next_block_boundary <= last_cache_position

398 + else 0,

399 + # Never run past the last cacheable block boundary mid-chunk.

400 + last_cache_position,

401 + # Fine-grained hits: the prompt's partial-tail entry can only be

402 + # registered by a chunk ending exactly at its last hash boundary.

403 + tail_boundary

404 + if last_cache_position < tail_boundary < request.num_prompt_tokens

405 + else 0,

406 + # Marconi shared-prefix junction, block-floored (a sub-block

407 + # junction's state is not separately cacheable): cache its state

408 + # so sibling requests sharing the prefix can reuse it.

409 + start + (request.shared_prefix_boundary - start) // block_size * block_size

410 + if start < request.shared_prefix_boundary < end

411 + else 0,

412 + )

413 + # Stop at the earliest mandatory position strictly inside the chunk.

414 + end = min((s for s in stops if start < s < end), default=end)

415 + return max(end - start, 0)

416

417 def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:

418 self.current_step += 1

@@ -718,7 +702,6 @@ class Scheduler(SchedulerInterface):

702 num_external_computed_tokens = 0

703 load_kv_async = False

704 connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0

721 - num_uncached_common_prefix_tokens = 0

705

706 # Get already-cached tokens.

707 if request.num_computed_tokens == 0:

@@ -733,8 +716,7 @@ class Scheduler(SchedulerInterface):

716 ):

717 computed, per_group_hits = (

718 self.kv_cache_manager.coordinator.find_longest_cache_hit_per_group(

736 - request.block_hashes,

737 - request.num_tokens - 1,

719 + request.block_hashes, request.num_tokens - 1

720 )

721 )

722 new_computed_blocks = (

@@ -750,6 +732,9 @@ class Scheduler(SchedulerInterface):

732 # the last block) is transferred unconditionally by

733 # _apply_prefix_caching in nixl/worker.py.

734 num_new_local_computed_tokens = max(per_group_hits)

735 + # The per-group lookup does not detect an uncached shared

736 + # prefix, so there is no junction to pin in this path.

737 + request.shared_prefix_boundary = 0

738 if self.kv_cache_manager.log_stats:

739 assert self.kv_cache_manager.prefix_cache_stats is not None

740 self.kv_cache_manager.prefix_cache_stats.record(

@@ -758,17 +743,15 @@ class Scheduler(SchedulerInterface):

743 preempted=request.num_preemptions > 0,

744 )

745 else:

761 - new_computed_blocks, num_new_local_computed_tokens = (

762 - self.kv_cache_manager.get_computed_blocks(request)

763 - )

764 -

765 - # In case of hybrid models, obtain hint for Marconi-style APC logic

766 - if self.has_mamba_layers:

767 - num_uncached_common_prefix_tokens = getattr(

768 - self.kv_cache_manager.coordinator,

769 - "num_uncached_common_prefix_tokens",

770 - 0,

771 - )

746 + (

747 + new_computed_blocks,

748 + num_new_local_computed_tokens,

749 + # Junction to pin (Marconi-style APC) so its

750 + # sparse-retention state (Mamba block / sliding-window

751 + # tail) survives retention and serves a later hit; 0

752 + # if no uncached shared prefix was detected.

753 + request.shared_prefix_boundary,

754 + ) = self.kv_cache_manager.get_computed_blocks(request)

755

756 # Get externally-cached tokens if using a KVConnector.

757 if self.connector is not None:

@@ -906,7 +889,6 @@ class Scheduler(SchedulerInterface):

889 num_new_tokens,

890 num_new_local_computed_tokens,

891 num_external_computed_tokens,

909 - num_uncached_common_prefix_tokens,

892 )

893 if num_new_tokens == 0:

894 break

vllm/v1/core/single_type_kv_cache_manager.py

+37 -27

@@ -419,6 +419,13 @@ class SingleTypeKVCacheManager(ABC):

419 if num_cached_blocks >= num_full_blocks:

420 return

421

422 + # Token boundaries whose reachable tail must be retained under sparse

423 + # retention: the replay boundary (``num_prompt - 1``, capped by

424 + # ``get_computed_blocks``) and any detected shared-prefix junction.

425 + reachable_boundaries = [request.num_prompt_tokens - 1]

426 + if request.shared_prefix_boundary:

427 + reachable_boundaries.append(request.shared_prefix_boundary)

428 +

429 block_mask = self.reachable_block_mask(

430 start_block=num_cached_blocks,

431 end_block=num_full_blocks,

@@ -426,7 +433,7 @@ class SingleTypeKVCacheManager(ABC):

433 kv_cache_spec=self.kv_cache_spec,

434 use_eagle=self.use_eagle,

435 retention_interval=retention_interval,

429 - num_prompt_tokens=request.num_prompt_tokens,

436 + reachable_boundaries=reachable_boundaries,

437 )

438 self.block_pool.cache_full_blocks(

439 request=request,

@@ -449,14 +456,15 @@ class SingleTypeKVCacheManager(ABC):

456 kv_cache_spec: KVCacheSpec,

457 use_eagle: bool,

458 retention_interval: int | None = None,

452 - num_prompt_tokens: int | None = None,

459 + reachable_boundaries: Sequence[int] = (),

460 ) -> list[bool] | None:

461 """Per-block mask for ``cache_full_blocks``. ``None`` means cache

462 every (non-null) block — the default for full attention.

463

457 - Subclasses with sparse hit semantics (SWA) override this to skip

458 - blocks that can never serve a hit at any alignment-aligned prefix

459 - length.

464 + Subclasses with sparse hit semantics (SWA / Mamba) override this to skip

465 + blocks that can never serve a hit at any alignment-aligned prefix length.

466 + ``reachable_boundaries`` are token positions whose reachable tail must be

467 + retained; the base (dense) policy ignores them.

468 """

469 return None

470

@@ -962,7 +970,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

970 kv_cache_spec: KVCacheSpec,

971 use_eagle: bool,

972 retention_interval: int | None = None,

965 - num_prompt_tokens: int | None = None,

973 + reachable_boundaries: Sequence[int] = (),

974 ) -> list[bool] | None:

975 assert isinstance(kv_cache_spec, SlidingWindowSpec)

976 if alignment_tokens is None:

@@ -1002,18 +1010,16 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

1010 if i >= shift and (i - shift) % per_segment >= per_segment - need:

1011 mask[i - start_block] = True

1012

1005 - # (2) Replay-boundary tail. ``get_computed_blocks`` caps hits at

1006 - # ``num_prompt - 1`` (to recompute the last token's logits), so an exact

1007 - # prompt replay can only land on the latest *fine*-aligned boundary.

1008 - # Sparse retention would otherwise skip it, so keep its tail explicitly.

1009 - if retention_interval is not None and num_prompt_tokens is not None:

1010 - latest = (num_prompt_tokens - 1) // alignment_tokens * alignment_tokens

1011 - prompt_end_block = latest // block_size + shift

1012 - for i in range(

1013 - max(start_block, prompt_end_block - need),

1014 - min(end_block, prompt_end_block),

1015 - ):

1016 - mask[i - start_block] = True

1013 + # (2) Reachable-boundary tails: the replay boundary (``num_prompt - 1``,

1014 + # capped by ``get_computed_blocks``) and any shared-prefix junction. Both

1015 + # land before segments would cover them under sparse retention, so keep

1016 + # the ``need``-block tail ending on each boundary explicitly.

1017 + if retention_interval is not None:

1018 + for boundary_tokens in reachable_boundaries:

1019 + aligned = boundary_tokens // alignment_tokens * alignment_tokens

1020 + end = aligned // block_size + shift

1021 + for j in range(max(start_block, end - need), min(end_block, end)):

1022 + mask[j - start_block] = True

1023

1024 return mask

1025

@@ -1321,15 +1327,19 @@ class MambaManager(SingleTypeKVCacheManager):

1327 kv_cache_spec: KVCacheSpec,

1328 use_eagle: bool,

1329 retention_interval: int | None = None,

1324 - num_prompt_tokens: int | None = None,

1330 + reachable_boundaries: Sequence[int] = (),

1331 ) -> list[bool] | None:

1332 """Sparse Mamba state-snapshot retention.

1333

1334 ``retention_interval``:

1335

1336 ``None`` -> dense (cache every block; default, unchanged behavior)

1331 - ``0`` -> keep only the latest replay boundary

1337 + ``0`` -> keep only the ``reachable_boundaries`` states

1338 ``> 0`` -> keep one state per ``retention_interval``-sized segment

1339 +

1340 + ``reachable_boundaries`` are proven reuse points (the replay boundary and

1341 + any cross-request shared-prefix junction, Marconi-style APC); their

1342 + boundary state is always kept so sparse retention does not defeat reuse.

1343 """

1344 if retention_interval is None or alignment_tokens is None:

1345 # Dense caching (default) or no alignment constraint imposed.

@@ -1354,13 +1364,13 @@ class MambaManager(SingleTypeKVCacheManager):

1364 for i in range(first_boundary - start_block, len(mask), per_segment):

1365 mask[i] = True

1366

1357 - # (2) Replay boundary. ``get_computed_blocks`` caps hits at

1358 - # ``num_prompt - 1``, so an exact prompt replay lands on the latest

1359 - # fine-aligned boundary. Sparse retention would otherwise skip its

1360 - # state, so keep it explicitly.

1361 - if num_prompt_tokens is not None:

1362 - latest = (num_prompt_tokens - 1) // alignment_tokens * alignment_tokens

1363 - boundary_block = latest // block_size - 1

1367 + # (2) Reachable-boundary states: the replay boundary (``num_prompt - 1``,

1368 + # capped by ``get_computed_blocks``) and any shared-prefix junction, both

1369 + # of which segments would otherwise skip under sparse retention. A Mamba

1370 + # hit needs exactly the single state block ending on the boundary.

1371 + for boundary_tokens in reachable_boundaries:

1372 + aligned = boundary_tokens // alignment_tokens * alignment_tokens

1373 + boundary_block = aligned // block_size - 1

1374 if start_block <= boundary_block < end_block:

1375 mask[boundary_block - start_block] = True

1376

vllm/v1/request.py

@@ -182,6 +182,11 @@ class Request:

182 # True if this request is scheduled as a non-final prefill chunk.

183 self.is_prefill_chunk = False

184

185 + # Block-aligned token position of a proven shared prefix worth pinning

186 + # in the (sparse) prefix cache; 0 means none. Set at admission for

187 + # hybrid/Mamba models when a shared prefix is detected (Marconi-style).

188 + self.shared_prefix_boundary = 0

189 +

190 # The number of NaNs in logits. A value greater than 0

191 # indicates that the output is corrupted

192 self.num_nans_in_logits = 0

vllm/v1/simple_kv_offload/manager.py

+1 -1

@@ -262,7 +262,7 @@ class SimpleCPUOffloadScheduler:

262 max_hit_len = request.num_tokens - 1 - num_computed_tokens

263 if max_hit_len <= 0:

264 return 0, False

265 - cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(

265 + cpu_hit_blocks, hit_length, _ = self.cpu_coordinator.find_longest_cache_hit(

266 remaining_hashes, max_hit_len

267 )

268