[KV Offload] Support partial secondary-tier load results (#50321)
Signed-off-by: Moein Khazraee <moein@nvidia.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Moein Khazraee committed
Aug 4, 2026 at 20:48 UTC
41a7e7da0c893ff5476d8e6a6a7d2444bb9846d8
3 files changed
+88
-6
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+47
@@ -403,6 +403,53 @@ class TestTieringOffloadingManager:
403
# Next lookup should succeed
404
assert count_hits(self.manager, blocks) == 3
405
406
+ @pytest.mark.parametrize(
407
+ ("successful_indices", "expected_results"),
408
+ [
409
+ (
410
+ (0, 2),
411
+ [LookupResult.HIT, LookupResult.MISS, LookupResult.HIT],
412
+ ),
413
+ (
414
+ None,
415
+ [LookupResult.MISS, LookupResult.MISS, LookupResult.MISS],
416
+ ),
417
+ ],
418
+ ids=["partial", "legacy-full-failure"],
419
+ )
420
+ def test_failed_promotion_keeps_only_successful_blocks(
421
+ self, manager_setup, successful_indices, expected_results
422
+ ):
423
+ blocks = to_keys(range(3))
424
+ for block in blocks:
425
+ self.secondary_tier1.blocks[block] = True
426
+
427
+ def submit_partial(job_metadata: JobMetadata) -> None:
428
+ successful_keys = (
429
+ None
430
+ if successful_indices is None
431
+ else tuple(blocks[i] for i in successful_indices)
432
+ )
433
+ self.secondary_tier1.completed_jobs.append(
434
+ JobResult(
435
+ job_id=job_metadata.job_id,
436
+ success=False,
437
+ successful_keys=successful_keys,
438
+ )
439
+ )
440
+
441
+ self.secondary_tier1.submit_load = submit_partial
442
+
443
+ for block in blocks:
444
+ assert self.manager.lookup(block, _CTX) is LookupResult.RETRY
445
+
446
+ self._simulate_on_schedule_end()
447
+ self._simulate_on_schedule_end()
448
+
449
+ assert [
450
+ self.primary_tier.lookup(block, _CTX) for block in blocks
451
+ ] == expected_results
452
+
453
def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup):
454
"""Resolved lookups report one sync delay sample on allocation."""
455
self._start_request()
vllm/v1/kv_offload/tiering/base.py
+6
-1
@@ -54,10 +54,15 @@ class JobMetadata:
54
55
@dataclass
56
class JobResult:
57
- """Result of an async transfer job (successful or failed)."""
57
+ """Result of an async transfer job."""
58
59
job_id: JobId
60
+ # True if all keys succeeded; False if all or some failed.
61
success: bool
62
+ # Only applicable to promotion jobs. On partial failure, identifies the
63
+ # keys that were successfully loaded. None means all keys share the fate
64
+ # indicated by `success`. Must be a subset of the job's original keys.
65
+ successful_keys: Collection[OffloadKey] | None = None
66
67
68
class ParentManager(ABC):
vllm/v1/kv_offload/tiering/manager.py
+35
-5
@@ -49,6 +49,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
49
from vllm.v1.kv_offload.tiering.base import (
50
JobId,
51
JobMetadata,
52
+ JobResult,
53
ParentManager,
54
SecondaryTierManager,
55
TieringOffloadingMetrics,
@@ -243,6 +244,39 @@ class TieringOffloadingManager(OffloadingManager):
244
self._processed_jobs_this_step = True
245
self._process_finished_jobs()
246
247
+ def _complete_promotion(
248
+ self, job_metadata: JobMetadata, completed_job: JobResult
249
+ ) -> None:
250
+ successful_keys = completed_job.successful_keys
251
+ failed_keys: Collection[OffloadKey]
252
+ if completed_job.success:
253
+ successful_keys = job_metadata.keys
254
+ failed_keys = ()
255
+ elif successful_keys:
256
+ failed_keys_set = set(job_metadata.keys)
257
+ assert failed_keys_set.issuperset(successful_keys), (
258
+ f"Finished promotion job_id {completed_job.job_id} "
259
+ "reported unknown successful keys"
260
+ )
261
+ failed_keys_set.difference_update(successful_keys)
262
+ failed_keys = failed_keys_set
263
+ else:
264
+ successful_keys = ()
265
+ failed_keys = job_metadata.keys
266
+
267
+ if successful_keys:
268
+ self.primary_tier.complete_write(
269
+ successful_keys,
270
+ job_metadata.req_context,
271
+ True,
272
+ )
273
+ if failed_keys:
274
+ self.primary_tier.complete_write(
275
+ failed_keys,
276
+ job_metadata.req_context,
277
+ False,
278
+ )
279
+
280
def _process_finished_jobs(self):
281
"""
282
Unconditionally poll all secondary tiers for completed jobs.
@@ -266,11 +300,7 @@ class TieringOffloadingManager(OffloadingManager):
300
if job_metadata.is_promotion:
301
# secondary→primary transfer (promotion) completed.
302
# Make blocks available in primary tier.
269
- self.primary_tier.complete_write(
270
- job_metadata.keys,
271
- job_metadata.req_context,
272
- completed_job.success,
273
- )
303
+ self._complete_promotion(job_metadata, completed_job)
304
else:
305
# primary→secondary transfer completed.
306
# Decrement ref_cnt on primary blocks.