| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | import logging |
| 6 | from contextlib import nullcontext |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS |
| 11 | from tests.models.registry import HF_EXAMPLE_MODELS |
| 12 | from tests.utils import ( |
| 13 | compare_two_settings, |
| 14 | create_new_process_for_each_test, |
| 15 | multi_gpu_test, |
| 16 | ) |
| 17 | from vllm.config import ( |
| 18 | CompilationMode, |
| 19 | ) |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.transformers_utils.repo_utils import hf_api |
| 22 | from vllm.utils.flashinfer import has_flashinfer |
| 23 | |
| 24 | NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" |
| 25 | NVFP4_HF_OVERRIDES = { |
| 26 | "num_hidden_layers": 4, |
| 27 | "hidden_size": 512, |
| 28 | "intermediate_size": 800, |
| 29 | "num_attention_heads": 4, |
| 30 | "num_key_value_heads": 1, |
| 31 | } |
| 32 | |
| 33 | |
| 34 | @multi_gpu_test(num_gpus=2) |
| 35 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") |
| 36 | def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn): |
| 37 | """Trained-model outputs must agree with unfused TP, including padded decode.""" |
| 38 | # A snapshot path pins both servers and the comparison helper's tokenizer. |
| 39 | model = hf_api().snapshot_download( |
| 40 | "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca" |
| 41 | ) |
| 42 | common_args = [ |
| 43 | "--dtype", |
| 44 | "bfloat16", |
| 45 | "--load-format", |
| 46 | "auto", |
| 47 | "--tensor-parallel-size", |
| 48 | "2", |
| 49 | "--distributed-executor-backend", |
| 50 | "mp", |
| 51 | "--max-model-len", |
| 52 | "512", |
| 53 | "--max-num-batched-tokens", |
| 54 | "512", |
| 55 | "--max-num-seqs", |
| 56 | "8", |
| 57 | "--kv-cache-memory-bytes", |
| 58 | str(64 * 1024**2), |
| 59 | "--gpu-memory-utilization", |
| 60 | "0.01", |
| 61 | "--attention-backend", |
| 62 | "TRITON_ATTN", |
| 63 | "--disable-custom-all-reduce", |
| 64 | ] |
| 65 | |
| 66 | def settings(enabled): |
| 67 | return [ |
| 68 | *common_args, |
| 69 | "--compilation-config", |
| 70 | json.dumps( |
| 71 | { |
| 72 | "mode": CompilationMode.VLLM_COMPILE, |
| 73 | "cudagraph_mode": "NONE", |
| 74 | "splitting_ops": [], |
| 75 | "inductor_compile_config": {"force_disable_caches": True}, |
| 76 | "pass_config": { |
| 77 | "enable_sp": enabled, |
| 78 | "fuse_gemm_comms": enabled, |
| 79 | "fuse_allreduce_rms": False, |
| 80 | "sp_min_token_num": 1, |
| 81 | }, |
| 82 | } |
| 83 | ), |
| 84 | ] |
| 85 | |
| 86 | env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"} |
| 87 | with caplog_mp_spawn(logging.DEBUG) as logs: |
| 88 | compare_two_settings( |
| 89 | model, |
| 90 | settings(False), |
| 91 | settings(True), |
| 92 | env1=env, |
| 93 | env2=env, |
| 94 | method="generate", |
| 95 | force_v1_runner=True, |
| 96 | ) |
| 97 | matches = [ |
| 98 | int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text) |
| 99 | ] |
| 100 | print(f"AsyncTP rewrite counts: {matches}") |
| 101 | assert len(matches) >= 2 and all(count > 0 for count in matches), ( |
| 102 | f"Expected positive AsyncTP rewrites on both TP workers, found {matches}" |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | @create_new_process_for_each_test() |
| 107 | @pytest.mark.parametrize( |
| 108 | "model_id", |
| 109 | ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Llama-3.2-1B-Instruct-FP8"], |
| 110 | ) |
| 111 | @pytest.mark.parametrize("tp_size", [2]) |
| 112 | @pytest.mark.parametrize("async_tp_enabled", [True]) |
| 113 | @pytest.mark.parametrize("distributed_backend", ["mp"]) |
| 114 | def test_async_tp_pass_correctness( |
| 115 | model_id: str, |
| 116 | tp_size: int, |
| 117 | async_tp_enabled: bool, |
| 118 | distributed_backend: str, |
| 119 | num_gpus_available: int, |
| 120 | monkeypatch, |
| 121 | ): |
| 122 | # Disable FlashInfer FP8 scaled_mm kernel as it is incompatible with |
| 123 | # async TP patterns. No-op on H100 (kernel requires CC >= 100). |
| 124 | monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel") |
| 125 | |
| 126 | model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) |
| 127 | model_info.check_transformers_version(on_fail="skip") |
| 128 | model_info.check_available_online(on_fail="skip") |
| 129 | |
| 130 | pp_size = 1 |
| 131 | if num_gpus_available < tp_size: |
| 132 | pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs") |
| 133 | |
| 134 | common_args = [ |
| 135 | "--dtype", |
| 136 | "bfloat16", |
| 137 | "--max-model-len", |
| 138 | "2048", |
| 139 | "--max-num-seqs", |
| 140 | "8", |
| 141 | ] |
| 142 | |
| 143 | compilation_config = { |
| 144 | "mode": CompilationMode.VLLM_COMPILE, |
| 145 | "compile_sizes": [2, 4, 8], |
| 146 | "splitting_ops": [], |
| 147 | "pass_config": {"fuse_gemm_comms": async_tp_enabled}, |
| 148 | } |
| 149 | |
| 150 | async_tp_args = [ |
| 151 | *common_args, |
| 152 | "--tensor-parallel-size", |
| 153 | str(tp_size), |
| 154 | "--distributed-executor-backend", |
| 155 | distributed_backend, |
| 156 | "--compilation_config", |
| 157 | json.dumps(compilation_config), |
| 158 | ] |
| 159 | |
| 160 | tp_args = [ |
| 161 | *common_args, |
| 162 | "--tensor-parallel-size", |
| 163 | str(tp_size), |
| 164 | "--distributed-executor-backend", |
| 165 | "mp", |
| 166 | ] |
| 167 | |
| 168 | compare_two_settings( |
| 169 | model_id, |
| 170 | async_tp_args, |
| 171 | tp_args, |
| 172 | method="generate", |
| 173 | force_v1_runner=True, |
| 174 | ) |
| 175 | |
| 176 | |
| 177 | @create_new_process_for_each_test() |
| 178 | def test_async_tp_pass_nvfp4_correctness(num_gpus_available: int, caplog_mp_spawn): |
| 179 | is_blackwell = current_platform.is_cuda() and ( |
| 180 | current_platform.is_device_capability_family(100) |
| 181 | ) |
| 182 | is_rocm_cdna = False |
| 183 | if current_platform.is_rocm(): |
| 184 | from vllm.platforms.rocm import on_cdna |
| 185 | |
| 186 | is_rocm_cdna = on_cdna() |
| 187 | if not (is_blackwell or is_rocm_cdna): |
| 188 | pytest.skip("NVFP4 AsyncTP requires Blackwell or ROCm CDNA") |
| 189 | if is_blackwell and not has_flashinfer(): |
| 190 | pytest.skip("FlashInfer is required for the NVFP4 AsyncTP path") |
| 191 | |
| 192 | tp_size = 2 |
| 193 | if num_gpus_available < tp_size: |
| 194 | pytest.skip(f"Need at least {tp_size} GPUs") |
| 195 | |
| 196 | common_args = [ |
| 197 | "--dtype", |
| 198 | "bfloat16", |
| 199 | "--max-model-len", |
| 200 | "2048", |
| 201 | "--max-num-seqs", |
| 202 | "8", |
| 203 | "--load-format", |
| 204 | "dummy", |
| 205 | *(["--linear-backend", "flashinfer_cutlass"] if is_blackwell else []), |
| 206 | "--hf-overrides", |
| 207 | json.dumps(NVFP4_HF_OVERRIDES), |
| 208 | ] |
| 209 | |
| 210 | compilation_config = { |
| 211 | "mode": CompilationMode.VLLM_COMPILE, |
| 212 | "compile_sizes": [2, 4, 8], |
| 213 | "splitting_ops": [], |
| 214 | "pass_config": { |
| 215 | "enable_sp": True, |
| 216 | "fuse_gemm_comms": True, |
| 217 | "fuse_allreduce_rms": False, |
| 218 | "sp_min_token_num": 1, |
| 219 | }, |
| 220 | } |
| 221 | if is_rocm_cdna: |
| 222 | compilation_config["inductor_compile_config"] = {"force_disable_caches": True} |
| 223 | |
| 224 | async_tp_args = [ |
| 225 | *common_args, |
| 226 | "--tensor-parallel-size", |
| 227 | str(tp_size), |
| 228 | "--distributed-executor-backend", |
| 229 | "mp", |
| 230 | "--compilation_config", |
| 231 | json.dumps(compilation_config), |
| 232 | ] |
| 233 | |
| 234 | tp_args = [ |
| 235 | *common_args, |
| 236 | "--tensor-parallel-size", |
| 237 | str(tp_size), |
| 238 | "--distributed-executor-backend", |
| 239 | "mp", |
| 240 | ] |
| 241 | |
| 242 | env = {"VLLM_DISABLE_COMPILE_CACHE": "1"} if is_rocm_cdna else None |
| 243 | log_context = caplog_mp_spawn(logging.DEBUG) if is_rocm_cdna else nullcontext() |
| 244 | with log_context as logs: |
| 245 | compare_two_settings( |
| 246 | NVFP4_MODEL_ID, |
| 247 | async_tp_args, |
| 248 | tp_args, |
| 249 | env1=env, |
| 250 | env2=env, |
| 251 | method="generate", |
| 252 | force_v1_runner=True, |
| 253 | ) |
| 254 | if is_rocm_cdna: |
| 255 | assert logs is not None |
| 256 | matches_by_rank = { |
| 257 | rank: [ |
| 258 | int(count) |
| 259 | for line in logs.text.splitlines() |
| 260 | if line.startswith(f"(Worker_TP{rank} pid=") |
| 261 | for count in FUSION_LOG_PATTERNS["async_tp"].findall(line) |
| 262 | ] |
| 263 | for rank in range(tp_size) |
| 264 | } |
| 265 | assert all( |
| 266 | any(count > 0 for count in matches) for matches in matches_by_rank.values() |
| 267 | ), ( |
| 268 | "Expected positive NVFP4 AsyncTP rewrites on both TP workers, " |
| 269 | f"found {matches_by_rank}" |
| 270 | ) |