| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.config import PassConfig |
| 8 | from vllm.platforms import current_platform |
| 9 | |
| 10 | from ...utils import multi_gpu_test |
| 11 | from .common import ( |
| 12 | INDUCTOR_GRAPH_PARTITION, |
| 13 | AttentionBackendCase, |
| 14 | Matches, |
| 15 | custom_ops_combos, |
| 16 | is_blackwell, |
| 17 | nvfp4_kernel_exposes_input_quant_key, |
| 18 | ) |
| 19 | from .models import ( |
| 20 | FLASHINFER_ATTN, |
| 21 | FLASHINFER_MLA_ATTN, |
| 22 | ROCM_AITER_MLA_ATTN, |
| 23 | ROCM_AITER_UNIFIED_ATTN, |
| 24 | ROCM_ATTN, |
| 25 | TRITON_ATTN, |
| 26 | deepseek_coder_v2_lite_fp8, |
| 27 | deepseek_r1_fp4, |
| 28 | deepseek_v3_fp8, |
| 29 | gpt_oss_20b, |
| 30 | llama3_8b, |
| 31 | llama3_8b_fp4, |
| 32 | llama3_8b_fp8, |
| 33 | llama4_scout_fp4, |
| 34 | llama4_scout_fp8, |
| 35 | qwen3_a3b, |
| 36 | qwen3_a3b_fp8, |
| 37 | ) |
| 38 | |
| 39 | pytestmark = pytest.mark.skipif( |
| 40 | not current_platform.is_cuda_alike(), reason="Only test CUDA/ROCm" |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet |
| 45 | # supported. ROCm carries only the payloads its two backends support. |
| 46 | TP2_FP8_MODELS = ( |
| 47 | [ |
| 48 | llama3_8b_fp8, |
| 49 | llama4_scout_fp8, |
| 50 | qwen3_a3b_fp8, |
| 51 | deepseek_coder_v2_lite_fp8, |
| 52 | deepseek_v3_fp8, |
| 53 | ] |
| 54 | if current_platform.is_cuda() |
| 55 | else [llama3_8b_fp8, qwen3_a3b_fp8, deepseek_coder_v2_lite_fp8] |
| 56 | ) |
| 57 | |
| 58 | TP2_FP8_ATTN_BACKENDS = ( |
| 59 | [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN] |
| 60 | if current_platform.is_cuda() |
| 61 | else [ROCM_AITER_UNIFIED_ATTN, ROCM_AITER_MLA_ATTN] |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | @multi_gpu_test(num_gpus=2) |
| 66 | @pytest.mark.parametrize( |
| 67 | "model_name, matches_fn, model_kwargs, hf_overrides", TP2_FP8_MODELS |
| 68 | ) |
| 69 | @pytest.mark.parametrize("attn_backend", TP2_FP8_ATTN_BACKENDS) |
| 70 | @pytest.mark.parametrize("n_layers", [4]) |
| 71 | @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) |
| 72 | @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) |
| 73 | def test_tp2_ar_rms_fp8_fusions( |
| 74 | model_name: str, |
| 75 | matches_fn: Callable[[int], Matches], |
| 76 | model_kwargs: dict, |
| 77 | hf_overrides: Callable[[int], dict], |
| 78 | attn_backend: AttentionBackendCase, |
| 79 | n_layers: int, |
| 80 | custom_ops: str, |
| 81 | inductor_graph_partition: bool, |
| 82 | run_e2e_fusion_test, |
| 83 | monkeypatch, |
| 84 | ): |
| 85 | matches = matches_fn(n_layers) |
| 86 | |
| 87 | model_name_lower = model_name.lower() |
| 88 | # DeepSeek-Coder on AITER MLA is the one static-FP8 payload here, validated |
| 89 | # against the native QuantFP8 pattern; every other DeepSeek stays group |
| 90 | # quantized as on CUDA. |
| 91 | rocm_static_fp8_mla = ( |
| 92 | current_platform.is_rocm() |
| 93 | and "deepseek-coder" in model_name_lower |
| 94 | and attn_backend.backend.name == "ROCM_AITER_MLA" |
| 95 | ) |
| 96 | block_fp8 = "qwen" in model_name_lower or "deepseek" in model_name_lower |
| 97 | if rocm_static_fp8_mla: |
| 98 | block_fp8 = False |
| 99 | if custom_ops != "-quant_fp8,-rms_norm": |
| 100 | pytest.skip("ROCm static-FP8 MLA parity uses native quant and RMSNorm") |
| 101 | # AITER's MoE GEMM rejects this reduced dummy model shape; keep AITER |
| 102 | # for MLA and the all-reduce, and let the MoE run on Triton. |
| 103 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0") |
| 104 | if block_fp8 and "-quant_fp8" in custom_ops: |
| 105 | # This is why config forces +quant_fp8 by default |
| 106 | pytest.skip("native QuantFP8 matching not supported for group quant") |
| 107 | |
| 108 | # Reduce size of model and skip weight loading time |
| 109 | model_kwargs["hf_overrides"] = hf_overrides(n_layers) |
| 110 | model_kwargs["load_format"] = "dummy" |
| 111 | model_kwargs["max_model_len"] = 1024 |
| 112 | model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} |
| 113 | model_kwargs["disable_custom_all_reduce"] = False |
| 114 | if current_platform.is_rocm() and ( |
| 115 | "llama-3" in model_name_lower or "qwen" in model_name_lower |
| 116 | ): |
| 117 | model_kwargs.update( |
| 118 | gpu_memory_utilization=0.01, |
| 119 | kv_cache_memory_bytes=64 * 1024**2, |
| 120 | max_num_batched_tokens=1024, |
| 121 | max_num_seqs=32, |
| 122 | ) |
| 123 | |
| 124 | compilation_config = dict( |
| 125 | use_inductor_graph_partition=inductor_graph_partition, |
| 126 | custom_ops=custom_ops.split(","), |
| 127 | pass_config=PassConfig( |
| 128 | fuse_norm_quant=True, |
| 129 | fuse_act_quant=False, |
| 130 | fuse_attn_quant=True, |
| 131 | enable_qk_norm_rope_fusion=True, |
| 132 | fuse_allreduce_rms=True, |
| 133 | ), |
| 134 | ) |
| 135 | |
| 136 | matches_check = [ |
| 137 | "rms_quant_fusion", |
| 138 | "norm_rope_fusion", |
| 139 | "attn_quant_fusion", |
| 140 | "ar_rms_fusion", |
| 141 | ] |
| 142 | |
| 143 | use_aiter = current_platform.is_rocm() |
| 144 | if use_aiter: |
| 145 | matches = matches._replace(aiter_ar_rms_fusion=matches.ar_rms_fusion) |
| 146 | # The AITER all-reduce pass registers the AR+RMS(+group-quant) patterns |
| 147 | # before its AR+RMS-only ones, so the standalone AITER RMS-quant pass |
| 148 | # does not own the post-AR sites here. ROCm's contract is native |
| 149 | # QK-norm+RoPE and attention-output quant fusion plus AITER AR+RMS. |
| 150 | matches_check = [ |
| 151 | "norm_rope_fusion", |
| 152 | "attn_quant_fusion", |
| 153 | "aiter_ar_rms_fusion", |
| 154 | ] |
| 155 | |
| 156 | run_e2e_fusion_test( |
| 157 | model_name, |
| 158 | matches, |
| 159 | model_kwargs, |
| 160 | attn_backend, |
| 161 | compilation_config, |
| 162 | matches_check, |
| 163 | tp_size=2, |
| 164 | use_aiter=use_aiter, |
| 165 | ) |
| 166 | |
| 167 | |
| 168 | @multi_gpu_test(num_gpus=2) |
| 169 | @pytest.mark.parametrize( |
| 170 | "model_name, matches_fn, model_kwargs, hf_overrides", |
| 171 | [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4], |
| 172 | ) |
| 173 | @pytest.mark.parametrize( |
| 174 | "attn_backend", |
| 175 | [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN], |
| 176 | ) |
| 177 | @pytest.mark.parametrize("n_layers", [4]) |
| 178 | @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) |
| 179 | @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) |
| 180 | @pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") |
| 181 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") |
| 182 | def test_tp2_ar_rms_fp4_fusions( |
| 183 | model_name: str, |
| 184 | matches_fn: Callable[[int], Matches], |
| 185 | model_kwargs: dict, |
| 186 | hf_overrides: Callable[[int], dict], |
| 187 | attn_backend: AttentionBackendCase, |
| 188 | n_layers: int, |
| 189 | custom_ops: str, |
| 190 | inductor_graph_partition: bool, |
| 191 | run_e2e_fusion_test, |
| 192 | ): |
| 193 | if nvfp4_kernel_exposes_input_quant_key(): |
| 194 | pytest.skip( |
| 195 | "NVFP4 kernel exposes input_quant_key; manual fusion fires " |
| 196 | "instead of compiler pass-based fusion" |
| 197 | ) |
| 198 | |
| 199 | matches = matches_fn(n_layers) |
| 200 | |
| 201 | # Reduce size of model and skip weight loading time |
| 202 | model_kwargs["hf_overrides"] = hf_overrides(n_layers) |
| 203 | model_kwargs["load_format"] = "dummy" |
| 204 | model_kwargs["max_model_len"] = 1024 |
| 205 | model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} |
| 206 | model_kwargs["disable_custom_all_reduce"] = False |
| 207 | |
| 208 | compilation_config = dict( |
| 209 | use_inductor_graph_partition=inductor_graph_partition, |
| 210 | custom_ops=custom_ops.split(","), |
| 211 | pass_config=PassConfig( |
| 212 | fuse_act_quant=False, |
| 213 | fuse_attn_quant=True, |
| 214 | fuse_allreduce_rms=True, |
| 215 | ), |
| 216 | ) |
| 217 | |
| 218 | matches_check = [ |
| 219 | "attn_quant_fusion", |
| 220 | "ar_rms_fusion", |
| 221 | ] |
| 222 | |
| 223 | run_e2e_fusion_test( |
| 224 | model_name, |
| 225 | matches, |
| 226 | model_kwargs, |
| 227 | attn_backend, |
| 228 | compilation_config, |
| 229 | matches_check, |
| 230 | tp_size=2, |
| 231 | ) |
| 232 | |
| 233 | |
| 234 | @multi_gpu_test(num_gpus=2) |
| 235 | @pytest.mark.parametrize( |
| 236 | "model_name, matches_fn, model_kwargs, hf_overrides, model_impl", |
| 237 | [ |
| 238 | (*llama3_8b, "auto"), |
| 239 | (*llama3_8b, "transformers"), |
| 240 | (*qwen3_a3b, "auto"), |
| 241 | (*gpt_oss_20b, "auto"), |
| 242 | ], |
| 243 | ) |
| 244 | @pytest.mark.parametrize( |
| 245 | "attn_backend", |
| 246 | [ |
| 247 | TRITON_ATTN, |
| 248 | FLASHINFER_ATTN, |
| 249 | ROCM_ATTN, |
| 250 | ROCM_AITER_UNIFIED_ATTN, |
| 251 | ], |
| 252 | ) |
| 253 | @pytest.mark.parametrize("n_layers", [4]) |
| 254 | @pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm"))) |
| 255 | @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) |
| 256 | @pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Only test CUDA/ROCm") |
| 257 | def test_tp2_ar_rms_fusions( |
| 258 | model_name: str, |
| 259 | matches_fn: Callable[[int], Matches], |
| 260 | model_kwargs: dict, |
| 261 | hf_overrides: Callable[[int], dict], |
| 262 | model_impl: str, |
| 263 | attn_backend: AttentionBackendCase, |
| 264 | n_layers: int, |
| 265 | custom_ops: str, |
| 266 | inductor_graph_partition: bool, |
| 267 | run_e2e_fusion_test, |
| 268 | ): |
| 269 | if model_impl == "transformers" and not current_platform.is_rocm(): |
| 270 | pytest.skip("Transformers 3D AR+RMS regression is ROCm-only") |
| 271 | |
| 272 | matches = matches_fn(n_layers) |
| 273 | if model_impl == "transformers": |
| 274 | # Transformers add+RMSNorm canonicalization exposes every generic |
| 275 | # AR+RMS fusion site, including the final norm. |
| 276 | matches = matches._replace(aiter_ar_rms_fusion=matches.ar_rms_fusion) |
| 277 | |
| 278 | # Reduce size of model and skip weight loading time |
| 279 | model_kwargs["hf_overrides"] = hf_overrides(n_layers) |
| 280 | model_kwargs["load_format"] = "dummy" |
| 281 | model_kwargs["model_impl"] = model_impl |
| 282 | model_kwargs["max_model_len"] = 1024 |
| 283 | model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} |
| 284 | model_kwargs["disable_custom_all_reduce"] = False |
| 285 | |
| 286 | compilation_config = dict( |
| 287 | use_inductor_graph_partition=inductor_graph_partition, |
| 288 | custom_ops=custom_ops.split(","), |
| 289 | pass_config=PassConfig( |
| 290 | enable_qk_norm_rope_fusion=True, |
| 291 | fuse_allreduce_rms=True, |
| 292 | ), |
| 293 | ) |
| 294 | |
| 295 | matches_check = [ |
| 296 | "norm_rope_fusion", |
| 297 | ] |
| 298 | |
| 299 | if current_platform.is_rocm(): |
| 300 | matches_check.append("aiter_ar_rms_fusion") |
| 301 | else: |
| 302 | matches_check.append("ar_rms_fusion") |
| 303 | |
| 304 | run_e2e_fusion_test( |
| 305 | model_name, |
| 306 | matches, |
| 307 | model_kwargs, |
| 308 | attn_backend, |
| 309 | compilation_config, |
| 310 | matches_check, |
| 311 | tp_size=2, |
| 312 | use_aiter=current_platform.is_rocm(), |
| 313 | ) |