[Bug][VLM] Fix shared memory race condition in ShmPointerMMData broadcast for multi-GPU VLM serving (#21655)
This commit is contained in:
@@ -1725,6 +1725,19 @@ def wrap_shm_features(obj):
|
|||||||
return obj
|
return obj
|
||||||
|
|
||||||
|
|
||||||
|
def has_shm_features(recv_reqs):
|
||||||
|
"""Return True if any request in the list contains ShmPointerMMData."""
|
||||||
|
for req in recv_reqs:
|
||||||
|
if hasattr(req, "batch"):
|
||||||
|
if has_shm_features(req.batch):
|
||||||
|
return True
|
||||||
|
elif hasattr(req, "mm_inputs") and req.mm_inputs:
|
||||||
|
for item in req.mm_inputs.get("mm_items", []):
|
||||||
|
if isinstance(item.feature, ShmPointerMMData):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def unwrap_shm_features(obj):
|
def unwrap_shm_features(obj):
|
||||||
"""
|
"""
|
||||||
Restore ShmPointerMMData wrappers back into standard torch.Tensors.
|
Restore ShmPointerMMData wrappers back into standard torch.Tensors.
|
||||||
|
|||||||
@@ -142,7 +142,11 @@ from sglang.srt.managers.io_struct import (
|
|||||||
UpdateWeightsFromIPCReqInput,
|
UpdateWeightsFromIPCReqInput,
|
||||||
UpdateWeightsFromTensorReqInput,
|
UpdateWeightsFromTensorReqInput,
|
||||||
)
|
)
|
||||||
from sglang.srt.managers.mm_utils import init_mm_embedding_cache, unwrap_shm_features
|
from sglang.srt.managers.mm_utils import (
|
||||||
|
has_shm_features,
|
||||||
|
init_mm_embedding_cache,
|
||||||
|
unwrap_shm_features,
|
||||||
|
)
|
||||||
from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
|
from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
|
||||||
from sglang.srt.managers.overlap_utils import FutureMap
|
from sglang.srt.managers.overlap_utils import FutureMap
|
||||||
from sglang.srt.managers.prefill_delayer import (
|
from sglang.srt.managers.prefill_delayer import (
|
||||||
@@ -1501,6 +1505,27 @@ class Scheduler(
|
|||||||
# 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
|
||||||
|
# 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.tp_size > 1
|
||||||
|
and self.model_config.is_multimodal
|
||||||
|
and has_shm_features(recv_reqs)
|
||||||
|
):
|
||||||
|
barrier(group=self.tp_cpu_group)
|
||||||
for req in recv_reqs:
|
for req in recv_reqs:
|
||||||
unwrap_shm_features(req)
|
unwrap_shm_features(req)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user