From 75f0e4c1520dd6a9b7fd3f4220c71dd8cfc280e6 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 11 Aug 2025 10:55:19 -0700 Subject: [PATCH 01/21] init flashinfer all reduce Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/__init__.py | 7 +- tensorrt_llm/_torch/distributed/ops.py | 99 ++++++++++++++++++++- tensorrt_llm/_torch/modules/linear.py | 15 +++- tensorrt_llm/functional.py | 33 +++++++ 4 files changed, 149 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/__init__.py b/tensorrt_llm/_torch/distributed/__init__.py index 85ca1a4ae5e8..c3d0972f19db 100644 --- a/tensorrt_llm/_torch/distributed/__init__.py +++ b/tensorrt_llm/_torch/distributed/__init__.py @@ -1,8 +1,9 @@ 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, + FlashInferAllReduce, FlashInferAllReduceParams, MoEAllReduce, + MoEAllReduceParams, allgather, reducescatter, alltoall_helix, userbuffers_allreduce_finalize) __all__ = [ @@ -20,4 +21,6 @@ "PPComm", "MPIDist", "Distributed", + "FlashInferAllReduce", + "FlashInferAllReduceParams", ] diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index f53be7a44229..b43a7751a618 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -10,11 +10,21 @@ 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 +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 +716,90 @@ def forward( nranks=self.mapping.tp_size, eps=all_reduce_params.eps, ) + + +class FlashInferAllReduce(nn.Module): + + def __init__(self, mapping: Mapping, + strategy: flashinfer_comm.AllReduceStrategyType, + hidden_dim: int, dtype: torch.dtype): + super().__init__() + self.mapping = mapping + self.max_message_size = 70 * 1024 * 1024 # 70MB + self.dtype = dtype + self.hidden_dim = hidden_dim + self.strategy = strategy + + # flashinfer all reduce is better for message sizes < ~70MB + # num tokens * hidden dim * 2 < 70 *1024*1024 + # TODO: 2 is for bf16,fp16. need to add fp8, fp4 + self.max_num_tokens = self.max_message_size // (self.hidden_dim * 2) + + self.workspace = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( + rank=self.mapping.rank, + tp_size=self.mapping.tp_size, + max_token_num=self.max_num_tokens, + hidden_dim=self.hidden_dim, + group=None, + ) + self._flag_value = 1 + + def forward( + self, + input: torch.Tensor, + *, + all_reduce_params: Optional[FlashInferAllReduceParams] = None, + ) -> torch.Tensor: + + input = input.contiguous() + num_token = input.size(0) + hidden_dim = input.size(1) + + output_buffer = torch.empty(num_token * hidden_dim) + + if all_reduce_params is None: + all_reduce_params = FlashInferAllReduceParams( + strategy=self.strategy, + fusion_op=AllReduceFusionOp.NONE, + config_mode=0, + ) + + flashinfer_comm.trtllm_custom_all_reduce( + inp=input.ravel(), + out=output_buffer, + tp_size=self.mapping.tp_size, + tp_rank=self.mapping.rank, + token_num=num_token, + fusion_op_code=all_reduce_params.fusion_op, + strategy_code=all_reduce_params.strategy, + config_code=all_reduce_params.config_mode, + launch_with_pdl=False, + flag_value=self._flag_value, + peer_comm_buffer_ptrs=torch.tensor(self.workspace[0], + dtype=torch.int64), + peer_barrier_ptrs_in=torch.tensor(self.workspace[2], + dtype=torch.int64), + peer_barrier_ptrs_out=torch.tensor(self.workspace[3], + dtype=torch.int64), + bias=None, + residual=None, + weight=None, + weight_pre_residual_norm=None, + eps=None, + intermediate_buffer=None, + lamport_peer_comm_buffer_ptrs_0=None, + lamport_peer_comm_buffer_ptrs_1=None, + lamport_peer_comm_buffer_ptrs_2=None, + # lamport_peer_comm_buffer_ptrs_0=torch.tensor( + # self.workspace[4], dtype=torch.int64 + # ), + # lamport_peer_comm_buffer_ptrs_1=torch.tensor( + # self.workspace[5], dtype=torch.int64 + # ), + # lamport_peer_comm_buffer_ptrs_2=torch.tensor( + # self.workspace[6], dtype=torch.int64 + # ), + ) + self._flag_value += 1 + + return output_buffer.reshape(num_token, hidden_dim) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 47bedcd8a953..8f26edffe203 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -15,7 +15,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AllReduceStrategy) + AllReduceStrategy, FlashInferAllReduce) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.functional import \ preprocess_weights_for_mixed_gemm @@ -1824,6 +1824,8 @@ def __init__( use_cute_dsl_nvfp4_blockscaling_mm: bool = False, disable_deep_gemm: bool = False, ): + import flashinfer.comm as flashinfer_comm + from ..distributed import AllReduce super().__init__() @@ -1866,6 +1868,11 @@ def __init__( self.all_reduce = AllReduce(mapping=self.mapping, strategy=allreduce_strategy, dtype=self.dtype) if reduce_output else None + self.flash_infer_all_reduce = FlashInferAllReduce( + mapping=self.mapping, + hidden_dim=self.out_features, + strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, + 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,7 +2007,11 @@ 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 = self.all_reduce( + # output, + # all_reduce_params=all_reduce_params, + # ) + output = self.flash_infer_all_reduce( output, all_reduce_params=all_reduce_params, ) diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 282febd262e1..460bff63990e 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,37 @@ 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.AllReduceConfigMode] = 0, + 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, From 39b1b7319c9c34263ef35999b1736bc6bd316717 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 11 Aug 2025 16:53:03 -0700 Subject: [PATCH 02/21] add debug traces Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 14 ++++++++++++-- tensorrt_llm/_torch/modules/linear.py | 10 +++++----- tensorrt_llm/functional.py | 4 +++- 3 files changed, 20 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index b43a7751a618..b9fc3f1fa5a9 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -743,6 +743,12 @@ def __init__(self, mapping: Mapping, group=None, ) self._flag_value = 1 + self._enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" + + if self._enable_debug: + print( + f"FlashInferAllReduce.forward: {self.mapping.rank} {self.mapping.tp_size} {self.max_num_tokens} {self.hidden_dim}" + ) def forward( self, @@ -756,12 +762,16 @@ def forward( hidden_dim = input.size(1) output_buffer = torch.empty(num_token * hidden_dim) + if self._enable_debug: + print( + f"FlashInferAllReduce.forward: {num_token} {hidden_dim} {all_reduce_params}" + ) if all_reduce_params is None: all_reduce_params = FlashInferAllReduceParams( strategy=self.strategy, - fusion_op=AllReduceFusionOp.NONE, - config_mode=0, + fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, + config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, ) flashinfer_comm.trtllm_custom_all_reduce( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 8f26edffe203..fc83c75f0f20 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1826,8 +1826,6 @@ def __init__( ): import flashinfer.comm as flashinfer_comm - from ..distributed import AllReduce - super().__init__() self.has_bias = bias self.dtype = dtype @@ -1865,14 +1863,16 @@ def __init__( self.in_features = local_in_features self.out_features = local_out_features - self.all_reduce = AllReduce(mapping=self.mapping, - strategy=allreduce_strategy, - dtype=self.dtype) if reduce_output else None + # self.all_reduce = AllReduce(mapping=self.mapping, + # strategy=allreduce_strategy, + # dtype=self.dtype) if reduce_output else None + self.flash_infer_all_reduce = FlashInferAllReduce( mapping=self.mapping, hidden_dim=self.out_features, strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, 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 diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 460bff63990e..095e921ed36f 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -3989,7 +3989,9 @@ def __init__(self, fusion_op: Optional[ flashinfer_comm.AllReduceFusionOp] = flashinfer_comm. AllReduceFusionOp.NONE, - config_mode: Optional[flashinfer_comm.AllReduceConfigMode] = 0, + config_mode: Optional[ + flashinfer_comm.AllReduceConfigMode] = flashinfer_comm. + AllReduceStrategyConfig.USE_MEMCPY, bias: Optional[Tensor] = None, residual: Optional[Tensor] = None, norm_weight: Optional[Tensor] = None, From 48119cd898f97302121e20063dec61adf69cea97 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 12 Aug 2025 11:27:49 -0700 Subject: [PATCH 03/21] Simplify Signed-off-by: Shreyas Misra --- .../_torch/distributed/communicator.py | 3 + tensorrt_llm/_torch/distributed/ops.py | 68 ++++++++++++------- tensorrt_llm/_torch/models/modeling_llama.py | 21 +++++- tensorrt_llm/_torch/modules/linear.py | 24 +++---- tensorrt_llm/functional.py | 2 +- 5 files changed, 72 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 07f2b4227f99..849c62ba78ed 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -765,6 +765,9 @@ def pp_broadcast(self, obj, root=0): device=torch.device("cpu")) return ret[0] + def allgather(self, obj, root=0): + raise NotImplementedError('allgather is not implemented for TorchDist') + # TODO: rename to PPCommNCCL class PPComm: diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index b9fc3f1fa5a9..64126db96dd4 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -17,6 +17,8 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.plugin.plugin import CustomAllReduceHelper +from .communicator import TorchDist + try: import flashinfer.comm as flashinfer_comm except ImportError: @@ -27,6 +29,30 @@ _thread_local = threading.local() +_enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" + + +def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, hidden_dim: int): + if not hasattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}'): + setattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}', {}) + + flashinfer_allreduce_workspaces = getattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}') + + if mapping not in flashinfer_allreduce_workspaces: + dist = TorchDist(mapping) + flashinfer_allreduce_workspaces[mapping] = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( + rank=mapping.rank, + tp_size=mapping.tp_size, + max_token_num=max_num_tokens, + hidden_dim=hidden_dim, + group=None, + ) + if _enable_debug: + print( + f"FlashInferAllReduce.init: {mapping.rank} {mapping.tp_size} {max_num_tokens} {hidden_dim}" + ) + return flashinfer_allreduce_workspaces[mapping] + def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): @@ -725,30 +751,24 @@ def __init__(self, mapping: Mapping, hidden_dim: int, dtype: torch.dtype): super().__init__() self.mapping = mapping - self.max_message_size = 70 * 1024 * 1024 # 70MB + # self.max_message_size = 70 * 1024 * 1024 # 70MB self.dtype = dtype - self.hidden_dim = hidden_dim + self.hidden_dim = 8192 self.strategy = strategy # flashinfer all reduce is better for message sizes < ~70MB # num tokens * hidden dim * 2 < 70 *1024*1024 # TODO: 2 is for bf16,fp16. need to add fp8, fp4 - self.max_num_tokens = self.max_message_size // (self.hidden_dim * 2) + self.max_num_tokens = 8192 #self.max_message_size // (self.hidden_dim * 2) - self.workspace = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( - rank=self.mapping.rank, - tp_size=self.mapping.tp_size, - max_token_num=self.max_num_tokens, - hidden_dim=self.hidden_dim, - group=None, - ) + self.workspace = get_flashinfer_allreduce_workspace(self.mapping, self.max_num_tokens, self.hidden_dim) self._flag_value = 1 - self._enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" - if self._enable_debug: - print( - f"FlashInferAllReduce.forward: {self.mapping.rank} {self.mapping.tp_size} {self.max_num_tokens} {self.hidden_dim}" - ) + self.all_reduce_params = FlashInferAllReduceParams( + strategy=self.strategy, + fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, + config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, + ) def forward( self, @@ -762,17 +782,13 @@ def forward( hidden_dim = input.size(1) output_buffer = torch.empty(num_token * hidden_dim) - if self._enable_debug: + if _enable_debug: print( - f"FlashInferAllReduce.forward: {num_token} {hidden_dim} {all_reduce_params}" + f"FlashInferAllReduce.forward: {num_token} {hidden_dim} {self.all_reduce_params}" ) - if all_reduce_params is None: - all_reduce_params = FlashInferAllReduceParams( - strategy=self.strategy, - fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, - ) + # if all_reduce_params is None: + flashinfer_comm.trtllm_custom_all_reduce( inp=input.ravel(), @@ -780,9 +796,9 @@ def forward( tp_size=self.mapping.tp_size, tp_rank=self.mapping.rank, token_num=num_token, - fusion_op_code=all_reduce_params.fusion_op, - strategy_code=all_reduce_params.strategy, - config_code=all_reduce_params.config_mode, + fusion_op_code=self.all_reduce_params.fusion_op, + strategy_code=self.all_reduce_params.strategy, + config_code=self.all_reduce_params.config_mode, launch_with_pdl=False, flag_value=self._flag_value, peer_comm_buffer_ptrs=torch.tensor(self.workspace[0], diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 0f510341eb6e..f9793c765a8c 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -10,8 +10,10 @@ from transformers.modeling_utils import load_sharded_checkpoint from transformers.models.llama4.modeling_llama4 import Llama4MultiModalProjector +import flashinfer.comm as flashinfer_comm + from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp, - AllReduceParams, MoEAllReduce) + AllReduceParams, MoEAllReduce, FlashInferAllReduce, FlashInferAllReduceParams) from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ BaseWeightMapper from tensorrt_llm._utils import get_sm_version @@ -686,6 +688,12 @@ def __init__( or self.mapping.tp_size == 1 or self.enable_attention_dp) + self.flash_infer_all_reduce = FlashInferAllReduce( + mapping=self.mapping, + hidden_dim=8192, + strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, + dtype=config.torch_dtype) + def forward( self, position_ids: torch.IntTensor, @@ -706,9 +714,16 @@ def forward( attn_metadata=attn_metadata, attention_mask=self.attention_mask, all_reduce_params=AllReduceParams( - enable_allreduce=not self.disable_attn_allreduce), + enable_allreduce=False), **kwargs, ) + hidden_states = self.flash_infer_all_reduce( + hidden_states, + all_reduce_params=FlashInferAllReduceParams( + strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, + fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, + config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, + )) # Fully Connected if self.PRE_MLP_FUSION: if self.is_nvfp4 or self.is_fp8_quant: @@ -746,7 +761,7 @@ def forward( hidden_states = self.mlp( hidden_states, final_all_reduce_params=AllReduceParams( - enable_allreduce=not self.disable_mlp_allreduce), + enable_allreduce=False), **kwargs, ) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index fc83c75f0f20..1d2d455e19f1 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -15,7 +15,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AllReduceStrategy, FlashInferAllReduce) + AllReduceStrategy, FlashInferAllReduceParams) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.functional import \ preprocess_weights_for_mixed_gemm @@ -1824,7 +1824,8 @@ def __init__( use_cute_dsl_nvfp4_blockscaling_mm: bool = False, disable_deep_gemm: bool = False, ): - import flashinfer.comm as flashinfer_comm + + from ..distributed import AllReduce super().__init__() self.has_bias = bias @@ -1863,15 +1864,10 @@ def __init__( self.in_features = local_in_features self.out_features = local_out_features - # self.all_reduce = AllReduce(mapping=self.mapping, - # strategy=allreduce_strategy, - # dtype=self.dtype) if reduce_output else None - - self.flash_infer_all_reduce = FlashInferAllReduce( - mapping=self.mapping, - hidden_dim=self.out_features, - strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, - dtype=self.dtype) if reduce_output else None + # Needed for embedding layer + self.all_reduce = AllReduce(mapping=self.mapping, + strategy=allreduce_strategy, + dtype=self.dtype) if reduce_output else None self._weights_created = False self.reduce_output = reduce_output @@ -2007,11 +2003,7 @@ 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, - # ) - output = self.flash_infer_all_reduce( + output = self.all_reduce( output, all_reduce_params=all_reduce_params, ) diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 095e921ed36f..a3605a7b44bc 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -3990,7 +3990,7 @@ def __init__(self, flashinfer_comm.AllReduceFusionOp] = flashinfer_comm. AllReduceFusionOp.NONE, config_mode: Optional[ - flashinfer_comm.AllReduceConfigMode] = flashinfer_comm. + flashinfer_comm.AllReduceStrategyConfig] = flashinfer_comm. AllReduceStrategyConfig.USE_MEMCPY, bias: Optional[Tensor] = None, residual: Optional[Tensor] = None, From caa03ac51e3a81d13dbdb1a2e8185934194ccead Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 12 Aug 2025 12:26:37 -0700 Subject: [PATCH 04/21] working impl; non performant --- tensorrt_llm/_torch/distributed/ops.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 64126db96dd4..f56047655068 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -6,6 +6,7 @@ import torch from torch import nn +import torch.distributed as dist from tensorrt_llm._utils import mpi_comm, mpi_disabled from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer @@ -39,13 +40,12 @@ def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, hi flashinfer_allreduce_workspaces = getattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}') if mapping not in flashinfer_allreduce_workspaces: - dist = TorchDist(mapping) flashinfer_allreduce_workspaces[mapping] = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( rank=mapping.rank, tp_size=mapping.tp_size, max_token_num=max_num_tokens, hidden_dim=hidden_dim, - group=None, + group=dist.group.WORLD, ) if _enable_debug: print( @@ -743,6 +743,7 @@ def forward( eps=all_reduce_params.eps, ) +_RANK_DIST_INITIALIZED = False class FlashInferAllReduce(nn.Module): @@ -750,12 +751,18 @@ def __init__(self, mapping: Mapping, strategy: flashinfer_comm.AllReduceStrategyType, hidden_dim: int, dtype: torch.dtype): super().__init__() + global _RANK_DIST_INITIALIZED + self.mapping = mapping # self.max_message_size = 70 * 1024 * 1024 # 70MB self.dtype = dtype self.hidden_dim = 8192 self.strategy = strategy + if not _RANK_DIST_INITIALIZED: + TorchDist(mapping) + _RANK_DIST_INITIALIZED = True + # flashinfer all reduce is better for message sizes < ~70MB # num tokens * hidden dim * 2 < 70 *1024*1024 # TODO: 2 is for bf16,fp16. need to add fp8, fp4 @@ -781,10 +788,10 @@ def forward( num_token = input.size(0) hidden_dim = input.size(1) - output_buffer = torch.empty(num_token * hidden_dim) + output_buffer = torch.empty(num_token * hidden_dim, dtype=input.dtype, device=input.device) if _enable_debug: print( - f"FlashInferAllReduce.forward: {num_token} {hidden_dim} {self.all_reduce_params}" + f"FlashInferAllReduce.forward: {num_token} {hidden_dim}" ) # if all_reduce_params is None: From baca3be1206236f383265093050e2eee7ad55c93 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 13 Aug 2025 09:16:27 -0700 Subject: [PATCH 05/21] Revert "Simplify" This reverts commit 29665dddb1c6f40be5e8e3a0a8f2e2d4d838110a. --- .../_torch/distributed/communicator.py | 3 -- tensorrt_llm/_torch/distributed/ops.py | 40 +++++++++++-------- tensorrt_llm/_torch/models/modeling_llama.py | 22 ++-------- tensorrt_llm/_torch/modules/linear.py | 24 +++++++---- tensorrt_llm/functional.py | 2 +- 5 files changed, 44 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 849c62ba78ed..07f2b4227f99 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -765,9 +765,6 @@ def pp_broadcast(self, obj, root=0): device=torch.device("cpu")) return ret[0] - def allgather(self, obj, root=0): - raise NotImplementedError('allgather is not implemented for TorchDist') - # TODO: rename to PPCommNCCL class PPComm: diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index f56047655068..eca9bf4b4c98 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -18,8 +18,6 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from .communicator import TorchDist - try: import flashinfer.comm as flashinfer_comm except ImportError: @@ -754,9 +752,9 @@ def __init__(self, mapping: Mapping, global _RANK_DIST_INITIALIZED self.mapping = mapping - # self.max_message_size = 70 * 1024 * 1024 # 70MB + self.max_message_size = 70 * 1024 * 1024 # 70MB self.dtype = dtype - self.hidden_dim = 8192 + self.hidden_dim = hidden_dim self.strategy = strategy if not _RANK_DIST_INITIALIZED: @@ -766,16 +764,22 @@ def __init__(self, mapping: Mapping, # flashinfer all reduce is better for message sizes < ~70MB # num tokens * hidden dim * 2 < 70 *1024*1024 # TODO: 2 is for bf16,fp16. need to add fp8, fp4 - self.max_num_tokens = 8192 #self.max_message_size // (self.hidden_dim * 2) + self.max_num_tokens = self.max_message_size // (self.hidden_dim * 2) - self.workspace = get_flashinfer_allreduce_workspace(self.mapping, self.max_num_tokens, self.hidden_dim) + self.workspace = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( + rank=self.mapping.rank, + tp_size=self.mapping.tp_size, + max_token_num=self.max_num_tokens, + hidden_dim=self.hidden_dim, + group=None, + ) self._flag_value = 1 + self._enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" - self.all_reduce_params = FlashInferAllReduceParams( - strategy=self.strategy, - fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, - ) + if self._enable_debug: + print( + f"FlashInferAllReduce.forward: {self.mapping.rank} {self.mapping.tp_size} {self.max_num_tokens} {self.hidden_dim}" + ) def forward( self, @@ -794,8 +798,12 @@ def forward( f"FlashInferAllReduce.forward: {num_token} {hidden_dim}" ) - # if all_reduce_params is None: - + if all_reduce_params is None: + all_reduce_params = FlashInferAllReduceParams( + strategy=self.strategy, + fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, + config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, + ) flashinfer_comm.trtllm_custom_all_reduce( inp=input.ravel(), @@ -803,9 +811,9 @@ def forward( tp_size=self.mapping.tp_size, tp_rank=self.mapping.rank, token_num=num_token, - fusion_op_code=self.all_reduce_params.fusion_op, - strategy_code=self.all_reduce_params.strategy, - config_code=self.all_reduce_params.config_mode, + fusion_op_code=all_reduce_params.fusion_op, + strategy_code=all_reduce_params.strategy, + config_code=all_reduce_params.config_mode, launch_with_pdl=False, flag_value=self._flag_value, peer_comm_buffer_ptrs=torch.tensor(self.workspace[0], diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index f9793c765a8c..c656a6d81479 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -10,10 +10,8 @@ from transformers.modeling_utils import load_sharded_checkpoint from transformers.models.llama4.modeling_llama4 import Llama4MultiModalProjector -import flashinfer.comm as flashinfer_comm - from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp, - AllReduceParams, MoEAllReduce, FlashInferAllReduce, FlashInferAllReduceParams) + AllReduceParams, MoEAllReduce) from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ BaseWeightMapper from tensorrt_llm._utils import get_sm_version @@ -688,12 +686,6 @@ def __init__( or self.mapping.tp_size == 1 or self.enable_attention_dp) - self.flash_infer_all_reduce = FlashInferAllReduce( - mapping=self.mapping, - hidden_dim=8192, - strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, - dtype=config.torch_dtype) - def forward( self, position_ids: torch.IntTensor, @@ -714,23 +706,15 @@ def forward( attn_metadata=attn_metadata, attention_mask=self.attention_mask, all_reduce_params=AllReduceParams( - enable_allreduce=False), + enable_allreduce=not self.disable_attn_allreduce), **kwargs, ) - hidden_states = self.flash_infer_all_reduce( - hidden_states, - all_reduce_params=FlashInferAllReduceParams( - strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, - fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, - )) # Fully Connected if self.PRE_MLP_FUSION: if self.is_nvfp4 or self.is_fp8_quant: scale = self.mlp.gate_up_proj.input_scale else: scale = None - all_reduce_output = self.all_reduce( hidden_states, all_reduce_params=AllReduceParams( @@ -761,7 +745,7 @@ def forward( hidden_states = self.mlp( hidden_states, final_all_reduce_params=AllReduceParams( - enable_allreduce=False), + enable_allreduce=not self.disable_mlp_allreduce), **kwargs, ) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 1d2d455e19f1..fc83c75f0f20 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -15,7 +15,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AllReduceStrategy, FlashInferAllReduceParams) + AllReduceStrategy, FlashInferAllReduce) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.functional import \ preprocess_weights_for_mixed_gemm @@ -1824,8 +1824,7 @@ def __init__( use_cute_dsl_nvfp4_blockscaling_mm: bool = False, disable_deep_gemm: bool = False, ): - - from ..distributed import AllReduce + import flashinfer.comm as flashinfer_comm super().__init__() self.has_bias = bias @@ -1864,10 +1863,15 @@ def __init__( self.in_features = local_in_features self.out_features = local_out_features - # Needed for embedding layer - self.all_reduce = AllReduce(mapping=self.mapping, - strategy=allreduce_strategy, - dtype=self.dtype) if reduce_output else None + # self.all_reduce = AllReduce(mapping=self.mapping, + # strategy=allreduce_strategy, + # dtype=self.dtype) if reduce_output else None + + self.flash_infer_all_reduce = FlashInferAllReduce( + mapping=self.mapping, + hidden_dim=self.out_features, + strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, + dtype=self.dtype) if reduce_output else None self._weights_created = False self.reduce_output = reduce_output @@ -2003,7 +2007,11 @@ 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 = self.all_reduce( + # output, + # all_reduce_params=all_reduce_params, + # ) + output = self.flash_infer_all_reduce( output, all_reduce_params=all_reduce_params, ) diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index a3605a7b44bc..095e921ed36f 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -3990,7 +3990,7 @@ def __init__(self, flashinfer_comm.AllReduceFusionOp] = flashinfer_comm. AllReduceFusionOp.NONE, config_mode: Optional[ - flashinfer_comm.AllReduceStrategyConfig] = flashinfer_comm. + flashinfer_comm.AllReduceConfigMode] = flashinfer_comm. AllReduceStrategyConfig.USE_MEMCPY, bias: Optional[Tensor] = None, residual: Optional[Tensor] = None, From 5e09b599a85761c7912172908b4699939afe62a5 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 13 Aug 2025 10:01:37 -0700 Subject: [PATCH 06/21] working in linear; 256 token heuristic Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 52 +++++++++++------------- tensorrt_llm/_torch/modules/embedding.py | 1 + tensorrt_llm/_torch/modules/linear.py | 41 +++++++++++-------- tensorrt_llm/functional.py | 2 +- 4 files changed, 50 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index eca9bf4b4c98..db929806bf11 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -18,6 +18,8 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.plugin.plugin import CustomAllReduceHelper +from .communicator import TorchDist + try: import flashinfer.comm as flashinfer_comm except ImportError: @@ -752,9 +754,9 @@ def __init__(self, mapping: Mapping, global _RANK_DIST_INITIALIZED self.mapping = mapping - self.max_message_size = 70 * 1024 * 1024 # 70MB + # self.max_message_size = 70 * 1024 * 1024 # 70MB self.dtype = dtype - self.hidden_dim = hidden_dim + self.hidden_dim = 8192 self.strategy = strategy if not _RANK_DIST_INITIALIZED: @@ -764,22 +766,16 @@ def __init__(self, mapping: Mapping, # flashinfer all reduce is better for message sizes < ~70MB # num tokens * hidden dim * 2 < 70 *1024*1024 # TODO: 2 is for bf16,fp16. need to add fp8, fp4 - self.max_num_tokens = self.max_message_size // (self.hidden_dim * 2) - - self.workspace = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( - rank=self.mapping.rank, - tp_size=self.mapping.tp_size, - max_token_num=self.max_num_tokens, - hidden_dim=self.hidden_dim, - group=None, - ) + self.max_num_tokens = 8192 #self.max_message_size // (self.hidden_dim * 2) self._flag_value = 1 - self._enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" - if self._enable_debug: - print( - f"FlashInferAllReduce.forward: {self.mapping.rank} {self.mapping.tp_size} {self.max_num_tokens} {self.hidden_dim}" - ) + self.workspace = get_flashinfer_allreduce_workspace(self.mapping, self.max_num_tokens, self.hidden_dim) + + self.all_reduce_params = FlashInferAllReduceParams( + strategy=self.strategy, + fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, + config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, + ) def forward( self, @@ -792,28 +788,28 @@ def forward( num_token = input.size(0) hidden_dim = input.size(1) - output_buffer = torch.empty(num_token * hidden_dim, dtype=input.dtype, device=input.device) + output_buffer = torch.empty_like(input) if _enable_debug: print( f"FlashInferAllReduce.forward: {num_token} {hidden_dim}" ) - if all_reduce_params is None: - all_reduce_params = FlashInferAllReduceParams( - strategy=self.strategy, - fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, - ) + # if all_reduce_params is None: + # all_reduce_params = FlashInferAllReduceParams( + # strategy=self.strategy, + # fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, + # config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, + # ) flashinfer_comm.trtllm_custom_all_reduce( - inp=input.ravel(), + inp=input, out=output_buffer, tp_size=self.mapping.tp_size, tp_rank=self.mapping.rank, token_num=num_token, - fusion_op_code=all_reduce_params.fusion_op, - strategy_code=all_reduce_params.strategy, - config_code=all_reduce_params.config_mode, + fusion_op_code=self.all_reduce_params.fusion_op, + strategy_code=self.all_reduce_params.strategy, + config_code=self.all_reduce_params.config_mode, launch_with_pdl=False, flag_value=self._flag_value, peer_comm_buffer_ptrs=torch.tensor(self.workspace[0], @@ -843,4 +839,4 @@ def forward( ) self._flag_value += 1 - return output_buffer.reshape(num_token, hidden_dim) + return output_buffer 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 fc83c75f0f20..c5097465dc3b 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -15,7 +15,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AllReduceStrategy, FlashInferAllReduce) + AllReduceStrategy, FlashInferAllReduceParams) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.functional import \ preprocess_weights_for_mixed_gemm @@ -1823,7 +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 = False, ): + from ..distributed import AllReduce, FlashInferAllReduce import flashinfer.comm as flashinfer_comm super().__init__() @@ -1863,15 +1865,17 @@ def __init__( self.in_features = local_in_features self.out_features = local_out_features - # self.all_reduce = AllReduce(mapping=self.mapping, - # strategy=allreduce_strategy, - # dtype=self.dtype) if reduce_output else None + self.all_reduce = AllReduce(mapping=self.mapping, + strategy=allreduce_strategy, + dtype=self.dtype) if reduce_output else None - self.flash_infer_all_reduce = FlashInferAllReduce( - mapping=self.mapping, - hidden_dim=self.out_features, - strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, - dtype=self.dtype) if reduce_output else None + self.use_flashinfer_allreduce = use_flashinfer_allreduce + if use_flashinfer_allreduce: + self.flash_infer_all_reduce = FlashInferAllReduce( + mapping=self.mapping, + hidden_dim=self.out_features, + strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, + dtype=self.dtype) if reduce_output else None self._weights_created = False self.reduce_output = reduce_output @@ -2007,14 +2011,17 @@ 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, - # ) - output = self.flash_infer_all_reduce( - output, - all_reduce_params=all_reduce_params, - ) + + if self.use_flashinfer_allreduce and output.size(0) <= 256: + output = self.flash_infer_all_reduce( + output, + all_reduce_params=all_reduce_params, + ) + 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/functional.py b/tensorrt_llm/functional.py index 095e921ed36f..a3605a7b44bc 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -3990,7 +3990,7 @@ def __init__(self, flashinfer_comm.AllReduceFusionOp] = flashinfer_comm. AllReduceFusionOp.NONE, config_mode: Optional[ - flashinfer_comm.AllReduceConfigMode] = flashinfer_comm. + flashinfer_comm.AllReduceStrategyConfig] = flashinfer_comm. AllReduceStrategyConfig.USE_MEMCPY, bias: Optional[Tensor] = None, residual: Optional[Tensor] = None, From 20d16628d40467089f0482548ed6ac3c1f8997ef Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 13 Aug 2025 15:12:22 -0700 Subject: [PATCH 07/21] use global flag Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 52 ++++++++++++++++---------- tensorrt_llm/_torch/modules/linear.py | 4 +- 2 files changed, 35 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index db929806bf11..25b7bf618e62 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -744,6 +744,7 @@ def forward( ) _RANK_DIST_INITIALIZED = False +_GLOBAL_FLAG_VALUE = 1 class FlashInferAllReduce(nn.Module): @@ -767,22 +768,44 @@ def __init__(self, mapping: Mapping, # num tokens * hidden dim * 2 < 70 *1024*1024 # TODO: 2 is for bf16,fp16. need to add fp8, fp4 self.max_num_tokens = 8192 #self.max_message_size // (self.hidden_dim * 2) - self._flag_value = 1 self.workspace = get_flashinfer_allreduce_workspace(self.mapping, self.max_num_tokens, self.hidden_dim) self.all_reduce_params = FlashInferAllReduceParams( - strategy=self.strategy, + strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, ) + # each linear layer will have its own pointers to the workspace. + # The workspace is common + # Can we use common pointers? + self.peer_comm_buffer_ptrs = torch.tensor( + self.workspace[0], dtype=torch.int64 + ) + self.peer_barrier_ptrs_in = torch.tensor( + self.workspace[2], dtype=torch.int64 + ) + self.peer_barrier_ptrs_out = torch.tensor( + self.workspace[3], dtype=torch.int64 + ) + self.lamport_peer_comm_buffer_ptrs_0 = torch.tensor( + self.workspace[4], dtype=torch.int64 + ) + self.lamport_peer_comm_buffer_ptrs_1 = torch.tensor( + self.workspace[5], dtype=torch.int64 + ) + self.lamport_peer_comm_buffer_ptrs_2 = torch.tensor( + self.workspace[6], dtype=torch.int64 + ) + def forward( self, input: torch.Tensor, *, all_reduce_params: Optional[FlashInferAllReduceParams] = None, ) -> torch.Tensor: + global _GLOBAL_FLAG_VALUE input = input.contiguous() num_token = input.size(0) @@ -811,13 +834,10 @@ def forward( strategy_code=self.all_reduce_params.strategy, config_code=self.all_reduce_params.config_mode, launch_with_pdl=False, - flag_value=self._flag_value, - peer_comm_buffer_ptrs=torch.tensor(self.workspace[0], - dtype=torch.int64), - peer_barrier_ptrs_in=torch.tensor(self.workspace[2], - dtype=torch.int64), - peer_barrier_ptrs_out=torch.tensor(self.workspace[3], - dtype=torch.int64), + flag_value=_GLOBAL_FLAG_VALUE, + peer_comm_buffer_ptrs=self.peer_comm_buffer_ptrs, + peer_barrier_ptrs_in=self.peer_barrier_ptrs_in, + peer_barrier_ptrs_out=self.peer_barrier_ptrs_out, bias=None, residual=None, weight=None, @@ -827,16 +847,10 @@ def forward( lamport_peer_comm_buffer_ptrs_0=None, lamport_peer_comm_buffer_ptrs_1=None, lamport_peer_comm_buffer_ptrs_2=None, - # lamport_peer_comm_buffer_ptrs_0=torch.tensor( - # self.workspace[4], dtype=torch.int64 - # ), - # lamport_peer_comm_buffer_ptrs_1=torch.tensor( - # self.workspace[5], dtype=torch.int64 - # ), - # lamport_peer_comm_buffer_ptrs_2=torch.tensor( - # self.workspace[6], dtype=torch.int64 - # ), + # lamport_peer_comm_buffer_ptrs_0=self.lamport_peer_comm_buffer_ptrs_0, + # lamport_peer_comm_buffer_ptrs_1=self.lamport_peer_comm_buffer_ptrs_1, + # lamport_peer_comm_buffer_ptrs_2=self.lamport_peer_comm_buffer_ptrs_2, ) - self._flag_value += 1 + _GLOBAL_FLAG_VALUE += 1 return output_buffer diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index c5097465dc3b..051e1f5823e0 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -28,6 +28,7 @@ from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..utils import Fp4QuantizedTensor +import flashinfer.comm as flashinfer_comm class WeightMode(str, enum.Enum): # weight of a vanilla layer @@ -1826,7 +1827,6 @@ def __init__( use_flashinfer_allreduce: bool = False, ): from ..distributed import AllReduce, FlashInferAllReduce - import flashinfer.comm as flashinfer_comm super().__init__() self.has_bias = bias @@ -2012,7 +2012,7 @@ def forward( bias = None if fuse_bias else bias output = self.apply_linear(input, bias, lora_params, layer_idx) - if self.use_flashinfer_allreduce and output.size(0) <= 256: + if self.use_flashinfer_allreduce and output.size(0) <= 150: output = self.flash_infer_all_reduce( output, all_reduce_params=all_reduce_params, From 8084e901c8a101d47074c68484f6af13584653a3 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 29 Sep 2025 15:08:15 -0700 Subject: [PATCH 08/21] add vllm custom ar Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/__init__.py | 6 +- tensorrt_llm/_torch/distributed/ops.py | 276 ++++++++++++++------ tensorrt_llm/_torch/modules/linear.py | 9 +- 3 files changed, 210 insertions(+), 81 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/__init__.py b/tensorrt_llm/_torch/distributed/__init__.py index c3d0972f19db..8e1d3ce40173 100644 --- a/tensorrt_llm/_torch/distributed/__init__.py +++ b/tensorrt_llm/_torch/distributed/__init__.py @@ -2,8 +2,9 @@ from .communicator import Distributed, MPIDist, PPComm, TorchDist from .ops import (AllReduce, AllReduceParams, AllReduceStrategy, - FlashInferAllReduce, FlashInferAllReduceParams, MoEAllReduce, - MoEAllReduceParams, allgather, reducescatter, alltoall_helix, + FlashInferAllReduce, FlashInferAllReduceParams, + FlashInferVLLMAllReduce, MoEAllReduce, MoEAllReduceParams, + allgather, alltoall_helix, reducescatter, userbuffers_allreduce_finalize) __all__ = [ @@ -23,4 +24,5 @@ "Distributed", "FlashInferAllReduce", "FlashInferAllReduceParams", + "FlashInferVLLMAllReduce", ] diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 25b7bf618e62..7796de7ad7e5 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -5,8 +5,8 @@ from typing import List, Optional, Tuple, Union import torch -from torch import nn import torch.distributed as dist +from torch import nn from tensorrt_llm._utils import mpi_comm, mpi_disabled from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer @@ -33,20 +33,25 @@ _enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" -def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, hidden_dim: int): - if not hasattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}'): - setattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}', {}) - - flashinfer_allreduce_workspaces = getattr(_thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}') +def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, + hidden_dim: int): + if not hasattr(_thread_local, + f'flashinfer_allreduce_workspaces_{mapping.pp_rank}'): + setattr(_thread_local, + f'flashinfer_allreduce_workspaces_{mapping.pp_rank}', {}) + + flashinfer_allreduce_workspaces = getattr( + _thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}') if mapping not in flashinfer_allreduce_workspaces: - flashinfer_allreduce_workspaces[mapping] = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( - rank=mapping.rank, - tp_size=mapping.tp_size, - max_token_num=max_num_tokens, - hidden_dim=hidden_dim, - group=dist.group.WORLD, - ) + flashinfer_allreduce_workspaces[ + mapping] = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( + rank=mapping.rank, + tp_size=mapping.tp_size, + max_token_num=max_num_tokens, + hidden_dim=hidden_dim, + group=dist.group.WORLD, + ) if _enable_debug: print( f"FlashInferAllReduce.init: {mapping.rank} {mapping.tp_size} {max_num_tokens} {hidden_dim}" @@ -743,9 +748,11 @@ def forward( eps=all_reduce_params.eps, ) + _RANK_DIST_INITIALIZED = False _GLOBAL_FLAG_VALUE = 1 + class FlashInferAllReduce(nn.Module): def __init__(self, mapping: Mapping, @@ -755,49 +762,41 @@ def __init__(self, mapping: Mapping, global _RANK_DIST_INITIALIZED self.mapping = mapping - # self.max_message_size = 70 * 1024 * 1024 # 70MB self.dtype = dtype - self.hidden_dim = 8192 + self.hidden_dim = hidden_dim # Use the passed parameter instead of hardcoded value self.strategy = strategy if not _RANK_DIST_INITIALIZED: TorchDist(mapping) _RANK_DIST_INITIALIZED = True - # flashinfer all reduce is better for message sizes < ~70MB - # num tokens * hidden dim * 2 < 70 *1024*1024 - # TODO: 2 is for bf16,fp16. need to add fp8, fp4 - self.max_num_tokens = 8192 #self.max_message_size // (self.hidden_dim * 2) + self.max_num_tokens = 8192 - self.workspace = get_flashinfer_allreduce_workspace(self.mapping, self.max_num_tokens, self.hidden_dim) + self.workspace = get_flashinfer_allreduce_workspace( + self.mapping, self.max_num_tokens, self.hidden_dim) self.all_reduce_params = FlashInferAllReduceParams( - strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, + strategy= + strategy, # Use the passed strategy instead of hardcoded TWOSHOT fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, + # config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, ) # each linear layer will have its own pointers to the workspace. # The workspace is common # Can we use common pointers? - self.peer_comm_buffer_ptrs = torch.tensor( - self.workspace[0], dtype=torch.int64 - ) - self.peer_barrier_ptrs_in = torch.tensor( - self.workspace[2], dtype=torch.int64 - ) - self.peer_barrier_ptrs_out = torch.tensor( - self.workspace[3], dtype=torch.int64 - ) - self.lamport_peer_comm_buffer_ptrs_0 = torch.tensor( - self.workspace[4], dtype=torch.int64 - ) - self.lamport_peer_comm_buffer_ptrs_1 = torch.tensor( - self.workspace[5], dtype=torch.int64 - ) - self.lamport_peer_comm_buffer_ptrs_2 = torch.tensor( - self.workspace[6], dtype=torch.int64 - ) + self.peer_comm_buffer_ptrs = torch.tensor(self.workspace[0], + dtype=torch.int64) + self.peer_barrier_ptrs_in = torch.tensor(self.workspace[2], + dtype=torch.int64) + self.peer_barrier_ptrs_out = torch.tensor(self.workspace[3], + dtype=torch.int64) + self.lamport_peer_comm_buffer_ptrs_0 = torch.tensor(self.workspace[4], + dtype=torch.int64) + self.lamport_peer_comm_buffer_ptrs_1 = torch.tensor(self.workspace[5], + dtype=torch.int64) + self.lamport_peer_comm_buffer_ptrs_2 = torch.tensor(self.workspace[6], + dtype=torch.int64) def forward( self, @@ -811,46 +810,171 @@ def forward( num_token = input.size(0) hidden_dim = input.size(1) + # Validate input dimensions + if hidden_dim != self.hidden_dim: + logger.warning( + f"Input hidden_dim {hidden_dim} doesn't match expected {self.hidden_dim}" + ) + + if num_token > self.max_num_tokens: + logger.warning( + f"Input num_token {num_token} exceeds max_num_tokens {self.max_num_tokens}" + ) + output_buffer = torch.empty_like(input) if _enable_debug: - print( - f"FlashInferAllReduce.forward: {num_token} {hidden_dim}" + print(f"FlashInferAllReduce.forward: {num_token} {hidden_dim}") + + # Use provided params or default + if all_reduce_params is None: + all_reduce_params = self.all_reduce_params + + # Validate required parameters + if not hasattr(all_reduce_params, + 'fusion_op') or all_reduce_params.fusion_op is None: + logger.warning("fusion_op is None, using NONE as default") + all_reduce_params.fusion_op = flashinfer_comm.AllReduceFusionOp.NONE + + if not hasattr(all_reduce_params, + 'strategy') or all_reduce_params.strategy is None: + logger.warning("strategy is None, using TWOSHOT as default") + all_reduce_params.strategy = flashinfer_comm.AllReduceStrategyType.TWOSHOT + + if not hasattr(all_reduce_params, + 'config_mode') or all_reduce_params.config_mode is None: + logger.warning("config_mode is None, using USE_MEMCPY as default") + all_reduce_params.config_mode = flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY + + try: + flashinfer_comm.trtllm_custom_all_reduce( + inp=input, + out=output_buffer, + tp_size=self.mapping.tp_size, + tp_rank=self.mapping.rank, + token_num=num_token, + fusion_op_code=all_reduce_params.fusion_op, + strategy_code=all_reduce_params.strategy, + config_code=all_reduce_params.config_mode, + launch_with_pdl=False, + flag_value=_GLOBAL_FLAG_VALUE, + peer_comm_buffer_ptrs=self.peer_comm_buffer_ptrs, + peer_barrier_ptrs_in=self.peer_barrier_ptrs_in, + peer_barrier_ptrs_out=self.peer_barrier_ptrs_out, + bias=all_reduce_params.bias, + residual=all_reduce_params.residual, + weight=all_reduce_params.norm_weight, + weight_pre_residual_norm=all_reduce_params. + norm_pre_residual_weight, + eps=all_reduce_params.eps, + intermediate_buffer=None, + # Enable Lamport communication buffers for better synchronization + lamport_peer_comm_buffer_ptrs_0=self. + lamport_peer_comm_buffer_ptrs_0, + lamport_peer_comm_buffer_ptrs_1=self. + lamport_peer_comm_buffer_ptrs_1, + lamport_peer_comm_buffer_ptrs_2=self. + lamport_peer_comm_buffer_ptrs_2, ) + except Exception as e: + logger.error(f"FlashInferAllReduce failed: {e}") + raise - # if all_reduce_params is None: - # all_reduce_params = FlashInferAllReduceParams( - # strategy=self.strategy, - # fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - # config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, - # ) - - flashinfer_comm.trtllm_custom_all_reduce( - inp=input, - out=output_buffer, - tp_size=self.mapping.tp_size, - tp_rank=self.mapping.rank, - token_num=num_token, - fusion_op_code=self.all_reduce_params.fusion_op, - strategy_code=self.all_reduce_params.strategy, - config_code=self.all_reduce_params.config_mode, - launch_with_pdl=False, - flag_value=_GLOBAL_FLAG_VALUE, - peer_comm_buffer_ptrs=self.peer_comm_buffer_ptrs, - peer_barrier_ptrs_in=self.peer_barrier_ptrs_in, - peer_barrier_ptrs_out=self.peer_barrier_ptrs_out, - bias=None, - residual=None, - weight=None, - weight_pre_residual_norm=None, - eps=None, - intermediate_buffer=None, - lamport_peer_comm_buffer_ptrs_0=None, - lamport_peer_comm_buffer_ptrs_1=None, - lamport_peer_comm_buffer_ptrs_2=None, - # lamport_peer_comm_buffer_ptrs_0=self.lamport_peer_comm_buffer_ptrs_0, - # lamport_peer_comm_buffer_ptrs_1=self.lamport_peer_comm_buffer_ptrs_1, - # lamport_peer_comm_buffer_ptrs_2=self.lamport_peer_comm_buffer_ptrs_2, - ) _GLOBAL_FLAG_VALUE += 1 return output_buffer + + +class FlashInferVLLMAllReduce(nn.Module): + + def __init__(self, mapping: Mapping, dtype: torch.dtype): + super().__init__() + global _RANK_DIST_INITIALIZED + + self.mapping = mapping + self.dtype = dtype + + if not _RANK_DIST_INITIALIZED: + TorchDist(mapping) + _RANK_DIST_INITIALIZED = True + + # Create IPC workspace for vLLM allreduce + max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 + self.meta_ptrs = flashinfer_comm.create_shared_buffer( + flashinfer_comm.vllm_meta_size() + max_size, dist.group.WORLD) + + # 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 = flashinfer_comm.create_shared_buffer( + max_size, dist.group.WORLD) + + # Initialize custom allreduce + self.fa = flashinfer_comm.vllm_init_custom_ar( + ipc_tensors=self.meta_ptrs, + rank_data=self.rank_data, + rank=mapping.rank, + full_nvlink=True) + + # Register buffer - this is crucial! + flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs) + + # Create registration buffer for allreduce operations + self.reg_buffer_size = max_size + + 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 + """ + input_tensor = input.contiguous() + output_tensor = torch.empty_like(input_tensor) + + try: + # Perform vLLM custom allreduce + flashinfer_comm.vllm_all_reduce( + self.fa, + input_tensor, + output_tensor, + self.buffer_ptrs[self.mapping.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 destroy(self): + """Clean up vLLM resources""" + try: + # Clean up vLLM resources + flashinfer_comm.vllm_dispose(self.fa) + + # Free shared buffers + flashinfer_comm.free_shared_buffer(self.buffer_ptrs, + dist.group.WORLD) + flashinfer_comm.free_shared_buffer(self.meta_ptrs, dist.group.WORLD) + except Exception as e: + logger.error(f"Error during FlashInferVLLMAllReduce cleanup: {e}") + + def __del__(self): + """Destructor to ensure cleanup""" + try: + self.destroy() + except: + pass diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 051e1f5823e0..06005965fa35 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from typing import Dict, List, Optional, Union +import flashinfer.comm as flashinfer_comm import torch import torch.nn.functional as F from torch import nn @@ -15,7 +16,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AllReduceStrategy, FlashInferAllReduceParams) + AllReduceStrategy) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.functional import \ preprocess_weights_for_mixed_gemm @@ -28,7 +29,6 @@ from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..utils import Fp4QuantizedTensor -import flashinfer.comm as flashinfer_comm class WeightMode(str, enum.Enum): # weight of a vanilla layer @@ -1876,6 +1876,9 @@ def __init__( hidden_dim=self.out_features, strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, dtype=self.dtype) if reduce_output else None + # self.flash_infer_vllm_all_reduce = FlashInferVLLMAllReduce( + # mapping=self.mapping, + # dtype=self.dtype) if reduce_output else None self._weights_created = False self.reduce_output = reduce_output @@ -2012,7 +2015,7 @@ def forward( bias = None if fuse_bias else bias output = self.apply_linear(input, bias, lora_params, layer_idx) - if self.use_flashinfer_allreduce and output.size(0) <= 150: + if self.use_flashinfer_allreduce and output.size(0) == 150: output = self.flash_infer_all_reduce( output, all_reduce_params=all_reduce_params, From cd2b7bb90545b6ea0cd7b9db7bf9fe7430b0466d Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 1 Oct 2025 12:28:19 -0700 Subject: [PATCH 09/21] add env flags Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 18 +-------------- tensorrt_llm/_torch/modules/linear.py | 31 +++++++++++++++++++------- 2 files changed, 24 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 7796de7ad7e5..19ac2853a0ee 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -829,22 +829,6 @@ def forward( if all_reduce_params is None: all_reduce_params = self.all_reduce_params - # Validate required parameters - if not hasattr(all_reduce_params, - 'fusion_op') or all_reduce_params.fusion_op is None: - logger.warning("fusion_op is None, using NONE as default") - all_reduce_params.fusion_op = flashinfer_comm.AllReduceFusionOp.NONE - - if not hasattr(all_reduce_params, - 'strategy') or all_reduce_params.strategy is None: - logger.warning("strategy is None, using TWOSHOT as default") - all_reduce_params.strategy = flashinfer_comm.AllReduceStrategyType.TWOSHOT - - if not hasattr(all_reduce_params, - 'config_mode') or all_reduce_params.config_mode is None: - logger.warning("config_mode is None, using USE_MEMCPY as default") - all_reduce_params.config_mode = flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY - try: flashinfer_comm.trtllm_custom_all_reduce( inp=input, @@ -855,7 +839,7 @@ def forward( fusion_op_code=all_reduce_params.fusion_op, strategy_code=all_reduce_params.strategy, config_code=all_reduce_params.config_mode, - launch_with_pdl=False, + launch_with_pdl=True, flag_value=_GLOBAL_FLAG_VALUE, peer_comm_buffer_ptrs=self.peer_comm_buffer_ptrs, peer_barrier_ptrs_in=self.peer_barrier_ptrs_in, diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 06005965fa35..cb0f2fea5500 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1824,9 +1824,10 @@ 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 = False, + use_flashinfer_allreduce: bool = True, ): - from ..distributed import AllReduce, FlashInferAllReduce + from ..distributed import (AllReduce, FlashInferAllReduce, + FlashInferVLLMAllReduce) super().__init__() self.has_bias = bias @@ -1870,15 +1871,24 @@ def __init__( dtype=self.dtype) if reduce_output else None self.use_flashinfer_allreduce = use_flashinfer_allreduce - if use_flashinfer_allreduce: + self.flashinfer_trtllm = self.use_flashinfer_allreduce and os.getenv( + "_USE_FLASHINFER_TRTLLM_ALLREDUCE", "0") == "1" + self.flashinfer_vllm = self.use_flashinfer_allreduce and os.getenv( + "_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1" + print(f"flashinfer_trtllm: {self.flashinfer_trtllm}") + print(f"flashinfer_vllm: {self.flashinfer_vllm}") + + if self.flashinfer_trtllm: self.flash_infer_all_reduce = FlashInferAllReduce( mapping=self.mapping, hidden_dim=self.out_features, strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, dtype=self.dtype) if reduce_output else None - # self.flash_infer_vllm_all_reduce = FlashInferVLLMAllReduce( - # mapping=self.mapping, - # dtype=self.dtype) if reduce_output else None + + 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 @@ -2015,10 +2025,15 @@ def forward( bias = None if fuse_bias else bias output = self.apply_linear(input, bias, lora_params, layer_idx) - if self.use_flashinfer_allreduce and output.size(0) == 150: + if self.flashinfer_trtllm and output.size(0) == 150: output = self.flash_infer_all_reduce( output, - all_reduce_params=all_reduce_params, + all_reduce_params=None, + ) + elif self.flashinfer_vllm and output.size(0) == 150: + output = self.flash_infer_all_reduce( + output, + all_reduce_params=None, ) else: output = self.all_reduce( From 8b0ece89eb68e736e47db8708377ee3b7ea76b26 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 1 Oct 2025 15:38:38 -0700 Subject: [PATCH 10/21] only initialize once per mapping Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 83 ++++++++++++++++---------- tensorrt_llm/_torch/modules/linear.py | 4 +- 2 files changed, 54 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 19ac2853a0ee..fadf6e7b1210 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -2,6 +2,7 @@ import os import platform import threading +from types import SimpleNamespace from typing import List, Optional, Tuple, Union import torch @@ -59,6 +60,44 @@ def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, return flashinfer_allreduce_workspaces[mapping] +def get_flashinfer_vllm_allreduce_workspace(mapping: Mapping): + if not hasattr(_thread_local, + f'flashinfer_vllm_allreduce_workspaces_{mapping.pp_rank}'): + setattr(_thread_local, + f'flashinfer_vllm_allreduce_workspaces_{mapping.pp_rank}', {}) + + flashinfer_vllm_allreduce_workspaces = getattr( + _thread_local, + f'flashinfer_vllm_allreduce_workspaces_{mapping.pp_rank}') + + if mapping not in flashinfer_vllm_allreduce_workspaces: + max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 + meta_ptrs = flashinfer_comm.create_shared_buffer( + flashinfer_comm.vllm_meta_size() + max_size, dist.group.WORLD) + + # Create rank data buffer (8MB as in test) + rank_data = torch.empty(8 * 1024 * 1024, + dtype=torch.uint8, + device=f"cuda:{mapping.local_rank}") + + # Create buffer pointers for IPC communication + buffer_ptrs = flashinfer_comm.create_shared_buffer( + max_size, dist.group.WORLD) + + # Initialize custom allreduce + fa = flashinfer_comm.vllm_init_custom_ar(ipc_tensors=meta_ptrs, + rank_data=rank_data, + rank=mapping.rank, + full_nvlink=True) + + # Register buffer - this is crucial! + flashinfer_comm.vllm_register_buffer(fa, buffer_ptrs) + flashinfer_vllm_allreduce_workspaces[mapping] = SimpleNamespace( + fa=fa, buffer_ptrs=buffer_ptrs, meta_ptrs=meta_ptrs) + + return flashinfer_vllm_allreduce_workspaces[mapping] + + def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): setattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}', {}) @@ -883,27 +922,8 @@ def __init__(self, mapping: Mapping, dtype: torch.dtype): # Create IPC workspace for vLLM allreduce max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - self.meta_ptrs = flashinfer_comm.create_shared_buffer( - flashinfer_comm.vllm_meta_size() + max_size, dist.group.WORLD) - - # 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 = flashinfer_comm.create_shared_buffer( - max_size, dist.group.WORLD) - # Initialize custom allreduce - self.fa = flashinfer_comm.vllm_init_custom_ar( - ipc_tensors=self.meta_ptrs, - rank_data=self.rank_data, - rank=mapping.rank, - full_nvlink=True) - - # Register buffer - this is crucial! - flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs) + self.workspace = get_flashinfer_vllm_allreduce_workspace(mapping) # Create registration buffer for allreduce operations self.reg_buffer_size = max_size @@ -930,10 +950,10 @@ def forward( try: # Perform vLLM custom allreduce flashinfer_comm.vllm_all_reduce( - self.fa, + self.workspace.fa, input_tensor, output_tensor, - self.buffer_ptrs[self.mapping.rank], + self.workspace.buffer_ptrs[self.mapping.rank], self.reg_buffer_size, 36, # CTA upper bounds: 36 as mentioned in the API ) @@ -947,18 +967,19 @@ def destroy(self): """Clean up vLLM resources""" try: # Clean up vLLM resources - flashinfer_comm.vllm_dispose(self.fa) + flashinfer_comm.vllm_dispose(self.workspace.fa) # Free shared buffers - flashinfer_comm.free_shared_buffer(self.buffer_ptrs, + flashinfer_comm.free_shared_buffer(self.workspace.buffer_ptrs, + dist.group.WORLD) + flashinfer_comm.free_shared_buffer(self.workspace.meta_ptrs, dist.group.WORLD) - flashinfer_comm.free_shared_buffer(self.meta_ptrs, dist.group.WORLD) except Exception as e: logger.error(f"Error during FlashInferVLLMAllReduce cleanup: {e}") - def __del__(self): - """Destructor to ensure cleanup""" - try: - self.destroy() - except: - pass + # def __del__(self): + # """Destructor to ensure cleanup""" + # try: + # self.destroy() + # except: + # pass diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index cb0f2fea5500..44731853a178 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1875,8 +1875,8 @@ def __init__( "_USE_FLASHINFER_TRTLLM_ALLREDUCE", "0") == "1" self.flashinfer_vllm = self.use_flashinfer_allreduce and os.getenv( "_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1" - print(f"flashinfer_trtllm: {self.flashinfer_trtllm}") - print(f"flashinfer_vllm: {self.flashinfer_vllm}") + # print(f"flashinfer_trtllm: {self.flashinfer_trtllm}") + # print(f"flashinfer_vllm: {self.flashinfer_vllm}") if self.flashinfer_trtllm: self.flash_infer_all_reduce = FlashInferAllReduce( From c7c1ba40f65a4d69f6875aafaeecc6b04e946ae0 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 1 Oct 2025 16:58:12 -0700 Subject: [PATCH 11/21] refactor Signed-off-by: Shreyas Misra --- .../_torch/distributed/communicator.py | 42 ++++++++++ tensorrt_llm/_torch/distributed/ops.py | 82 ++----------------- .../_torch/pyexecutor/model_engine.py | 4 +- 3 files changed, 51 insertions(+), 77 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 07f2b4227f99..4dc508f4c161 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -15,6 +15,8 @@ except Exception: MPI = None # deferred; functions will error if used when ENABLE_MULTI_DEVICE is True +import flashinfer.comm as flashinfer_comm + from tensorrt_llm._utils import (mpi_allgather, mpi_barrier, mpi_comm, mpi_disabled, mpi_isend, mpi_isend_object, mpi_recv, mpi_recv_object, mpi_send, @@ -429,6 +431,32 @@ def wait(self): raise RuntimeError(f"Asynchronous operation failed: {e}") from e +class FlashInferVLLMComm: + + def __init__(self, mapping: Mapping, group): + max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 + self.meta_ptrs = flashinfer_comm.create_shared_buffer( + flashinfer_comm.vllm_meta_size() + max_size, group) + + # 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 = flashinfer_comm.create_shared_buffer(max_size, group) + + # Initialize custom allreduce + self.fa = flashinfer_comm.vllm_init_custom_ar( + ipc_tensors=self.meta_ptrs, + rank_data=self.rank_data, + rank=mapping.rank, + full_nvlink=True) + + # Register buffer - this is crucial! + flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs) + + class TorchDist(Distributed): @property @@ -452,6 +480,9 @@ def __init__(self, mapping: Mapping): init_pg(torch.distributed.group.WORLD, self.local_comm, torch_pybind11_abi()) + + self._flashinfer_vllm_comm = FlashInferVLLMComm(mapping, + self.cpu_tp_group) def setup_local_comm(self): self._get_cluster_info() @@ -812,6 +843,7 @@ def recv(self, tensor: torch.Tensor, src: Optional[int] = None): _pp_comm = None +_torch_dist_tp_comm = None def init_pp_comm(mapping): @@ -833,3 +865,13 @@ def pp_recv(tensor): def pp_send(tensor): """Send tensors to next pp rank.""" _pp_comm.send(tensor) + + +def init_torch_dist_tp_comm(mapping): + """Initialize TorchDistTPComm once at startup""" + global _torch_dist_tp_comm + _torch_dist_tp_comm = TorchDist(mapping) + + +def get_torch_dist_flashinfer_vllm_comm(): + return _torch_dist_tp_comm._flashinfer_vllm_comm diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index fadf6e7b1210..b3ae8a3288d0 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -2,7 +2,6 @@ import os import platform import threading -from types import SimpleNamespace from typing import List, Optional, Tuple, Union import torch @@ -19,7 +18,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from .communicator import TorchDist +from .communicator import TorchDist, get_torch_dist_flashinfer_vllm_comm try: import flashinfer.comm as flashinfer_comm @@ -60,44 +59,6 @@ def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, return flashinfer_allreduce_workspaces[mapping] -def get_flashinfer_vllm_allreduce_workspace(mapping: Mapping): - if not hasattr(_thread_local, - f'flashinfer_vllm_allreduce_workspaces_{mapping.pp_rank}'): - setattr(_thread_local, - f'flashinfer_vllm_allreduce_workspaces_{mapping.pp_rank}', {}) - - flashinfer_vllm_allreduce_workspaces = getattr( - _thread_local, - f'flashinfer_vllm_allreduce_workspaces_{mapping.pp_rank}') - - if mapping not in flashinfer_vllm_allreduce_workspaces: - max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - meta_ptrs = flashinfer_comm.create_shared_buffer( - flashinfer_comm.vllm_meta_size() + max_size, dist.group.WORLD) - - # Create rank data buffer (8MB as in test) - rank_data = torch.empty(8 * 1024 * 1024, - dtype=torch.uint8, - device=f"cuda:{mapping.local_rank}") - - # Create buffer pointers for IPC communication - buffer_ptrs = flashinfer_comm.create_shared_buffer( - max_size, dist.group.WORLD) - - # Initialize custom allreduce - fa = flashinfer_comm.vllm_init_custom_ar(ipc_tensors=meta_ptrs, - rank_data=rank_data, - rank=mapping.rank, - full_nvlink=True) - - # Register buffer - this is crucial! - flashinfer_comm.vllm_register_buffer(fa, buffer_ptrs) - flashinfer_vllm_allreduce_workspaces[mapping] = SimpleNamespace( - fa=fa, buffer_ptrs=buffer_ptrs, meta_ptrs=meta_ptrs) - - return flashinfer_vllm_allreduce_workspaces[mapping] - - def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): setattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}', {}) @@ -911,22 +872,10 @@ class FlashInferVLLMAllReduce(nn.Module): def __init__(self, mapping: Mapping, dtype: torch.dtype): super().__init__() - global _RANK_DIST_INITIALIZED - self.mapping = mapping self.dtype = dtype - if not _RANK_DIST_INITIALIZED: - TorchDist(mapping) - _RANK_DIST_INITIALIZED = True - - # Create IPC workspace for vLLM allreduce - max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - - self.workspace = get_flashinfer_vllm_allreduce_workspace(mapping) - - # Create registration buffer for allreduce operations - self.reg_buffer_size = max_size + self.reg_buffer_size = 8192 * 8192 * 2 def forward( self, @@ -947,13 +896,15 @@ def forward( input_tensor = input.contiguous() output_tensor = torch.empty_like(input_tensor) + workspace = get_torch_dist_flashinfer_vllm_comm() + try: # Perform vLLM custom allreduce flashinfer_comm.vllm_all_reduce( - self.workspace.fa, + workspace.fa, input_tensor, output_tensor, - self.workspace.buffer_ptrs[self.mapping.rank], + workspace.buffer_ptrs[self.mapping.rank], self.reg_buffer_size, 36, # CTA upper bounds: 36 as mentioned in the API ) @@ -962,24 +913,3 @@ def forward( raise return output_tensor - - def destroy(self): - """Clean up vLLM resources""" - try: - # Clean up vLLM resources - flashinfer_comm.vllm_dispose(self.workspace.fa) - - # Free shared buffers - flashinfer_comm.free_shared_buffer(self.workspace.buffer_ptrs, - dist.group.WORLD) - flashinfer_comm.free_shared_buffer(self.workspace.meta_ptrs, - dist.group.WORLD) - except Exception as e: - logger.error(f"Error during FlashInferVLLMAllReduce cleanup: {e}") - - # def __del__(self): - # """Destructor to ensure cleanup""" - # try: - # self.destroy() - # except: - # pass diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 95a04e5c582e..dd1eadb5a454 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -32,7 +32,7 @@ from ..compilation.backend import Backend from ..compilation.utils import capture_piecewise_cuda_graph from ..distributed import MPIDist -from ..distributed.communicator import init_pp_comm +from ..distributed.communicator import init_pp_comm, init_torch_dist_tp_comm from ..expert_statistic import ExpertStatistic from ..metadata import KVCacheParams from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader @@ -153,6 +153,8 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) + if mapping.has_tp(): + init_torch_dist_tp_comm(mapping) self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) From 0eb4e7a7b066b924abbdc15e8499918ab5223336 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 2 Oct 2025 12:06:03 -0700 Subject: [PATCH 12/21] cleanup Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/communicator.py | 11 ++++++++--- tensorrt_llm/_torch/distributed/ops.py | 7 +++---- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 4dc508f4c161..385b7c389205 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from functools import wraps from typing import Optional +import os import numpy as np import torch @@ -441,7 +442,8 @@ def __init__(self, mapping: Mapping, group): # 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}") + device=f"cuda:{mapping.local_rank}", + pin_memory=True) # Create buffer pointers for IPC communication self.buffer_ptrs = flashinfer_comm.create_shared_buffer(max_size, group) @@ -481,8 +483,11 @@ def __init__(self, mapping: Mapping): init_pg(torch.distributed.group.WORLD, self.local_comm, torch_pybind11_abi()) - self._flashinfer_vllm_comm = FlashInferVLLMComm(mapping, - self.cpu_tp_group) + if os.getenv("_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1": + self._flashinfer_vllm_comm = FlashInferVLLMComm( + mapping, self.cpu_tp_group) + else: + self._flashinfer_vllm_comm = None def setup_local_comm(self): self._get_cluster_info() diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index b3ae8a3288d0..350afb00faba 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -876,6 +876,7 @@ def __init__(self, mapping: Mapping, dtype: torch.dtype): self.dtype = dtype self.reg_buffer_size = 8192 * 8192 * 2 + self.workspace = get_torch_dist_flashinfer_vllm_comm() def forward( self, @@ -896,15 +897,13 @@ def forward( input_tensor = input.contiguous() output_tensor = torch.empty_like(input_tensor) - workspace = get_torch_dist_flashinfer_vllm_comm() - try: # Perform vLLM custom allreduce flashinfer_comm.vllm_all_reduce( - workspace.fa, + self.workspace.fa, input_tensor, output_tensor, - workspace.buffer_ptrs[self.mapping.rank], + self.workspace.buffer_ptrs[self.mapping.rank], self.reg_buffer_size, 36, # CTA upper bounds: 36 as mentioned in the API ) From 103a801ae71fff05d0536bd713b467c44c1285e2 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 2 Oct 2025 16:49:10 -0700 Subject: [PATCH 13/21] small fixes Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/communicator.py | 10 ++++++---- tensorrt_llm/_torch/modules/linear.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 385b7c389205..70645b6b0b02 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -442,8 +442,7 @@ def __init__(self, mapping: Mapping, group): # 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}", - pin_memory=True) + device=f"cuda:{mapping.local_rank}") # Create buffer pointers for IPC communication self.buffer_ptrs = flashinfer_comm.create_shared_buffer(max_size, group) @@ -875,8 +874,11 @@ def pp_send(tensor): def init_torch_dist_tp_comm(mapping): """Initialize TorchDistTPComm once at startup""" global _torch_dist_tp_comm - _torch_dist_tp_comm = TorchDist(mapping) + if _torch_dist_tp_comm is None: + _torch_dist_tp_comm = TorchDist(mapping) def get_torch_dist_flashinfer_vllm_comm(): - return _torch_dist_tp_comm._flashinfer_vllm_comm + if _torch_dist_tp_comm is not None: + return _torch_dist_tp_comm._flashinfer_vllm_comm + return None diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 44731853a178..4233f1b79078 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -2030,7 +2030,7 @@ def forward( output, all_reduce_params=None, ) - elif self.flashinfer_vllm and output.size(0) == 150: + elif self.flashinfer_vllm and output.size(0) <= 256: output = self.flash_infer_all_reduce( output, all_reduce_params=None, From 9f7eaabb755a170a2362e8f36ccba9f1e7e599ab Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 7 Oct 2025 10:31:17 -0700 Subject: [PATCH 14/21] add token threshold env var Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/modules/linear.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 4233f1b79078..0a52f8f320c9 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1875,6 +1875,8 @@ def __init__( "_USE_FLASHINFER_TRTLLM_ALLREDUCE", "0") == "1" 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")) # print(f"flashinfer_trtllm: {self.flashinfer_trtllm}") # print(f"flashinfer_vllm: {self.flashinfer_vllm}") @@ -2030,7 +2032,7 @@ def forward( output, all_reduce_params=None, ) - elif self.flashinfer_vllm and output.size(0) <= 256: + elif self.flashinfer_vllm and output.size(0) <= self.flashinfer_token_threshold: output = self.flash_infer_all_reduce( output, all_reduce_params=None, From 419a3c761fa6dac738343ae9ee7d57e9c8bccd71 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 8 Oct 2025 09:31:16 -0700 Subject: [PATCH 15/21] fix pre commit Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/communicator.py | 4 ++-- tensorrt_llm/_torch/modules/linear.py | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 70645b6b0b02..73f52661acd8 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -1,9 +1,9 @@ import math +import os import pickle # nosec B403 from abc import ABC, abstractmethod from functools import wraps from typing import Optional -import os import numpy as np import torch @@ -481,7 +481,7 @@ def __init__(self, mapping: Mapping): init_pg(torch.distributed.group.WORLD, self.local_comm, torch_pybind11_abi()) - + if os.getenv("_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1": self._flashinfer_vllm_comm = FlashInferVLLMComm( mapping, self.cpu_tp_group) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 0a52f8f320c9..fb8766082748 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1875,8 +1875,8 @@ def __init__( "_USE_FLASHINFER_TRTLLM_ALLREDUCE", "0") == "1" 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")) + self.flashinfer_token_threshold = int( + os.getenv("_FLASHINFER_TOKEN_THRESHOLD", "256")) # print(f"flashinfer_trtllm: {self.flashinfer_trtllm}") # print(f"flashinfer_vllm: {self.flashinfer_vllm}") @@ -2032,7 +2032,8 @@ def forward( output, all_reduce_params=None, ) - elif self.flashinfer_vllm and output.size(0) <= self.flashinfer_token_threshold: + elif self.flashinfer_vllm and output.size( + 0) <= self.flashinfer_token_threshold: output = self.flash_infer_all_reduce( output, all_reduce_params=None, From a1cf28f9581381fbf95c6bf2f96cb9dcdefac21c Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 8 Oct 2025 10:17:42 -0700 Subject: [PATCH 16/21] support mpi and torch dist; cleanup Signed-off-by: Shreyas Misra --- .../_torch/distributed/communicator.py | 49 ------ tensorrt_llm/_torch/distributed/ops.py | 152 +----------------- tensorrt_llm/_torch/flashinfer_utils.py | 90 +++++++++++ tensorrt_llm/_torch/modules/linear.py | 22 +-- .../_torch/pyexecutor/model_engine.py | 5 +- .../_torch/pyexecutor/py_executor_creator.py | 5 + 6 files changed, 101 insertions(+), 222 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 73f52661acd8..07f2b4227f99 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -1,5 +1,4 @@ import math -import os import pickle # nosec B403 from abc import ABC, abstractmethod from functools import wraps @@ -16,8 +15,6 @@ except Exception: MPI = None # deferred; functions will error if used when ENABLE_MULTI_DEVICE is True -import flashinfer.comm as flashinfer_comm - from tensorrt_llm._utils import (mpi_allgather, mpi_barrier, mpi_comm, mpi_disabled, mpi_isend, mpi_isend_object, mpi_recv, mpi_recv_object, mpi_send, @@ -432,32 +429,6 @@ def wait(self): raise RuntimeError(f"Asynchronous operation failed: {e}") from e -class FlashInferVLLMComm: - - def __init__(self, mapping: Mapping, group): - max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - self.meta_ptrs = flashinfer_comm.create_shared_buffer( - flashinfer_comm.vllm_meta_size() + max_size, group) - - # 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 = flashinfer_comm.create_shared_buffer(max_size, group) - - # Initialize custom allreduce - self.fa = flashinfer_comm.vllm_init_custom_ar( - ipc_tensors=self.meta_ptrs, - rank_data=self.rank_data, - rank=mapping.rank, - full_nvlink=True) - - # Register buffer - this is crucial! - flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs) - - class TorchDist(Distributed): @property @@ -482,12 +453,6 @@ def __init__(self, mapping: Mapping): init_pg(torch.distributed.group.WORLD, self.local_comm, torch_pybind11_abi()) - if os.getenv("_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1": - self._flashinfer_vllm_comm = FlashInferVLLMComm( - mapping, self.cpu_tp_group) - else: - self._flashinfer_vllm_comm = None - def setup_local_comm(self): self._get_cluster_info() @@ -847,7 +812,6 @@ def recv(self, tensor: torch.Tensor, src: Optional[int] = None): _pp_comm = None -_torch_dist_tp_comm = None def init_pp_comm(mapping): @@ -869,16 +833,3 @@ def pp_recv(tensor): def pp_send(tensor): """Send tensors to next pp rank.""" _pp_comm.send(tensor) - - -def init_torch_dist_tp_comm(mapping): - """Initialize TorchDistTPComm once at startup""" - global _torch_dist_tp_comm - if _torch_dist_tp_comm is None: - _torch_dist_tp_comm = TorchDist(mapping) - - -def get_torch_dist_flashinfer_vllm_comm(): - if _torch_dist_tp_comm is not None: - return _torch_dist_tp_comm._flashinfer_vllm_comm - return None diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 350afb00faba..39b5f3695a07 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -5,7 +5,6 @@ from typing import List, Optional, Tuple, Union import torch -import torch.distributed as dist from torch import nn from tensorrt_llm._utils import mpi_comm, mpi_disabled @@ -18,7 +17,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from .communicator import TorchDist, get_torch_dist_flashinfer_vllm_comm +from ..flashinfer_utils import current_flashinfer_allreduce_workspace try: import flashinfer.comm as flashinfer_comm @@ -30,34 +29,6 @@ _thread_local = threading.local() -_enable_debug = os.environ.get("_FLASHINFER_DEBUG", "0") == "1" - - -def get_flashinfer_allreduce_workspace(mapping: Mapping, max_num_tokens: int, - hidden_dim: int): - if not hasattr(_thread_local, - f'flashinfer_allreduce_workspaces_{mapping.pp_rank}'): - setattr(_thread_local, - f'flashinfer_allreduce_workspaces_{mapping.pp_rank}', {}) - - flashinfer_allreduce_workspaces = getattr( - _thread_local, f'flashinfer_allreduce_workspaces_{mapping.pp_rank}') - - if mapping not in flashinfer_allreduce_workspaces: - flashinfer_allreduce_workspaces[ - mapping] = flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce( - rank=mapping.rank, - tp_size=mapping.tp_size, - max_token_num=max_num_tokens, - hidden_dim=hidden_dim, - group=dist.group.WORLD, - ) - if _enable_debug: - print( - f"FlashInferAllReduce.init: {mapping.rank} {mapping.tp_size} {max_num_tokens} {hidden_dim}" - ) - return flashinfer_allreduce_workspaces[mapping] - def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): @@ -749,125 +720,6 @@ def forward( ) -_RANK_DIST_INITIALIZED = False -_GLOBAL_FLAG_VALUE = 1 - - -class FlashInferAllReduce(nn.Module): - - def __init__(self, mapping: Mapping, - strategy: flashinfer_comm.AllReduceStrategyType, - hidden_dim: int, dtype: torch.dtype): - super().__init__() - global _RANK_DIST_INITIALIZED - - self.mapping = mapping - self.dtype = dtype - self.hidden_dim = hidden_dim # Use the passed parameter instead of hardcoded value - self.strategy = strategy - - if not _RANK_DIST_INITIALIZED: - TorchDist(mapping) - _RANK_DIST_INITIALIZED = True - - self.max_num_tokens = 8192 - - self.workspace = get_flashinfer_allreduce_workspace( - self.mapping, self.max_num_tokens, self.hidden_dim) - - self.all_reduce_params = FlashInferAllReduceParams( - strategy= - strategy, # Use the passed strategy instead of hardcoded TWOSHOT - fusion_op=flashinfer_comm.AllReduceFusionOp.NONE, - # config_mode=flashinfer_comm.AllReduceStrategyConfig.USE_MEMCPY, - ) - - # each linear layer will have its own pointers to the workspace. - # The workspace is common - # Can we use common pointers? - self.peer_comm_buffer_ptrs = torch.tensor(self.workspace[0], - dtype=torch.int64) - self.peer_barrier_ptrs_in = torch.tensor(self.workspace[2], - dtype=torch.int64) - self.peer_barrier_ptrs_out = torch.tensor(self.workspace[3], - dtype=torch.int64) - self.lamport_peer_comm_buffer_ptrs_0 = torch.tensor(self.workspace[4], - dtype=torch.int64) - self.lamport_peer_comm_buffer_ptrs_1 = torch.tensor(self.workspace[5], - dtype=torch.int64) - self.lamport_peer_comm_buffer_ptrs_2 = torch.tensor(self.workspace[6], - dtype=torch.int64) - - def forward( - self, - input: torch.Tensor, - *, - all_reduce_params: Optional[FlashInferAllReduceParams] = None, - ) -> torch.Tensor: - global _GLOBAL_FLAG_VALUE - - input = input.contiguous() - num_token = input.size(0) - hidden_dim = input.size(1) - - # Validate input dimensions - if hidden_dim != self.hidden_dim: - logger.warning( - f"Input hidden_dim {hidden_dim} doesn't match expected {self.hidden_dim}" - ) - - if num_token > self.max_num_tokens: - logger.warning( - f"Input num_token {num_token} exceeds max_num_tokens {self.max_num_tokens}" - ) - - output_buffer = torch.empty_like(input) - if _enable_debug: - print(f"FlashInferAllReduce.forward: {num_token} {hidden_dim}") - - # Use provided params or default - if all_reduce_params is None: - all_reduce_params = self.all_reduce_params - - try: - flashinfer_comm.trtllm_custom_all_reduce( - inp=input, - out=output_buffer, - tp_size=self.mapping.tp_size, - tp_rank=self.mapping.rank, - token_num=num_token, - fusion_op_code=all_reduce_params.fusion_op, - strategy_code=all_reduce_params.strategy, - config_code=all_reduce_params.config_mode, - launch_with_pdl=True, - flag_value=_GLOBAL_FLAG_VALUE, - peer_comm_buffer_ptrs=self.peer_comm_buffer_ptrs, - peer_barrier_ptrs_in=self.peer_barrier_ptrs_in, - peer_barrier_ptrs_out=self.peer_barrier_ptrs_out, - bias=all_reduce_params.bias, - residual=all_reduce_params.residual, - weight=all_reduce_params.norm_weight, - weight_pre_residual_norm=all_reduce_params. - norm_pre_residual_weight, - eps=all_reduce_params.eps, - intermediate_buffer=None, - # Enable Lamport communication buffers for better synchronization - lamport_peer_comm_buffer_ptrs_0=self. - lamport_peer_comm_buffer_ptrs_0, - lamport_peer_comm_buffer_ptrs_1=self. - lamport_peer_comm_buffer_ptrs_1, - lamport_peer_comm_buffer_ptrs_2=self. - lamport_peer_comm_buffer_ptrs_2, - ) - except Exception as e: - logger.error(f"FlashInferAllReduce failed: {e}") - raise - - _GLOBAL_FLAG_VALUE += 1 - - return output_buffer - - class FlashInferVLLMAllReduce(nn.Module): def __init__(self, mapping: Mapping, dtype: torch.dtype): @@ -876,7 +728,7 @@ def __init__(self, mapping: Mapping, dtype: torch.dtype): self.dtype = dtype self.reg_buffer_size = 8192 * 8192 * 2 - self.workspace = get_torch_dist_flashinfer_vllm_comm() + self.workspace = current_flashinfer_allreduce_workspace() def forward( self, diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 5b150665b5b9..b5f083e43f43 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -2,6 +2,11 @@ import platform import traceback +import torch + +from tensorrt_llm._utils import mpi_barrier, mpi_comm, mpi_disabled, torch_comm +from tensorrt_llm.mapping import Mapping + from ..logger import logger IS_FLASHINFER_AVAILABLE = False @@ -18,6 +23,8 @@ def get_env_enable_pdl(): if platform.system() != "Windows": try: import flashinfer + import flashinfer.comm as flashinfer_comm + from flashinfer.comm.cuda_ipc import cudart logger.info(f"flashinfer is available: {flashinfer.__version__}") IS_FLASHINFER_AVAILABLE = True except ImportError: @@ -25,3 +32,86 @@ 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" + ) + + if not mpi_disabled(): + new_group = mpi_comm().group.Incl(mapping.tp_group) + self.tp_comm = mpi_comm().Create_group(new_group) + else: + self.tp_comm = torch_comm() + + self.mapping = mapping + + max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 + self.meta_ptrs = self.create_shared_buffer( + 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 = self.create_shared_buffer(max_size) + + # Initialize custom allreduce + self.fa = flashinfer_comm.vllm_init_custom_ar( + ipc_tensors=self.meta_ptrs, + rank_data=self.rank_data, + rank=mapping.rank, + full_nvlink=True) + + # Register buffer - this is crucial! + flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs) + + def create_shared_buffer(self, size_in_bytes: int): + # Allocate local memory and get IPC handle + _, pointer = cudart.cudaMalloc(size_in_bytes) + _, handle = cudart.cudaIpcGetMemHandle(pointer) + + tp_rank = self.mapping.tp_rank + handles = self.tp_comm.allgather(handle) + + # Open IPC handles + pointers = [] + for i, h in enumerate(handles): + if i == tp_rank: + pointers.append(pointer.value) + else: + _, ptr = cudart.cudaIpcOpenMemHandle(h) + pointers.append(ptr.value) + + if mpi_disabled(): + self.tp_comm.barrier() + else: + mpi_barrier() + + return pointers + + def free_shared_buffer(self): + pass + + def destroy(self): + pass + + +flashinfer_allreduce_workspace = None + + +def init_flashinfer_allreduce_workspace(mapping: Mapping): + global flashinfer_allreduce_workspace + 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/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index fb8766082748..162c0ae2a16a 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -7,7 +7,6 @@ from dataclasses import dataclass from typing import Dict, List, Optional, Union -import flashinfer.comm as flashinfer_comm import torch import torch.nn.functional as F from torch import nn @@ -1826,8 +1825,7 @@ def __init__( disable_deep_gemm: bool = False, use_flashinfer_allreduce: bool = True, ): - from ..distributed import (AllReduce, FlashInferAllReduce, - FlashInferVLLMAllReduce) + from ..distributed import AllReduce, FlashInferVLLMAllReduce super().__init__() self.has_bias = bias @@ -1871,21 +1869,10 @@ def __init__( dtype=self.dtype) if reduce_output else None self.use_flashinfer_allreduce = use_flashinfer_allreduce - self.flashinfer_trtllm = self.use_flashinfer_allreduce and os.getenv( - "_USE_FLASHINFER_TRTLLM_ALLREDUCE", "0") == "1" 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")) - # print(f"flashinfer_trtllm: {self.flashinfer_trtllm}") - # print(f"flashinfer_vllm: {self.flashinfer_vllm}") - - if self.flashinfer_trtllm: - self.flash_infer_all_reduce = FlashInferAllReduce( - mapping=self.mapping, - hidden_dim=self.out_features, - strategy=flashinfer_comm.AllReduceStrategyType.TWOSHOT, - dtype=self.dtype) if reduce_output else None if self.flashinfer_vllm: self.flash_infer_all_reduce = FlashInferVLLMAllReduce( @@ -2027,12 +2014,7 @@ def forward( bias = None if fuse_bias else bias output = self.apply_linear(input, bias, lora_params, layer_idx) - if self.flashinfer_trtllm and output.size(0) == 150: - output = self.flash_infer_all_reduce( - output, - all_reduce_params=None, - ) - elif self.flashinfer_vllm and output.size( + if self.flashinfer_vllm and output.size( 0) <= self.flashinfer_token_threshold: output = self.flash_infer_all_reduce( output, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index dd1eadb5a454..98b6179b7182 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -32,7 +32,7 @@ from ..compilation.backend import Backend from ..compilation.utils import capture_piecewise_cuda_graph from ..distributed import MPIDist -from ..distributed.communicator import init_pp_comm, init_torch_dist_tp_comm +from ..distributed.communicator import init_pp_comm from ..expert_statistic import ExpertStatistic from ..metadata import KVCacheParams from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader @@ -153,8 +153,7 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - if mapping.has_tp(): - init_torch_dist_tp_comm(mapping) + self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index ff5807519768..469619f1b778 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -29,6 +29,7 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist, TorchDist +from ..flashinfer_utils import init_flashinfer_allreduce_workspace from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, get_spec_resource_manager) from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, @@ -296,6 +297,10 @@ def create_py_executor( else: dist = MPIDist(mapping=mapping) + # TODO: convert to arg + if os.getenv("_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1": + init_flashinfer_allreduce_workspace(mapping) + cache_transceiver_config = None if llm_args.cache_transceiver_config is not None: cache_transceiver_config = PybindMirror.maybe_to_pybind( From d985d103cd10713b8c9fb0aa4e6097e148f275d5 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 8 Oct 2025 12:32:06 -0700 Subject: [PATCH 17/21] few fixes Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/__init__.py | 3 +- tensorrt_llm/_torch/distributed/ops.py | 2 +- tensorrt_llm/_torch/flashinfer_utils.py | 52 ++++----------------- 3 files changed, 10 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/__init__.py b/tensorrt_llm/_torch/distributed/__init__.py index 8e1d3ce40173..0cb5d48adaf3 100644 --- a/tensorrt_llm/_torch/distributed/__init__.py +++ b/tensorrt_llm/_torch/distributed/__init__.py @@ -2,7 +2,7 @@ from .communicator import Distributed, MPIDist, PPComm, TorchDist from .ops import (AllReduce, AllReduceParams, AllReduceStrategy, - FlashInferAllReduce, FlashInferAllReduceParams, + FlashInferAllReduceParams, FlashInferVLLMAllReduce, MoEAllReduce, MoEAllReduceParams, allgather, alltoall_helix, reducescatter, userbuffers_allreduce_finalize) @@ -22,7 +22,6 @@ "PPComm", "MPIDist", "Distributed", - "FlashInferAllReduce", "FlashInferAllReduceParams", "FlashInferVLLMAllReduce", ] diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 39b5f3695a07..d75b61fdcdf4 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -755,7 +755,7 @@ def forward( self.workspace.fa, input_tensor, output_tensor, - self.workspace.buffer_ptrs[self.mapping.rank], + self.workspace.buffer_ptrs_ipc.peer_ptrs[self.mapping.rank], self.reg_buffer_size, 36, # CTA upper bounds: 36 as mentioned in the API ) diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index b5f083e43f43..26072764669f 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -4,8 +4,8 @@ import torch -from tensorrt_llm._utils import mpi_barrier, mpi_comm, mpi_disabled, torch_comm from tensorrt_llm.mapping import Mapping +from tensorrt_llm._ipc_utils import IpcMemory from ..logger import logger @@ -42,17 +42,11 @@ def __init__(self, mapping: Mapping): "flashinfer is not installed properly, please try pip install or building from source codes" ) - if not mpi_disabled(): - new_group = mpi_comm().group.Incl(mapping.tp_group) - self.tp_comm = mpi_comm().Create_group(new_group) - else: - self.tp_comm = torch_comm() - self.mapping = mapping max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - self.meta_ptrs = self.create_shared_buffer( - flashinfer_comm.vllm_meta_size() + max_size) + print(f"Opening IPC memory for meta with size {flashinfer_comm.vllm_meta_size() + max_size}") + 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, @@ -60,47 +54,16 @@ def __init__(self, mapping: Mapping): device=f"cuda:{mapping.local_rank}") # Create buffer pointers for IPC communication - self.buffer_ptrs = self.create_shared_buffer(max_size) + 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_tensors=self.meta_ptrs_ipc.peer_ptrs, rank_data=self.rank_data, rank=mapping.rank, full_nvlink=True) - # Register buffer - this is crucial! - flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs) - - def create_shared_buffer(self, size_in_bytes: int): - # Allocate local memory and get IPC handle - _, pointer = cudart.cudaMalloc(size_in_bytes) - _, handle = cudart.cudaIpcGetMemHandle(pointer) - - tp_rank = self.mapping.tp_rank - handles = self.tp_comm.allgather(handle) - - # Open IPC handles - pointers = [] - for i, h in enumerate(handles): - if i == tp_rank: - pointers.append(pointer.value) - else: - _, ptr = cudart.cudaIpcOpenMemHandle(h) - pointers.append(ptr.value) - - if mpi_disabled(): - self.tp_comm.barrier() - else: - mpi_barrier() - - return pointers - - def free_shared_buffer(self): - pass - - def destroy(self): - pass + flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs_ipc.peer_ptrs) flashinfer_allreduce_workspace = None @@ -108,7 +71,8 @@ def destroy(self): def init_flashinfer_allreduce_workspace(mapping: Mapping): global flashinfer_allreduce_workspace - flashinfer_allreduce_workspace = FlashInferAllReduceWorkspace(mapping) + if flashinfer_allreduce_workspace is None: + flashinfer_allreduce_workspace = FlashInferAllReduceWorkspace(mapping) def current_flashinfer_allreduce_workspace(): From 9bf98b609323d713da385475ca2c8eb5974be8c0 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 8 Oct 2025 12:38:15 -0700 Subject: [PATCH 18/21] fix Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/flashinfer_utils.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 26072764669f..1e44931e245f 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -4,8 +4,8 @@ import torch -from tensorrt_llm.mapping import Mapping from tensorrt_llm._ipc_utils import IpcMemory +from tensorrt_llm.mapping import Mapping from ..logger import logger @@ -24,7 +24,6 @@ def get_env_enable_pdl(): try: import flashinfer import flashinfer.comm as flashinfer_comm - from flashinfer.comm.cuda_ipc import cudart logger.info(f"flashinfer is available: {flashinfer.__version__}") IS_FLASHINFER_AVAILABLE = True except ImportError: @@ -45,8 +44,12 @@ def __init__(self, mapping: Mapping): self.mapping = mapping max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - print(f"Opening IPC memory for meta with size {flashinfer_comm.vllm_meta_size() + max_size}") - self.meta_ptrs_ipc = IpcMemory(mapping, flashinfer_comm.vllm_meta_size() + max_size) + print( + f"Opening IPC memory for meta with size {flashinfer_comm.vllm_meta_size() + max_size}" + ) + 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, @@ -63,7 +66,8 @@ def __init__(self, mapping: Mapping): rank=mapping.rank, full_nvlink=True) - flashinfer_comm.vllm_register_buffer(self.fa, self.buffer_ptrs_ipc.peer_ptrs) + flashinfer_comm.vllm_register_buffer(self.fa, + self.buffer_ptrs_ipc.peer_ptrs) flashinfer_allreduce_workspace = None From 3bd858d36eb12a0ddd69f30ceea9ff6ae6da8a59 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 8 Oct 2025 13:18:16 -0700 Subject: [PATCH 19/21] cuda graph support init Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 48 ++++++---- tensorrt_llm/_torch/flashinfer_utils.py | 96 ++++++++++++++----- .../_torch/pyexecutor/cuda_graph_runner.py | 11 ++- .../_torch/pyexecutor/model_engine.py | 6 ++ .../_torch/pyexecutor/py_executor_creator.py | 5 - 5 files changed, 120 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index d75b61fdcdf4..69b0c7ef1a14 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -730,6 +730,29 @@ def __init__(self, mapping: Mapping, dtype: torch.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.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, @@ -746,21 +769,10 @@ def forward( Returns: torch.Tensor: Reduced tensor """ - input_tensor = input.contiguous() - output_tensor = torch.empty_like(input_tensor) - - try: - # Perform vLLM custom allreduce - flashinfer_comm.vllm_all_reduce( - self.workspace.fa, - input_tensor, - output_tensor, - self.workspace.buffer_ptrs_ipc.peer_ptrs[self.mapping.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 + 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 1e44931e245f..65905aab4f0a 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -1,3 +1,4 @@ +import contextlib import os import platform import traceback @@ -5,6 +6,7 @@ import torch from tensorrt_llm._ipc_utils import IpcMemory +from tensorrt_llm._utils import mpi_comm, mpi_disabled, torch_comm from tensorrt_llm.mapping import Mapping from ..logger import logger @@ -44,30 +46,80 @@ def __init__(self, mapping: Mapping): self.mapping = mapping max_size = 8192 * 8192 * 2 # 2 bytes for bfloat16 - print( + logger.info( f"Opening IPC memory for meta with size {flashinfer_comm.vllm_meta_size() + max_size}" ) - 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.rank, - full_nvlink=True) - - flashinfer_comm.vllm_register_buffer(self.fa, - self.buffer_ptrs_ipc.peer_ptrs) + 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.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 + self._graph_registered = False + logger.info( + f"FlashInferAllReduceWorkspace initialized for rank {mapping.rank}") + + @contextlib.contextmanager + def capture(self): + try: + self._is_capturing = True + logger.info( + f"Rank {self.mapping.rank}: Starting CUDA graph capture") + yield + finally: + self._is_capturing = False + # Register graph buffers after capture if not already done + if not self._graph_registered: + self.register_graph_buffers() + logger.info( + f"Rank {self.mapping.rank}: Finished CUDA graph capture") + + def register_graph_buffers(self): + # add error handling + handle, offsets = flashinfer_comm.get_graph_buffer_ipc_meta(self.fa) + logger.info( + f"Rank {self.mapping.rank}: Registering {len(handle)} graph buffer(s)" + ) + + # broadcast + if mpi_disabled(): + allgather = torch_comm().tp_allgather + else: + comm = mpi_comm().Split( + self.mapping.pp_rank * self.mapping.cp_size + + self.mapping.cp_rank, self.mapping.tp_rank) + allgather = comm.allgather + + handles = allgather(handle) + offsets = allgather(offsets) + + flashinfer_comm.vllm_register_graph_buffers(self.fa, handles, offsets) + self._graph_registered = True + logger.info( + f"Rank {self.mapping.rank}: Registered {len(handle)} graph buffer(s)" + ) flashinfer_allreduce_workspace = None 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 98b6179b7182..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 @@ -173,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/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 469619f1b778..ff5807519768 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -29,7 +29,6 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist, TorchDist -from ..flashinfer_utils import init_flashinfer_allreduce_workspace from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, get_spec_resource_manager) from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, @@ -297,10 +296,6 @@ def create_py_executor( else: dist = MPIDist(mapping=mapping) - # TODO: convert to arg - if os.getenv("_USE_FLASHINFER_VLLM_ALLREDUCE", "0") == "1": - init_flashinfer_allreduce_workspace(mapping) - cache_transceiver_config = None if llm_args.cache_transceiver_config is not None: cache_transceiver_config = PybindMirror.maybe_to_pybind( From 62b1ff66fded8f21db4123838daabe170d973076 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 8 Oct 2025 15:21:54 -0700 Subject: [PATCH 20/21] running in mem access Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 4 ++- tensorrt_llm/_torch/flashinfer_utils.py | 40 ++++++++++++++----------- 2 files changed, 25 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 69b0c7ef1a14..892dd0fb22db 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -736,14 +736,16 @@ def custom_all_reduce(self, input: torch.Tensor, registered: bool = False): try: if registered: + print(f"Rank {self.mapping.tp_rank}: Registered all reduce") flashinfer_comm.vllm_all_reduce(self.workspace.fa, input_tensor, output_tensor, 0, 0, 36) else: + print(f"Rank {self.mapping.tp_rank}: Unregistered all reduce", self.workspace.buffer_ptrs_ipc.peer_ptrs[self.mapping.tp_rank]) flashinfer_comm.vllm_all_reduce( self.workspace.fa, input_tensor, output_tensor, - self.workspace.buffer_ptrs_ipc.peer_ptrs[self.mapping.rank], + 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 ) diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 65905aab4f0a..08b0f27e034b 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -66,7 +66,7 @@ def __init__(self, mapping: Mapping): self.fa = flashinfer_comm.vllm_init_custom_ar( ipc_tensors=self.meta_ptrs_ipc.peer_ptrs, rank_data=self.rank_data, - rank=mapping.rank, + rank=mapping.tp_rank, full_nvlink=True) flashinfer_comm.vllm_register_buffer(self.fa, @@ -79,14 +79,14 @@ def __init__(self, mapping: Mapping): self._is_capturing = False self._graph_registered = False logger.info( - f"FlashInferAllReduceWorkspace initialized for rank {mapping.rank}") + f"FlashInferAllReduceWorkspace initialized for rank {mapping.tp_rank}") @contextlib.contextmanager def capture(self): try: self._is_capturing = True logger.info( - f"Rank {self.mapping.rank}: Starting CUDA graph capture") + f"Rank {self.mapping.tp_rank}: Starting CUDA graph capture") yield finally: self._is_capturing = False @@ -94,31 +94,35 @@ def capture(self): if not self._graph_registered: self.register_graph_buffers() logger.info( - f"Rank {self.mapping.rank}: Finished CUDA graph capture") + f"Rank {self.mapping.tp_rank}: Finished CUDA graph capture") def register_graph_buffers(self): # add error handling - handle, offsets = flashinfer_comm.get_graph_buffer_ipc_meta(self.fa) + handle, offsets = flashinfer_comm.vllm_get_graph_buffer_ipc_meta(self.fa) logger.info( - f"Rank {self.mapping.rank}: Registering {len(handle)} graph buffer(s)" + f"Rank {self.mapping.tp_rank}: Registering {len(handle)} graph buffer(s)" ) - # broadcast - if mpi_disabled(): - allgather = torch_comm().tp_allgather - else: - comm = mpi_comm().Split( - self.mapping.pp_rank * self.mapping.cp_size + - self.mapping.cp_rank, self.mapping.tp_rank) - allgather = comm.allgather + world_size = self.mapping.tp_size + tp_rank = self.mapping.tp_rank + tp_group_ranks = sorted(self.mapping.tp_group) - handles = allgather(handle) - offsets = allgather(offsets) + all_data = [[None, None] for _ in range(world_size)] + all_data[tp_rank] = [handle, offsets] - flashinfer_comm.vllm_register_graph_buffers(self.fa, handles, 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) self._graph_registered = True logger.info( - f"Rank {self.mapping.rank}: Registered {len(handle)} graph buffer(s)" + f"Rank {self.mapping.tp_rank}: Registered {len(handle)} graph buffer(s)" ) From 0ed45feaaa973497e5761aed55d4802dfbdc8b46 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 9 Oct 2025 07:43:27 -0700 Subject: [PATCH 21/21] working w cuda graphs Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/distributed/ops.py | 5 ++--- tensorrt_llm/_torch/flashinfer_utils.py | 30 +++++++++---------------- 2 files changed, 12 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 892dd0fb22db..7ee55ea4d261 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -736,16 +736,15 @@ def custom_all_reduce(self, input: torch.Tensor, registered: bool = False): try: if registered: - print(f"Rank {self.mapping.tp_rank}: Registered all reduce") flashinfer_comm.vllm_all_reduce(self.workspace.fa, input_tensor, output_tensor, 0, 0, 36) else: - print(f"Rank {self.mapping.tp_rank}: Unregistered all reduce", self.workspace.buffer_ptrs_ipc.peer_ptrs[self.mapping.tp_rank]) flashinfer_comm.vllm_all_reduce( self.workspace.fa, input_tensor, output_tensor, - self.workspace.buffer_ptrs_ipc.peer_ptrs[self.mapping.tp_rank], + 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 ) diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 08b0f27e034b..a76d4ae9caba 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -6,7 +6,7 @@ import torch from tensorrt_llm._ipc_utils import IpcMemory -from tensorrt_llm._utils import mpi_comm, mpi_disabled, torch_comm +from tensorrt_llm._utils import mpi_comm from tensorrt_llm.mapping import Mapping from ..logger import logger @@ -77,41 +77,34 @@ def __init__(self, mapping: Mapping): raise self._is_capturing = False - self._graph_registered = False logger.info( - f"FlashInferAllReduceWorkspace initialized for rank {mapping.tp_rank}") + f"FlashInferAllReduceWorkspace initialized for rank {mapping.tp_rank}" + ) @contextlib.contextmanager def capture(self): try: self._is_capturing = True - logger.info( - f"Rank {self.mapping.tp_rank}: Starting CUDA graph capture") yield finally: self._is_capturing = False # Register graph buffers after capture if not already done - if not self._graph_registered: - self.register_graph_buffers() - logger.info( - f"Rank {self.mapping.tp_rank}: Finished CUDA graph capture") + self.register_graph_buffers() def register_graph_buffers(self): # add error handling - handle, offsets = flashinfer_comm.vllm_get_graph_buffer_ipc_meta(self.fa) - logger.info( - f"Rank {self.mapping.tp_rank}: Registering {len(handle)} graph buffer(s)" - ) + handle, offsets = flashinfer_comm.vllm_get_graph_buffer_ipc_meta( + self.fa) world_size = self.mapping.tp_size tp_rank = self.mapping.tp_rank - tp_group_ranks = sorted(self.mapping.tp_group) 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) + 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) @@ -119,11 +112,8 @@ def register_graph_buffers(self): 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) - self._graph_registered = True - logger.info( - f"Rank {self.mapping.tp_rank}: Registered {len(handle)} graph buffer(s)" - ) + flashinfer_comm.vllm_register_graph_buffers(self.fa, handles, + offsets_list) flashinfer_allreduce_workspace = None