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

[compile] Enable sequence parallelism matching w/o custom ops enabled (#27126)

Signed-off-by: angelayi <yiangela7@gmail.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Signed-off-by: ProExpertProg <lgovedic@redhat.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <luka.govedic@gmail.com>

Angela Yi committed Nov 15, 2025 at 03:46 UTC f36292dbee27a5ebe0e7115c061b82f6f5372dcf
6 files changed +472 -444
.buildkite/test-pipeline.yaml
+8 -6
@@ -478,10 +478,11 @@ steps:
478 - vllm/
479 - tests/compile
480 commands:
481 + # fp8 kv scales not supported on sm89, tested on Blackwell instead
482 - pytest -v -s compile/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
483 # Limit to no custom ops to reduce running time
484 # Wrap with quotes to escape yaml and avoid starting -k string with a -
484 - - "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and -quant_fp8'"
485 + - "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
486
487 - label: Cudagraph test
488 timeout_in_minutes: 20
@@ -925,7 +926,7 @@ steps:
926 - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
927 - pytest -v -s tests/kernels/moe/test_flashinfer.py
928
928 -- label: Blackwell Fusion Tests # 30 min
929 +- label: Blackwell Fusion & Compile Tests # 30 min
930 timeout_in_minutes: 40
931 working_dir: "/vllm-workspace/"
932 gpu: b200
@@ -946,7 +947,9 @@ steps:
947 - pytest -v -s tests/compile/test_fusion_all_reduce.py
948 # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
949 # Wrap with quotes to escape yaml
949 - - "pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and Llama-3.1 and -quant_fp8 and -rms_norm'"
950 + - "pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
951 + # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
952 + - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile
953
954 - label: Blackwell Fusion E2E Tests # 30 min
955 timeout_in_minutes: 40
@@ -969,8 +972,6 @@ steps:
972 - nvidia-smi
973 # Run all e2e fusion tests
974 - pytest -v -s tests/compile/test_fusions_e2e.py
972 - # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
973 - - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile
975
976 - label: Blackwell GPT-OSS Eval
977 timeout_in_minutes: 60
@@ -1266,7 +1267,8 @@ steps:
1267 - pytest -v -s tests/compile/test_async_tp.py
1268 - pytest -v -s tests/compile/test_sequence_parallelism.py
1269 - pytest -v -s tests/compile/test_fusion_all_reduce.py
1269 - - pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
1270 + - "pytest -v -s tests/compile/test_fusions_e2e.py -k 'not Llama-4'"
1271 + - pytest -v -s tests/distributed/test_sequence_parallel.py
1272 - pytest -v -s tests/distributed/test_context_parallel.py
1273 - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048
1274 - pytest -v -s tests/v1/distributed/test_dbo.py
tests/compile/test_fusions_e2e.py
+190 -38
@@ -20,13 +20,22 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
20
21 from ..utils import flat_product, multi_gpu_test
22
23 +is_blackwell = lambda: current_platform.is_device_capability(100)
24 +"""Are we running on Blackwell, a lot of tests depend on it"""
25 +
26 +
27 +class Matches(NamedTuple):
28 + attention_fusion: int = 0
29 + allreduce_fusion: int = 0
30 + sequence_parallel: int = 0
31 + async_tp: int = 0
32 +
33
34 class ModelBackendTestCase(NamedTuple):
35 model_name: str
36 model_kwargs: dict[str, Any]
37 backend: AttentionBackendEnum
28 - attention_fusions: int
29 - allreduce_fusions: int | None = None
38 + matches: Matches
39
40
41 MODELS_FP8: list[ModelBackendTestCase] = []
@@ -38,17 +47,33 @@ if current_platform.is_cuda():
47 ModelBackendTestCase(
48 # Use smaller model for L40s in CI
49 model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
41 - model_kwargs=dict(max_model_len=1024),
42 - backend=AttentionBackendEnum.TRITON_ATTN,
43 - attention_fusions=32,
44 - allreduce_fusions=65,
50 + # TODO while llama4 is broken, use FLASHINFER for llama3 on Blackwell
51 + # so FI attention+fp8_quant is at least tested once
52 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
53 + backend=AttentionBackendEnum.FLASHINFER
54 + if is_blackwell()
55 + else AttentionBackendEnum.TRITON_ATTN,
56 + matches=Matches(
57 + attention_fusion=32,
58 + allreduce_fusion=65,
59 + sequence_parallel=65,
60 + async_tp=128,
61 + ),
62 ),
63 ModelBackendTestCase(
64 model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
65 model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
49 - backend=AttentionBackendEnum.FLASHINFER,
50 - attention_fusions=48,
51 - allreduce_fusions=96,
66 + # TODO FlashInfer attn broken on Hopper with kvcache=fp8:
67 + # https://github.com/vllm-project/vllm/issues/28568
68 + # TODO FlashInfer attn broken on Blackwell for llama4:
69 + # https://github.com/vllm-project/vllm/issues/28604
70 + backend=AttentionBackendEnum.TRITON_ATTN,
71 + matches=Matches(
72 + attention_fusion=48,
73 + allreduce_fusion=96,
74 + sequence_parallel=96,
75 + async_tp=95, # mlp is moe, no fusion there
76 + ),
77 ),
78 ]
79
@@ -57,8 +82,12 @@ if current_platform.is_cuda():
82 model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
83 model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
84 backend=AttentionBackendEnum.FLASHINFER,
60 - attention_fusions=32,
61 - allreduce_fusions=65,
85 + matches=Matches(
86 + attention_fusion=32,
87 + allreduce_fusion=65,
88 + sequence_parallel=65,
89 + async_tp=128,
90 + ),
91 ),
92 ]
93
@@ -68,15 +97,23 @@ if current_platform.is_cuda():
97 model_name="meta-llama/Llama-3.1-8B-Instruct",
98 model_kwargs=dict(max_model_len=1024),
99 backend=AttentionBackendEnum.TRITON_ATTN,
71 - attention_fusions=0,
72 - allreduce_fusions=65,
100 + matches=Matches(
101 + attention_fusion=0,
102 + allreduce_fusion=65,
103 + sequence_parallel=65,
104 + async_tp=128,
105 + ),
106 ),
107 ModelBackendTestCase(
108 model_name="Qwen/Qwen3-30B-A3B",
109 model_kwargs=dict(max_model_len=1024),
110 backend=AttentionBackendEnum.TRITON_ATTN,
78 - attention_fusions=0,
79 - allreduce_fusions=97,
111 + matches=Matches(
112 + attention_fusion=0,
113 + allreduce_fusion=97,
114 + sequence_parallel=97,
115 + async_tp=96, # MLP is MoE, half the fusions of dense
116 + ),
117 ),
118 ]
119
@@ -86,19 +123,19 @@ elif current_platform.is_rocm():
123 model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
124 model_kwargs=dict(max_model_len=1024),
125 backend=AttentionBackendEnum.TRITON_ATTN,
89 - attention_fusions=32,
126 + matches=Matches(attention_fusion=32),
127 ),
128 ModelBackendTestCase(
129 model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
130 model_kwargs=dict(max_model_len=1024),
131 backend=AttentionBackendEnum.ROCM_ATTN,
95 - attention_fusions=32,
132 + matches=Matches(attention_fusion=32),
133 ),
134 ModelBackendTestCase(
135 model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
136 model_kwargs=dict(max_model_len=1024),
137 backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
101 - attention_fusions=32,
138 + matches=Matches(attention_fusion=32),
139 ),
140 ]
141
@@ -106,8 +143,7 @@ CUSTOM_OPS_FP8 = ["-quant_fp8", "+quant_fp8"]
143
144
145 @pytest.mark.parametrize(
109 - "model_name, model_kwargs, backend, "
110 - "attention_fusions, allreduce_fusions, custom_ops",
146 + "model_name, model_kwargs, backend, matches, custom_ops",
147 # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
148 list(flat_product(MODELS_FP8, CUSTOM_OPS_FP8))
149 # quant_fp4 only has the custom impl
@@ -118,15 +154,14 @@ def test_attn_quant(
154 model_name: str,
155 model_kwargs: dict[str, Any],
156 backend: AttentionBackendEnum,
121 - attention_fusions: int,
122 - allreduce_fusions: int,
157 + matches: Matches,
158 custom_ops: str,
159 inductor_graph_partition: bool,
160 caplog_mp_spawn,
161 monkeypatch,
162 ):
163 if backend == AttentionBackendEnum.FLASHINFER and (
129 - not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
164 + not is_blackwell() or not has_flashinfer()
165 ):
166 pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
167 if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
@@ -169,12 +204,12 @@ def test_attn_quant(
204 with caplog_mp_spawn(logging.DEBUG) as log_holder:
205 run_model(compilation_config, model_name, **model_kwargs)
206
172 - matches = re.findall(
207 + log_matches = re.findall(
208 r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
209 log_holder.text,
210 )
176 - assert len(matches) == 1, log_holder.text
177 - assert int(matches[0]) == attention_fusions
211 + assert len(log_matches) == 1, log_holder.text
212 + assert int(log_matches[0]) == matches.attention_fusion
213
214
215 CUSTOM_OPS_RMS_NORM = ["-rms_norm", "+rms_norm"]
@@ -187,8 +222,7 @@ def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]:
222
223 @multi_gpu_test(num_gpus=2)
224 @pytest.mark.parametrize(
190 - "model_name, model_kwargs, backend, "
191 - "attention_fusions, allreduce_fusions, custom_ops",
225 + "model_name, model_kwargs, backend, matches, custom_ops",
226 # Toggle RMSNorm and QuantFP8 for FP8 models
227 list(
228 flat_product(
@@ -209,8 +243,7 @@ def test_tp2_attn_quant_allreduce_rmsnorm(
243 model_name: str,
244 model_kwargs: dict,
245 backend: AttentionBackendEnum,
212 - attention_fusions: int,
213 - allreduce_fusions: int,
246 + matches: Matches,
247 custom_ops: str,
248 inductor_graph_partition: bool,
249 caplog_mp_spawn,
@@ -219,6 +252,13 @@ def test_tp2_attn_quant_allreduce_rmsnorm(
252 if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
253 pytest.skip("Inductor graph partition requires torch>=2.9")
254
255 + if "fp4" in model_name.lower() and not is_blackwell():
256 + pytest.skip("NVFP4 quant requires Blackwell")
257 +
258 + if backend == AttentionBackendEnum.FLASHINFER and not is_blackwell():
259 + # FlashInfer attn fusion requires Blackwell
260 + matches = matches._replace(attention_fusion=0)
261 +
262 custom_ops_list = custom_ops.split(",") if custom_ops else []
263
264 if inductor_graph_partition:
@@ -258,23 +298,135 @@ def test_tp2_attn_quant_allreduce_rmsnorm(
298 run_model(
299 compilation_config, model_name, tensor_parallel_size=2, **model_kwargs
300 )
261 - matches = re.findall(
301 + log_matches = re.findall(
302 + r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
303 + log_holder.text,
304 + )
305 + assert len(log_matches) == 2, log_holder.text
306 +
307 + assert int(log_matches[0]) == matches.attention_fusion
308 + assert int(log_matches[1]) == matches.attention_fusion
309 +
310 + log_matches = re.findall(
311 + r"collective_fusion.py:\d+] Replaced (\d+) patterns",
312 + log_holder.text,
313 + )
314 + assert len(log_matches) == 2, log_holder.text
315 +
316 + assert int(log_matches[0]) == matches.allreduce_fusion
317 + assert int(log_matches[1]) == matches.allreduce_fusion
318 +
319 +
320 +@multi_gpu_test(num_gpus=2)
321 +@pytest.mark.parametrize(
322 + "model_name, model_kwargs, backend, matches, custom_ops",
323 + # Toggle RMSNorm and QuantFP8 for FP8 models
324 + list(
325 + flat_product(
326 + MODELS_FP8, custom_ops_product(CUSTOM_OPS_FP8, CUSTOM_OPS_RMS_NORM)
327 + )
328 + )
329 + # Toggle RMSNorm for FP4 models and unquant models
330 + + list(flat_product(MODELS_FP4 + MODELS, CUSTOM_OPS_RMS_NORM)),
331 +)
332 +@pytest.mark.parametrize("inductor_graph_partition", [True, False])
333 +@pytest.mark.skipif(
334 + not current_platform.is_cuda(),
335 + reason="sequence parallel only tested on CUDA",
336 +)
337 +def test_tp2_attn_quant_async_tp(
338 + model_name: str,
339 + model_kwargs: dict,
340 + backend: AttentionBackendEnum,
341 + matches: Matches,
342 + custom_ops: str,
343 + inductor_graph_partition: bool,
344 + caplog_mp_spawn,
345 + monkeypatch,
346 +):
347 + if is_blackwell():
348 + # TODO: https://github.com/vllm-project/vllm/issues/27893
349 + pytest.skip("Blackwell is not supported for AsyncTP pass")
350 +
351 + if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
352 + pytest.skip("Inductor graph partition requires torch>=2.9")
353 +
354 + if "fp4" in model_name.lower() and not is_blackwell():
355 + pytest.skip("NVFP4 quant requires Blackwell")
356 +
357 + if backend == AttentionBackendEnum.FLASHINFER:
358 + if not has_flashinfer():
359 + pytest.skip("FlashInfer backend requires flashinfer installed")
360 + if not is_blackwell():
361 + # FlashInfer attn fusion requires Blackwell
362 + matches = matches._replace(attention_fusion=0)
363 +
364 + custom_ops_list = custom_ops.split(",") if custom_ops else []
365 +
366 + if inductor_graph_partition:
367 + mode = CUDAGraphMode.FULL_AND_PIECEWISE
368 + splitting_ops: list[str] | None = None
369 + else:
370 + mode = CUDAGraphMode.FULL_DECODE_ONLY
371 + splitting_ops = []
372 +
373 + # Disable, compile cache to make sure custom passes run.
374 + # Otherwise, we can't verify fusion happened through the logs.
375 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
376 +
377 + # To capture subprocess logs, we need to know whether spawn or fork is used.
378 + # Force spawn as it is more general.
379 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
380 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)
381 +
382 + compilation_config = CompilationConfig(
383 + # Testing properties
384 + use_inductor_graph_partition=inductor_graph_partition,
385 + cudagraph_mode=mode,
386 + custom_ops=custom_ops_list,
387 + splitting_ops=splitting_ops,
388 + # Common
389 + level=CompilationMode.VLLM_COMPILE,
390 + pass_config=PassConfig(
391 + enable_attn_fusion=True,
392 + enable_noop=True,
393 + enable_sequence_parallelism=True,
394 + enable_async_tp=True,
395 + ),
396 + # Inductor caches custom passes by default as well via uuid
397 + inductor_compile_config={"force_disable_caches": True},
398 + )
399 +
400 + with caplog_mp_spawn(logging.DEBUG) as log_holder:
401 + run_model(
402 + compilation_config, model_name, tensor_parallel_size=2, **model_kwargs
403 + )
404 + log_matches = re.findall(
405 r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
406 log_holder.text,
407 )
265 - assert len(matches) == 2, log_holder.text
408 + assert len(log_matches) == 2, log_holder.text
409 +
410 + assert int(log_matches[0]) == matches.attention_fusion
411 + assert int(log_matches[1]) == matches.attention_fusion
412 +
413 + log_matches = re.findall(
414 + r"sequence_parallelism.py:\d+] Replaced (\d+) patterns",
415 + log_holder.text,
416 + )
417 + assert len(log_matches) == 2, log_holder.text
418
267 - assert int(matches[0]) == attention_fusions
268 - assert int(matches[1]) == attention_fusions
419 + assert int(log_matches[0]) == matches.sequence_parallel
420 + assert int(log_matches[1]) == matches.sequence_parallel
421
270 - matches = re.findall(
422 + log_matches = re.findall(
423 r"collective_fusion.py:\d+] Replaced (\d+) patterns",
424 log_holder.text,
425 )
274 - assert len(matches) == 2, log_holder.text
426 + assert len(log_matches) == 2, log_holder.text
427
276 - assert int(matches[0]) == allreduce_fusions
277 - assert int(matches[1]) == allreduce_fusions
428 + assert int(log_matches[0]) == matches.async_tp
429 + assert int(log_matches[1]) == matches.async_tp
430
431
432 def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):
tests/compile/test_sequence_parallelism.py
+123 -139
@@ -5,15 +5,15 @@ import pytest
5 import torch
6
7 import vllm.envs as envs
8 -from vllm.compilation.fix_functionalization import FixFunctionalizationPass
8 from vllm.compilation.fusion import RMSNormQuantFusionPass
10 -from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
9 +from vllm.compilation.fx_utils import find_auto_fn
10 from vllm.compilation.noop_elimination import NoOpEliminationPass
11 from vllm.compilation.post_cleanup import PostCleanupPass
12 from vllm.compilation.sequence_parallelism import SequenceParallelismPass
13 from vllm.compilation.vllm_inductor_pass import VllmInductorPass
14 from vllm.config import (
15 CompilationConfig,
16 + CUDAGraphMode,
17 DeviceConfig,
18 ModelConfig,
19 PassConfig,
@@ -27,6 +27,7 @@ from vllm.distributed.parallel_state import (
27 initialize_model_parallel,
28 )
29 from vllm.model_executor.layers.layernorm import RMSNorm
30 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
31 from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
32 from vllm.platforms import current_platform
33 from vllm.utils.system_utils import update_environment_variables
@@ -43,172 +44,157 @@ prompts = [
44 ]
45
46
46 -class TestModel(torch.nn.Module):
47 - def __init__(self, hidden_size=16, intermediate_size=32):
47 +class TestAllReduceRMSNormModel(torch.nn.Module):
48 + def __init__(self, hidden_size=16, eps=1e-6):
49 super().__init__()
50 self.hidden_size = hidden_size
50 - self.intermediate_size = intermediate_size
51 - self.gate_proj = torch.nn.Parameter(
52 - torch.empty((intermediate_size, hidden_size))
53 - )
54 - self.norm = RMSNorm(intermediate_size, 1e-05)
55 - # Initialize weights
56 - torch.nn.init.normal_(self.gate_proj, std=0.02)
51 + self.eps = eps
52 + self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]
53 + self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)]
54
58 - def forward(self, hidden_states, residual):
59 - """
60 - Forward pass implementing the operations in the FX graph
55 + def forward(self, x):
56 + z = torch.relu(x)
57 + x = resid = tensor_model_parallel_all_reduce(z)
58 + y = self.norm[0](x)
59
62 - Args:
63 - hidden_states: Input tensor
64 - residual: Residual tensor from previous layer
60 + z2 = torch.mm(y, self.w[0])
61 + x2 = tensor_model_parallel_all_reduce(z2)
62
66 - Returns:
67 - Tuple containing the output tensor
68 - """
69 - # Reshape input
70 - view = hidden_states.reshape(-1, self.hidden_size)
63 + y2, resid = self.norm[1](x2, resid)
64
72 - # matrix multiplication
73 - permute = self.gate_proj.permute(1, 0)
74 - mm = torch.mm(view, permute)
65 + z3 = torch.mm(y2, self.w[1])
66 + x3 = tensor_model_parallel_all_reduce(z3)
67
76 - # Tensor parallel all-reduce
77 - all_reduce = tensor_model_parallel_all_reduce(mm)
68 + y3, resid = self.norm[2](x3, resid)
69
79 - # layer normalization
80 - norm_output, residual_output = self.norm(all_reduce, residual)
70 + z4 = torch.mm(y3, self.w[2])
71 + x4 = tensor_model_parallel_all_reduce(z4)
72
82 - return norm_output, residual_output
73 + y4, resid = self.norm[3](x4, resid)
74 + return y4
75
76 def ops_in_model_before(self):
77 return [torch.ops.vllm.all_reduce.default]
78
79 def ops_in_model_after(self):
80 return [
89 - torch.ops.vllm.reduce_scatter.default,
81 torch.ops.vllm.all_gather.default,
82 + torch.ops.vllm.reduce_scatter.default,
83 ]
84
85 def ops_in_model(self):
94 - return [torch.ops._C.fused_add_rms_norm.default]
86 + if RMSNorm.enabled():
87 + return [
88 + torch.ops._C.rms_norm.default,
89 + torch.ops._C.fused_add_rms_norm.default,
90 + ]
91 + else:
92 + return []
93
94
97 -class TestQuantModel(torch.nn.Module):
98 - def __init__(self, hidden_size=16, intermediate_size=32):
95 +class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
96 + def __init__(self, hidden_size=16, eps=1e-6):
97 super().__init__()
100 - self.hidden_size = hidden_size
101 - self.intermediate_size = intermediate_size
98 self.vllm_config = get_current_vllm_config()
103 - self.gate_proj = torch.nn.Parameter(
104 - torch.empty((intermediate_size, hidden_size)), requires_grad=False
99 + self.hidden_size = hidden_size
100 + self.eps = eps
101 + self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]
102 + self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
103 + self.w = [
104 + torch.rand(hidden_size, hidden_size)
105 + .to(dtype=current_platform.fp8_dtype())
106 + .t()
107 + for _ in range(3)
108 + ]
109 +
110 + self.fp8_linear = Fp8LinearOp(
111 + act_quant_static=True,
112 + act_quant_group_shape=GroupShape.PER_TENSOR,
113 )
106 - self.norm = RMSNorm(intermediate_size, 1e-05)
107 - # Initialize weights
108 - torch.nn.init.normal_(self.gate_proj, std=0.02)
109 -
110 - self.fp8_linear = Fp8LinearOp(act_quant_static=True)
111 -
112 - self.scale = torch.rand(1, dtype=torch.float32)
113 - # Create a weight that is compatible with torch._scaled_mm,
114 - # which expects a column-major layout.
115 - self.w = torch.rand(hidden_size, intermediate_size).to(dtype=FP8_DTYPE).t()
116 - self.wscale = torch.rand(1, dtype=torch.float32)
117 -
118 - def forward(self, hidden_states, residual):
119 - """
120 - Forward pass implementing the operations in the FX graph
121 -
122 - Args:
123 - hidden_states: Input tensor
124 - residual: Residual tensor from previous layer
125 -
126 - Returns:
127 - Tuple containing the output tensor
128 - """
129 - # Reshape input
130 - view = hidden_states.reshape(-1, self.hidden_size)
131 -
132 - # matrix multiplication
133 - permute = self.gate_proj.permute(1, 0)
134 - mm = torch.mm(view, permute)
135 -
136 - # Tensor parallel all-reduce
137 - all_reduce = tensor_model_parallel_all_reduce(mm)
138 -
139 - # layer normalization
140 - norm_output, residual_output = self.norm(all_reduce, residual)
141 -
142 - # scaled_mm with static input quantization
143 - fp8_linear_result = self.fp8_linear.apply(
144 - norm_output,
145 - self.w,
146 - self.wscale,
147 - input_scale=self.scale.to(norm_output.device),
114 +
115 + self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
116 +
117 + def forward(self, hidden_states):
118 + # avoid having graph input be an arg to a pattern directly
119 + z = torch.relu(hidden_states)
120 + x = resid = tensor_model_parallel_all_reduce(z)
121 + y = self.norm[0](x)
122 +
123 + z2 = self.fp8_linear.apply(
124 + y, self.w[0], self.wscale[0], input_scale=self.scale[0]
125 )
126
150 - return fp8_linear_result, residual_output
127 + x2 = tensor_model_parallel_all_reduce(z2)
128 + y2, resid = self.norm[1](x2, resid)
129
152 - def ops_in_model_before(self):
153 - ops_to_remove = [torch.ops.vllm.all_reduce.default] # Always removed by SP
154 - # The following are only removed if fusion happens
155 - if (
156 - self.vllm_config
157 - and self.vllm_config.compilation_config.pass_config.enable_fusion
158 - ):
159 - ops_to_remove.extend(
160 - [
161 - torch.ops._C.fused_add_rms_norm.default,
162 - torch.ops._C.static_scaled_fp8_quant.default,
163 - ]
164 - )
165 - return ops_to_remove
130 + z3 = self.fp8_linear.apply(
131 + y2, self.w[1], self.wscale[1], input_scale=self.scale[1]
132 + )
133 +
134 + x3 = tensor_model_parallel_all_reduce(z3)
135 + y3, resid = self.norm[2](x3, resid) # use resid here
136 +
137 + z4 = self.fp8_linear.apply(
138 + y3, self.w[2], self.wscale[2], input_scale=self.scale[2]
139 + )
140 + x4 = tensor_model_parallel_all_reduce(z4)
141 + y4, resid = self.norm[3](x4, resid) # use resid here
142 + return y4
143
144 def ops_in_model_after(self):
168 - ops_to_add = [
169 - torch.ops.vllm.reduce_scatter.default,
145 + return [
146 torch.ops.vllm.all_gather.default,
147 + torch.ops.vllm.reduce_scatter.default,
148 + ]
149 +
150 + def ops_in_model_before(self):
151 + return [
152 + torch.ops.vllm.all_reduce.default,
153 ]
172 - # The following is only added if fusion happens
173 - if (
174 - self.vllm_config
175 - and self.vllm_config.compilation_config.pass_config.enable_fusion
176 - ):
177 - ops_to_add.append(torch.ops._C.fused_add_rms_norm_static_fp8_quant.default)
178 - return ops_to_add
154
155 def ops_in_model(self):
181 - if (
182 - self.vllm_config
183 - and self.vllm_config.compilation_config.pass_config.enable_fusion
184 - ):
185 - # If fusion happens, the fused op is the one
186 - # we check for (de)functionalization
156 + if self.vllm_config.compilation_config.pass_config.enable_fusion:
157 return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
188 - else:
189 - # If no fusion, the original ops are checked
158 + elif RMSNorm.enabled():
159 return [
160 torch.ops._C.fused_add_rms_norm.default,
192 - # TODO functionalization pass does not handle this yet
193 - # torch.ops._C.static_scaled_fp8_quant.default,
161 ]
162 + elif self.fp8_linear.quant_fp8.enabled():
163 + return [
164 + torch.ops._C.static_scaled_fp8_quant.default,
165 + ]
166 + else:
167 + return []
168
169
170 @multi_gpu_test(num_gpus=2)
198 -@pytest.mark.parametrize("test_model_cls", [TestModel, TestQuantModel])
171 +@pytest.mark.parametrize(
172 + "test_model_cls, custom_ops",
173 + [
174 + (TestAllReduceRMSNormModel, "+rms_norm"),
175 + (TestAllReduceRMSNormModel, "-rms_norm"),
176 + (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,+quant_fp8"),
177 + (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,-quant_fp8"),
178 + (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,+quant_fp8"),
179 + (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,-quant_fp8"),
180 + ],
181 +)
182 @pytest.mark.parametrize("batch_size", [8])
183 @pytest.mark.parametrize("seq_len", [16])
184 @pytest.mark.parametrize("hidden_size", [16])
185 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
186 @pytest.mark.parametrize("enable_fusion", [True, False])
187 +@pytest.mark.parametrize("dynamic", [False, True])
188 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
189 def test_sequence_parallelism_pass(
190 test_model_cls: type[torch.nn.Module],
191 + custom_ops: str,
192 batch_size: int,
193 seq_len: int,
194 hidden_size: int,
195 dtype: torch.dtype,
196 enable_fusion: bool,
197 + dynamic: bool,
198 ):
199 num_processes = 2
200
@@ -220,11 +206,13 @@ def test_sequence_parallelism_pass(
206 args=(
207 num_processes,
208 test_model_cls,
209 + custom_ops,
210 batch_size,
211 seq_len,
212 hidden_size,
213 dtype,
214 enable_fusion,
215 + dynamic,
216 ),
217 nprocs=nprocs,
218 )
@@ -236,11 +224,13 @@ def sequence_parallelism_pass_on_test_model(
224 local_rank: int,
225 world_size: int,
226 test_model_cls: type[torch.nn.Module],
227 + custom_ops: str,
228 batch_size: int,
229 seq_len: int,
230 hidden_size: int,
231 dtype: torch.dtype,
232 enable_fusion: bool,
233 + dynamic: bool,
234 ):
235 current_platform.seed_everything(0)
236
@@ -264,12 +254,16 @@ def sequence_parallelism_pass_on_test_model(
254 initialize_model_parallel(tensor_model_parallel_size=world_size)
255
256 # configure vllm config for SequenceParallelismPass
257 + custom_ops_list = custom_ops.split(",") if custom_ops else []
258 compilation_config = CompilationConfig(
259 + splitting_ops=[], # avoid automatic rms_norm enablement
260 + cudagraph_mode=CUDAGraphMode.NONE, # avoid piecewise warnings
261 + custom_ops=custom_ops_list,
262 pass_config=PassConfig(
263 enable_sequence_parallelism=True,
264 enable_fusion=enable_fusion,
265 enable_noop=True,
272 - )
266 + ),
267 ) # NoOp needed for fusion
268 device_config = DeviceConfig(device=torch.device("cuda"))
269
@@ -289,7 +283,6 @@ def sequence_parallelism_pass_on_test_model(
283 with set_current_vllm_config(vllm_config):
284 noop_pass = NoOpEliminationPass(vllm_config)
285 sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
292 - func_pass = FixFunctionalizationPass(vllm_config)
286 cleanup_pass = PostCleanupPass(vllm_config)
287 assert (
288 sequence_parallelism_pass.compilation_config.splitting_ops
@@ -310,38 +303,29 @@ def sequence_parallelism_pass_on_test_model(
303
304 passes_for_backend.append(cleanup_pass)
305
313 - backend_no_func = TestBackend(*passes_for_backend)
314 - backend_func = TestBackend(*passes_for_backend, func_pass)
306 + backend = TestBackend(*passes_for_backend)
307
316 - model = test_model_cls(hidden_size, hidden_size * 2)
308 + model = test_model_cls(hidden_size)
309
310 hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
319 - residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
311
321 - compiled_model_no_func = torch.compile(model, backend=backend_no_func)
322 - compiled_model_no_func(hidden_states, residual)
323 - compiled_model_func = torch.compile(model, backend=backend_func)
324 - compiled_model_func(hidden_states, residual)
312 + if dynamic:
313 + torch._dynamo.mark_dynamic(hidden_states, 0)
314 +
315 + compiled_model = torch.compile(model, backend=backend)
316 + compiled_model(hidden_states)
317
326 - assert sequence_parallelism_pass.matched_count == 1
318 + assert sequence_parallelism_pass.matched_count == 4
319
320 # In pre-nodes, all reduce should be there,
321 # reduce scatter and all gather should not
330 - backend_no_func.check_before_ops(model.ops_in_model_before())
322 + for op in model.ops_in_model_before():
323 + assert backend.op_count(op, before=True) == 4
324
325 # In post-nodes, reduce scatter and all gather should be there,
326 # all reduce should not
334 - backend_no_func.check_after_ops(model.ops_in_model_after())
327 + for op in model.ops_in_model_after():
328 + assert backend.op_count(op, before=False) == 4
329
336 - # check if the functionalization pass is applied
330 for op in model.ops_in_model():
338 - find_auto_fn(backend_no_func.graph_post_pass.nodes, op)
339 - assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None
340 -
341 - # make sure the ops were all de-functionalized
342 - found = dict()
343 - for node in backend_func.graph_post_pass.nodes:
344 - for op in model.ops_in_model():
345 - if is_func(node, op):
346 - found[op] = True
347 - assert all(found[op] for op in model.ops_in_model())
331 + find_auto_fn(backend.graph_post_pass.nodes, op)
tests/distributed/test_sequence_parallel.py
+14 -1
@@ -18,6 +18,7 @@ import pytest
18 from vllm.config.compilation import CompilationMode
19 from vllm.config.model import RunnerOption
20 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
@@ -161,6 +162,7 @@ def _compare_sp(
162 test_options: SPTestOptions,
163 num_gpus_available: int,
164 use_inductor_graph_partition: bool,
165 + enable_async_tp: bool,
166 *,
167 method: Literal["generate", "encode"],
168 is_multimodal: bool,
@@ -244,10 +246,10 @@ def _compare_sp(
246
247 compilation_config = {
248 "mode": CompilationMode.VLLM_COMPILE,
247 - "custom_ops": ["+rms_norm"],
249 "compile_sizes": [4, 8],
250 "pass_config": {
251 "enable_sequence_parallelism": True,
252 + "enable_async_tp": enable_async_tp,
253 "enable_fusion": enable_fusion,
254 "enable_noop": True,
255 },
@@ -307,6 +309,7 @@ SP_TEST_MODELS = [
309 ],
310 )
311 @pytest.mark.parametrize("use_inductor_graph_partition", [True, False])
312 +@pytest.mark.parametrize("enable_async_tp", [False]) # TODO: enable async TP
313 @create_new_process_for_each_test()
314 def test_tp_sp_generation(
315 model_id: str,
@@ -316,10 +319,19 @@ def test_tp_sp_generation(
319 test_options: SPTestOptions,
320 num_gpus_available,
321 use_inductor_graph_partition: bool,
322 + enable_async_tp: bool,
323 ):
324 if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
325 pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
326
327 + # Skip FP8 SP-only test on sm89 (compute capability 8.9)
328 + if (
329 + "fp8" in model_id.lower()
330 + and current_platform.get_device_capability() < (9, 0)
331 + and (not enable_async_tp)
332 + ):
333 + pytest.skip("FP8 reduction support begins with sm90 capable devices.")
334 +
335 _compare_sp(
336 model_id,
337 parallel_setup,
@@ -328,6 +340,7 @@ def test_tp_sp_generation(
340 test_options,
341 num_gpus_available,
342 use_inductor_graph_partition,
343 + enable_async_tp=enable_async_tp,
344 method="generate",
345 is_multimodal=False,
346 )
vllm/compilation/sequence_parallelism.py
+111 -258
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import functools
5 +
6 import torch
7 import torch._inductor.pattern_matcher as pm
8 import torch.fx as fx
@@ -10,98 +12,28 @@ from vllm.config import VllmConfig
12 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
13 from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
14 from vllm.logger import init_logger
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + kFp8StaticTensorSym,
17 +)
18 from vllm.platforms import current_platform
19
20 from .inductor_pass import enable_fake_mode
21 +from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
22 +from .noop_elimination import NoOpEliminationPass
23 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
24
25 logger = init_logger(__name__)
26
27
21 -class _RMSNormAndQuantOpHelper:
22 - """Base helper for RMSNorm and RMSNorm + Quantization functionalization."""
28 +def get_first_out_wrapper(fn):
29 + @functools.wraps(fn)
30 + def wrapper(*args):
31 + return fn(*args)[0]
32
24 - def __init__(
25 - self,
26 - epsilon: float,
27 - dtype: torch.dtype,
28 - device: str,
29 - quant_op: torch._ops.OpOverload | None = None,
30 - **kwargs,
31 - ):
32 - self.epsilon = epsilon
33 - self.dtype = dtype
34 - self.device = device
35 - self.quant_op = quant_op
36 -
37 - def _functional_rmsnorm(self, result_buffer, input_tensor, weight_tensor):
38 - return torch.ops.higher_order.auto_functionalized(
39 - torch.ops._C.rms_norm.default,
40 - result=result_buffer,
41 - input=input_tensor,
42 - weight=weight_tensor,
43 - epsilon=self.epsilon,
44 - )
33 + return wrapper
34
46 - def _functional_fused_add_rmsnorm(
47 - self, input_tensor, residual_tensor, weight_tensor
48 - ):
49 - return torch.ops.higher_order.auto_functionalized(
50 - torch.ops._C.fused_add_rms_norm.default,
51 - input=input_tensor,
52 - residual=residual_tensor,
53 - weight=weight_tensor,
54 - epsilon=self.epsilon,
55 - )
35
57 - def _functional_rmsnorm_then_quant(
58 - self,
59 - rmsnorm_result_buffer,
60 - quant_result_buffer,
61 - input_tensor,
62 - weight_tensor,
63 - scale_tensor,
64 - ):
65 - if self.quant_op is None:
66 - raise RuntimeError(
67 - "_RMSNormAndQuantOpHelper was not initialized with a quant_op."
68 - )
69 - rmsnorm_out_tuple = self._functional_rmsnorm(
70 - rmsnorm_result_buffer, input_tensor, weight_tensor
71 - )
72 - quant_out_tuple = torch.ops.higher_order.auto_functionalized(
73 - self.quant_op,
74 - result=quant_result_buffer,
75 - input=rmsnorm_out_tuple[1],
76 - scale=scale_tensor,
77 - )
78 - return quant_out_tuple
79 -
80 - def _functional_fused_add_rmsnorm_then_quant(
81 - self,
82 - quant_result_buffer,
83 - input_tensor,
84 - residual_tensor,
85 - weight_tensor,
86 - scale_tensor,
87 - ):
88 - if self.quant_op is None:
89 - raise RuntimeError(
90 - "_RMSNormAndQuantOpHelper was not initialized with a quant_op."
91 - )
92 - fused_add_rmsnorm_out_tuple = self._functional_fused_add_rmsnorm(
93 - input_tensor, residual_tensor, weight_tensor
94 - )
95 - quant_out_tuple = torch.ops.higher_order.auto_functionalized(
96 - self.quant_op,
97 - result=quant_result_buffer,
98 - input=fused_add_rmsnorm_out_tuple[1],
99 - scale=scale_tensor,
100 - )
101 - return quant_out_tuple, fused_add_rmsnorm_out_tuple[2]
102 -
103 -
104 -class _SequenceParallelPatternHelper(_RMSNormAndQuantOpHelper):
36 +class _SequenceParallelPatternHelper:
37 """Helper for sequence parallelism patterns."""
38
39 def __init__(
@@ -109,10 +41,10 @@ class _SequenceParallelPatternHelper(_RMSNormAndQuantOpHelper):
41 epsilon: float,
42 dtype: torch.dtype,
43 device: str,
112 - quant_op: torch._ops.OpOverload | None = None,
113 - **kwargs,
44 ):
115 - super().__init__(epsilon, dtype, device, quant_op=quant_op, **kwargs)
45 + self.epsilon = epsilon
46 + self.dtype = dtype
47 + self.device = device
48 self.tp_group = get_tp_group()
49 self.tp_size = get_tensor_model_parallel_world_size()
50
@@ -131,36 +63,34 @@ class _SequenceParallelPatternHelper(_RMSNormAndQuantOpHelper):
63
64
65 class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
66 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str):
67 + super().__init__(epsilon, dtype, device)
68 + self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
69 +
70 def get_inputs(self):
71 input = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
136 - permute = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
72 arg3_1 = torch.empty([4], device=self.device, dtype=self.dtype)
73
139 - return [input, permute, arg3_1]
74 + return [input, arg3_1]
75
76 def register(self, pm_pass: PatternMatcherPass):
77 def pattern(
78 input: torch.Tensor,
144 - permute: torch.Tensor,
79 arg3_1: torch.Tensor,
80 ):
81 all_reduce = self._all_reduce(input)
148 - rmsnorm = self._functional_rmsnorm(permute, all_reduce, arg3_1)
82 + rmsnorm = self.rmsnorm_matcher(all_reduce, arg3_1)
83
150 - return rmsnorm[1], all_reduce
84 + return rmsnorm, all_reduce
85
86 def replacement(
87 input: torch.Tensor,
154 - permute: torch.Tensor,
88 arg3_1: torch.Tensor,
89 ):
90 reduce_scatter = self._reduce_scatter(input)
91
159 - rmsnorm_result = torch.empty_like(reduce_scatter)
160 - rmsnorm = self._functional_rmsnorm(rmsnorm_result, reduce_scatter, arg3_1)
161 -
162 - all_gather = self._all_gather(rmsnorm[1])
163 -
92 + rmsnorm = self.rmsnorm_matcher(reduce_scatter, arg3_1)
93 + all_gather = self._all_gather(rmsnorm)
94 return all_gather, reduce_scatter
95
96 pm.register_replacement(
@@ -169,6 +99,10 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
99
100
101 class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
102 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str):
103 + super().__init__(epsilon, dtype, device)
104 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
105 +
106 def get_inputs(self):
107 mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
108
@@ -188,67 +122,34 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
122 rms_norm_weights: torch.Tensor,
123 ) -> tuple[torch.Tensor, torch.Tensor]:
124 all_reduce = self._all_reduce(mm_1)
191 - rmsnorm = self._functional_fused_add_rmsnorm(
192 - all_reduce, residual, rms_norm_weights
193 - )
194 - return rmsnorm[1], rmsnorm[2]
125 + rmsnorm = self.rmsnorm_matcher(all_reduce, rms_norm_weights, residual)
126 + return rmsnorm[0], rmsnorm[1]
127
128 def replacement(
129 residual: torch.Tensor,
130 mm_1: torch.Tensor,
131 rms_norm_weights: torch.Tensor,
132 ) -> tuple[torch.Tensor, torch.Tensor]:
133 + # pattern matcher replaces from top-to-bottom,
134 + # so residual is still the full size here.
135 + # once the seqpar pattern with the previous rmsnorm is replaced
136 reduce_scatter = self._reduce_scatter(mm_1)
202 - rmsnorm = self._functional_fused_add_rmsnorm(
203 - reduce_scatter, residual, rms_norm_weights
204 - )
205 - all_gather = self._all_gather(rmsnorm[1])
206 - return all_gather, rmsnorm[2]
137 + residual = residual[0 : reduce_scatter.size(0), ...]
138 + rmsnorm = self.rmsnorm_matcher(reduce_scatter, rms_norm_weights, residual)
139 + all_gather = self._all_gather(rmsnorm[0])
140 + # shape of residual changes but that's fine,
141 + # next node is already slicing it, now becomes a noop
142 + return all_gather, rmsnorm[1]
143
144 pm.register_replacement(
145 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
146 )
211 -
212 -
213 -class LastAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
214 - def get_inputs(self):
215 - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
216 -
217 - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
218 - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype)
219 -
220 - return [
221 - residual,
222 - mm_1,
223 - rms_norm_weights,
224 - ]
225 -
226 - def register(self, pm_pass: PatternMatcherPass):
227 - def pattern(
228 - residual: torch.Tensor,
229 - mm_1: torch.Tensor,
230 - rms_norm_weights: torch.Tensor,
231 - ) -> tuple[torch.Tensor, torch.Tensor]:
232 - all_reduce = self._all_reduce(mm_1)
233 - rmsnorm = self._functional_fused_add_rmsnorm(
234 - all_reduce, residual, rms_norm_weights
235 - )
236 - return rmsnorm[1]
237 -
238 - def replacement(
239 - residual: torch.Tensor,
240 - mm_1: torch.Tensor,
241 - rms_norm_weights: torch.Tensor,
242 - ) -> tuple[torch.Tensor, torch.Tensor]:
243 - reduce_scatter = self._reduce_scatter(mm_1)
244 - rmsnorm = self._functional_fused_add_rmsnorm(
245 - reduce_scatter, residual, rms_norm_weights
246 - )
247 - normalized = self._all_gather(rmsnorm[1])
248 - return normalized
249 -
147 pm.register_replacement(
251 - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
148 + get_first_out_wrapper(pattern),
149 + get_first_out_wrapper(replacement),
150 + self.get_inputs(),
151 + pm.fwd_only,
152 + pm_pass,
153 )
154
155
@@ -257,52 +158,41 @@ FP8_DTYPE = current_platform.fp8_dtype()
158
159 class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
160 def __init__(
260 - self, epsilon: float, dtype: torch.dtype, device: str, op: torch._ops.OpOverload
161 + self,
162 + epsilon: float,
163 + dtype: torch.dtype,
164 + device: str,
165 ):
262 - super().__init__(epsilon, dtype, device, quant_op=op)
166 + super().__init__(epsilon, dtype, device)
167 + self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
168 + self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
169
170 def get_inputs(self):
171 input = torch.zeros([1, 8, 4], device=self.device, dtype=self.dtype)
266 - rmsnorm_result = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
267 - quant_result = torch.empty([1, 8, 4], device=self.device, dtype=FP8_DTYPE)
172 weight = torch.empty([4], device=self.device, dtype=self.dtype)
173 scale = torch.tensor(1.0, device=self.device, dtype=torch.float32)
270 - return [input, rmsnorm_result, quant_result, weight, scale]
174 + return [input, weight, scale]
175
176 def register(self, pm_pass: PatternMatcherPass):
177 def pattern(
178 input: torch.Tensor,
275 - rmsnorm_result: torch.Tensor,
276 - quant_result: torch.Tensor,
179 weight: torch.Tensor,
180 scale: torch.Tensor,
181 ):
182 all_reduce = self._all_reduce(input)
281 - static_fp8 = self._functional_rmsnorm_then_quant(
282 - rmsnorm_result, quant_result, all_reduce, weight, scale
283 - )
284 - return static_fp8[1], all_reduce
183 + rms = self.rmsnorm_matcher(all_reduce, weight)
184 + quant, _ = self.quant_matcher(rms, scale)
185 + return quant, all_reduce
186
187 def replacement(
188 input: torch.Tensor,
288 - rmsnorm_result: torch.Tensor,
289 - quant_result: torch.Tensor,
189 weight: torch.Tensor,
190 scale: torch.Tensor,
191 ):
192 reduce_scatter = self._reduce_scatter(input)
294 -
295 - rmsnorm_result = torch.empty_like(
296 - reduce_scatter, dtype=rmsnorm_result.dtype
297 - )
298 - quant_result = torch.empty_like(
299 - rmsnorm_result, # Output of RMSNorm
300 - dtype=quant_result.dtype,
301 - )
302 - static_fp8 = self._functional_rmsnorm_then_quant(
303 - rmsnorm_result, quant_result, reduce_scatter, weight, scale
304 - )
305 - all_gather = self._all_gather(static_fp8[1])
193 + rms = self.rmsnorm_matcher(reduce_scatter, weight)
194 + quant, _ = self.quant_matcher(rms, scale)
195 + all_gather = self._all_gather(quant)
196
197 return all_gather, reduce_scatter
198
@@ -312,118 +202,64 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
202
203
204 class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
315 - def __init__(
316 - self, epsilon: float, dtype: torch.dtype, device: str, op: torch._ops.OpOverload
317 - ):
318 - super().__init__(epsilon, dtype, device, quant_op=op)
205 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str):
206 + super().__init__(epsilon, dtype, device)
207 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
208 + self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
209
210 def get_inputs(self):
211 mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
322 -
212 residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
213 rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype)
325 - result = torch.empty([4, 4], device=self.device, dtype=FP8_DTYPE)
214 scale = torch.empty([1, 1], device=self.device, dtype=torch.float32)
215
328 - return [
329 - result,
330 - residual,
331 - mm_1,
332 - rms_norm_weights,
333 - scale,
334 - ]
216 + return [residual, mm_1, rms_norm_weights, scale]
217
218 def register(self, pm_pass: PatternMatcherPass):
219 def pattern(
338 - result: torch.Tensor,
220 residual: torch.Tensor,
221 mm_1: torch.Tensor,
222 rms_norm_weights: torch.Tensor,
223 scale: torch.Tensor,
224 ) -> tuple[torch.Tensor, torch.Tensor]:
225 all_reduce = self._all_reduce(mm_1)
345 - static_fp8, rmsnorm_residual_out = (
346 - self._functional_fused_add_rmsnorm_then_quant( # noqa: E501
347 - result, all_reduce, residual, rms_norm_weights, scale
348 - )
226 + rms, residual_out = self.rmsnorm_matcher(
227 + all_reduce, rms_norm_weights, residual
228 )
350 - return static_fp8[1], rmsnorm_residual_out
229 + quant, _ = self.quant_matcher(rms, scale)
230 + return quant, residual_out
231
232 def replacement(
353 - result: torch.Tensor,
233 residual: torch.Tensor,
234 mm_1: torch.Tensor,
235 rms_norm_weights: torch.Tensor,
236 scale: torch.Tensor,
237 ) -> tuple[torch.Tensor, torch.Tensor]:
238 + # pattern matcher replaces from top-to-bottom,
239 + # so residual is still the full size here.
240 + # add a temporary slice which will become a noop
241 + # once the seqpar pattern with the previous rmsnorm is replaced
242 reduce_scatter = self._reduce_scatter(mm_1)
360 - quant_result_buf = torch.empty_like(reduce_scatter, dtype=result.dtype)
361 - static_fp8, rmsnorm_residual_out = (
362 - self._functional_fused_add_rmsnorm_then_quant( # noqa: E501
363 - quant_result_buf, reduce_scatter, residual, rms_norm_weights, scale
364 - )
243 + residual = residual[0 : reduce_scatter.size(0), ...]
244 + rms, residual_out = self.rmsnorm_matcher(
245 + reduce_scatter, rms_norm_weights, residual
246 )
366 - all_gather = self._all_gather(static_fp8[1])
367 - return all_gather, rmsnorm_residual_out
247 + quant, _ = self.quant_matcher(rms, scale)
248 + all_gather = self._all_gather(quant)
249 + # shape of residual changes but that's fine,
250 + # next node is already slicing it, now becomes a noop
251 + return all_gather, residual_out
252
253 pm.register_replacement(
254 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
255 )
256
373 -
374 -class LastAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
375 - def __init__(
376 - self, epsilon: float, dtype: torch.dtype, device: str, op: torch._ops.OpOverload
377 - ):
378 - super().__init__(epsilon, dtype, device, quant_op=op)
379 -
380 - def get_inputs(self):
381 - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
382 -
383 - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
384 - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype)
385 - result = torch.empty([4, 4], device=self.device, dtype=FP8_DTYPE)
386 - scale = torch.empty([1, 1], device=self.device, dtype=torch.float32)
387 -
388 - return [
389 - result,
390 - residual,
391 - mm_1,
392 - rms_norm_weights,
393 - scale,
394 - ]
395 -
396 - def register(self, pm_pass: PatternMatcherPass):
397 - def pattern(
398 - result: torch.Tensor,
399 - residual: torch.Tensor,
400 - mm_1: torch.Tensor,
401 - rms_norm_weights: torch.Tensor,
402 - scale: torch.Tensor,
403 - ) -> tuple[torch.Tensor, torch.Tensor]:
404 - all_reduce = self._all_reduce(mm_1)
405 - static_fp8, _ = self._functional_fused_add_rmsnorm_then_quant(
406 - result, all_reduce, residual, rms_norm_weights, scale
407 - )
408 - return static_fp8[1]
409 -
410 - def replacement(
411 - result: torch.Tensor,
412 - residual: torch.Tensor,
413 - mm_1: torch.Tensor,
414 - rms_norm_weights: torch.Tensor,
415 - scale: torch.Tensor,
416 - ) -> tuple[torch.Tensor, torch.Tensor]:
417 - reduce_scatter = self._reduce_scatter(mm_1)
418 - quant_result_buf = torch.empty_like(reduce_scatter, dtype=result.dtype)
419 - static_fp8, _ = self._functional_fused_add_rmsnorm_then_quant(
420 - quant_result_buf, reduce_scatter, residual, rms_norm_weights, scale
421 - )
422 - normalized = self._all_gather(static_fp8[1])
423 - return normalized
424 -
257 pm.register_replacement(
426 - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
258 + get_first_out_wrapper(pattern),
259 + get_first_out_wrapper(replacement),
260 + self.get_inputs(),
261 + pm.fwd_only,
262 + pm_pass,
263 )
264
265
@@ -445,27 +281,45 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
281 GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can
282 significantly reduce communication overhead and improve overall model
283 performance.
284 +
285 +
286 + This pass splits up the residual tensor across TP ranks and hence divides its size.
287 + Because the pattern matcher starts at the end of the graph, the replacement
288 + contains a slice that temporarily conforms the input residual to the correct size.
289 + After all patterns have been matched, we use a NoOpEliminationPass to clean up
290 + what have now become no-op slices.
291 +
292 + Note that an older version of the pass did not need this as it operated only on
293 + custom rms_norm and fused_rms_norm_add custom ops which did not complain about
294 + mismatched shapes during replacement. So this approach has the same assumption that
295 + correctness is only maintained if all rms_norm operations are split across ranks.
296 +
297 + Correctness-wise, this is approach strictly better than before - before,
298 + the graph was incorrect semantically and shape-wise during the pass.
299 + With this approach there's only semantic incorrectness during the pass.
300 + Both approaches restore a correct graph once all patterns are matched.
301 """
302
303 @enable_fake_mode
304 def __init__(self, config: VllmConfig):
305 super().__init__(config)
306
307 + # Used to cleanup redundant views created temporarily
308 + # to circumvent residual shape change issues
309 + self.noop_cleanup = NoOpEliminationPass(config)
310 + self.noop_cleanup.pass_name = f"{self.pass_name}.{self.noop_cleanup.pass_name}"
311 +
312 self.patterns: PatternMatcherPass = PatternMatcherPass(
313 pass_name="sequence_parallelism_pass"
314 )
315
316 for epsilon in [1e-5, 1e-6]:
317 # RMSNorm + Static FP8 quantization patterns
460 - fp8_quant_op = torch.ops._C.static_scaled_fp8_quant.default
318 FirstAllReduceRMSNormStaticFP8Pattern(
462 - epsilon, self.model_dtype, self.device, fp8_quant_op
319 + epsilon, self.model_dtype, self.device
320 ).register(self.patterns)
321 MiddleAllReduceRMSNormStaticFP8Pattern(
465 - epsilon, self.model_dtype, self.device, fp8_quant_op
466 - ).register(self.patterns)
467 - LastAllReduceRMSNormStaticFP8Pattern(
468 - epsilon, self.model_dtype, self.device, fp8_quant_op
322 + epsilon, self.model_dtype, self.device
323 ).register(self.patterns)
324
325 # Normal RMSNorm patterns
@@ -477,9 +331,6 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
331 epsilon, self.model_dtype, self.device
332 ).register(self.patterns)
333
480 - LastAllReduceRMSNormPattern(
481 - epsilon, self.model_dtype, self.device
482 - ).register(self.patterns)
334 self.dump_patterns(config, self.patterns)
335
336 def is_applicable(self, shape: int | None) -> bool:
@@ -508,3 +359,5 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
359 def __call__(self, graph: fx.Graph):
360 self.matched_count = self.patterns.apply(graph)
361 logger.debug("Replaced %s patterns", self.matched_count)
362 + # Clean up reshape nodes
363 + self.noop_cleanup(graph)
vllm/config/vllm.py
+26 -2
@@ -445,8 +445,6 @@ class VllmConfig:
445 # and requires it to be enabled.
446 if self.compilation_config.pass_config.enable_async_tp:
447 self.compilation_config.pass_config.enable_sequence_parallelism = True
448 - if self.compilation_config.pass_config.enable_sequence_parallelism:
449 - self.compilation_config.custom_ops.append("+rms_norm")
448
449 if current_platform.support_static_graph_mode():
450 # if cudagraph_mode is not explicitly set by users, set default
@@ -620,6 +618,32 @@ class VllmConfig:
618 if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:
619 self.compilation_config.set_splitting_ops_for_v1()
620
621 + if self.compilation_config.pass_config.enable_sequence_parallelism:
622 + # With pipeline parallelism or dynamo partitioning,
623 + # native rms norm tracing errors due to incorrect residual shape.
624 + # Use custom rms norm to unblock. In the future,
625 + # the pass will operate on higher-level IR to avoid the issue.
626 + # TODO: https://github.com/vllm-project/vllm/issues/27894
627 + is_fullgraph = (
628 + self.compilation_config.use_inductor_graph_partition
629 + or len(self.compilation_config.splitting_ops) == 0
630 + )
631 + if self.parallel_config.pipeline_parallel_size > 1 or not is_fullgraph:
632 + if "-rms_norm" not in self.compilation_config.custom_ops:
633 + self.compilation_config.custom_ops.append("+rms_norm")
634 + else:
635 + regime = (
636 + "Dynamo partition"
637 + if not is_fullgraph
638 + else "pipeline parallelism"
639 + )
640 + logger.warning_once(
641 + "Sequence parallelism not supported with"
642 + "native rms_norm when using %s, "
643 + "this will likely lead to an error.",
644 + regime,
645 + )
646 +
647 # final check of cudagraph mode after all possible updates
648 if current_platform.is_cuda_alike():
649 if (