From 3d8fc9a0cad331f60b7c635e6809dd89eb409f9d Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Tue, 17 Mar 2026 11:59:02 -0700 Subject: [PATCH] Revert "[Nvidia] Add trtllm mnnvl allreduce with unified flashinfer allreduce fusion api" (#20792) --- docs/advanced_features/server_arguments.md | 3 +- python/sglang/srt/layers/communicator.py | 2 +- .../srt/layers/flashinfer_comm_fusion.py | 271 +++--------------- python/sglang/srt/server_args.py | 51 +--- 4 files changed, 53 insertions(+), 274 deletions(-) diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 98de77dce..730a22745 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -314,7 +314,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--moe-a2a-backend` | Select the backend for all-to-all communication for expert parallelism. | `none` | `none`, `deepep`, `mooncake`, `mori`, `nixl`, `ascend_fuseep`| | `--moe-runner-backend` | Choose the runner backend for MoE. | `auto` | `auto`, `deep_gemm`, `triton`, `triton_kernel`, `flashinfer_trtllm`, `flashinfer_trtllm_routed`, `flashinfer_cutlass`, `flashinfer_mxfp4`, `flashinfer_cutedsl`, `cutlass` | | `--flashinfer-mxfp4-moe-precision` | Choose the computation precision of flashinfer mxfp4 moe | `default` | `default`, `bf16` | -| `--flashinfer-allreduce-fusion-backend` | Enable FlashInfer allreduce fusion (fused allreduce + Residual + RMSNorm) and choose backend. When not set, the feature is disabled. Options: `auto` (choose best), `trtllm` (SM90/100, single-node only), `mnnvl` (SM100, single/multi-node). Backend support table (SM100/SM90, single/multi-node) is in `sglang.srt.layers.flashinfer_comm_fusion`. | `None` | `auto`, `trtllm`, `mnnvl` | +| `--enable-flashinfer-allreduce-fusion` | Enable FlashInfer allreduce fusion with Residual RMSNorm. | `False` | bool flag (set to enable) | | `--enable-aiter-allreduce-fusion` | Enable aiter allreduce fusion with Residual RMSNorm. | `False` | bool flag (set to enable) | | `--deepep-mode` | Select the mode when enable DeepEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch. | `auto` | `normal`, `low_latency`, `auto` | | `--ep-num-redundant-experts` | Allocate this number of redundant experts in expert parallel. | `0` | Type: int | @@ -563,7 +563,6 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--enable-flashinfer-trtllm-moe` | NOTE: --enable-flashinfer-trtllm-moe is deprecated. Please set `--moe-runner-backend` to 'flashinfer_trtllm' instead. | `None` | N/A | | `--enable-triton-kernel-moe` | NOTE: --enable-triton-kernel-moe is deprecated. Please set `--moe-runner-backend` to 'triton_kernel' instead. | `None` | N/A | | `--enable-flashinfer-mxfp4-moe` | NOTE: --enable-flashinfer-mxfp4-moe is deprecated. Please set `--moe-runner-backend` to 'flashinfer_mxfp4' instead. | `None` | N/A | -| `--enable-flashinfer-allreduce-fusion` | NOTE: --enable-flashinfer-allreduce-fusion is deprecated. Please set `--flashinfer-allreduce-fusion-backend=auto` instead. | `None` | N/A | | `--crash-on-nan` | Crash the server on nan logprobs. | `False` | Type: str | | `--hybrid-kvcache-ratio` | Mix ratio in [0,1] between uniform and hybrid kv buffers (0.0 = pure uniform: swa_size / full_size = 1)(1.0 = pure hybrid: swa_size / full_size = local_attention_size / context_length) | `None` | Optional[float] | | `--load-watch-interval` | The interval of load watching in seconds. | `0.1` | Type: float | diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 5eb69fe05..b354d0a04 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -100,7 +100,7 @@ def apply_flashinfer_allreduce_fusion(batch_size: int): and batch_size > 0 and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE and not is_dp_attention_enabled() - and get_global_server_args().flashinfer_allreduce_fusion_backend is not None + and get_global_server_args().enable_flashinfer_allreduce_fusion and not is_flashinfer_allreduce_unavailable() ) diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 28bdd9fcc..c45ff977f 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -2,28 +2,18 @@ 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_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - get_tp_group, ) -from sglang.srt.distributed.parallel_state import in_the_same_node_as -from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import is_flashinfer_available 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 _workspace_manager = None -_mnnvl_comm_backend = None -_AllReduceFusionPattern = None -_create_allreduce_fusion_workspace = None -_allreduce_fusion = None _flashinfer_allreduce_unavailable = False if is_flashinfer_available(): @@ -47,102 +37,12 @@ if is_flashinfer_available(): "implementation" ) - try: - # Try to import the unified allreduce API - from flashinfer.comm.allreduce import ( - allreduce_fusion, - create_allreduce_fusion_workspace, - ) - - # AllReduceFusionPattern might be in trtllm_ar or allreduce module - try: - from flashinfer.comm.allreduce import AllReduceFusionPattern - except ImportError: - from flashinfer.comm.trtllm_ar import AllReduceFusionPattern - - _AllReduceFusionPattern = AllReduceFusionPattern - _create_allreduce_fusion_workspace = create_allreduce_fusion_workspace - _allreduce_fusion = allreduce_fusion - except ImportError: - # Fall back to legacy API if unified API is not available - _AllReduceFusionPattern = None - _create_allreduce_fusion_workspace = None - _allreduce_fusion = None - logger.warning( - "FlashInfer unified allreduce API not available, using legacy API" - ) - - 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. - Uses torch.distributed.broadcast_object_list under the hood. - - Returns the broadcasted object on every rank. - """ - obj_list = [data] - # broadcast_object_list mutates obj_list in-place - dist.broadcast_object_list(obj_list, src=root, group=self._group) - return obj_list[0] - - def barrier(self): - """ - Synchronize all ranks in this communicator. - """ - 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 (fused allreduce + Residual + RMSNorm) backend support -# for --flashinfer-allreduce-fusion-backend: -# -# Feature / Framework | SM100 | SM90 | Single Node | Multi-Node | -# --------------------- | ----- | ---- | ----------- | ---------- | -# TRT-LLM AllReduce | Yes | Yes | Yes | No | -# MNNVL AllReduce | Yes | No | Yes | Yes | -# -# With backend "auto": trtllm is used on single-node, mnnvl on single or multi-node (SM100 only). -# Multi-node + Hopper is unsupported (trtllm would be chosen but does not support multi-node). - def is_flashinfer_allreduce_unavailable() -> bool: return _flashinfer_allreduce_unavailable class FlashInferWorkspaceManager: - """ - Workspace manager using FlashInfer's unified allreduce API. - Wraps FlashInfer's create_allreduce_fusion_workspace() for automatic backend selection. - """ - def __init__(self): self.workspace = None self.world_size = None @@ -151,10 +51,6 @@ class FlashInferWorkspaceManager: self.hidden_dim = None self.dtype = None self.initialized = False - # Max size ever requested (not cleared on cleanup) so we only grow and minimize recreates - self._max_token_num_seen: Optional[int] = None - self._max_hidden_dim_seen: Optional[int] = None - self._logged_init = False def initialize( self, @@ -162,86 +58,27 @@ class FlashInferWorkspaceManager: 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, + dtype: torch.dtype, use_oneshot: Optional[bool] = None, ): - """Initialize workspace using FlashInfer's unified API""" - # Track max size ever requested so we can create with at least that (only grow, minimize recreates) - 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 one - if self.initialized and self.world_size == world_size: - self.cleanup() - + """Initialize workspace""" if _flashinfer_comm is None: logger.warning( "FlashInfer comm not available, skipping workspace initialization" ) return - # Determine GPUs per node for MNNVL backend - # FlashInfer will use this to determine topology internally - gpus_per_node = None - if group is not None: - gpus_per_node = sum(in_the_same_node_as(group, source_rank=0)) - - # Create comm backend for MNNVL if needed - comm_backend = None - if _mnnvl_comm_backend is not None and group is not None: - comm_backend = _mnnvl_comm_backend(group) - + self.cleanup() try: - # Create with at least the max size we've ever been asked for (only grow, fewer recreates) - 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) - # Use FlashInfer's unified API to create workspace - create_kw = dict( - backend=backend, + self.workspace = _flashinfer_comm.create_allreduce_fusion_workspace( + backend="trtllm", 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, - comm_backend=comm_backend, + max_token_num=max_token_num, + hidden_dim=hidden_dim, + dtype=dtype, + force_oneshot_support=bool(use_oneshot), ) - if use_oneshot is not None: - create_kw["force_oneshot_support"] = bool(use_oneshot) - self.workspace = _create_allreduce_fusion_workspace(**create_kw) - self.world_size = world_size - self.rank = rank - 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 workspace initialized for rank {rank}, " - f"world_size {world_size}, backend: {backend_name}" - ) - self._logged_init = True - else: - logger.debug( - f"FlashInfer workspace re-initialized for rank {rank}, " - f"world_size {world_size}, backend: {backend_name}" - ) except Exception as e: global _flashinfer_allreduce_unavailable _flashinfer_allreduce_unavailable = True @@ -251,6 +88,20 @@ class FlashInferWorkspaceManager: ) self.workspace = None self.initialized = False + return + + self.world_size = world_size + self.rank = rank + self.max_token_num = max_token_num + self.hidden_dim = hidden_dim + self.dtype = dtype + self.initialized = True + + backend = getattr(self.workspace, "backend", "unknown") + logger.info( + f"FlashInfer workspace initialized for rank {rank}, " + f"world_size {world_size}, backend {backend}" + ) def is_buffer_size_sufficient( self, @@ -271,22 +122,13 @@ class FlashInferWorkspaceManager: ) except Exception as e: logger.debug(f"FlashInfer workspace size check failed: {e}") - # Fallback: some backends (e.g. MNNVL) may use a different API; reuse if within our allocated size - 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.""" + """Clean up workspace""" if self.workspace is not None: try: - if hasattr(self.workspace, "destroy"): - self.workspace.destroy() + self.workspace.destroy() except Exception as e: logger.warning(f"Failed to cleanup FlashInfer workspace: {e}") finally: @@ -297,7 +139,6 @@ class FlashInferWorkspaceManager: self.max_token_num = None self.hidden_dim = None self.dtype = None - self._logged_init = False _workspace_manager = FlashInferWorkspaceManager() @@ -306,9 +147,7 @@ _workspace_manager = FlashInferWorkspaceManager() def ensure_workspace_initialized( max_token_num: int = 2048, hidden_dim: int = 4096, - use_fp32_lamport: bool = False, - dtype: Optional[torch.dtype] = None, - group: Optional[ProcessGroup] = None, + dtype: torch.dtype = torch.float16, token_num: Optional[int] = None, use_oneshot: Optional[bool] = None, ): @@ -325,7 +164,6 @@ def ensure_workspace_initialized( rank = get_tensor_model_parallel_rank() token_num = token_num or max_token_num - effective_dtype = dtype or torch.bfloat16 if ( not _workspace_manager.initialized @@ -334,20 +172,16 @@ def ensure_workspace_initialized( or not _workspace_manager.is_buffer_size_sufficient( token_num=token_num, hidden_dim=hidden_dim, - dtype=effective_dtype, + dtype=dtype, use_oneshot=use_oneshot, ) ): - backend = get_global_server_args().flashinfer_allreduce_fusion_backend or "auto" _workspace_manager.initialize( world_size=world_size, rank=rank, max_token_num=max_token_num, hidden_dim=hidden_dim, - backend=backend, - use_fp32_lamport=use_fp32_lamport, dtype=dtype, - group=group, use_oneshot=use_oneshot, ) @@ -384,8 +218,7 @@ def flashinfer_allreduce_residual_rmsnorm( fp32_acc: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Use FlashInfer's unified fused allreduce + residual + RMS norm operation. - Automatically selects between IPC and MNNVL backends based on topology and hardware. + Use FlashInfer's fused allreduce + residual + RMS norm operation Args: input_tensor: Input tensor that needs allreduce @@ -400,7 +233,7 @@ def flashinfer_allreduce_residual_rmsnorm( Returns: Tuple[torch.Tensor, torch.Tensor]: (norm_output, residual_output) """ - if not is_flashinfer_available(): + if not is_flashinfer_available() or _flashinfer_comm is None: logger.debug( "FlashInfer not available, falling back to standard implementation" ) @@ -420,57 +253,37 @@ def flashinfer_allreduce_residual_rmsnorm( logger.debug("Non-contiguous tensors, skipping FlashInfer allreduce fusion") return None, None - # Get TP group for workspace initialization - try: - group = get_tp_group().cpu_group - except Exception: - group = 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, - group=group, token_num=input_tensor.shape[0], use_oneshot=use_oneshot, ): logger.debug("FlashInfer workspace not available") return None, None - 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) - try: - if _AllReduceFusionPattern is None or _allreduce_fusion is None: - return None, None - - _allreduce_fusion( - input=input_tensor, - workspace=_workspace_manager.workspace, - pattern=_AllReduceFusionPattern.kARResidualRMSNorm, - launch_with_pdl=trigger_completion_at_end, - use_oneshot=use_oneshot, - fp32_acc=fp32_acc, - residual_in=residual, - residual_out=residual_out, - norm_out=norm_out, - rms_gamma=weight, - rms_eps=eps, - ) - except Exception as e: - logger.warning(f"FlashInfer allreduce fusion failed: {e}") - return None, None + _flashinfer_comm.allreduce_fusion( + 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, + ) return norm_out, residual_out def cleanup_flashinfer_workspace(): - """Clean up FlashInfer workspace""" global _workspace_manager if _workspace_manager is not None: _workspace_manager.cleanup() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e98dd6082..e39f4810b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -513,9 +513,6 @@ class ServerArgs: ] = "none" moe_runner_backend: str = "auto" flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default" - flashinfer_allreduce_fusion_backend: Optional[ - Literal["auto", "trtllm", "mnnvl"] - ] = None enable_flashinfer_allreduce_fusion: bool = False enable_aiter_allreduce_fusion: bool = False deepep_mode: Literal["auto", "normal", "low_latency"] = "auto" @@ -900,18 +897,6 @@ class ServerArgs: ) self.tool_call_parser = deprecated_tool_call_parsers[self.tool_call_parser] - # When user passes --enable-flashinfer-allreduce-fusion, enable with auto backend - if ( - self.enable_flashinfer_allreduce_fusion - and self.flashinfer_allreduce_fusion_backend is None - ): - logger.warning( - "--enable-flashinfer-allreduce-fusion is deprecated. " - "Please use --flashinfer-allreduce-fusion-backend=auto instead." - ) - self.flashinfer_allreduce_fusion_backend = "auto" - self.enable_flashinfer_allreduce_fusion = False - if self.enable_nan_detection: logger.warning( "--enable-nan-detection is deprecated. " @@ -1661,7 +1646,7 @@ class ServerArgs: if is_blackwell_supported(): # workaround for https://github.com/flashinfer-ai/flashinfer/issues/2006 if not self.enable_dp_attention and self.nnodes == 1: - self.flashinfer_allreduce_fusion_backend = "auto" + self.enable_flashinfer_allreduce_fusion = True logger.info( "Enable FlashInfer AllReduce Fusion on sm100 for GptOssForCausalLM" ) @@ -1995,15 +1980,16 @@ class ServerArgs: "Overlap scheduler is disabled when using sparse head for embedding model." ) - # FlashInfer allreduce fusion: auto-enable when single-node (any SM90/100) or multi-node + Blackwell. - # See sglang.srt.layers.flashinfer_comm_fusion for backend support table (TRT-LLM vs MNNVL, SM90/100, single/multi-node). + # TRTLLM AllReduce Fusion supports SM90/100, enable it by default + # for models with explicit support (DeepseekV3, GptOss, Glm4Moe, Qwen3Moe) + # TODO: currently, it is only supported in the single node scenario. https://github.com/flashinfer-ai/flashinfer/issues/2006 # TODO: there is currently a bug on H20 device specifically, https://github.com/flashinfer-ai/flashinfer/issues/2204 device_name = get_device_name() is_h20_device = ( device_name and "H20" in device_name and "H200" not in device_name ) if ( - self.flashinfer_allreduce_fusion_backend is None + not self.enable_flashinfer_allreduce_fusion and model_arch in [ "DeepseekV3ForCausalLM", @@ -2015,11 +2001,11 @@ class ServerArgs: ] and (is_sm90_supported() or is_sm100_supported()) and not self.enable_dp_attention + and self.nnodes == 1 and not is_h20_device - and (self.nnodes == 1 or is_sm100_supported()) and self.moe_a2a_backend == "none" ): - self.flashinfer_allreduce_fusion_backend = "auto" + self.enable_flashinfer_allreduce_fusion = True def _handle_mamba_radix_cache( self, @@ -4636,21 +4622,10 @@ class ServerArgs: default=ServerArgs.flashinfer_mxfp4_moe_precision, help="Choose the computation precision of flashinfer mxfp4 moe", ) - parser.add_argument( - "--flashinfer-allreduce-fusion-backend", - type=str, - choices=["auto", "trtllm", "mnnvl"], - default=None, - help="Enable FlashInfer allreduce fusion and choose backend. When not set, the feature is disabled. " - "Options: 'auto' (choose best), 'trtllm' (SM90/100, single-node only), 'mnnvl' (SM100, single/multi-node). " - "Fuses allreduce with Residual + RMSNorm for supported MoE models.", - ) parser.add_argument( "--enable-flashinfer-allreduce-fusion", - action=DeprecatedStoreTrueAction, - new_flag="--flashinfer-allreduce-fusion-backend=auto", - help="(Deprecated: use --flashinfer-allreduce-fusion-backend=auto) " - "Enable FlashInfer allreduce fusion with Residual RMSNorm.", + action="store_true", + help="Enable FlashInfer allreduce fusion with Residual RMSNorm.", ) parser.add_argument( "--enable-aiter-allreduce-fusion", @@ -5784,14 +5759,6 @@ class ServerArgs: f"Invalid value: '{self.served_model_name}'" ) - # FlashInfer allreduce fusion: mnnvl backend requires Blackwell (SM100) - if self.flashinfer_allreduce_fusion_backend == "mnnvl": - if not is_sm100_supported(): - raise ValueError( - "FlashInfer allreduce fusion backend 'mnnvl' is only supported on Blackwell GPUs (SM100). " - "On Hopper (SM90) or other architectures, use --flashinfer-allreduce-fusion-backend=trtllm or --flashinfer-allreduce-fusion-backend=auto instead." - ) - # Check LoRA self.check_lora_server_args()