From 7b0fb43c7a7e5a6d8584e9ebdd5bc6b12e2f0fcd Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Sun, 22 Feb 2026 00:07:16 +0800 Subject: [PATCH] [FlashInfer] Switch FlashInfer allreduce fusion to unified API (#18341) --- .../srt/layers/flashinfer_comm_fusion.py | 150 ++++++++++++------ 1 file changed, 100 insertions(+), 50 deletions(-) diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index ecc89bb6d..9d1f91900 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -2,9 +2,11 @@ import logging from typing import Optional, Tuple import torch -import torch.distributed as dist -from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) from sglang.srt.utils import is_flashinfer_available from sglang.srt.utils.custom_op import register_custom_op @@ -17,7 +19,15 @@ if is_flashinfer_available(): try: import flashinfer.comm as comm - _flashinfer_comm = comm + if hasattr(comm, "allreduce_fusion") and hasattr( + comm, "create_allreduce_fusion_workspace" + ): + _flashinfer_comm = comm + else: + logger.warning( + "flashinfer.comm unified allreduce_fusion API is not available, " + "falling back to standard implementation" + ) except ImportError: logger.warning( "flashinfer.comm is not available, falling back to standard " @@ -27,10 +37,12 @@ if is_flashinfer_available(): class FlashInferWorkspaceManager: def __init__(self): - self.workspace_tensor = None - self.ipc_handles = None + self.workspace = None self.world_size = None self.rank = None + self.max_token_num = None + self.hidden_dim = None + self.dtype = None self.initialized = False def initialize( @@ -39,13 +51,10 @@ class FlashInferWorkspaceManager: rank: int, max_token_num: int, hidden_dim: int, - group=None, - use_fp32_lamport: bool = False, + dtype: torch.dtype, + use_oneshot: Optional[bool] = None, ): """Initialize workspace""" - if self.initialized and self.world_size == world_size: - return - if _flashinfer_comm is None: logger.warning( "FlashInfer comm not available, skipping workspace " "initialization" @@ -53,47 +62,82 @@ class FlashInferWorkspaceManager: return self.cleanup() - - self.ipc_handles, self.workspace_tensor = ( - comm.trtllm_create_ipc_workspace_for_all_reduce_fusion( - rank, - world_size, - max_token_num, - hidden_dim, - group=group, - use_fp32_lamport=use_fp32_lamport, + try: + self.workspace = _flashinfer_comm.create_allreduce_fusion_workspace( + backend="trtllm", + world_size=world_size, + rank=rank, + max_token_num=max_token_num, + hidden_dim=hidden_dim, + dtype=dtype, + force_oneshot_support=bool(use_oneshot), ) - ) + except Exception as e: + logger.warning(f"Failed to initialize FlashInfer workspace: {e}") + 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}" + f"world_size {world_size}, backend {backend}" ) + 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}") + return False + def cleanup(self): """Clean up workspace""" - if self.initialized and self.ipc_handles is not None: + if self.workspace is not None: try: - _flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce( - self.ipc_handles, group=dist.group.WORLD - ) + self.workspace.destroy() except Exception as e: logger.warning(f"Failed to cleanup FlashInfer workspace: {e}") finally: - self.workspace_tensor = None - self.ipc_handles = None + self.workspace = None self.initialized = False + self.world_size = None + self.rank = None + self.max_token_num = None + self.hidden_dim = None + self.dtype = None _workspace_manager = FlashInferWorkspaceManager() def ensure_workspace_initialized( - max_token_num: int = 2048, hidden_dim: int = 4096, use_fp32_lamport: bool = False + max_token_num: int = 2048, + hidden_dim: int = 4096, + dtype: torch.dtype = torch.float16, + token_num: Optional[int] = None, + use_oneshot: Optional[bool] = None, ): """Ensure workspace is initialized""" if not is_flashinfer_available() or _flashinfer_comm is None: @@ -103,18 +147,27 @@ def ensure_workspace_initialized( if world_size <= 1: return False - rank = dist.get_rank() + rank = get_tensor_model_parallel_rank() + token_num = token_num or max_token_num if ( not _workspace_manager.initialized or _workspace_manager.world_size != world_size + or _workspace_manager.rank != rank + or not _workspace_manager.is_buffer_size_sufficient( + token_num=token_num, + hidden_dim=hidden_dim, + dtype=dtype, + use_oneshot=use_oneshot, + ) ): _workspace_manager.initialize( world_size=world_size, rank=rank, max_token_num=max_token_num, hidden_dim=hidden_dim, - use_fp32_lamport=use_fp32_lamport, + dtype=dtype, + use_oneshot=use_oneshot, ) return _workspace_manager.initialized @@ -177,42 +230,39 @@ def flashinfer_allreduce_residual_rmsnorm( 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, ): logger.debug("FlashInfer workspace not available") return None, None - token_num, hidden_dim = input_tensor.shape - residual_out = torch.empty_like(residual) norm_out = torch.empty_like(input_tensor) - _flashinfer_comm.trtllm_allreduce_fusion( - allreduce_in=input_tensor, - world_size=world_size, - world_rank=dist.get_rank(), - token_num=token_num, - hidden_dim=hidden_dim, - workspace_ptrs=_workspace_manager.workspace_tensor, + _flashinfer_comm.allreduce_fusion( + input=input_tensor, + workspace=_workspace_manager.workspace, + pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, launch_with_pdl=True, - use_oneshot=use_oneshot, - trigger_completion_at_end=trigger_completion_at_end, - fp32_acc=fp32_acc, - pattern_code=(_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm), - allreduce_out=None, - residual_in=residual, residual_out=residual_out, norm_out=norm_out, - quant_out=None, - scale_out=None, + residual_in=residual, rms_gamma=weight, rms_eps=eps, - scale_factor=None, - layout_code=None, + use_oneshot=use_oneshot, + fp32_acc=fp32_acc, ) return norm_out, residual_out