diff --git a/tensorrt_llm/_torch/distributed/__init__.py b/tensorrt_llm/_torch/distributed/__init__.py index 85ca1a4ae5e8..0cb5d48adaf3 100644 --- a/tensorrt_llm/_torch/distributed/__init__.py +++ b/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", ] diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index f53be7a44229..7ee55ea4d261 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/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) diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 5b150665b5b9..a76d4ae9caba 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/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,6 +25,7 @@ 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: @@ -25,3 +33,99 @@ def get_env_enable_pdl(): 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 diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 0f510341eb6e..c656a6d81479 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/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( diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index badce20b44ca..01ccc27208c6 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/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: diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 47bedcd8a953..162c0ae2a16a 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/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: diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index c8abf9654862..df67ab4e143a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 95a04e5c582e..d92109a838a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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, diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 282febd262e1..a3605a7b44bc 100755 --- a/tensorrt_llm/functional.py +++ b/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,