[FlashInfer] Switch FlashInfer allreduce fusion to unified API (#18341)

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
 
+        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.workspace = None
-        self.ipc_handles = 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,
+        dtype: torch.dtype,
-        use_fp32_lamport: bool = False,
+        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()
-
+        try:
-        self.ipc_handles, self.workspace_tensor = (
+            self.workspace = _flashinfer_comm.create_allreduce_fusion_workspace(
-            comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
+                backend="trtllm",
-                rank,
+                world_size=world_size,
-                world_size,
+                rank=rank,
-                max_token_num,
+                max_token_num=max_token_num,
-                hidden_dim,
+                hidden_dim=hidden_dim,
-                group=group,
+                dtype=dtype,
-                use_fp32_lamport=use_fp32_lamport,
+                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.workspace.destroy()
-                    self.ipc_handles, group=dist.group.WORLD
-                )
             except Exception as e:
                 logger.warning(f"Failed to cleanup FlashInfer workspace: {e}")
             finally:
-                self.workspace_tensor = None
+                self.workspace = None
-                self.ipc_handles = 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(
+    _flashinfer_comm.allreduce_fusion(
-        allreduce_in=input_tensor,
+        input=input_tensor,
-        world_size=world_size,
+        workspace=_workspace_manager.workspace,
-        world_rank=dist.get_rank(),
+        pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
-        token_num=token_num,
-        hidden_dim=hidden_dim,
-        workspace_ptrs=_workspace_manager.workspace_tensor,
         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,
+        residual_in=residual,
-        scale_out=None,
         rms_gamma=weight,
         rms_eps=eps,
-        scale_factor=None,
+        use_oneshot=use_oneshot,
-        layout_code=None,
+        fp32_acc=fp32_acc,
     )
 
     return norm_out, residual_out