diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 086a2a463..661ef4336 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/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}. " diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 2d5d58981..85a34d505 100644 --- a/python/sglang/srt/utils/common.py +++ b/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. diff --git a/test/registered/distributed/test_flashinfer_fusion_preflight.py b/test/registered/distributed/test_flashinfer_fusion_preflight.py new file mode 100644 index 000000000..82b00e12d --- /dev/null +++ b/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)