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, 00:00 UTC
Add NVFP4 all-gather GEMM fusion for AsyncTP (#41882)
Signed-off-by: roG0d <baonudesifeizhai@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
baonudesifeizhai committed
May 9, 2026 at 21:13 UTC
bc5fdc1e6a718aec4b76b47b69ec5e7cb3c2412a
7 files changed
+605
-6
tests/compile/correctness_e2e/test_async_tp.py
+73
@@ -13,6 +13,17 @@ from tests.utils import (
13
from vllm.config import (
14
CompilationMode,
15
)
16
+from vllm.platforms import current_platform
17
+from vllm.utils.flashinfer import has_flashinfer
18
+
19
+NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
20
+NVFP4_HF_OVERRIDES = {
21
+ "num_hidden_layers": 4,
22
+ "hidden_size": 512,
23
+ "intermediate_size": 800,
24
+ "num_attention_heads": 4,
25
+ "num_key_value_heads": 1,
26
+}
27
28
29
@create_new_process_for_each_test()
@@ -82,3 +93,65 @@ def test_async_tp_pass_correctness(
93
]
94
95
compare_two_settings(model_id, async_tp_args, tp_args, method="generate")
96
+
97
+
98
+@create_new_process_for_each_test()
99
+def test_async_tp_pass_nvfp4_correctness(num_gpus_available: int, monkeypatch):
100
+ if (
101
+ not current_platform.is_cuda()
102
+ or not current_platform.is_device_capability_family(100)
103
+ ):
104
+ pytest.skip("NVFP4 requires Blackwell")
105
+ if not has_flashinfer():
106
+ pytest.skip("FlashInfer is required for the NVFP4 AsyncTP path")
107
+
108
+ monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", "flashinfer-cutlass")
109
+
110
+ tp_size = 2
111
+ if num_gpus_available < tp_size:
112
+ pytest.skip(f"Need at least {tp_size} GPUs")
113
+
114
+ common_args = [
115
+ "--dtype",
116
+ "bfloat16",
117
+ "--max-model-len",
118
+ "2048",
119
+ "--max-num-seqs",
120
+ "8",
121
+ "--load-format",
122
+ "dummy",
123
+ "--hf-overrides",
124
+ json.dumps(NVFP4_HF_OVERRIDES),
125
+ ]
126
+
127
+ compilation_config = {
128
+ "mode": CompilationMode.VLLM_COMPILE,
129
+ "compile_sizes": [2, 4, 8],
130
+ "splitting_ops": [],
131
+ "pass_config": {
132
+ "enable_sp": True,
133
+ "fuse_gemm_comms": True,
134
+ "fuse_allreduce_rms": False,
135
+ "sp_min_token_num": 1,
136
+ },
137
+ }
138
+
139
+ async_tp_args = [
140
+ *common_args,
141
+ "--tensor-parallel-size",
142
+ str(tp_size),
143
+ "--distributed-executor-backend",
144
+ "mp",
145
+ "--compilation_config",
146
+ json.dumps(compilation_config),
147
+ ]
148
+
149
+ tp_args = [
150
+ *common_args,
151
+ "--tensor-parallel-size",
152
+ str(tp_size),
153
+ "--distributed-executor-backend",
154
+ "mp",
155
+ ]
156
+
157
+ compare_two_settings(NVFP4_MODEL_ID, async_tp_args, tp_args, method="generate")
tests/compile/correctness_e2e/test_sequence_parallel.py
+44
-5
@@ -21,12 +21,14 @@ from vllm.logger import init_logger
21
from vllm.platforms import current_platform
22
from vllm.utils.torch_utils import is_torch_equal_or_newer
23
24
-from ...models.registry import HF_EXAMPLE_MODELS
24
+from ...models.registry import HF_EXAMPLE_MODELS, _HfExamplesInfo
25
from ...utils import compare_two_settings, create_new_process_for_each_test
26
27
logger = init_logger("test_sequence_parallel")
28
29
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
30
+NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
31
+NVFP4_MODEL_INFO = _HfExamplesInfo(NVFP4_MODEL_ID)
32
33
34
class ParallelSetup(NamedTuple):
@@ -41,6 +43,7 @@ class ParallelSetup(NamedTuple):
43
class SPTestOptions(NamedTuple):
44
multi_node_only: bool
45
load_format: str | None = None
46
+ model_info: _HfExamplesInfo | None = None
47
48
49
@dataclass
@@ -170,6 +173,7 @@ def _compare_sp(
173
*,
174
method: Literal["generate", "encode"],
175
is_multimodal: bool,
176
+ dtype: str = "float16",
177
):
178
(
179
tp_size,
@@ -180,14 +184,15 @@ def _compare_sp(
184
chunked_prefill,
185
) = parallel_setup
186
183
- multi_node_only, load_format = test_options
187
+ multi_node_only = test_options.multi_node_only
188
+ load_format = test_options.load_format
189
185
- model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
190
+ model_info = test_options.model_info or HF_EXAMPLE_MODELS.find_hf_info(model_id)
191
model_info.check_transformers_version(on_fail="skip")
192
193
trust_remote_code = model_info.trust_remote_code
194
tokenizer_mode = model_info.tokenizer_mode
190
- hf_overrides = model_info.hf_overrides
195
+ hf_overrides = dict(model_info.hf_overrides)
196
require_embed_inputs = model_info.require_embed_inputs
197
198
if load_format == "dummy":
@@ -220,7 +225,7 @@ def _compare_sp(
225
common_args = [
226
# use half precision for speed and memory savings in CI environment
227
"--dtype",
223
- "float16",
228
+ dtype,
229
"--max-model-len",
230
"2048",
231
"--max-num-seqs",
@@ -352,3 +357,37 @@ def test_tp_sp_generation(
357
method="generate",
358
is_multimodal=False,
359
)
360
+
361
+
362
+@create_new_process_for_each_test()
363
+def test_tp_sp_nvfp4_generation(num_gpus_available: int):
364
+ if (
365
+ not current_platform.is_cuda()
366
+ or not current_platform.is_device_capability_family(100)
367
+ ):
368
+ pytest.skip("NVFP4 requires Blackwell")
369
+
370
+ _compare_sp(
371
+ NVFP4_MODEL_ID,
372
+ ParallelSetup(
373
+ tp_size=2,
374
+ pp_size=1,
375
+ fuse_norm_quant=True,
376
+ fuse_act_quant=True,
377
+ eager_mode=True,
378
+ chunked_prefill=False,
379
+ ),
380
+ "mp",
381
+ "auto",
382
+ SPTestOptions(
383
+ multi_node_only=False,
384
+ load_format="dummy",
385
+ model_info=NVFP4_MODEL_INFO,
386
+ ),
387
+ num_gpus_available,
388
+ use_inductor_graph_partition=False,
389
+ fuse_gemm_comms=False,
390
+ method="generate",
391
+ is_multimodal=False,
392
+ dtype="bfloat16",
393
+ )
tests/compile/fullgraph/test_toy_llama.py
+2
-1
@@ -17,7 +17,6 @@ import pytest
17
import torch
18
from torch import nn
19
20
-from vllm.compilation.counter import compilation_counter
20
from vllm.compilation.decorators import support_torch_compile
21
from vllm.config import (
22
CompilationConfig,
@@ -340,6 +339,8 @@ def run_model(llama_config, compile_config: CompilationConfig) -> torch.Tensor:
339
def test_toy_llama(
340
backend: str, use_inductor_graph_partition: bool, monkeypatch, tmp_path
341
):
342
+ from vllm.compilation.counter import compilation_counter
343
+
344
# We disable the vLLM compile cache into a new tmp dir for 1 reason:
345
# 1. To make sure we can properly track the number of Inductor compilations.
346
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
tests/compile/fusions_e2e/test_tp2_async_tp.py
+65
@@ -13,11 +13,13 @@ from .common import (
13
AttentionBackendCase,
14
Matches,
15
custom_ops_combos,
16
+ is_blackwell,
17
)
18
from .models import (
19
FLASHINFER_ATTN,
20
TRITON_ATTN,
21
llama3_8b,
22
+ llama3_8b_fp4,
23
llama3_8b_fp8,
24
llama4_scout_fp8,
25
qwen3_a3b,
@@ -90,6 +92,69 @@ def test_tp2_async_tp_fp8_fusions(
92
)
93
94
95
+@multi_gpu_test(num_gpus=2)
96
+@pytest.mark.parametrize(
97
+ "model_name, matches_fn, model_kwargs, hf_overrides",
98
+ [llama3_8b_fp4],
99
+)
100
+@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
101
+@pytest.mark.parametrize("n_layers", [4])
102
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
103
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
104
+@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
105
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
106
+def test_tp2_async_tp_nvfp4_fusions(
107
+ model_name: str,
108
+ matches_fn: Callable[[int], Matches],
109
+ model_kwargs: dict,
110
+ hf_overrides: Callable[[int], dict],
111
+ attn_backend: AttentionBackendCase,
112
+ n_layers: int,
113
+ custom_ops: str,
114
+ inductor_graph_partition: bool,
115
+ run_e2e_fusion_test,
116
+):
117
+ # NVFP4 currently wires the all-gather + GEMM path only.
118
+ matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2)
119
+
120
+ # Reduce size of model and skip weight loading time
121
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
122
+ model_kwargs["load_format"] = "dummy"
123
+ model_kwargs["max_model_len"] = 1024
124
+ model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
125
+
126
+ compilation_config = dict(
127
+ use_inductor_graph_partition=inductor_graph_partition,
128
+ custom_ops=custom_ops.split(","),
129
+ pass_config=PassConfig(
130
+ fuse_act_quant=True,
131
+ fuse_attn_quant=True,
132
+ enable_sp=True,
133
+ fuse_gemm_comms=True,
134
+ fuse_allreduce_rms=False,
135
+ # Override threshold for testing (models have small hidden_size)
136
+ sp_min_token_num=512,
137
+ ),
138
+ )
139
+
140
+ matches_check = [
141
+ "act_quant_fusion",
142
+ "attn_quant_fusion",
143
+ "sequence_parallel",
144
+ "async_tp",
145
+ ]
146
+
147
+ run_e2e_fusion_test(
148
+ model_name,
149
+ matches,
150
+ model_kwargs,
151
+ attn_backend,
152
+ compilation_config,
153
+ matches_check,
154
+ tp_size=2,
155
+ )
156
+
157
+
158
@multi_gpu_test(num_gpus=2)
159
@pytest.mark.parametrize(
160
"model_name, matches_fn, model_kwargs, hf_overrides",
vllm/compilation/passes/fusion/collective_fusion.py
+243
@@ -74,6 +74,36 @@ def _flashinfer_scaled_mm_out(
74
)
75
76
77
+def _flashinfer_fp4_mm_out(
78
+ A: torch.Tensor,
79
+ B: torch.Tensor,
80
+ *,
81
+ scale_a: torch.Tensor,
82
+ scale_b: torch.Tensor,
83
+ out: torch.Tensor,
84
+ alpha: torch.Tensor,
85
+ out_dtype: torch.dtype | None = None,
86
+ use_8x4_sf_layout: bool = False,
87
+ backend: str = "cutlass",
88
+) -> None:
89
+ from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm_out
90
+
91
+ assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, (
92
+ "FlashInfer FP4 symm_mem adapter expects 2D inputs and output"
93
+ )
94
+ flashinfer_scaled_fp4_mm_out(
95
+ A,
96
+ B,
97
+ scale_a,
98
+ scale_b,
99
+ alpha,
100
+ out=out,
101
+ out_dtype=out_dtype or out.dtype,
102
+ use_8x4_sf_layout=use_8x4_sf_layout,
103
+ backend=backend,
104
+ )
105
+
106
+
107
def fused_flashinfer_scaled_matmul_reduce_scatter_fake(
108
A: torch.Tensor,
109
B: torch.Tensor,
@@ -197,6 +227,90 @@ def fused_all_gather_flashinfer_scaled_matmul(
227
return outputs[0]
228
229
230
+def fused_all_gather_flashinfer_fp4_matmul_fake(
231
+ A_shard: torch.Tensor,
232
+ B: torch.Tensor,
233
+ A_scale_shard: torch.Tensor,
234
+ B_scale: torch.Tensor,
235
+ alpha: torch.Tensor,
236
+ gather_dim: int,
237
+ group_name: str,
238
+ out_dtype: torch.dtype | None = None,
239
+ view_a_scale_as_fp8: bool = False,
240
+ use_8x4_sf_layout: bool = False,
241
+ backend: str = "cutlass",
242
+) -> torch.Tensor:
243
+ world_size = c10d._resolve_process_group(group_name).size()
244
+ output_shape = list(A_shard.shape)
245
+ output_shape[gather_dim] *= world_size
246
+ output_shape[-1] = B.shape[1]
247
+ return torch.empty(
248
+ output_shape,
249
+ dtype=out_dtype or torch.bfloat16,
250
+ device=A_shard.device,
251
+ )
252
+
253
+
254
+def fused_all_gather_flashinfer_fp4_matmul(
255
+ A_shard: torch.Tensor,
256
+ B: torch.Tensor,
257
+ A_scale_shard: torch.Tensor,
258
+ B_scale: torch.Tensor,
259
+ alpha: torch.Tensor,
260
+ gather_dim: int,
261
+ group_name: str,
262
+ out_dtype: torch.dtype | None = None,
263
+ view_a_scale_as_fp8: bool = False,
264
+ use_8x4_sf_layout: bool = False,
265
+ backend: str = "cutlass",
266
+) -> torch.Tensor:
267
+ assert gather_dim == 0, (
268
+ "FlashInfer FP4 symm_mem adapter currently only supports gather_dim=0"
269
+ )
270
+ assert A_shard.ndim == 2 and A_scale_shard.ndim == 2 and B.ndim == 2, (
271
+ "FlashInfer FP4 symm_mem adapter expects 2D inputs"
272
+ )
273
+ if view_a_scale_as_fp8:
274
+ A_scale_shard = A_scale_shard.view(torch.float8_e4m3fn)
275
+
276
+ group = c10d._resolve_process_group(group_name)
277
+ world_size = group.size()
278
+ output = A_shard.new_empty(
279
+ A_shard.shape[0] * world_size,
280
+ B.shape[1],
281
+ dtype=out_dtype or torch.bfloat16,
282
+ )
283
+ output_shards = output.chunk(world_size)
284
+
285
+ A = A_shard.new_empty(A_shard.shape[0] * world_size, A_shard.shape[1])
286
+ A_scale = A_scale_shard.new_empty(
287
+ A_scale_shard.shape[0] * world_size,
288
+ A_scale_shard.shape[1],
289
+ )
290
+
291
+ def fp4_shard_consumer(shards: list[torch.Tensor], rank: int) -> None:
292
+ _flashinfer_fp4_mm_out(
293
+ shards[0],
294
+ B,
295
+ scale_a=shards[1],
296
+ scale_b=B_scale,
297
+ alpha=alpha,
298
+ out=output_shards[rank],
299
+ out_dtype=out_dtype,
300
+ use_8x4_sf_layout=use_8x4_sf_layout,
301
+ backend=backend,
302
+ )
303
+
304
+ torch.distributed._symmetric_memory._pipelined_multi_all_gather_and_consume(
305
+ [A_shard, A_scale_shard],
306
+ fp4_shard_consumer,
307
+ [A, A_scale],
308
+ group_name,
309
+ False,
310
+ )
311
+ return output
312
+
313
+
314
direct_register_custom_op(
315
op_name="fused_flashinfer_scaled_matmul_reduce_scatter",
316
op_func=fused_flashinfer_scaled_matmul_reduce_scatter,
@@ -209,6 +323,12 @@ direct_register_custom_op(
323
fake_impl=fused_all_gather_flashinfer_scaled_matmul_fake,
324
)
325
326
+direct_register_custom_op(
327
+ op_name="fused_all_gather_flashinfer_fp4_matmul",
328
+ op_func=fused_all_gather_flashinfer_fp4_matmul,
329
+ fake_impl=fused_all_gather_flashinfer_fp4_matmul_fake,
330
+)
331
+
332
333
class BasePattern:
334
def __init__(self, dtype: torch.dtype, device: str | None) -> None:
@@ -682,6 +802,101 @@ class FlashInferAllGatherBMMFP8Pattern(
802
return _replacement
803
804
805
+class FlashInferAllGatherFP4Pattern(
806
+ BasePattern, VllmPatternReplacement[..., torch.Tensor]
807
+):
808
+ def __init__(
809
+ self,
810
+ dtype: torch.dtype,
811
+ device: str | None,
812
+ backend: str,
813
+ use_8x4_sf_layout: bool,
814
+ a_scale_view: str,
815
+ ) -> None:
816
+ super().__init__(dtype, device)
817
+ self.backend = backend
818
+ self.use_8x4_sf_layout = use_8x4_sf_layout
819
+ self.a_scale_view = a_scale_view
820
+
821
+ def get_inputs(self) -> list[torch.Tensor]:
822
+ a_shard_2d = torch.empty([8, 8], device=self.device, dtype=torch.uint8)
823
+ b_2d = torch.empty([8, 16], device=self.device, dtype=torch.uint8)
824
+ a_scale_shard = torch.empty([128, 4], device=self.device, dtype=torch.int32)
825
+ b_scale = torch.empty([4, 128], device=self.device, dtype=torch.uint8)
826
+ alpha = torch.empty([], device=self.device, dtype=torch.float32)
827
+ return [
828
+ a_shard_2d,
829
+ b_2d,
830
+ a_scale_shard,
831
+ b_scale,
832
+ alpha,
833
+ ]
834
+
835
+ @property
836
+ def pattern(self) -> Callable[..., torch.Tensor]:
837
+ def _pattern(
838
+ a_shard_2d: torch.Tensor,
839
+ b_2d: torch.Tensor,
840
+ a_scale_shard: torch.Tensor,
841
+ b_scale: torch.Tensor,
842
+ alpha: torch.Tensor,
843
+ ) -> torch.Tensor:
844
+ all_gather_a = torch.ops.vllm.all_gather.default(
845
+ a_shard_2d,
846
+ dim=0,
847
+ world_size=self.tp_size,
848
+ group_name=self.tp.unique_name,
849
+ )
850
+ all_gather_a_scale = torch.ops.vllm.all_gather.default(
851
+ a_scale_shard,
852
+ dim=0,
853
+ world_size=self.tp_size,
854
+ group_name=self.tp.unique_name,
855
+ )
856
+ a_scale = all_gather_a_scale
857
+ if self.a_scale_view in ("float8", "float8_uint8"):
858
+ a_scale = torch.ops.aten.view.dtype(a_scale, torch.float8_e4m3fn)
859
+ if self.a_scale_view in ("uint8", "float8_uint8"):
860
+ a_scale = torch.ops.aten.view.dtype(a_scale, torch.uint8)
861
+ return torch.ops.vllm.flashinfer_mm_fp4.default(
862
+ all_gather_a,
863
+ b_2d,
864
+ a_scale,
865
+ b_scale,
866
+ alpha,
867
+ self.dtype,
868
+ self.use_8x4_sf_layout,
869
+ self.backend,
870
+ )
871
+
872
+ return _pattern
873
+
874
+ @property
875
+ def replacement(self) -> Callable[..., torch.Tensor]:
876
+ def _replacement(
877
+ a_shard_2d: torch.Tensor,
878
+ b_2d: torch.Tensor,
879
+ a_scale_shard: torch.Tensor,
880
+ b_scale: torch.Tensor,
881
+ alpha: torch.Tensor,
882
+ ) -> torch.Tensor:
883
+ return torch.ops.vllm.fused_all_gather_flashinfer_fp4_matmul.default(
884
+ a_shard_2d,
885
+ b_2d,
886
+ a_scale_shard,
887
+ b_scale,
888
+ alpha,
889
+ 0,
890
+ self.tp.device_group.group_name,
891
+ self.dtype,
892
+ self.a_scale_view in ("float8", "float8_uint8"),
893
+ self.use_8x4_sf_layout,
894
+ self.backend,
895
+ )
896
+
897
+ return _replacement
898
+
899
+
900
class AsyncTPPass(VllmFusionPatternMatcherPass):
901
@enable_fake_mode
902
def __init__(self, config: VllmConfig) -> None:
@@ -718,6 +933,34 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
933
self.register(
934
FlashInferBMMFP8ReduceScatterPattern(self.model_dtype, self.device)
935
)
936
+ if hasattr(torch.ops.vllm, "flashinfer_mm_fp4"):
937
+ for backend in ("cutlass", "cudnn"):
938
+ for a_scale_view in ("float8_uint8", "uint8"):
939
+ self.register(
940
+ FlashInferAllGatherFP4Pattern(
941
+ self.model_dtype,
942
+ self.device,
943
+ backend,
944
+ use_8x4_sf_layout=False,
945
+ a_scale_view=a_scale_view,
946
+ )
947
+ )
948
+ for use_8x4_sf_layout in (False, True):
949
+ for a_scale_view in ("float8",):
950
+ self.register(
951
+ FlashInferAllGatherFP4Pattern(
952
+ self.model_dtype,
953
+ self.device,
954
+ "trtllm",
955
+ use_8x4_sf_layout=use_8x4_sf_layout,
956
+ a_scale_view=a_scale_view,
957
+ )
958
+ )
959
+ # NVFP4 reduce-scatter does not need scale communication: FP4
960
+ # scales are consumed by the local GEMM and only BF16 partial
961
+ # outputs are reduced. Keep this PR scoped to the all-gather
962
+ # path; reduce-scatter needs a dedicated FP4 producer rather
963
+ # than the existing FP8-style helper.
964
965
self.dump_patterns(config, self.pm_pass)
966
vllm/compilation/passes/fusion/sequence_parallelism.py
+136
@@ -8,6 +8,7 @@ from typing import Any
8
import torch
9
import torch._inductor.pattern_matcher as pm
10
import torch.fx as fx
11
+from torch._higher_order_ops.auto_functionalize import auto_functionalized
12
from torch._inductor.pattern_matcher import PatternMatcherPass
13
14
import vllm.ir.ops
@@ -27,6 +28,10 @@ from .matcher_utils import MatcherQuantFP8
28
29
logger = init_logger(__name__)
30
31
+if hasattr(torch.ops._C, "scaled_fp4_quant"):
32
+ SCALED_FP4_QUANT_OUT_OVERLOAD = torch.ops._C.scaled_fp4_quant.out
33
+ SCALED_FP4_QUANT_DEFAULT_OVERLOAD = torch.ops._C.scaled_fp4_quant.default
34
+
35
# Min hidden size per device capability for sequence parallelism
36
# Only apply sequence parallelism for models with hidden_size >= threshold
37
SP_MIN_HIDDEN_SIZE: dict[int, int] = {
@@ -332,6 +337,129 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
337
)
338
339
340
+class FirstAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper):
341
+ def get_inputs(self) -> list[torch.Tensor]:
342
+ input = self.empty([8, 16])
343
+ weight = self.empty([16])
344
+ input_global_scale = self.empty_f32([1, 1])
345
+ quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8)
346
+ output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
347
+ return [input, weight, input_global_scale, quant_output, output_scale]
348
+
349
+ def register(self, pm_pass: PatternMatcherPass) -> None:
350
+ def pattern(
351
+ input: torch.Tensor,
352
+ weight: torch.Tensor,
353
+ input_global_scale: torch.Tensor,
354
+ quant_output: torch.Tensor,
355
+ output_scale: torch.Tensor,
356
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
357
+ all_reduce = self._all_reduce(input)
358
+ rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon)
359
+ quant = auto_functionalized(
360
+ SCALED_FP4_QUANT_OUT_OVERLOAD,
361
+ input=rms,
362
+ input_scale=input_global_scale,
363
+ is_sf_swizzled_layout=True,
364
+ output=quant_output,
365
+ output_scale=output_scale,
366
+ )
367
+ return quant[1], all_reduce, quant[2]
368
+
369
+ def replacement(
370
+ input: torch.Tensor,
371
+ weight: torch.Tensor,
372
+ input_global_scale: torch.Tensor,
373
+ quant_output: torch.Tensor,
374
+ output_scale: torch.Tensor,
375
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
376
+ reduce_scatter = self._reduce_scatter(input)
377
+ rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon)
378
+ rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]])
379
+ quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD(
380
+ rms,
381
+ input_global_scale,
382
+ True,
383
+ )
384
+ return (
385
+ self._all_gather(quant[0]),
386
+ reduce_scatter,
387
+ self._all_gather(quant[1]),
388
+ )
389
+
390
+ pm.register_replacement(
391
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
392
+ )
393
+
394
+
395
+class MiddleAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper):
396
+ def get_inputs(self) -> list[torch.Tensor]:
397
+ mm_1 = self.empty([8, 16])
398
+ residual = self.empty([8, 16])
399
+ rms_norm_weights = self.empty([16])
400
+ input_global_scale = self.empty_f32([1, 1])
401
+ quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8)
402
+ output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
403
+ return [
404
+ residual,
405
+ mm_1,
406
+ rms_norm_weights,
407
+ input_global_scale,
408
+ quant_output,
409
+ output_scale,
410
+ ]
411
+
412
+ def register(self, pm_pass: PatternMatcherPass) -> None:
413
+ def pattern(
414
+ residual: torch.Tensor,
415
+ mm_1: torch.Tensor,
416
+ rms_norm_weights: torch.Tensor,
417
+ input_global_scale: torch.Tensor,
418
+ quant_output: torch.Tensor,
419
+ output_scale: torch.Tensor,
420
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
421
+ all_reduce = self._all_reduce(mm_1)
422
+ rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
423
+ all_reduce, residual, rms_norm_weights, self.epsilon
424
+ )
425
+ quant = auto_functionalized(
426
+ SCALED_FP4_QUANT_OUT_OVERLOAD,
427
+ input=rms,
428
+ input_scale=input_global_scale,
429
+ is_sf_swizzled_layout=True,
430
+ output=quant_output,
431
+ output_scale=output_scale,
432
+ )
433
+ return quant[1], residual_out, quant[2]
434
+
435
+ def replacement(
436
+ residual: torch.Tensor,
437
+ mm_1: torch.Tensor,
438
+ rms_norm_weights: torch.Tensor,
439
+ input_global_scale: torch.Tensor,
440
+ quant_output: torch.Tensor,
441
+ output_scale: torch.Tensor,
442
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
443
+ # Keep this slice in sync with the non-quantized SP replacement:
444
+ # once the previous SP pattern fires, it becomes a no-op.
445
+ reduce_scatter = self._reduce_scatter(mm_1)
446
+ residual = residual[0 : reduce_scatter.size(0), ...]
447
+ rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
448
+ reduce_scatter, residual, rms_norm_weights, self.epsilon
449
+ )
450
+ rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]])
451
+ quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD(
452
+ rms,
453
+ input_global_scale,
454
+ True,
455
+ )
456
+ return self._all_gather(quant[0]), residual_out, self._all_gather(quant[1])
457
+
458
+ pm.register_replacement(
459
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
460
+ )
461
+
462
+
463
class SequenceParallelismPass(VllmPatternMatcherPass):
464
"""
465
This pass enables sequence parallelism for models.
@@ -404,6 +532,14 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
532
epsilon, self.model_dtype, self.device
533
).register(self.patterns)
534
535
+ if "SCALED_FP4_QUANT_OUT_OVERLOAD" in globals():
536
+ FirstAllReduceRMSNormStaticNVFP4Pattern(
537
+ epsilon, self.model_dtype, self.device
538
+ ).register(self.patterns)
539
+ MiddleAllReduceRMSNormStaticNVFP4Pattern(
540
+ epsilon, self.model_dtype, self.device
541
+ ).register(self.patterns)
542
+
543
# Normal RMSNorm patterns
544
FirstAllReduceRMSNormPattern(
545
epsilon, self.model_dtype, self.device
vllm/utils/flashinfer.py
+42
@@ -685,6 +685,47 @@ def flashinfer_scaled_fp4_mm(
685
)
686
687
688
+def flashinfer_scaled_fp4_mm_out(
689
+ a: torch.Tensor,
690
+ b: torch.Tensor,
691
+ block_scale_a: torch.Tensor,
692
+ block_scale_b: torch.Tensor,
693
+ alpha: torch.Tensor,
694
+ out: torch.Tensor,
695
+ out_dtype: torch.dtype | None,
696
+ use_8x4_sf_layout: bool,
697
+ backend: str,
698
+) -> torch.Tensor:
699
+ assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2
700
+ assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2
701
+ assert a.stride(-1) == 1
702
+ assert a.shape[1] == b.shape[0]
703
+ assert out.shape == (a.shape[0], b.shape[1])
704
+ assert out.device.type == "cuda"
705
+
706
+ if backend in ("cutlass", "cudnn"):
707
+ if block_scale_a.dtype != torch.uint8:
708
+ block_scale_a = block_scale_a.view(torch.uint8)
709
+ if block_scale_b.dtype != torch.uint8:
710
+ block_scale_b = block_scale_b.view(torch.uint8)
711
+
712
+ from flashinfer import mm_fp4 as flashinfer_mm_fp4_
713
+
714
+ flashinfer_mm_fp4_(
715
+ a,
716
+ b,
717
+ block_scale_a,
718
+ block_scale_b,
719
+ alpha,
720
+ out_dtype or out.dtype,
721
+ out=out,
722
+ block_size=16,
723
+ use_8x4_sf_layout=use_8x4_sf_layout,
724
+ backend=backend,
725
+ )
726
+ return out
727
+
728
+
729
def flashinfer_scaled_fp8_mm(
730
a: torch.Tensor,
731
b: torch.Tensor,
@@ -864,6 +905,7 @@ __all__ = [
905
"can_use_trtllm_attention",
906
"use_trtllm_attention",
907
"flashinfer_scaled_fp4_mm",
908
+ "flashinfer_scaled_fp4_mm_out",
909
"flashinfer_scaled_fp8_mm",
910
"flashinfer_scaled_fp8_mm_out",
911
"flashinfer_quant_nvfp4_8x4_sf_layout",