[Core][KV Events] Echo session_id on GPU BlockStored events (#51381)
Signed-off-by: liuguangli <liuguangli35@gmail.com> Signed-off-by: jinx-machinepulse <2679622408@qq.com> Co-authored-by: liuguangli <liuguangli35@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
guangli_Liu committed
Sep 4, 2026 at 10:22 UTC
d9e2b5238d4a42341db3374226890a801357c646
6 files changed
+148
examples/features/kv_events/kv_events_subscriber.py
+1
@@ -45,6 +45,7 @@ class BlockStored(KVCacheEvent):
45
kv_cache_spec_kind: str | None = None
46
kv_cache_spec_sliding_window: int | None = None
47
locality: str | None = None
48
+ session_id: str | None = None
49
50
51
class BlockRemoved(KVCacheEvent):
tests/distributed/test_kv_cache_events.py
+65
@@ -46,10 +46,33 @@ class _LegacyBlockRemoved(
46
group_idx: int | None = None
47
48
49
+class _PreSessionBlockStored(
50
+ msgspec.Struct,
51
+ omit_defaults=True, # type: ignore[call-arg]
52
+ gc=False, # type: ignore[call-arg]
53
+ tag="BlockStored", # type: ignore[call-arg]
54
+):
55
+ """BlockStored wire schema before session_id was added."""
56
+
57
+ block_hashes: list[bytes]
58
+ parent_block_hash: bytes | None
59
+ token_ids: list[int]
60
+ block_size: int
61
+ lora_id: int | None
62
+ medium: str | None
63
+ lora_name: str | None
64
+ extra_keys: list[tuple[Any, ...] | None] | None = None
65
+ group_idx: int | None = None
66
+ kv_cache_spec_kind: str | None = None
67
+ kv_cache_spec_sliding_window: int | None = None
68
+ locality: str | None = None
69
+
70
+
71
def _make_block_stored(
72
group_idx: int | None = None,
73
kv_cache_spec_sliding_window: int | None = None,
74
locality: str | None = None,
75
+ session_id: str | None = None,
76
) -> BlockStored:
77
return BlockStored(
78
block_hashes=[_FAKE_HASH],
@@ -62,6 +85,7 @@ def _make_block_stored(
85
group_idx=group_idx,
86
kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
87
locality=locality,
88
+ session_id=session_id,
89
)
90
91
@@ -105,6 +129,12 @@ def test_block_stored_hash_same_for_equal_group_idx():
129
assert hash(event_a) == hash(event_b)
130
131
132
+def test_block_stored_hash_differs_by_session_id():
133
+ event_a = _make_block_stored(session_id="session-a")
134
+ event_b = _make_block_stored(session_id="session-b")
135
+ assert hash(event_a) != hash(event_b)
136
+
137
+
138
@pytest.mark.parametrize("group_idx", [1, 2, 3])
139
def test_block_removed_hash_differs_by_group_idx(group_idx: int):
140
"""BlockRemoved events that differ only in group_idx must hash differently."""
@@ -183,3 +213,38 @@ def test_block_removed_locality_is_wire_compatible():
213
new_payload = msgspec.msgpack.encode(_make_block_removed(locality="REMOTE"))
214
assert msgspec.msgpack.decode(new_payload)["locality"] == "REMOTE"
215
assert msgspec.msgpack.decode(new_payload, type=_LegacyBlockRemoved).medium == "GPU"
216
+
217
+
218
+def test_block_stored_session_id_is_wire_compatible():
219
+ pre_session = _PreSessionBlockStored(
220
+ block_hashes=[_FAKE_HASH],
221
+ parent_block_hash=None,
222
+ token_ids=[1, 2, 3, 4],
223
+ block_size=4,
224
+ lora_id=None,
225
+ medium="GPU",
226
+ lora_name=None,
227
+ group_idx=2,
228
+ kv_cache_spec_sliding_window=128,
229
+ locality="LOCAL",
230
+ )
231
+ pre_session_payload = msgspec.msgpack.encode(pre_session)
232
+ assert (
233
+ msgspec.msgpack.encode(
234
+ _make_block_stored(
235
+ group_idx=2,
236
+ kv_cache_spec_sliding_window=128,
237
+ locality="LOCAL",
238
+ )
239
+ )
240
+ == pre_session_payload
241
+ )
242
+ assert (
243
+ msgspec.msgpack.decode(pre_session_payload, type=BlockStored).session_id is None
244
+ )
245
+
246
+ new_payload = msgspec.msgpack.encode(_make_block_stored(session_id="session-1"))
247
+ assert msgspec.msgpack.decode(new_payload)["session_id"] == "session-1"
248
+ assert (
249
+ msgspec.msgpack.decode(new_payload, type=_PreSessionBlockStored).medium == "GPU"
250
+ )
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py
+4
@@ -33,6 +33,7 @@ def make_request(
33
prompt_token_ids: list[int],
34
hash_block_size: int,
35
hash_fn: Callable,
36
+ session_id: str | None = None,
37
) -> Request:
38
sampling_params = SamplingParams(max_tokens=17)
39
sampling_params.update_from_generation_config({}, eos_token_id=100)
@@ -42,6 +43,7 @@ def make_request(
43
sampling_params=sampling_params,
44
pooling_params=None,
45
block_hasher=get_request_block_hasher(hash_block_size, hash_fn),
46
+ session_id=session_id,
47
)
48
49
@@ -123,6 +125,7 @@ def test_cache_partial_block_kv_cache_events():
125
prompt_token_ids=list(range(hash_block_size * 2)),
126
hash_block_size=hash_block_size,
127
hash_fn=sha256,
128
+ session_id="agent-session-partial",
129
)
130
131
block = pool.get_new_blocks(1)[0]
@@ -148,6 +151,7 @@ def test_cache_partial_block_kv_cache_events():
151
assert stored_event.token_ids == req.all_token_ids[hash_block_size:]
152
assert stored_event.block_size == 4
153
assert stored_event.group_idx == kv_cache_group_id
154
+ assert stored_event.session_id == "agent-session-partial"
155
156
duplicate_entry_hash = pool.cache_partial_block(
157
request=req,
tests/v1/core/test_prefix_caching.py
+69
@@ -74,6 +74,7 @@ def make_request(
74
prompt_logprobs: int | None = None,
75
cache_salt: str | None = None,
76
lora_request: LoRARequest | None = None,
77
+ session_id: str | None = None,
78
):
79
mm_features = []
80
if mm_positions is not None:
@@ -99,6 +100,7 @@ def make_request(
100
lora_request=lora_request,
101
cache_salt=cache_salt,
102
block_hasher=get_request_block_hasher(block_size, hash_fn),
103
+ session_id=session_id,
104
)
105
106
@@ -2378,6 +2380,71 @@ def test_block_stored_event_group_idx(group_id: int):
2380
)
2381
2382
2383
+def test_block_stored_event_session_id():
2384
+ block_size = 4
2385
+ session_id = "agent-session-1"
2386
+ manager = make_kv_cache_manager(
2387
+ make_kv_cache_config(block_size, num_blocks=5),
2388
+ max_model_len=8192,
2389
+ enable_caching=True,
2390
+ enable_kv_cache_events=True,
2391
+ hash_block_size=block_size,
2392
+ )
2393
+ req = make_request(
2394
+ "req_session_id",
2395
+ prompt_token_ids=list(range(block_size * 2)),
2396
+ block_size=block_size,
2397
+ hash_fn=sha256,
2398
+ session_id=session_id,
2399
+ )
2400
+
2401
+ manager.allocate_slots(req, req.num_tokens)
2402
+ events = manager.take_events()
2403
+
2404
+ assert len(events) == 1
2405
+ assert isinstance(events[0], BlockStored)
2406
+ assert events[0].session_id == session_id
2407
+
2408
+
2409
+def test_session_id_does_not_affect_prefix_cache_identity():
2410
+ block_size = 4
2411
+ prompt_token_ids = list(range(block_size * 2 + 1))
2412
+ manager = make_kv_cache_manager(
2413
+ make_kv_cache_config(block_size, num_blocks=8),
2414
+ max_model_len=8192,
2415
+ enable_caching=True,
2416
+ hash_block_size=block_size,
2417
+ )
2418
+ first_req = make_request(
2419
+ "req_session_a",
2420
+ prompt_token_ids,
2421
+ block_size,
2422
+ sha256,
2423
+ session_id="session-a",
2424
+ )
2425
+ computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(first_req)
2426
+ allocated_blocks = manager.allocate_slots(
2427
+ first_req,
2428
+ first_req.num_tokens,
2429
+ num_computed_tokens,
2430
+ computed_blocks,
2431
+ )
2432
+ assert allocated_blocks is not None
2433
+
2434
+ second_req = make_request(
2435
+ "req_session_b",
2436
+ prompt_token_ids,
2437
+ block_size,
2438
+ sha256,
2439
+ session_id="session-b",
2440
+ )
2441
+ assert second_req.block_hashes == first_req.block_hashes
2442
+
2443
+ cached_blocks, num_cached_tokens, _ = manager.get_computed_blocks(second_req)
2444
+ assert num_cached_tokens == block_size * 2
2445
+ assert cached_blocks.get_block_ids() == (allocated_blocks.get_block_ids()[0][:2],)
2446
+
2447
+
2448
def test_block_stored_event_group_idx_multiple_groups():
2449
"""
2450
Test BlockStored events for separate HMA groups that each carry the
@@ -2574,6 +2641,7 @@ def test_emit_cached_block_events():
2641
prompt_token_ids=list(range(num_tokens)),
2642
block_size=block_size,
2643
hash_fn=sha256,
2644
+ session_id="agent-session-reuse",
2645
)
2646
assert len(req.block_hashes) >= num_cached_blocks
2647
@@ -2611,6 +2679,7 @@ def test_emit_cached_block_events():
2679
assert event.medium == MEDIUM_GPU
2680
assert event.lora_id is None
2681
assert event.lora_name is None
2682
+ assert event.session_id == "agent-session-reuse"
2683
2684
2685
def test_emit_cached_block_events_disabled():
vllm/distributed/kv_events.py
+7
@@ -78,6 +78,12 @@ class BlockStored(KVCacheEvent):
78
locality: str | None = None
79
"""Secondary offloading tier identifier, if generated by one."""
80
ownership: str | None = None
81
+ """Session that triggered this store or reuse report.
82
+
83
+ This identifies the request context that emitted the event, not exclusive
84
+ ownership of the underlying block, which may be shared across sessions.
85
+ """
86
+ session_id: str | None = None
87
88
def __hash__(self) -> int:
89
return hash(
@@ -94,6 +100,7 @@ class BlockStored(KVCacheEvent):
100
self.kv_cache_spec_sliding_window,
101
self.locality,
102
self.ownership,
103
+ self.session_id,
104
)
105
)
106
vllm/v1/core/block_pool.py
+2
@@ -368,6 +368,7 @@ class BlockPool:
368
lora_name=request.lora_request.name if request.lora_request else None,
369
extra_keys=extra_keys_list if extra_keys_list else None,
370
group_idx=kv_cache_group_id,
371
+ session_id=request.session_id,
372
)
373
374
def emit_cached_block_events(
@@ -539,6 +540,7 @@ class BlockPool:
540
else None,
541
extra_keys=[extra_keys],
542
group_idx=kv_cache_group_id,
543
+ session_id=request.session_id,
544
)
545
)
546
return block_hash_with_group_id