| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import logging |
| 4 | from collections import defaultdict |
| 5 | |
| 6 | import pytest |
| 7 | import regex as re |
| 8 | |
| 9 | from vllm import LLM, SamplingParams |
| 10 | from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode |
| 11 | |
| 12 | from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches |
| 13 | |
| 14 | |
| 15 | def run_model( |
| 16 | compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs |
| 17 | ): |
| 18 | """Run a model with the given compilation config for E2E fusion tests.""" |
| 19 | compilation_config = ( |
| 20 | compile_config |
| 21 | if isinstance(compile_config, CompilationConfig) |
| 22 | else CompilationConfig(mode=compile_config) |
| 23 | ) |
| 24 | |
| 25 | prompts = [ |
| 26 | "Hello, my name is", |
| 27 | "The president of the United States is", |
| 28 | "The capital of France is", |
| 29 | "The future of AI is", |
| 30 | ] |
| 31 | sampling_params = SamplingParams(temperature=0) |
| 32 | # Allow override from model_kwargs |
| 33 | model_kwargs = {"tensor_parallel_size": 1, **model_kwargs} |
| 34 | model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs} |
| 35 | |
| 36 | # No cudagraphs by default |
| 37 | if compilation_config.cudagraph_mode is None: |
| 38 | compilation_config.cudagraph_mode = CUDAGraphMode.NONE |
| 39 | llm = LLM( |
| 40 | model=model, |
| 41 | compilation_config=compilation_config, |
| 42 | **model_kwargs, |
| 43 | ) |
| 44 | outputs = llm.generate(prompts, sampling_params) |
| 45 | |
| 46 | # Print the outputs. |
| 47 | for output in outputs: |
| 48 | prompt = output.prompt |
| 49 | generated_text = output.outputs[0].text |
| 50 | print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") |
| 51 | |
| 52 | # Get the compile ranges endpoints after vllm config post init |
| 53 | # in order to compute compile ranges correctly |
| 54 | compilation_config.compile_ranges_endpoints = ( |
| 55 | llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints |
| 56 | ) |
| 57 | |
| 58 | # Fetch match table from each worker via RPC and sum across workers. |
| 59 | worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc( |
| 60 | "get_compilation_match_table" |
| 61 | ) |
| 62 | combined: defaultdict[str, int] = defaultdict(int) |
| 63 | for table in worker_tables: |
| 64 | for k, v in table.items(): |
| 65 | combined[k] += v |
| 66 | return dict(combined) |
| 67 | |
| 68 | |
| 69 | @pytest.fixture |
| 70 | def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): |
| 71 | def run( |
| 72 | model_name: str, |
| 73 | matches: Matches, |
| 74 | model_kwargs: dict, |
| 75 | attn_backend: AttentionBackendCase, |
| 76 | compilation_config: dict, |
| 77 | matches_check: list[str], |
| 78 | use_deepgemm: bool = False, |
| 79 | use_aiter: bool = False, |
| 80 | tp_size: int = 1, |
| 81 | ): |
| 82 | monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0") |
| 83 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") |
| 84 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1" if use_aiter else "0") |
| 85 | from vllm._aiter_ops import rocm_aiter_ops |
| 86 | |
| 87 | rocm_aiter_ops.refresh_env_variables() |
| 88 | |
| 89 | # Filter here to reduce code duplication |
| 90 | backend_name = attn_backend.backend.name.lower() |
| 91 | requires_mla = "deepseek" in model_name.lower() |
| 92 | is_mla = "mla" in backend_name |
| 93 | |
| 94 | if requires_mla != is_mla: |
| 95 | pytest.skip( |
| 96 | f"Incompatible model '{model_name}' and " |
| 97 | f"attention backend '{attn_backend.backend.name}'" |
| 98 | ) |
| 99 | |
| 100 | if backend_name == "rocm_attn" and model_name == "openai/gpt-oss-20b": |
| 101 | pytest.skip( |
| 102 | "ROCM_ATTN does not support attention sinks (required by gpt-oss-20b)" |
| 103 | ) |
| 104 | |
| 105 | if attn_backend.backend.name == "FLASHINFER": |
| 106 | from vllm.utils.flashinfer import supports_trtllm_attention |
| 107 | |
| 108 | if not supports_trtllm_attention(): |
| 109 | matches = matches._replace(attn_quant_fusion=0) |
| 110 | |
| 111 | # Disable, compile cache to make sure custom passes run. |
| 112 | # Otherwise, we can't verify fusion happened through the logs. |
| 113 | monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") |
| 114 | |
| 115 | # To capture subprocess logs, we need to know whether spawn or fork is used. |
| 116 | # Force spawn as it is more general. |
| 117 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 118 | |
| 119 | model_kwargs = {**attn_backend.model_kwargs, **model_kwargs} |
| 120 | model_kwargs["attention_config"] = {"backend": attn_backend.backend.name} |
| 121 | model_kwargs["tensor_parallel_size"] = tp_size |
| 122 | |
| 123 | # Cap warmup memory: tests use small max_model_len (1024) but the |
| 124 | # engine default max_num_batched_tokens is 16384. Warming up large |
| 125 | # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM. |
| 126 | model_kwargs.setdefault("max_num_batched_tokens", 8192) |
| 127 | |
| 128 | # Always compile the full graph instead of piecewise |
| 129 | if not compilation_config["use_inductor_graph_partition"]: |
| 130 | compilation_config["splitting_ops"] = [] |
| 131 | |
| 132 | full_compilation_config = CompilationConfig( |
| 133 | cudagraph_mode=CUDAGraphMode.NONE, |
| 134 | mode=CompilationMode.VLLM_COMPILE, |
| 135 | inductor_compile_config={"force_disable_caches": True}, |
| 136 | **compilation_config, |
| 137 | ) |
| 138 | |
| 139 | with caplog_mp_spawn(logging.DEBUG) as log_holder: |
| 140 | match_table = run_model(full_compilation_config, model_name, **model_kwargs) |
| 141 | |
| 142 | num_compile_ranges = len(full_compilation_config.get_compile_ranges()) |
| 143 | assert num_compile_ranges in [1, 2, 3] |
| 144 | |
| 145 | print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}") |
| 146 | print("Fusion results:") |
| 147 | |
| 148 | # Iterate through all so printing happens before asserting |
| 149 | log_matches_dict = {} |
| 150 | for match_name, pattern in FUSION_LOG_PATTERNS.items(): |
| 151 | log_matches_dict[match_name] = list(pattern.findall(log_holder.text)) |
| 152 | print(f"- {match_name}={','.join(log_matches_dict[match_name])}") |
| 153 | |
| 154 | # Now check the matches |
| 155 | for match_name in matches_check: |
| 156 | log_matches = list(int(ms) for ms in log_matches_dict[match_name]) |
| 157 | |
| 158 | # AR+RMS skips the largest range; SP skips the smallest. |
| 159 | # When both are enabled, AR+RMS activation count is |
| 160 | # model-dependent (hidden_size affects threshold), so derive |
| 161 | # from log data. |
| 162 | if ( |
| 163 | match_name == "ar_rms_fusion" |
| 164 | and "sequence_parallel" in matches_check |
| 165 | and num_compile_ranges >= 2 |
| 166 | ): |
| 167 | assert ( |
| 168 | len(log_matches) >= tp_size and len(log_matches) % tp_size == 0 |
| 169 | ), ( |
| 170 | f"Expected multiple of {tp_size} ar_rms log entries, " |
| 171 | f"found {len(log_matches)}" |
| 172 | ) |
| 173 | num_ranges_activated = len(log_matches) // tp_size |
| 174 | elif ( |
| 175 | match_name in ("ar_rms_fusion", "sequence_parallel") |
| 176 | and num_compile_ranges >= 2 |
| 177 | ): |
| 178 | num_ranges_activated = num_compile_ranges - 1 |
| 179 | else: |
| 180 | num_ranges_activated = num_compile_ranges |
| 181 | |
| 182 | # TODO: Remove log counting in unit tests |
| 183 | # once all matchers implement VllmFusionPatternMatcherPass |
| 184 | n_expected = tp_size * num_ranges_activated |
| 185 | if match_name not in ( |
| 186 | "attn_quant_fusion", |
| 187 | "act_quant_fusion", |
| 188 | "norm_rope_fusion", |
| 189 | ): |
| 190 | assert len(log_matches) == n_expected, ( |
| 191 | f"Could not find {n_expected} {match_name} " |
| 192 | f"(found {len(log_matches)}) in:\n {log_holder.text}" |
| 193 | ) |
| 194 | |
| 195 | expected_matches = getattr(matches, match_name) |
| 196 | |
| 197 | if match_name == "norm_rope_fusion": |
| 198 | # Opaque LayerName lets the combined pass consume short-range |
| 199 | # sites before the standalone pass; count either owner. |
| 200 | assert len(log_matches) >= n_expected |
| 201 | assert all(m in (0, expected_matches) for m in log_matches) |
| 202 | assert sum(log_matches) == expected_matches * n_expected, ( |
| 203 | f"{match_name} expected {expected_matches * n_expected} " |
| 204 | f"sites across both passes, found: {log_matches}" |
| 205 | ) |
| 206 | elif match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check: |
| 207 | # AR+rms+quant takes precedence over rms+quant if activated. |
| 208 | # That means we get full matching where ar+rms+quant was not |
| 209 | # activated, and less where it was (only the smallest range). |
| 210 | assert sum(m == expected_matches for m in log_matches) == tp_size * ( |
| 211 | num_ranges_activated - 1 |
| 212 | ), "Expecting full rms+quant fusion where ar+rms+quant not activated" |
| 213 | |
| 214 | assert all( |
| 215 | expected_matches - matches.ar_rms_fusion <= m <= expected_matches |
| 216 | for m in log_matches |
| 217 | ), ( |
| 218 | f"Expecting at least {expected_matches - matches.ar_rms_fusion} " |
| 219 | f"where ar+rms+quant was activated" |
| 220 | ) |
| 221 | elif ( |
| 222 | match_name == "async_tp" |
| 223 | and "sequence_parallel" in matches_check |
| 224 | and num_compile_ranges >= 2 |
| 225 | ): |
| 226 | # AsyncTP only finds patterns on ranges where SP ran. |
| 227 | n_sp_ranges = num_compile_ranges - 1 |
| 228 | assert ( |
| 229 | sum(m == expected_matches for m in log_matches) |
| 230 | == tp_size * n_sp_ranges |
| 231 | ), ( |
| 232 | f"Expecting {expected_matches} async_tp on " |
| 233 | f"{tp_size * n_sp_ranges} SP-range entries, " |
| 234 | f"found: {log_matches}" |
| 235 | ) |
| 236 | assert sum(m == 0 for m in log_matches) == tp_size, ( |
| 237 | f"Expecting 0 async_tp on {tp_size} small-range entries " |
| 238 | f"(no SP), found: {log_matches}" |
| 239 | ) |
| 240 | elif ( |
| 241 | match_name == "ar_rms_fusion" |
| 242 | and "sequence_parallel" in matches_check |
| 243 | and num_compile_ranges >= 2 |
| 244 | ): |
| 245 | # SP consumes allreduce patterns first, so AR+RMS finds |
| 246 | # full matches only on the smallest range (no SP). |
| 247 | assert sum(m == expected_matches for m in log_matches) == tp_size, ( |
| 248 | f"Expecting {expected_matches} ar_rms on " |
| 249 | f"{tp_size} small-range entries, found: {log_matches}" |
| 250 | ) |
| 251 | assert sum(m == 0 for m in log_matches) == tp_size * ( |
| 252 | num_ranges_activated - 1 |
| 253 | ), ( |
| 254 | f"Expecting 0 ar_rms on " |
| 255 | f"{tp_size * (num_ranges_activated - 1)} large-range " |
| 256 | f"entries (SP took precedence), found: {log_matches}" |
| 257 | ) |
| 258 | |
| 259 | elif match_name == "act_quant_fusion": |
| 260 | actual_match = match_table.get("activation_quant_fusion_pass", 0) |
| 261 | assert actual_match == expected_matches * n_expected, ( |
| 262 | f"Could not find {expected_matches * n_expected} " |
| 263 | f"{match_name} (found {actual_match})." |
| 264 | ) |
| 265 | elif match_name == "attn_quant_fusion": |
| 266 | actual_match = match_table.get( |
| 267 | "attn_quant_fusion", 0 |
| 268 | ) + match_table.get("mla_attn_quant_fusion", 0) |
| 269 | assert actual_match == expected_matches * n_expected, ( |
| 270 | f"Could not find {expected_matches * n_expected} " |
| 271 | f"{match_name} (found {actual_match})." |
| 272 | ) |
| 273 | else: |
| 274 | expected_matches_list = [expected_matches] * n_expected |
| 275 | assert sorted(log_matches) == expected_matches_list, ( |
| 276 | f"{match_name} expected: {expected_matches_list}, " |
| 277 | f"found: {sorted(log_matches)}" |
| 278 | ) |
| 279 | |
| 280 | if match_name == "ar_rms_fusion" and num_compile_ranges >= 2: |
| 281 | log_matches = re.findall( |
| 282 | r"pass_manager.py:\d+] Skipping " |
| 283 | r".*AllReduceFusionPass.* with compile range", |
| 284 | log_holder.text, |
| 285 | ) |
| 286 | |
| 287 | n_expected = tp_size * (num_compile_ranges - num_ranges_activated) |
| 288 | assert len(log_matches) == n_expected, ( |
| 289 | f'Could not find {n_expected} "Skipping AllReduceFusionPass" ' |
| 290 | f"(found {len(log_matches)}) in:\n {log_holder.text}" |
| 291 | ) |
| 292 | |
| 293 | if match_name == "sequence_parallel" and num_compile_ranges >= 2: |
| 294 | log_matches = re.findall( |
| 295 | r"pass_manager.py:\d+] Skipping " |
| 296 | r".*SequenceParallelismPass.* with compile range", |
| 297 | log_holder.text, |
| 298 | ) |
| 299 | |
| 300 | n_expected = tp_size * (num_compile_ranges - num_ranges_activated) |
| 301 | assert len(log_matches) == n_expected, ( |
| 302 | f'Could not find {n_expected} "Skipping SequenceParallelismPass" ' |
| 303 | f"(found {len(log_matches)}) in:\n {log_holder.text}" |
| 304 | ) |
| 305 | |
| 306 | return run |