@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c3868bbbe4
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 4, 10:00 UTC

[compile] Add FlashInfer FP8 async TP fusion and preserve allreduce fusion ordering #27893 (#39505)

[compile] Add FlashInfer FP8 async TP fusion and preserve allreduce fusion ordering #27893 (#39505) Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com> Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Signed-off-by: roG0d <baonudesifeizhai@gmail.com>

baonudesifeizhai committed May 1, 2026 at 01:08 UTC c3868bbbe4b160d89adf339bcc069f6956314345
6 files changed +398 -31
tests/compile/conftest.py
+17 -3
@@ -24,10 +24,24 @@ def mock_cuda_platform():
24 def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None):
25 mock_platform = MagicMock()
26 mock_platform.is_cuda.return_value = is_cuda
27 - if capability is not None:
28 - mock_platform.get_device_capability.return_value = DeviceCapability(
29 - *capability
27 + device_capability = (
28 + DeviceCapability(*capability) if capability is not None else None
29 + )
30 + mock_platform.get_device_capability.return_value = device_capability
31 +
32 + def is_device_capability_family(
33 + requested_capability: int, device_id: int = 0
34 + ) -> bool:
35 + current_capability = mock_platform.get_device_capability(
36 + device_id=device_id
37 )
38 + if current_capability is None:
39 + return False
40 + return current_capability.major == (requested_capability // 10)
41 +
42 + mock_platform.is_device_capability_family.side_effect = (
43 + is_device_capability_family
44 + )
45 with patch("vllm.platforms.current_platform", mock_platform):
46 yield mock_platform
47
tests/compile/fusions_e2e/conftest.py
+6
@@ -97,6 +97,12 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
97 f"attention backend '{attn_backend.backend.name}'"
98 )
99
100 + if attn_backend.backend.name == "FLASHINFER":
101 + from vllm.utils.flashinfer import supports_trtllm_attention
102 +
103 + if not supports_trtllm_attention():
104 + matches = matches._replace(attn_quant_fusion=0)
105 +
106 # TODO: remove this after finishing migration from envs to model kwargs
107 if model_name == "openai/gpt-oss-20b":
108 from .common import is_blackwell
tests/compile/fusions_e2e/test_tp2_async_tp.py
-11
@@ -13,7 +13,6 @@ from .common import (
13 AttentionBackendCase,
14 Matches,
15 custom_ops_combos,
16 - is_blackwell,
16 )
17 from .models import (
18 FLASHINFER_ATTN,
@@ -46,14 +45,9 @@ def test_tp2_async_tp_fp8_fusions(
45 custom_ops: str,
46 inductor_graph_partition: bool,
47 run_e2e_fusion_test,
49 - monkeypatch,
48 ):
49 matches = matches_fn(n_layers)
50
53 - if is_blackwell():
54 - # Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns
55 - monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel")
56 -
51 # Reduce size of model and skip weight loading time
52 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
53 model_kwargs["load_format"] = "dummy"
@@ -173,14 +167,9 @@ def test_tp2_sp_ar_rms_fp8_fusions(
167 custom_ops: str,
168 inductor_graph_partition: bool,
169 run_e2e_fusion_test,
176 - monkeypatch,
170 ):
171 matches = matches_fn(n_layers)
172
180 - if is_blackwell():
181 - # Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns
182 - monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel")
183 -
173 # Reduce size of model and skip weight loading time
174 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
175 model_kwargs["load_format"] = "dummy"
vllm/compilation/passes/fusion/collective_fusion.py
+333 -16
@@ -1,8 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Callable
5 +from contextlib import suppress
6 +
7 import torch
8 import torch._inductor.pattern_matcher as pm
9 +import torch.distributed.distributed_c10d as c10d
10 import torch.fx as fx
11 from torch._inductor.pattern_matcher import PatternMatcherPass
12 from torch.distributed._symmetric_memory import enable_symm_mem_for_group
@@ -15,15 +19,197 @@ from vllm.distributed.parallel_state import (
19 )
20 from vllm.logger import init_logger
21 from vllm.platforms import current_platform
22 +from vllm.utils.torch_utils import direct_register_custom_op
23
24 from ..inductor_pass import enable_fake_mode
20 -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
25 +from ..vllm_inductor_pass import (
26 + VllmFusionPatternMatcherPass,
27 + VllmInductorPass,
28 + VllmPatternMatcherPass,
29 + VllmPatternReplacement,
30 +)
31
32 FP8_DTYPE = current_platform.fp8_dtype()
33
34 logger = init_logger(__name__)
35
36
37 +def _flashinfer_scaled_mm_out(
38 + A: torch.Tensor,
39 + B: torch.Tensor,
40 + *,
41 + scale_a: torch.Tensor,
42 + scale_b: torch.Tensor,
43 + out: torch.Tensor,
44 + bias: torch.Tensor | None = None,
45 + scale_result: torch.Tensor | None = None,
46 + out_dtype: torch.dtype | None = None,
47 + use_fast_accum: bool = False,
48 +) -> None:
49 + # Import lazily to avoid a circular import during module initialization
50 + # when docs or other tooling import the pass without FlashInfer.
51 + from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm_out
52 +
53 + assert bias is None, "FlashInfer symm_mem adapter does not support bias"
54 + assert scale_result is None, (
55 + "FlashInfer symm_mem adapter does not support result scaling"
56 + )
57 + assert not use_fast_accum, (
58 + "FlashInfer symm_mem adapter does not support use_fast_accum"
59 + )
60 + assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, (
61 + "FlashInfer symm_mem adapter expects 2D inputs and output"
62 + )
63 + assert scale_a.numel() == 1 and scale_b.numel() == 1, (
64 + "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales"
65 + )
66 +
67 + flashinfer_scaled_fp8_mm_out(
68 + A,
69 + B,
70 + scale_a,
71 + scale_b,
72 + out=out,
73 + out_dtype=out_dtype or out.dtype,
74 + )
75 +
76 +
77 +def fused_flashinfer_scaled_matmul_reduce_scatter_fake(
78 + A: torch.Tensor,
79 + B: torch.Tensor,
80 + A_scale: torch.Tensor,
81 + B_scale: torch.Tensor,
82 + reduce_op: str,
83 + orig_scatter_dim: int,
84 + scatter_dim_after_maybe_reshape: int,
85 + group_name: str,
86 + output_shape: list[int],
87 + out_dtype: torch.dtype | None = None,
88 +) -> torch.Tensor:
89 + world_size = c10d._resolve_process_group(group_name).size()
90 + result_shape = list(output_shape)
91 + result_shape[orig_scatter_dim] //= world_size
92 + return torch.empty(
93 + result_shape,
94 + dtype=out_dtype or torch.bfloat16,
95 + device=A.device,
96 + )
97 +
98 +
99 +def fused_flashinfer_scaled_matmul_reduce_scatter(
100 + A: torch.Tensor,
101 + B: torch.Tensor,
102 + A_scale: torch.Tensor,
103 + B_scale: torch.Tensor,
104 + reduce_op: str,
105 + orig_scatter_dim: int,
106 + scatter_dim_after_maybe_reshape: int,
107 + group_name: str,
108 + output_shape: list[int],
109 + out_dtype: torch.dtype | None = None,
110 +) -> torch.Tensor:
111 + assert orig_scatter_dim == 0 and scatter_dim_after_maybe_reshape == 0, (
112 + "FlashInfer symm_mem adapter currently only supports scatter_dim=0"
113 + )
114 + world_size = c10d._resolve_process_group(group_name).size()
115 + assert A.ndim == 2 and B.ndim == 2, "FlashInfer symm_mem adapter expects 2D inputs"
116 + assert A.is_contiguous(), "FlashInfer symm_mem adapter expects contiguous A"
117 + assert A_scale.numel() == 1 and B_scale.numel() == 1, (
118 + "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales"
119 + )
120 + assert A.shape[0] % world_size == 0, (
121 + "FlashInfer symm_mem adapter expects M divisible by world size"
122 + )
123 +
124 + kwargs = {
125 + "scale_b": B_scale,
126 + "bias": None,
127 + "scale_result": None,
128 + "out_dtype": out_dtype,
129 + "use_fast_accum": False,
130 + }
131 + return torch.distributed._symmetric_memory._fused_scaled_matmul_reduce_scatter_impl(
132 + mm_out_op=_flashinfer_scaled_mm_out,
133 + A=A,
134 + B=B,
135 + A_scale=A_scale,
136 + kwargs=kwargs,
137 + out_dtype=out_dtype,
138 + reduce_op=reduce_op,
139 + orig_scatter_dim=orig_scatter_dim,
140 + scatter_dim_after_maybe_reshape=scatter_dim_after_maybe_reshape,
141 + group_name=group_name,
142 + output_shape=output_shape,
143 + )
144 +
145 +
146 +def fused_all_gather_flashinfer_scaled_matmul_fake(
147 + A_shard: torch.Tensor,
148 + B: torch.Tensor,
149 + A_scale: torch.Tensor,
150 + B_scale: torch.Tensor,
151 + gather_dim: int,
152 + group_name: str,
153 + out_dtype: torch.dtype | None = None,
154 +) -> torch.Tensor:
155 + world_size = c10d._resolve_process_group(group_name).size()
156 + output_shape = list(A_shard.shape)
157 + output_shape[gather_dim] *= world_size
158 + output_shape[-1] = B.shape[1]
159 + return torch.empty(
160 + output_shape,
161 + dtype=out_dtype or torch.bfloat16,
162 + device=A_shard.device,
163 + )
164 +
165 +
166 +def fused_all_gather_flashinfer_scaled_matmul(
167 + A_shard: torch.Tensor,
168 + B: torch.Tensor,
169 + A_scale: torch.Tensor,
170 + B_scale: torch.Tensor,
171 + gather_dim: int,
172 + group_name: str,
173 + out_dtype: torch.dtype | None = None,
174 +) -> torch.Tensor:
175 + assert gather_dim == 0, (
176 + "FlashInfer symm_mem adapter currently only supports gather_dim=0"
177 + )
178 + _, outputs = torch.distributed._symmetric_memory._fused_all_gather_matmul_impl(
179 + mm_out_op=_flashinfer_scaled_mm_out,
180 + A_shard=A_shard,
181 + Bs=[B],
182 + A_scale=A_scale,
183 + kwargs_list=[
184 + {
185 + "scale_b": B_scale,
186 + "bias": None,
187 + "scale_result": None,
188 + "out_dtype": out_dtype,
189 + "use_fast_accum": False,
190 + }
191 + ],
192 + out_dtypes=[out_dtype],
193 + gather_dim=gather_dim,
194 + group_name=group_name,
195 + return_A=False,
196 + )
197 + return outputs[0]
198 +
199 +
200 +direct_register_custom_op(
201 + op_name="fused_flashinfer_scaled_matmul_reduce_scatter",
202 + op_func=fused_flashinfer_scaled_matmul_reduce_scatter,
203 + fake_impl=fused_flashinfer_scaled_matmul_reduce_scatter_fake,
204 +)
205 +
206 +direct_register_custom_op(
207 + op_name="fused_all_gather_flashinfer_scaled_matmul",
208 + op_func=fused_all_gather_flashinfer_scaled_matmul,
209 + fake_impl=fused_all_gather_flashinfer_scaled_matmul_fake,
210 +)
211 +
212 +
213 class BasePattern:
214 def __init__(self, dtype: torch.dtype, device: str | None) -> None:
215 self.dtype = dtype
@@ -371,39 +557,169 @@ class AllGatherCutlassScaledMMPattern(BasePattern):
557 )
558
559
374 -class AsyncTPPass(VllmPatternMatcherPass):
560 +class FlashInferBMMFP8ReduceScatterPattern(
561 + BasePattern, VllmPatternReplacement[..., torch.Tensor]
562 +):
563 + def get_inputs(self) -> list[torch.Tensor]:
564 + a_2d = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
565 + b_2d = (
566 + torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
567 + .contiguous()
568 + .transpose(0, 1)
569 + )
570 + a_scale = torch.empty([1], device=self.device, dtype=torch.float32)
571 + b_scale = torch.empty([1], device=self.device, dtype=torch.float32)
572 + return [a_2d, b_2d, a_scale, b_scale]
573 +
574 + @property
575 + def pattern(self) -> Callable[..., torch.Tensor]:
576 + def _pattern(
577 + a_2d: torch.Tensor,
578 + b_2d: torch.Tensor,
579 + a_scale: torch.Tensor,
580 + b_scale: torch.Tensor,
581 + ) -> torch.Tensor:
582 + bmm = torch.ops.vllm.bmm_fp8.default(
583 + torch.ops.aten.unsqueeze.default(a_2d, 0),
584 + torch.ops.aten.unsqueeze.default(b_2d, 0),
585 + a_scale,
586 + b_scale,
587 + self.dtype,
588 + "auto",
589 + )
590 + output = torch.ops.aten.reshape.default(bmm, list(bmm.shape[1:]))
591 + return torch.ops.vllm.reduce_scatter.default(
592 + output,
593 + dim=0,
594 + world_size=self.tp_size,
595 + group_name=self.tp.unique_name,
596 + )
597 +
598 + return _pattern
599 +
600 + @property
601 + def replacement(self) -> Callable[..., torch.Tensor]:
602 + def _replacement(
603 + a_2d: torch.Tensor,
604 + b_2d: torch.Tensor,
605 + a_scale: torch.Tensor,
606 + b_scale: torch.Tensor,
607 + ) -> torch.Tensor:
608 + return torch.ops.vllm.fused_flashinfer_scaled_matmul_reduce_scatter.default(
609 + a_2d,
610 + b_2d,
611 + a_scale,
612 + b_scale,
613 + "sum",
614 + 0,
615 + 0,
616 + self.tp.device_group.group_name,
617 + [a_2d.shape[0], b_2d.shape[1]],
618 + self.dtype,
619 + )
620 +
621 + return _replacement
622 +
623 +
624 +class FlashInferAllGatherBMMFP8Pattern(
625 + BasePattern, VllmPatternReplacement[..., torch.Tensor]
626 +):
627 + def get_inputs(self) -> list[torch.Tensor]:
628 + a_shard_2d = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
629 + b_2d = (
630 + torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
631 + .contiguous()
632 + .transpose(0, 1)
633 + )
634 + a_scale = torch.empty([1], device=self.device, dtype=torch.float32)
635 + b_scale = torch.empty([1], device=self.device, dtype=torch.float32)
636 + return [a_shard_2d, b_2d, a_scale, b_scale]
637 +
638 + @property
639 + def pattern(self) -> Callable[..., torch.Tensor]:
640 + def _pattern(
641 + a_shard_2d: torch.Tensor,
642 + b_2d: torch.Tensor,
643 + a_scale: torch.Tensor,
644 + b_scale: torch.Tensor,
645 + ) -> torch.Tensor:
646 + all_gather = torch.ops.vllm.all_gather.default(
647 + a_shard_2d,
648 + dim=0,
649 + world_size=self.tp_size,
650 + group_name=self.tp.unique_name,
651 + )
652 + return torch.ops.vllm.bmm_fp8.default(
653 + torch.ops.aten.unsqueeze.default(all_gather, 0),
654 + torch.ops.aten.unsqueeze.default(b_2d, 0),
655 + a_scale,
656 + b_scale,
657 + self.dtype,
658 + "auto",
659 + )
660 +
661 + return _pattern
662 +
663 + @property
664 + def replacement(self) -> Callable[..., torch.Tensor]:
665 + def _replacement(
666 + a_shard_2d: torch.Tensor,
667 + b_2d: torch.Tensor,
668 + a_scale: torch.Tensor,
669 + b_scale: torch.Tensor,
670 + ) -> torch.Tensor:
671 + fused = torch.ops.vllm.fused_all_gather_flashinfer_scaled_matmul.default(
672 + a_shard_2d,
673 + b_2d,
674 + a_scale,
675 + b_scale,
676 + 0,
677 + self.tp.device_group.group_name,
678 + self.dtype,
679 + )
680 + return torch.ops.aten.unsqueeze.default(fused, 0)
681 +
682 + return _replacement
683 +
684 +
685 +class AsyncTPPass(VllmFusionPatternMatcherPass):
686 @enable_fake_mode
687 def __init__(self, config: VllmConfig) -> None:
377 - super().__init__(config)
688 + super().__init__(config, pass_name="async_tp_pass")
689
379 - # Enable symmetric memory for the TP process group
690 enable_symm_mem_for_group(get_tp_group().device_group.group_name)
381 - self.patterns: PatternMatcherPass = PatternMatcherPass(
382 - pass_name="async_tp_pass"
383 - )
384 - GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.patterns)
691 + GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass)
692
386 - AllGatherGEMMPattern(self.model_dtype, self.device).register(self.patterns)
693 + AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass)
694
695 # These fusions are enabled only for bfloat16 models because
696 # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
697 # only supports bfloat16 as the output dtype.
698 if self.model_dtype == torch.bfloat16:
699 ScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
393 - self.patterns
700 + self.pm_pass
701 )
702 AllGatherScaledMMPattern(self.model_dtype, self.device).register(
396 - self.patterns
703 + self.pm_pass
704 )
705
706 CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
400 - self.patterns
707 + self.pm_pass
708 )
709 AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
403 - self.patterns
710 + self.pm_pass
711 )
405 -
406 - self.dump_patterns(config, self.patterns)
712 + with suppress(ImportError):
713 + import vllm.utils.flashinfer # noqa: F401
714 + if hasattr(torch.ops.vllm, "bmm_fp8"):
715 + self.register(
716 + FlashInferAllGatherBMMFP8Pattern(self.model_dtype, self.device)
717 + )
718 + self.register(
719 + FlashInferBMMFP8ReduceScatterPattern(self.model_dtype, self.device)
720 + )
721 +
722 + self.dump_patterns(config, self.pm_pass)
723
724 def is_applicable_for_range(self, compile_range: Range) -> bool:
725 # This pass is applied on top of the sequence parallelism pass,
@@ -416,5 +732,6 @@ class AsyncTPPass(VllmPatternMatcherPass):
732
733 @VllmInductorPass.time_and_log
734 def __call__(self, graph: fx.Graph) -> None:
419 - self.matched_count = self.patterns.apply(graph)
735 + self.matched_count = self.pm_pass.apply(graph)
736 + VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
737 logger.debug("Replaced %s patterns", self.matched_count)
vllm/compilation/passes/fusion/sequence_parallelism.py
+9 -1
@@ -31,6 +31,7 @@ logger = init_logger(__name__)
31 # Only apply sequence parallelism for models with hidden_size >= threshold
32 SP_MIN_HIDDEN_SIZE: dict[int, int] = {
33 90: 8192, # H100: only for models with hidden_size >= 8192
34 + 100: 8192, # Blackwell family: only for models with hidden_size >= 8192
35 }
36
37 # Min size per GPU per device capability for sequence parallelism
@@ -38,6 +39,8 @@ SP_MIN_HIDDEN_SIZE: dict[int, int] = {
39 # This ensures the threshold scales appropriately with tensor parallelism
40 SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = {
41 90: 8, # 8MB per GPU for H100
42 + # Use a more conservative threshold on Blackwell so TP8 starts later.
43 + 100: 32,
44 }
45
46
@@ -67,7 +70,12 @@ def get_sequence_parallelism_threshold(
70 capability = current_platform.get_device_capability()
71 if capability is None:
72 return None
70 - device_capability = capability.to_int()
73 +
74 + # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket.
75 + if current_platform.is_device_capability_family(100):
76 + device_capability = 100
77 + else:
78 + device_capability = capability.to_int()
79
80 # Check if device has configured thresholds
81 min_hidden_size = SP_MIN_HIDDEN_SIZE.get(device_capability)
vllm/utils/flashinfer.py
+33
@@ -715,6 +715,38 @@ def flashinfer_scaled_fp8_mm(
715 return output
716
717
718 +def flashinfer_scaled_fp8_mm_out(
719 + a: torch.Tensor,
720 + b: torch.Tensor,
721 + scale_a: torch.Tensor,
722 + scale_b: torch.Tensor,
723 + out: torch.Tensor,
724 + out_dtype: torch.dtype | None = None,
725 +) -> torch.Tensor:
726 + assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2
727 + assert a.shape[1] == b.shape[0]
728 + assert out.shape == (a.shape[0], b.shape[1])
729 + assert scale_a.numel() == 1 and scale_b.numel() == 1
730 + assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn
731 + assert out.device.type == "cuda"
732 + assert a.is_contiguous()
733 +
734 + from flashinfer import bmm_fp8 as bmm_fp8_
735 +
736 + bmm_fp8_(
737 + a.unsqueeze(0),
738 + # FlashInfer expects the weight in the same column-major view layout
739 + # consumed by flashinfer_scaled_fp8_mm, so keep the transposed view.
740 + b.unsqueeze(0),
741 + scale_a,
742 + scale_b,
743 + out_dtype or out.dtype,
744 + out.unsqueeze(0),
745 + "auto",
746 + )
747 + return out
748 +
749 +
750 def flashinfer_quant_nvfp4_8x4_sf_layout(
751 a: torch.Tensor, a_global_sf: torch.Tensor
752 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -833,6 +865,7 @@ __all__ = [
865 "use_trtllm_attention",
866 "flashinfer_scaled_fp4_mm",
867 "flashinfer_scaled_fp8_mm",
868 + "flashinfer_scaled_fp8_mm_out",
869 "flashinfer_quant_nvfp4_8x4_sf_layout",
870 "flashinfer_fp8_blockscale_gemm",
871 "should_use_flashinfer_for_blockscale_fp8_gemm",