From 32685874f335282496087d15c3d67e67102cdda8 Mon Sep 17 00:00:00 2001 From: Shu Wang Date: Mon, 15 Jun 2026 22:19:15 -0500 Subject: [PATCH] Reenable MNNVL backend for FlashInfer allreduce fusion (#23402) --- python/sglang/srt/layers/communicator.py | 2 +- .../srt/layers/flashinfer_comm_fusion.py | 293 +++++++++++++++--- .../sglang/srt/model_executor/model_runner.py | 2 +- python/sglang/srt/server_args.py | 61 +++- ...wen3_30b_a3b_instruct_2507_logprob_diff.py | 1 + .../test_deepseek_v32_fp4_mtp_tp.py | 10 +- .../layers/test_flashinfer_comm_fusion.py | 181 +++++++++++ 7 files changed, 481 insertions(+), 69 deletions(-) create mode 100644 test/registered/unit/layers/test_flashinfer_comm_fusion.py diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index da94a3ffa..318ee9703 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -175,7 +175,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().enable_flashinfer_allreduce_fusion + and get_global_server_args().flashinfer_allreduce_fusion_backend is not None 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 8bb26d512..9f0e0b63c 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -3,6 +3,8 @@ 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_tensor_model_parallel_rank, @@ -16,21 +18,69 @@ from sglang.srt.distributed import ( get_moe_tp_group, 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 ( 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 +_mnnvl_non_blackwell_fallback_logged = False + + +def _mnnvl_supported(is_multi_node: bool) -> bool: + """Whether the mnnvl backend is usable on the current system. + + mnnvl runs on Blackwell (SM10x) for both single- and multi-node, and on + SM90 for single-node only. Multi-node mnnvl on non-Blackwell is not + supported and must fall back to trtllm. + """ + 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.""" + global _mnnvl_non_blackwell_fallback_logged + + if backend == "auto": + # Prefer mnnvl wherever it is supported (any Blackwell system, or SM90 + # single-node); fall back to trtllm otherwise. + return "mnnvl" if _mnnvl_supported(is_multi_node) else "trtllm" + + if backend == "mnnvl" and not _mnnvl_supported(is_multi_node): + if not _mnnvl_non_blackwell_fallback_logged: + logger.info( + "FlashInfer allreduce fusion: forcing trtllm backend " + "(mnnvl requires a Blackwell system, or SM90 single-node)." + ) + _mnnvl_non_blackwell_fallback_logged = True + return "trtllm" + 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(): @@ -41,6 +91,7 @@ if is_flashinfer_available(): 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 @@ -58,11 +109,12 @@ if is_flashinfer_available(): "flashinfer.comm unified allreduce_fusion API is not available, " "falling back to standard implementation" ) - except ImportError: + except (ImportError, AttributeError) as e: _flashinfer_allreduce_unavailable = True logger.warning( - "flashinfer.comm is not available, falling back to standard " - "implementation" + "flashinfer.comm allreduce_fusion API is not available (%s), " + "falling back to standard implementation", + e, ) try: @@ -102,6 +154,61 @@ if is_flashinfer_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 | +# +# mnnvl runs on any Blackwell GPU (SM10x) for both single- and multi-node, and +# on SM90 for single-node only. auto resolves to mnnvl wherever it is supported +# and to trtllm otherwise. An explicit mnnvl request on an unsupported +# configuration (e.g. SM90 multi-node) falls back to trtllm (see +# _resolve_backend). + def is_flashinfer_allreduce_unavailable() -> bool: return _flashinfer_allreduce_unavailable @@ -269,6 +376,11 @@ def _preflight_check_workspace_memory( 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 @@ -278,6 +390,10 @@ class FlashInferWorkspaceManager: 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, @@ -285,13 +401,39 @@ class FlashInferWorkspaceManager: rank: int, max_token_num: int, hidden_dim: int, - dtype: torch.dtype, + 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""" - if _flashinfer_comm is 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" ) @@ -299,7 +441,6 @@ class FlashInferWorkspaceManager: self.cleanup() - global _flashinfer_allreduce_unavailable if not _preflight_check_workspace_memory( world_size=world_size, max_token_num=max_token_num, @@ -312,56 +453,82 @@ class FlashInferWorkspaceManager: 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: - kwargs = dict( - backend="trtllm", + 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=max_token_num, - hidden_dim=hidden_dim, - dtype=dtype, - force_oneshot_support=bool(use_oneshot), + max_token_num=alloc_token_num, + hidden_dim=alloc_hidden_dim, + dtype=dtype or torch.bfloat16, + gpus_per_node=gpus_per_node, ) - create_workspace = _flashinfer_comm.create_allreduce_fusion_workspace - if _flashinfer_create_workspace_supports_group: - # Pin the symmetric-memory rendezvous to the actual subgroup. - # Older FlashInfer releases only support comm_backend. - kwargs["group"] = device_group if ( - _TorchDistBackend is not None - and _flashinfer_create_workspace_supports_comm_backend - and device_group is not None - and cpu_group is not None + _flashinfer_create_workspace_supports_comm_backend + and comm_backend is not None ): - kwargs["comm_backend"] = _TorchDistBackend( - device_group=device_group, cpu_group=cpu_group + 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}" ) - self.workspace = create_workspace(**kwargs) except Exception as e: _flashinfer_allreduce_unavailable = True logger.warning( - f"Failed to initialize FlashInfer workspace: {e}. " + f"Failed to initialize FlashInfer workspace (backend={backend}): {e}. " "Disabling flashinfer allreduce fusion permanently." ) self.workspace = None self.initialized = False return - self.world_size = world_size - self.rank = rank - self.group = (device_group, cpu_group) - 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}, " - f"max_token_num {max_token_num}, hidden_dim {hidden_dim}" - ) - def is_buffer_size_sufficient( self, token_num: int, @@ -381,13 +548,23 @@ class FlashInferWorkspaceManager: ) 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""" + """Clean up workspace.""" if self.workspace is not None: try: - self.workspace.destroy() + if hasattr(self.workspace, "destroy"): + self.workspace.destroy() except Exception as e: logger.warning(f"Failed to cleanup FlashInfer workspace: {e}") finally: @@ -399,6 +576,7 @@ class FlashInferWorkspaceManager: self.max_token_num = None self.hidden_dim = None self.dtype = None + self._logged_init = False _attn_tp_workspace_manager = FlashInferWorkspaceManager() @@ -445,12 +623,13 @@ def _sync_allreduce_unavailable_across_tp(): def ensure_workspace_initialized( max_token_num: int = 2048, hidden_dim: int = 4096, - dtype: torch.dtype = torch.float16, + 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""" + """Ensure workspace is initialized.""" if _flashinfer_allreduce_unavailable: return False @@ -484,6 +663,11 @@ def ensure_workspace_initialized( 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_global_server_args() + backend = resolve_flashinfer_allreduce_fusion_backend(server_args) + if backend is None: + return False if ( not workspace_manager.initialized @@ -493,7 +677,7 @@ def ensure_workspace_initialized( or not workspace_manager.is_buffer_size_sufficient( token_num=token_num, hidden_dim=hidden_dim, - dtype=dtype, + dtype=effective_dtype, use_oneshot=use_oneshot, ) ): @@ -502,6 +686,9 @@ def ensure_workspace_initialized( 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, @@ -545,7 +732,9 @@ def flashinfer_allreduce_residual_rmsnorm( use_attn_tp_group: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Use FlashInfer's fused allreduce + residual + RMS norm operation + 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 @@ -570,9 +759,6 @@ def flashinfer_allreduce_residual_rmsnorm( if use_attn_tp_group: world_size = get_attn_tensor_model_parallel_world_size() else: - # If MoE expert parallel world size > 1, use expert parallel group - # Otherwise, use tensor parallel group - # The two values cannot be larger than 1 at the same time if get_moe_expert_parallel_world_size() > 1: world_size = get_moe_expert_parallel_world_size() else: @@ -594,6 +780,7 @@ def flashinfer_allreduce_residual_rmsnorm( 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, @@ -602,10 +789,14 @@ def flashinfer_allreduce_residual_rmsnorm( 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) - workspace_manager = _get_workspace_manager(use_attn_tp_group) kwargs = dict( input=input_tensor, workspace=workspace_manager.workspace, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 30f2c2d61..e9ce8579e 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2463,7 +2463,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): (broadcasts, barriers) inside the graph capture context, which can deadlock with custom_all_reduce.register_graph_buffers. """ - if not self.server_args.enable_flashinfer_allreduce_fusion: + if self.server_args.flashinfer_allreduce_fusion_backend is None: return from sglang.srt.layers.communicator import FUSE_ALLREDUCE_MAX_BATCH_SIZE diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e4edefa98..7a1e2b02c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -659,6 +659,9 @@ class ServerArgs: flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default" enable_flashinfer_allreduce_fusion: bool = False enforce_disable_flashinfer_allreduce_fusion: bool = False + flashinfer_allreduce_fusion_backend: Optional[ + Literal["auto", "trtllm", "mnnvl"] + ] = None enable_aiter_allreduce_fusion: bool = False deepep_mode: Literal["auto", "normal", "low_latency"] = "auto" deepep_dispatcher_output_dtype: Literal["auto", "bf16", "fp8", "int8", "nvfp4"] = ( @@ -1200,6 +1203,17 @@ 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 # Deprecated attention-backend alias: "compressed" -> "dsv4". for attr in ( "attention_backend", @@ -2777,13 +2791,14 @@ class ServerArgs: "Overlap scheduler is disabled when using sparse head for embedding model." ) - # TRTLLM AllReduce Fusion supports SM90/100, enable it by default - # for models with explicit support (DeepseekV3, GptOss, Glm4Moe, - # MistralLarge3, Qwen3/Qwen3Next/Qwen3.5 MoE families) - # TODO: currently, it is only supported in the single node scenario. https://github.com/flashinfer-ai/flashinfer/issues/2006 - + # Auto-enable FlashInfer AllReduce Fusion on SM100 only, for models with + # explicit support (DeepseekV3, GptOss, Glm4Moe, MistralLarge3, + # Qwen3/Qwen3Next/Qwen3.5 MoE families). SM90 is not auto-enabled because + # auto resolves to mnnvl, which requires a working NVLink multicast fabric + # that SM90 nodes do not reliably have; SM90 users can opt in explicitly + # via --flashinfer-allreduce-fusion-backend. if ( - not self.enable_flashinfer_allreduce_fusion + self.flashinfer_allreduce_fusion_backend is None and model_arch in [ "DeepseekV3ForCausalLM", @@ -2800,20 +2815,19 @@ class ServerArgs: "InternS2PreviewForConditionalGeneration", "Qwen3_5ForConditionalGeneration", ] - and (is_sm90_supported() or is_sm100_supported()) + and is_sm100_supported() and self.tp_size > 1 and not self.enable_dp_attention - and self.nnodes == 1 and self.moe_a2a_backend == "none" ): - self.enable_flashinfer_allreduce_fusion = True + self.flashinfer_allreduce_fusion_backend = "auto" logger.info( - f"Auto-enabling FlashInfer AllReduce Fusion on SM90/SM10X for {model_arch}" + f"Auto-enabling FlashInfer AllReduce Fusion on SM10X for {model_arch}" ) # Apply enforce_disable_flashinfer_allreduce_fusion after all model-specific adjustments if self.enforce_disable_flashinfer_allreduce_fusion: - self.enable_flashinfer_allreduce_fusion = False + self.flashinfer_allreduce_fusion_backend = None logger.info( "FlashInfer allreduce fusion is forcibly disabled " "via --enforce-disable-flashinfer-allreduce-fusion." @@ -4402,11 +4416,11 @@ class ServerArgs: ) self.enable_aiter_allreduce_fusion = False - if self.enable_flashinfer_allreduce_fusion: + if self.flashinfer_allreduce_fusion_backend is not None: logger.warning( - "Disable --enable-flashinfer-allreduce-fusion because deterministic inference is enabled." + "Disable --flashinfer-allreduce-fusion-backend because deterministic inference is enabled." ) - self.enable_flashinfer_allreduce_fusion = False + self.flashinfer_allreduce_fusion_backend = None # Check sampling backend if self.sampling_backend != "ascend": @@ -6256,10 +6270,27 @@ 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. " + "Defaults to auto. " + "'auto': choose mnnvl on SM90 single-node systems and " + "SM100/SM103 single-node or multi-node systems; choose trtllm otherwise. " + "'trtllm': available on single-node systems only. " + "'mnnvl': available on SM90 single-node systems and SM100/SM103 " + "single-node or multi-node systems via MNNVL fabric. " + "Fuses allreduce with Residual + RMSNorm for supported MoE models." + ), + ) parser.add_argument( "--enable-flashinfer-allreduce-fusion", action="store_true", - help="Enable FlashInfer allreduce fusion with Residual RMSNorm.", + help="(Deprecated: use --flashinfer-allreduce-fusion-backend=auto) " + "Enable FlashInfer allreduce fusion with Residual RMSNorm.", ) parser.add_argument( "--enforce-disable-flashinfer-allreduce-fusion", diff --git a/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py b/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py index 58e9ffb87..b9b20a920 100644 --- a/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py @@ -82,6 +82,7 @@ class TestLoRAQwen3_30B_A3B_Instruct_2507_LogprobDiff(CustomTestCase): lora_paths={"my_lora": adapter_path}, lora_backend=LORA_BACKEND, attention_backend="flashinfer", + flashinfer_allreduce_fusion_backend="trtllm", moe_runner_backend=MOE_RUNNER_BACKEND, experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS, prefill_attention_backend=PREFILL_ATTENTION_BACKEND, diff --git a/test/registered/models_e2e/test_deepseek_v32_fp4_mtp_tp.py b/test/registered/models_e2e/test_deepseek_v32_fp4_mtp_tp.py index c0eaaf230..c7dd6120b 100644 --- a/test/registered/models_e2e/test_deepseek_v32_fp4_mtp_tp.py +++ b/test/registered/models_e2e/test_deepseek_v32_fp4_mtp_tp.py @@ -46,7 +46,15 @@ class TestDeepseekV32FP4TPSpec(GSM8KMixin, DefaultServerBase): gsm8k_accept_length_thres = 2.7 def test_z_bs_1_speed(self): - args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048) + args = BenchArgs( + port=int(self.base_url.split(":")[-1]), + max_new_tokens=2048, + prompt=( + "Human: Think carefully before answering. Build a fully functional FastAPI todo server. " + "Start with a short design plan, then output the complete Python code, then show how to run it " + "and test three endpoints.\n\nAssistant:" + ), + ) acc_length, speed = send_one_prompt(args) print(f"{acc_length=:.2f} {speed=:.2f}") diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py new file mode 100644 index 000000000..a0efc037e --- /dev/null +++ b/test/registered/unit/layers/test_flashinfer_comm_fusion.py @@ -0,0 +1,181 @@ +import types +import unittest +from unittest.mock import patch + +import torch + +from sglang.srt.layers import flashinfer_comm_fusion as fusion +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-b200") +register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-gb300") + + +class _FakeWorkspace: + def __init__(self, backend, world_size): + self.backend = backend + self.world_size = world_size + + def is_buffer_size_sufficient(self, **_kwargs): + return True + + +class _FakeFlashInferComm: + class AllReduceFusionPattern: + kARResidualRMSNorm = object() + + def __init__(self): + self.calls = [] + + def create_allreduce_fusion_workspace(self, **kwargs): + self.calls.append(kwargs) + return _FakeWorkspace(kwargs["backend"], kwargs["world_size"]) + + def allreduce_fusion( + self, + *, + input, + workspace, + residual_out, + norm_out, + residual_in, + rms_gamma, + rms_eps, + **_kwargs, + ): + allreduced = input * workspace.world_size + expected_residual = allreduced + residual_in + variance = expected_residual.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + expected_norm = ( + expected_residual.to(torch.float32) + * torch.rsqrt(variance + rms_eps) + * rms_gamma.to(torch.float32) + ).to(input.dtype) + residual_out.copy_(expected_residual) + norm_out.copy_(expected_norm) + + +def _torch_allreduce_residual_rmsnorm_baseline( + input_tensor, residual, weight, world_size, eps +): + allreduced = input_tensor * world_size + residual_out = allreduced + residual + variance = residual_out.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + norm_out = ( + residual_out.to(torch.float32) + * torch.rsqrt(variance + eps) + * weight.to(torch.float32) + ).to(input_tensor.dtype) + return norm_out, residual_out + + +class TestFlashInferCommFusion(unittest.TestCase): + def test_auto_backend_resolves_by_arch(self): + single_node = types.SimpleNamespace( + flashinfer_allreduce_fusion_backend="auto", nnodes=1 + ) + multi_node = types.SimpleNamespace( + flashinfer_allreduce_fusion_backend="auto", nnodes=2 + ) + + # Blackwell: mnnvl regardless of node count. + with patch.object(fusion, "is_sm100_supported", return_value=True): + self.assertEqual( + fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), "mnnvl" + ) + self.assertEqual( + fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node), "mnnvl" + ) + + # SM90: mnnvl on single-node, trtllm fallback on multi-node. + with ( + patch.object(fusion, "is_sm100_supported", return_value=False), + patch.object(fusion, "is_sm90_supported", return_value=True), + ): + self.assertEqual( + fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), "mnnvl" + ) + self.assertEqual( + fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node), "trtllm" + ) + + # Pre-SM90: trtllm everywhere. + with ( + patch.object(fusion, "is_sm100_supported", return_value=False), + patch.object(fusion, "is_sm90_supported", return_value=False), + ): + self.assertEqual( + fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), + "trtllm", + ) + + def test_allreduce_fusion_backends_match_torch_baseline(self): + fake_comm = _FakeFlashInferComm() + original_comm = fusion._flashinfer_comm + original_create = fusion._create_allreduce_fusion_workspace + original_manager = fusion._attn_tp_workspace_manager + original_unavailable = fusion._flashinfer_allreduce_unavailable + try: + fusion._flashinfer_comm = fake_comm + fusion._create_allreduce_fusion_workspace = ( + fake_comm.create_allreduce_fusion_workspace + ) + fusion._flashinfer_allreduce_unavailable = False + + for backend in ("trtllm", "mnnvl"): + with self.subTest(backend=backend): + world_size = 4 + manager = fusion.FlashInferWorkspaceManager() + manager.workspace = _FakeWorkspace(backend, world_size) + manager.initialized = True + fusion._attn_tp_workspace_manager = manager + if not torch.cuda.is_available(): + self.skipTest("FlashInfer allreduce custom op is CUDA-only") + device = torch.device("cuda") + torch.manual_seed(0) + input_tensor = torch.randn(4, 8, dtype=torch.float32, device=device) + residual = torch.randn(4, 8, dtype=torch.float32, device=device) + weight = torch.randn(8, dtype=torch.float32, device=device) + eps = 1e-6 + + expected_norm, expected_residual = ( + _torch_allreduce_residual_rmsnorm_baseline( + input_tensor, residual, weight, world_size, eps + ) + ) + + with ( + patch.object( + fusion, "is_flashinfer_available", return_value=True + ), + patch.object( + fusion, + "get_attn_tensor_model_parallel_world_size", + return_value=world_size, + ), + patch.object( + fusion, "ensure_workspace_initialized", return_value=True + ), + ): + norm_out, residual_out = ( + fusion.flashinfer_allreduce_residual_rmsnorm( + input_tensor=input_tensor, + residual=residual, + weight=weight, + eps=eps, + max_token_num=8, + ) + ) + + torch.testing.assert_close(norm_out, expected_norm) + torch.testing.assert_close(residual_out, expected_residual) + finally: + fusion._flashinfer_comm = original_comm + fusion._create_allreduce_fusion_workspace = original_create + fusion._attn_tp_workspace_manager = original_manager + fusion._flashinfer_allreduce_unavailable = original_unavailable + + +if __name__ == "__main__": + unittest.main()