@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cc07dad789
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 06:00 UTC

[HMA] [KVEvent] Enable GPU-side KV events for HMA (#37688)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Or Ozeri <or@ozery.com>

Martin Hickey committed Apr 12, 2026 at 08:01 UTC cc07dad789e9ae3f02ab6a4794800b8c1cc240c6
8 files changed +300 -28
examples/online_serving/kv_events_subscriber.py
+3
@@ -43,10 +43,13 @@ class BlockStored(KVCacheEvent):
43 prompt embeddings data, etc. for that specific block.
44 """
45
46 + group_idx: int | None = None
47 +
48
49 class BlockRemoved(KVCacheEvent):
50 block_hashes: list[ExternalBlockHash]
51 medium: str | None
52 + group_idx: int | None = None
53
54
55 class AllBlocksCleared(KVCacheEvent):
tests/distributed/test_kv_cache_events.py new
+74
@@ -0,0 +1,74 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.distributed.kv_events import BlockRemoved, BlockStored
7 +
8 +# Minimal ExternalBlockHash for testing (bytes are a valid ExternalBlockHash).
9 +_FAKE_HASH: bytes = b"\xab" * 32
10 +
11 +
12 +def _make_block_stored(group_idx: int | None = None) -> BlockStored:
13 + return BlockStored(
14 + block_hashes=[_FAKE_HASH],
15 + parent_block_hash=None,
16 + token_ids=[1, 2, 3, 4],
17 + block_size=4,
18 + lora_id=None,
19 + medium="GPU",
20 + lora_name=None,
21 + group_idx=group_idx,
22 + )
23 +
24 +
25 +def _make_block_removed(group_idx: int | None = None) -> BlockRemoved:
26 + return BlockRemoved(
27 + block_hashes=[_FAKE_HASH],
28 + medium="GPU",
29 + group_idx=group_idx,
30 + )
31 +
32 +
33 +def test_block_stored_default_group_idx_is_none():
34 + """group_idx defaults to None when not provided."""
35 + event = _make_block_stored()
36 + assert event.group_idx is None
37 +
38 +
39 +def test_block_removed_default_group_idx_is_none():
40 + """group_idx defaults to None when not provided."""
41 + event = _make_block_removed()
42 + assert event.group_idx is None
43 +
44 +
45 +@pytest.mark.parametrize("group_idx", [1, 2, 3])
46 +def test_block_stored_hash_differs_by_group_idx(group_idx: int):
47 + """BlockStored events that differ only in group_idx must hash differently."""
48 + other_group_idx = group_idx + 1
49 + event_a = _make_block_stored(group_idx=group_idx)
50 + event_b = _make_block_stored(group_idx=other_group_idx)
51 + assert hash(event_a) != hash(event_b)
52 +
53 +
54 +def test_block_stored_hash_same_for_equal_group_idx():
55 + """Two BlockStored events with identical fields produce the same hash."""
56 + event_a = _make_block_stored(group_idx=1)
57 + event_b = _make_block_stored(group_idx=1)
58 + assert hash(event_a) == hash(event_b)
59 +
60 +
61 +@pytest.mark.parametrize("group_idx", [1, 2, 3])
62 +def test_block_removed_hash_differs_by_group_idx(group_idx: int):
63 + """BlockRemoved events that differ only in group_idx must hash differently."""
64 + other_group_idx = group_idx + 1
65 + event_a = _make_block_removed(group_idx=group_idx)
66 + event_b = _make_block_removed(group_idx=other_group_idx)
67 + assert hash(event_a) != hash(event_b)
68 +
69 +
70 +def test_block_removed_hash_same_for_equal_group_idx():
71 + """Two BlockRemoved events with identical fields produce the same hash."""
72 + event_a = _make_block_removed(group_idx=1)
73 + event_b = _make_block_removed(group_idx=1)
74 + assert hash(event_a) == hash(event_b)
tests/v1/core/test_kv_cache_utils.py
+28
@@ -10,6 +10,7 @@ import torch
10
11 import vllm.v1.core.kv_cache_utils as kv_cache_utils
12 from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
13 +from vllm.config.kv_events import KVEventsConfig
14 from vllm.lora.request import LoRARequest
15 from vllm.multimodal.inputs import (
16 MultiModalFeatureSpec,
@@ -2137,3 +2138,30 @@ def test_unify_hybrid_kv_cache_specs():
2138
2139 with pytest.raises(ValueError):
2140 kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
2141 +
2142 +
2143 +def test_hma_not_disabled_when_kv_events_enabled():
2144 + """
2145 + Test enabling KV events must not force disable_hybrid_kv_cache_manager to True.
2146 +
2147 + This test guards against that regression by verifying that a VllmConfig
2148 + with kv_events_config set still resolves disable_hybrid_kv_cache_manager
2149 + to False (i.e. HMA remains enabled) when no other condition requires it
2150 + to be disabled.
2151 + """
2152 + model_config = ModelConfig(max_model_len=16)
2153 + kv_events_config = KVEventsConfig(
2154 + enable_kv_cache_events=True,
2155 + publisher="null",
2156 + )
2157 +
2158 + # Leave disable_hybrid_kv_cache_manager as None (the default) so that
2159 + # VllmConfig.__post_init__ resolves it automatically.
2160 + vllm_config = VllmConfig(
2161 + model_config=model_config,
2162 + kv_events_config=kv_events_config,
2163 + )
2164 +
2165 + assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, (
2166 + "kv_events_config must not force-disable the hybrid KV cache manager."
2167 + )
tests/v1/core/test_prefix_caching.py
+150 -1
@@ -1970,6 +1970,7 @@ def test_null_parent_block_hash():
1970 block_size = 1
1971 num_cached_blocks = 2
1972 num_full_blocks = 4
1973 + kv_cache_group_id = 0
1974
1975 pool = BlockPool(
1976 num_gpu_blocks=8,
@@ -2002,7 +2003,7 @@ def test_null_parent_block_hash():
2003 num_cached_blocks=num_cached_blocks,
2004 num_full_blocks=num_full_blocks,
2005 block_size=block_size,
2005 - kv_cache_group_id=0,
2006 + kv_cache_group_id=kv_cache_group_id,
2007 )
2008
2009 events = pool.take_events()
@@ -2021,6 +2022,7 @@ def test_null_parent_block_hash():
2022 for h in req.block_hashes[num_cached_blocks:num_full_blocks]
2023 ]
2024 assert event.block_hashes == expected_new_hashes
2025 + assert event.group_idx == kv_cache_group_id
2026
2027 # Ensure we didn't accidentally assign a hash to the null block.
2028 assert pool.null_block.block_hash is None
@@ -2087,6 +2089,153 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int):
2089 assert block_stored_event.block_size == block_size
2090
2091
2092 +@pytest.mark.parametrize("group_id", [0, 1, 2])
2093 +def test_block_stored_event_group_idx(group_id: int):
2094 + """Test BlockStored events emitted by cache_full_blocks carry the correct
2095 + group_idx."""
2096 + block_size = 4
2097 + num_tokens = block_size * 2
2098 +
2099 + pool = BlockPool(
2100 + num_gpu_blocks=5,
2101 + enable_caching=True,
2102 + hash_block_size=block_size,
2103 + enable_kv_cache_events=True,
2104 + )
2105 +
2106 + req = make_request(
2107 + "req_grp_idx",
2108 + prompt_token_ids=list(range(num_tokens)),
2109 + block_size=block_size,
2110 + hash_fn=sha256,
2111 + )
2112 +
2113 + blocks = pool.get_new_blocks(2)
2114 + pool.cache_full_blocks(
2115 + request=req,
2116 + blocks=blocks,
2117 + num_cached_blocks=0,
2118 + num_full_blocks=2,
2119 + block_size=block_size,
2120 + kv_cache_group_id=group_id,
2121 + )
2122 +
2123 + events = pool.take_events()
2124 + assert len(events) == 1
2125 + assert isinstance(events[0], BlockStored)
2126 + assert events[0].group_idx == group_id
2127 +
2128 +
2129 +def test_block_stored_event_group_idx_multiple_groups():
2130 + """
2131 + Test BlockStored events for separate HMA groups that each carry the
2132 + correct group_idx.
2133 +
2134 + Simulates the HMA scenario where full-attention blocks (group 0) and
2135 + sliding-window blocks (group 1) are cached independently and must be
2136 + distinguishable by consumers doing HMA-aware prefix-cache routing.
2137 + """
2138 + block_size = 4
2139 + num_tokens = block_size * 2
2140 +
2141 + # null block + 4 usable (2 per group)
2142 + pool = BlockPool(
2143 + num_gpu_blocks=5,
2144 + enable_caching=True,
2145 + hash_block_size=block_size,
2146 + enable_kv_cache_events=True,
2147 + )
2148 +
2149 + req = make_request(
2150 + "req_multi_grp",
2151 + prompt_token_ids=list(range(num_tokens)),
2152 + block_size=block_size,
2153 + hash_fn=sha256,
2154 + )
2155 +
2156 + # Cache blocks for group 0 (full-attention)
2157 + blocks_grp0 = pool.get_new_blocks(2)
2158 + pool.cache_full_blocks(
2159 + request=req,
2160 + blocks=blocks_grp0,
2161 + num_cached_blocks=0,
2162 + num_full_blocks=2,
2163 + block_size=block_size,
2164 + kv_cache_group_id=0,
2165 + )
2166 +
2167 + # Cache blocks for group 1 (sliding-window)
2168 + blocks_grp1 = pool.get_new_blocks(2)
2169 + pool.cache_full_blocks(
2170 + request=req,
2171 + blocks=blocks_grp1,
2172 + num_cached_blocks=0,
2173 + num_full_blocks=2,
2174 + block_size=block_size,
2175 + kv_cache_group_id=1,
2176 + )
2177 +
2178 + events = pool.take_events()
2179 + assert len(events) == 2
2180 + assert isinstance(events[0], BlockStored)
2181 + assert events[0].group_idx == 0
2182 + assert isinstance(events[1], BlockStored)
2183 + assert events[1].group_idx == 1
2184 +
2185 +
2186 +@pytest.mark.parametrize("group_id", [0, 1, 2])
2187 +def test_block_removed_event_group_idx(group_id: int):
2188 + """
2189 + Test BlockRemoved events emitted on eviction carry the group_idx extracted
2190 + from the evicted block's BlockHashWithGroupId via get_group_id().
2191 + """
2192 + block_size = 4
2193 + num_tokens = block_size * 2
2194 +
2195 + # null block + 4 usable; allocate all 4, cache 2, free all, re-allocate
2196 + # all 4 so the 2 cached blocks are forced through _maybe_evict_cached_block.
2197 + pool = BlockPool(
2198 + num_gpu_blocks=5,
2199 + enable_caching=True,
2200 + hash_block_size=block_size,
2201 + enable_kv_cache_events=True,
2202 + )
2203 +
2204 + req = make_request(
2205 + "req_evict_grp",
2206 + prompt_token_ids=list(range(num_tokens)),
2207 + block_size=block_size,
2208 + hash_fn=sha256,
2209 + )
2210 +
2211 + # Allocate all usable blocks and cache the first two for the target group.
2212 + all_blocks = pool.get_new_blocks(4)
2213 + pool.cache_full_blocks(
2214 + request=req,
2215 + blocks=all_blocks,
2216 + num_cached_blocks=0,
2217 + num_full_blocks=2,
2218 + block_size=block_size,
2219 + kv_cache_group_id=group_id,
2220 + )
2221 +
2222 + # Drain the BlockStored events so only eviction events remain later.
2223 + pool.take_events()
2224 +
2225 + # Return all blocks to the free queue so they become eviction candidates.
2226 + pool.free_blocks(all_blocks)
2227 +
2228 + # Re-allocate all blocks; the two with hashes trigger BlockRemoved events.
2229 + pool.get_new_blocks(4)
2230 +
2231 + events = pool.take_events()
2232 + removed_events = [e for e in events if isinstance(e, BlockRemoved)]
2233 +
2234 + assert len(removed_events) == 2
2235 + for event in removed_events:
2236 + assert event.group_idx == group_id
2237 +
2238 +
2239 def test_eagle_enabled_removes_last_block():
2240 """Verify Eagle does NOT remove blocks when request
2241 length is divisible by block size."""
tests/v1/engine/test_engine_core_client.py
+31 -19
@@ -936,6 +936,13 @@ async def test_engine_core_client_future_utility_async(
936 client.shutdown()
937
938
939 +@pytest.mark.parametrize(
940 + "model_name,num_groups",
941 + [
942 + ("meta-llama/Llama-3.2-1B-Instruct", 1),
943 + ("google/gemma-3-1b-it", 7),
944 + ],
945 +)
946 @pytest.mark.parametrize(
947 "multiprocessing_mode,publisher_config",
948 [(True, "tcp"), (False, "inproc")],
@@ -944,12 +951,14 @@ async def test_engine_core_client_future_utility_async(
951 def test_kv_cache_events(
952 multiprocessing_mode: bool,
953 publisher_config,
954 + model_name: str,
955 + num_groups: int,
956 ):
957 block_size = 16
958 num_blocks = 2
959
960 engine_args = EngineArgs(
952 - model=MODEL_NAME,
961 + model=model_name,
962 enforce_eager=True,
963 enable_prefix_caching=True,
964 block_size=block_size,
@@ -985,26 +994,29 @@ def test_kv_cache_events(
994 assert result is not None, "No message received"
995
996 seq, received = result
988 -
997 assert seq == 0, "Sequence number mismatch"
990 - assert len(received.events) == 1, "We should have exactly one BlockStored event"
991 - event = received.events[0]
992 - assert isinstance(event, BlockStored), "We should have a BlockStored event"
993 - assert len(event.block_hashes) == num_blocks, (
994 - "We should have a BlockStored event with 2 block_hashes"
995 - )
996 - assert event.block_size == block_size, (
997 - "Block size should be the same as the block size"
998 - )
999 - assert event.parent_block_hash is None, "Parent block hash should be None"
1000 - assert event.lora_id is None, "Lora id should be None"
1001 - assert event.lora_name is None, "Lora name should be None"
1002 - assert len(event.token_ids) == num_blocks * block_size, (
1003 - "Token ids should be the same as the custom tokens"
1004 - )
1005 - assert event.token_ids == custom_tokens, (
1006 - "Token ids should be the same as the custom tokens"
998 + assert len(received.events) == num_groups, (
999 + f"Expected {num_groups} BlockStored event(s), got {len(received.events)}"
1000 )
1001 +
1002 + for index, event in enumerate(received.events):
1003 + assert isinstance(event, BlockStored), "We should have a BlockStored event"
1004 + assert len(event.block_hashes) == num_blocks, (
1005 + "We should have a BlockStored event with 2 block_hashes"
1006 + )
1007 + assert event.block_size == block_size, (
1008 + "Block size should be the same as the block size"
1009 + )
1010 + assert event.parent_block_hash is None, "Parent block hash should be None"
1011 + assert event.lora_id is None, "Lora id should be None"
1012 + assert event.lora_name is None, "Lora name should be None"
1013 + assert len(event.token_ids) == num_blocks * block_size, (
1014 + "Token ids should be the same as the custom tokens"
1015 + )
1016 + assert event.token_ids == custom_tokens, (
1017 + "Token ids should be the same as the custom tokens"
1018 + )
1019 + assert event.group_idx == index
1020 finally:
1021 client.shutdown()
1022 subscriber.close()
vllm/config/vllm.py
-3
@@ -1229,9 +1229,6 @@ class VllmConfig:
1229 if not current_platform.support_hybrid_kv_cache():
1230 # Hybrid KV cache manager is not supported on non-GPU platforms.
1231 need_disable_hybrid_kv_cache_manager = True
1232 - if self.kv_events_config is not None:
1233 - # Hybrid KV cache manager is not compatible with KV events.
1234 - need_disable_hybrid_kv_cache_manager = True
1232 if (
1233 self.model_config is not None
1234 and self.model_config.attention_chunk_size is not None
vllm/distributed/kv_events.py
+11 -1
@@ -67,6 +67,8 @@ class BlockStored(KVCacheEvent):
67 KV cache consumers to reconstruct block hashes.
68 """
69
70 + group_idx: int | None = None
71 +
72 def __hash__(self) -> int:
73 return hash(
74 (
@@ -77,6 +79,7 @@ class BlockStored(KVCacheEvent):
79 self.lora_id,
80 self.medium,
81 tuple(self.extra_keys) if self.extra_keys else None,
82 + self.group_idx,
83 )
84 )
85
@@ -84,9 +87,16 @@ class BlockStored(KVCacheEvent):
87 class BlockRemoved(KVCacheEvent):
88 block_hashes: list[ExternalBlockHash]
89 medium: str | None
90 + group_idx: int | None = None
91
92 def __hash__(self) -> int:
89 - return hash((tuple(self.block_hashes), self.medium))
93 + return hash(
94 + (
95 + tuple(self.block_hashes),
96 + self.medium,
97 + self.group_idx,
98 + )
99 + )
100
101
102 class AllBlocksCleared(KVCacheEvent):
vllm/v1/core/block_pool.py
+3 -4
@@ -22,6 +22,7 @@ from vllm.v1.core.kv_cache_utils import (
22 KVCacheBlock,
23 generate_block_hash_extra_keys,
24 get_block_hash,
25 + get_group_id,
26 make_block_hash_with_group_id,
27 maybe_convert_block_hash,
28 )
@@ -314,6 +315,7 @@ class BlockPool:
315 if request.lora_request
316 else None,
317 extra_keys=extra_keys_list if extra_keys_list else None,
318 + group_idx=kv_cache_group_id,
319 )
320 )
321
@@ -377,14 +379,11 @@ class BlockPool:
379 block.reset_hash()
380
381 if self.enable_kv_cache_events:
380 - # FIXME (Chen): Not sure whether we should return `hash_value`
381 - # or `(hash_value, group_id)` here. But it's fine now because
382 - # we disable hybrid kv cache manager when kv cache event is
383 - # enabled, so there is only one group.
382 self.kv_event_queue.append(
383 BlockRemoved(
384 block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
385 medium=MEDIUM_GPU,
386 + group_idx=get_group_id(block_hash),
387 )
388 )
389 return True