address review: clarity

Signed-off-by: nicklucche <[email protected]>

Author: NickLucche

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

@@ -155,7 +155,8 @@ def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
             self.connector_worker: Optional[NixlConnectorWorker] = None
         elif role == KVConnectorRole.WORKER:
             self.connector_scheduler = None
-            self.connector_worker = NixlConnectorWorker(str(self.engine_id))
+            self.connector_worker = NixlConnectorWorker(
+                vllm_config, str(self.engine_id))
 
     ############################################################
     # Scheduler Side Methods
@@ -349,13 +350,17 @@ def request_finished(
 class NixlConnectorWorker:
     """Implementation of Worker side methods"""
 
-    def __init__(self, engine_id: str):
+    def __init__(self, vllm_config: VllmConfig, engine_id: str):
         if NixlWrapper is None:
             logger.error("NIXL is not available")
             raise RuntimeError("NIXL is not available")
         logger.info("Initializing NIXL wrapper")
         logger.info("Initializing NIXL worker %s", engine_id)
 
+        # Config.
+        self.vllm_config = vllm_config
+        self.block_size = vllm_config.cache_config.block_size
+
         # Agent.
         self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), None)
         # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
@@ -454,41 +459,39 @@ def _nixl_handshake(self, host: str, port: int):
         # a hack to keep us moving. We will switch when moving to etcd
         # or where we have a single ZMQ socket in the scheduler.
 
-        def handshake(sock, rank: int) -> NixlAgentMetadata:
+        def handshake(path: str, rank: int) -> NixlAgentMetadata:
             # Send query for the request.
-            sock.send(GET_META_MSG)
-            metadata_bytes = sock.recv()
-            decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
-            metadata = decoder.decode(metadata_bytes)
-            got_metadata_time = time.perf_counter()
-
-            # Register Remote agent.
-            self.add_remote_agent(metadata, rank)
-            setup_agent_time = time.perf_counter()
-
-            logger.debug("NIXL handshake: get metadata took: %s",
-                         got_metadata_time - start_time)
-            logger.debug("NIXL handshake: add agent took: %s",
-                         setup_agent_time - got_metadata_time)
-            return metadata
+            with zmq_ctx(zmq.REQ, path) as sock:
+                sock.send(GET_META_MSG)
+                metadata_bytes = sock.recv()
+                decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
+                metadata = decoder.decode(metadata_bytes)
+                got_metadata_time = time.perf_counter()
+
+                # Register Remote agent.
+                self.add_remote_agent(metadata, rank)
+                setup_agent_time = time.perf_counter()
+
+                logger.debug("NIXL handshake: get metadata took: %s",
+                             got_metadata_time - start_time)
+                logger.debug("NIXL handshake: add agent took: %s",
+                             setup_agent_time - got_metadata_time)
+                return metadata
 
         # Handshake with remote agent-rank0 first to get the tp_size of remote
         path = f"tcp://{host}:{port}"
         logger.debug("Querying master rank metadata on path: %s", path)
-        with zmq_ctx(zmq.REQ, path) as sock:
-            metadata = handshake(sock, 0)
+        metadata = handshake(path, 0)
 
         # Handshake only with the other TP remote the current local rank will
         # pull from. With homogeneous TP it happens to be the same rank_i.
-        d_workers_per_p_worker = self._tp_size[
-            self.engine_id] // metadata.tp_size
-        p_remote_rank = self.rank // d_workers_per_p_worker
+        tp_rate = self._tp_size[self.engine_id] // metadata.tp_size
+        p_remote_rank = self.rank // tp_rate
         if p_remote_rank > 0:
             path = f"tcp://{host}:{port + p_remote_rank}"
             logger.debug("Querying metadata on path: %s at remote rank %s",
                          path, p_remote_rank)
-            with zmq_ctx(zmq.REQ, path) as sock:
-                metadata = handshake(sock, p_remote_rank)
+            _ = handshake(path, p_remote_rank)
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
@@ -503,17 +506,17 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             self.num_blocks = first_kv_cache.shape[0]
             block_rank = 2  # [block_size, latent_dim]
             block_shape = first_kv_cache.shape[-block_rank:]
-            self.block_size, kv_latent_dim = block_shape
+            block_size, kv_latent_dim = block_shape
             self.slot_size_bytes = kv_elem_size * kv_latent_dim
         else:
             # [2 (k and v), num_blocks, block_size, kv_heads, head_dim]
             self.num_blocks = first_kv_cache.shape[1]
             block_rank = 3  # [block_size, kv_heads, head_dim]
             block_shape = first_kv_cache.shape[-block_rank:]
-            self.block_size, n_kv_heads, head_dim = block_shape
+            block_size, n_kv_heads, head_dim = block_shape
             # head size in bytes.
             self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim
-
+        assert block_size == self.block_size
         # TODO(tms): self.block_len needs to be per-layer for sliding window,
         # hybrid attn, etc
         # block size in bytes