[diffusion] IPC all-to-all: keep exported buffers out of expandable segments

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

@@ -30,6 +30,23 @@ class _Unsupported(RuntimeError):
     """This topology cannot run the transport -- an expected outcome, not a bug."""
 
 
+def ipc_shareable_zeros(*shape: int, dtype: torch.dtype) -> torch.Tensor:
+    """Zeros for a buffer exported over CUDA IPC, kept out of expandable segments.
+
+    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.
+    """
+    if "expandable_segments:True" not in torch._C._accelerator_getAllocatorSettings():
+        return torch.zeros(*shape, dtype=dtype, device="cuda")
+    torch._C._accelerator_setAllocatorSettings("expandable_segments:False")
+    try:
+        return torch.zeros(*shape, dtype=dtype, device="cuda")
+    finally:
+        torch._C._accelerator_setAllocatorSettings("expandable_segments:True")
+
+
 def _peer_cuda_device(group, rank: int, device: int) -> int:
     """Return the peer's local CUDA ordinal for a two-rank same-host group."""
     world_size = dist.get_world_size(group=group)
@@ -160,9 +177,9 @@ def init(self, group):
         # kernel-level dereference of peer mappings needs explicit peer access
         ctypes.CDLL("libcudart.so").cudaDeviceEnablePeerAccess(peer_dev, 0)
         self.ops = load_ipc_a2a_sync()
-        self.flag = torch.zeros(1, dtype=torch.int32, device="cuda")
+        self.flag = ipc_shareable_zeros(1, dtype=torch.int32)
         self.my_seq = torch.zeros(1, 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_flag = self._share(self.flag, group)
@@ -183,7 +200,7 @@ def get_staging(self, n_local, n_peer, dtype, group):
         if pair is None:
             if torch.cuda.is_current_stream_capturing():
                 return None
-            local = torch.zeros(2, n_local, dtype=dtype, device="cuda")
+            local = ipc_shareable_zeros(2, n_local, dtype=dtype)
             peer = self._share(local, group)
             pair = (local, peer)
             if len(self.staging) >= self.max_buffers:

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

@@ -30,6 +30,7 @@
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import (
     _Unsupported,
+    ipc_shareable_zeros,
 )
 
 logger = logging.getLogger(__name__)
@@ -253,7 +254,7 @@ class _PipelinePool:
     def __init__(self, state, dtype, in_numel: int, out_numel: int, groups: int):
         world = state.world
         self.sizes = (in_numel, out_numel, groups)
-        zeros = lambda n, dt=dtype: torch.zeros(n, dtype=dt, device="cuda")
+        zeros = lambda n, dt=dtype: ipc_shareable_zeros(n, dtype=dt)
         self.inb = state._share(zeros(in_numel), state.group)
         self.outb = state._share(zeros(out_numel), state.group)
         # fin[p]: groups peer p has delivered; fout[g, p]: calls whose group g

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

@@ -7,6 +7,7 @@
     _peer_cuda_device,
     _Unsupported,
     ipc_a2a_ready,
+    ipc_shareable_zeros,
 )
 
 _IPC = "sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a"
@@ -137,3 +138,23 @@ def test_drop_a2a_staging_buffers_clears_the_ulysses_cache():
         usp.drop_a2a_staging_buffers()
         usp.drop_a2a_staging_buffers()  # idempotent on an empty cache
     assert usp._A2A_STAGING_BUFFERS == {}
+
+
+@pytest.mark.parametrize("expandable", [False, True])
+def test_ipc_buffers_are_allocated_outside_expandable_segments(expandable):
+    settings = "expandable_segments:True" if expandable else ""
+    calls = []
+    with (
+        patch(
+            f"{_IPC}.torch._C._accelerator_getAllocatorSettings",
+            return_value=settings,
+        ),
+        patch(
+            f"{_IPC}.torch._C._accelerator_setAllocatorSettings",
+            side_effect=lambda conf: calls.append(conf),
+        ),
+        patch(f"{_IPC}.torch.zeros", side_effect=lambda *a, **k: calls.append("alloc")),
+    ):
+        ipc_shareable_zeros(2, 8, dtype=None)
+    expected = ["expandable_segments:False", "alloc", "expandable_segments:True"]
+    assert calls == (expected if expandable else ["alloc"])