flashinfer all reduce

tensorrt_llm/_torch/distributed/__init__.py

@@ -1,8 +1,10 @@
 from tensorrt_llm.functional import AllReduceFusionOp
 
 from .communicator import Distributed, MPIDist, PPComm, TorchDist
-from .ops import (AllReduce, AllReduceParams, AllReduceStrategy, MoEAllReduce,
-                  MoEAllReduceParams, allgather, alltoall_helix, reducescatter,
+from .ops import (AllReduce, AllReduceParams, AllReduceStrategy,
+                  FlashInferAllReduceParams,
+                  FlashInferVLLMAllReduce, MoEAllReduce, MoEAllReduceParams,
+                  allgather, alltoall_helix, reducescatter,
                   userbuffers_allreduce_finalize)
 
 __all__ = [
@@ -20,4 +22,6 @@
     "PPComm",
     "MPIDist",
     "Distributed",
+    "FlashInferAllReduceParams",
+    "FlashInferVLLMAllReduce",
 ]

tensorrt_llm/_torch/distributed/ops.py

@@ -10,11 +10,23 @@
 from tensorrt_llm._utils import mpi_comm, mpi_disabled
 from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer
 from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams,
-                                     AllReduceStrategy, MoEAllReduceParams)
+                                     AllReduceStrategy,
+                                     FlashInferAllReduceParams,
+                                     MoEAllReduceParams)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.plugin.plugin import CustomAllReduceHelper
 
+from ..flashinfer_utils import current_flashinfer_allreduce_workspace
+
+try:
+    import flashinfer.comm as flashinfer_comm
+except ImportError:
+    print(
+        "FlashInfer comm module not found. Follow readme to install Flashinfer >=2.8.0."
+    )
+    exit(1)
+
 _thread_local = threading.local()
 
 
@@ -706,3 +718,62 @@ def forward(
                 nranks=self.mapping.tp_size,
                 eps=all_reduce_params.eps,
             )
+
+
+class FlashInferVLLMAllReduce(nn.Module):
+
+    def __init__(self, mapping: Mapping, dtype: torch.dtype):
+        super().__init__()
+        self.mapping = mapping
+        self.dtype = dtype
+
+        self.reg_buffer_size = 8192 * 8192 * 2
+        self.workspace = current_flashinfer_allreduce_workspace()
+
+    def custom_all_reduce(self, input: torch.Tensor, registered: bool = False):
+        input_tensor = input.contiguous()
+        output_tensor = torch.empty_like(input_tensor)
+
+        try:
+            if registered:
+                flashinfer_comm.vllm_all_reduce(self.workspace.fa, input_tensor,
+                                                output_tensor, 0, 0, 36)
+            else:
+                flashinfer_comm.vllm_all_reduce(
+                    self.workspace.fa,
+                    input_tensor,
+                    output_tensor,
+                    self.workspace.buffer_ptrs_ipc.peer_ptrs[
+                        self.mapping.tp_rank],
+                    self.reg_buffer_size,
+                    36,  # CTA upper bounds: 36 as mentioned in the API
+                )
+        except Exception as e:
+            logger.error(f"FlashInferVLLMAllReduce failed: {e}")
+            raise
+
+        return output_tensor
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        *,
+        all_reduce_params: Optional[FlashInferAllReduceParams] = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass for vLLM custom AllReduce.
+
+        Args:
+            input (torch.Tensor): Input tensor to be reduced
+            all_reduce_params (Optional[FlashInferAllReduceParams]): Parameters for fused operations
+
+        Returns:
+            torch.Tensor: Reduced tensor
+        """
+        if self.workspace._is_capturing:
+            if torch.cuda.is_current_stream_capturing():
+                return self.custom_all_reduce(input, registered=True)
+            else:
+                return torch.empty_like(input)
+        else:
+            return self.custom_all_reduce(input, registered=False)

tensorrt_llm/_torch/flashinfer_utils.py

@@ -1,7 +1,14 @@
+import contextlib
 import os
 import platform
 import traceback
 
+import torch
+
+from tensorrt_llm._ipc_utils import IpcMemory
+from tensorrt_llm._utils import mpi_comm
+from tensorrt_llm.mapping import Mapping
+
 from ..logger import logger
 
 IS_FLASHINFER_AVAILABLE = False
@@ -18,10 +25,107 @@ def get_env_enable_pdl():
 if platform.system() != "Windows":
     try:
         import flashinfer
+        import flashinfer.comm as flashinfer_comm
         logger.info(f"flashinfer is available: {flashinfer.__version__}")
         IS_FLASHINFER_AVAILABLE = True
     except ImportError:
         traceback.print_exc()
         print(
             "flashinfer is not installed properly, please try pip install or building from source codes"
         )
+
+
+class FlashInferAllReduceWorkspace:
+
+    def __init__(self, mapping: Mapping):
+        if not IS_FLASHINFER_AVAILABLE:
+            raise RuntimeError(
+                "flashinfer is not installed properly, please try pip install or building from source codes"
+            )
+
+        self.mapping = mapping
+
+        max_size = 8192 * 8192 * 2  # 2 bytes for bfloat16
+        logger.info(
+            f"Opening IPC memory for meta with size {flashinfer_comm.vllm_meta_size() + max_size}"
+        )
+        try:
+            self.meta_ptrs_ipc = IpcMemory(
+                mapping,
+                flashinfer_comm.vllm_meta_size() + max_size)
+
+            # Create rank data buffer (8MB as in test)
+            self.rank_data = torch.empty(8 * 1024 * 1024,
+                                         dtype=torch.uint8,
+                                         device=f"cuda:{mapping.local_rank}")
+
+            # Create buffer pointers for IPC communication
+            self.buffer_ptrs_ipc = IpcMemory(mapping, max_size)
+
+            # Initialize custom allreduce
+            self.fa = flashinfer_comm.vllm_init_custom_ar(
+                ipc_tensors=self.meta_ptrs_ipc.peer_ptrs,
+                rank_data=self.rank_data,
+                rank=mapping.tp_rank,
+                full_nvlink=True)
+
+            flashinfer_comm.vllm_register_buffer(self.fa,
+                                                 self.buffer_ptrs_ipc.peer_ptrs)
+        except Exception:
+            traceback.print_exc()
+            logger.error(f"Error initializing FlashInferAllReduceWorkspace")
+            raise
+
+        self._is_capturing = False
+        logger.info(
+            f"FlashInferAllReduceWorkspace initialized for rank {mapping.tp_rank}"
+        )
+
+    @contextlib.contextmanager
+    def capture(self):
+        try:
+            self._is_capturing = True
+            yield
+        finally:
+            self._is_capturing = False
+            # Register graph buffers after capture if not already done
+            self.register_graph_buffers()
+
+    def register_graph_buffers(self):
+        # add error handling
+        handle, offsets = flashinfer_comm.vllm_get_graph_buffer_ipc_meta(
+            self.fa)
+
+        world_size = self.mapping.tp_size
+        tp_rank = self.mapping.tp_rank
+
+        all_data = [[None, None] for _ in range(world_size)]
+        all_data[tp_rank] = [handle, offsets]
+
+        comm = mpi_comm().Split(
+            self.mapping.pp_rank * self.mapping.cp_size + self.mapping.cp_rank,
+            self.mapping.tp_rank)
+
+        for i in range(world_size):
+            all_data[i] = comm.bcast(all_data[i], i)
+
+        handles = [d[0] for d in all_data]
+        offsets_list = [d[1] for d in all_data]
+
+        flashinfer_comm.vllm_register_graph_buffers(self.fa, handles,
+                                                    offsets_list)
+
+
+flashinfer_allreduce_workspace = None
+
+
+def init_flashinfer_allreduce_workspace(mapping: Mapping):
+    global flashinfer_allreduce_workspace
+    if flashinfer_allreduce_workspace is None:
+        flashinfer_allreduce_workspace = FlashInferAllReduceWorkspace(mapping)
+
+
+def current_flashinfer_allreduce_workspace():
+    global flashinfer_allreduce_workspace
+    assert flashinfer_allreduce_workspace is not None, "You must call `init_flashinfer_allreduce_workspace` first"
+    return flashinfer_allreduce_workspace

tensorrt_llm/_torch/models/modeling_llama.py

@@ -715,7 +715,6 @@ def forward(
                 scale = self.mlp.gate_up_proj.input_scale
             else:
                 scale = None
-
             all_reduce_output = self.all_reduce(
                 hidden_states,
                 all_reduce_params=AllReduceParams(

tensorrt_llm/_torch/modules/embedding.py

@@ -64,6 +64,7 @@ def __init__(
             tensor_parallel_mode=tensor_parallel_mode,
             gather_output=gather_output,
             use_custom_cublas_mm=use_custom_cublas_mm,
+            use_flashinfer_allreduce=False,
         )
 
         if tensor_parallel_mode == TensorParallelMode.ROW:

tensorrt_llm/_torch/modules/linear.py

@@ -1823,8 +1823,9 @@ def __init__(
         use_cute_dsl_blockscaling_mm: bool = False,
         use_cute_dsl_nvfp4_blockscaling_mm: bool = False,
         disable_deep_gemm: bool = False,
+        use_flashinfer_allreduce: bool = True,
     ):
-        from ..distributed import AllReduce
+        from ..distributed import AllReduce, FlashInferVLLMAllReduce
 
         super().__init__()
         self.has_bias = bias
@@ -1866,6 +1867,18 @@ def __init__(
         self.all_reduce = AllReduce(mapping=self.mapping,
                                     strategy=allreduce_strategy,
                                     dtype=self.dtype) if reduce_output else None
+
+        self.use_flashinfer_allreduce = use_flashinfer_allreduce
+        self.flashinfer_vllm = self.use_flashinfer_allreduce and os.getenv(
+            "_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1"
+        self.flashinfer_token_threshold = int(
+            os.getenv("_FLASHINFER_TOKEN_THRESHOLD", "256"))
+
+        if self.flashinfer_vllm:
+            self.flash_infer_all_reduce = FlashInferVLLMAllReduce(
+                mapping=self.mapping,
+                dtype=self.dtype) if reduce_output else None
+
         self._weights_created = False
         self.reduce_output = reduce_output
         self.use_custom_cublas_mm = use_custom_cublas_mm
@@ -2000,10 +2013,18 @@ def forward(
                     bias, all_reduce_params)
                 bias = None if fuse_bias else bias
                 output = self.apply_linear(input, bias, lora_params, layer_idx)
-                output = self.all_reduce(
-                    output,
-                    all_reduce_params=all_reduce_params,
-                )
+
+                if self.flashinfer_vllm and output.size(
+                        0) <= self.flashinfer_token_threshold:
+                    output = self.flash_infer_all_reduce(
+                        output,
+                        all_reduce_params=None,
+                    )
+                else:
+                    output = self.all_reduce(
+                        output,
+                        all_reduce_params=all_reduce_params,
+                    )
             else:
                 output = self.apply_linear(input, bias, lora_params, layer_idx)
         elif self.tp_mode == TensorParallelMode.COLUMN:

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -244,12 +244,21 @@ def _setup_spec_decoding_and_forward(key: Tuple[int, int, int],
                 capture_inputs['attn_metadata'].use_spec_decoding = True
             return forward_fn(capture_inputs)
 
+        def _maybe_add_flashinfer_allreduce_context():
+            if engine._enable_flashinfer_allreduce:
+                from ..flashinfer_utils import \
+                    current_flashinfer_allreduce_workspace
+                return current_flashinfer_allreduce_workspace().capture()
+            else:
+                return contextlib.nullcontext()
+
         # We have to do warm up runs to initialize PyTorch's
         # internal states according to the docs:
         # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graph-semantics
         # This also lets us initialize states in the attn_metadata.
         graph = torch.cuda.CUDAGraph()
-        with with_multi_stream(True), piecewise_cuda_graph(False):
+        with with_multi_stream(True), piecewise_cuda_graph(
+                False), _maybe_add_flashinfer_allreduce_context():
             for _ in range(self.WARMUP_STEPS):
                 _setup_spec_decoding_and_forward(key, forward_fn,
                                                  capture_inputs)

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -34,6 +34,7 @@
 from ..distributed import MPIDist
 from ..distributed.communicator import init_pp_comm
 from ..expert_statistic import ExpertStatistic
+from ..flashinfer_utils import init_flashinfer_allreduce_workspace
 from ..metadata import KVCacheParams
 from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader
 from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids
@@ -153,6 +154,7 @@ def __init__(
         self.mapping = mapping
         if mapping.has_pp():
             init_pp_comm(mapping)
+
         self.dist = dist
         if dist is not None:
             ExpertStatistic.create(self.dist.rank)
@@ -172,6 +174,11 @@ def __init__(
         self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures(
         )
 
+        self._enable_flashinfer_allreduce = False
+        if os.getenv("_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1":
+            init_flashinfer_allreduce_workspace(mapping)
+            self._enable_flashinfer_allreduce = True
+
         if model is None:
             loader = ModelLoader(
                 pytorch_backend_config=pytorch_backend_config,

tensorrt_llm/functional.py

@@ -26,6 +26,8 @@
 import tensorrt as trt
 # isort: on
 
+import flashinfer.comm as flashinfer_comm
+
 from . import graph_rewriting as gw
 from ._common import default_net, default_trtnet, precision
 from ._utils import (QuantModeWrapper, bf16_array, bool_array,
@@ -3978,6 +3980,39 @@ def is_valid(self):
             return (self.expanded_idx_to_permuted_idx is not None)
 
 
+class FlashInferAllReduceParams(AllReduceParams):
+
+    def __init__(self,
+                 strategy: Optional[
+                     flashinfer_comm.AllReduceStrategyType] = flashinfer_comm.
+                 AllReduceStrategyType.TWOSHOT,
+                 fusion_op: Optional[
+                     flashinfer_comm.AllReduceFusionOp] = flashinfer_comm.
+                 AllReduceFusionOp.NONE,
+                 config_mode: Optional[
+                     flashinfer_comm.AllReduceStrategyConfig] = flashinfer_comm.
+                 AllReduceStrategyConfig.USE_MEMCPY,
+                 bias: Optional[Tensor] = None,
+                 residual: Optional[Tensor] = None,
+                 norm_weight: Optional[Tensor] = None,
+                 scale: Optional[Tensor] = None,
+                 norm_pre_residual_weight: Optional[Tensor] = None,
+                 eps: float = 1e-06,
+                 enable_allreduce: bool = True):
+        super().__init__(
+            bias=bias,
+            residual=residual,
+            norm_weight=norm_weight,
+            scale=scale,
+            norm_pre_residual_weight=norm_pre_residual_weight,
+            eps=eps,
+            enable_allreduce=enable_allreduce,
+        )
+        self.strategy = strategy
+        self.fusion_op = fusion_op
+        self.config_mode = config_mode
+
+
 def create_allreduce_plugin(
     network: trt.INetworkDefinition,
     tensor: trt.ITensor,