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,27 +235,18 @@ class SchedulerRequestReceiver:
# so that ShmPointerMMData metadata (not full tensor data) is what
# gets serialized during broadcast_pyobj.
if recv_reqs:
# Barrier for the non-DP-attention path only: there is a single
# broadcast_pyobj on tp_cpu_group where the source rank returns
# the original objects immediately while other ranks are still in
# pickle.loads (-> __setstate__ -> shm_open). Without a barrier
# the source can call materialize() / shm_unlink before others
# open the segment. recv_reqs is consistent across all ranks
# here (same broadcast), so the guard is deadlock-free.
#
# Under DP-attention no barrier is needed: the control_reqs
# broadcast on tp_cpu_group (step 3) is a collective that forces
# every rank to complete the earlier attn_tp / attn_cp work_reqs
# 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)
if self.model_config.is_multimodal and has_shm_features(recv_reqs):
# The broadcast source returns with its original objects while
# peer ranks may still be unpickling ShmPointerMMData
# (-> shm_open). Synchronize the same CPU groups that carried
# SHM-backed work requests before materialize() unlinks them.
if self.server_args.enable_dp_attention:
if self.ps.attn_tp_size > 1:
barrier(group=self.attn_tp_cpu_group)
if self.ps.attn_cp_size > 1:
barrier(group=self.attn_cp_cpu_group)
elif self.ps.tp_size > 1:
barrier(group=self.tp_cpu_group)
for req in recv_reqs:
unwrap_shm_features(req)