| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import contextlib |
| 5 | import importlib |
| 6 | import multiprocessing as mp |
| 7 | import os |
| 8 | import queue |
| 9 | import traceback |
| 10 | from functools import lru_cache |
| 11 | from types import SimpleNamespace |
| 12 | from typing import Literal |
| 13 | from unittest.mock import patch |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | import torch.distributed as dist |
| 18 | |
| 19 | import vllm.envs as envs |
| 20 | from vllm import LLM, SamplingParams |
| 21 | from vllm.distributed import cleanup_dist_env_and_memory |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.transformers_utils.repo_utils import hf_api |
| 24 | from vllm.utils.network_utils import get_open_port |
| 25 | |
| 26 | from ..utils import multi_gpu_test |
| 27 | |
| 28 | pytestmark = pytest.mark.skipif( |
| 29 | not current_platform.is_rocm(), |
| 30 | reason="ROCm-only quick-reduce tests", |
| 31 | ) |
| 32 | |
| 33 | MB = 1024 * 1024 |
| 34 | WORLD_SIZE = 2 |
| 35 | QUANT_LEVELS = ["FP", "INT8", "INT6", "INT4"] |
| 36 | |
| 37 | |
| 38 | def _log(message: str) -> None: |
| 39 | print(f"[rocm_quick_reduce] {message}", flush=True) |
| 40 | |
| 41 | |
| 42 | def _reload_envs(): |
| 43 | return importlib.reload(envs) |
| 44 | |
| 45 | |
| 46 | def _make_quick_allreduce( |
| 47 | *, |
| 48 | disabled: bool = False, |
| 49 | world_size: int = 2, |
| 50 | quant_level: str = "FP", |
| 51 | use_fp16_kernels: bool = False, |
| 52 | qr_max_size: int = 64 * MB, |
| 53 | ): |
| 54 | from vllm.distributed.device_communicators.quick_all_reduce import ( |
| 55 | QuickAllReduce, |
| 56 | QuickReduceRegime, |
| 57 | ) |
| 58 | |
| 59 | qar = QuickAllReduce.__new__(QuickAllReduce) |
| 60 | qar.disabled = disabled |
| 61 | qar.world_size = world_size |
| 62 | qar.use_fp16_kernels = use_fp16_kernels |
| 63 | qar.qr_quant_level = QuickReduceRegime[quant_level] |
| 64 | qar.qr_max_size = qr_max_size |
| 65 | qar.qr_min_size = None |
| 66 | qar.qr_quantization_min_size = None |
| 67 | return qar |
| 68 | |
| 69 | |
| 70 | def _quick_allreduce_worker( |
| 71 | rank: int, |
| 72 | port: int, |
| 73 | quant_level: str, |
| 74 | dtype_name: str, |
| 75 | cast_bf16: bool, |
| 76 | ): |
| 77 | os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level |
| 78 | os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1" if cast_bf16 else "0" |
| 79 | _log( |
| 80 | f"worker start: rank={rank} quant={quant_level} " |
| 81 | f"dtype={dtype_name} cast_bf16={cast_bf16}" |
| 82 | ) |
| 83 | |
| 84 | device = torch.device(f"cuda:{rank}") |
| 85 | torch.accelerator.set_device_index(device) |
| 86 | dist.init_process_group( |
| 87 | backend="gloo", |
| 88 | init_method=f"tcp://127.0.0.1:{port}", |
| 89 | rank=rank, |
| 90 | world_size=WORLD_SIZE, |
| 91 | ) |
| 92 | |
| 93 | qar = None |
| 94 | try: |
| 95 | from vllm.distributed.device_communicators.quick_all_reduce import ( |
| 96 | QuickAllReduce, |
| 97 | ) |
| 98 | |
| 99 | qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank) |
| 100 | assert not qar.disabled |
| 101 | |
| 102 | num_elements = 8 * MB if dtype_name == "float16" else 4 * MB |
| 103 | |
| 104 | dtype = getattr(torch, dtype_name) |
| 105 | inp = torch.ones(num_elements, dtype=dtype, device=device) |
| 106 | |
| 107 | assert qar.should_quick_allreduce(inp) |
| 108 | if cast_bf16: |
| 109 | assert qar.use_fp16_kernels |
| 110 | |
| 111 | out = qar.quick_all_reduce(inp) |
| 112 | assert torch.allclose(out, inp * WORLD_SIZE, atol=2.5, rtol=0.1) |
| 113 | _log( |
| 114 | f"worker complete: rank={rank} quant={quant_level} " |
| 115 | f"dtype={dtype_name} num_elements={num_elements} " |
| 116 | f"use_fp16_kernels={qar.use_fp16_kernels}" |
| 117 | ) |
| 118 | finally: |
| 119 | if qar is not None: |
| 120 | qar.close() |
| 121 | if dist.is_initialized(): |
| 122 | dist.destroy_process_group() |
| 123 | |
| 124 | |
| 125 | def _run_two_gpu_quick_allreduce_test( |
| 126 | *, |
| 127 | quant_level: str, |
| 128 | dtype_name: str, |
| 129 | cast_bf16: bool, |
| 130 | ): |
| 131 | _log( |
| 132 | f"launch 2-GPU case: quant={quant_level} " |
| 133 | f"dtype={dtype_name} cast_bf16={cast_bf16}" |
| 134 | ) |
| 135 | ctx = mp.get_context("spawn") |
| 136 | port = get_open_port() |
| 137 | procs = [] |
| 138 | |
| 139 | for rank in range(WORLD_SIZE): |
| 140 | proc = ctx.Process( |
| 141 | target=_quick_allreduce_worker, |
| 142 | args=(rank, port, quant_level, dtype_name, cast_bf16), |
| 143 | ) |
| 144 | proc.start() |
| 145 | procs.append(proc) |
| 146 | |
| 147 | for proc in procs: |
| 148 | proc.join(timeout=60) |
| 149 | assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}" |
| 150 | _log( |
| 151 | f"finished 2-GPU case: quant={quant_level} " |
| 152 | f"dtype={dtype_name} cast_bf16={cast_bf16}" |
| 153 | ) |
| 154 | |
| 155 | |
| 156 | CUDAGRAPH_WORLD_SIZE = 2 |
| 157 | CUDAGRAPH_ROUNDS = 10 |
| 158 | CUDAGRAPH_NUM_ELEMENTS = 1 << 21 # 2M fp16 = 4 MB, above the quick-reduce thresholds |
| 159 | |
| 160 | |
| 161 | def _quick_allreduce_cudagraph_worker( |
| 162 | rank: int, |
| 163 | world_size: int, |
| 164 | port: int, |
| 165 | quant_level: str, |
| 166 | ): |
| 167 | # FP keeps the all-reduce bit-exact for small fp16 integers, so every |
| 168 | # replayed round can be checked exactly. |
| 169 | os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level |
| 170 | os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0" |
| 171 | _log(f"cudagraph worker start: rank={rank} quant={quant_level}") |
| 172 | |
| 173 | device = torch.device(f"cuda:{rank}") |
| 174 | torch.accelerator.set_device_index(device) |
| 175 | dist.init_process_group( |
| 176 | backend="gloo", |
| 177 | init_method=f"tcp://127.0.0.1:{port}", |
| 178 | rank=rank, |
| 179 | world_size=world_size, |
| 180 | ) |
| 181 | |
| 182 | qar = None |
| 183 | try: |
| 184 | from vllm.distributed.device_communicators.quick_all_reduce import ( |
| 185 | QuickAllReduce, |
| 186 | ) |
| 187 | |
| 188 | qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank) |
| 189 | assert not qar.disabled |
| 190 | |
| 191 | N = CUDAGRAPH_NUM_ELEMENTS |
| 192 | inp = torch.empty(N, dtype=torch.float16, device=device) |
| 193 | out = torch.empty(N, dtype=torch.float16, device=device) |
| 194 | assert qar.should_quick_allreduce(inp) |
| 195 | |
| 196 | # Every rank contributes the same value v in a round, so the true |
| 197 | # cross-rank all-reduce sum is simply world_size * v. |
| 198 | def expected(v): |
| 199 | return float(world_size * v) |
| 200 | |
| 201 | if rank == 0: |
| 202 | print( |
| 203 | f"[repro] world_size={world_size} elems={N} regime={quant_level} fp16", |
| 204 | flush=True, |
| 205 | ) |
| 206 | |
| 207 | # Warmup, then capture a graph with EXACTLY ONE quick-reduce (isolated |
| 208 | # qr, so it is the sole writer of its flag slot -- the condition that |
| 209 | # triggers the stale-flag bug). |
| 210 | inp.fill_(1.0) |
| 211 | qar.quick_all_reduce(inp, out=out) |
| 212 | torch.accelerator.synchronize() |
| 213 | dist.barrier() |
| 214 | |
| 215 | g = torch.cuda.CUDAGraph() |
| 216 | with torch.cuda.graph(g): |
| 217 | qar.quick_all_reduce(inp, out=out) |
| 218 | torch.accelerator.synchronize() |
| 219 | dist.barrier() |
| 220 | |
| 221 | for v in range(CUDAGRAPH_ROUNDS): |
| 222 | inp.fill_(float(v)) # in-place: same value on every rank |
| 223 | dist.barrier() |
| 224 | g.replay() |
| 225 | torch.accelerator.synchronize() |
| 226 | dist.barrier() |
| 227 | got = out.float() |
| 228 | expect = expected(v) |
| 229 | if rank == 0: |
| 230 | print(f"round {v}: got={got[:10]}, expected={expect}", flush=True) |
| 231 | mismatch = ~torch.isclose(got, torch.full_like(got, expect)) |
| 232 | num_mismatch = int(mismatch.sum().item()) |
| 233 | assert num_mismatch == 0, ( |
| 234 | f"rank={rank} round={v} expected={expect} " |
| 235 | f"mismatched {num_mismatch}/{got.numel()} elements; " |
| 236 | f"unique wrong values={torch.unique(got[mismatch])[:8].tolist()}" |
| 237 | ) |
| 238 | _log(f"cudagraph worker complete: rank={rank} rounds={CUDAGRAPH_ROUNDS}") |
| 239 | finally: |
| 240 | if qar is not None: |
| 241 | qar.close() |
| 242 | if dist.is_initialized(): |
| 243 | dist.destroy_process_group() |
| 244 | |
| 245 | |
| 246 | def _run_cudagraph_replay_test(*, world_size: int, quant_level: str): |
| 247 | _log(f"launch {world_size}-GPU cudagraph replay case: quant={quant_level}") |
| 248 | ctx = mp.get_context("spawn") |
| 249 | port = get_open_port() |
| 250 | procs = [] |
| 251 | |
| 252 | for rank in range(world_size): |
| 253 | proc = ctx.Process( |
| 254 | target=_quick_allreduce_cudagraph_worker, |
| 255 | args=(rank, world_size, port, quant_level), |
| 256 | ) |
| 257 | proc.start() |
| 258 | procs.append(proc) |
| 259 | |
| 260 | for proc in procs: |
| 261 | proc.join(timeout=120) |
| 262 | assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}" |
| 263 | _log(f"finished {world_size}-GPU cudagraph replay case: quant={quant_level}") |
| 264 | |
| 265 | |
| 266 | MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" |
| 267 | E2E_PREFILL_TOKENS = 1024 |
| 268 | E2E_MAX_MODEL_LEN = 1536 |
| 269 | E2E_GPU_MEMORY_UTILIZATION = 0.3 |
| 270 | E2E_KV_CACHE_MEMORY_BYTES = 2 << 30 |
| 271 | |
| 272 | _BACKGROUND_LINE = ( |
| 273 | "Background filler: this archived operations memo repeats a routine status " |
| 274 | "line so the distributed test uses a realistically long prefill." |
| 275 | ) |
| 276 | _BACKGROUND_BLOCK = " ".join([_BACKGROUND_LINE] * 48) |
| 277 | |
| 278 | |
| 279 | def _build_prompt(*, fact_block: str, question: str) -> str: |
| 280 | return ( |
| 281 | "Read the archived operations memo below. Most of the memo is filler. " |
| 282 | "Use only the fact block near the end when answering.\n" |
| 283 | f"{_BACKGROUND_BLOCK}\n" |
| 284 | "Fact block:\n" |
| 285 | f"{fact_block}\n" |
| 286 | f"Question: {question}\n" |
| 287 | "Answer in one short sentence." |
| 288 | ) |
| 289 | |
| 290 | |
| 291 | E2E_PROMPTS = [ |
| 292 | _build_prompt( |
| 293 | fact_block=( |
| 294 | "- Festival city: Oslo\n- Mascot animal: otter\n- Welcome drink: tea" |
| 295 | ), |
| 296 | question="Which city hosts the festival, and what animal is the mascot?", |
| 297 | ), |
| 298 | _build_prompt( |
| 299 | fact_block=( |
| 300 | "- Meeting day: Tuesday\n" |
| 301 | "- Planned snack: apricot cake\n" |
| 302 | "- Backup room: Cedar" |
| 303 | ), |
| 304 | question="What day is the meeting, and what snack is planned?", |
| 305 | ), |
| 306 | ] |
| 307 | RECORDED_RESPONSE_TEXTS = ( |
| 308 | " The city hosting the festival is Oslo, and the mascot is an otter.", |
| 309 | " The meeting is on Tuesday and the snack planned is apricot cake.", |
| 310 | ) |
| 311 | REQUIRED_WORDS = (("oslo", "otter"), ("tuesday", "apricot")) |
| 312 | |
| 313 | |
| 314 | def _log_prompt_summaries() -> None: |
| 315 | for i, prompt in enumerate(E2E_PROMPTS): |
| 316 | prompt_lines = prompt.splitlines() |
| 317 | fact_block = [line for line in prompt_lines if line.startswith("- ")] |
| 318 | fact_summary = "; ".join(line.removeprefix("- ") for line in fact_block) |
| 319 | _log(f"prompt {i} facts: {fact_summary}") |
| 320 | |
| 321 | |
| 322 | @lru_cache(maxsize=1) |
| 323 | def _get_model_path() -> str: |
| 324 | try: |
| 325 | path = hf_api().snapshot_download(repo_id=MODEL_NAME, local_files_only=True) |
| 326 | _log(f"using cached model snapshot: {path}") |
| 327 | return path |
| 328 | except Exception: |
| 329 | path = hf_api().snapshot_download(repo_id=MODEL_NAME) |
| 330 | _log(f"downloaded model snapshot: {path}") |
| 331 | return path |
| 332 | |
| 333 | |
| 334 | def _get_hidden_size(model_config) -> int: |
| 335 | hidden_size = getattr(model_config, "hidden_size", None) |
| 336 | if hidden_size is None and hasattr(model_config, "text_config"): |
| 337 | hidden_size = getattr(model_config.text_config, "hidden_size", None) |
| 338 | assert isinstance(hidden_size, int) |
| 339 | return hidden_size |
| 340 | |
| 341 | |
| 342 | def _check_tp_allreduce_uses_quick_reduce( |
| 343 | self, |
| 344 | num_tokens: int, |
| 345 | dtype_name: str = "float16", |
| 346 | ) -> dict[str, int | bool]: |
| 347 | from vllm.distributed.communication_op import tensor_model_parallel_all_reduce |
| 348 | from vllm.distributed.parallel_state import get_tp_group |
| 349 | |
| 350 | assert self.device is not None |
| 351 | qr_comm = get_tp_group().device_communicator.qr_comm |
| 352 | assert qr_comm is not None |
| 353 | assert not qr_comm.disabled |
| 354 | |
| 355 | hidden_size = _get_hidden_size(self.model_runner.model.config) |
| 356 | dtype = getattr(torch, dtype_name) |
| 357 | sample = torch.full( |
| 358 | (num_tokens, hidden_size), |
| 359 | fill_value=float(self.rank + 1), |
| 360 | dtype=dtype, |
| 361 | device=self.device, |
| 362 | ) |
| 363 | |
| 364 | assert qr_comm.should_quick_allreduce(sample) |
| 365 | |
| 366 | expected = sample.clone() |
| 367 | reduced = tensor_model_parallel_all_reduce(sample) |
| 368 | dist.all_reduce(expected, group=get_tp_group().device_group) |
| 369 | torch.testing.assert_close(reduced, expected, atol=2.5, rtol=0.1) |
| 370 | |
| 371 | stats = { |
| 372 | "rank": self.rank, |
| 373 | "hidden_size": hidden_size, |
| 374 | "num_tokens": num_tokens, |
| 375 | "use_fp16_kernels": qr_comm.use_fp16_kernels, |
| 376 | } |
| 377 | _log( |
| 378 | "worker quick-reduce check: " |
| 379 | f"rank={self.rank} hidden_size={hidden_size} " |
| 380 | f"num_tokens={num_tokens} use_fp16_kernels={qr_comm.use_fp16_kernels}" |
| 381 | ) |
| 382 | return stats |
| 383 | |
| 384 | |
| 385 | def _check_quick_reduce_disabled(self) -> int: |
| 386 | from vllm.distributed.parallel_state import get_tp_group |
| 387 | |
| 388 | qr_comm = get_tp_group().device_communicator.qr_comm |
| 389 | assert qr_comm is not None |
| 390 | assert qr_comm.disabled |
| 391 | _log(f"worker confirmed quick reduce is disabled: rank={self.rank}") |
| 392 | return self.rank |
| 393 | |
| 394 | |
| 395 | def _collect_generations(outputs) -> list[tuple[tuple[int, ...], str]]: |
| 396 | return [ |
| 397 | (tuple(output.outputs[0].token_ids), output.outputs[0].text) |
| 398 | for output in outputs |
| 399 | ] |
| 400 | |
| 401 | |
| 402 | def _shutdown_llm(llm: LLM | None) -> None: |
| 403 | if llm is None: |
| 404 | cleanup_dist_env_and_memory() |
| 405 | return |
| 406 | |
| 407 | with contextlib.suppress(Exception): |
| 408 | llm.llm_engine.engine_core.shutdown() |
| 409 | |
| 410 | del llm |
| 411 | cleanup_dist_env_and_memory() |
| 412 | |
| 413 | |
| 414 | def _log_generations( |
| 415 | label: str, |
| 416 | generations: list[tuple[tuple[int, ...], str]], |
| 417 | ) -> None: |
| 418 | for i, (token_ids, text) in enumerate(generations): |
| 419 | _log(f"{label} prompt {i} token ids: {list(token_ids)}") |
| 420 | _log(f"{label} prompt {i} text: {text!r}") |
| 421 | |
| 422 | |
| 423 | def _assert_required_words( |
| 424 | label: str, |
| 425 | generations: list[tuple[tuple[int, ...], str]], |
| 426 | ) -> None: |
| 427 | for i, (_, text) in enumerate(generations): |
| 428 | lowered = text.lower() |
| 429 | missing = [word for word in REQUIRED_WORDS[i] if word not in lowered] |
| 430 | assert not missing, ( |
| 431 | f"{label} prompt {i} is missing required words {missing}. " |
| 432 | f"Observed text: {text!r}" |
| 433 | ) |
| 434 | |
| 435 | |
| 436 | def _collect_soft_mismatches( |
| 437 | baseline_generations: list[tuple[tuple[int, ...], str]], |
| 438 | quick_reduce_generations: list[tuple[tuple[int, ...], str]], |
| 439 | ) -> list[str]: |
| 440 | mismatches = [] |
| 441 | |
| 442 | for i, (_, text) in enumerate(baseline_generations): |
| 443 | expected = RECORDED_RESPONSE_TEXTS[i] |
| 444 | if text != expected: |
| 445 | mismatches.append( |
| 446 | f"baseline prompt {i} drifted from the recorded response.\n" |
| 447 | f"expected={expected!r}\nactual={text!r}" |
| 448 | ) |
| 449 | |
| 450 | for i, (_, text) in enumerate(quick_reduce_generations): |
| 451 | expected = RECORDED_RESPONSE_TEXTS[i] |
| 452 | if text != expected: |
| 453 | mismatches.append( |
| 454 | f"quick-reduce prompt {i} drifted from the recorded response.\n" |
| 455 | f"expected={expected!r}\nactual={text!r}" |
| 456 | ) |
| 457 | |
| 458 | for i, ((_, baseline_text), (_, quick_reduce_text)) in enumerate( |
| 459 | zip(baseline_generations, quick_reduce_generations) |
| 460 | ): |
| 461 | if baseline_text != quick_reduce_text: |
| 462 | mismatches.append( |
| 463 | f"baseline and quick-reduce responses differ for prompt {i}.\n" |
| 464 | f"baseline={baseline_text!r}\nquick_reduce={quick_reduce_text!r}" |
| 465 | ) |
| 466 | |
| 467 | return mismatches |
| 468 | |
| 469 | |
| 470 | def _run_generation( |
| 471 | *, |
| 472 | backend: Literal["mp", "ray"], |
| 473 | quant_mode: str, |
| 474 | expect_quick_reduce: bool, |
| 475 | ) -> list[tuple[tuple[int, ...], str]]: |
| 476 | llm = None |
| 477 | monkeypatch = pytest.MonkeyPatch() |
| 478 | with monkeypatch.context() as m: |
| 479 | m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 480 | m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode) |
| 481 | model_path = _get_model_path() |
| 482 | _log( |
| 483 | f"starting generation: backend={backend} quant={quant_mode} " |
| 484 | f"gpu_memory_utilization={E2E_GPU_MEMORY_UTILIZATION} " |
| 485 | f"kv_cache_bytes={E2E_KV_CACHE_MEMORY_BYTES} model={model_path}" |
| 486 | ) |
| 487 | |
| 488 | try: |
| 489 | llm = LLM( |
| 490 | model=model_path, |
| 491 | tokenizer=model_path, |
| 492 | tensor_parallel_size=2, |
| 493 | distributed_executor_backend=backend, |
| 494 | dtype="half", |
| 495 | enforce_eager=True, |
| 496 | max_model_len=E2E_MAX_MODEL_LEN, |
| 497 | max_num_seqs=len(E2E_PROMPTS), |
| 498 | gpu_memory_utilization=E2E_GPU_MEMORY_UTILIZATION, |
| 499 | kv_cache_memory_bytes=E2E_KV_CACHE_MEMORY_BYTES, |
| 500 | seed=0, |
| 501 | ) |
| 502 | |
| 503 | if not expect_quick_reduce: |
| 504 | assert llm.collective_rpc(_check_quick_reduce_disabled) == [0, 1] |
| 505 | |
| 506 | if expect_quick_reduce: |
| 507 | worker_stats = llm.collective_rpc( |
| 508 | _check_tp_allreduce_uses_quick_reduce, |
| 509 | args=(E2E_PREFILL_TOKENS,), |
| 510 | ) |
| 511 | assert [stat["rank"] for stat in worker_stats] == [0, 1] |
| 512 | worker_summary = "; ".join( |
| 513 | "rank={rank} hidden_size={hidden_size} num_tokens={num_tokens} " |
| 514 | "use_fp16_kernels={use_fp16_kernels}".format(**stat) |
| 515 | for stat in worker_stats |
| 516 | ) |
| 517 | _log(f"{backend} quick-reduce worker checks: {worker_summary}") |
| 518 | |
| 519 | outputs = llm.generate( |
| 520 | E2E_PROMPTS, |
| 521 | SamplingParams( |
| 522 | temperature=0.0, |
| 523 | max_tokens=20, |
| 524 | stop=["\nAnswer:", " Answer:"], |
| 525 | ), |
| 526 | use_tqdm=False, |
| 527 | ) |
| 528 | generations = _collect_generations(outputs) |
| 529 | assert all(text.strip() for _, text in generations) |
| 530 | _log_generations(f"{backend} {quant_mode}", generations) |
| 531 | return generations |
| 532 | finally: |
| 533 | _shutdown_llm(llm) |
| 534 | |
| 535 | |
| 536 | def _run_quick_reduce_llm_e2e_in_subprocess( |
| 537 | *, |
| 538 | backend: Literal["mp", "ray"], |
| 539 | ) -> str | None: |
| 540 | _log(f"running LLM e2e: backend={backend}") |
| 541 | _log_prompt_summaries() |
| 542 | baseline_outputs = _run_generation( |
| 543 | backend=backend, |
| 544 | quant_mode="NONE", |
| 545 | expect_quick_reduce=False, |
| 546 | ) |
| 547 | quick_reduce_outputs = _run_generation( |
| 548 | backend=backend, |
| 549 | quant_mode="FP", |
| 550 | expect_quick_reduce=True, |
| 551 | ) |
| 552 | |
| 553 | _assert_required_words("baseline", baseline_outputs) |
| 554 | _assert_required_words("quick-reduce", quick_reduce_outputs) |
| 555 | |
| 556 | mismatches = _collect_soft_mismatches(baseline_outputs, quick_reduce_outputs) |
| 557 | if mismatches: |
| 558 | details = "\n\n".join(mismatches) |
| 559 | _log(f"soft response mismatch:\n{details}") |
| 560 | return details |
| 561 | |
| 562 | _log(f"LLM e2e backend={backend} matched the recorded responses exactly") |
| 563 | return None |
| 564 | |
| 565 | |
| 566 | def _quick_reduce_llm_e2e_worker( |
| 567 | result_queue: mp.Queue, |
| 568 | backend: Literal["mp", "ray"], |
| 569 | ) -> None: |
| 570 | try: |
| 571 | xfail_reason = _run_quick_reduce_llm_e2e_in_subprocess(backend=backend) |
| 572 | except Exception: |
| 573 | result_queue.put({"status": "error", "reason": traceback.format_exc()}) |
| 574 | raise |
| 575 | else: |
| 576 | if xfail_reason is not None: |
| 577 | result_queue.put({"status": "xfail", "reason": xfail_reason}) |
| 578 | else: |
| 579 | result_queue.put({"status": "ok"}) |
| 580 | |
| 581 | |
| 582 | def run_quick_reduce_llm_e2e( |
| 583 | *, |
| 584 | backend: Literal["mp", "ray"], |
| 585 | ) -> None: |
| 586 | ctx = mp.get_context("spawn") |
| 587 | result_queue = ctx.Queue() |
| 588 | proc = ctx.Process( |
| 589 | target=_quick_reduce_llm_e2e_worker, |
| 590 | args=(result_queue, backend), |
| 591 | ) |
| 592 | proc.start() |
| 593 | proc.join(timeout=600) |
| 594 | |
| 595 | try: |
| 596 | result = result_queue.get(timeout=5) |
| 597 | except queue.Empty as exc: |
| 598 | if proc.exitcode != 0: |
| 599 | raise AssertionError( |
| 600 | f"quick-reduce llm e2e subprocess failed for backend={backend} " |
| 601 | f"with exit code {proc.exitcode} and produced no result" |
| 602 | ) from exc |
| 603 | raise AssertionError( |
| 604 | f"quick-reduce llm e2e subprocess produced no result for backend={backend}" |
| 605 | ) from exc |
| 606 | |
| 607 | if result["status"] == "xfail": |
| 608 | pytest.xfail(result["reason"]) |
| 609 | if result["status"] == "error": |
| 610 | raise AssertionError( |
| 611 | f"quick-reduce llm e2e subprocess failed for backend={backend}:\n" |
| 612 | f"{result['reason']}" |
| 613 | ) |
| 614 | |
| 615 | assert proc.exitcode == 0, ( |
| 616 | f"quick-reduce llm e2e subprocess failed for backend={backend} " |
| 617 | f"with exit code {proc.exitcode}" |
| 618 | ) |
| 619 | |
| 620 | |
| 621 | def test_quick_reduce_regime_values(): |
| 622 | from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime |
| 623 | |
| 624 | assert QuickReduceRegime.FP.value == 0 |
| 625 | assert QuickReduceRegime.INT8.value == 1 |
| 626 | assert QuickReduceRegime.INT6.value == 2 |
| 627 | assert QuickReduceRegime.INT4.value == 3 |
| 628 | assert QuickReduceRegime.INT3.value == 4 |
| 629 | assert QuickReduceRegime.NONE.value == 5 |
| 630 | |
| 631 | |
| 632 | def test_quick_reduce_regime_names(): |
| 633 | from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime |
| 634 | |
| 635 | assert set(QuickReduceRegime.__members__) == { |
| 636 | "FP", |
| 637 | "INT8", |
| 638 | "INT6", |
| 639 | "INT4", |
| 640 | "INT3", |
| 641 | "NONE", |
| 642 | } |
| 643 | |
| 644 | |
| 645 | @pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"]) |
| 646 | def test_quick_reduce_quantization_env_var(monkeypatch, quant_level): |
| 647 | monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_level) |
| 648 | |
| 649 | reloaded_envs = _reload_envs() |
| 650 | assert quant_level == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION |
| 651 | |
| 652 | |
| 653 | def test_quick_reduce_quantization_default(monkeypatch): |
| 654 | monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", raising=False) |
| 655 | |
| 656 | reloaded_envs = _reload_envs() |
| 657 | assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION == "NONE" |
| 658 | |
| 659 | |
| 660 | @pytest.mark.parametrize("cast_bf16", [True, False]) |
| 661 | def test_quick_reduce_cast_bf16_to_fp16_env_var(monkeypatch, cast_bf16): |
| 662 | monkeypatch.setenv( |
| 663 | "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1" if cast_bf16 else "0" |
| 664 | ) |
| 665 | |
| 666 | reloaded_envs = _reload_envs() |
| 667 | assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is cast_bf16 |
| 668 | |
| 669 | |
| 670 | def test_quick_reduce_cast_bf16_to_fp16_default(monkeypatch): |
| 671 | monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", raising=False) |
| 672 | |
| 673 | reloaded_envs = _reload_envs() |
| 674 | assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is True |
| 675 | |
| 676 | |
| 677 | @pytest.mark.parametrize("max_mb", [128, 512, 2048, None]) |
| 678 | def test_quick_reduce_max_size_env_var(monkeypatch, max_mb): |
| 679 | if max_mb is None: |
| 680 | monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False) |
| 681 | else: |
| 682 | monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", str(max_mb)) |
| 683 | |
| 684 | reloaded_envs = _reload_envs() |
| 685 | assert max_mb == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB |
| 686 | |
| 687 | |
| 688 | def test_quick_reduce_max_size_default(monkeypatch): |
| 689 | monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False) |
| 690 | |
| 691 | reloaded_envs = _reload_envs() |
| 692 | assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is None |
| 693 | |
| 694 | |
| 695 | @pytest.mark.parametrize( |
| 696 | ("gcn_arch_name", "expected"), |
| 697 | [ |
| 698 | ("gfx942", True), |
| 699 | ("gfx950", True), |
| 700 | ("gfx90a", False), |
| 701 | ("", False), |
| 702 | ], |
| 703 | ) |
| 704 | def test_quick_allreduce_rocm_arch_available(gcn_arch_name, expected): |
| 705 | from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce |
| 706 | |
| 707 | qar = QuickAllReduce.__new__(QuickAllReduce) |
| 708 | qar.disabled = True |
| 709 | |
| 710 | with ( |
| 711 | patch( |
| 712 | "vllm.distributed.device_communicators.quick_all_reduce.current_platform." |
| 713 | "is_rocm", |
| 714 | return_value=True, |
| 715 | ), |
| 716 | patch( |
| 717 | "torch.cuda.get_device_properties", |
| 718 | return_value=SimpleNamespace(gcnArchName=gcn_arch_name), |
| 719 | ), |
| 720 | ): |
| 721 | assert qar._rocm_arch_available() is expected |
| 722 | |
| 723 | |
| 724 | def test_quick_allreduce_rocm_arch_available_handles_probe_failure(): |
| 725 | from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce |
| 726 | |
| 727 | qar = QuickAllReduce.__new__(QuickAllReduce) |
| 728 | qar.disabled = True |
| 729 | |
| 730 | with ( |
| 731 | patch( |
| 732 | "vllm.distributed.device_communicators.quick_all_reduce.current_platform." |
| 733 | "is_rocm", |
| 734 | return_value=True, |
| 735 | ), |
| 736 | patch("torch.cuda.get_device_properties", side_effect=RuntimeError), |
| 737 | ): |
| 738 | assert qar._rocm_arch_available() is False |
| 739 | |
| 740 | |
| 741 | def test_quick_allreduce_rejects_disabled(): |
| 742 | qar = _make_quick_allreduce(disabled=True) |
| 743 | |
| 744 | inp = torch.zeros(1024, dtype=torch.float16) |
| 745 | assert qar.should_quick_allreduce(inp) is False |
| 746 | |
| 747 | |
| 748 | def test_quick_allreduce_rejects_unsupported_dtype(): |
| 749 | qar = _make_quick_allreduce() |
| 750 | |
| 751 | inp = torch.zeros(1024 * 1024, dtype=torch.float32) |
| 752 | assert qar.should_quick_allreduce(inp) is False |
| 753 | |
| 754 | |
| 755 | def test_quick_allreduce_rejects_non_aligned_input(): |
| 756 | qar = _make_quick_allreduce() |
| 757 | |
| 758 | inp = torch.zeros(5, dtype=torch.float16) |
| 759 | assert qar.should_quick_allreduce(inp) is False |
| 760 | |
| 761 | |
| 762 | def test_quick_allreduce_rejects_non_contiguous_input(): |
| 763 | qar = _make_quick_allreduce() |
| 764 | |
| 765 | inp = torch.zeros((1024, 1024), dtype=torch.float16)[:, ::2] |
| 766 | assert qar.should_quick_allreduce(inp) is False |
| 767 | |
| 768 | |
| 769 | def test_quick_allreduce_rejects_input_smaller_than_threshold(): |
| 770 | qar = _make_quick_allreduce() |
| 771 | |
| 772 | inp = torch.zeros((MB // 2) - 8, dtype=torch.float16) |
| 773 | assert qar.should_quick_allreduce(inp) is False |
| 774 | |
| 775 | |
| 776 | def test_quick_allreduce_accepts_input_at_threshold(): |
| 777 | qar = _make_quick_allreduce() |
| 778 | |
| 779 | inp = torch.zeros(MB // 2, dtype=torch.float16) |
| 780 | assert qar.should_quick_allreduce(inp) is True |
| 781 | |
| 782 | |
| 783 | def test_quick_allreduce_rejects_input_larger_than_max_size(): |
| 784 | qar = _make_quick_allreduce(qr_max_size=1 * MB) |
| 785 | |
| 786 | inp = torch.zeros(MB, dtype=torch.float16) |
| 787 | assert qar.should_quick_allreduce(inp) is False |
| 788 | |
| 789 | |
| 790 | def test_quick_allreduce_bf16_uses_fp16_threshold_when_cast_enabled(): |
| 791 | inp = torch.zeros(MB // 2, dtype=torch.bfloat16) |
| 792 | |
| 793 | without_cast = _make_quick_allreduce(use_fp16_kernels=False) |
| 794 | with_cast = _make_quick_allreduce(use_fp16_kernels=True) |
| 795 | |
| 796 | assert without_cast.should_quick_allreduce(inp) is False |
| 797 | assert with_cast.should_quick_allreduce(inp) is True |
| 798 | |
| 799 | |
| 800 | def test_quick_allreduce_supported_world_sizes(): |
| 801 | from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce |
| 802 | |
| 803 | assert QuickAllReduce._SUPPORTED_WORLD_SIZES == [2, 4, 8] |
| 804 | |
| 805 | |
| 806 | def test_quick_allreduce_supported_dtypes(): |
| 807 | from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce |
| 808 | |
| 809 | assert [torch.float16, torch.bfloat16] == QuickAllReduce._SUPPORTED_DTYPES |
| 810 | |
| 811 | |
| 812 | def test_quick_allreduce_min_size_table(): |
| 813 | from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce |
| 814 | |
| 815 | for dtype in [torch.float16, torch.bfloat16]: |
| 816 | for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES: |
| 817 | min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)] |
| 818 | assert len(min_sizes) == 5 |
| 819 | assert all(size > 0 for size in min_sizes) |
| 820 | |
| 821 | |
| 822 | def test_qr_max_size(): |
| 823 | from vllm import _custom_ops as ops |
| 824 | |
| 825 | max_size = ops.qr_max_size() |
| 826 | assert isinstance(max_size, int) |
| 827 | assert max_size > 0 |
| 828 | |
| 829 | |
| 830 | @pytest.mark.skipif( |
| 831 | current_platform.device_count() < WORLD_SIZE, |
| 832 | reason="requires 2 ROCm GPUs", |
| 833 | ) |
| 834 | @pytest.mark.parametrize("quant_level", QUANT_LEVELS) |
| 835 | def test_quick_allreduce_two_gpu_correctness(quant_level): |
| 836 | _log(f"two-GPU correctness case: quant={quant_level}") |
| 837 | _run_two_gpu_quick_allreduce_test( |
| 838 | quant_level=quant_level, |
| 839 | dtype_name="float16", |
| 840 | cast_bf16=False, |
| 841 | ) |
| 842 | |
| 843 | |
| 844 | @multi_gpu_test(num_gpus=CUDAGRAPH_WORLD_SIZE) |
| 845 | def test_quick_allreduce_cudagraph_replay(): |
| 846 | # Regression test for the stale flag_color bug: a quick-reduce captured in a |
| 847 | # CUDA graph must return correct results on every replay, not the data from |
| 848 | # a previous round. |
| 849 | _log("cudagraph replay case") |
| 850 | _run_cudagraph_replay_test( |
| 851 | world_size=CUDAGRAPH_WORLD_SIZE, |
| 852 | quant_level="FP", |
| 853 | ) |
| 854 | |
| 855 | |
| 856 | @pytest.mark.skipif( |
| 857 | current_platform.device_count() < WORLD_SIZE, |
| 858 | reason="requires 2 ROCm GPUs", |
| 859 | ) |
| 860 | def test_quick_allreduce_bf16_cast_mode(): |
| 861 | _log("BF16 cast case") |
| 862 | _run_two_gpu_quick_allreduce_test( |
| 863 | quant_level="FP", |
| 864 | dtype_name="bfloat16", |
| 865 | cast_bf16=True, |
| 866 | ) |
| 867 | |
| 868 | |
| 869 | @pytest.mark.skipif( |
| 870 | current_platform.device_count() < WORLD_SIZE, |
| 871 | reason="requires 2 ROCm GPUs", |
| 872 | ) |
| 873 | def test_quick_allreduce_llm_e2e(): |
| 874 | _log("LLM e2e case: backend=mp") |
| 875 | run_quick_reduce_llm_e2e(backend="mp") |
| 876 | |
| 877 | |
| 878 | @pytest.mark.skipif( |
| 879 | current_platform.device_count() < WORLD_SIZE, |
| 880 | reason="requires 2 ROCm GPUs", |
| 881 | ) |
| 882 | def test_quick_allreduce_llm_e2e_ray(): |
| 883 | _log("LLM e2e case: backend=ray") |
| 884 | run_quick_reduce_llm_e2e(backend="ray") |