diff --git a/examples/disaggregated_serving_xpyd/router.py b/examples/disaggregated_serving_xpyd/router.py index b836550d..8b2b2822 100644 --- a/examples/disaggregated_serving_xpyd/router.py +++ b/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", "") diff --git a/vllm_fl/distributed/device_communicators/flagcx.py b/vllm_fl/distributed/device_communicators/flagcx.py index ca641e2e..10bb6943 100644 --- a/vllm_fl/distributed/device_communicators/flagcx.py +++ b/vllm_fl/distributed/device_communicators/flagcx.py @@ -98,7 +98,7 @@ def __init__( if self.rank == 0: # get the unique id from NCCL - self.unique_id = self.flagcx.flagcxGetUniqueId().contents + self.unique_id = self.flagcx.flagcxGetUniqueId() else: # construct an empty unique id self.unique_id = flagcxUniqueId() @@ -134,7 +134,7 @@ def __init__( with device_ctx: self.comm = self.flagcx.flagcxCommInitRank( - self.world_size, ctypes.byref(self.unique_id), self.rank) + self.world_size, self.unique_id, self.rank) stream = current_stream() # A small all_reduce for warmup. diff --git a/vllm_fl/distributed/kv_transfer/flagcx_connector.py b/vllm_fl/distributed/kv_transfer/flagcx_connector.py index 5e0641cb..61ebf37b 100644 --- a/vllm_fl/distributed/kv_transfer/flagcx_connector.py +++ b/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 diff --git a/vllm_fl/ops/fused_moe/fused_moe_utils.py b/vllm_fl/ops/fused_moe/fused_moe_utils.py index 43d529fc..f494a6cb 100644 --- a/vllm_fl/ops/fused_moe/fused_moe_utils.py +++ b/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