fix: match P/D requests by transfer_id in FlagCX connector (#315)

### PR Category
Core

### PR Type
Bug Fixes

### Description
1. In disaggregated (P/D) serving, the Prefill and Decode engines assign
**different** `request_id`s to the same logical request. The FlagCX
connector previously used `request_id` as the key for its send/receive
bookkeeping, so the Decode side's `req_id` could never be found in the
Prefill side's `reqs_need_send` map, causing KV transfers to silently
fail to match (`Request %s not found in reqs_need_send`).

2. This PR introduces a stable, cross-engine `transfer_id` that both
sides agree on, rekeys all connector bookkeeping by it, and hardens the
transfer path with proper timeouts.

3. This PR will use the FlagCX v0.13.0 interface, while also being
compatible with FlagCX v0.9.0 within CICD.

### Related Issues
<!-- Link any related issues: Fixes #issue, Closes #issue, or Related to
#issue -->

### Changes
- **`flagcx_connector.py`**: added a `TransferId` type and threaded
`transfer_id` through `RecvReqMeta`, `SendBlockMeta`, `SendReqMeta`, and
the connector metadata. `reqs_to_send` / `req_blocks` now carry
`(transfer_id, block_ids)`; the worker looks up send metadata by
`transfer_id` while still reporting completion/timeout under the Prefill
`request_id` (`p_req_id`) that vLLM expects.
- **`flagcx_connector.py`**: guarded against a missing `transfer_id` —
the scheduler logs a warning and skips the request instead of
mis-matching.
- **`flagcx_connector.py`**: replaced unbounded / hardcoded 60s waits
with bounded waits driven by `_abort_request_timeout` —
`send_meta.ready.wait()` uses a deadline and raises a clear
`RuntimeError` on timeout, and the receive socket `RCVTIMEO` now scales
with `_abort_request_timeout`.
- **`examples/disaggregated_serving_xpyd/router.py`**: inject
`transfer_id = f"fgx-{request_id}"` into `kv_transfer_params` for both
the prefill and decode legs so both engines share the same transfer key.
- **`vllm_fl/distributed/device_communicators/flagcx.py`**: fixed ctypes
usage for `flagcxGetUniqueId()` / `flagcxCommInitRank()` — pass the
unique id object directly instead of `.contents` / `ctypes.byref(...)`.
- **`vllm_fl/ops/fused_moe/fused_moe_utils.py`**: provide a default MoE
backend priority (`TRITON`, `BATCHED_TRITON`) for out-of-tree platforms.

### Testing
<!-- How has this change been tested? Include test commands, hardware
used, etc. -->
-

### Checklist
- [x] I have run the existing tests and they pass
- [x] I have added tests for my changes (if applicable)
- [x] I have updated the documentation (if applicable)

---------

Co-authored-by: ceci3 <ceci3@users.noreply.github.com>

Author: 2 people

README.md

@@ -99,7 +99,7 @@ In theory, vllm-plugin-FL can support all models available in vLLM, as long as n
     ```sh
     git clone https://github.com/flagos-ai/FlagCX.git
     cd FlagCX
-    git checkout -b v0.9.0
+    git checkout -b v0.13.0
     git submodule update --init --recursive
     ```
 

examples/disaggregated_serving_xpyd/router.py

@@ -124,7 +124,11 @@ async def lifespan(app):
 
 async def send_to_prefill(prefill_client, endpoint, req_data, request_id):
     data = req_data.copy()
-    data["kv_transfer_params"] = {"do_remote_decode": True, "do_remote_prefill": False}
+    data["kv_transfer_params"] = {
+        "do_remote_decode": True,
+        "do_remote_prefill": False,
+        "transfer_id": f"fgx-{request_id}",
+    }
     data["stream"] = False
     data["max_tokens"] = 1
     if "max_completion_tokens" in data:
@@ -153,6 +157,7 @@ async def stream_from_decode(
         "do_remote_decode": False,
         "remote_host": prefill_client["remote_host"],
         "remote_port": prefill_client["side_channel_port"],
+        "transfer_id": f"fgx-{request_id}",
     }
     headers = {"X-Request-Id": request_id}
     api_key = os.environ.get("OPENAI_API_KEY", "")

vllm_fl/distributed/device_communicators/flagcx.py

@@ -96,9 +96,17 @@ def __init__(
         self.available = True
         self.disabled = False
 
+        self._legacy_unique_id_api = (
+            hasattr(self.flagcx, "handler")
+            and not hasattr(self.flagcx, "devHandle")
+        )
+
         if self.rank == 0:
             # get the unique id from NCCL
-            self.unique_id = self.flagcx.flagcxGetUniqueId().contents
+            if self._legacy_unique_id_api:
+                self.unique_id = self.flagcx.flagcxGetUniqueId().contents
+            else:
+                self.unique_id = self.flagcx.flagcxGetUniqueId()
         else:
             # construct an empty unique id
             self.unique_id = flagcxUniqueId()
@@ -133,8 +141,12 @@ def __init__(
             device_ctx = torch.cuda.device(self.device)
 
         with device_ctx:
-            self.comm = self.flagcx.flagcxCommInitRank(
-                self.world_size, ctypes.byref(self.unique_id), self.rank)
+            if self._legacy_unique_id_api:
+                self.comm = self.flagcx.flagcxCommInitRank(
+                    self.world_size, ctypes.pointer(self.unique_id), self.rank)
+            else:
+                self.comm = self.flagcx.flagcxCommInitRank(
+                    self.world_size, self.unique_id, self.rank)
 
             stream = current_stream()
             # A small all_reduce for warmup.

vllm_fl/distributed/kv_transfer/flagcx_connector.py

@@ -74,6 +74,7 @@
 
 EngineId = str
 ReqId = str
+TransferId = str
 
 TRANS_DONE = b"trans_done"
 TRANS_ERROR = b"trans_error"
@@ -109,8 +110,7 @@ class FlagCXAgentMetadata(
     remote_port: int
     remote_tp_size: int
     remote_tp_rank: int
-    # req_id -> per KV-cache-group block ids on the Decode (receiver) side.
-    req_blocks: dict[ReqId, list[list[int]]]
+    req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]]
     kv_caches_base_addr: list[int]
     block_lens: list[int]
     kv_block_lens: list[int]
@@ -124,11 +124,14 @@ class RecvReqMeta:
     local_block_ids: list[list[int]]
     remote_host: str
     remote_port: int
+    transfer_id: TransferId
     remote_tp_size: int = 0
 
 
 @dataclass
 class SendBlockMeta:
+    p_req_id: ReqId
+    transfer_id: TransferId
     local_block_ids: list[list[int]]
     ready: threading.Event
     expire_time: float = float("inf")
@@ -138,7 +141,7 @@ class SendBlockMeta:
 
 @dataclass
 class SendReqMeta:
-    reqs: dict[ReqId, SendBlockMeta]
+    reqs: dict[TransferId, SendBlockMeta]
     lock: threading.Lock
 
 
@@ -158,7 +161,9 @@ class FlagCXConnectorMetadata(KVConnectorMetadata):
     def __init__(self):
         self.reqs_to_recv: dict[ReqId, RecvReqMeta] = {}
         # Per-request, per KV-cache-group block ids.
-        self.reqs_to_send: dict[ReqId, list[list[int]]] = {}
+        self.reqs_to_send: dict[
+            ReqId, tuple[TransferId, list[list[int]]]
+        ] = {}
 
     def add_new_req(
         self,
@@ -167,15 +172,17 @@ def add_new_req(
         kv_transfer_params: dict[str, Any],
         load_remote_cache: bool = True,
     ):
+        transfer_id = kv_transfer_params["transfer_id"]
         if load_remote_cache:
             self.reqs_to_recv[request_id] = RecvReqMeta(
                 local_block_ids=local_block_ids,
                 remote_host=kv_transfer_params["remote_host"],
                 remote_port=kv_transfer_params["remote_port"],
+                transfer_id=transfer_id,
                 remote_tp_size=kv_transfer_params.get("remote_tp_size", 0),
             )
         else:
-            self.reqs_to_send[request_id] = local_block_ids
+            self.reqs_to_send[request_id] = (transfer_id, local_block_ids)
 
 
 class FlagCXConnector(KVConnectorBase_V1, SupportsHMA):
@@ -334,7 +341,9 @@ def __init__(
         ]
 
         self._reqs_need_recv: dict[ReqId, tuple[Request, list[list[int]]]] = {}
-        self._reqs_need_send: dict[ReqId, list[list[int]]] = {}
+        self._reqs_need_send: dict[
+            ReqId, tuple["Request", list[list[int]]]
+        ] = {}
 
     def get_sw_clipped_blocks(
         self,
@@ -415,7 +424,9 @@ def update_state_after_alloc(
 
         if params.get("do_remote_prefill"):
             assert self.kv_role != "kv_producer"
-            if all(p in params for p in ("remote_host", "remote_port")):
+            if all(
+                p in params for p in ("remote_host", "remote_port", "transfer_id")
+            ):
                 unhashed_block_ids = (
                     blocks.get_unhashed_block_ids_all_groups()
                     if num_external_tokens > 0
@@ -431,7 +442,10 @@ def update_state_after_alloc(
             params["do_remote_prefill"] = False
 
         elif params.get("do_remote_decode"):
-            self._reqs_need_send[request.request_id] = []
+            if not params.get("transfer_id"):
+                logger.warning("Missing transfer_id in KVTransferParams: %s", params)
+            else:
+                self._reqs_need_send[request.request_id] = (request, [])
 
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
@@ -449,11 +463,12 @@ def build_connector_meta(
             self._reqs_need_recv.clear()
 
         if self.kv_role != "kv_consumer":
-            for req_id, block_ids in self._reqs_need_send.items():
+            for req_id, (req, block_ids) in self._reqs_need_send.items():
+                assert req.kv_transfer_params is not None
                 meta.add_new_req(
                     request_id=req_id,
                     local_block_ids=block_ids,
-                    kv_transfer_params={},
+                    kv_transfer_params=req.kv_transfer_params,
                     load_remote_cache=False,
                 )
             self._reqs_need_send.clear()
@@ -464,7 +479,7 @@ def request_finished(
         self, request: "Request", block_ids: tuple[list[int], ...]
     ) -> tuple[bool, dict[str, Any] | None]:
         params = request.kv_transfer_params
-        if not params:
+        if not params or not params.get("transfer_id"):
             return False, None
 
         if params.get("do_remote_prefill"):
@@ -483,8 +498,9 @@ def request_finished(
         delay_free_blocks = any(len(group) > 0 for group in block_ids)
 
         if delay_free_blocks:
-            self._reqs_need_send[request.request_id] = self.get_sw_clipped_blocks(
-                block_ids
+            self._reqs_need_send[request.request_id] = (
+                request,
+                self.get_sw_clipped_blocks(block_ids),
             )
 
         return delay_free_blocks, dict(
@@ -493,6 +509,7 @@ def request_finished(
             remote_host=self.side_channel_host,
             remote_port=self.side_channel_port,
             remote_tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
+            transfer_id=params["transfer_id"],
         )
 
 
@@ -920,19 +937,30 @@ def _send_kv_to_decode(self, meta: FlagCXAgentMetadata) -> None:
 
         ready_reqs: list[tuple[ReqId, SendBlockMeta]] = []
         with self.reqs_need_send.lock:
-            for req_id in meta.req_blocks:
-                send_meta = self.reqs_need_send.reqs.get(req_id)
+            for d_req_id, (transfer_id, _) in meta.req_blocks.items():
+                send_meta = self.reqs_need_send.reqs.get(transfer_id)
                 if send_meta is None:
-                    logger.warning("Request %s not found in reqs_need_send", req_id)
-                    return
+                    send_meta = SendBlockMeta(
+                        p_req_id="",
+                        transfer_id=transfer_id,
+                        local_block_ids=[],
+                        ready=threading.Event(),
+                    )
+                    self.reqs_need_send.reqs[transfer_id] = send_meta
                 # Mark it as not expired. We will send it now.
                 send_meta.expire_time = float("inf")
                 send_meta.need_send = need
-                ready_reqs.append((req_id, send_meta))
+                ready_reqs.append((d_req_id, send_meta))
 
         # Wait until the scheduler has committed each request's blocks.
-        for _, send_meta in ready_reqs:
-            send_meta.ready.wait()
+        deadline = time.perf_counter() + self._abort_request_timeout
+        for d_req_id, send_meta in ready_reqs:
+            remaining = deadline - time.perf_counter()
+            if remaining <= 0 or not send_meta.ready.wait(timeout=remaining):
+                raise RuntimeError(
+                    "Timed out waiting for prefill blocks of transfer "
+                    f"{send_meta.transfer_id} (decode request {d_req_id})."
+                )
 
         remote_session = f"{meta.remote_hostname}:{meta.remote_port}"
         conn = self._get_conn(remote_session)
@@ -980,11 +1008,15 @@ def _send_kv_to_decode(self, meta: FlagCXAgentMetadata) -> None:
 
         finished: list[ReqId] = []
         with self.reqs_need_send.lock:
-            for req_id, send_meta in ready_reqs:
+            for _, send_meta in ready_reqs:
                 send_meta.sent += 1
                 if send_meta.sent >= max(send_meta.need_send, 1):
-                    self.reqs_need_send.reqs.pop(req_id, None)
-                    finished.append(req_id)
+                    self.reqs_need_send.reqs.pop(send_meta.transfer_id, None)
+                    if not send_meta.p_req_id:
+                        raise RuntimeError(
+                            f"Missing Prefill request ID for {send_meta.transfer_id}."
+                        )
+                    finished.append(send_meta.p_req_id)
         if finished:
             with self.finished_sending_reqs.lock:
                 self.finished_sending_reqs.set.update(finished)
@@ -1011,7 +1043,7 @@ def _build_transfer_params(
         group_specs = self.kv_cache_config.kv_cache_groups
 
         for d_req_id, send_meta in ready_reqs:
-            remote_block_ids_per_group = agent_meta.req_blocks[d_req_id]
+            _, remote_block_ids_per_group = agent_meta.req_blocks[d_req_id]
             if not remote_block_ids_per_group or all(
                 len(g) == 0 for g in remote_block_ids_per_group
             ):
@@ -1129,7 +1161,7 @@ def _receiver_loop_fn(self, loop: asyncio.AbstractEventLoop):
     async def _receive_kv(
         self,
         path: str,
-        req_blocks: dict[ReqId, list[list[int]]],
+        req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]],
     ):
         req_ids = list(req_blocks.keys())
 
@@ -1152,7 +1184,9 @@ async def _receive_kv(
         sock: zmq.asyncio.Socket = make_zmq_socket(
             self.async_zmq_ctx, path, zmq.REQ, bind=False, linger=0
         )
-        sock.setsockopt(zmq.RCVTIMEO, 60000)
+        sock.setsockopt(
+            zmq.RCVTIMEO, (self._abort_request_timeout + 60) * 1000
+        )
 
         try:
             await sock.send(encoded_data)
@@ -1193,13 +1227,21 @@ def start_load_kv(self, metadata: FlagCXConnectorMetadata):
 
         if self.kv_role != "kv_consumer":
             with self.reqs_need_send.lock:
-                for req_id, block_ids in metadata.reqs_to_send.items():
-                    send_meta = self.reqs_need_send.reqs.get(req_id)
+                for p_req_id, (
+                    transfer_id,
+                    block_ids,
+                ) in metadata.reqs_to_send.items():
+                    send_meta = self.reqs_need_send.reqs.get(transfer_id)
                     if send_meta is None:
                         send_meta = SendBlockMeta(
-                            local_block_ids=[], ready=threading.Event()
+                            p_req_id=p_req_id,
+                            transfer_id=transfer_id,
+                            local_block_ids=[],
+                            ready=threading.Event(),
                         )
-                        self.reqs_need_send.reqs[req_id] = send_meta
+                        self.reqs_need_send.reqs[transfer_id] = send_meta
+                    else:
+                        send_meta.p_req_id = p_req_id
                     # Non-empty means request_finished() has committed the
                     # per-group block ids; arm the send.
                     if block_ids:
@@ -1223,7 +1265,9 @@ async def _group_kv_pull(
         loop; populates ``_pull_pending`` and launches the per-peer pulls
         without an intervening await, so it is atomic w.r.t. other coroutines.
         """
-        kv_pulls: dict[str, dict[ReqId, list[list[int]]]] = defaultdict(dict)
+        kv_pulls: dict[
+            str, dict[ReqId, tuple[TransferId, list[list[int]]]]
+        ] = defaultdict(dict)
         for req_id, meta in reqs_to_recv.items():
             remote_tp_size = meta.remote_tp_size or self.tp_size
             target_p_ranks = self.kv_topo.handshake_target_ranks(remote_tp_size)
@@ -1232,7 +1276,10 @@ async def _group_kv_pull(
                 path = make_zmq_path(
                     "tcp", meta.remote_host, meta.remote_port + p_rank
                 )
-                kv_pulls[path][req_id] = meta.local_block_ids
+                kv_pulls[path][req_id] = (
+                    meta.transfer_id,
+                    meta.local_block_ids,
+                )
         for path, req_blocks in kv_pulls.items():
             asyncio.ensure_future(self._receive_kv(path, req_blocks))
 
@@ -1261,15 +1308,17 @@ def get_finished(self) -> tuple[set[str] | None, set[str] | None]:
         now = time.perf_counter()
         with self.reqs_need_send.lock:
             expired = [
-                rid
-                for rid, sm in self.reqs_need_send.reqs.items()
+                transfer_id
+                for transfer_id, sm in self.reqs_need_send.reqs.items()
                 if sm.expire_time < now
             ]
-            for rid in expired:
-                logger.warning("Request %s send timed out, freeing blocks", rid)
-                del self.reqs_need_send.reqs[rid]
-            if expired:
-                finished_sending.update(expired)
+            for transfer_id in expired:
+                send_meta = self.reqs_need_send.reqs.pop(transfer_id)
+                logger.warning(
+                    "Transfer %s send timed out, freeing blocks", transfer_id
+                )
+                if send_meta.p_req_id:
+                    finished_sending.add(send_meta.p_req_id)
 
         return finished_sending or None, finished_recving or None
 

vllm_fl/ops/fused_moe/fused_moe_utils.py

@@ -85,6 +85,11 @@ def _move_to_back(
         _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU]
     elif current_platform.is_cpu():
         _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU]
+    elif current_platform.is_out_of_tree():
+        _AVAILABLE_BACKENDS = [
+            UnquantizedMoeBackend.TRITON,
+            UnquantizedMoeBackend.BATCHED_TRITON,
+        ]
     return _AVAILABLE_BACKENDS
 
 ## Adopt from select_unquantized_moe_backend