chore: import upstream snapshot with attribution

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -0,0 +1,872 @@
+import inspect
+import logging
+from typing import Optional, Tuple
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.distributed import (
+    get_attn_tp_group,
+    get_moe_ep_group,
+    get_moe_tp_group,
+    get_tp_group,
+)
+from sglang.srt.distributed.parallel_state import in_the_same_node_as
+from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.utils import (
+    ceil_align,
+    get_cuda_driver_bindings,
+    is_flashinfer_available,
+    is_sm90_supported,
+    is_sm100_supported,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+logger = logging.getLogger(__name__)
+
+# FlashInfer allreduce fusion: set when flashinfer is available (see block below)
+_flashinfer_comm = None
+_TorchDistBackend = None
+_mnnvl_comm_backend = None
+_create_allreduce_fusion_workspace = None
+_flashinfer_allreduce_unavailable = False
+_flashinfer_create_workspace_supports_group = False
+_flashinfer_create_workspace_supports_comm_backend = False
+_flashinfer_allreduce_supports_trigger_completion = False
+
+
+def _mnnvl_supported(is_multi_node: bool) -> bool:
+    """Whether the mnnvl backend is usable on the current system."""
+    if is_sm100_supported():
+        return True
+    return is_sm90_supported() and not is_multi_node
+
+
+def _resolve_backend(backend: str, is_multi_node: bool = False) -> str:
+    """Resolve the requested FlashInfer allreduce fusion backend."""
+    if not (is_sm90_supported() or is_sm100_supported()):
+        raise ValueError(
+            "FlashInfer allreduce fusion requires SM90 or SM10X NVIDIA GPUs."
+        )
+
+    if backend == "auto":
+        if is_multi_node:
+            if is_sm100_supported():
+                return "mnnvl"
+            raise ValueError(
+                "FlashInfer allreduce fusion does not support multi-node on "
+                "non-Blackwell systems."
+            )
+        if is_sm100_supported():
+            return "mnnvl"
+        return "trtllm"
+
+    if backend == "trtllm" and is_multi_node:
+        raise ValueError(
+            "FlashInfer allreduce fusion trtllm backend supports single-node only."
+        )
+
+    if backend == "mnnvl" and not _mnnvl_supported(is_multi_node):
+        raise ValueError(
+            "FlashInfer allreduce fusion mnnvl backend requires a Blackwell "
+            "system, or SM90 single-node."
+        )
+    return backend
+
+
+def resolve_flashinfer_allreduce_fusion_backend(server_args) -> Optional[str]:
+    backend = getattr(server_args, "flashinfer_allreduce_fusion_backend", None)
+    if backend is None:
+        return None
+    is_multi_node = getattr(server_args, "nnodes", 1) > 1
+    return _resolve_backend(backend, is_multi_node)
+
+
+if is_flashinfer_available():
+    try:
+        import flashinfer.comm as comm
+
+        if hasattr(comm, "allreduce_fusion") and hasattr(
+            comm, "create_allreduce_fusion_workspace"
+        ):
+            _flashinfer_comm = comm
+            _create_allreduce_fusion_workspace = comm.create_allreduce_fusion_workspace
+            workspace_params = inspect.signature(
+                comm.create_allreduce_fusion_workspace
+            ).parameters
+            allreduce_params = inspect.signature(comm.allreduce_fusion).parameters
+            _flashinfer_create_workspace_supports_group = "group" in workspace_params
+            _flashinfer_create_workspace_supports_comm_backend = (
+                "comm_backend" in workspace_params
+            )
+            _flashinfer_allreduce_supports_trigger_completion = (
+                "trigger_completion_at_end" in allreduce_params
+            )
+        else:
+            _flashinfer_allreduce_unavailable = True
+            logger.warning(
+                "flashinfer.comm unified allreduce_fusion API is not available, "
+                "falling back to standard implementation"
+            )
+    except (ImportError, AttributeError) as e:
+        _flashinfer_allreduce_unavailable = True
+        logger.warning(
+            "flashinfer.comm allreduce_fusion API is not available (%s), "
+            "falling back to standard implementation",
+            e,
+        )
+
+    try:
+        from flashinfer.comm.mnnvl import TorchDistBackend
+
+        class _FixedTorchDistBackend(TorchDistBackend):
+            """Workaround for FlashInfer TorchDistBackend issues.
+
+            1. bcast fix: TorchDistBackend.bcast passes the in-group rank
+               directly as `src` to broadcast_object_list, which expects a
+               global rank.
+            2. Graph-capture fix: initialize with NCCL device_group (so
+               the backend derives correct device_idx / GPU mapping), but
+               broadcast via GLOO cpu_group (to avoid NCCL collectives
+               that interfere with CUDA graph capture).
+            """
+
+            def __init__(self, device_group, cpu_group):
+                super().__init__(group=device_group)
+                self._cpu_group = cpu_group
+
+            def bcast(self, data, root):
+                import torch.distributed as dist
+
+                group_ranks = dist.get_process_group_ranks(self._cpu_group)
+                global_root = group_ranks[root]
+                object_list = [data]
+                dist.broadcast_object_list(
+                    object_list, src=global_root, group=self._cpu_group
+                )
+                return object_list[0]
+
+        _TorchDistBackend = _FixedTorchDistBackend
+    except ImportError:
+        logger.debug(
+            "flashinfer.comm.mnnvl.TorchDistBackend is not available, "
+            "allreduce fusion will use the default process group"
+        )
+
+    try:
+        from flashinfer.comm.mnnvl import CommBackend
+
+        class TorchDistributedCommBackend(CommBackend):
+            """
+            Use torch distributed instead of MPI to set up flashinfer MNNVL
+            workspaces during initialization.
+            """
+
+            def __init__(self, group: ProcessGroup):
+                self._group = group
+
+            def Get_rank(self) -> int:
+                return self._group.rank()
+
+            def Get_size(self) -> int:
+                return self._group.size()
+
+            def allgather(self, data: int):
+                gathered = [None] * self.Get_size()
+                dist.all_gather_object(gathered, data, group=self._group)
+                return gathered
+
+            def bcast(self, data, root: int = 0):
+                """Broadcast a picklable Python object from root to all ranks."""
+                obj_list = [data]
+                dist.broadcast_object_list(obj_list, src=root, group=self._group)
+                return obj_list[0]
+
+            def barrier(self):
+                dist.barrier(group=self._group)
+
+            def Split(self, color: int, key: int):
+                # No need to split; we already use the proper group.
+                return self._group
+
+        _mnnvl_comm_backend = TorchDistributedCommBackend
+    except ImportError:
+        _mnnvl_comm_backend = None
+
+
+# FlashInfer allreduce fusion backend support matrix for
+# --flashinfer-allreduce-fusion-backend:
+#
+#   Backend   | SM103 | SM100 | SM90        | Single-Node | Multi-Node |
+#   --------- | ----- | ----- | ----------- | ----------- | ---------- |
+#   trtllm    | Yes   | Yes   | Yes         | Yes         | No         |
+#   mnnvl     | Yes   | Yes   | Single-node | Yes         | Blackwell  |
+#
+# FlashInfer allreduce fusion requires SM90 or SM10X. auto resolves to mnnvl
+# on Blackwell (SM100/SM103) systems (single- and multi-node) and to trtllm on
+# SM90 single-node systems. SM90 multi-node and non-SM90/SM10X configurations
+# are rejected. Either mnnvl or trtllm can be requested explicitly on
+# single-node systems, and mnnvl additionally on Blackwell multi-node.
+
+
+def is_flashinfer_allreduce_unavailable() -> bool:
+    return _flashinfer_allreduce_unavailable
+
+
+def _make_flashinfer_workspace_allocation_prop(cuda_driver):
+    from flashinfer.comm.mnnvl import is_mnnvl_fabric_supported
+
+    handle_types = cuda_driver.CUmemAllocationHandleType
+    if is_mnnvl_fabric_supported(torch.cuda.current_device()):
+        handle_type = handle_types.CU_MEM_HANDLE_TYPE_FABRIC
+    else:
+        handle_type = handle_types.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:
+    """
+    Manages FlashInfer's unified allreduce workspace.
+    Supports trtllm and mnnvl backends via create_allreduce_fusion_workspace().
+    """
+
+    def __init__(self):
+        self.workspace = None
+        self.world_size = None
+        self.rank = None
+        self.group = None
+        self.max_token_num = None
+        self.hidden_dim = None
+        self.dtype = None
+        self.initialized = False
+        # Track max sizes ever requested so the workspace only grows (fewer recreates)
+        self._max_token_num_seen: Optional[int] = None
+        self._max_hidden_dim_seen: Optional[int] = None
+        self._logged_init = False
+
+    def initialize(
+        self,
+        world_size: int,
+        rank: int,
+        max_token_num: int,
+        hidden_dim: int,
+        backend: str = "auto",
+        group: Optional[ProcessGroup] = None,
+        use_fp32_lamport: bool = False,
+        dtype: Optional[torch.dtype] = None,
+        use_oneshot: Optional[bool] = None,
+        device_group: Optional["torch.distributed.ProcessGroup"] = None,
+        cpu_group: Optional["torch.distributed.ProcessGroup"] = None,
+    ):
+        """Initialize workspace using FlashInfer's unified API."""
+        global _flashinfer_allreduce_unavailable
+
+        # Track the high-water mark so allocations only grow
+        self._max_token_num_seen = max(max_token_num, self._max_token_num_seen or 0)
+        self._max_hidden_dim_seen = max(hidden_dim, self._max_hidden_dim_seen or 0)
+
+        # Reuse existing workspace if it already covers this problem size
+        if (
+            self.initialized
+            and self.world_size == world_size
+            and self.is_buffer_size_sufficient(
+                token_num=max_token_num,
+                hidden_dim=hidden_dim,
+                dtype=dtype or torch.bfloat16,
+                use_oneshot=use_oneshot,
+            )
+        ):
+            return
+
+        # Same world_size but buffer too small: free old workspace before creating new
+        if self.initialized and self.world_size == world_size:
+            self.cleanup()
+
+        if _flashinfer_comm is None or _create_allreduce_fusion_workspace is None:
+            logger.warning(
+                "FlashInfer comm not available, skipping workspace initialization"
+            )
+            return
+
+        self.cleanup()
+
+        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
+
+        # Determine GPUs per node for MNNVL topology detection
+        gpus_per_node = None
+        node_pg = cpu_group if cpu_group is not None else group
+        if node_pg is not None:
+            gpus_per_node = sum(in_the_same_node_as(node_pg, source_rank=0))
+        comm_backend = None
+        if (
+            _TorchDistBackend is not None
+            and device_group is not None
+            and cpu_group is not None
+        ):
+            comm_backend = _TorchDistBackend(
+                device_group=device_group, cpu_group=cpu_group
+            )
+        elif _mnnvl_comm_backend is not None and group is not None:
+            comm_backend = _mnnvl_comm_backend(group)
+
+        try:
+            alloc_token_num = max(max_token_num, self._max_token_num_seen or 0)
+            alloc_hidden_dim = max(hidden_dim, self._max_hidden_dim_seen or 0)
+            create_kw = dict(
+                backend=backend,
+                world_size=world_size,
+                rank=rank,
+                max_token_num=alloc_token_num,
+                hidden_dim=alloc_hidden_dim,
+                dtype=dtype or torch.bfloat16,
+                gpus_per_node=gpus_per_node,
+            )
+            if (
+                _flashinfer_create_workspace_supports_comm_backend
+                and comm_backend is not None
+            ):
+                create_kw["comm_backend"] = comm_backend
+            if _flashinfer_create_workspace_supports_group:
+                # Pin the symmetric-memory rendezvous to the actual
+                # subgroup. Without this, flashinfer >=0.6.10 falls back
+                # to WORLD and TP/EP/CP subgroup peers get addressed
+                # incorrectly (kernel hangs in cuda-graph warmup).
+                create_kw["group"] = device_group
+            if use_oneshot is not None:
+                create_kw["force_oneshot_support"] = bool(use_oneshot)
+            if use_fp32_lamport:
+                create_kw["use_fp32_lamport"] = True
+            self.workspace = _create_allreduce_fusion_workspace(**create_kw)
+            self.world_size = world_size
+            self.rank = rank
+            self.group = (device_group, cpu_group)
+            self.max_token_num = alloc_token_num
+            self.hidden_dim = alloc_hidden_dim
+            self.dtype = dtype or torch.bfloat16
+            self.initialized = True
+
+            backend_name = getattr(self.workspace, "backend", "unknown")
+            if not self._logged_init:
+                logger.info(
+                    f"FlashInfer AllReduce Fusion enabled and workspace initialized: "
+                    f"backend={backend_name}, rank={rank}, world_size={world_size}, "
+                    f"max_token_num={self.max_token_num}, hidden_dim={self.hidden_dim}"
+                )
+                self._logged_init = True
+            else:
+                logger.debug(
+                    f"FlashInfer workspace re-initialized: backend={backend_name}, "
+                    f"rank={rank}, world_size={world_size}"
+                )
+        except Exception as e:
+            _flashinfer_allreduce_unavailable = True
+            logger.warning(
+                f"Failed to initialize FlashInfer workspace (backend={backend}): {e}. "
+                "Disabling flashinfer allreduce fusion permanently."
+            )
+            self.workspace = None
+            self.initialized = False
+            return
+
+    def is_buffer_size_sufficient(
+        self,
+        token_num: int,
+        hidden_dim: int,
+        dtype: torch.dtype,
+        use_oneshot: Optional[bool] = None,
+    ) -> bool:
+        if not self.initialized or self.workspace is None:
+            return False
+        try:
+            return self.workspace.is_buffer_size_sufficient(
+                tp_size=self.world_size,
+                num_tokens=token_num,
+                hidden_dim=hidden_dim,
+                dtype=dtype,
+                use_oneshot=use_oneshot,
+            )
+        except Exception as e:
+            logger.debug(f"FlashInfer workspace size check failed: {e}")
+            # Fallback: some backends may not implement is_buffer_size_sufficient;
+            # reuse if within our allocated dimensions.
+            if (
+                self.max_token_num is not None
+                and self.hidden_dim is not None
+                and token_num <= self.max_token_num
+                and hidden_dim <= self.hidden_dim
+            ):
+                return True
+            return False
+
+    def cleanup(self):
+        """Clean up workspace."""
+        if self.workspace is not None:
+            try:
+                if hasattr(self.workspace, "destroy"):
+                    self.workspace.destroy()
+            except Exception as e:
+                logger.warning(f"Failed to cleanup FlashInfer workspace: {e}")
+            finally:
+                self.workspace = None
+                self.initialized = False
+                self.world_size = None
+                self.rank = None
+                self.group = None
+                self.max_token_num = None
+                self.hidden_dim = None
+                self.dtype = None
+                self._logged_init = False
+
+
+def _get_workspace_manager(use_attn_tp_group: bool) -> FlashInferWorkspaceManager:
+    """The per-group fusion workspace manager; the instances live on
+    ``ctx.resources`` (one per comm group, created lazily)."""
+    from sglang.srt.runtime_context import get_resources
+
+    buffers = get_resources().buffers
+    name = (
+        "flashinfer_fusion_attn_tp_workspace"
+        if use_attn_tp_group
+        else "flashinfer_fusion_moe_tp_workspace"
+    )
+    manager = buffers.get(name)
+    if manager is None:
+        manager = FlashInferWorkspaceManager()
+        buffers[name] = manager
+    return manager
+
+
+def _sync_allreduce_unavailable_across_tp():
+    """Synchronize _flashinfer_allreduce_unavailable across all TP ranks.
+
+    If workspace initialization fails on any rank, all ranks must agree to
+    disable fusion. Otherwise ranks diverge during CUDA graph capture: some
+    use FlashInfer fusion (skipping custom allreduce), others fall back to
+    standard allreduce (calling register_buffer collectives), causing a hang
+    in register_graph_buffers.
+    """
+    global _flashinfer_allreduce_unavailable
+    try:
+        import torch.distributed as dist
+
+        tp_group = get_tp_group()
+        if tp_group.world_size <= 1:
+            return
+        flag = torch.tensor(
+            [1 if _flashinfer_allreduce_unavailable else 0],
+            dtype=torch.int32,
+        )
+        dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=tp_group.cpu_group)
+        if flag.item() > 0 and not _flashinfer_allreduce_unavailable:
+            _flashinfer_allreduce_unavailable = True
+            logger.warning(
+                "FlashInfer allreduce fusion disabled globally because "
+                "workspace initialization failed on at least one rank."
+            )
+    except Exception as e:
+        logger.debug(f"Failed to sync flashinfer unavailable flag: {e}")
+
+
+def ensure_workspace_initialized(
+    max_token_num: int = 2048,
+    hidden_dim: int = 4096,
+    use_fp32_lamport: bool = False,
+    dtype: Optional[torch.dtype] = None,
+    token_num: Optional[int] = None,
+    use_oneshot: Optional[bool] = None,
+    use_attn_tp_group: bool = True,
+):
+    """Ensure workspace is initialized."""
+    if _flashinfer_allreduce_unavailable:
+        return False
+
+    if not is_flashinfer_available() or _flashinfer_comm is None:
+        return False
+
+    if use_attn_tp_group:
+        world_size = get_parallel().attn_tp_size
+        rank = get_parallel().attn_tp_rank
+        coordinator = get_attn_tp_group()
+    else:
+        if get_parallel().moe_ep_size > 1:
+            world_size = get_parallel().moe_ep_size
+            rank = get_parallel().moe_ep_rank
+            coordinator = get_moe_ep_group()
+        else:
+            world_size = get_parallel().moe_tp_size
+            rank = get_parallel().moe_tp_rank
+            coordinator = get_moe_tp_group()
+
+    # Always pass the coordinator's groups: flashinfer >=0.6.10 reads the
+    # rendezvous group from `group=...` (falling back to WORLD when None),
+    # so leaving it None silently rendezvouses on WORLD and the kernel ends
+    # up addressing the wrong peers in TP/EP/CP subgroup setups.
+    device_group = coordinator.device_group
+    cpu_group = coordinator.cpu_group
+
+    if world_size <= 1:
+        return False
+
+    workspace_manager = _get_workspace_manager(use_attn_tp_group)
+    token_num = token_num or max_token_num
+    group_key = (device_group, cpu_group)
+    effective_dtype = dtype or torch.bfloat16
+    server_args = get_server_args()
+    backend = resolve_flashinfer_allreduce_fusion_backend(server_args)
+    if backend is None:
+        return False
+
+    if (
+        not workspace_manager.initialized
+        or workspace_manager.world_size != world_size
+        or workspace_manager.rank != rank
+        or workspace_manager.group != group_key
+        or not workspace_manager.is_buffer_size_sufficient(
+            token_num=token_num,
+            hidden_dim=hidden_dim,
+            dtype=effective_dtype,
+            use_oneshot=use_oneshot,
+        )
+    ):
+        workspace_manager.initialize(
+            world_size=world_size,
+            rank=rank,
+            max_token_num=max_token_num,
+            hidden_dim=hidden_dim,
+            backend=backend,
+            group=cpu_group,
+            use_fp32_lamport=use_fp32_lamport,
+            dtype=dtype,
+            use_oneshot=use_oneshot,
+            device_group=device_group,
+            cpu_group=cpu_group,
+        )
+
+        _sync_allreduce_unavailable_across_tp()
+
+    return workspace_manager.initialized
+
+
+def fake_flashinfer_allreduce_residual_rmsnorm(
+    input_tensor: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    max_token_num: int = 16384,
+    use_oneshot: Optional[bool] = None,
+    trigger_completion_at_end: bool = False,
+    fp32_acc: bool = False,
+    use_attn_tp_group: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    residual_out = torch.empty_like(residual)
+    norm_out = torch.empty_like(input_tensor)
+    return norm_out, residual_out
+
+
+@register_custom_op(
+    mutates_args=["input_tensor", "residual", "weight"],
+    fake_impl=fake_flashinfer_allreduce_residual_rmsnorm,
+)
+def flashinfer_allreduce_residual_rmsnorm(
+    input_tensor: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    max_token_num: int = 2048,
+    use_oneshot: Optional[bool] = None,
+    trigger_completion_at_end: bool = False,
+    fp32_acc: bool = False,
+    use_attn_tp_group: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Use FlashInfer's unified fused allreduce + residual + RMS norm operation.
+    Automatically selects between trtllm and mnnvl backends based on topology
+    and hardware (controlled by --flashinfer-allreduce-fusion-backend).
+
+    Args:
+        input_tensor: Input tensor that needs allreduce
+        residual: Residual tensor
+        weight: RMS norm weight
+        eps: RMS norm epsilon
+        max_token_num: Maximum token number
+        use_oneshot: Whether to use oneshot mode
+        trigger_completion_at_end: Whether to trigger completion at end
+        fp32_acc: Whether to use fp32 precision
+        use_attn_tp_group: If True, use attention TP group; otherwise use MoE TP group
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: (norm_output, residual_output)
+    """
+    if not is_flashinfer_available() or _flashinfer_comm is None:
+        logger.debug(
+            "FlashInfer not available, falling back to standard implementation"
+        )
+        return None, None
+
+    if use_attn_tp_group:
+        world_size = get_parallel().attn_tp_size
+    else:
+        if get_parallel().moe_ep_size > 1:
+            world_size = get_parallel().moe_ep_size
+        else:
+            world_size = get_parallel().moe_tp_size
+
+    if world_size <= 1:
+        logger.debug("Single GPU, no need for allreduce fusion")
+        return None, None
+
+    assert input_tensor.shape[0] <= max_token_num
+    if (
+        not input_tensor.is_contiguous()
+        or not residual.is_contiguous()
+        or not weight.is_contiguous()
+    ):
+        logger.debug("Non-contiguous tensors, skipping FlashInfer allreduce fusion")
+        return None, None
+
+    if not ensure_workspace_initialized(
+        max_token_num=max_token_num,
+        hidden_dim=input_tensor.shape[-1],
+        use_fp32_lamport=(input_tensor.dtype == torch.float32),
+        dtype=input_tensor.dtype,
+        token_num=input_tensor.shape[0],
+        use_oneshot=use_oneshot,
+        use_attn_tp_group=use_attn_tp_group,
+    ):
+        logger.debug("FlashInfer workspace not available")
+        return None, None
+
+    workspace_manager = _get_workspace_manager(use_attn_tp_group)
+    if workspace_manager.workspace is None:
+        logger.debug("FlashInfer workspace is None")
+        return None, None
+
+    residual_out = torch.empty_like(residual)
+    norm_out = torch.empty_like(input_tensor)
+
+    kwargs = dict(
+        input=input_tensor,
+        workspace=workspace_manager.workspace,
+        pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
+        launch_with_pdl=True,
+        residual_out=residual_out,
+        norm_out=norm_out,
+        residual_in=residual,
+        rms_gamma=weight,
+        rms_eps=eps,
+        use_oneshot=use_oneshot,
+        fp32_acc=fp32_acc,
+    )
+    if _flashinfer_allreduce_supports_trigger_completion:
+        kwargs["trigger_completion_at_end"] = trigger_completion_at_end
+    _flashinfer_comm.allreduce_fusion(**kwargs)
+
+    return norm_out, residual_out
+
+
+def pre_initialize_workspaces(
+    max_token_num: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+    use_oneshot: Optional[bool] = None,
+):
+    """Pre-initialize flashinfer workspaces before CUDA graph capture.
+
+    This must be called before graph capture to avoid collective operations
+    (broadcasts, barriers) inside the graph capture context, which can
+    deadlock with custom_all_reduce.register_graph_buffers.
+    """
+    if _flashinfer_allreduce_unavailable or _flashinfer_comm is None:
+        return
+
+    # Initialize MoE workspace
+    ensure_workspace_initialized(
+        max_token_num=max_token_num,
+        hidden_dim=hidden_dim,
+        dtype=dtype,
+        use_oneshot=use_oneshot,
+        use_attn_tp_group=False,
+    )
+
+    # Initialize attention workspace
+    ensure_workspace_initialized(
+        max_token_num=max_token_num,
+        hidden_dim=hidden_dim,
+        dtype=dtype,
+        use_oneshot=use_oneshot,
+        use_attn_tp_group=True,
+    )
+
+
+def cleanup_flashinfer_workspace():
+    from sglang.srt.runtime_context import get_resources
+
+    buffers = get_resources().buffers
+    for name in (
+        "flashinfer_fusion_attn_tp_workspace",
+        "flashinfer_fusion_moe_tp_workspace",
+    ):
+        manager = buffers.get(name)
+        if manager is not None:
+            manager.cleanup()