@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f067737b2a

[Bugfix][TieredOffloading] : Return HIT_PENDING when KV promotion is triggered (#51840)

Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>

Varun Sundar Rabindranath committed Aug 11, 2026 at 22:07 UTC f067737b2a4ba4385f30256baa297200d06300c1

2 files changed +15 -15

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+14 -14

@@ -481,7 +481,7 @@ class TestTieringOffloadingManager:

481 # Lookup each block to initiate promotion for all of them

482 for block in blocks:

483 result = self.manager.lookup(block, _CTX)

484 - assert result is LookupResult.RETRY # promotion initiated

484 + assert result is LookupResult.HIT_PENDING # promotion initiated

485

486 # End of step 1: flushes deferred submit_load() calls

487 self._simulate_on_schedule_end()

@@ -533,7 +533,7 @@ class TestTieringOffloadingManager:

533 self.secondary_tier1.submit_load = submit_partial

534

535 for block in blocks:

536 - assert self.manager.lookup(block, _CTX) is LookupResult.RETRY

536 + assert self.manager.lookup(block, _CTX) is LookupResult.HIT_PENDING

537

538 self._simulate_on_schedule_end()

539 self._simulate_on_schedule_end()

@@ -589,7 +589,7 @@ class TestTieringOffloadingManager:

589 self.secondary_tier1.blocks[block] = True

590

591 # First lookup finds the block in a secondary tier and defers.

592 - assert self.manager.lookup(block, _CTX) is LookupResult.RETRY

592 + assert self.manager.lookup(block, _CTX) is LookupResult.HIT_PENDING

593

594 # Promotion is not treated as unresolved secondary lookup time.

595 self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id])

@@ -648,7 +648,7 @@ class TestTieringOffloadingManager:

648 stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY

649 )

650

651 - assert self.manager.lookup(block, ctx) is LookupResult.RETRY

651 + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING

652

653 stats = self.manager.get_stats()

654 assert stats is not None

@@ -771,11 +771,11 @@ class TestTieringOffloadingManager:

771 ctx_a = ReqContext(req_id="req_a")

772 ctx_b = ReqContext(req_id="req_b")

773

774 - # All lookups return RETRY: secondary hit triggers promotion

775 - assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.RETRY

776 - assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.RETRY

777 - assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.RETRY

778 - assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.RETRY

774 + # All lookups return HIT_PENDING: secondary hit triggers promotion

775 + assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.HIT_PENDING

776 + assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.HIT_PENDING

777 + assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.HIT_PENDING

778 + assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.HIT_PENDING

779

780 # submit_load must not fire during lookup - only at end of step

781 self.secondary_tier1.submit_load.assert_not_called()

@@ -812,9 +812,9 @@ class TestTieringOffloadingManager:

812 result_a = self.manager.lookup(shared_block, ctx_a)

813 result_b = self.manager.lookup(shared_block, ctx_b)

814

815 - # First lookup triggers promotion (RETRY), second finds block

816 - # already in primary with write in-flight (HIT_PENDING).

817 - assert result_a is LookupResult.RETRY

815 + # Both lookups find the shared_block and trigger promotion

816 + # returning HIT_PENDING.

817 + assert result_a is LookupResult.HIT_PENDING

818 assert result_b is LookupResult.HIT_PENDING

819

820 self._simulate_on_schedule_end()

@@ -1098,7 +1098,7 @@ class TestTieringOffloadingManager:

1098 self.secondary_tier1.blocks[promo_block] = True

1099 assert (

1100 self.manager.lookup(promo_block, ReqContext(req_id="pending"))

1101 - is LookupResult.RETRY

1101 + is LookupResult.HIT_PENDING

1102 )

1103 assert self.manager._pending_load_submissions

1104

@@ -1209,7 +1209,7 @@ class TestTieringOffloadingManager:

1209 self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup)

1210

1211 ctx = ReqContext(req_id="r2", load_tier_filter=load_tier_filter)

1212 - assert self.manager.lookup(blocks[0], ctx) is LookupResult.RETRY

1212 + assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT_PENDING

1213 self.secondary_tier1.lookup.assert_called()

1214

1215

vllm/v1/kv_offload/tiering/manager.py

+1 -1

@@ -401,7 +401,7 @@ class TieringOffloadingManager(OffloadingManager):

401 lookup_duration,

402 )

403 promoted = self._initiate_promotion(i, key, req_context)

404 - return LookupResult.MISS if not promoted else LookupResult.RETRY

404 + return LookupResult.MISS if not promoted else LookupResult.HIT_PENDING

405 if result is LookupResult.RETRY:

406 any_retry = True

407 self._metrics.on_lookup(