@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 445ded18c1
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 19:30 UTC

[ROCm][CI] Extend ROCm quick reduce coverage (#40990)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed May 26, 2026 at 08:57 UTC 445ded18c1184a5a44d0f41010d614adbd107ca7
4 files changed +869 -38
.buildkite/test-amd.yaml
+20 -4
@@ -2703,19 +2703,35 @@ steps:
2703 optional: true
2704 working_dir: "/vllm-workspace/"
2705 source_file_dependencies:
2706 + - csrc/custom_quickreduce.cu
2707 + - csrc/ops.h
2708 + - csrc/torch_bindings.cpp
2709 - vllm/distributed/
2707 - - vllm/v1/distributed/
2710 + - vllm/model_executor/layers/
2711 + - vllm/entrypoints/llm.py
2712 + - vllm/config/parallel.py
2713 - vllm/model_executor/layers/fused_moe/
2714 + - vllm/v1/engine/
2715 + - vllm/v1/executor/
2716 + - vllm/v1/worker/
2717 + - vllm/v1/distributed/
2718 - vllm/v1/attention/backends/
2719 - vllm/v1/attention/selector.py
2711 - - tests/distributed/test_context_parallel.py
2712 - - tests/v1/distributed/test_dbo.py
2713 - - examples/features/data_parallel/data_parallel_offline.py
2720 - vllm/_aiter_ops.py
2721 + - vllm/_custom_ops.py
2722 - vllm/platforms/rocm.py
2723 + - vllm/envs.py
2724 + - examples/offline_inference/data_parallel.py
2725 + - tests/distributed/test_context_parallel.py
2726 + - tests/distributed/test_rocm_quick_reduce.py
2727 + - tests/distributed/test_quick_all_reduce.py
2728 + - tests/v1/distributed/test_dbo.py
2729 + - tests/utils.py
2730 commands:
2731 - pytest -v -s tests/distributed/test_context_parallel.py
2732 - pytest -v -s tests/v1/distributed/test_dbo.py
2733 + - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
2734 + - pytest -v -s tests/distributed/test_quick_all_reduce.py
2735
2736 #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
2737
tests/distributed/test_quick_all_reduce.py
+73 -3
@@ -25,16 +25,34 @@ from ..utils import (
25 ensure_model_parallel_initialized,
26 init_test_distributed_environment,
27 multi_process_parallel,
28 + set_random_seed,
29 )
30
30 -torch.manual_seed(42)
31 -random.seed(44)
31 +
32 +def on_gfx942() -> bool:
33 + if current_platform.is_rocm():
34 + from vllm.platforms.rocm import on_gfx942 as rocm_on_gfx942
35 +
36 + return rocm_on_gfx942()
37 + return False
38 +
39 +
40 +set_random_seed(42)
41 +_test_size_rng = random.Random(44)
42 # Size over 8MB is sufficient for custom quick allreduce.
33 -test_sizes = [random.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)]
43 +test_sizes = [
44 + _test_size_rng.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)
45 +]
46 for i, v in enumerate(test_sizes):
47 test_sizes[i] -= v % 8
48
49
50 +def _assert_quickreduce(fa, inp):
51 + assert fa is not None
52 + assert not fa.disabled
53 + assert fa.should_quick_allreduce(inp)
54 +
55 +
56 @pytest.fixture
57 def envs_cache_disabled():
58 disable_envs_cache()
@@ -216,11 +234,14 @@ def graph_quickreduce(
234 ):
235 with monkeypatch.context() as m:
236 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
237 + m.delenv("HIP_VISIBLE_DEVICES", raising=False)
238 + m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
239 device = torch.device(f"cuda:{rank}")
240 torch.accelerator.set_device_index(device)
241 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
242 ensure_model_parallel_initialized(tp_size, pp_size)
243 group = get_tp_group().device_group
244 + fa = get_tp_group().device_communicator.qr_comm
245
246 # A small all_reduce for warmup.
247 # this is needed because device communicators might be created lazily
@@ -246,6 +267,8 @@ def graph_quickreduce(
267 device_idx = torch.accelerator.current_device_index()
268 inp1 = torch.randint(1, 23, (sz,), dtype=dtype, device=device_idx)
269 inp2 = torch.randint(-23, 1, (sz,), dtype=dtype, device=device_idx)
270 + _assert_quickreduce(fa, inp1)
271 + _assert_quickreduce(fa, inp2)
272
273 torch.accelerator.synchronize()
274 graph = torch.cuda.CUDAGraph()
@@ -270,6 +293,8 @@ def eager_quickreduce(
293 ):
294 with monkeypatch.context() as m:
295 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
296 + m.delenv("HIP_VISIBLE_DEVICES", raising=False)
297 + m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
298 device = torch.device(f"cuda:{rank}")
299 torch.accelerator.set_device_index(device)
300
@@ -281,12 +306,42 @@ def eager_quickreduce(
306 inp = torch.tensor(
307 [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.float16, device=device
308 )
309 + _assert_quickreduce(fa, inp)
310 out = fa.quick_all_reduce(inp)
311 torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
312
313 inp = torch.tensor(
314 [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
315 )
316 + _assert_quickreduce(fa, inp)
317 + out = fa.quick_all_reduce(inp)
318 + torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
319 +
320 +
321 +@ray.remote(num_gpus=1, max_calls=1)
322 +def bf16_cast_quickreduce(
323 + monkeypatch: pytest.MonkeyPatch,
324 + tp_size,
325 + pp_size,
326 + rank,
327 + distributed_init_port,
328 +):
329 + with monkeypatch.context() as m:
330 + m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
331 + m.delenv("HIP_VISIBLE_DEVICES", raising=False)
332 + m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
333 + m.setenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1")
334 + device = torch.device(f"cuda:{rank}")
335 + torch.accelerator.set_device_index(device)
336 + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
337 +
338 + sz = 16 * 1024 * 1024
339 + fa = get_tp_group().device_communicator.qr_comm
340 + inp = torch.tensor(
341 + [1.0 * (i % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
342 + )
343 + _assert_quickreduce(fa, inp)
344 + assert fa.use_fp16_kernels
345 out = fa.quick_all_reduce(inp)
346 torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
347
@@ -308,12 +363,27 @@ def test_custom_quick_allreduce(
363 world_size = tp_size * pipeline_parallel_size
364 if world_size > torch.accelerator.device_count():
365 pytest.skip("Not enough GPUs to run the test.")
366 + if test_target is graph_quickreduce and on_gfx942():
367 + pytest.xfail(
368 + "CUDA graph capture with quick reduce hits "
369 + "hipErrorStreamCaptureInvalidated on gfx942"
370 + )
371
372 monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
373
374 multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
375
376
377 +@pytest.mark.skipif(
378 + not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
379 +)
380 +def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch):
381 + if torch.accelerator.device_count() < 2:
382 + pytest.skip("Not enough GPUs to run the test.")
383 + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "FP")
384 + multi_process_parallel(monkeypatch, 2, 1, bf16_cast_quickreduce)
385 +
386 +
387 def qr_variable_input(rank, world_size):
388 """
389 When the tensor parallelism is set to 4 or 8, frequent changes
tests/distributed/test_rocm_quick_reduce.py new
+750
@@ -0,0 +1,750 @@
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 +from huggingface_hub import snapshot_download
19 +
20 +import vllm.envs as envs
21 +from vllm import LLM, SamplingParams
22 +from vllm.distributed import cleanup_dist_env_and_memory
23 +from vllm.platforms import current_platform
24 +from vllm.utils.network_utils import get_open_port
25 +
26 +pytestmark = pytest.mark.skipif(
27 + not current_platform.is_rocm(),
28 + reason="ROCm-only quick-reduce tests",
29 +)
30 +
31 +MB = 1024 * 1024
32 +WORLD_SIZE = 2
33 +QUANT_LEVELS = ["FP", "INT8", "INT6", "INT4"]
34 +
35 +
36 +def _log(message: str) -> None:
37 + print(f"[rocm_quick_reduce] {message}", flush=True)
38 +
39 +
40 +def _reload_envs():
41 + return importlib.reload(envs)
42 +
43 +
44 +def _make_quick_allreduce(
45 + *,
46 + disabled: bool = False,
47 + world_size: int = 2,
48 + quant_level: str = "FP",
49 + use_fp16_kernels: bool = False,
50 + qr_max_size: int = 64 * MB,
51 +):
52 + from vllm.distributed.device_communicators.quick_all_reduce import (
53 + QuickAllReduce,
54 + QuickReduceRegime,
55 + )
56 +
57 + qar = QuickAllReduce.__new__(QuickAllReduce)
58 + qar.disabled = disabled
59 + qar.world_size = world_size
60 + qar.use_fp16_kernels = use_fp16_kernels
61 + qar.qr_quant_level = QuickReduceRegime[quant_level]
62 + qar.qr_max_size = qr_max_size
63 + return qar
64 +
65 +
66 +def _quick_allreduce_worker(
67 + rank: int,
68 + port: int,
69 + quant_level: str,
70 + dtype_name: str,
71 + cast_bf16: bool,
72 +):
73 + os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level
74 + os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1" if cast_bf16 else "0"
75 + _log(
76 + f"worker start: rank={rank} quant={quant_level} "
77 + f"dtype={dtype_name} cast_bf16={cast_bf16}"
78 + )
79 +
80 + device = torch.device(f"cuda:{rank}")
81 + torch.accelerator.set_device_index(device)
82 + dist.init_process_group(
83 + backend="gloo",
84 + init_method=f"tcp://127.0.0.1:{port}",
85 + rank=rank,
86 + world_size=WORLD_SIZE,
87 + )
88 +
89 + qar = None
90 + try:
91 + from vllm.distributed.device_communicators.quick_all_reduce import (
92 + QuickAllReduce,
93 + )
94 +
95 + qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank)
96 + assert not qar.disabled
97 +
98 + num_elements = 8 * MB if dtype_name == "float16" else 4 * MB
99 +
100 + dtype = getattr(torch, dtype_name)
101 + inp = torch.ones(num_elements, dtype=dtype, device=device)
102 +
103 + assert qar.should_quick_allreduce(inp)
104 + if cast_bf16:
105 + assert qar.use_fp16_kernels
106 +
107 + out = qar.quick_all_reduce(inp)
108 + assert torch.allclose(out, inp * WORLD_SIZE, atol=2.5, rtol=0.1)
109 + _log(
110 + f"worker complete: rank={rank} quant={quant_level} "
111 + f"dtype={dtype_name} num_elements={num_elements} "
112 + f"use_fp16_kernels={qar.use_fp16_kernels}"
113 + )
114 + finally:
115 + if qar is not None:
116 + qar.close()
117 + if dist.is_initialized():
118 + dist.destroy_process_group()
119 +
120 +
121 +def _run_two_gpu_quick_allreduce_test(
122 + *,
123 + quant_level: str,
124 + dtype_name: str,
125 + cast_bf16: bool,
126 +):
127 + _log(
128 + f"launch 2-GPU case: quant={quant_level} "
129 + f"dtype={dtype_name} cast_bf16={cast_bf16}"
130 + )
131 + ctx = mp.get_context("spawn")
132 + port = get_open_port()
133 + procs = []
134 +
135 + for rank in range(WORLD_SIZE):
136 + proc = ctx.Process(
137 + target=_quick_allreduce_worker,
138 + args=(rank, port, quant_level, dtype_name, cast_bf16),
139 + )
140 + proc.start()
141 + procs.append(proc)
142 +
143 + for proc in procs:
144 + proc.join(timeout=60)
145 + assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}"
146 + _log(
147 + f"finished 2-GPU case: quant={quant_level} "
148 + f"dtype={dtype_name} cast_bf16={cast_bf16}"
149 + )
150 +
151 +
152 +MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
153 +E2E_PREFILL_TOKENS = 1024
154 +E2E_MAX_MODEL_LEN = 1536
155 +E2E_GPU_MEMORY_UTILIZATION = 0.3
156 +E2E_KV_CACHE_MEMORY_BYTES = 2 << 30
157 +
158 +_BACKGROUND_LINE = (
159 + "Background filler: this archived operations memo repeats a routine status "
160 + "line so the distributed test uses a realistically long prefill."
161 +)
162 +_BACKGROUND_BLOCK = " ".join([_BACKGROUND_LINE] * 48)
163 +
164 +
165 +def _build_prompt(*, fact_block: str, question: str) -> str:
166 + return (
167 + "Read the archived operations memo below. Most of the memo is filler. "
168 + "Use only the fact block near the end when answering.\n"
169 + f"{_BACKGROUND_BLOCK}\n"
170 + "Fact block:\n"
171 + f"{fact_block}\n"
172 + f"Question: {question}\n"
173 + "Answer in one short sentence."
174 + )
175 +
176 +
177 +E2E_PROMPTS = [
178 + _build_prompt(
179 + fact_block=(
180 + "- Festival city: Oslo\n- Mascot animal: otter\n- Welcome drink: tea"
181 + ),
182 + question="Which city hosts the festival, and what animal is the mascot?",
183 + ),
184 + _build_prompt(
185 + fact_block=(
186 + "- Meeting day: Tuesday\n"
187 + "- Planned snack: apricot cake\n"
188 + "- Backup room: Cedar"
189 + ),
190 + question="What day is the meeting, and what snack is planned?",
191 + ),
192 +]
193 +RECORDED_RESPONSE_TEXTS = (
194 + " The city hosting the festival is Oslo, and the mascot is an otter.",
195 + " The meeting is on Tuesday and the snack planned is apricot cake.",
196 +)
197 +REQUIRED_WORDS = (("oslo", "otter"), ("tuesday", "apricot"))
198 +
199 +
200 +def _log_prompt_summaries() -> None:
201 + for i, prompt in enumerate(E2E_PROMPTS):
202 + prompt_lines = prompt.splitlines()
203 + fact_block = [line for line in prompt_lines if line.startswith("- ")]
204 + fact_summary = "; ".join(line.removeprefix("- ") for line in fact_block)
205 + _log(f"prompt {i} facts: {fact_summary}")
206 +
207 +
208 +@lru_cache(maxsize=1)
209 +def _get_model_path() -> str:
210 + try:
211 + path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True)
212 + _log(f"using cached model snapshot: {path}")
213 + return path
214 + except Exception:
215 + path = snapshot_download(repo_id=MODEL_NAME)
216 + _log(f"downloaded model snapshot: {path}")
217 + return path
218 +
219 +
220 +def _get_hidden_size(model_config) -> int:
221 + hidden_size = getattr(model_config, "hidden_size", None)
222 + if hidden_size is None and hasattr(model_config, "text_config"):
223 + hidden_size = getattr(model_config.text_config, "hidden_size", None)
224 + assert isinstance(hidden_size, int)
225 + return hidden_size
226 +
227 +
228 +def _check_tp_allreduce_uses_quick_reduce(
229 + self,
230 + num_tokens: int,
231 + dtype_name: str = "float16",
232 +) -> dict[str, int | bool]:
233 + from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
234 + from vllm.distributed.parallel_state import get_tp_group
235 +
236 + assert self.device is not None
237 + qr_comm = get_tp_group().device_communicator.qr_comm
238 + assert qr_comm is not None
239 + assert not qr_comm.disabled
240 +
241 + hidden_size = _get_hidden_size(self.model_runner.model.config)
242 + dtype = getattr(torch, dtype_name)
243 + sample = torch.full(
244 + (num_tokens, hidden_size),
245 + fill_value=float(self.rank + 1),
246 + dtype=dtype,
247 + device=self.device,
248 + )
249 +
250 + assert qr_comm.should_quick_allreduce(sample)
251 +
252 + expected = sample.clone()
253 + reduced = tensor_model_parallel_all_reduce(sample)
254 + dist.all_reduce(expected, group=get_tp_group().device_group)
255 + torch.testing.assert_close(reduced, expected, atol=2.5, rtol=0.1)
256 +
257 + stats = {
258 + "rank": self.rank,
259 + "hidden_size": hidden_size,
260 + "num_tokens": num_tokens,
261 + "use_fp16_kernels": qr_comm.use_fp16_kernels,
262 + }
263 + _log(
264 + "worker quick-reduce check: "
265 + f"rank={self.rank} hidden_size={hidden_size} "
266 + f"num_tokens={num_tokens} use_fp16_kernels={qr_comm.use_fp16_kernels}"
267 + )
268 + return stats
269 +
270 +
271 +def _check_quick_reduce_disabled(self) -> int:
272 + from vllm.distributed.parallel_state import get_tp_group
273 +
274 + qr_comm = get_tp_group().device_communicator.qr_comm
275 + assert qr_comm is not None
276 + assert qr_comm.disabled
277 + _log(f"worker confirmed quick reduce is disabled: rank={self.rank}")
278 + return self.rank
279 +
280 +
281 +def _collect_generations(outputs) -> list[tuple[tuple[int, ...], str]]:
282 + return [
283 + (tuple(output.outputs[0].token_ids), output.outputs[0].text)
284 + for output in outputs
285 + ]
286 +
287 +
288 +def _shutdown_llm(llm: LLM | None) -> None:
289 + if llm is None:
290 + cleanup_dist_env_and_memory()
291 + return
292 +
293 + with contextlib.suppress(Exception):
294 + llm.llm_engine.engine_core.shutdown()
295 +
296 + del llm
297 + cleanup_dist_env_and_memory()
298 +
299 +
300 +def _log_generations(
301 + label: str,
302 + generations: list[tuple[tuple[int, ...], str]],
303 +) -> None:
304 + for i, (token_ids, text) in enumerate(generations):
305 + _log(f"{label} prompt {i} token ids: {list(token_ids)}")
306 + _log(f"{label} prompt {i} text: {text!r}")
307 +
308 +
309 +def _assert_required_words(
310 + label: str,
311 + generations: list[tuple[tuple[int, ...], str]],
312 +) -> None:
313 + for i, (_, text) in enumerate(generations):
314 + lowered = text.lower()
315 + missing = [word for word in REQUIRED_WORDS[i] if word not in lowered]
316 + assert not missing, (
317 + f"{label} prompt {i} is missing required words {missing}. "
318 + f"Observed text: {text!r}"
319 + )
320 +
321 +
322 +def _collect_soft_mismatches(
323 + baseline_generations: list[tuple[tuple[int, ...], str]],
324 + quick_reduce_generations: list[tuple[tuple[int, ...], str]],
325 +) -> list[str]:
326 + mismatches = []
327 +
328 + for i, (_, text) in enumerate(baseline_generations):
329 + expected = RECORDED_RESPONSE_TEXTS[i]
330 + if text != expected:
331 + mismatches.append(
332 + f"baseline prompt {i} drifted from the recorded response.\n"
333 + f"expected={expected!r}\nactual={text!r}"
334 + )
335 +
336 + for i, (_, text) in enumerate(quick_reduce_generations):
337 + expected = RECORDED_RESPONSE_TEXTS[i]
338 + if text != expected:
339 + mismatches.append(
340 + f"quick-reduce prompt {i} drifted from the recorded response.\n"
341 + f"expected={expected!r}\nactual={text!r}"
342 + )
343 +
344 + for i, ((_, baseline_text), (_, quick_reduce_text)) in enumerate(
345 + zip(baseline_generations, quick_reduce_generations)
346 + ):
347 + if baseline_text != quick_reduce_text:
348 + mismatches.append(
349 + f"baseline and quick-reduce responses differ for prompt {i}.\n"
350 + f"baseline={baseline_text!r}\nquick_reduce={quick_reduce_text!r}"
351 + )
352 +
353 + return mismatches
354 +
355 +
356 +def _run_generation(
357 + *,
358 + backend: Literal["mp", "ray"],
359 + quant_mode: str,
360 + expect_quick_reduce: bool,
361 +) -> list[tuple[tuple[int, ...], str]]:
362 + llm = None
363 + monkeypatch = pytest.MonkeyPatch()
364 + with monkeypatch.context() as m:
365 + m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
366 + m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
367 + model_path = _get_model_path()
368 + _log(
369 + f"starting generation: backend={backend} quant={quant_mode} "
370 + f"gpu_memory_utilization={E2E_GPU_MEMORY_UTILIZATION} "
371 + f"kv_cache_bytes={E2E_KV_CACHE_MEMORY_BYTES} model={model_path}"
372 + )
373 +
374 + try:
375 + llm = LLM(
376 + model=model_path,
377 + tokenizer=model_path,
378 + tensor_parallel_size=2,
379 + distributed_executor_backend=backend,
380 + dtype="half",
381 + enforce_eager=True,
382 + max_model_len=E2E_MAX_MODEL_LEN,
383 + max_num_seqs=len(E2E_PROMPTS),
384 + gpu_memory_utilization=E2E_GPU_MEMORY_UTILIZATION,
385 + kv_cache_memory_bytes=E2E_KV_CACHE_MEMORY_BYTES,
386 + seed=0,
387 + )
388 +
389 + if not expect_quick_reduce:
390 + assert llm.collective_rpc(_check_quick_reduce_disabled) == [0, 1]
391 +
392 + if expect_quick_reduce:
393 + worker_stats = llm.collective_rpc(
394 + _check_tp_allreduce_uses_quick_reduce,
395 + args=(E2E_PREFILL_TOKENS,),
396 + )
397 + assert [stat["rank"] for stat in worker_stats] == [0, 1]
398 + worker_summary = "; ".join(
399 + "rank={rank} hidden_size={hidden_size} num_tokens={num_tokens} "
400 + "use_fp16_kernels={use_fp16_kernels}".format(**stat)
401 + for stat in worker_stats
402 + )
403 + _log(f"{backend} quick-reduce worker checks: {worker_summary}")
404 +
405 + outputs = llm.generate(
406 + E2E_PROMPTS,
407 + SamplingParams(
408 + temperature=0.0,
409 + max_tokens=20,
410 + stop=["\nAnswer:", " Answer:"],
411 + ),
412 + use_tqdm=False,
413 + )
414 + generations = _collect_generations(outputs)
415 + assert all(text.strip() for _, text in generations)
416 + _log_generations(f"{backend} {quant_mode}", generations)
417 + return generations
418 + finally:
419 + _shutdown_llm(llm)
420 +
421 +
422 +def _run_quick_reduce_llm_e2e_in_subprocess(
423 + *,
424 + backend: Literal["mp", "ray"],
425 +) -> str | None:
426 + _log(f"running LLM e2e: backend={backend}")
427 + _log_prompt_summaries()
428 + baseline_outputs = _run_generation(
429 + backend=backend,
430 + quant_mode="NONE",
431 + expect_quick_reduce=False,
432 + )
433 + quick_reduce_outputs = _run_generation(
434 + backend=backend,
435 + quant_mode="FP",
436 + expect_quick_reduce=True,
437 + )
438 +
439 + _assert_required_words("baseline", baseline_outputs)
440 + _assert_required_words("quick-reduce", quick_reduce_outputs)
441 +
442 + mismatches = _collect_soft_mismatches(baseline_outputs, quick_reduce_outputs)
443 + if mismatches:
444 + details = "\n\n".join(mismatches)
445 + _log(f"soft response mismatch:\n{details}")
446 + return details
447 +
448 + _log(f"LLM e2e backend={backend} matched the recorded responses exactly")
449 + return None
450 +
451 +
452 +def _quick_reduce_llm_e2e_worker(
453 + result_queue: mp.Queue,
454 + backend: Literal["mp", "ray"],
455 +) -> None:
456 + try:
457 + xfail_reason = _run_quick_reduce_llm_e2e_in_subprocess(backend=backend)
458 + except Exception:
459 + result_queue.put({"status": "error", "reason": traceback.format_exc()})
460 + raise
461 + else:
462 + if xfail_reason is not None:
463 + result_queue.put({"status": "xfail", "reason": xfail_reason})
464 + else:
465 + result_queue.put({"status": "ok"})
466 +
467 +
468 +def run_quick_reduce_llm_e2e(
469 + *,
470 + backend: Literal["mp", "ray"],
471 +) -> None:
472 + ctx = mp.get_context("spawn")
473 + result_queue = ctx.Queue()
474 + proc = ctx.Process(
475 + target=_quick_reduce_llm_e2e_worker,
476 + args=(result_queue, backend),
477 + )
478 + proc.start()
479 + proc.join(timeout=600)
480 +
481 + try:
482 + result = result_queue.get(timeout=5)
483 + except queue.Empty as exc:
484 + if proc.exitcode != 0:
485 + raise AssertionError(
486 + f"quick-reduce llm e2e subprocess failed for backend={backend} "
487 + f"with exit code {proc.exitcode} and produced no result"
488 + ) from exc
489 + raise AssertionError(
490 + f"quick-reduce llm e2e subprocess produced no result for backend={backend}"
491 + ) from exc
492 +
493 + if result["status"] == "xfail":
494 + pytest.xfail(result["reason"])
495 + if result["status"] == "error":
496 + raise AssertionError(
497 + f"quick-reduce llm e2e subprocess failed for backend={backend}:\n"
498 + f"{result['reason']}"
499 + )
500 +
501 + assert proc.exitcode == 0, (
502 + f"quick-reduce llm e2e subprocess failed for backend={backend} "
503 + f"with exit code {proc.exitcode}"
504 + )
505 +
506 +
507 +def test_quick_reduce_regime_values():
508 + from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
509 +
510 + assert QuickReduceRegime.FP.value == 0
511 + assert QuickReduceRegime.INT8.value == 1
512 + assert QuickReduceRegime.INT6.value == 2
513 + assert QuickReduceRegime.INT4.value == 3
514 + assert QuickReduceRegime.NONE.value == 4
515 +
516 +
517 +def test_quick_reduce_regime_names():
518 + from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
519 +
520 + assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
521 +
522 +
523 +@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
524 +def test_quick_reduce_quantization_env_var(monkeypatch, quant_level):
525 + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_level)
526 +
527 + reloaded_envs = _reload_envs()
528 + assert quant_level == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION
529 +
530 +
531 +def test_quick_reduce_quantization_default(monkeypatch):
532 + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", raising=False)
533 +
534 + reloaded_envs = _reload_envs()
535 + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION == "NONE"
536 +
537 +
538 +@pytest.mark.parametrize("cast_bf16", [True, False])
539 +def test_quick_reduce_cast_bf16_to_fp16_env_var(monkeypatch, cast_bf16):
540 + monkeypatch.setenv(
541 + "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1" if cast_bf16 else "0"
542 + )
543 +
544 + reloaded_envs = _reload_envs()
545 + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is cast_bf16
546 +
547 +
548 +def test_quick_reduce_cast_bf16_to_fp16_default(monkeypatch):
549 + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", raising=False)
550 +
551 + reloaded_envs = _reload_envs()
552 + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is True
553 +
554 +
555 +@pytest.mark.parametrize("max_mb", [128, 512, 2048, None])
556 +def test_quick_reduce_max_size_env_var(monkeypatch, max_mb):
557 + if max_mb is None:
558 + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False)
559 + else:
560 + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", str(max_mb))
561 +
562 + reloaded_envs = _reload_envs()
563 + assert max_mb == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
564 +
565 +
566 +def test_quick_reduce_max_size_default(monkeypatch):
567 + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False)
568 +
569 + reloaded_envs = _reload_envs()
570 + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is None
571 +
572 +
573 +@pytest.mark.parametrize(
574 + ("gcn_arch_name", "expected"),
575 + [
576 + ("gfx942", True),
577 + ("gfx950", True),
578 + ("gfx90a", False),
579 + ("", False),
580 + ],
581 +)
582 +def test_quick_allreduce_rocm_arch_available(gcn_arch_name, expected):
583 + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
584 +
585 + qar = QuickAllReduce.__new__(QuickAllReduce)
586 + qar.disabled = True
587 +
588 + with (
589 + patch(
590 + "vllm.distributed.device_communicators.quick_all_reduce.current_platform."
591 + "is_rocm",
592 + return_value=True,
593 + ),
594 + patch(
595 + "torch.cuda.get_device_properties",
596 + return_value=SimpleNamespace(gcnArchName=gcn_arch_name),
597 + ),
598 + ):
599 + assert qar._rocm_arch_available() is expected
600 +
601 +
602 +def test_quick_allreduce_rocm_arch_available_handles_probe_failure():
603 + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
604 +
605 + qar = QuickAllReduce.__new__(QuickAllReduce)
606 + qar.disabled = True
607 +
608 + with (
609 + patch(
610 + "vllm.distributed.device_communicators.quick_all_reduce.current_platform."
611 + "is_rocm",
612 + return_value=True,
613 + ),
614 + patch("torch.cuda.get_device_properties", side_effect=RuntimeError),
615 + ):
616 + assert qar._rocm_arch_available() is False
617 +
618 +
619 +def test_quick_allreduce_rejects_disabled():
620 + qar = _make_quick_allreduce(disabled=True)
621 +
622 + inp = torch.zeros(1024, dtype=torch.float16)
623 + assert qar.should_quick_allreduce(inp) is False
624 +
625 +
626 +def test_quick_allreduce_rejects_unsupported_dtype():
627 + qar = _make_quick_allreduce()
628 +
629 + inp = torch.zeros(1024 * 1024, dtype=torch.float32)
630 + assert qar.should_quick_allreduce(inp) is False
631 +
632 +
633 +def test_quick_allreduce_rejects_non_aligned_input():
634 + qar = _make_quick_allreduce()
635 +
636 + inp = torch.zeros(5, dtype=torch.float16)
637 + assert qar.should_quick_allreduce(inp) is False
638 +
639 +
640 +def test_quick_allreduce_rejects_non_contiguous_input():
641 + qar = _make_quick_allreduce()
642 +
643 + inp = torch.zeros((1024, 1024), dtype=torch.float16)[:, ::2]
644 + assert qar.should_quick_allreduce(inp) is False
645 +
646 +
647 +def test_quick_allreduce_rejects_input_smaller_than_threshold():
648 + qar = _make_quick_allreduce()
649 +
650 + inp = torch.zeros((MB // 2) - 8, dtype=torch.float16)
651 + assert qar.should_quick_allreduce(inp) is False
652 +
653 +
654 +def test_quick_allreduce_accepts_input_at_threshold():
655 + qar = _make_quick_allreduce()
656 +
657 + inp = torch.zeros(MB // 2, dtype=torch.float16)
658 + assert qar.should_quick_allreduce(inp) is True
659 +
660 +
661 +def test_quick_allreduce_rejects_input_larger_than_max_size():
662 + qar = _make_quick_allreduce(qr_max_size=1 * MB)
663 +
664 + inp = torch.zeros(MB, dtype=torch.float16)
665 + assert qar.should_quick_allreduce(inp) is False
666 +
667 +
668 +def test_quick_allreduce_bf16_uses_fp16_threshold_when_cast_enabled():
669 + inp = torch.zeros(MB // 2, dtype=torch.bfloat16)
670 +
671 + without_cast = _make_quick_allreduce(use_fp16_kernels=False)
672 + with_cast = _make_quick_allreduce(use_fp16_kernels=True)
673 +
674 + assert without_cast.should_quick_allreduce(inp) is False
675 + assert with_cast.should_quick_allreduce(inp) is True
676 +
677 +
678 +def test_quick_allreduce_supported_world_sizes():
679 + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
680 +
681 + assert QuickAllReduce._SUPPORTED_WORLD_SIZES == [2, 4, 8]
682 +
683 +
684 +def test_quick_allreduce_supported_dtypes():
685 + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
686 +
687 + assert [torch.float16, torch.bfloat16] == QuickAllReduce._SUPPORTED_DTYPES
688 +
689 +
690 +def test_quick_allreduce_min_size_table():
691 + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
692 +
693 + for dtype in [torch.float16, torch.bfloat16]:
694 + for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
695 + min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
696 + assert len(min_sizes) == 4
697 + assert all(size > 0 for size in min_sizes)
698 +
699 +
700 +def test_qr_max_size():
701 + from vllm import _custom_ops as ops
702 +
703 + max_size = ops.qr_max_size()
704 + assert isinstance(max_size, int)
705 + assert max_size > 0
706 +
707 +
708 +@pytest.mark.skipif(
709 + current_platform.device_count() < WORLD_SIZE,
710 + reason="requires 2 ROCm GPUs",
711 +)
712 +@pytest.mark.parametrize("quant_level", QUANT_LEVELS)
713 +def test_quick_allreduce_two_gpu_correctness(quant_level):
714 + _log(f"two-GPU correctness case: quant={quant_level}")
715 + _run_two_gpu_quick_allreduce_test(
716 + quant_level=quant_level,
717 + dtype_name="float16",
718 + cast_bf16=False,
719 + )
720 +
721 +
722 +@pytest.mark.skipif(
723 + current_platform.device_count() < WORLD_SIZE,
724 + reason="requires 2 ROCm GPUs",
725 +)
726 +def test_quick_allreduce_bf16_cast_mode():
727 + _log("BF16 cast case")
728 + _run_two_gpu_quick_allreduce_test(
729 + quant_level="FP",
730 + dtype_name="bfloat16",
731 + cast_bf16=True,
732 + )
733 +
734 +
735 +@pytest.mark.skipif(
736 + current_platform.device_count() < WORLD_SIZE,
737 + reason="requires 2 ROCm GPUs",
738 +)
739 +def test_quick_allreduce_llm_e2e():
740 + _log("LLM e2e case: backend=mp")
741 + run_quick_reduce_llm_e2e(backend="mp")
742 +
743 +
744 +@pytest.mark.skipif(
745 + current_platform.device_count() < WORLD_SIZE,
746 + reason="requires 2 ROCm GPUs",
747 +)
748 +def test_quick_allreduce_llm_e2e_ray():
749 + _log("LLM e2e case: backend=ray")
750 + run_quick_reduce_llm_e2e(backend="ray")
tests/utils.py
+26 -31
@@ -1363,43 +1363,38 @@ def multi_process_parallel(
1363 ) -> None:
1364 import ray
1365
1366 - # Using ray helps debugging the error when it failed
1367 - # as compared to multiprocessing.
1368 - # NOTE: We need to set working_dir for distributed tests,
1369 - # otherwise we may get import errors on ray workers
1370 - # NOTE: Force ray not to use gitignore file as excluding, otherwise
1371 - # it will not move .so files to working dir.
1372 - # So we have to manually add some of large directories
1373 - os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1"
1366 + # Using ray helps debugging the error when it failed as compared to
1367 + # multiprocessing. For local Ray workers, putting the repo root on
1368 + # PYTHONPATH is enough and avoids uploading the full source tree, which
1369 + # exceeds Ray's working_dir package size limit on CI.
1370 + env_vars = {
1371 + "PYTHONPATH": os.pathsep.join(
1372 + filter(None, [str(VLLM_PATH), os.environ.get("PYTHONPATH")])
1373 + ),
1374 + **{env_var: "1" for env_var in current_platform.ray_noset_device_env_vars},
1375 + }
1376 ray.init(
1377 runtime_env={
1376 - "working_dir": VLLM_PATH,
1377 - "excludes": [
1378 - "build",
1379 - ".git",
1380 - "cmake-build-*",
1381 - "shellcheck",
1382 - "dist",
1383 - "ep_kernels_workspace",
1384 - ],
1378 + "env_vars": env_vars,
1379 }
1380 )
1381
1382 distributed_init_port = get_open_port()
1389 - refs = []
1390 - for rank in range(tp_size * pp_size):
1391 - refs.append(
1392 - test_target.remote(
1393 - monkeypatch,
1394 - tp_size,
1395 - pp_size,
1396 - rank,
1397 - distributed_init_port,
1398 - ),
1399 - )
1400 - ray.get(refs)
1401 -
1402 - ray.shutdown()
1383 + try:
1384 + refs = []
1385 + for rank in range(tp_size * pp_size):
1386 + refs.append(
1387 + test_target.remote(
1388 + monkeypatch,
1389 + tp_size,
1390 + pp_size,
1391 + rank,
1392 + distributed_init_port,
1393 + ),
1394 + )
1395 + ray.get(refs)
1396 + finally:
1397 + ray.shutdown()
1398
1399
1400 @contextmanager