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

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 11, 23:45 UTC

[NIXL] Small cleanup of unused variables (#29618)

Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Nicolò Lucchesi committed Dec 5, 2025 at 18:17 UTC 78c44fd722fe3ce010b0ba9c5e7349fe4094b39d

2 files changed +8 -14

tests/v1/kv_connector/unit/test_nixl_connector.py

+1 -1

@@ -1307,7 +1307,7 @@ def test_shutdown_cleans_up_resources(dist_init):

1307 patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem_agent,

1308 patch.object(nixl_wrapper, "deregister_memory") as mock_dereg,

1309 ):

1310 - worker._recving_transfers = {"req1": [(123, time.perf_counter())]}

1310 + worker._recving_transfers = {"req1": [123]}

1311 worker.src_xfer_side_handle = 456

1312 worker.dst_xfer_side_handles = {"engine1": 789}

1313 worker._remote_agents = {"engine1": {0: "agent1"}}

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

+7 -13

@@ -55,7 +55,7 @@ if TYPE_CHECKING:

55 from vllm.v1.kv_cache_interface import KVCacheConfig

56 from vllm.v1.request import Request

58 -Transfer = tuple[int, float] # (xfer_handle, start_time)

58 +TransferHandle = int

59 EngineId = str

60 ReqId = str

@@ -874,7 +874,7 @@ class NixlConnectorWorker:

874 # In progress transfers.

875 # [req_id -> list[handle]]

876 self._recving_metadata: dict[ReqId, ReqMeta] = {}

877 - self._recving_transfers = defaultdict[ReqId, list[Transfer]](list)

877 + self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list)

878 # Track the expiration time of requests that are waiting to be sent.

879 self._reqs_to_send: dict[ReqId, float] = {}

880 # Set of requests that have been part of a batch, regardless of status.

@@ -1201,14 +1201,11 @@ class NixlConnectorWorker:

1201 # Enable different block lengths for different layers when MLA is used.

1202 self.block_len_per_layer = list[int]()

1203 self.slot_size_per_layer = list[int]() # HD bytes in kv terms

1204 - self.device_id = self.tp_rank

1204 for layer_name, cache_or_caches in xfer_buffers.items():

1205 cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]

1206

1207 for cache in cache_list:

1208 base_addr = cache.data_ptr()

1210 - if not self.use_host_buffer and current_platform.is_cuda_alike():

1211 - self.device_id = cache.device.index

1209 if base_addr in seen_base_addresses:

1210 continue

1211

@@ -1251,8 +1248,7 @@ class NixlConnectorWorker:

1248 "All kv cache tensors must have the same size"

1249 )

1250 # Need to make sure the device ID is non-negative for NIXL,

1254 - # Torch uses -1 to indicate CPU tensors while NIXL uses explicit

1255 - # memory type.

1251 + # Torch uses -1 to indicate CPU tensors.

1252 self.device_id = max(cache.get_device(), 0)

1253 caches_data.append(

1254 (base_addr, curr_tensor_size_bytes, self.device_id, "")

@@ -1842,9 +1838,7 @@ class NixlConnectorWorker:

1838 self._reqs_to_send.pop(req_id, None)

1839 return notified_req_ids

1840

1845 - def _pop_done_transfers(

1846 - self, transfers: dict[str, list[tuple[int, float]]]

1847 - ) -> set[str]:

1841 + def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:

1842 """

1843 Pop completed xfers by checking for DONE state.

1844 Args:

@@ -1855,7 +1849,7 @@ class NixlConnectorWorker:

1849 done_req_ids: set[str] = set()

1850 for req_id, handles in list(transfers.items()):

1851 in_progress = False

1858 - for handle, xfer_start_time in handles:

1852 + for handle in handles:

1853 try:

1854 xfer_state = self.nixl_wrapper.check_xfer_state(handle)

1855 if xfer_state == "DONE":

@@ -2120,7 +2114,7 @@ class NixlConnectorWorker:

2114 self.nixl_wrapper.transfer(handle)

2115

2116 # Use handle to check completion in future step().

2123 - self._recving_transfers[request_id].append((handle, time.perf_counter()))

2117 + self._recving_transfers[request_id].append(handle)

2118 except Exception:

2119 logger.exception(

2120 "NIXL transfer setup/initiation failed for request %s. "

@@ -2251,7 +2245,7 @@ class NixlConnectorWorker:

2245 """Shutdown the connector worker."""

2246 self._handshake_initiation_executor.shutdown(wait=False)

2247 for handles in self._recving_transfers.values():

2254 - for handle, _ in handles:

2248 + for handle in handles:

2249 self.nixl_wrapper.release_xfer_handle(handle)

2250 self._recving_transfers.clear()

2251 if self.src_xfer_side_handle: