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 8, 00:00 UTC
[feat] Add per-block extra_keys to KV events (#33304)
Signed-off-by: zhongdaor-nv <zhongdaor@nvidia.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
zhongdaor-nv committed
Feb 20, 2026 at 21:11 UTC
a0fe7ea2f052bb44820bc06a5635456b8d1383af
6 files changed
+100
-21
examples/online_serving/kv_events_subscriber.py
+6
@@ -37,6 +37,12 @@ class BlockStored(KVCacheEvent):
37
medium: str | None
38
lora_name: str | None
39
40
+ extra_keys: list[tuple[Any, ...] | None] | None = None
41
+ """Extra keys used in block hash computation, one entry per block in
42
+ block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
43
+ prompt embeddings data, etc. for that specific block.
44
+ """
45
+
46
47
class BlockRemoved(KVCacheEvent):
48
block_hashes: list[ExternalBlockHash]
tests/v1/core/test_kv_cache_utils.py
+48
-12
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import hashlib
4
import importlib
5
from collections.abc import Callable
6
from typing import Any
@@ -498,14 +499,41 @@ def test_generate_block_hash_extra_keys_prompt_embeds():
499
# Test with prompt embeds for the first block
500
extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
501
expected_embeds = prompt_embeds[0:5]
501
- expected_bytes = kv_cache_utils.tensor_data(expected_embeds).tobytes()
502
- assert extra_keys == (expected_bytes,)
502
+ expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()
503
+ assert extra_keys == (expected_hash,)
504
505
# Test with prompt embeds for the second block
506
extra_keys, _ = generate_block_hash_extra_keys(request, 5, 10, 0)
507
expected_embeds = prompt_embeds[5:10]
507
- expected_bytes = kv_cache_utils.tensor_data(expected_embeds).tobytes()
508
- assert extra_keys == (expected_bytes,)
508
+ expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()
509
+ assert extra_keys == (expected_hash,)
510
+
511
+
512
+def test_generate_block_hash_extra_keys_prompt_embeds_cached(monkeypatch):
513
+ prompt_embeds = torch.randn(10, 3)
514
+ request = make_request(
515
+ request_id="0",
516
+ prompt_token_ids=None,
517
+ mm_positions=None,
518
+ mm_hashes=None,
519
+ prompt_embeds=prompt_embeds,
520
+ block_size=20,
521
+ )
522
+
523
+ num_tensor_data_calls = 0
524
+ original_tensor_data = kv_cache_utils.tensor_data
525
+
526
+ def counting_tensor_data(tensor: torch.Tensor):
527
+ nonlocal num_tensor_data_calls
528
+ num_tensor_data_calls += 1
529
+ return original_tensor_data(tensor)
530
+
531
+ monkeypatch.setattr(kv_cache_utils, "tensor_data", counting_tensor_data)
532
+
533
+ extra_keys_1, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
534
+ extra_keys_2, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
535
+ assert extra_keys_1 == extra_keys_2
536
+ assert num_tensor_data_calls == 1
537
538
539
def test_generate_block_hash_extra_keys_different_prompt_embeds():
@@ -1858,22 +1886,26 @@ def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]
1886
block_hashes = request.block_hashes
1887
assert len(block_hashes) == 2
1888
1861
- block1_embeds_bytes = tensor_data(prompt_embeds[:block_size]).tobytes()
1889
+ block1_embeds_hash = hashlib.sha256(
1890
+ tensor_data(prompt_embeds[:block_size])
1891
+ ).digest()
1892
expected_hash1 = hash_fn(
1893
(
1894
kv_cache_utils.NONE_HASH,
1895
tuple(prompt_token_ids[:block_size]),
1866
- (block1_embeds_bytes,),
1896
+ (block1_embeds_hash,),
1897
)
1898
)
1899
assert block_hashes[0] == expected_hash1
1900
1871
- block2_embeds_bytes = tensor_data(prompt_embeds[block_size:num_tokens]).tobytes()
1901
+ block2_embeds_hash = hashlib.sha256(
1902
+ tensor_data(prompt_embeds[block_size:num_tokens])
1903
+ ).digest()
1904
expected_hash2 = hash_fn(
1905
(
1906
block_hashes[0],
1907
tuple(prompt_token_ids[block_size:num_tokens]),
1876
- (block2_embeds_bytes,),
1908
+ (block2_embeds_hash,),
1909
)
1910
)
1911
assert block_hashes[1] == expected_hash2
@@ -1903,22 +1935,26 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes
1935
block_hashes = request.block_hashes
1936
assert len(block_hashes) == 2
1937
1906
- block1_embeds_bytes = tensor_data(prompt_embeds[:block_size]).tobytes()
1938
+ block1_embeds_hash = hashlib.sha256(
1939
+ tensor_data(prompt_embeds[:block_size])
1940
+ ).digest()
1941
expected_hash1 = hash_fn(
1942
(
1943
kv_cache_utils.NONE_HASH,
1944
tuple(prompt_token_ids[:block_size]),
1911
- ("hash1", block1_embeds_bytes),
1945
+ ("hash1", block1_embeds_hash),
1946
)
1947
)
1948
assert block_hashes[0] == expected_hash1
1949
1916
- block2_embeds_bytes = tensor_data(prompt_embeds[block_size:num_tokens]).tobytes()
1950
+ block2_embeds_hash = hashlib.sha256(
1951
+ tensor_data(prompt_embeds[block_size:num_tokens])
1952
+ ).digest()
1953
expected_hash2 = hash_fn(
1954
(
1955
block_hashes[0],
1956
tuple(prompt_token_ids[block_size:num_tokens]),
1921
- ("hash2", block2_embeds_bytes),
1957
+ ("hash2", block2_embeds_hash),
1958
)
1959
)
1960
assert block_hashes[1] == expected_hash2
vllm/distributed/kv_events.py
+8
@@ -60,6 +60,13 @@ class BlockStored(KVCacheEvent):
60
medium: str | None
61
lora_name: str | None
62
63
+ extra_keys: list[tuple[Any, ...] | None] | None = None
64
+ """Extra keys used in block hash computation, one entry per block in
65
+ block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
66
+ prompt embedding hashes, etc. for that specific block. Exposed for external
67
+ KV cache consumers to reconstruct block hashes.
68
+ """
69
+
70
def __hash__(self) -> int:
71
return hash(
72
(
@@ -69,6 +76,7 @@ class BlockStored(KVCacheEvent):
76
self.block_size,
77
self.lora_id,
78
self.medium,
79
+ tuple(self.extra_keys) if self.extra_keys else None,
80
)
81
)
82
vllm/v1/core/block_pool.py
+23
-3
@@ -20,6 +20,7 @@ from vllm.v1.core.kv_cache_utils import (
20
ExternalBlockHash,
21
FreeKVCacheBlockQueue,
22
KVCacheBlock,
23
+ generate_block_hash_extra_keys,
24
get_block_hash,
25
make_block_hash_with_group_id,
26
maybe_convert_block_hash,
@@ -279,13 +280,31 @@ class BlockPool:
280
block_hashes[num_cached_blocks - 1]
281
)
282
283
+ # Calculate token range for the blocks being cached
284
+ start_token_idx = num_cached_blocks * block_size
285
+ end_token_idx = num_full_blocks * block_size
286
+
287
+ # Generate extra keys for each block individually.
288
+ # Each block may have different extra_keys (e.g., different MM
289
+ # features, or cache_salt only for the first block).
290
+ # Skip null blocks to match the length of new_hashes.
291
+ extra_keys_list: list[tuple[Any, ...] | None] = []
292
+ curr_mm_idx = 0
293
+ for i in range(num_cached_blocks, num_full_blocks):
294
+ if blocks[i].is_null:
295
+ continue
296
+ block_start = i * block_size
297
+ block_end = block_start + block_size
298
+ extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
299
+ request, block_start, block_end, curr_mm_idx
300
+ )
301
+ extra_keys_list.append(extra_keys)
302
+
303
self.kv_event_queue.append(
304
BlockStored(
305
block_hashes=new_hashes,
306
parent_block_hash=parent_block_hash,
286
- token_ids=request.all_token_ids[
287
- num_cached_blocks * block_size : num_full_blocks * block_size
288
- ],
307
+ token_ids=request.all_token_ids[start_token_idx:end_token_idx],
308
block_size=block_size,
309
lora_id=request.lora_request.adapter_id
310
if request.lora_request
@@ -294,6 +313,7 @@ class BlockPool:
313
lora_name=request.lora_request.name
314
if request.lora_request
315
else None,
316
+ extra_keys=extra_keys_list if extra_keys_list else None,
317
)
318
)
319
vllm/v1/core/kv_cache_utils.py
+12
-6
@@ -3,6 +3,7 @@
3
"""KV-Cache Utilities."""
4
5
import copy
6
+import hashlib
7
import os
8
from collections import defaultdict
9
from collections.abc import Callable, Iterable, Iterator, Sequence
@@ -475,14 +476,19 @@ def _gen_prompt_embeds_extra_hash_keys(
476
end_token_idx: The end token index of the block.
477
478
Returns:
478
- Return prompt embeddings data of the request if it has prompt embeds.
479
- Return empty list otherwise.
479
+ Return a stable hash of the block prompt embeddings if prompt embeds
480
+ are present. Return empty list otherwise.
481
"""
482
if request.prompt_embeds is None:
483
return []
483
- block_prompt_embeds = request.prompt_embeds[start_token_idx:end_token_idx]
484
- embeds_bytes = tensor_data(block_prompt_embeds).tobytes()
485
- return [embeds_bytes]
484
+ block_range = (start_token_idx, end_token_idx)
485
+ embeds_hash = request._prompt_embeds_per_block_hashes.get(block_range)
486
+ if embeds_hash is None:
487
+ block_prompt_embeds = request.prompt_embeds[start_token_idx:end_token_idx]
488
+ # Hash prompt embeds once per block and cache on request
489
+ embeds_hash = hashlib.sha256(tensor_data(block_prompt_embeds)).digest()
490
+ request._prompt_embeds_per_block_hashes[block_range] = embeds_hash
491
+ return [embeds_hash]
492
493
494
def generate_block_hash_extra_keys(
@@ -490,7 +496,7 @@ def generate_block_hash_extra_keys(
496
) -> tuple[tuple[Any, ...] | None, int]:
497
"""Generate extra keys for the block hash. The extra keys can come from
498
the multi-modal inputs, request specific metadata (e.g., LoRA names), and
493
- data from prompt embeddings.
499
+ hashed data from prompt embeddings.
500
501
Args:
502
request: The request object.
vllm/v1/request.py
+3
@@ -114,6 +114,9 @@ class Request:
114
115
self.prompt_token_ids = prompt_token_ids
116
self.prompt_embeds = prompt_embeds
117
+ # Cache per-block prompt-embed hashes to avoid rehashing the same
118
+ # tensor slices when generating extra keys.
119
+ self._prompt_embeds_per_block_hashes: dict[tuple[int, int], bytes] = {}
120
self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
121
prompt_token_ids, prompt_embeds
122
)