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 6, 09:30 UTC
fix: disambiguate multimodal prefix cache keys (#36708)
Signed-off-by: tianshu.yu <tianshuyu.formal@gmail.com>
tianshu-Michael-yu committed
Mar 19, 2026 at 19:33 UTC
269bf46d99f1df74e4d779f9c52c74002e057a17
3 files changed
+29
-16
tests/v1/core/test_kv_cache_utils.py
+10
-8
@@ -447,12 +447,12 @@ def test_generate_block_hash_extra_keys():
447
448
# Test with no extra keys
449
extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 5, 0)
450
- assert extra_keys == ("hash1",)
450
+ assert extra_keys == (("hash1", 0),)
451
assert next_mm_idx == 1
452
453
# Test with partial overlap
454
extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 3, 8, 0)
455
- assert extra_keys == ("hash1",)
455
+ assert extra_keys == (("hash1", -3),)
456
assert next_mm_idx == 1
457
458
# Test with no overlap
@@ -462,7 +462,7 @@ def test_generate_block_hash_extra_keys():
462
463
# Test with multiple extra keys
464
extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 15, 0)
465
- assert extra_keys == ("hash1", "hash2")
465
+ assert extra_keys == (("hash1", 0), ("hash2", 10))
466
assert next_mm_idx == 2
467
468
@@ -513,7 +513,7 @@ def test_generate_block_hash_extra_keys_cache_salt():
513
514
# Test with no extra keys
515
extra_keys, next_mm_idx = generate_block_hash_extra_keys(request_mm, 0, 5, 0)
516
- assert extra_keys == ("hash1", "salt")
516
+ assert extra_keys == (("hash1", 0), "salt")
517
assert next_mm_idx == 1
518
519
@@ -637,8 +637,10 @@ def test_request_block_hasher(hash_fn):
637
638
block_hashes = request.block_hashes
639
assert len(block_hashes) == 2
640
- assert block_hashes[0] == hash_fn((kv_cache_utils.NONE_HASH, (0, 1, 2), ("hash1",)))
641
- assert block_hashes[1] == hash_fn((block_hashes[0], (3, 4, 5), ("hash2",)))
640
+ assert block_hashes[0] == hash_fn(
641
+ (kv_cache_utils.NONE_HASH, (0, 1, 2), (("hash1", 0),))
642
+ )
643
+ assert block_hashes[1] == hash_fn((block_hashes[0], (3, 4, 5), (("hash2", 0),)))
644
645
646
@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
@@ -1973,7 +1975,7 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes
1975
(
1976
kv_cache_utils.NONE_HASH,
1977
tuple(prompt_token_ids[:block_size]),
1976
- ("hash1", block1_embeds_hash),
1978
+ (("hash1", 0), block1_embeds_hash),
1979
)
1980
)
1981
assert block_hashes[0] == expected_hash1
@@ -1985,7 +1987,7 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes
1987
(
1988
block_hashes[0],
1989
tuple(prompt_token_ids[block_size:num_tokens]),
1988
- ("hash2", block2_embeds_hash),
1990
+ (("hash2", 0), block2_embeds_hash),
1991
)
1992
)
1993
assert block_hashes[1] == expected_hash2
tests/v1/core/test_prefix_caching.py
+12
-4
@@ -1570,20 +1570,24 @@ def test_mm_prefix_caching():
1570
block_hashes = req0.block_hashes
1571
assert len(block_hashes) == 3
1572
assert block_hashes[0] == sha256(
1573
- (kv_cache_utils.NONE_HASH, tuple(all_token_ids[:block_size]), ("aaa",))
1573
+ (
1574
+ kv_cache_utils.NONE_HASH,
1575
+ tuple(all_token_ids[:block_size]),
1576
+ (("aaa", 11),),
1577
+ )
1578
)
1579
assert block_hashes[1] == sha256(
1580
(
1581
block_hashes[0],
1582
tuple(all_token_ids[block_size : block_size * 2]),
1579
- ("aaa", "bbb"),
1583
+ (("aaa", -5), ("bbb", 14)),
1584
)
1585
)
1586
assert block_hashes[2] == sha256(
1587
(
1588
block_hashes[1],
1589
tuple(all_token_ids[block_size * 2 : block_size * 3]),
1586
- ("bbb",),
1590
+ (("bbb", -2),),
1591
)
1592
)
1593
@@ -1603,7 +1607,11 @@ def test_mm_prefix_caching():
1607
assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
1608
assert len(block_hashes) == 4
1609
assert block_hashes[3] == sha256(
1606
- (block_hashes[2], tuple(all_token_ids[3 * block_size :] + [8] * 5), ("ccc",))
1610
+ (
1611
+ block_hashes[2],
1612
+ tuple(all_token_ids[3 * block_size :] + [8] * 5),
1613
+ (("ccc", 0),),
1614
+ )
1615
)
1616
1617
# Cache hit.
vllm/v1/core/kv_cache_utils.py
+7
-4
@@ -413,7 +413,7 @@ def _gen_mm_extra_hash_keys(
413
# We do not need to check all mm inputs if the start token index is out of
414
# range. This usually happens in the late prefill phase and decoding phase.
415
last_pos = mm_features[-1].mm_position
416
- if last_pos.offset + last_pos.length < start_token_idx:
416
+ if last_pos.offset + last_pos.length <= start_token_idx:
417
return extra_keys, start_mm_idx
418
419
# Support start_mm_idx == -1 to indicate the last mm input.
@@ -428,13 +428,16 @@ def _gen_mm_extra_hash_keys(
428
offset = mm_feature.mm_position.offset
429
length = mm_feature.mm_position.length
430
if end_token_idx > offset:
431
- if start_token_idx > offset + length:
431
+ if start_token_idx >= offset + length:
432
# This block has passed the current mm input.
433
curr_mm_idx += 1
434
continue
435
436
- # The block contains the current mm input.
437
- extra_keys.append(mm_feature.identifier)
436
+ # The block contains the current mm input. Include its offset
437
+ # relative to the start of the block so prefix-cache keys stay
438
+ # distinct when the same MM item appears at different positions
439
+ # within otherwise-identical placeholder blocks.
440
+ extra_keys.append((mm_feature.identifier, offset - start_token_idx))
441
442
if end_token_idx >= offset + length:
443
# If this block contains the end of the current mm input,