Fix DP-attention SHM feature finalization race (#29543)

Co-authored-by: Yinghai Lu <yinghai@meta.com>
This commit is contained in:
Lianmin Zheng
2026-06-27 21:36:58 -07:00
committed by GitHub
co-authored by Yinghai Lu
parent ae09b8302f
commit 5747ed3b19
@@ -235,26 +235,17 @@ class SchedulerRequestReceiver:
# so that ShmPointerMMData metadata (not full tensor data) is what # so that ShmPointerMMData metadata (not full tensor data) is what
# gets serialized during broadcast_pyobj. # gets serialized during broadcast_pyobj.
if recv_reqs: if recv_reqs:
# Barrier for the non-DP-attention path only: there is a single if self.model_config.is_multimodal and has_shm_features(recv_reqs):
# broadcast_pyobj on tp_cpu_group where the source rank returns # The broadcast source returns with its original objects while
# the original objects immediately while other ranks are still in # peer ranks may still be unpickling ShmPointerMMData
# pickle.loads (-> __setstate__ -> shm_open). Without a barrier # (-> shm_open). Synchronize the same CPU groups that carried
# the source can call materialize() / shm_unlink before others # SHM-backed work requests before materialize() unlinks them.
# open the segment. recv_reqs is consistent across all ranks if self.server_args.enable_dp_attention:
# here (same broadcast), so the guard is deadlock-free. if self.ps.attn_tp_size > 1:
# barrier(group=self.attn_tp_cpu_group)
# Under DP-attention no barrier is needed: the control_reqs if self.ps.attn_cp_size > 1:
# broadcast on tp_cpu_group (step 3) is a collective that forces barrier(group=self.attn_cp_cpu_group)
# every rank to complete the earlier attn_tp / attn_cp work_reqs elif self.ps.tp_size > 1:
# deserializations (steps 1-2, which call shm_open) before any
# rank returns from step 3. POSIX guarantees shm_unlink only
# removes the name; already-open handles stay valid.
if (
not self.server_args.enable_dp_attention
and self.ps.tp_size > 1
and self.model_config.is_multimodal
and has_shm_features(recv_reqs)
):
barrier(group=self.tp_cpu_group) barrier(group=self.tp_cpu_group)
for req in recv_reqs: for req in recv_reqs:
unwrap_shm_features(req) unwrap_shm_features(req)