[Core][KV Events] Echo session_id on GPU BlockStored events

examples/features/kv_events/kv_events_subscriber.py

@@ -45,6 +45,7 @@ class BlockStored(KVCacheEvent):
     kv_cache_spec_kind: str | None = None
     kv_cache_spec_sliding_window: int | None = None
     locality: str | None = None
+    session_id: str | None = None
 
 
 class BlockRemoved(KVCacheEvent):

tests/distributed/test_kv_cache_events.py

@@ -46,10 +46,33 @@ class _LegacyBlockRemoved(
     group_idx: int | None = None
 
 
+class _PreSessionBlockStored(
+    msgspec.Struct,
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,  # type: ignore[call-arg]
+    tag="BlockStored",  # type: ignore[call-arg]
+):
+    """BlockStored wire schema before session_id was added."""
+
+    block_hashes: list[bytes]
+    parent_block_hash: bytes | None
+    token_ids: list[int]
+    block_size: int
+    lora_id: int | None
+    medium: str | None
+    lora_name: str | None
+    extra_keys: list[tuple[Any, ...] | None] | None = None
+    group_idx: int | None = None
+    kv_cache_spec_kind: str | None = None
+    kv_cache_spec_sliding_window: int | None = None
+    locality: str | None = None
+
+
 def _make_block_stored(
     group_idx: int | None = None,
     kv_cache_spec_sliding_window: int | None = None,
     locality: str | None = None,
+    session_id: str | None = None,
 ) -> BlockStored:
     return BlockStored(
         block_hashes=[_FAKE_HASH],
@@ -62,6 +85,7 @@ def _make_block_stored(
         group_idx=group_idx,
         kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
         locality=locality,
+        session_id=session_id,
     )
 
 
@@ -105,6 +129,12 @@ def test_block_stored_hash_same_for_equal_group_idx():
     assert hash(event_a) == hash(event_b)
 
 
+def test_block_stored_hash_differs_by_session_id():
+    event_a = _make_block_stored(session_id="session-a")
+    event_b = _make_block_stored(session_id="session-b")
+    assert hash(event_a) != hash(event_b)
+
+
 @pytest.mark.parametrize("group_idx", [1, 2, 3])
 def test_block_removed_hash_differs_by_group_idx(group_idx: int):
     """BlockRemoved events that differ only in group_idx must hash differently."""
@@ -183,3 +213,38 @@ def test_block_removed_locality_is_wire_compatible():
     new_payload = msgspec.msgpack.encode(_make_block_removed(locality="REMOTE"))
     assert msgspec.msgpack.decode(new_payload)["locality"] == "REMOTE"
     assert msgspec.msgpack.decode(new_payload, type=_LegacyBlockRemoved).medium == "GPU"
+
+
+def test_block_stored_session_id_is_wire_compatible():
+    pre_session = _PreSessionBlockStored(
+        block_hashes=[_FAKE_HASH],
+        parent_block_hash=None,
+        token_ids=[1, 2, 3, 4],
+        block_size=4,
+        lora_id=None,
+        medium="GPU",
+        lora_name=None,
+        group_idx=2,
+        kv_cache_spec_sliding_window=128,
+        locality="LOCAL",
+    )
+    pre_session_payload = msgspec.msgpack.encode(pre_session)
+    assert (
+        msgspec.msgpack.encode(
+            _make_block_stored(
+                group_idx=2,
+                kv_cache_spec_sliding_window=128,
+                locality="LOCAL",
+            )
+        )
+        == pre_session_payload
+    )
+    assert (
+        msgspec.msgpack.decode(pre_session_payload, type=BlockStored).session_id is None
+    )
+
+    new_payload = msgspec.msgpack.encode(_make_block_stored(session_id="session-1"))
+    assert msgspec.msgpack.decode(new_payload)["session_id"] == "session-1"
+    assert (
+        msgspec.msgpack.decode(new_payload, type=_PreSessionBlockStored).medium == "GPU"
+    )

tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py

@@ -33,6 +33,7 @@ def make_request(
     prompt_token_ids: list[int],
     hash_block_size: int,
     hash_fn: Callable,
+    session_id: str | None = None,
 ) -> Request:
     sampling_params = SamplingParams(max_tokens=17)
     sampling_params.update_from_generation_config({}, eos_token_id=100)
@@ -42,6 +43,7 @@ def make_request(
         sampling_params=sampling_params,
         pooling_params=None,
         block_hasher=get_request_block_hasher(hash_block_size, hash_fn),
+        session_id=session_id,
     )
 
 
@@ -123,6 +125,7 @@ def test_cache_partial_block_kv_cache_events():
         prompt_token_ids=list(range(hash_block_size * 2)),
         hash_block_size=hash_block_size,
         hash_fn=sha256,
+        session_id="agent-session-partial",
     )
 
     block = pool.get_new_blocks(1)[0]
@@ -148,6 +151,7 @@ def test_cache_partial_block_kv_cache_events():
     assert stored_event.token_ids == req.all_token_ids[hash_block_size:]
     assert stored_event.block_size == 4
     assert stored_event.group_idx == kv_cache_group_id
+    assert stored_event.session_id == "agent-session-partial"
 
     duplicate_entry_hash = pool.cache_partial_block(
         request=req,

tests/v1/core/test_prefix_caching.py

@@ -73,6 +73,7 @@ def make_request(
     prompt_logprobs: int | None = None,
     cache_salt: str | None = None,
     lora_request: LoRARequest | None = None,
+    session_id: str | None = None,
 ):
     mm_features = []
     if mm_positions is not None:
@@ -98,6 +99,7 @@ def make_request(
         lora_request=lora_request,
         cache_salt=cache_salt,
         block_hasher=get_request_block_hasher(block_size, hash_fn),
+        session_id=session_id,
     )
 
 
@@ -2294,6 +2296,71 @@ def test_block_stored_event_group_idx(group_id: int):
     )
 
 
+def test_block_stored_event_session_id():
+    block_size = 4
+    session_id = "agent-session-1"
+    manager = make_kv_cache_manager(
+        make_kv_cache_config(block_size, num_blocks=5),
+        max_model_len=8192,
+        enable_caching=True,
+        enable_kv_cache_events=True,
+        hash_block_size=block_size,
+    )
+    req = make_request(
+        "req_session_id",
+        prompt_token_ids=list(range(block_size * 2)),
+        block_size=block_size,
+        hash_fn=sha256,
+        session_id=session_id,
+    )
+
+    manager.allocate_slots(req, req.num_tokens)
+    events = manager.take_events()
+
+    assert len(events) == 1
+    assert isinstance(events[0], BlockStored)
+    assert events[0].session_id == session_id
+
+
+def test_session_id_does_not_affect_prefix_cache_identity():
+    block_size = 4
+    prompt_token_ids = list(range(block_size * 2 + 1))
+    manager = make_kv_cache_manager(
+        make_kv_cache_config(block_size, num_blocks=8),
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+    first_req = make_request(
+        "req_session_a",
+        prompt_token_ids,
+        block_size,
+        sha256,
+        session_id="session-a",
+    )
+    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(first_req)
+    allocated_blocks = manager.allocate_slots(
+        first_req,
+        first_req.num_tokens,
+        num_computed_tokens,
+        computed_blocks,
+    )
+    assert allocated_blocks is not None
+
+    second_req = make_request(
+        "req_session_b",
+        prompt_token_ids,
+        block_size,
+        sha256,
+        session_id="session-b",
+    )
+    assert second_req.block_hashes == first_req.block_hashes
+
+    cached_blocks, num_cached_tokens, _ = manager.get_computed_blocks(second_req)
+    assert num_cached_tokens == block_size * 2
+    assert cached_blocks.get_block_ids() == (allocated_blocks.get_block_ids()[0][:2],)
+
+
 def test_block_stored_event_group_idx_multiple_groups():
     """
     Test BlockStored events for separate HMA groups that each carry the
@@ -2490,6 +2557,7 @@ def test_emit_cached_block_events():
         prompt_token_ids=list(range(num_tokens)),
         block_size=block_size,
         hash_fn=sha256,
+        session_id="agent-session-reuse",
     )
     assert len(req.block_hashes) >= num_cached_blocks
 
@@ -2527,6 +2595,7 @@ def test_emit_cached_block_events():
     assert event.medium == MEDIUM_GPU
     assert event.lora_id is None
     assert event.lora_name is None
+    assert event.session_id == "agent-session-reuse"
 
 
 def test_emit_cached_block_events_disabled():

vllm/distributed/kv_events.py

@@ -75,6 +75,12 @@ class BlockStored(KVCacheEvent):
     kv_cache_spec_sliding_window: int | None = None
     locality: str | None = None
     """LOCAL or REMOTE relative to the publisher; None means unspecified."""
+    session_id: str | None = None
+    """Session that triggered this store or reuse report.
+
+    This identifies the request context that emitted the event, not exclusive
+    ownership of the underlying block, which may be shared across sessions.
+    """
 
     def __hash__(self) -> int:
         return hash(
@@ -90,6 +96,7 @@ def __hash__(self) -> int:
                 self.kv_cache_spec_kind,
                 self.kv_cache_spec_sliding_window,
                 self.locality,
+                self.session_id,
             )
         )
 

vllm/v1/core/block_pool.py

@@ -368,6 +368,7 @@ def _build_block_stored_event(
             lora_name=request.lora_request.name if request.lora_request else None,
             extra_keys=extra_keys_list if extra_keys_list else None,
             group_idx=kv_cache_group_id,
+            session_id=request.session_id,
         )
 
     def emit_cached_block_events(
@@ -539,6 +540,7 @@ def cache_partial_block(
                     else None,
                     extra_keys=[extra_keys],
                     group_idx=kv_cache_group_id,
+                    session_id=request.session_id,
                 )
             )
         return block_hash_with_group_id