minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[BugFix][EPD] Fix Mooncake source-MR lifecycle for multi-TP /send (#32071)

Browse Source
This commit is contained in:
Zheng Wengang
2026-07-27 15:06:01 +08:00
committed by GitHub
parent 082b2a10b6
commit 3d3ba4f746
1 changed files with 19 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/encode_server.py
+19 -15
View File
@@ -1813,16 +1813,15 @@ class MMEncoder:
f"(shape={mm_data.shape}, element_size={self._element_size})" f"(shape={mm_data.shape}, element_size={self._element_size})"
) )
# MR was registered once in _run_forward and is shared across all # Request-level shared MR, registered lazily on the first /send;
# sibling-TP /send calls; # deregistration is deferred to _cleanup_inflight_encode_state.
mr_already_registered = ( fwd_state = self._forward_results.setdefault(req_id, {})
self._forward_results.get(req_id, {}).get("mr_ptr") mr_already_registered = fwd_state.get("mr_ptr") == embedding.data_ptr()
== embedding.data_ptr()
)
if not mr_already_registered: if not mr_already_registered:
self.engine.register(embedding.data_ptr(), embedding.nbytes) self.engine.register(embedding.data_ptr(), embedding.nbytes)
self._forward_results[req_id]["mr_ptr"] = embedding.data_ptr()
_t_xfer_start = time.monotonic() _t_xfer_start = time.monotonic()
await asyncio.to_thread( xfer_ret = await asyncio.to_thread(
self.engine.transfer_sync, self.engine.transfer_sync,
session_id, session_id,
embedding.data_ptr(), embedding.data_ptr(),
@@ -1834,18 +1833,20 @@ class MMEncoder:
encoder_metrics_collector.observe_transfer( encoder_metrics_collector.observe_transfer(
xfer_ms / 1000.0, backend="mooncake" xfer_ms / 1000.0, backend="mooncake"
) )
if not mr_already_registered: if xfer_ret < 0:
self.engine.deregister(embedding.data_ptr()) raise InternalError(
# Only emit at INFO when transfer is slow or fell back f"Mooncake transfer_sync failed for {req_id} "
# to per-/send register; f"(session={session_id}, nbytes={embedding.nbytes}, "
f"ret={xfer_ret})"
)
# Only emit at INFO when transfer is slow or the MR was
# registered lazily by this /send;
if xfer_ms > 200.0 or not mr_already_registered: if xfer_ms > 200.0 or not mr_already_registered:
logger.info( logger.info(
f"[{req_id}] mooncake transfer_sync={xfer_ms:.1f}ms " f"[{req_id}] mooncake transfer_sync={xfer_ms:.1f}ms "
f"nbytes={embedding.nbytes} shared_mr={mr_already_registered}" f"nbytes={embedding.nbytes} shared_mr={mr_already_registered}"
) )
mm_data.embedding = None
# Send ack/data # Send ack/data
if url is not None: if url is not None:
endpoint = NetworkAddress.parse(url).to_tcp() endpoint = NetworkAddress.parse(url).to_tcp()
@@ -2010,8 +2011,6 @@ class MMEncoder:
task.cancel() task.cancel()
# Also clean up embedding data and forward state # Also clean up embedding data and forward state
mm_data = self.embedding_to_send.pop(req_id, None) mm_data = self.embedding_to_send.pop(req_id, None)
if mm_data is not None:
mm_data.cached_embedding = None
# Release the rkey after all /send calls have completed. # Release the rkey after all /send calls have completed.
forward_state = self._forward_results.pop(req_id, None) forward_state = self._forward_results.pop(req_id, None)
if forward_state is not None: if forward_state is not None:
@@ -2023,6 +2022,11 @@ class MMEncoder:
logger.warning( logger.warning(
f"Shared-MR deregister failed for {req_id}: {dereg_err}" f"Shared-MR deregister failed for {req_id}: {dereg_err}"
) )
forward_state.pop("embedding", None)
# Release the embedding only after the MR is deregistered.
if mm_data is not None:
mm_data.embedding = None
mm_data.cached_embedding = None
self._forward_ready_events.pop(req_id, None) self._forward_ready_events.pop(req_id, None)
def _schedule_inflight_encode_cleanup(self, req_id: str): def _schedule_inflight_encode_cleanup(self, req_id: str):