[Fix] diffusion: Keep the N-rank IPC all-to-all's flags and staging out of expandable segments

python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py

@@ -36,7 +36,8 @@ def ipc_shareable_zeros(*shape: int, dtype: torch.dtype) -> torch.Tensor:
     Freeing an expandable segment that was exported over IPC breaks the CUDA
     allocator on fabric-handle systems (GB200/GB300): empty_cache raises
     "std::get: wrong index for variant" or "bad optional access" on some ranks,
-    which then diverge and hang the next collective.
+    which then diverge and hang the next collective. Importing an fd-handle
+    segment needs pidfd_getfd, which containers without ptrace rights deny.
     """
     if "expandable_segments:True" not in torch._C._accelerator_getAllocatorSettings():
         return torch.zeros(*shape, dtype=dtype, device="cuda")

python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a_multi.py

@@ -444,13 +444,18 @@ def init(self, group) -> None:
         self.ops = load_ipc_a2a_sync()
         self.memops = memops
         # flags[p] is written by peer p; seq[p] counts the calls sent to p
-        self.flags = torch.zeros(self.world, dtype=torch.int32, device="cuda")
+        self.flags = ipc_shareable_zeros(self.world, dtype=torch.int32)
         self.seq = torch.zeros(self.world, dtype=torch.int32, device="cuda")
-        self.timed_out = torch.zeros(1, dtype=torch.int32, device="cuda")
+        self.timed_out = ipc_shareable_zeros(1, dtype=torch.int32)
         self.budget_ns = int(envs.SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS * 1e6)
         self.max_buffers = envs.SGLANG_DIFFUSION_IPC_A2A_MAX_BUFFERS
-        self.peer_flags = self._share(self.flags, group)
-        self.peer_timed_out = self._share(self.timed_out, group)
+        try:
+            self.peer_flags = self._share(self.flags, group)
+            self.peer_timed_out = self._share(self.timed_out, group)
+        except RuntimeError as e:
+            raise RuntimeError(
+                f"sharing the signal flags over CUDA IPC failed: {e}"
+            ) from e
         self.inited = True
 
     def get_staging(self, chunk: int, dtype: torch.dtype) -> list[torch.Tensor] | None:
@@ -462,7 +467,7 @@ def get_staging(self, chunk: int, dtype: torch.dtype) -> list[torch.Tensor] | No
         if buffers is None:
             if torch.cuda.is_current_stream_capturing():
                 return None
-            local = torch.zeros(2, self.world, chunk, dtype=dtype, device="cuda")
+            local = ipc_shareable_zeros(2, self.world, chunk, dtype=dtype)
             buffers = self._share(local, self.group)
             self._insert(key, buffers)
         return buffers

python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_multi_gpu.py

@@ -10,6 +10,8 @@
 
 from __future__ import annotations
 
+import ctypes
+import errno
 import os
 import subprocess
 import sys
@@ -45,6 +47,35 @@ def _attend_rows(q, k, v):
     ).to(q.dtype)
 
 
+def _deny_pidfd_getfd() -> None:
+    """Fail pidfd_getfd with EPERM in this process, as containers without ptrace
+    rights do; importing an expandable-segment IPC handle needs it."""
+
+    class Insn(ctypes.Structure):  # struct sock_filter
+        _fields_ = [
+            ("code", ctypes.c_uint16),
+            ("jt", ctypes.c_uint8),
+            ("jf", ctypes.c_uint8),
+            ("k", ctypes.c_uint32),
+        ]
+
+    class Prog(ctypes.Structure):  # struct sock_fprog
+        _fields_ = [("len", ctypes.c_uint16), ("filter", ctypes.POINTER(Insn))]
+
+    insns = (Insn * 4)(
+        Insn(0x20, 0, 0, 0),  # load the syscall number
+        Insn(0x15, 0, 1, 438),  # pidfd_getfd, the same number on every arch
+        Insn(0x06, 0, 0, 0x00050000 | errno.EPERM),  # SECCOMP_RET_ERRNO
+        Insn(0x06, 0, 0, 0x7FFF0000),  # SECCOMP_RET_ALLOW
+    )
+    libc = ctypes.CDLL(None, use_errno=True)
+    # PR_SET_NO_NEW_PRIVS, then PR_SET_SECCOMP with SECCOMP_MODE_FILTER
+    if libc.prctl(38, 1, 0, 0, 0) or libc.prctl(
+        22, 2, ctypes.byref(Prog(len(insns), insns)), 0, 0
+    ):
+        raise OSError(ctypes.get_errno(), "could not install the seccomp filter")
+
+
 def _worker() -> int:
     import torch.distributed as dist
 
@@ -326,7 +357,7 @@ def slow_attend(q, k, v):
 
 
 class TestIpcA2AMultiGpu(CustomTestCase):
-    def _run(self, world: int, port: int):
+    def _run(self, world: int, port: int, *worker_args: str, **env: str):
         if not current_platform.is_cuda():
             self.skipTest("CUDA-IPC transport is unavailable on this platform")
         if torch.cuda.device_count() < world:
@@ -340,10 +371,12 @@ def _run(self, world: int, port: int):
                 f"--master-port={port}",
                 __file__,
                 "--worker",
+                *worker_args,
             ],
             capture_output=True,
             text=True,
             timeout=1200,
+            env={**os.environ, **env},
         )
         print(proc.stdout[-4000:])
         if proc.returncode != 0:
@@ -358,8 +391,18 @@ def test_two_ranks_match_sequential_bitwise(self):
     def test_four_ranks_match_sequential_bitwise(self):
         self._run(4, 29537)
 
+    def test_two_ranks_engage_under_expandable_segments_without_pidfd_getfd(self):
+        self._run(
+            2,
+            29547,
+            "--deny-pidfd-getfd",
+            PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True",
+        )
+
 
 if __name__ == "__main__":
     if "--worker" in sys.argv:
+        if "--deny-pidfd-getfd" in sys.argv:
+            _deny_pidfd_getfd()
         raise SystemExit(_worker())
     unittest.main()

python/sglang/multimodal_gen/test/unit/test_ipc_a2a_lifecycle.py

@@ -1,6 +1,7 @@
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
+import torch
 
 from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import (
     IpcA2AState,
@@ -9,9 +10,14 @@
     ipc_a2a_ready,
     ipc_shareable_zeros,
 )
+from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a_multi import (
+    IpcA2AMultiState,
+    _PipelinePool,
+)
 from sglang.multimodal_gen.runtime.platforms import current_platform
 
 _IPC = "sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a"
+_MULTI = f"{_IPC}_multi"
 
 # ipc_a2a_ready() only sets up the transport on CUDA. Elsewhere the platform may
 # lack is_full_nvlink and torch may lack the accelerator allocator-settings API.
@@ -167,3 +173,42 @@ def test_ipc_buffers_are_allocated_outside_expandable_segments(expandable):
         ipc_shareable_zeros(2, 8, dtype=None)
     expected = ["expandable_segments:False", "alloc", "expandable_segments:True"]
     assert calls == (expected if expandable else ["alloc"])
+
+
+def test_multi_rank_transport_exports_only_ipc_shareable_buffers():
+    # under expandable segments a plain allocation exports a handle whose import
+    # needs pidfd_getfd, so one such buffer sends the whole group to NCCL
+    state = IpcA2AMultiState()
+    shareable, exported = [], []
+
+    def shareable_zeros(*shape, dtype):
+        shareable.append(MagicMock())
+        return shareable[-1]
+
+    def share(t, group):
+        exported.append(t)
+        return [t, t]
+
+    agreed = MagicMock()
+    agreed.item.return_value = 1
+    with (
+        patch(f"{_MULTI}.ipc_shareable_zeros", side_effect=shareable_zeros),
+        patch(f"{_MULTI}.torch.zeros"),
+        patch(f"{_MULTI}.torch.tensor", return_value=agreed),
+        patch(f"{_MULTI}.torch.cuda.current_device", return_value=0),
+        patch(f"{_MULTI}.torch.cuda.can_device_access_peer", return_value=True),
+        patch(f"{_MULTI}.torch.cuda.is_current_stream_capturing", return_value=False),
+        patch(f"{_MULTI}.dist.get_rank", return_value=0),
+        patch(f"{_MULTI}.dist.get_world_size", return_value=2),
+        patch(f"{_MULTI}.dist.all_reduce"),
+        patch(f"{_MULTI}._member_devices", return_value=[0, 1]),
+        patch(f"{_MULTI}._StreamMemOps"),
+        patch(f"{_MULTI}.load_ipc_a2a_sync"),
+        patch("ctypes.CDLL"),
+        patch.object(state, "_share", side_effect=share),
+    ):
+        state.init(object())
+        state.get_staging(8, torch.bfloat16)
+        _PipelinePool(state, torch.bfloat16, 8, 8, 2)
+    # flags, timeout flag, staging, pipeline receive/output boxes and counters
+    assert [any(t is s for s in shareable) for t in exported] == [True] * 7