# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib import importlib import multiprocessing as mp import os import queue import traceback from functools import lru_cache from types import SimpleNamespace from typing import Literal from unittest.mock import patch import pytest import torch import torch.distributed as dist from huggingface_hub import snapshot_download import vllm.envs as envs from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_port pytestmark = pytest.mark.skipif( not current_platform.is_rocm(), reason="ROCm-only quick-reduce tests", ) MB = 1024 * 1024 WORLD_SIZE = 2 QUANT_LEVELS = ["FP", "INT8", "INT6", "INT4"] def _log(message: str) -> None: print(f"[rocm_quick_reduce] {message}", flush=True) def _reload_envs(): return importlib.reload(envs) def _make_quick_allreduce( *, disabled: bool = False, world_size: int = 2, quant_level: str = "FP", use_fp16_kernels: bool = False, qr_max_size: int = 64 * MB, ): from vllm.distributed.device_communicators.quick_all_reduce import ( QuickAllReduce, QuickReduceRegime, ) qar = QuickAllReduce.__new__(QuickAllReduce) qar.disabled = disabled qar.world_size = world_size qar.use_fp16_kernels = use_fp16_kernels qar.qr_quant_level = QuickReduceRegime[quant_level] qar.qr_max_size = qr_max_size return qar def _quick_allreduce_worker( rank: int, port: int, quant_level: str, dtype_name: str, cast_bf16: bool, ): os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1" if cast_bf16 else "0" _log( f"worker start: rank={rank} quant={quant_level} " f"dtype={dtype_name} cast_bf16={cast_bf16}" ) device = torch.device(f"cuda:{rank}") torch.accelerator.set_device_index(device) dist.init_process_group( backend="gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=WORLD_SIZE, ) qar = None try: from vllm.distributed.device_communicators.quick_all_reduce import ( QuickAllReduce, ) qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank) assert not qar.disabled num_elements = 8 * MB if dtype_name == "float16" else 4 * MB dtype = getattr(torch, dtype_name) inp = torch.ones(num_elements, dtype=dtype, device=device) assert qar.should_quick_allreduce(inp) if cast_bf16: assert qar.use_fp16_kernels out = qar.quick_all_reduce(inp) assert torch.allclose(out, inp * WORLD_SIZE, atol=2.5, rtol=0.1) _log( f"worker complete: rank={rank} quant={quant_level} " f"dtype={dtype_name} num_elements={num_elements} " f"use_fp16_kernels={qar.use_fp16_kernels}" ) finally: if qar is not None: qar.close() if dist.is_initialized(): dist.destroy_process_group() def _run_two_gpu_quick_allreduce_test( *, quant_level: str, dtype_name: str, cast_bf16: bool, ): _log( f"launch 2-GPU case: quant={quant_level} " f"dtype={dtype_name} cast_bf16={cast_bf16}" ) ctx = mp.get_context("spawn") port = get_open_port() procs = [] for rank in range(WORLD_SIZE): proc = ctx.Process( target=_quick_allreduce_worker, args=(rank, port, quant_level, dtype_name, cast_bf16), ) proc.start() procs.append(proc) for proc in procs: proc.join(timeout=60) assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}" _log( f"finished 2-GPU case: quant={quant_level} " f"dtype={dtype_name} cast_bf16={cast_bf16}" ) MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" E2E_PREFILL_TOKENS = 1024 E2E_MAX_MODEL_LEN = 1536 E2E_GPU_MEMORY_UTILIZATION = 0.3 E2E_KV_CACHE_MEMORY_BYTES = 2 << 30 _BACKGROUND_LINE = ( "Background filler: this archived operations memo repeats a routine status " "line so the distributed test uses a realistically long prefill." ) _BACKGROUND_BLOCK = " ".join([_BACKGROUND_LINE] * 48) def _build_prompt(*, fact_block: str, question: str) -> str: return ( "Read the archived operations memo below. Most of the memo is filler. " "Use only the fact block near the end when answering.\n" f"{_BACKGROUND_BLOCK}\n" "Fact block:\n" f"{fact_block}\n" f"Question: {question}\n" "Answer in one short sentence." ) E2E_PROMPTS = [ _build_prompt( fact_block=( "- Festival city: Oslo\n- Mascot animal: otter\n- Welcome drink: tea" ), question="Which city hosts the festival, and what animal is the mascot?", ), _build_prompt( fact_block=( "- Meeting day: Tuesday\n" "- Planned snack: apricot cake\n" "- Backup room: Cedar" ), question="What day is the meeting, and what snack is planned?", ), ] RECORDED_RESPONSE_TEXTS = ( " The city hosting the festival is Oslo, and the mascot is an otter.", " The meeting is on Tuesday and the snack planned is apricot cake.", ) REQUIRED_WORDS = (("oslo", "otter"), ("tuesday", "apricot")) def _log_prompt_summaries() -> None: for i, prompt in enumerate(E2E_PROMPTS): prompt_lines = prompt.splitlines() fact_block = [line for line in prompt_lines if line.startswith("- ")] fact_summary = "; ".join(line.removeprefix("- ") for line in fact_block) _log(f"prompt {i} facts: {fact_summary}") @lru_cache(maxsize=1) def _get_model_path() -> str: try: path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True) _log(f"using cached model snapshot: {path}") return path except Exception: path = snapshot_download(repo_id=MODEL_NAME) _log(f"downloaded model snapshot: {path}") return path def _get_hidden_size(model_config) -> int: hidden_size = getattr(model_config, "hidden_size", None) if hidden_size is None and hasattr(model_config, "text_config"): hidden_size = getattr(model_config.text_config, "hidden_size", None) assert isinstance(hidden_size, int) return hidden_size def _check_tp_allreduce_uses_quick_reduce( self, num_tokens: int, dtype_name: str = "float16", ) -> dict[str, int | bool]: from vllm.distributed.communication_op import tensor_model_parallel_all_reduce from vllm.distributed.parallel_state import get_tp_group assert self.device is not None qr_comm = get_tp_group().device_communicator.qr_comm assert qr_comm is not None assert not qr_comm.disabled hidden_size = _get_hidden_size(self.model_runner.model.config) dtype = getattr(torch, dtype_name) sample = torch.full( (num_tokens, hidden_size), fill_value=float(self.rank + 1), dtype=dtype, device=self.device, ) assert qr_comm.should_quick_allreduce(sample) expected = sample.clone() reduced = tensor_model_parallel_all_reduce(sample) dist.all_reduce(expected, group=get_tp_group().device_group) torch.testing.assert_close(reduced, expected, atol=2.5, rtol=0.1) stats = { "rank": self.rank, "hidden_size": hidden_size, "num_tokens": num_tokens, "use_fp16_kernels": qr_comm.use_fp16_kernels, } _log( "worker quick-reduce check: " f"rank={self.rank} hidden_size={hidden_size} " f"num_tokens={num_tokens} use_fp16_kernels={qr_comm.use_fp16_kernels}" ) return stats def _check_quick_reduce_disabled(self) -> int: from vllm.distributed.parallel_state import get_tp_group qr_comm = get_tp_group().device_communicator.qr_comm assert qr_comm is not None assert qr_comm.disabled _log(f"worker confirmed quick reduce is disabled: rank={self.rank}") return self.rank def _collect_generations(outputs) -> list[tuple[tuple[int, ...], str]]: return [ (tuple(output.outputs[0].token_ids), output.outputs[0].text) for output in outputs ] def _shutdown_llm(llm: LLM | None) -> None: if llm is None: cleanup_dist_env_and_memory() return with contextlib.suppress(Exception): llm.llm_engine.engine_core.shutdown() del llm cleanup_dist_env_and_memory() def _log_generations( label: str, generations: list[tuple[tuple[int, ...], str]], ) -> None: for i, (token_ids, text) in enumerate(generations): _log(f"{label} prompt {i} token ids: {list(token_ids)}") _log(f"{label} prompt {i} text: {text!r}") def _assert_required_words( label: str, generations: list[tuple[tuple[int, ...], str]], ) -> None: for i, (_, text) in enumerate(generations): lowered = text.lower() missing = [word for word in REQUIRED_WORDS[i] if word not in lowered] assert not missing, ( f"{label} prompt {i} is missing required words {missing}. " f"Observed text: {text!r}" ) def _collect_soft_mismatches( baseline_generations: list[tuple[tuple[int, ...], str]], quick_reduce_generations: list[tuple[tuple[int, ...], str]], ) -> list[str]: mismatches = [] for i, (_, text) in enumerate(baseline_generations): expected = RECORDED_RESPONSE_TEXTS[i] if text != expected: mismatches.append( f"baseline prompt {i} drifted from the recorded response.\n" f"expected={expected!r}\nactual={text!r}" ) for i, (_, text) in enumerate(quick_reduce_generations): expected = RECORDED_RESPONSE_TEXTS[i] if text != expected: mismatches.append( f"quick-reduce prompt {i} drifted from the recorded response.\n" f"expected={expected!r}\nactual={text!r}" ) for i, ((_, baseline_text), (_, quick_reduce_text)) in enumerate( zip(baseline_generations, quick_reduce_generations) ): if baseline_text != quick_reduce_text: mismatches.append( f"baseline and quick-reduce responses differ for prompt {i}.\n" f"baseline={baseline_text!r}\nquick_reduce={quick_reduce_text!r}" ) return mismatches def _run_generation( *, backend: Literal["mp", "ray"], quant_mode: str, expect_quick_reduce: bool, ) -> list[tuple[tuple[int, ...], str]]: llm = None monkeypatch = pytest.MonkeyPatch() with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode) model_path = _get_model_path() _log( f"starting generation: backend={backend} quant={quant_mode} " f"gpu_memory_utilization={E2E_GPU_MEMORY_UTILIZATION} " f"kv_cache_bytes={E2E_KV_CACHE_MEMORY_BYTES} model={model_path}" ) try: llm = LLM( model=model_path, tokenizer=model_path, tensor_parallel_size=2, distributed_executor_backend=backend, dtype="half", enforce_eager=True, max_model_len=E2E_MAX_MODEL_LEN, max_num_seqs=len(E2E_PROMPTS), gpu_memory_utilization=E2E_GPU_MEMORY_UTILIZATION, kv_cache_memory_bytes=E2E_KV_CACHE_MEMORY_BYTES, seed=0, ) if not expect_quick_reduce: assert llm.collective_rpc(_check_quick_reduce_disabled) == [0, 1] if expect_quick_reduce: worker_stats = llm.collective_rpc( _check_tp_allreduce_uses_quick_reduce, args=(E2E_PREFILL_TOKENS,), ) assert [stat["rank"] for stat in worker_stats] == [0, 1] worker_summary = "; ".join( "rank={rank} hidden_size={hidden_size} num_tokens={num_tokens} " "use_fp16_kernels={use_fp16_kernels}".format(**stat) for stat in worker_stats ) _log(f"{backend} quick-reduce worker checks: {worker_summary}") outputs = llm.generate( E2E_PROMPTS, SamplingParams( temperature=0.0, max_tokens=20, stop=["\nAnswer:", " Answer:"], ), use_tqdm=False, ) generations = _collect_generations(outputs) assert all(text.strip() for _, text in generations) _log_generations(f"{backend} {quant_mode}", generations) return generations finally: _shutdown_llm(llm) def _run_quick_reduce_llm_e2e_in_subprocess( *, backend: Literal["mp", "ray"], ) -> str | None: _log(f"running LLM e2e: backend={backend}") _log_prompt_summaries() baseline_outputs = _run_generation( backend=backend, quant_mode="NONE", expect_quick_reduce=False, ) quick_reduce_outputs = _run_generation( backend=backend, quant_mode="FP", expect_quick_reduce=True, ) _assert_required_words("baseline", baseline_outputs) _assert_required_words("quick-reduce", quick_reduce_outputs) mismatches = _collect_soft_mismatches(baseline_outputs, quick_reduce_outputs) if mismatches: details = "\n\n".join(mismatches) _log(f"soft response mismatch:\n{details}") return details _log(f"LLM e2e backend={backend} matched the recorded responses exactly") return None def _quick_reduce_llm_e2e_worker( result_queue: mp.Queue, backend: Literal["mp", "ray"], ) -> None: try: xfail_reason = _run_quick_reduce_llm_e2e_in_subprocess(backend=backend) except Exception: result_queue.put({"status": "error", "reason": traceback.format_exc()}) raise else: if xfail_reason is not None: result_queue.put({"status": "xfail", "reason": xfail_reason}) else: result_queue.put({"status": "ok"}) def run_quick_reduce_llm_e2e( *, backend: Literal["mp", "ray"], ) -> None: ctx = mp.get_context("spawn") result_queue = ctx.Queue() proc = ctx.Process( target=_quick_reduce_llm_e2e_worker, args=(result_queue, backend), ) proc.start() proc.join(timeout=600) try: result = result_queue.get(timeout=5) except queue.Empty as exc: if proc.exitcode != 0: raise AssertionError( f"quick-reduce llm e2e subprocess failed for backend={backend} " f"with exit code {proc.exitcode} and produced no result" ) from exc raise AssertionError( f"quick-reduce llm e2e subprocess produced no result for backend={backend}" ) from exc if result["status"] == "xfail": pytest.xfail(result["reason"]) if result["status"] == "error": raise AssertionError( f"quick-reduce llm e2e subprocess failed for backend={backend}:\n" f"{result['reason']}" ) assert proc.exitcode == 0, ( f"quick-reduce llm e2e subprocess failed for backend={backend} " f"with exit code {proc.exitcode}" ) def test_quick_reduce_regime_values(): from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime assert QuickReduceRegime.FP.value == 0 assert QuickReduceRegime.INT8.value == 1 assert QuickReduceRegime.INT6.value == 2 assert QuickReduceRegime.INT4.value == 3 assert QuickReduceRegime.NONE.value == 4 def test_quick_reduce_regime_names(): from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"} @pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"]) def test_quick_reduce_quantization_env_var(monkeypatch, quant_level): monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_level) reloaded_envs = _reload_envs() assert quant_level == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION def test_quick_reduce_quantization_default(monkeypatch): monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", raising=False) reloaded_envs = _reload_envs() assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION == "NONE" @pytest.mark.parametrize("cast_bf16", [True, False]) def test_quick_reduce_cast_bf16_to_fp16_env_var(monkeypatch, cast_bf16): monkeypatch.setenv( "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1" if cast_bf16 else "0" ) reloaded_envs = _reload_envs() assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is cast_bf16 def test_quick_reduce_cast_bf16_to_fp16_default(monkeypatch): monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", raising=False) reloaded_envs = _reload_envs() assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is True @pytest.mark.parametrize("max_mb", [128, 512, 2048, None]) def test_quick_reduce_max_size_env_var(monkeypatch, max_mb): if max_mb is None: monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False) else: monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", str(max_mb)) reloaded_envs = _reload_envs() assert max_mb == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB def test_quick_reduce_max_size_default(monkeypatch): monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False) reloaded_envs = _reload_envs() assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is None @pytest.mark.parametrize( ("gcn_arch_name", "expected"), [ ("gfx942", True), ("gfx950", True), ("gfx90a", False), ("", False), ], ) def test_quick_allreduce_rocm_arch_available(gcn_arch_name, expected): from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce qar = QuickAllReduce.__new__(QuickAllReduce) qar.disabled = True with ( patch( "vllm.distributed.device_communicators.quick_all_reduce.current_platform." "is_rocm", return_value=True, ), patch( "torch.cuda.get_device_properties", return_value=SimpleNamespace(gcnArchName=gcn_arch_name), ), ): assert qar._rocm_arch_available() is expected def test_quick_allreduce_rocm_arch_available_handles_probe_failure(): from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce qar = QuickAllReduce.__new__(QuickAllReduce) qar.disabled = True with ( patch( "vllm.distributed.device_communicators.quick_all_reduce.current_platform." "is_rocm", return_value=True, ), patch("torch.cuda.get_device_properties", side_effect=RuntimeError), ): assert qar._rocm_arch_available() is False def test_quick_allreduce_rejects_disabled(): qar = _make_quick_allreduce(disabled=True) inp = torch.zeros(1024, dtype=torch.float16) assert qar.should_quick_allreduce(inp) is False def test_quick_allreduce_rejects_unsupported_dtype(): qar = _make_quick_allreduce() inp = torch.zeros(1024 * 1024, dtype=torch.float32) assert qar.should_quick_allreduce(inp) is False def test_quick_allreduce_rejects_non_aligned_input(): qar = _make_quick_allreduce() inp = torch.zeros(5, dtype=torch.float16) assert qar.should_quick_allreduce(inp) is False def test_quick_allreduce_rejects_non_contiguous_input(): qar = _make_quick_allreduce() inp = torch.zeros((1024, 1024), dtype=torch.float16)[:, ::2] assert qar.should_quick_allreduce(inp) is False def test_quick_allreduce_rejects_input_smaller_than_threshold(): qar = _make_quick_allreduce() inp = torch.zeros((MB // 2) - 8, dtype=torch.float16) assert qar.should_quick_allreduce(inp) is False def test_quick_allreduce_accepts_input_at_threshold(): qar = _make_quick_allreduce() inp = torch.zeros(MB // 2, dtype=torch.float16) assert qar.should_quick_allreduce(inp) is True def test_quick_allreduce_rejects_input_larger_than_max_size(): qar = _make_quick_allreduce(qr_max_size=1 * MB) inp = torch.zeros(MB, dtype=torch.float16) assert qar.should_quick_allreduce(inp) is False def test_quick_allreduce_bf16_uses_fp16_threshold_when_cast_enabled(): inp = torch.zeros(MB // 2, dtype=torch.bfloat16) without_cast = _make_quick_allreduce(use_fp16_kernels=False) with_cast = _make_quick_allreduce(use_fp16_kernels=True) assert without_cast.should_quick_allreduce(inp) is False assert with_cast.should_quick_allreduce(inp) is True def test_quick_allreduce_supported_world_sizes(): from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce assert QuickAllReduce._SUPPORTED_WORLD_SIZES == [2, 4, 8] def test_quick_allreduce_supported_dtypes(): from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce assert [torch.float16, torch.bfloat16] == QuickAllReduce._SUPPORTED_DTYPES def test_quick_allreduce_min_size_table(): from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce for dtype in [torch.float16, torch.bfloat16]: for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES: min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)] assert len(min_sizes) == 4 assert all(size > 0 for size in min_sizes) def test_qr_max_size(): from vllm import _custom_ops as ops max_size = ops.qr_max_size() assert isinstance(max_size, int) assert max_size > 0 @pytest.mark.skipif( current_platform.device_count() < WORLD_SIZE, reason="requires 2 ROCm GPUs", ) @pytest.mark.parametrize("quant_level", QUANT_LEVELS) def test_quick_allreduce_two_gpu_correctness(quant_level): _log(f"two-GPU correctness case: quant={quant_level}") _run_two_gpu_quick_allreduce_test( quant_level=quant_level, dtype_name="float16", cast_bf16=False, ) @pytest.mark.skipif( current_platform.device_count() < WORLD_SIZE, reason="requires 2 ROCm GPUs", ) def test_quick_allreduce_bf16_cast_mode(): _log("BF16 cast case") _run_two_gpu_quick_allreduce_test( quant_level="FP", dtype_name="bfloat16", cast_bf16=True, ) @pytest.mark.skipif( current_platform.device_count() < WORLD_SIZE, reason="requires 2 ROCm GPUs", ) def test_quick_allreduce_llm_e2e(): _log("LLM e2e case: backend=mp") run_quick_reduce_llm_e2e(backend="mp") @pytest.mark.skipif( current_platform.device_count() < WORLD_SIZE, reason="requires 2 ROCm GPUs", ) def test_quick_allreduce_llm_e2e_ray(): _log("LLM e2e case: backend=ray") run_quick_reduce_llm_e2e(backend="ray")