@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d4f123cc48

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 16:30 UTC

[Kernel] FlashInfer: switch allreduce fusion to unified API (#33985)

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>

Mohammad Miadh Angkad committed Feb 9, 2026 at 23:43 UTC d4f123cc48c374f7aad48cd808d797c71711ebc7

3 files changed +80 -114

benchmarks/kernels/benchmark_fused_collective.py

+54 -69

@@ -5,7 +5,7 @@

5 Benchmark for FlashInfer fused collective operations vs standard operations.

7 This benchmark compares:

8 -1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant)

8 +1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)

9 2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations

11 Usage with torchrun:

@@ -24,7 +24,6 @@ import torch.distributed as dist # type: ignore

25 from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config

26 from vllm.distributed import (

27 - get_tp_group,

27 tensor_model_parallel_all_reduce,

28 )

29 from vllm.distributed.parallel_state import (

@@ -52,11 +51,12 @@ logger = init_logger(__name__)

51 try:

52 import flashinfer.comm as flashinfer_comm # type: ignore

55 - if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"):

54 + if not (

55 + hasattr(flashinfer_comm, "allreduce_fusion")

56 + and hasattr(flashinfer_comm, "create_allreduce_fusion_workspace")

57 + ):

58 flashinfer_comm = None

57 - logger.warning(

58 - "FlashInfer comm module found but missing trtllm_allreduce_fusion"

59 - )

59 + logger.warning("FlashInfer comm module found but missing allreduce_fusion API")

60 except ImportError:

61 flashinfer_comm = None

62 logger.warning("FlashInfer not found, only benchmarking standard operations")

@@ -75,7 +75,7 @@ _FI_MAX_SIZES = {

75 }

77 # Global workspace tensor for FlashInfer

78 -_FI_WORKSPACE_TENSOR = None

78 +_FI_WORKSPACE = None

81 def setup_flashinfer_workspace(

@@ -83,10 +83,10 @@ def setup_flashinfer_workspace(

83 rank: int,

84 hidden_dim: int,

85 max_token_num: int,

86 - use_fp32_lamport: bool = False,

86 + dtype: torch.dtype,

87 ):

88 """Setup FlashInfer workspace for fused allreduce operations."""

89 - global _FI_WORKSPACE_TENSOR

89 + global _FI_WORKSPACE

91 if flashinfer_comm is None:

92 return None, None

@@ -96,33 +96,29 @@ def setup_flashinfer_workspace(

96 return None, None

98 try:

99 - # Create IPC workspace

100 - ipc_handles, workspace_tensor = (

101 - flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(

102 - tp_rank=rank,

103 - tp_size=world_size,

104 - max_token_num=max_token_num,

105 - hidden_dim=hidden_dim,

106 - group=get_tp_group().device_group,

107 - use_fp32_lamport=use_fp32_lamport,

108 - )

99 + workspace = flashinfer_comm.create_allreduce_fusion_workspace(

100 + backend="trtllm",

101 + world_size=world_size,

102 + rank=rank,

103 + max_token_num=max_token_num,

104 + hidden_dim=hidden_dim,

105 + dtype=dtype,

106 )

107

111 - _FI_WORKSPACE_TENSOR = workspace_tensor

112 - return ipc_handles, workspace_tensor

108 + _FI_WORKSPACE = workspace

109 + return workspace

110 except Exception as e:

111 logger.error("Failed to setup FlashInfer workspace: %s", e)

115 - return None, None

112 + return None

113

114

118 -def cleanup_flashinfer_workspace(ipc_handles):

115 +def cleanup_flashinfer_workspace(workspace):

116 """Cleanup FlashInfer workspace."""

120 - if flashinfer_comm is None or ipc_handles is None:

117 + if flashinfer_comm is None or workspace is None:

118 return

119

120 try:

124 - group = get_tp_group().device_group

125 - flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group)

121 + workspace.destroy()

122 except Exception as e:

123 logger.error("Failed to cleanup FlashInfer workspace: %s", e)

124

@@ -132,25 +128,15 @@ class FlashInferFusedAllReduceParams:

128

129 def __init__(

130 self,

135 - rank: int,

136 - world_size: int,

137 - use_fp32_lamport: bool = False,

131 max_token_num: int = 1024,

132 ):

140 - self.rank = rank

141 - self.world_size = world_size

142 - self.use_fp32_lamport = use_fp32_lamport

143 - self.trigger_completion_at_end = True

133 self.launch_with_pdl = True

134 self.fp32_acc = True

135 self.max_token_num = max_token_num

136

137 def get_trtllm_fused_allreduce_kwargs(self):

138 return {

150 - "world_rank": self.rank,

151 - "world_size": self.world_size,

139 "launch_with_pdl": self.launch_with_pdl,

153 - "trigger_completion_at_end": self.trigger_completion_at_end,

140 "fp32_acc": self.fp32_acc,

141 }

142

@@ -165,7 +151,7 @@ def flashinfer_fused_allreduce_rmsnorm(

151 norm_out: torch.Tensor | None = None,

152 ):

153 """FlashInfer fused allreduce + rmsnorm operation."""

168 - if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:

154 + if flashinfer_comm is None or _FI_WORKSPACE is None:

155 raise RuntimeError("FlashInfer not available or workspace not initialized")

156

157 if norm_out is None:

@@ -174,18 +160,15 @@ def flashinfer_fused_allreduce_rmsnorm(

160 else:

161 residual_out = input_tensor

162

177 - flashinfer_comm.trtllm_allreduce_fusion(

178 - allreduce_in=input_tensor,

179 - token_num=input_tensor.shape[0],

163 + flashinfer_comm.allreduce_fusion(

164 + input=input_tensor,

165 + workspace=_FI_WORKSPACE,

166 + pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,

167 residual_in=residual,

168 residual_out=residual_out,

169 norm_out=norm_out,

170 rms_gamma=rms_gamma,

171 rms_eps=rms_eps,

185 - hidden_dim=input_tensor.shape[-1],

186 - workspace_ptrs=_FI_WORKSPACE_TENSOR,

187 - pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,

188 - allreduce_out=None,

172 quant_out=None,

173 scale_out=None,

174 layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,

@@ -207,7 +190,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(

190 quant_out: torch.Tensor | None = None,

191 ):

192 """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""

210 - if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:

193 + if flashinfer_comm is None or _FI_WORKSPACE is None:

194 raise RuntimeError("FlashInfer not available or workspace not initialized")

195

196 if norm_out is None:

@@ -216,18 +199,15 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(

199 else:

200 residual_out = input_tensor

201

219 - flashinfer_comm.trtllm_allreduce_fusion(

220 - allreduce_in=input_tensor,

221 - token_num=input_tensor.shape[0],

202 + flashinfer_comm.allreduce_fusion(

203 + input=input_tensor,

204 + workspace=_FI_WORKSPACE,

205 + pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,

206 residual_in=residual,

207 residual_out=residual_out,

208 norm_out=norm_out,

209 rms_gamma=rms_gamma,

210 rms_eps=rms_eps,

227 - hidden_dim=input_tensor.shape[-1],

228 - workspace_ptrs=_FI_WORKSPACE_TENSOR,

229 - pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,

230 - allreduce_out=None,

211 quant_out=quant_out,

212 scale_out=None,

213 layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,

@@ -250,7 +230,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(

230 norm_out: torch.Tensor | None = None,

231 ):

232 """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""

253 - if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:

233 + if flashinfer_comm is None or _FI_WORKSPACE is None:

234 raise RuntimeError("FlashInfer not available or workspace not initialized")

235

236 if norm_out is None:

@@ -259,18 +239,15 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(

239 else:

240 residual_out = input_tensor

241

262 - flashinfer_comm.trtllm_allreduce_fusion(

263 - allreduce_in=input_tensor,

264 - token_num=input_tensor.shape[0],

242 + flashinfer_comm.allreduce_fusion(

243 + input=input_tensor,

244 + workspace=_FI_WORKSPACE,

245 + pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,

246 residual_in=residual,

247 residual_out=residual_out,

248 norm_out=norm_out,

249 rms_gamma=rms_gamma,

250 rms_eps=rms_eps,

270 - hidden_dim=input_tensor.shape[-1],

271 - workspace_ptrs=_FI_WORKSPACE_TENSOR,

272 - pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,

273 - allreduce_out=None,

251 quant_out=quant_out,

252 scale_out=output_scale,

253 layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,

@@ -1040,23 +1017,31 @@ def main():

1017 configs = list(itertools.product(args.num_tokens, dtypes, residual_options))

1018

1019 # Setup FlashInfer workspace if available

1043 - ipc_handles = None

1020 + workspace = None

1021 allreduce_params = None

1022

1023 if flashinfer_comm is not None:

1024 # Use the largest hidden dimension for workspace setup

1025 + max_element_size = max(torch.finfo(dt).bits // 8 for dt in dtypes)

1026 + workspace_dtype = (

1027 + torch.float32

1028 + if max_element_size == 4

1029 + else (torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16)

1030 + )

1031 max_num_token = _FI_MAX_SIZES.get(world_size) // (

1049 - args.hidden_dim * world_size * 2

1032 + args.hidden_dim * max_element_size

1033 )

1034

1052 - ipc_handles, workspace_tensor = setup_flashinfer_workspace(

1053 - world_size, rank, args.hidden_dim, max_num_token

1035 + workspace = setup_flashinfer_workspace(

1036 + world_size,

1037 + rank,

1038 + args.hidden_dim,

1039 + max_num_token,

1040 + dtype=workspace_dtype,

1041 )

1042

1056 - if workspace_tensor is not None:

1043 + if workspace is not None:

1044 allreduce_params = FlashInferFusedAllReduceParams(

1058 - rank=rank,

1059 - world_size=world_size,

1045 max_token_num=max_num_token,

1046 )

1047

@@ -1119,8 +1104,8 @@ def main():

1104

1105 finally:

1106 # Cleanup

1122 - if ipc_handles is not None:

1123 - cleanup_flashinfer_workspace(ipc_handles)

1107 + if workspace is not None:

1108 + cleanup_flashinfer_workspace(workspace)

1109

1110 dist.barrier()

1111

tests/compile/passes/distributed/test_fusion_all_reduce.py

+3 -2

@@ -202,9 +202,10 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):

202 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")

203 @pytest.mark.skipif(

204 not find_spec("flashinfer")

205 - or not has_module_attribute("flashinfer.comm", "trtllm_allreduce_fusion"),

205 + or not has_module_attribute("flashinfer.comm", "allreduce_fusion")

206 + or not has_module_attribute("flashinfer.comm", "create_allreduce_fusion_workspace"),

207 reason="flashinfer is not found or flashinfer "

207 - "is not compiled with trtllm_allreduce_fusion",

208 + "is not compiled with allreduce_fusion",

209 )

210 def test_all_reduce_fusion_pass_replace(

211 test_model: torch.nn.Module,

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

+23 -43

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import contextlib

4 from importlib.util import find_spec

5 from types import ModuleType

@@ -36,7 +37,9 @@ if find_spec("flashinfer"):

37 try:

38 import flashinfer.comm as _flashinfer_comm

39 - if hasattr(_flashinfer_comm, "trtllm_allreduce_fusion"):

40 + if hasattr(_flashinfer_comm, "allreduce_fusion") and hasattr(

41 + _flashinfer_comm, "create_allreduce_fusion_workspace"

42 + ):

43 flashinfer_comm = _flashinfer_comm

44 except ImportError:

45 pass

@@ -79,7 +82,7 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {

84 if flashinfer_comm is not None:

82 - _FI_WORKSPACE_TENSOR = None

85 + _FI_WORKSPACE = None

86 MiB = 1024 * 1024

88 def call_trtllm_fused_allreduce_norm(

@@ -87,10 +90,8 @@ if flashinfer_comm is not None:

90 residual: torch.Tensor,

91 rms_gamma: torch.Tensor,

92 rms_eps: float,

90 - world_rank: int,

93 world_size: int,

94 launch_with_pdl: bool,

93 - trigger_completion_at_end: bool,

95 fp32_acc: bool,

96 max_token_num: int,

97 pattern_code: int,

@@ -121,7 +122,7 @@ if flashinfer_comm is not None:

122 max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB

123 )

124

124 - assert _FI_WORKSPACE_TENSOR is not None, (

125 + assert _FI_WORKSPACE is not None, (

126 "Flashinfer must be enabled when using flashinfer"

127 )

128 if norm_out is None:

@@ -134,24 +135,18 @@ if flashinfer_comm is not None:

135 residual_out = allreduce_in

136 # For the sizes that are smaller than the max size,

137 # we only use flashinfer one shot allreduce

137 - flashinfer_comm.trtllm_allreduce_fusion(

138 - allreduce_in=allreduce_in,

139 - token_num=allreduce_in.shape[0],

138 + flashinfer_comm.allreduce_fusion(

139 + input=allreduce_in,

140 + workspace=_FI_WORKSPACE,

141 + pattern=pattern_code,

142 residual_in=residual,

143 residual_out=residual_out,

144 norm_out=norm_out,

145 rms_gamma=rms_gamma,

146 rms_eps=rms_eps,

145 - world_rank=world_rank,

146 - world_size=world_size,

147 - hidden_dim=allreduce_in.shape[-1],

148 - workspace_ptrs=_FI_WORKSPACE_TENSOR,

147 launch_with_pdl=launch_with_pdl,

148 use_oneshot=use_oneshot,

151 - trigger_completion_at_end=trigger_completion_at_end,

149 fp32_acc=fp32_acc,

153 - pattern_code=pattern_code,

154 - allreduce_out=None,

150 quant_out=quant_out,

151 scale_out=scale_out,

152 # in vllm we only support swizzled layout

@@ -164,10 +159,8 @@ if flashinfer_comm is not None:

159 residual: torch.Tensor,

160 rms_gamma: torch.Tensor,

161 rms_eps: float,

167 - world_rank: int,

162 world_size: int,

163 launch_with_pdl: bool,

170 - trigger_completion_at_end: bool,

164 fp32_acc: bool,

165 max_token_num: int,

166 pattern_code: int,

@@ -200,25 +193,18 @@ class FlashInferFusedAllReduceParams:

193

194 def __init__(

195 self,

203 - rank: int,

196 world_size: int,

205 - use_fp32_lamport: bool = False,

197 max_token_num: int = 1024,

198 ) -> None:

208 - self.rank = rank

199 self.world_size = world_size

210 - self.use_fp32_lamport = use_fp32_lamport

211 - self.trigger_completion_at_end = True

200 self.launch_with_pdl = True

201 self.fp32_acc = True

202 self.max_token_num = max_token_num

203

204 def get_trtllm_fused_allreduce_kwargs(self) -> dict[str, bool | int]:

205 return {

218 - "world_rank": self.rank,

206 "world_size": self.world_size,

207 "launch_with_pdl": self.launch_with_pdl,

221 - "trigger_completion_at_end": self.trigger_completion_at_end,

208 "fp32_acc": self.fp32_acc,

209 "max_token_num": self.max_token_num,

210 }

@@ -712,7 +698,6 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

698 self.hidden_dim = config.model_config.get_hidden_size()

699 self.group = get_tp_group().device_group

700 rank = get_tensor_model_parallel_rank()

715 - use_fp32_lamport = self.model_dtype == torch.float32

701 if flashinfer_comm is None:

702 logger.warning(

703 "Flashinfer is not installed or comm module not found, "

@@ -730,7 +715,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

715 self.tp_size,

716 )

717 return

733 - element_size = 4 if use_fp32_lamport else 2

718 + element_size = torch.tensor([], dtype=self.model_dtype).element_size()

719 self.max_token_num = max_size // (self.hidden_dim * element_size)

720 # take the min to save workspace size and we'll never use more

721 # than max_num_batched_tokens anyways

@@ -744,23 +729,19 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

729 scope="global",

730 )

731

747 - self.ipc_handles, workspace_tensor = (

748 - flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(

749 - tp_rank=rank,

750 - tp_size=self.tp_size,

751 - max_token_num=self.max_token_num,

752 - hidden_dim=self.hidden_dim,

753 - group=self.group,

754 - use_fp32_lamport=use_fp32_lamport,

755 - )

732 + self.workspace = flashinfer_comm.create_allreduce_fusion_workspace(

733 + backend="trtllm",

734 + world_size=self.tp_size,

735 + rank=rank,

736 + max_token_num=self.max_token_num,

737 + hidden_dim=self.hidden_dim,

738 + dtype=self.model_dtype,

739 )

740

758 - global _FI_WORKSPACE_TENSOR

759 - _FI_WORKSPACE_TENSOR = workspace_tensor

741 + global _FI_WORKSPACE

742 + _FI_WORKSPACE = self.workspace

743 self.allreduce_params = FlashInferFusedAllReduceParams(

761 - rank=rank,

744 world_size=self.tp_size,

763 - use_fp32_lamport=use_fp32_lamport,

745 max_token_num=self.max_token_num,

746 )

747

@@ -832,7 +813,6 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

813 def __del__(self) -> None:

814 if getattr(self, "disabled", True):

815 return

835 - if flashinfer_comm is not None:

836 - flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(

837 - self.ipc_handles, self.group

838 - )

816 + if getattr(self, "workspace", None) is not None:

817 + with contextlib.suppress(Exception):

818 + self.workspace.destroy()