Fix flashinfer workspace OOM (#24172)

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -18,7 +18,11 @@ from sglang.srt.distributed import (
     get_tp_group,
 )
 from sglang.srt.environ import envs
-from sglang.srt.utils import is_flashinfer_available
+from sglang.srt.utils import (
+    ceil_align,
+    get_cuda_driver_bindings,
+    is_flashinfer_available,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
 logger = logging.getLogger(__name__)
@@ -157,6 +161,175 @@ def is_flashinfer_allreduce_unavailable() -> bool:
     return _flashinfer_allreduce_unavailable
 
 
+def _make_flashinfer_workspace_allocation_prop(cuda_driver):
+    if _should_force_posix_fd_transport():
+        handle_type = (
+            cuda_driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
+        )
+    else:
+        from flashinfer.comm.mnnvl import is_mnnvl_fabric_supported
+
+        if is_mnnvl_fabric_supported(torch.cuda.current_device()):
+            handle_type = (
+                cuda_driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC
+            )
+        else:
+            handle_type = (
+                cuda_driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
+            )
+
+    prop = cuda_driver.CUmemAllocationProp()
+    prop.requestedHandleTypes = handle_type
+    prop.type = cuda_driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
+    prop.location = cuda_driver.CUmemLocation()
+    prop.location.type = cuda_driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+    prop.location.id = torch.cuda.current_device()
+    prop.allocFlags.gpuDirectRDMACapable = 1
+    return prop
+
+
+def _flashinfer_trtllm_workspace_allocation_sizes(
+    cuda_driver,
+    prop,
+    world_size: int,
+    max_token_num: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+) -> list[int]:
+    """Mirror FlashInfer TRTLLM SymmDeviceMemory local allocation sizes."""
+    elem_size = 4 if dtype == torch.float32 else 2
+    buffer_size = world_size * max_token_num * hidden_dim * 2
+    flag_size = world_size * 256 * 4
+
+    max_comm_size = 2147483647 & ~((1 << 21) - 1)
+    lamport_comm_size = min(
+        world_size * max_token_num * hidden_dim * elem_size,
+        max_comm_size,
+    )
+    lamport_buffer_size = lamport_comm_size * 3
+
+    # trtllm_create_ipc_workspace_for_all_reduce_fusion rounds each logical
+    # buffer to 2 MiB before passing it to SymmDeviceMemory.
+    buffer_sizes = (
+        ceil_align(size, 1 << 21)
+        for size in (buffer_size, flag_size, lamport_buffer_size)
+    )
+
+    signal_pad_size = 2048
+    allocation_sizes = []
+    for buffer_size in buffer_sizes:
+        err, alloc_granularity = cuda_driver.cuMemGetAllocationGranularity(
+            prop,
+            cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED,
+        )
+        if err != cuda_driver.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(
+                "cuMemGetAllocationGranularity failed for FlashInfer "
+                f"workspace preflight: {err}"
+            )
+
+        allocation_size = ceil_align(buffer_size + signal_pad_size, alloc_granularity)
+
+        mc_prop = cuda_driver.CUmulticastObjectProp()
+        mc_prop.numDevices = world_size
+        mc_prop.size = allocation_size
+        mc_prop.handleTypes = prop.requestedHandleTypes
+
+        err, mc_granularity = cuda_driver.cuMulticastGetGranularity(
+            mc_prop,
+            cuda_driver.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED,
+        )
+        if err != cuda_driver.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(
+                "cuMulticastGetGranularity failed for FlashInfer "
+                f"workspace preflight: {err}"
+            )
+
+        allocation_size = ceil_align(allocation_size, mc_granularity)
+        allocation_sizes.append(allocation_size)
+    return allocation_sizes
+
+
+def _probe_cumem_create_sequence(cuda_driver, allocation_sizes, prop) -> bool:
+    handles = []
+    try:
+        for allocation_size in allocation_sizes:
+            err, handle = cuda_driver.cuMemCreate(allocation_size, prop, 0)
+            if err != cuda_driver.CUresult.CUDA_SUCCESS:
+                return False
+            handles.append(handle)
+        return True
+    finally:
+        for handle in reversed(handles):
+            cuda_driver.cuMemRelease(handle)
+
+
+def _preflight_check_workspace_memory(
+    world_size: int,
+    max_token_num: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+    cpu_group: Optional["torch.distributed.ProcessGroup"] = None,
+) -> bool:
+    """Collectively decide whether to enter FlashInfer workspace creation.
+
+    FlashInfer TRTLLM workspaces allocate several SymmDeviceMemory buffers and
+    then exchange handles across ranks. If one rank fails local cuMemCreate and
+    exits while peers enter handle exchange, peers can hang until the watchdog
+    aborts. Probe the same handle type and allocation sequence first, then vote
+    on a CPU group so all ranks proceed or skip together.
+    """
+    import torch.distributed as dist
+
+    group = cpu_group
+    if group is None:
+        tp_group = get_tp_group()
+        if tp_group.world_size <= 1:
+            return True
+        group = tp_group.cpu_group
+
+    allocation_sizes = []
+    try:
+        cuda_driver = get_cuda_driver_bindings()
+        prop = _make_flashinfer_workspace_allocation_prop(cuda_driver)
+        allocation_sizes = _flashinfer_trtllm_workspace_allocation_sizes(
+            cuda_driver,
+            prop,
+            world_size,
+            max_token_num,
+            hidden_dim,
+            dtype,
+        )
+        local_ok = _probe_cumem_create_sequence(cuda_driver, allocation_sizes, prop)
+    except Exception as e:
+        logger.warning(
+            "FlashInfer workspace preflight probe failed (%s). "
+            "Skipping allreduce fusion.",
+            e,
+        )
+        local_ok = False
+
+    flag = torch.tensor([1 if local_ok else 0], dtype=torch.int32)
+    dist.all_reduce(flag, op=dist.ReduceOp.BAND, group=group)
+
+    logger.debug(
+        "FlashInfer workspace preflight [rank %s]: probe=%.2f GB, "
+        "local_probe=%s, vote=%s",
+        dist.get_rank(group=group),
+        sum(allocation_sizes) / 1e9,
+        "OK" if local_ok else "FAIL",
+        "PROCEED" if flag.item() == 1 else "SKIP",
+    )
+    if flag.item() == 0:
+        logger.warning(
+            "FlashInfer workspace preflight: cuMemCreate probe failed on at "
+            "least one rank. Skipping allreduce fusion to avoid cross-rank "
+            "desync inside the flashinfer collective."
+        )
+        return False
+    return True
+
+
 class FlashInferWorkspaceManager:
     def __init__(self):
         self.workspace = None
@@ -187,6 +360,20 @@ class FlashInferWorkspaceManager:
             return
 
         self.cleanup()
+
+        global _flashinfer_allreduce_unavailable
+        if not _preflight_check_workspace_memory(
+            world_size=world_size,
+            max_token_num=max_token_num,
+            hidden_dim=hidden_dim,
+            dtype=dtype,
+            cpu_group=cpu_group,
+        ):
+            _flashinfer_allreduce_unavailable = True
+            self.workspace = None
+            self.initialized = False
+            return
+
         try:
             kwargs = dict(
                 backend="trtllm",
@@ -210,7 +397,6 @@ class FlashInferWorkspaceManager:
                     **kwargs
                 )
         except Exception as e:
-            global _flashinfer_allreduce_unavailable
             _flashinfer_allreduce_unavailable = True
             logger.warning(
                 f"Failed to initialize FlashInfer workspace: {e}. "

python/sglang/srt/utils/common.py

@@ -3650,6 +3650,15 @@ def check_cuda_result(raw_output):
     return results
 
 
+def get_cuda_driver_bindings():
+    try:
+        from cuda.bindings import driver as cuda_driver
+    except ImportError:
+        from cuda import cuda as cuda_driver
+
+    return cuda_driver
+
+
 def get_physical_device_id(pytorch_device_id: int) -> int:
     """
     Convert PyTorch logical device ID to physical device ID.

test/registered/distributed/test_flashinfer_fusion_preflight.py

@@ -0,0 +1,163 @@
+"""Distributed tests for FlashInfer allreduce-fusion workspace preflight."""
+
+import multiprocessing as mp
+import os
+import socket
+import unittest
+
+import torch
+
+from sglang.srt.utils import get_cuda_driver_bindings, is_flashinfer_available
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, suite="stage-b-test-2-gpu-large")
+
+WORLD_SIZE = 2
+
+
+def _get_free_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _run_rank(rank, world_size, port, scenario, result_q):
+    held = None
+    cuda_driver = None
+    try:
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = str(port)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+        os.environ["LOCAL_RANK"] = str(rank)
+
+        torch.cuda.set_device(rank)
+
+        import torch.distributed as dist
+
+        dist.init_process_group(
+            backend="gloo",
+            rank=rank,
+            world_size=world_size,
+        )
+        cpu_group = dist.group.WORLD
+
+        from sglang.srt.layers.flashinfer_comm_fusion import (
+            _make_flashinfer_workspace_allocation_prop,
+            _preflight_check_workspace_memory,
+        )
+
+        probe_kwargs = dict(
+            world_size=8,
+            max_token_num=2048,
+            hidden_dim=12288,
+            dtype=torch.bfloat16,
+            cpu_group=cpu_group,
+        )
+
+        if scenario == "rank0_starved" and rank == 0:
+            cuda_driver = get_cuda_driver_bindings()
+            prop = _make_flashinfer_workspace_allocation_prop(cuda_driver)
+
+            free, _total = torch.cuda.mem_get_info(rank)
+            target = max(free - (1 << 30), 0)
+            granularity_flag = (
+                cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED
+            )
+            err, gran = cuda_driver.cuMemGetAllocationGranularity(
+                prop,
+                granularity_flag,
+            )
+            assert err == cuda_driver.CUresult.CUDA_SUCCESS, err
+            aligned = (target // gran) * gran
+            assert aligned > 0, "not enough free memory to starve the preflight"
+            err, held = cuda_driver.cuMemCreate(aligned, prop, 0)
+            assert err == cuda_driver.CUresult.CUDA_SUCCESS, (err, aligned)
+
+        decision = _preflight_check_workspace_memory(**probe_kwargs)
+        result_q.put((rank, "ok", bool(decision)))
+    except Exception as e:  # pragma: no cover - debug path
+        result_q.put((rank, "err", repr(e)))
+    finally:
+        if held is not None:
+            cuda_driver.cuMemRelease(held)
+        try:
+            import torch.distributed as dist
+
+            if dist.is_initialized():
+                dist.destroy_process_group()
+        except Exception:
+            pass
+
+
+def _spawn_and_collect(scenario, world_size=WORLD_SIZE):
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+    port = _get_free_port()
+    procs = []
+    for rank in range(world_size):
+        proc = ctx.Process(
+            target=_run_rank,
+            args=(rank, world_size, port, scenario, q),
+        )
+        proc.start()
+        procs.append(proc)
+
+    try:
+        results = {}
+        for _ in range(world_size):
+            rank, status, payload = q.get(timeout=300)
+            results[rank] = (status, payload)
+
+        for proc in procs:
+            proc.join(timeout=60)
+            assert proc.exitcode == 0, f"rank exited with {proc.exitcode}"
+    finally:
+        for proc in procs:
+            if proc.is_alive():
+                proc.terminate()
+                proc.join(timeout=10)
+
+    return results
+
+
+class TestFlashInferPreflightDistributed(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available() or torch.cuda.device_count() < WORLD_SIZE:
+            raise unittest.SkipTest(
+                f"Need {WORLD_SIZE} CUDA devices, got {torch.cuda.device_count()}"
+            )
+        if not is_flashinfer_available():
+            raise unittest.SkipTest("FlashInfer is not available")
+        try:
+            from sglang.srt.layers.flashinfer_comm_fusion import (
+                _make_flashinfer_workspace_allocation_prop,
+            )
+
+            cuda_driver = get_cuda_driver_bindings()
+            _make_flashinfer_workspace_allocation_prop(cuda_driver)
+        except Exception as e:
+            raise unittest.SkipTest(
+                f"FlashInfer preflight dependencies unavailable: {e}"
+            )
+
+    def test_happy_path_votes_proceed(self):
+        results = _spawn_and_collect("normal")
+        for rank, (status, payload) in results.items():
+            self.assertEqual(status, "ok", f"rank {rank}: {payload}")
+            self.assertTrue(payload, f"rank {rank} voted SKIP unexpectedly")
+
+    def test_starved_rank_broadcasts_skip(self):
+        results = _spawn_and_collect("rank0_starved")
+        for rank, (status, payload) in results.items():
+            self.assertEqual(status, "ok", f"rank {rank}: {payload}")
+            self.assertFalse(
+                payload,
+                f"rank {rank} voted PROCEED but rank 0 was starved",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)