flip protocol; fix scheduling order bug

Signed-off-by: Will Eaton <weaton@redhat.com>

Author: wseaton

vllm/config.py

@@ -1515,7 +1515,7 @@ class CacheConfig:
 
     transfer_handshake_metadata: Optional[dict[int, dict[
         int, KVConnectorHandshakeMetadata]]] = field(default=None, init=False)
-    """Metadata for the KV connector handshake."""
+    """Metadata for the KV connector handshake. Structure: dp_rank -> tp_rank -> metadata"""
 
     def compute_hash(self) -> str:
         """

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

@@ -489,8 +489,9 @@ def _nixl_handshake(self, host: str, port: int):
         start_time = time.perf_counter()
         logger.debug("Starting NIXL handshake with %s:%s", host, port)
 
-        # TODO: make the scheme dynamic, and/or implement https on both sides.
-        url = build_uri("http", host, port, path="get_kv_connector_metadata")
+        # Use the new endpoint scheme to filter by dp_rank and tp_rank
+        # Default to dp_rank 0 and use current tp_rank for optimal filtering
+        url = build_uri("http", host, port, path=f"get_kv_connector_metadata/0/{self.tp_rank}")
         logger.debug("Querying metadata on path: %s", url)
 
         try:
@@ -509,20 +510,29 @@ def _nixl_handshake(self, host: str, port: int):
             logger.warning("Remote server returned None metadata, skipping handshake")
             raise RuntimeError("Remote server returned None metadata")
 
-        remote_tp_size = len(res.keys())
-        # Default case is that the remote TP size is 1, so we can
-        # directly access the metadata.
-        tp_data = res.get(str(self.tp_rank), {}).get("0", {})
-        metadata_bytes = tp_data.get("agent_metadata", None)
-
-        # 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.
-        tp_ratio = self._tp_size[self.engine_id] // remote_tp_size
-        p_remote_rank = self.tp_rank // tp_ratio
-        if p_remote_rank > 0:
-            metadata_bytes = res.get(str(p_remote_rank),
-                                     {}).get("0",
-                                             {}).get("agent_metadata", None)
+        # With filtered response from new endpoint, we get: {dp_rank: {tp_rank: metadata}}
+        # Since we filtered by dp_rank=0 and tp_rank=self.tp_rank, extract directly
+        if "0" in res and str(self.tp_rank) in res["0"]:
+            tp_data = res["0"][str(self.tp_rank)]
+            metadata_bytes = tp_data.get("agent_metadata", None)
+            p_remote_rank = self.tp_rank  # Use current tp_rank for filtered response
+        else:
+            # Fallback to unfiltered endpoint for heterogeneous TP cases
+            url_fallback = build_uri("http", host, port, path="get_kv_connector_metadata")
+            logger.debug("Using fallback unfiltered endpoint: %s", url_fallback)
+            req = Request(url_fallback)
+            with urlopen(req, timeout=5.0) as response:
+                response_data = response.read().decode('utf-8')
+                res = json.loads(response_data)
+            
+            dp_data = res.get("0", {})
+            remote_tp_size = len(dp_data.keys()) if dp_data else 1
+            
+            # Handle heterogeneous TP mapping
+            tp_ratio = self._tp_size[self.engine_id] // remote_tp_size
+            p_remote_rank = self.tp_rank // tp_ratio
+            tp_data = dp_data.get(str(p_remote_rank), {})
+            metadata_bytes = tp_data.get("agent_metadata", None)
 
         if metadata_bytes is not None:
             # Reconstruct NixlAgentMetadata from JSON response
@@ -962,6 +972,7 @@ def start_load_kv(self, metadata: NixlConnectorMetadata):
         Start loading by triggering non-blocking nixl_xfer.
         We check for these trnxs to complete in each step().
         """
+        
         for req_id, meta in metadata.requests.items():
             logger.debug(
                 "start_load_kv for request %s from remote engine %s. "

vllm/entrypoints/openai/api_server.py

@@ -869,8 +869,29 @@ async def show_server_info(raw_request: Request):
         return JSONResponse(content=server_info)
 
     @router.get("/get_kv_connector_metadata")
-    async def get_kv_connector_metadata(raw_request: Request):
+    @router.get("/get_kv_connector_metadata/{dp_rank}")
+    @router.get("/get_kv_connector_metadata/{dp_rank}/{tp_rank}")
+    async def get_kv_connector_metadata(raw_request: Request, dp_rank: int = None, tp_rank: int = None):
         kv_connector_metadata = raw_request.app.state.vllm_config.cache_config.transfer_handshake_metadata
+        
+        if kv_connector_metadata is None:
+            return JSONResponse(content=None)
+        
+        # Filter by dp_rank if specified
+        if dp_rank is not None:
+            if dp_rank not in kv_connector_metadata:
+                return JSONResponse(content={})
+            dp_data = kv_connector_metadata[dp_rank]
+            
+            # Filter by tp_rank if also specified
+            if tp_rank is not None:
+                if tp_rank not in dp_data:
+                    return JSONResponse(content={})
+                return JSONResponse(content={dp_rank: {tp_rank: dp_data[tp_rank]}})
+            else:
+                return JSONResponse(content={dp_rank: dp_data})
+        
+        # Return all metadata if no filtering
         return JSONResponse(content=kv_connector_metadata)
 
     @router.post("/reset_prefix_cache")

vllm/v1/core/sched/scheduler.py

@@ -719,11 +719,7 @@ def update_from_output(
         for request in self.running:
             req_id = request.request_id
             num_tokens_scheduled = num_scheduled_tokens.get(req_id, 0)
-            if num_tokens_scheduled == 0:
-                # The request was not scheduled in this step.
-                new_running.append(request)
-                continue
-
+            
             # Check if this request is pending handshake and needs to reschedule
             if (pending_handshake_req_ids
                     and req_id in pending_handshake_req_ids):
@@ -734,6 +730,11 @@ def update_from_output(
                 num_tokens_to_reschedule -= request.num_computed_tokens
                 new_running.append(request)
                 continue
+            
+            if num_tokens_scheduled == 0:
+                # The request was not scheduled in this step.
+                new_running.append(request)
+                continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index]

vllm/v1/engine/core.py

@@ -458,7 +458,11 @@ def _perform_handshake(
                 content = {}
                 for worker_dict in self.transfer_handshake_metadata:
                     if worker_dict is not None:
-                        content.update(worker_dict)
+                        # Deep merge the nested dictionaries instead of overwriting
+                        for dp_rank, tp_dict in worker_dict.items():
+                            if dp_rank not in content:
+                                content[dp_rank] = {}
+                            content[dp_rank].update(tp_dict)
                 handshake_message["transfer_handshake_metadata"] = content
 
             handshake_socket.send(msgspec.msgpack.encode(handshake_message))

vllm/v1/utils.py

@@ -543,10 +543,10 @@ def wait_for_engine_startup(
                     eng_index, txfer_metadata)
                 if cache_config.transfer_handshake_metadata is None:
                     cache_config.transfer_handshake_metadata = defaultdict(dict)
-                for tp_rank, dp_dict in txfer_metadata.items():
-                    for dp_rank, metadata in dp_dict.items():
-                        cache_config.transfer_handshake_metadata[tp_rank][
-                            dp_rank] = metadata
+                for dp_rank, tp_dict in txfer_metadata.items():
+                    for tp_rank, metadata in tp_dict.items():
+                        cache_config.transfer_handshake_metadata[dp_rank][
+                            tp_rank] = metadata
 
             start_pending[0 if local else 1] -= 1
             engine.state = CoreEngineState.READY

vllm/v1/worker/gpu_worker.py

@@ -252,7 +252,7 @@ def get_kv_connector_handshake_metadata(self) -> Optional[dict]:
 
         tp_rank = get_tp_group().rank_in_group
         dp_rank = self.vllm_config.parallel_config.data_parallel_rank_local
-        return {tp_rank: {dp_rank: msgspec.to_builtins(metadata)}}
+        return {dp_rank: {tp_rank: msgspec.to_builtins(metadata)}}
 
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         return self.model_runner.get_kv_cache_spec()