@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fc0f41d10a
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 15, 02:15 UTC

Integration SM100 FlashInfer fused allreduce RMSNorm (#20691)

Signed-off-by: ilmarkov <imarkov@redhat.com> Co-authored-by: ilmarkov <imarkov@redhat.com>

Ilya Markov committed Jul 12, 2025 at 03:58 UTC fc0f41d10aca510658a4d86c8bff2e6781d5d669
4 files changed +514 -6
tests/compile/test_fusion_all_reduce.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from importlib.util import find_spec
4 +
5 +import pytest
6 +import torch
7 +
8 +import vllm.envs as envs
9 +from vllm.compilation.collective_fusion import AllReduceFusionPass
10 +from vllm.config import (CompilationConfig, CompilationLevel, DeviceConfig,
11 + ModelConfig, PassConfig, VllmConfig)
12 +from vllm.distributed import tensor_model_parallel_all_reduce
13 +from vllm.distributed.parallel_state import (init_distributed_environment,
14 + initialize_model_parallel)
15 +from vllm.model_executor.layers.layernorm import RMSNorm
16 +from vllm.platforms import current_platform
17 +from vllm.utils import update_environment_variables
18 +
19 +from ..utils import multi_gpu_test
20 +from .backend import TestBackend
21 +
22 +
23 +class TestAllReduceRMSNormModel(torch.nn.Module):
24 +
25 + def __init__(self, hidden_size=16, eps=1e-6):
26 + super().__init__()
27 + self.hidden_size = hidden_size
28 + self.eps = eps
29 + self.norm = RMSNorm(hidden_size, eps)
30 +
31 + def forward(self, hidden_states, residual):
32 + view = hidden_states.reshape(-1, self.hidden_size)
33 + all_reduce = tensor_model_parallel_all_reduce(view)
34 + norm = self.norm(all_reduce)
35 + return norm
36 +
37 + def ops_in_model_before(self):
38 + return [torch.ops.vllm.all_reduce.default]
39 +
40 + def ops_in_model_after(self):
41 + return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
42 +
43 +
44 +class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):
45 +
46 + def __init__(self, hidden_size=16, eps=1e-6):
47 + super().__init__()
48 + self.hidden_size = hidden_size
49 + self.eps = eps
50 + self.norm = RMSNorm(hidden_size, eps)
51 +
52 + def forward(self, hidden_states, residual):
53 + view = hidden_states.reshape(-1, self.hidden_size)
54 + all_reduce = tensor_model_parallel_all_reduce(view)
55 + norm, _ = self.norm(all_reduce, residual)
56 + return norm
57 +
58 + def ops_in_model_before(self):
59 + return [torch.ops.vllm.all_reduce.default]
60 +
61 + def ops_in_model_after(self):
62 + return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
63 +
64 +
65 +@multi_gpu_test(num_gpus=2)
66 +@pytest.mark.parametrize(
67 + "test_model",
68 + [TestAllReduceRMSNormModel, TestAllReduceFusedAddRMSNormModel])
69 +@pytest.mark.parametrize("batch_size", [8])
70 +@pytest.mark.parametrize("seq_len", [8])
71 +@pytest.mark.parametrize("hidden_size", [4096])
72 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
73 +@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
74 + reason="Only test on CUDA")
75 +@pytest.mark.skipif(not find_spec("flashinfer"),
76 + reason="flashinfer is not installed")
77 +@pytest.mark.skipif(not current_platform.is_device_capability(100),
78 + reason="Only test on SM100")
79 +def test_all_reduce_fusion_pass_replace(test_model: torch.nn.Module,
80 + batch_size: int, seq_len: int,
81 + hidden_size: int, dtype: torch.dtype):
82 + num_processes = 2
83 +
84 + def run_torch_spawn(fn, nprocs):
85 + torch.multiprocessing.spawn(fn,
86 + args=(num_processes, test_model,
87 + batch_size, seq_len, hidden_size,
88 + dtype),
89 + nprocs=nprocs)
90 +
91 + run_torch_spawn(all_reduce_fusion_pass_on_test_model, num_processes)
92 +
93 +
94 +def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
95 + test_model_cls: torch.nn.Module,
96 + batch_size: int, seq_len: int,
97 + hidden_size: int, dtype: torch.dtype):
98 + current_platform.seed_everything(0)
99 +
100 + device = torch.device(f"cuda:{local_rank}")
101 + torch.cuda.set_device(device)
102 + torch.set_default_device(device)
103 + torch.set_default_dtype(dtype)
104 +
105 + update_environment_variables({
106 + 'RANK': str(local_rank),
107 + 'LOCAL_RANK': str(local_rank),
108 + 'WORLD_SIZE': str(world_size),
109 + 'MASTER_ADDR': 'localhost',
110 + 'MASTER_PORT': '12345',
111 + })
112 +
113 + init_distributed_environment()
114 + initialize_model_parallel(tensor_model_parallel_size=world_size)
115 +
116 + vllm_config = VllmConfig(
117 + compilation_config=CompilationConfig(level=CompilationLevel.PIECEWISE,
118 + custom_ops=["+rms_norm"],
119 + compile_sizes=[2, 4, 8]))
120 + vllm_config.compilation_config.pass_config = PassConfig(
121 + enable_fi_allreduce_fusion=True)
122 + vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
123 +
124 + # this is a fake model name to construct the model config
125 + # in the vllm_config, it's not really used.
126 + model_name = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
127 + vllm_config.model_config = ModelConfig(model=model_name,
128 + task="auto",
129 + tokenizer=model_name,
130 + tokenizer_mode="auto",
131 + trust_remote_code=True,
132 + dtype=dtype,
133 + seed=42)
134 +
135 + all_reduce_fusion_pass = AllReduceFusionPass(
136 + vllm_config, vllm_config.compilation_config.pass_config.
137 + fi_allreduce_fusion_max_token_num)
138 + backend = TestBackend(all_reduce_fusion_pass)
139 +
140 + model = test_model_cls(hidden_size)
141 +
142 + hidden_states = torch.randn((batch_size * seq_len, hidden_size),
143 + requires_grad=False)
144 + residual = torch.randn((batch_size * seq_len, hidden_size),
145 + requires_grad=False)
146 +
147 + compiled_model = torch.compile(model, backend=backend)
148 + compiled_model(hidden_states, residual)
149 +
150 + backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
151 + backend.check_after_ops(model.ops_in_model_after())
152 + del all_reduce_fusion_pass
vllm/compilation/collective_fusion.py
+352 -4
@@ -1,23 +1,39 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from importlib.util import find_spec
4 from typing import Optional
5
6 import torch
7 import torch._inductor.pattern_matcher as pm
8 import torch.fx as fx
9 +from torch._higher_order_ops.auto_functionalize import auto_functionalized
10 from torch._inductor.pattern_matcher import PatternMatcherPass
11 from torch.distributed._symmetric_memory import enable_symm_mem_for_group
12
13 from vllm.config import VllmConfig
12 -from vllm.distributed import get_tp_group
14 +from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
15 from vllm.distributed.parallel_state import (
14 - get_tensor_model_parallel_world_size)
16 + get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
17 from vllm.logger import init_logger
18 +from vllm.utils import direct_register_custom_op
19
20 from .vllm_inductor_pass import VllmInductorPass
21
22 +if find_spec("flashinfer"):
23 + import flashinfer.comm as flashinfer_comm
24 +
25 + flashinfer_comm = (flashinfer_comm if hasattr(
26 + flashinfer_comm, "trtllm_allreduce_fusion") else None)
27 +else:
28 + flashinfer_comm = None
29 +from vllm.platforms import current_platform
30 +
31 logger = init_logger(__name__)
32
33 +ALLREDUCE_OP = torch.ops.vllm.all_reduce.default
34 +RMS_OP = torch.ops._C.rms_norm.default
35 +RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
36 +
37
38 class BasePattern:
39
@@ -43,7 +59,8 @@ class GEMMReduceScatterPattern(BasePattern):
59 mm,
60 dim=0,
61 world_size=self.tp_size,
46 - group_name=self.tp.unique_name)
62 + group_name=self.tp.unique_name,
63 + )
64 return reduce_scatter
65
66 def replacement(mul: torch.Tensor, mm_weight: torch.Tensor):
@@ -79,7 +96,8 @@ class AllGatherGEMMPattern(BasePattern):
96 x,
97 dim=0,
98 world_size=self.tp_size,
82 - group_name=self.tp.unique_name)
99 + group_name=self.tp.unique_name,
100 + )
101
102 return torch.ops.aten.mm.default(all_gather, weight)
103
@@ -125,3 +143,333 @@ class AsyncTPPass(VllmInductorPass):
143 logger.debug("Replaced %s patterns", count)
144 self.dump_graph(graph, "after_async_tp_pass")
145 self.end_and_log()
146 +
147 +
148 +if flashinfer_comm is not None:
149 + _FI_WORKSPACE_TENSOR = None
150 +
151 + MiB = 1024 * 1024
152 + # Max size of the input tensor per world size
153 + # to use flashinfer fused allreduce
154 + _FI_MAX_SIZES = {
155 + 2: MiB, # 1MB
156 + 4: MiB, # 1MB
157 + 6: MiB // 2, # 512KB
158 + 8: MiB // 2, # 512KB
159 + }
160 +
161 + def call_trtllm_fused_allreduce_norm(
162 + allreduce_in: torch.Tensor,
163 + residual: torch.Tensor,
164 + rms_gamma: torch.Tensor,
165 + rms_eps: float,
166 + world_rank: int,
167 + world_size: int,
168 + launch_with_pdl: bool,
169 + trigger_completion_at_end: bool,
170 + fp32_acc: bool,
171 + max_token_num: int,
172 + norm_out: Optional[torch.Tensor] = None,
173 + ) -> None:
174 + use_flashinfer = allreduce_in.shape[0] * allreduce_in.shape[
175 + 1] * allreduce_in.element_size() <= min(
176 + _FI_MAX_SIZES[world_size],
177 + max_token_num * allreduce_in.shape[0] *
178 + allreduce_in.element_size(),
179 + )
180 + if use_flashinfer:
181 + assert (_FI_WORKSPACE_TENSOR is not None
182 + ), "Flashinfer must be enabled when using flashinfer"
183 + if norm_out is None:
184 + norm_out = allreduce_in
185 + residual_out = residual
186 + else:
187 + # return residual_out as allreduce_out with zeroed residual_in
188 + # as flashinfer does not support rms_norm
189 + # and allreduce_out together
190 + residual_out = allreduce_in
191 + # For the sizes that are smaller than the max size,
192 + # we only use flashinfer one shot allreduce
193 + flashinfer_comm.trtllm_allreduce_fusion(
194 + allreduce_in=allreduce_in,
195 + token_num=allreduce_in.shape[0],
196 + residual_in=residual,
197 + residual_out=residual_out,
198 + norm_out=norm_out,
199 + rms_gamma=rms_gamma,
200 + rms_eps=rms_eps,
201 + world_rank=world_rank,
202 + world_size=world_size,
203 + hidden_dim=allreduce_in.shape[-1],
204 + workspace_ptrs=_FI_WORKSPACE_TENSOR,
205 + launch_with_pdl=launch_with_pdl,
206 + use_oneshot=True,
207 + trigger_completion_at_end=trigger_completion_at_end,
208 + fp32_acc=fp32_acc,
209 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
210 + kARResidualRMSNorm,
211 + allreduce_out=None,
212 + quant_out=None,
213 + scale_out=None,
214 + layout_code=None,
215 + scale_factor=None,
216 + )
217 + else:
218 + allreduce_out = tensor_model_parallel_all_reduce(allreduce_in)
219 + if norm_out is None:
220 + torch.ops._C.fused_add_rms_norm(allreduce_out, residual,
221 + rms_gamma, rms_eps)
222 + else:
223 + torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma,
224 + rms_eps)
225 + allreduce_in.copy_(allreduce_out)
226 +
227 + def call_trtllm_fused_allreduce_norm_fake(
228 + allreduce_in: torch.Tensor,
229 + residual: torch.Tensor,
230 + rms_gamma: torch.Tensor,
231 + rms_eps: float,
232 + world_rank: int,
233 + world_size: int,
234 + launch_with_pdl: bool,
235 + trigger_completion_at_end: bool,
236 + fp32_acc: bool,
237 + max_token_num: int,
238 + norm_out: Optional[torch.Tensor] = None,
239 + ) -> None:
240 + pass
241 +
242 + direct_register_custom_op(
243 + op_name="flashinfer_trtllm_fused_allreduce_norm",
244 + op_func=call_trtllm_fused_allreduce_norm,
245 + mutates_args=[
246 + "allreduce_in",
247 + "residual",
248 + "norm_out",
249 + ],
250 + fake_impl=call_trtllm_fused_allreduce_norm_fake,
251 + dispatch_key=current_platform.dispatch_key,
252 + )
253 + flashinfer_trtllm_fused_allreduce_norm = (
254 + torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default)
255 +
256 +
257 +class FlashInferFusedAllReduceParams:
258 + """Parameters for FlashInfer fused allreduce operations."""
259 +
260 + def __init__(
261 + self,
262 + rank: int,
263 + world_size: int,
264 + use_fp32_lamport: bool = False,
265 + max_token_num: int = 1024,
266 + ):
267 + self.rank = rank
268 + self.world_size = world_size
269 + self.use_fp32_lamport = use_fp32_lamport
270 + self.trigger_completion_at_end = True
271 + self.launch_with_pdl = True
272 + self.fp32_acc = True
273 + self.use_oneshot = False
274 + self.max_token_num = max_token_num
275 +
276 + def get_trtllm_fused_allreduce_kwargs(self):
277 + return {
278 + "world_rank": self.rank,
279 + "world_size": self.world_size,
280 + "launch_with_pdl": self.launch_with_pdl,
281 + "trigger_completion_at_end": self.trigger_completion_at_end,
282 + "fp32_acc": self.fp32_acc,
283 + "max_token_num": self.max_token_num,
284 + }
285 +
286 +
287 +class AllReduceRMSNORMPattern(BasePattern):
288 +
289 + def __init__(
290 + self,
291 + epsilon: float,
292 + dtype: torch.dtype,
293 + device: str,
294 + allreduce_params: FlashInferFusedAllReduceParams,
295 + ):
296 + super().__init__(dtype, device)
297 + self.epsilon = epsilon
298 + self.allreduce_params = allreduce_params
299 +
300 + def get_inputs(self):
301 + input = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
302 + rms_result = torch.empty([1, 8, 4],
303 + device=self.device,
304 + dtype=self.dtype)
305 + weight = torch.empty([4], device=self.device, dtype=self.dtype)
306 +
307 + return [input, rms_result, weight]
308 +
309 + def register(self, pm_pass: PatternMatcherPass):
310 +
311 + def pattern(input: torch.Tensor, rms_result: torch.Tensor,
312 + weight: torch.Tensor):
313 + all_reduce_output = tensor_model_parallel_all_reduce(input)
314 + rms = auto_functionalized(
315 + RMS_OP,
316 + result=rms_result,
317 + input=all_reduce_output,
318 + weight=weight,
319 + epsilon=self.epsilon,
320 + )
321 + return rms[1], all_reduce_output
322 +
323 + def replacement(input: torch.Tensor, rms_result: torch.Tensor,
324 + weight: torch.Tensor):
325 + residual = torch.zeros_like(input)
326 + allreduce = auto_functionalized(
327 + torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default,
328 + allreduce_in=input,
329 + residual=residual,
330 + norm_out=rms_result,
331 + rms_gamma=weight,
332 + rms_eps=self.epsilon,
333 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
334 + )
335 +
336 + return allreduce[3], allreduce[1]
337 +
338 + pm.register_replacement(pattern, replacement, self.get_inputs(),
339 + pm.fwd_only, pm_pass)
340 +
341 +
342 +class AllReduceFusedAddRMSNormPattern(BasePattern):
343 +
344 + def __init__(
345 + self,
346 + epsilon: float,
347 + dtype: torch.dtype,
348 + device: str,
349 + allreduce_params: FlashInferFusedAllReduceParams,
350 + ):
351 + super().__init__(dtype, device)
352 + self.epsilon = epsilon
353 + self.allreduce_params = allreduce_params
354 +
355 + def get_inputs(self):
356 + input = torch.empty([4, 4], device=self.device, dtype=self.dtype)
357 + residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
358 + weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
359 + return [
360 + residual,
361 + input,
362 + weight,
363 + ]
364 +
365 + def register(self, pm_pass: PatternMatcherPass):
366 +
367 + def pattern(residual: torch.Tensor, input: torch.Tensor,
368 + weight: torch.Tensor):
369 + all_reduce_output = tensor_model_parallel_all_reduce(input)
370 + rms = auto_functionalized(
371 + RMS_ADD_OP,
372 + input=all_reduce_output,
373 + residual=residual,
374 + weight=weight,
375 + epsilon=self.epsilon,
376 + )
377 + return rms[1], rms[2]
378 +
379 + def replacement(residual: torch.Tensor, input: torch.Tensor,
380 + weight: torch.Tensor):
381 + allreduce = auto_functionalized(
382 + torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default,
383 + allreduce_in=input,
384 + residual=residual,
385 + rms_gamma=weight,
386 + rms_eps=self.epsilon,
387 + norm_out=None,
388 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
389 + )
390 + return allreduce[1], allreduce[2]
391 +
392 + pm.register_replacement(pattern, replacement, self.get_inputs(),
393 + pm.fwd_only, pm_pass)
394 +
395 +
396 +class AllReduceFusionPass(VllmInductorPass):
397 +
398 + def __init__(self, config: VllmConfig, max_token_num: int):
399 + super().__init__(config)
400 + self.disabled = True
401 + self.tp_size = get_tensor_model_parallel_world_size()
402 + if self.tp_size <= 1:
403 + return
404 + self.patterns: PatternMatcherPass = PatternMatcherPass(
405 + pass_name="all_reduce_fusion_pass")
406 + if config.model_config is None:
407 + return
408 + self.hidden_dim = config.model_config.get_hidden_size()
409 + self.group = get_tp_group().device_group
410 + rank = get_tensor_model_parallel_rank()
411 + use_fp32_lamport = self.model_dtype == torch.float32
412 + if flashinfer_comm is None:
413 + logger.warning(
414 + "Flashinfer is not installed, skipping allreduce fusion pass")
415 + return
416 + # Check if the world size is supported
417 + if self.tp_size not in _FI_MAX_SIZES:
418 + logger.warning(
419 + "Flashinfer allreduce fusion is not "
420 + "supported for world size %s",
421 + self.tp_size,
422 + )
423 + return
424 +
425 + self.ipc_handles, workspace_tensor = (
426 + flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
427 + tp_rank=rank,
428 + tp_size=self.tp_size,
429 + max_token_num=max_token_num,
430 + hidden_dim=self.hidden_dim,
431 + group=self.group,
432 + use_fp32_lamport=use_fp32_lamport,
433 + ))
434 +
435 + global _FI_WORKSPACE_TENSOR
436 + _FI_WORKSPACE_TENSOR = workspace_tensor
437 + self.allreduce_params = FlashInferFusedAllReduceParams(
438 + rank=rank,
439 + world_size=self.tp_size,
440 + use_fp32_lamport=use_fp32_lamport,
441 + max_token_num=max_token_num,
442 + )
443 +
444 + for epsilon in [1e-5, 1e-6]:
445 + AllReduceRMSNORMPattern(
446 + epsilon,
447 + self.model_dtype,
448 + self.device,
449 + self.allreduce_params,
450 + ).register(self.patterns)
451 + AllReduceFusedAddRMSNormPattern(
452 + epsilon,
453 + self.model_dtype,
454 + self.device,
455 + self.allreduce_params,
456 + ).register(self.patterns)
457 +
458 + self.disabled = False
459 +
460 + def __call__(self, graph: fx.Graph):
461 + if self.disabled:
462 + return
463 + self.begin()
464 + self.dump_graph(graph, "before_all_reduce_fusion_pass")
465 + count = self.patterns.apply(graph)
466 + logger.debug("Replaced %s patterns", count)
467 + self.dump_graph(graph, "after_all_reduce_fusion_pass")
468 + self.end_and_log()
469 +
470 + def __del__(self):
471 + if self.disabled:
472 + return
473 + if flashinfer_comm is not None:
474 + flashinfer_comm.trtllm_destroy_ipc_workspace(
475 + self.ipc_handles, self.group)
vllm/compilation/pass_manager.py
+6 -2
@@ -7,7 +7,7 @@ from vllm.config import VllmConfig
7 from vllm.logger import init_logger
8
9 from .activation_quant_fusion import ActivationQuantFusionPass
10 -from .collective_fusion import AsyncTPPass
10 +from .collective_fusion import AllReduceFusionPass, AsyncTPPass
11 from .fix_functionalization import FixFunctionalizationPass
12 from .fusion import FusionPass
13 from .fusion_attn import AttnFusionPass
@@ -62,7 +62,11 @@ class PostGradPassManager(CustomGraphPass):
62
63 if self.pass_config.enable_attn_fusion:
64 self.passes += [AttnFusionPass(config)]
65 -
65 + if self.pass_config.enable_fi_allreduce_fusion:
66 + self.passes += [
67 + AllReduceFusionPass(
68 + config, self.pass_config.fi_allreduce_fusion_max_token_num)
69 + ]
70 self.fix_functionalization = FixFunctionalizationPass(config)
71
72 def add(self, pass_: InductorPass):
vllm/config.py
+4
@@ -3962,6 +3962,10 @@ class PassConfig:
3962 """Whether to enable sequence parallelism."""
3963 enable_async_tp: bool = False
3964 """Whether to enable async TP."""
3965 + enable_fi_allreduce_fusion: bool = False
3966 + """Whether to enable flashinfer allreduce fusion."""
3967 + fi_allreduce_fusion_max_token_num: int = 1024
3968 + """Max number of tokens to used in flashinfer allreduce fusion."""
3969
3970 # TODO(luka) better pass enabling system.
3971