@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 765872e7ed

[Core][BugFix] Tag prefix-cache extra keys by source (#51899)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> Co-authored-by: Tai An <antai12232931@outlook.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Clinton Thomas committed Sep 29, 2026 at 21:24 UTC 765872e7ed5b6f372f0898a67e48d3365fa28139

7 files changed +131 -43

tests/v1/core/test_kv_cache_utils.py

+64 -16

@@ -54,6 +54,7 @@ from vllm.v1.core.kv_cache_utils import (

54 is_kv_cache_spec_uniform,

55 make_block_hash_with_group_id,

56 tensor_data,

57 + to_event_extra_keys,

58 )

59 from vllm.v1.hisparse.layout import (

60 create_hisparse_layout,

@@ -822,12 +823,12 @@ def test_generate_block_hash_extra_keys():

823

824 # Test with no extra keys

825 extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 5, 0)

825 - assert extra_keys == (("hash1", 0),)

826 + assert extra_keys == (("mm", "hash1", 0),)

827 assert next_mm_idx == 1

828

829 # Test with partial overlap

830 extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 3, 8, 0)

830 - assert extra_keys == (("hash1", -3),)

831 + assert extra_keys == (("mm", "hash1", -3),)

832 assert next_mm_idx == 1

833

834 # Test with no overlap

@@ -837,7 +838,7 @@ def test_generate_block_hash_extra_keys():

838

839 # Test with multiple extra keys

840 extra_keys, next_mm_idx = generate_block_hash_extra_keys(request, 0, 15, 0)

840 - assert extra_keys == (("hash1", 0), ("hash2", 10))

841 + assert extra_keys == (("mm", "hash1", 0), ("mm", "hash2", 10))

842 assert next_mm_idx == 2

843

844

@@ -865,9 +866,9 @@ def test_generate_block_hash_extra_keys_cache_salt():

866

867 # salt is added for the first token

868 extra_keys, _ = generate_block_hash_extra_keys(request, 0, 1, 0)

868 - assert extra_keys == ("salt",)

869 + assert extra_keys == (("cache_salt", "salt"),)

870 extra_keys, _ = generate_block_hash_extra_keys(request, 0, 10, 0)

870 - assert extra_keys == ("salt",)

871 + assert extra_keys == (("cache_salt", "salt"),)

872

873 # no salt added for other tokens

874 extra_keys, _ = generate_block_hash_extra_keys(request, 1, 2, 0)

@@ -888,7 +889,7 @@ def test_generate_block_hash_extra_keys_cache_salt():

889

890 # Test with no extra keys

891 extra_keys, next_mm_idx = generate_block_hash_extra_keys(request_mm, 0, 5, 0)

891 - assert extra_keys == (("hash1", 0), "salt")

892 + assert extra_keys == (("mm", "hash1", 0), ("cache_salt", "salt"))

893 assert next_mm_idx == 1

894

895

@@ -906,13 +907,13 @@ def test_generate_block_hash_extra_keys_prompt_embeds():

907 extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0)

908 expected_embeds = prompt_embeds[0:5]

909 expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()

909 - assert extra_keys == (expected_hash,)

910 + assert extra_keys == (("prompt_embeds", expected_hash),)

911

912 # Test with prompt embeds for the second block

913 extra_keys, _ = generate_block_hash_extra_keys(request, 5, 10, 0)

914 expected_embeds = prompt_embeds[5:10]

915 expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()

915 - assert extra_keys == (expected_hash,)

916 + assert extra_keys == (("prompt_embeds", expected_hash),)

917

918

919 def test_generate_block_hash_extra_keys_prompt_embeds_cached(monkeypatch):

@@ -976,13 +977,58 @@ def test_generate_block_hash_extra_keys_lora():

977 )

978

979 extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0)

979 - assert extra_keys == ("test_lora_adapter",)

980 + assert extra_keys == (("lora", "test_lora_adapter"),)

981

982 request.lora_request = None

983 extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0)

984 assert extra_keys is None

985

986

987 +@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])

988 +def test_lora_name_and_cache_salt_block_hashes_do_not_collide(hash_fn):

989 + """A cache_salt equal to a LoRA name must not share that LoRA's blocks."""

990 + lora_request = LoRARequest(

991 + lora_name="foo", lora_int_id=1, lora_path="/path/to/lora"

992 + )

993 + lora_req = Request(

994 + request_id="lora",

995 + prompt_token_ids=[0, 1, 2],

996 + sampling_params=SamplingParams(max_tokens=1),

997 + pooling_params=None,

998 + lora_request=lora_request,

999 + block_hasher=get_request_block_hasher(3, hash_fn),

1000 + )

1001 + salted_req = make_request("salted", [0, 1, 2], hash_fn=hash_fn, cache_salt="foo")

1002 +

1003 + assert lora_req.block_hashes[0] != salted_req.block_hashes[0]

1004 +

1005 +

1006 +def test_to_event_extra_keys_keeps_untagged_event_format():

1007 + """KV events keep publishing the extra-key shapes consumers already parse."""

1008 + request = make_request(

1009 + "0",

1010 + list(range(10)),

1011 + mm_positions=[PlaceholderRange(offset=2, length=3)],

1012 + mm_hashes=["hash1"],

1013 + cache_salt="salt",

1014 + )

1015 + request.lora_request = LoRARequest(

1016 + lora_name="adapter", lora_int_id=1, lora_path="/path/to/lora"

1017 + )

1018 +

1019 + extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0)

1020 +

1021 + assert extra_keys == (

1022 + ("lora", "adapter"),

1023 + ("mm", "hash1", 2),

1024 + ("cache_salt", "salt"),

1025 + )

1026 + event_keys = ("adapter", ("hash1", 2), "salt")

1027 + assert to_event_extra_keys([extra_keys, None]) == [event_keys, None]

1028 + assert to_event_extra_keys([]) is None

1029 + assert to_event_extra_keys(None) is None

1030 +

1031 +

1032 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])

1033 def test_hash_block_tokens(hash_fn):

1034 parent_block_hash = BlockHash(b"123")

@@ -1013,9 +1059,11 @@ def test_request_block_hasher(hash_fn):

1059 block_hashes = request.block_hashes

1060 assert len(block_hashes) == 2

1061 assert block_hashes[0] == hash_fn(

1016 - (kv_cache_utils.NONE_HASH, (0, 1, 2), (("hash1", 0),))

1062 + (kv_cache_utils.NONE_HASH, (0, 1, 2), (("mm", "hash1", 0),))

1063 + )

1064 + assert block_hashes[1] == hash_fn(

1065 + (block_hashes[0], (3, 4, 5), (("mm", "hash2", 0),))

1066 )

1018 - assert block_hashes[1] == hash_fn((block_hashes[0], (3, 4, 5), (("hash2", 0),)))

1067

1068

1069 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])

@@ -1043,7 +1091,7 @@ def test_request_block_hasher_incremental_append_with_multiple_mm_features(hash_

1091 )

1092

1093 expected_second_hash = hash_fn(

1046 - (incremental.block_hashes[0], (4, 5, 6, 7), (("A", 0), ("B", 2)))

1094 + (incremental.block_hashes[0], (4, 5, 6, 7), (("mm", "A", 0), ("mm", "B", 2)))

1095 )

1096 assert incremental.block_hashes[1] == expected_second_hash

1097 assert incremental.block_hashes == fresh.block_hashes

@@ -3657,7 +3705,7 @@ def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]

3705 (

3706 kv_cache_utils.NONE_HASH,

3707 tuple(prompt_token_ids[:block_size]),

3660 - (block1_embeds_hash,),

3708 + (("prompt_embeds", block1_embeds_hash),),

3709 )

3710 )

3711 assert block_hashes[0] == expected_hash1

@@ -3669,7 +3717,7 @@ def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]

3717 (

3718 block_hashes[0],

3719 tuple(prompt_token_ids[block_size:num_tokens]),

3672 - (block2_embeds_hash,),

3720 + (("prompt_embeds", block2_embeds_hash),),

3721 )

3722 )

3723 assert block_hashes[1] == expected_hash2

@@ -3706,7 +3754,7 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes

3754 (

3755 kv_cache_utils.NONE_HASH,

3756 tuple(prompt_token_ids[:block_size]),

3709 - (("hash1", 0), block1_embeds_hash),

3757 + (("mm", "hash1", 0), ("prompt_embeds", block1_embeds_hash)),

3758 )

3759 )

3760 assert block_hashes[0] == expected_hash1

@@ -3718,7 +3766,7 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes

3766 (

3767 block_hashes[0],

3768 tuple(prompt_token_ids[block_size:num_tokens]),

3721 - (("hash2", 0), block2_embeds_hash),

3769 + (("mm", "hash2", 0), ("prompt_embeds", block2_embeds_hash)),

3770 )

3771 )

3772 assert block_hashes[1] == expected_hash2

tests/v1/core/test_prefix_caching.py

+14 -6

@@ -2901,21 +2901,21 @@ def test_mm_prefix_caching():

2901 (

2902 kv_cache_utils.NONE_HASH,

2903 tuple(all_token_ids[:block_size]),

2904 - (("aaa", 11),),

2904 + (("mm", "aaa", 11),),

2905 )

2906 )

2907 assert block_hashes[1] == sha256(

2908 (

2909 block_hashes[0],

2910 tuple(all_token_ids[block_size : block_size * 2]),

2911 - (("aaa", -5), ("bbb", 14)),

2911 + (("mm", "aaa", -5), ("mm", "bbb", 14)),

2912 )

2913 )

2914 assert block_hashes[2] == sha256(

2915 (

2916 block_hashes[1],

2917 tuple(all_token_ids[block_size * 2 : block_size * 3]),

2918 - (("bbb", -2),),

2918 + (("mm", "bbb", -2),),

2919 )

2920 )

2921

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

2938 (

2939 block_hashes[2],

2940 tuple(all_token_ids[3 * block_size :] + [8] * 5),

2941 - (("ccc", 0),),

2941 + (("mm", "ccc", 0),),

2942 )

2943 )

2944

@@ -2984,7 +2984,11 @@ def test_cache_key_salting():

2984 block_hashes = req0.block_hashes

2985 assert len(block_hashes) == 3

2986 assert block_hashes[0] == sha256(

2987 - (kv_cache_utils.NONE_HASH, tuple(token_ids[:block_size]), ("salt1",))

2987 + (

2988 + kv_cache_utils.NONE_HASH,

2989 + tuple(token_ids[:block_size]),

2990 + (("cache_salt", "salt1"),),

2991 + )

2992 )

2993 assert block_hashes[1] == sha256(

2994 (block_hashes[0], tuple(token_ids[block_size : block_size * 2]), None)

@@ -3029,7 +3033,11 @@ def test_cache_key_salting():

3033 block_hashes = req2.block_hashes

3034 assert len(block_hashes) == 3

3035 assert block_hashes[0] == sha256(

3032 - (kv_cache_utils.NONE_HASH, tuple(token_ids[:block_size]), ("salt2",))

3036 + (

3037 + kv_cache_utils.NONE_HASH,

3038 + tuple(token_ids[:block_size]),

3039 + (("cache_salt", "salt2"),),

3040 + )

3041 )

3042 assert block_hashes[1] == sha256(

3043 (block_hashes[0], tuple(token_ids[block_size : block_size * 2]), None)

vllm/distributed/kv_events.py

+5 -3

@@ -68,9 +68,11 @@ class BlockStored(KVCacheEvent):

68 lora_name: str | None

70 """Extra keys used in block hash computation, one entry per block in

71 - block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,

72 - prompt embedding hashes, etc. for that specific block. Exposed for external

73 - KV cache consumers to reconstruct block hashes.

71 + block_hashes. Each entry contains, in order, the LoRA name (if

72 + `lora_name` is set), `(mm_identifier, offset)` pairs, the cache_salt

73 + string and prompt embedding digests for that specific block. These are

74 + published untagged; the block hash itself uses each value tagged with

75 + its source, e.g. `("lora", name)` or `("mm", mm_identifier, offset)`.

76 """

77 extra_keys: list[tuple[Any, ...] | None] | None = None

vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py

+2 -3

@@ -31,6 +31,7 @@ from vllm.v1.core.kv_cache_utils import (

31 BlockHash,

32 maybe_convert_block_hash,

33 resolve_block_hashes,

34 + to_event_extra_keys,

35 )

36 from vllm.v1.kv_cache_interface import (

37 KVCacheGroupSpec,

@@ -393,9 +394,7 @@ class OffloadingEventsTracker:

394 lora_id=meta.lora_id,

395 medium=_MEDIUM_TO_EVENT_STR[event.medium],

396 lora_name=meta.lora_name,

396 - extra_keys=(

397 - list(meta.extra_keys) if meta.extra_keys is not None else None

398 - ),

397 + extra_keys=to_event_extra_keys(meta.extra_keys),

398 group_idx=meta.group_idx,

399 kv_cache_spec_kind=meta.kv_cache_spec.kv_cache_spec_kind,

400 kv_cache_spec_sliding_window=(

vllm/v1/core/block_pool.py

+3 -2

@@ -25,6 +25,7 @@ from vllm.v1.core.kv_cache_utils import (

25 make_block_hash_with_group_id,

26 maybe_convert_block_hash,

27 resolve_block_hashes,

28 + to_event_extra_keys,

29 )

30 from vllm.v1.request import Request

@@ -367,7 +368,7 @@ class BlockPool:

368 lora_id=request.lora_request.adapter_id if request.lora_request else None,

369 medium=self.medium,

370 lora_name=request.lora_request.name if request.lora_request else None,

370 - extra_keys=extra_keys_list if extra_keys_list else None,

371 + extra_keys=to_event_extra_keys(extra_keys_list),

372 group_idx=kv_cache_group_id,

373 session_id=request.session_id,

374 )

@@ -557,7 +558,7 @@ class BlockPool:

558 lora_name=request.lora_request.name

559 if request.lora_request

560 else None,

560 - extra_keys=[extra_keys],

561 + extra_keys=to_event_extra_keys([extra_keys]),

562 group_idx=kv_cache_group_id,

563 session_id=request.session_id,

564 )

vllm/v1/core/kv_cache_utils.py

+41 -10

@@ -502,8 +502,8 @@ def _gen_mm_extra_hash_keys(

502 ) -> tuple[list[Any], int]:

503 """Generate extra keys related to MultiModal request for block hash

504 computation. For multi-modal inputs, the extra keys are

505 - (mm_hash, start_offset) that indicate a mm input contained in the

506 - block and its starting offset in the block tokens.

505 + ("mm", mm_hash, start_offset) tuples that indicate a mm input contained

506 + in the block and its starting offset in the block tokens.

507

508 Args:

509 request: The request object.

@@ -549,7 +549,7 @@ def _gen_mm_extra_hash_keys(

549 # relative to the start of the block so prefix-cache keys stay

550 # distinct when the same MM item appears at different positions

551 # within otherwise-identical placeholder blocks.

552 - extra_keys.append((mm_feature.identifier, offset - start_token_idx))

552 + extra_keys.append(("mm", mm_feature.identifier, offset - start_token_idx))

553

554 if end_token_idx >= offset + length:

555 # If this block contains the end of the current mm input,

@@ -565,7 +565,7 @@ def _gen_mm_extra_hash_keys(

565 return extra_keys, curr_mm_idx

566

567

568 -def _gen_lora_extra_hash_keys(request: Request) -> list[str]:

568 +def _gen_lora_extra_hash_keys(request: Request) -> list[tuple[str, str]]:

569 """Generate extra keys related to LoRA for block hash computation.

570

571 Args:

@@ -578,12 +578,12 @@ def _gen_lora_extra_hash_keys(request: Request) -> list[str]:

578 """

579 if not request.lora_request:

580 return []

581 - return [request.lora_request.lora_name]

581 + return [("lora", request.lora_request.lora_name)]

582

583

584 def _gen_prompt_embeds_extra_hash_keys(

585 request: Request, start_token_idx: int, end_token_idx: int

586 -) -> list[bytes]:

586 +) -> list[tuple[str, bytes]]:

587 """Generate extra keys related to prompt embeds for block hash computation.

588

589 Args:

@@ -605,7 +605,7 @@ def _gen_prompt_embeds_extra_hash_keys(

605 # Hash prompt embeds once per block and cache on request

606 embeds_hash = hashlib.sha256(tensor_data(block_prompt_embeds)).digest()

607 request._prompt_embeds_per_block_hashes[block_range] = embeds_hash

608 - return [embeds_hash]

608 + return [("prompt_embeds", embeds_hash)]

609

610

611 def generate_block_hash_extra_keys(

@@ -629,9 +629,11 @@ def generate_block_hash_extra_keys(

629 mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(

630 request, start_token_idx, end_token_idx, start_mm_idx

631 )

632 - lora_extra_keys: list[str] = _gen_lora_extra_hash_keys(request)

633 - cache_salt_keys: list[str] = (

634 - [request.cache_salt] if (start_token_idx == 0 and request.cache_salt) else []

632 + lora_extra_keys = _gen_lora_extra_hash_keys(request)

633 + cache_salt_keys: list[tuple[str, str]] = (

634 + [("cache_salt", request.cache_salt)]

635 + if (start_token_idx == 0 and request.cache_salt)

636 + else []

637 )

638 prompt_embeds_keys = _gen_prompt_embeds_extra_hash_keys(

639 request, start_token_idx, end_token_idx

@@ -647,6 +649,35 @@ def generate_block_hash_extra_keys(

649 return tuple(extra_keys), new_start_mm_idx

650

651

652 +def to_event_extra_keys(

653 + extra_keys: Iterable[tuple[Any, ...] | None] | None,

654 +) -> list[tuple[Any, ...] | None] | None:

655 + """Convert block-hash extra keys to the untagged per-block list published

656 + in KV events.

657 +

658 + External KV event consumers parse the pre-tagging shapes: bare LoRA names

659 + and cache salts, ``(mm_identifier, offset)`` pairs and bare prompt-embeds

660 + digests. Events keep that format until consumers handle the tagged keys.

661 +

662 + Args:

663 + extra_keys: One entry per block, each as returned by

664 + `generate_block_hash_extra_keys`, or None.

665 +

666 + Returns:

667 + One untagged entry per block, or None if there are no entries.

668 +

669 + """

670 + if not extra_keys:

671 + return None

672 + event_keys = [

673 + None

674 + if keys is None

675 + else tuple(key[1:] if key[0] == "mm" else key[1] for key in keys)

676 + for keys in extra_keys

677 + ]

678 + return event_keys or None

679 +

680 +

681 def hash_block_tokens(

682 hash_function: Callable[[Any], bytes],

683 parent_block_hash: BlockHash | None,

vllm/v1/simple_kv_offload/manager.py

+2 -3

@@ -34,6 +34,7 @@ from vllm.v1.core.kv_cache_utils import (

34 maybe_convert_block_hash,

35 resolve_block_hashes,

36 resolve_dcp_kv_block_size,

37 + to_event_extra_keys,

38 )

39 from vllm.v1.core.sched.output import SchedulerOutput

40 from vllm.v1.kv_cache_interface import (

@@ -1147,9 +1148,7 @@ class SimpleCPUOffloadScheduler:

1148 lora_id=meta.lora_id if meta else None,

1149 medium=self.kv_event_medium,

1150 lora_name=meta.lora_name if meta else None,

1150 - extra_keys=(

1151 - [extra_keys] if extra_keys is not None else None

1152 - ),

1151 + extra_keys=to_event_extra_keys(extra_keys and [extra_keys]),

1152 group_idx=group_idx,

1153 kv_cache_spec_kind=get_kv_cache_spec_kind(spec).value,

1154 kv_cache_spec_sliding_window=(