@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 14ce524249
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 10, 01:30 UTC

[CI] Breakup h200 tests (#30499)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Jan 15, 2026 at 23:23 UTC 14ce524249df0ccadab549ad0fa6a737bf0adb61
5 files changed +390 -296
.buildkite/test-pipeline.yaml
+53 -4
@@ -1047,6 +1047,47 @@ steps:
1047 # Run all e2e fusion tests
1048 - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
1049
1050 +- label: Hopper Fusion E2E Tests (H100) # 10min
1051 + timeout_in_minutes: 70
1052 + working_dir: "/vllm-workspace/"
1053 + gpu: h100
1054 + optional: true
1055 + source_file_dependencies:
1056 + - csrc/quantization/fp4/
1057 + - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1058 + - vllm/v1/attention/backends/flashinfer.py
1059 + - vllm/compilation/
1060 + # can affect pattern matching
1061 + - vllm/model_executor/layers/layernorm.py
1062 + - vllm/model_executor/layers/activation.py
1063 + - vllm/model_executor/layers/quantization/input_quant_fp8.py
1064 + - tests/compile/test_fusion_attn.py
1065 + commands:
1066 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1
1067 + - pytest -v -s tests/compile/test_fusion_attn.py
1068 +
1069 +- label: Hopper Fusion Distributed E2E Tests (2xH100) # 70min
1070 + timeout_in_minutes: 70
1071 + working_dir: "/vllm-workspace/"
1072 + gpu: h100
1073 + optional: true
1074 + num_gpus: 2
1075 + source_file_dependencies:
1076 + - csrc/quantization/fp4/
1077 + - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1078 + - vllm/v1/attention/backends/flashinfer.py
1079 + - vllm/compilation/
1080 + # can affect pattern matching
1081 + - vllm/model_executor/layers/layernorm.py
1082 + - vllm/model_executor/layers/activation.py
1083 + - vllm/model_executor/layers/quantization/input_quant_fp8.py
1084 + - tests/compile/distributed/test_fusions_e2e.py
1085 + commands:
1086 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1
1087 + # Run all e2e fusion tests
1088 + - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
1089 + - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1090 +
1091 - label: Blackwell GPT-OSS Eval
1092 timeout_in_minutes: 60
1093 working_dir: "/vllm-workspace/"
@@ -1346,6 +1387,18 @@ steps:
1387 - export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
1388 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
1389
1390 +- label: Sequence Parallel Tests (H100) # 60 min
1391 + timeout_in_minutes: 60
1392 + working_dir: "/vllm-workspace/"
1393 + gpu: h100
1394 + optional: true
1395 + num_gpus: 2
1396 + commands:
1397 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1
1398 + # Run sequence parallel tests
1399 + - pytest -v -s tests/distributed/test_sequence_parallel.py
1400 + - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
1401 +
1402 ##### H200 test #####
1403 - label: Distributed Tests (H200) # optional
1404 gpu: h200
@@ -1354,10 +1407,6 @@ steps:
1407 num_gpus: 2
1408 commands:
1409 - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
1357 - - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
1358 - - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1359 - - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1360 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
1410 - pytest -v -s tests/distributed/test_context_parallel.py
1411 - CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
1412 - pytest -v -s tests/v1/distributed/test_dbo.py
tests/compile/distributed/test_fusions_e2e.py
+15 -283
@@ -3,16 +3,26 @@
3
4 from __future__ import annotations
5
6 -import itertools
6 import logging
8 -from collections.abc import Iterable
9 -from typing import Any, NamedTuple
7 +from typing import Any
8
9 import pytest
10 import regex as re
11
12 +from tests.compile.fusion_test_utils import (
13 + CUSTOM_OPS_FP8,
14 + CUSTOM_OPS_QUANT_RMS_NORM,
15 + CUSTOM_OPS_RMS_NORM,
16 + MODELS,
17 + MODELS_FP4,
18 + MODELS_FP8,
19 + MODELS_GROUP_FP8,
20 + Matches,
21 + custom_ops_product,
22 + is_blackwell,
23 + run_model,
24 +)
25 from tests.v1.attention.utils import AttentionBackendEnum
15 -from vllm import LLM, SamplingParams
26 from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
27 from vllm.platforms import current_platform
28 from vllm.utils.flashinfer import has_flashinfer
@@ -20,228 +30,6 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
30
31 from ...utils import flat_product, multi_gpu_test
32
23 -is_blackwell = lambda: current_platform.is_device_capability_family(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 - rms_quant_norm_fusion: int = 0
31 - sequence_parallel: int = 0
32 - async_tp: int = 0
33 -
34 -
35 -class ModelBackendTestCase(NamedTuple):
36 - model_name: str
37 - model_kwargs: dict[str, Any]
38 - backend: AttentionBackendEnum
39 - matches: Matches
40 -
41 -
42 -MODELS_FP8: list[ModelBackendTestCase] = []
43 -MODELS_FP4: list[ModelBackendTestCase] = []
44 -MODELS_GROUP_FP8: list[ModelBackendTestCase] = []
45 -MODELS: list[ModelBackendTestCase] = [] # tp-only
46 -
47 -if current_platform.is_cuda():
48 - MODELS_FP8 = [
49 - ModelBackendTestCase(
50 - # Use smaller model for L40s in CI
51 - model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
52 - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
53 - backend=AttentionBackendEnum.TRITON_ATTN,
54 - matches=Matches(
55 - attention_fusion=32,
56 - allreduce_fusion=65,
57 - sequence_parallel=65,
58 - async_tp=128,
59 - ),
60 - ),
61 - ModelBackendTestCase(
62 - model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
63 - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
64 - # TODO FlashInfer attn broken on Hopper with kvcache=fp8:
65 - # https://github.com/vllm-project/vllm/issues/28568
66 - backend=AttentionBackendEnum.FLASHINFER
67 - if is_blackwell()
68 - else AttentionBackendEnum.TRITON_ATTN,
69 - matches=Matches(
70 - attention_fusion=48,
71 - allreduce_fusion=96,
72 - sequence_parallel=96,
73 - async_tp=95, # mlp is moe, no fusion there
74 - ),
75 - ),
76 - ]
77 -
78 - MODELS_FP4 = [
79 - ModelBackendTestCase(
80 - model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
81 - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
82 - backend=AttentionBackendEnum.FLASHINFER,
83 - matches=Matches(
84 - attention_fusion=32,
85 - allreduce_fusion=65,
86 - sequence_parallel=65,
87 - async_tp=128,
88 - ),
89 - ),
90 - ]
91 -
92 - # TP only
93 - MODELS = [
94 - ModelBackendTestCase(
95 - model_name="meta-llama/Llama-3.1-8B-Instruct",
96 - model_kwargs=dict(max_model_len=1024),
97 - backend=AttentionBackendEnum.TRITON_ATTN,
98 - matches=Matches(
99 - attention_fusion=0,
100 - allreduce_fusion=65,
101 - sequence_parallel=65,
102 - async_tp=128,
103 - ),
104 - ),
105 - ModelBackendTestCase(
106 - model_name="Qwen/Qwen3-30B-A3B",
107 - model_kwargs=dict(max_model_len=1024),
108 - backend=AttentionBackendEnum.TRITON_ATTN,
109 - matches=Matches(
110 - attention_fusion=0,
111 - allreduce_fusion=97,
112 - sequence_parallel=97,
113 - async_tp=96, # MLP is MoE, half the fusions of dense
114 - ),
115 - ),
116 - ]
117 -
118 -elif current_platform.is_rocm():
119 - MODELS_FP8 = [
120 - ModelBackendTestCase(
121 - model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
122 - model_kwargs=dict(max_model_len=1024),
123 - backend=AttentionBackendEnum.TRITON_ATTN,
124 - matches=Matches(attention_fusion=32),
125 - ),
126 - ModelBackendTestCase(
127 - model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
128 - model_kwargs=dict(max_model_len=1024),
129 - backend=AttentionBackendEnum.ROCM_ATTN,
130 - matches=Matches(attention_fusion=32),
131 - ),
132 - ModelBackendTestCase(
133 - model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
134 - model_kwargs=dict(max_model_len=1024),
135 - backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
136 - matches=Matches(attention_fusion=32),
137 - ),
138 - ]
139 -
140 -CUSTOM_OPS_FP8 = ["-quant_fp8", "+quant_fp8"]
141 -
142 -
143 -def has_cuda_graph_wrapper_metadata() -> bool:
144 - from importlib import import_module
145 -
146 - try:
147 - module = import_module("torch._inductor.utils")
148 - module.CUDAGraphWrapperMetadata # noqa B018
149 - except AttributeError:
150 - return False
151 - return True
152 -
153 -
154 -@pytest.mark.parametrize(
155 - "model_name, model_kwargs, backend, matches, custom_ops",
156 - # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
157 - list(flat_product(MODELS_FP8, CUSTOM_OPS_FP8))
158 - # quant_fp4 only has the custom impl
159 - + list(flat_product(MODELS_FP4, [""])),
160 -)
161 -@pytest.mark.parametrize(
162 - "inductor_graph_partition",
163 - [
164 - pytest.param(
165 - True,
166 - marks=pytest.mark.skipif(
167 - not has_cuda_graph_wrapper_metadata(),
168 - reason="This test requires"
169 - "torch._inductor.utils.CUDAGraphWrapperMetadata to run",
170 - ),
171 - ),
172 - False,
173 - ],
174 -)
175 -def test_attn_quant(
176 - model_name: str,
177 - model_kwargs: dict[str, Any],
178 - backend: AttentionBackendEnum,
179 - matches: Matches,
180 - custom_ops: str,
181 - inductor_graph_partition: bool,
182 - caplog_mp_spawn,
183 - monkeypatch,
184 -):
185 - if backend == AttentionBackendEnum.FLASHINFER and (
186 - not is_blackwell() or not has_flashinfer()
187 - ):
188 - pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
189 - if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
190 - pytest.skip("Inductor graph partition requires torch>=2.9")
191 -
192 - custom_ops_list = custom_ops.split(",") if custom_ops else []
193 -
194 - if inductor_graph_partition:
195 - mode = CUDAGraphMode.FULL_AND_PIECEWISE
196 - splitting_ops: list[str] | None = None
197 - else:
198 - # FIXME: Llama-4-Scout-17B-16E-Instruct-FP8 + FlashInfer + Blackwell end at
199 - # CUDAGraphMode.NONE here because it derives an attention backend that
200 - # does not support full cudagraphs
201 - mode = CUDAGraphMode.FULL_DECODE_ONLY
202 - splitting_ops = []
203 -
204 - # Disable, compile cache to make sure custom passes run.
205 - # Otherwise, we can't verify fusion happened through the logs.
206 - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
207 -
208 - # To capture subprocess logs, we need to know whether spawn or fork is used.
209 - # Force spawn as it is more general.
210 - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
211 -
212 - model_kwargs["attention_config"] = {"backend": backend.name}
213 -
214 - compilation_config = CompilationConfig(
215 - # Testing properties
216 - custom_ops=custom_ops_list,
217 - use_inductor_graph_partition=inductor_graph_partition,
218 - cudagraph_mode=mode,
219 - splitting_ops=splitting_ops,
220 - # Common
221 - mode=CompilationMode.VLLM_COMPILE,
222 - pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True),
223 - # Inductor caches custom passes by default as well via uuid
224 - inductor_compile_config={"force_disable_caches": True},
225 - )
226 -
227 - with caplog_mp_spawn(logging.DEBUG) as log_holder:
228 - run_model(compilation_config, model_name, **model_kwargs)
229 -
230 - log_matches = re.findall(
231 - r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
232 - log_holder.text,
233 - )
234 - assert len(log_matches) == 1, log_holder.text
235 - assert int(log_matches[0]) == matches.attention_fusion
236 -
237 -
238 -CUSTOM_OPS_RMS_NORM = ["-rms_norm", "+rms_norm"]
239 -
240 -
241 -def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]:
242 - for op_list in itertools.product(*custom_ops_lists):
243 - yield ",".join(op_list)
244 -
33
34 @multi_gpu_test(num_gpus=2)
35 @pytest.mark.parametrize(
@@ -421,7 +209,7 @@ def test_tp2_attn_quant_async_tp(
209 custom_ops=custom_ops_list,
210 splitting_ops=splitting_ops,
211 # Common
424 - level=CompilationMode.VLLM_COMPILE,
212 + mode=CompilationMode.VLLM_COMPILE,
213 pass_config=PassConfig(
214 fuse_attn_quant=True,
215 eliminate_noops=True,
@@ -464,62 +252,6 @@ def test_tp2_attn_quant_async_tp(
252 assert int(log_matches[1]) == matches.async_tp
253
254
467 -def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):
468 - compilation_config = (
469 - compile_config
470 - if isinstance(compile_config, CompilationConfig)
471 - else CompilationConfig(mode=compile_config)
472 - )
473 -
474 - prompts = [
475 - "Hello, my name is",
476 - "The president of the United States is",
477 - "The capital of France is",
478 - "The future of AI is",
479 - ]
480 - sampling_params = SamplingParams(temperature=0)
481 - # Allow override from model_kwargs
482 - model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}
483 - model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}
484 -
485 - # No cudagraphs by default
486 - if compilation_config.cudagraph_mode is None:
487 - compilation_config.cudagraph_mode = CUDAGraphMode.NONE
488 - llm = LLM(
489 - model=model,
490 - compilation_config=compilation_config,
491 - **model_kwargs,
492 - )
493 - outputs = llm.generate(prompts, sampling_params)
494 -
495 - # Print the outputs.
496 - for output in outputs:
497 - prompt = output.prompt
498 - generated_text = output.outputs[0].text
499 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
500 -
501 - # Get the compile ranges split points after vllm config post init
502 - # in order to compute compile ranges correctly
503 - compilation_config.compile_ranges_split_points = (
504 - llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points
505 - )
506 -
507 -
508 -if current_platform.is_cuda():
509 - MODELS_GROUP_FP8 = [
510 - ModelBackendTestCase(
511 - model_name="Qwen/Qwen3-30B-A3B-FP8",
512 - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
513 - backend=AttentionBackendEnum.TRITON_ATTN,
514 - matches=Matches(
515 - rms_quant_norm_fusion=48,
516 - ),
517 - ),
518 - ]
519 -
520 -CUSTOM_OPS_QUANT_RMS_NORM = ["+quant_fp8,+rms_norm"]
521 -
522 -
255 @pytest.mark.parametrize(
256 "model_name, model_kwargs, backend, matches, custom_ops",
257 # Test rms norm+group quant_fp8 fusion
tests/compile/fusion_test_utils.py new
+208
@@ -0,0 +1,208 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared utilities for fusion tests (e.g. test_fusion_attn.py)."""
4 +
5 +from __future__ import annotations
6 +
7 +import itertools
8 +from collections.abc import Iterable
9 +from typing import Any, NamedTuple
10 +
11 +from tests.v1.attention.utils import AttentionBackendEnum
12 +from vllm import LLM, SamplingParams
13 +from vllm.config import CompilationConfig, CUDAGraphMode
14 +from vllm.platforms import current_platform
15 +
16 +is_blackwell = lambda: current_platform.is_device_capability_family(100)
17 +"""Are we running on Blackwell, a lot of tests depend on it"""
18 +
19 +
20 +def has_cuda_graph_wrapper_metadata() -> bool:
21 + from importlib import import_module
22 +
23 + try:
24 + module = import_module("torch._inductor.utils")
25 + module.CUDAGraphWrapperMetadata # noqa B018
26 + except AttributeError:
27 + return False
28 + return True
29 +
30 +
31 +class Matches(NamedTuple):
32 + attention_fusion: int = 0
33 + allreduce_fusion: int = 0
34 + sequence_parallel: int = 0
35 + async_tp: int = 0
36 + rms_quant_norm_fusion: int = 0
37 +
38 +
39 +class ModelBackendTestCase(NamedTuple):
40 + model_name: str
41 + model_kwargs: dict[str, Any]
42 + backend: AttentionBackendEnum
43 + matches: Matches
44 +
45 +
46 +# E2E model test cases
47 +MODELS_FP8: list[ModelBackendTestCase] = []
48 +MODELS_FP4: list[ModelBackendTestCase] = []
49 +MODELS: list[ModelBackendTestCase] = [] # tp-only (unquantized)
50 +MODELS_GROUP_FP8: list[ModelBackendTestCase] = []
51 +
52 +if current_platform.is_cuda():
53 + MODELS_FP8 = [
54 + ModelBackendTestCase(
55 + # Use smaller model for L40s in CI
56 + model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
57 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
58 + backend=AttentionBackendEnum.TRITON_ATTN,
59 + matches=Matches(
60 + attention_fusion=32,
61 + allreduce_fusion=65,
62 + sequence_parallel=65,
63 + async_tp=128,
64 + ),
65 + ),
66 + ModelBackendTestCase(
67 + model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
68 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
69 + # TODO FlashInfer attn broken on Hopper with kvcache=fp8:
70 + # https://github.com/vllm-project/vllm/issues/28568
71 + backend=AttentionBackendEnum.FLASHINFER
72 + if is_blackwell()
73 + else AttentionBackendEnum.TRITON_ATTN,
74 + matches=Matches(
75 + attention_fusion=48,
76 + allreduce_fusion=96,
77 + sequence_parallel=96,
78 + async_tp=95, # mlp is moe, no fusion there
79 + ),
80 + ),
81 + ]
82 +
83 + MODELS_FP4 = [
84 + ModelBackendTestCase(
85 + model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
86 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
87 + backend=AttentionBackendEnum.FLASHINFER,
88 + matches=Matches(
89 + attention_fusion=32,
90 + allreduce_fusion=65,
91 + sequence_parallel=65,
92 + async_tp=128,
93 + ),
94 + ),
95 + ]
96 +
97 + # TP only (unquantized models)
98 + MODELS = [
99 + ModelBackendTestCase(
100 + model_name="meta-llama/Llama-3.1-8B-Instruct",
101 + model_kwargs=dict(max_model_len=1024),
102 + backend=AttentionBackendEnum.TRITON_ATTN,
103 + matches=Matches(
104 + attention_fusion=0,
105 + allreduce_fusion=65,
106 + sequence_parallel=65,
107 + async_tp=128,
108 + ),
109 + ),
110 + ModelBackendTestCase(
111 + model_name="Qwen/Qwen3-30B-A3B",
112 + model_kwargs=dict(max_model_len=1024),
113 + backend=AttentionBackendEnum.TRITON_ATTN,
114 + matches=Matches(
115 + attention_fusion=0,
116 + allreduce_fusion=97,
117 + sequence_parallel=97,
118 + async_tp=96, # MLP is MoE, half the fusions of dense
119 + ),
120 + ),
121 + ]
122 +
123 + MODELS_GROUP_FP8 = [
124 + ModelBackendTestCase(
125 + model_name="Qwen/Qwen3-30B-A3B-FP8",
126 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
127 + backend=AttentionBackendEnum.TRITON_ATTN,
128 + matches=Matches(
129 + rms_quant_norm_fusion=48,
130 + ),
131 + ),
132 + ]
133 +
134 +elif current_platform.is_rocm():
135 + MODELS_FP8 = [
136 + ModelBackendTestCase(
137 + model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
138 + model_kwargs=dict(max_model_len=1024),
139 + backend=AttentionBackendEnum.TRITON_ATTN,
140 + matches=Matches(attention_fusion=32),
141 + ),
142 + ModelBackendTestCase(
143 + model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
144 + model_kwargs=dict(max_model_len=1024),
145 + backend=AttentionBackendEnum.ROCM_ATTN,
146 + matches=Matches(attention_fusion=32),
147 + ),
148 + ModelBackendTestCase(
149 + model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
150 + model_kwargs=dict(max_model_len=1024),
151 + backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
152 + matches=Matches(attention_fusion=32),
153 + ),
154 + ]
155 +
156 +
157 +# Custom ops toggle lists for parametrization
158 +CUSTOM_OPS_FP8 = ["-quant_fp8", "+quant_fp8"]
159 +CUSTOM_OPS_RMS_NORM = ["-rms_norm", "+rms_norm"]
160 +CUSTOM_OPS_QUANT_RMS_NORM = ["+quant_fp8,+rms_norm"]
161 +
162 +
163 +def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]:
164 + """Generate all combinations of custom ops for parametrization."""
165 + for op_list in itertools.product(*custom_ops_lists):
166 + yield ",".join(op_list)
167 +
168 +
169 +def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):
170 + """Run a model with the given compilation config for E2E fusion tests."""
171 + compilation_config = (
172 + compile_config
173 + if isinstance(compile_config, CompilationConfig)
174 + else CompilationConfig(mode=compile_config)
175 + )
176 +
177 + prompts = [
178 + "Hello, my name is",
179 + "The president of the United States is",
180 + "The capital of France is",
181 + "The future of AI is",
182 + ]
183 + sampling_params = SamplingParams(temperature=0)
184 + # Allow override from model_kwargs
185 + model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}
186 + model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}
187 +
188 + # No cudagraphs by default
189 + if compilation_config.cudagraph_mode is None:
190 + compilation_config.cudagraph_mode = CUDAGraphMode.NONE
191 + llm = LLM(
192 + model=model,
193 + compilation_config=compilation_config,
194 + **model_kwargs,
195 + )
196 + outputs = llm.generate(prompts, sampling_params)
197 +
198 + # Print the outputs.
199 + for output in outputs:
200 + prompt = output.prompt
201 + generated_text = output.outputs[0].text
202 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
203 +
204 + # Get the compile ranges split points after vllm config post init
205 + # in order to compute compile ranges correctly
206 + compilation_config.compile_ranges_split_points = (
207 + llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points
208 + )
tests/compile/test_fusion_attn.py
+113 -8
@@ -1,12 +1,24 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import copy
4 +import logging
5 +from typing import Any
6
7 import pytest
8 +import regex as re
9 import torch._dynamo
10
11 from tests.compile.backend import LazyInitPass, TestBackend
9 -from tests.utils import flat_product
12 +from tests.compile.fusion_test_utils import (
13 + CUSTOM_OPS_FP8,
14 + MODELS_FP4,
15 + MODELS_FP8,
16 + Matches,
17 + has_cuda_graph_wrapper_metadata,
18 + is_blackwell,
19 + run_model,
20 +)
21 +from tests.utils import cuda_device_count_stateless, flat_product
22 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
23 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
24 from vllm.attention.layer import Attention
@@ -20,6 +32,7 @@ from vllm.config import (
32 CacheConfig,
33 CompilationConfig,
34 CompilationMode,
35 + CUDAGraphMode,
36 ModelConfig,
37 PassConfig,
38 SchedulerConfig,
@@ -35,6 +48,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
48 from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
49 from vllm.platforms import current_platform
50 from vllm.utils.flashinfer import has_flashinfer
51 +from vllm.utils.torch_utils import is_torch_equal_or_newer
52 from vllm.v1.attention.backend import AttentionMetadata
53 from vllm.v1.attention.backends.registry import AttentionBackendEnum
54 from vllm.v1.kv_cache_interface import AttentionSpec
@@ -241,8 +255,8 @@ class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
255 )
256
257
244 -MODELS_FP8: list[tuple[str, type]] = []
245 -MODELS_FP4: list[tuple[str, type]] = []
258 +PATTERN_TEST_MODELS_FP8: list[tuple[str, type]] = []
259 +PATTERN_TEST_MODELS_FP4: list[tuple[str, type]] = []
260 HEADS: list[tuple[int, int]] = []
261 SPLIT_ATTENTION: list[bool] = []
262 BACKENDS_FP8: list[AttentionBackendEnum] = []
@@ -250,13 +264,13 @@ BACKENDS_FP4: list[AttentionBackendEnum] = []
264
265 if current_platform.is_cuda():
266 HEADS = [(64, 8), (40, 8)]
253 - MODELS_FP8 = [
267 + PATTERN_TEST_MODELS_FP8 = [
268 (
269 "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
270 TestAttentionFp8StaticQuantPatternModel,
271 )
272 ]
259 - MODELS_FP4 = [
273 + PATTERN_TEST_MODELS_FP4 = [
274 (
275 "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
276 TestAttentionNvfp4QuantPatternModel,
@@ -267,7 +281,7 @@ if current_platform.is_cuda():
281
282 elif current_platform.is_rocm():
283 HEADS = [(32, 8), (40, 8)]
270 - MODELS_FP8 = [
284 + PATTERN_TEST_MODELS_FP8 = [
285 ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)
286 ]
287 BACKENDS = [
@@ -286,9 +300,13 @@ elif current_platform.is_rocm():
300 @pytest.mark.parametrize(
301 "backend, model_name, model_class, custom_ops",
302 # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
289 - list(flat_product(BACKENDS_FP8, MODELS_FP8, ["+quant_fp8", "-quant_fp8"]))
303 + list(
304 + flat_product(
305 + BACKENDS_FP8, PATTERN_TEST_MODELS_FP8, ["+quant_fp8", "-quant_fp8"]
306 + )
307 + )
308 # quant_fp4 only has the custom impl
291 - + list(flat_product(BACKENDS_FP4, MODELS_FP4, [""])),
309 + + list(flat_product(BACKENDS_FP4, PATTERN_TEST_MODELS_FP4, [""])),
310 )
311 @pytest.mark.skipif(
312 not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"
@@ -315,6 +333,8 @@ def test_attention_quant_pattern(
333 not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
334 ):
335 pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
336 + if "Llama-4-Scout" in model_name and cuda_device_count_stateless() < 2:
337 + pytest.skip("Llama-4-Scout requires at least 2 GPUs")
338
339 custom_ops_list = custom_ops.split(",") if custom_ops else []
340
@@ -483,3 +503,88 @@ def test_attention_quant_pattern(
503
504 # Check that results are close
505 torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)
506 +
507 +
508 +@pytest.mark.parametrize(
509 + "model_name, model_kwargs, backend, matches, custom_ops",
510 + # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
511 + list(flat_product(MODELS_FP8, CUSTOM_OPS_FP8))
512 + # quant_fp4 only has the custom impl
513 + + list(flat_product(MODELS_FP4, [""])),
514 +)
515 +@pytest.mark.parametrize(
516 + "inductor_graph_partition",
517 + [
518 + pytest.param(
519 + True,
520 + marks=pytest.mark.skipif(
521 + not has_cuda_graph_wrapper_metadata(),
522 + reason="This test requires"
523 + "torch._inductor.utils.CUDAGraphWrapperMetadata to run",
524 + ),
525 + ),
526 + False,
527 + ],
528 +)
529 +def test_attn_quant(
530 + model_name: str,
531 + model_kwargs: dict[str, Any],
532 + backend: AttentionBackendEnum,
533 + matches: Matches,
534 + custom_ops: str,
535 + inductor_graph_partition: bool,
536 + caplog_mp_spawn,
537 + monkeypatch,
538 +):
539 + if not current_platform.has_device_capability(90):
540 + pytest.skip("test_attn_quant requires H100 (SM90) or B200 (SM100) GPU")
541 + if backend == AttentionBackendEnum.FLASHINFER and (
542 + not is_blackwell() or not has_flashinfer()
543 + ):
544 + pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
545 + if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
546 + pytest.skip("Inductor graph partition requires torch>=2.9")
547 +
548 + custom_ops_list = custom_ops.split(",") if custom_ops else []
549 +
550 + if inductor_graph_partition:
551 + mode = CUDAGraphMode.FULL_AND_PIECEWISE
552 + splitting_ops: list[str] | None = None
553 + else:
554 + # FIXME: Llama-4-Scout-17B-16E-Instruct-FP8 + FlashInfer + Blackwell end at
555 + # CUDAGraphMode.NONE here because it derives an attention backend that
556 + # does not support full cudagraphs
557 + mode = CUDAGraphMode.FULL_DECODE_ONLY
558 + splitting_ops = []
559 +
560 + # Disable, compile cache to make sure custom passes run.
561 + # Otherwise, we can't verify fusion happened through the logs.
562 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
563 +
564 + # To capture subprocess logs, we need to know whether spawn or fork is used.
565 + # Force spawn as it is more general.
566 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
567 + model_kwargs["attention_config"] = {"backend": backend.name}
568 +
569 + compilation_config = CompilationConfig(
570 + # Testing properties
571 + custom_ops=custom_ops_list,
572 + use_inductor_graph_partition=inductor_graph_partition,
573 + cudagraph_mode=mode,
574 + splitting_ops=splitting_ops,
575 + # Common
576 + mode=CompilationMode.VLLM_COMPILE,
577 + pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True),
578 + # Inductor caches custom passes by default as well via uuid
579 + inductor_compile_config={"force_disable_caches": True},
580 + )
581 +
582 + with caplog_mp_spawn(logging.DEBUG) as log_holder:
583 + run_model(compilation_config, model_name, **model_kwargs)
584 +
585 + log_matches = re.findall(
586 + r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
587 + log_holder.text,
588 + )
589 + assert len(log_matches) == 1, log_holder.text
590 + assert int(log_matches[0]) == matches.attention_fusion
vllm/env_override.py
+1 -1
@@ -95,7 +95,7 @@ def memory_plan_reuse_patched(self):
95 # ===================================================
96 # This change monkeypatches get_graph_partition_signature in pytorch 2.9.0 to
97 # fix inductor partition + attention-nvfp4 quant fusion, tested in
98 -# `tests/compile/distributed/test_fusions_e2e.py::test_attn_quant`.
98 +# `tests/compile/test_fusion_attn.py::test_attn_quant`.
99 # For more context, see https://github.com/pytorch/pytorch/pull/165815.
100
101