From f65450e3dc9541eb311b8dd559b6460184864784 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 7 Jul 2025 01:27:40 +0000 Subject: [PATCH] updated Signed-off-by: Robert Shaw --- .../kv_transfer/kv_connector/v1/nixl_connector.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index 11c61f52ca7cc..325530cab69aa 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -329,14 +329,17 @@ class NixlConnectorWorker: self.block_size = vllm_config.cache_config.block_size # Agent. + import os + num_workers = 32 # setting num_workers on the prefiller causes no notifs to be recved??? # this is a hack to make sure we set num workers on the prefiller to 1. - NUM_WORKERS = 32 - logger.info(f"{NUM_WORKERS=}") + if os.getenv("VLLM_IS_PREFILL", "0") == "1": + num_workers = None + print(f"NUM_WORKERS: {num_workers=}") self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), None, num_workers=None, - num_shared_workers=NUM_WORKERS) + num_shared_workers=num_workers) # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}. self._remote_agents: dict[str, dict[int, str]] = defaultdict(dict)