import logging from typing import Optional, Tuple import torch 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 logger = logging.getLogger(__name__) _flashinfer_comm = None _workspace_manager = None 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 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 " "implementation" ) class FlashInferWorkspaceManager: def __init__(self): 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( self, world_size: int, rank: int, max_token_num: int, hidden_dim: int, dtype: torch.dtype, use_oneshot: Optional[bool] = None, ): """Initialize workspace""" if _flashinfer_comm is None: logger.warning( "FlashInfer comm not available, skipping workspace " "initialization" ) return self.cleanup() 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}, 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.workspace is not None: try: 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.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, 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: return False world_size = get_tensor_model_parallel_world_size() if world_size <= 1: return False 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, dtype=dtype, use_oneshot=use_oneshot, ) 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, ) -> 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, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Use FlashInfer's fused allreduce + residual + RMS norm operation 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 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 world_size = get_tensor_model_parallel_world_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], dtype=input_tensor.dtype, token_num=input_tensor.shape[0], use_oneshot=use_oneshot, ): logger.debug("FlashInfer workspace not available") return None, None residual_out = torch.empty_like(residual) norm_out = torch.empty_like(input_tensor) _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(): global _workspace_manager if _workspace_manager is not None: _workspace_manager.cleanup()