[Bugfix][KV Offload] Retain offload event metadata through batch translation (#57453)
Signed-off-by: Kapil Arya <kapila@nvidia.com> Signed-off-by: Kapil Arya <kapil.arya.17@gmail.com> Co-authored-by: Or Ozeri <or@ozery.com>
Kapil Arya committed
Sep 23, 2026 at 21:09 UTC
70fc359d257a984903b7ab9e123fbc45085c5254
2 files changed
+48
-2
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+38
@@ -506,6 +506,44 @@ def test_event_metadata_skips_non_full_attention_group(
506
assert events[0].block_size == 0
507
508
509
+@pytest.mark.parametrize("secondary_removal_in_same_batch", [False, True])
510
+def test_primary_removal_preserves_metadata_for_same_batch_kvcr_store(
511
+ secondary_removal_in_same_batch,
512
+):
513
+ """A removal must not drop the payload a later store in the same batch
514
+ still needs, and must still drop it once no residency is left."""
515
+ tracker = _tracker()
516
+ group_config = _group_config(block_size=64, blocks_per_chunk=2)
517
+ req = _request(block_hashes=[_hash(i) for i in range(4)], token_count=256)
518
+ key = _record_chunks(tracker, req, group_config, num_chunks=2)[1]
519
+ secondary_removal = _removed_event([key], ownership="kvcr")
520
+
521
+ batch = [
522
+ _removed_event([key]),
523
+ _stored_event([key], ownership="kvcr", removal_expected=True),
524
+ ]
525
+ if secondary_removal_in_same_batch:
526
+ batch.append(secondary_removal)
527
+ events = list(tracker.take_events(batch))
528
+ if not secondary_removal_in_same_batch:
529
+ events += tracker.take_events([secondary_removal])
530
+ removed, stored, final_removal = events
531
+
532
+ hashes = [_wire_hash(_hash(2)), _wire_hash(_hash(3))]
533
+ assert isinstance(removed, BlockRemoved)
534
+ assert removed.block_hashes == hashes
535
+ assert stored.block_size == 64
536
+ assert stored.token_ids == list(range(129, 257))
537
+ assert stored.parent_block_hash == _wire_hash(_hash(1))
538
+ assert stored.block_hashes == hashes
539
+ assert stored.ownership == "kvcr"
540
+ assert final_removal.block_hashes == hashes
541
+
542
+ # Nothing holds a residency now, so the payload is gone.
543
+ [after_eviction] = tracker.take_events([_stored_event([key], ownership="kvcr")])
544
+ assert after_eviction.block_size == 0
545
+
546
+
547
def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
548
tracker = _tracker()
549
block_hashes = [_hash(0), _hash(1)]
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+10
-2
@@ -247,12 +247,22 @@ class OffloadingEventsTracker:
247
the underlying :class:`OffloadingEvent` stream.
248
249
"""
250
+ removed_keys: set[OffloadKey] = set()
251
for event in events:
252
if event.removed:
253
+ if self.self_describing_enabled:
254
+ removed_keys.update(event.keys)
255
yield from self._take_removed_event(event)
256
else:
257
yield from self._take_stored_event(event)
258
259
+ # A primary removal can precede a queued secondary store in this batch.
260
+ # Keep its payload until all stores have registered their residencies.
261
+ for key in removed_keys:
262
+ meta = self._pending_event_metadata.get(key)
263
+ if meta is not None and not meta.active_residencies:
264
+ self._pending_event_metadata.pop(key)
265
+
266
def reset(self) -> None:
267
"""Drop all tracked state; pending payloads are stale after a
268
manager cache reset."""
@@ -407,8 +417,6 @@ class OffloadingEventsTracker:
417
maybe_convert_block_hash(h) for h in meta.block_hashes
418
)
419
meta.active_residencies.discard((event.medium, event.ownership))
410
- if not meta.active_residencies:
411
- self._pending_event_metadata.pop(key)
420
else:
421
if self.self_describing_enabled:
422
logger.warning_once(