Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
beefff2174 | ||
|
|
163266d0b2 | ||
|
|
a2fd28a7e1 | ||
|
|
7a80ac928f | ||
|
|
0ed11013b4 | ||
|
|
6d568b995a | ||
|
|
dfe9decbcb | ||
|
|
1324e6ff67 | ||
|
|
c1aba6d7ae | ||
|
|
200bef28c9 | ||
|
|
fd9820bbf9 |
@@ -10,6 +10,7 @@ from vllm.distributed.kv_transfer.kv_transfer_state import (
|
||||
get_kv_transfer_group,
|
||||
)
|
||||
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
|
||||
from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
|
||||
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
|
||||
|
||||
# Importing utils registers TestExampleConnector with the factory
|
||||
@@ -59,3 +60,29 @@ def test_kv_connector_mixin_clears_metadata():
|
||||
finally:
|
||||
# Ensure we clean up the global connector between tests
|
||||
ensure_kv_transfer_shutdown()
|
||||
|
||||
|
||||
def test_active_kv_connector_runs_lifecycle_hooks_for_empty_metadata():
|
||||
vllm_config = create_vllm_config()
|
||||
vllm_config.kv_transfer_config.kv_connector = "TestExampleConnector"
|
||||
vllm_config.kv_transfer_config.kv_role = "kv_both"
|
||||
vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "empty"
|
||||
|
||||
ensure_kv_transfer_initialized(vllm_config)
|
||||
|
||||
try:
|
||||
wrapped = get_kv_transfer_group()
|
||||
connector = ActiveKVConnector(vllm_config, {})
|
||||
scheduler_output = _make_empty_scheduler_output()
|
||||
|
||||
connector.pre_forward(scheduler_output)
|
||||
connector.post_forward(scheduler_output)
|
||||
|
||||
assert wrapped.call_record.get("bind_connector_metadata", 0) == 1
|
||||
assert wrapped.call_record.get("handle_preemptions", 0) == 1
|
||||
assert wrapped.call_record.get("start_load_kv", 0) == 1
|
||||
assert wrapped.call_record.get("wait_for_save", 0) == 1
|
||||
assert wrapped.call_record.get("get_finished", 0) == 1
|
||||
assert wrapped.call_record.get("clear_connector_metadata", 0) == 1
|
||||
finally:
|
||||
ensure_kv_transfer_shutdown()
|
||||
|
||||
@@ -521,7 +521,11 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
|
||||
def wait_for_save(self):
|
||||
assert self.connector_worker is not None
|
||||
assert isinstance(self._connector_metadata, NixlConnectorMetadata)
|
||||
if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks:
|
||||
if (
|
||||
self.connector_worker.use_host_buffer
|
||||
and self.connector_worker.copy_blocks
|
||||
and self._connector_metadata.reqs_to_save
|
||||
):
|
||||
self.connector_worker.save_kv_to_host(self._connector_metadata)
|
||||
|
||||
def shutdown(self):
|
||||
@@ -2466,6 +2470,16 @@ class NixlConnectorWorker:
|
||||
Start loading by triggering non-blocking nixl_xfer.
|
||||
We check for these trnxs to complete in each step().
|
||||
"""
|
||||
# skip the empty path
|
||||
if (
|
||||
not metadata.reqs_to_recv
|
||||
and not metadata.reqs_to_send
|
||||
and not metadata.reqs_in_batch
|
||||
and not metadata.reqs_not_processed
|
||||
and self._ready_requests.empty()
|
||||
):
|
||||
return
|
||||
|
||||
for req_id, meta in metadata.reqs_to_recv.items():
|
||||
meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
|
||||
meta.local_block_ids
|
||||
|
||||
Reference in New Issue
Block a user