@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fcf0687b27

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 13:15 UTC

[kv_offload+HMA][0/N]: Support block-level preemption handling (#34805)

Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

Or Ozeri committed Mar 18, 2026 at 08:49 UTC fcf0687b27b78c3b214504f5e9525f3f66a2d04a

7 files changed +27 -25

tests/v1/kv_connector/unit/test_multi_connector.py

+6 -4

@@ -231,10 +231,11 @@ def test_multi_example_connector_consistency():

231 ]

232 # First three events are from initialization (register_kv_caches,

233 # set_host_xfer_buffer_ops, get_handshake_metadata), then generate() events.

234 - assert events["storage1-WORKER"][:7] == [

234 + assert events["storage1-WORKER"][:8] == [

235 "register_kv_caches",

236 "set_host_xfer_buffer_ops",

237 "get_handshake_metadata",

238 + "handle_preemptions",

239 "bind_connector_metadata",

240 "start_load_kv",

241 "wait_for_layer_load",

@@ -246,10 +247,11 @@ def test_multi_example_connector_consistency():

247 "update_state_after_alloc num_blocks=[0] 0",

248 "build_connector_meta",

249 ]

249 - assert events["storage2-WORKER"][:7] == [

250 + assert events["storage2-WORKER"][:8] == [

251 "register_kv_caches",

252 "set_host_xfer_buffer_ops",

253 "get_handshake_metadata",

254 + "handle_preemptions",

255 "bind_connector_metadata",

256 "start_load_kv",

257 "wait_for_layer_load",

@@ -399,8 +401,8 @@ def test_multi_connector_handle_preemptions_integration():

401 # testing the delegation behavior of MultiConnector here.

402 # The connector attribute contains the KV connector.

403 assert scheduler.connector is not None, "Scheduler should have a connector"

402 - preempted_req_ids = {"req-1", "req-2", "req-3"}

403 - scheduler.connector.handle_preemptions(preempted_req_ids)

404 + connector_md = scheduler.connector.build_connector_meta(scheduler.schedule())

405 + scheduler.connector.handle_preemptions(connector_md)

406

407 # Verify both connectors received the handle_preemptions call

408 events = get_connector_events()

tests/v1/kv_connector/unit/test_offloading_connector.py

+1 -4

@@ -363,10 +363,7 @@ class RequestRunner:

363 assert kv_connector_metadata is not None

364 assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)

365

366 - if scheduler_output.preempted_req_ids:

367 - self.worker_connector.handle_preemptions(

368 - scheduler_output.preempted_req_ids

369 - )

366 + self.worker_connector.handle_preemptions(kv_connector_metadata)

367

368 self.worker_connector.bind_connector_metadata(kv_connector_metadata)

369 self.worker_connector.start_load_kv(self._dummy_ctx)

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+4 -4

@@ -25,8 +25,8 @@ The class provides the following primitives:

26 Worker-side: runs in each worker, loads/saves KV cache to/from

27 the Connector based on the metadata.

28 - handle_preemptions() - called if there are preempted requests,

29 - before their blocks are overwritten

28 + handle_preemptions() - called for handling preempted requests

29 + or request evicted blocks before they are overwritten

31 start_load_kv() - starts loading all KVs (maybe async)

32 wait_for_layer_load() - blocks until layer i load is done

@@ -288,9 +288,9 @@ class KVConnectorBase_V1(ABC):

288 """

289 return

290

291 - def handle_preemptions(self, preempted_req_ids: set[str]):

291 + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):

292 """

293 - Handle preempted requests BEFORE their blocks are overwritten.

293 + Handle preempted requests or evicted blocks BEFORE they are overwritten.

294 Needed for connectors which use async saves (e.g., OffloadingConnector)

295 """

296 return

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

+4 -3

@@ -315,10 +315,11 @@ class MultiConnector(KVConnectorBase_V1):

315 for c in self._connectors:

316 c.set_host_xfer_buffer_ops(copy_operation)

317

318 - def handle_preemptions(self, preempted_req_ids: set[str]):

318 + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):

319 """Handle preempted requests for all sub-connectors."""

320 - for c in self._connectors:

321 - c.handle_preemptions(preempted_req_ids)

320 + assert isinstance(kv_connector_metadata, MultiKVConnectorMetadata)

321 + for c, cm in zip(self._connectors, kv_connector_metadata.metadata):

322 + c.handle_preemptions(cm)

323

324 def get_finished_count(self) -> int | None:

325 # TODO(https://github.com/vllm-project/vllm/issues/33400)

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

+7 -4

@@ -111,6 +111,7 @@ class OffloadingConnectorStats(KVConnectorStats):

111 class OffloadingConnectorMetadata(KVConnectorMetadata):

112 reqs_to_load: dict[ReqId, TransferSpec]

113 reqs_to_store: dict[ReqId, TransferSpec]

114 + reqs_to_flush: set[str] | None = None

115

116

117 class OffloadingConnector(KVConnectorBase_V1):

@@ -146,9 +147,10 @@ class OffloadingConnector(KVConnectorBase_V1):

147 assert self.connector_worker is not None

148 self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)

149

149 - def handle_preemptions(self, preempted_req_ids: set[str]):

150 + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):

151 assert self.connector_worker is not None

151 - self.connector_worker.handle_preemptions(preempted_req_ids)

152 + assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)

153 + self.connector_worker.handle_preemptions(kv_connector_metadata)

154

155 def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:

156 assert self.connector_worker is not None

@@ -482,6 +484,7 @@ class OffloadingConnectorScheduler:

484 meta = OffloadingConnectorMetadata(

485 reqs_to_load=self._reqs_to_load,

486 reqs_to_store=self._get_reqs_to_store(scheduler_output),

487 + reqs_to_flush=scheduler_output.preempted_req_ids,

488 )

489 self._reqs_to_load = {}

490

@@ -619,13 +622,13 @@ class OffloadingConnectorWorker:

622 attn_backends = {cross_layer_name: attn_backend}

623 self._register_handlers(kv_caches, attn_backends)

624

622 - def handle_preemptions(self, preempted_req_ids: set[str]):

625 + def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):

626 for job_id, transfer_spec in self._unsubmitted_store_jobs:

627 success = self.worker.transfer_async(job_id, transfer_spec)

628 assert success

629 self._unsubmitted_store_jobs.clear()

630

628 - for req_id in preempted_req_ids:

631 + for req_id in kv_connector_metadata.reqs_to_flush or ():

632 job_ids = self._store_jobs.get(req_id)

633 if job_ids:

634 self.worker.wait(job_ids)

vllm/v1/worker/gpu/kv_connector.py

+1 -2

@@ -63,11 +63,10 @@ class ActiveKVConnector(KVConnector):

63 if self._disabled:

64 return

66 - if scheduler_output.preempted_req_ids:

67 - self.kv_connector.handle_preemptions(scheduler_output.preempted_req_ids)

66 kv_connector_metadata = scheduler_output.kv_connector_metadata

67 assert kv_connector_metadata is not None

68 self.kv_connector.bind_connector_metadata(kv_connector_metadata)

69 + self.kv_connector.handle_preemptions(kv_connector_metadata)

71 # TODO: sort out KV Connectors' use of forward_context

72 if is_forward_context_available():

vllm/v1/worker/gpu_model_runner.py

+4 -4

@@ -3594,10 +3594,10 @@ class GPUModelRunner(

3594 scheduled_spec_decode_tokens=spec_decode_tokens_copy,

3595 )

3596

3597 - if scheduler_output.preempted_req_ids and has_kv_transfer_group():

3598 - get_kv_transfer_group().handle_preemptions(

3599 - scheduler_output.preempted_req_ids

3600 - )

3597 + if has_kv_transfer_group():

3598 + kv_connector_metadata = scheduler_output.kv_connector_metadata

3599 + assert kv_connector_metadata is not None

3600 + get_kv_transfer_group().handle_preemptions(kv_connector_metadata)

3601

3602 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

3603 with (