utils.py 61.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import asyncio
5
import contextlib
6
import copy
7
import functools
8
import importlib
9
import itertools
10
import json
11
import os
12
import random
13
import signal
14
15
import subprocess
import sys
16
import tempfile
17
import time
18
import warnings
19
from collections.abc import Callable, Iterable
20
from contextlib import ExitStack, contextmanager, suppress
21
from multiprocessing import Process
22
from pathlib import Path
23
from typing import Any, Literal
24
from unittest.mock import patch
25

26
import anthropic
27
import cloudpickle
28
import httpx
29
import openai
30
import pytest
31
import requests
32
import torch
33
import torch.nn.functional as F
34
from openai.types.completion import Completion
35
from typing_extensions import ParamSpec
36

37
import vllm.envs as envs
38
from tests.models.utils import TextTextLogprobs
39
40
41
42
from vllm.distributed import (
    ensure_model_parallel_initialized,
    init_distributed_environment,
)
43
from vllm.engine.arg_utils import AsyncEngineArgs
44
from vllm.entrypoints.cli.serve import ServeSubcommand
45
from vllm.model_executor.kernels.linear import (
46
    FP8ScaledMMLinearKernel,
47
    init_fp8_linear_kernel,
48
49
50
51
52
53
)
from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    GroupShape,
    QuantKey,
)
54
from vllm.model_executor.model_loader import get_model_loader
55
from vllm.platforms import current_platform
56
from vllm.tokenizers import get_tokenizer
Cyrus Leung's avatar
Cyrus Leung committed
57
from vllm.utils.argparse_utils import FlexibleArgumentParser
58
from vllm.utils.mem_constants import GB_bytes
59
from vllm.utils.network_utils import get_open_port
60
61
62
63
from vllm.utils.torch_utils import (
    cuda_device_count_stateless,
    set_random_seed,  # noqa: F401 - re-exported for use in test files
)
64

65
66
FP8_DTYPE = current_platform.fp8_dtype()

67
if current_platform.is_rocm():
68
69
    import threading

70
71
72
73
74
75
    from amdsmi import (
        amdsmi_get_gpu_vram_usage,
        amdsmi_get_processor_handles,
        amdsmi_init,
        amdsmi_shut_down,
    )
76

77
78
    _amdsmi_lock = threading.Lock()

79
80
    @contextmanager
    def _nvml():
81
82
83
84
85
86
        with _amdsmi_lock:
            try:
                amdsmi_init()
                yield
            finally:
                amdsmi_shut_down()
87
elif current_platform.is_cuda():
88
89
90
91
92
93
    from vllm.third_party.pynvml import (
        nvmlDeviceGetHandleByIndex,
        nvmlDeviceGetMemoryInfo,
        nvmlInit,
        nvmlShutdown,
    )
94
95
96
97
98
99
100
101

    @contextmanager
    def _nvml():
        try:
            nvmlInit()
            yield
        finally:
            nvmlShutdown()
102
103
104
105
106
else:

    @contextmanager
    def _nvml():
        yield
107

108

109
110
VLLM_PATH = Path(__file__).parent.parent
"""Path to root of the vLLM repository."""
111

112
113
114
115
116
117
118
119
120
121
122
123
124
# ROCm: disable skinny GEMM to avoid non-deterministic results from
# atomic reductions in wvSplitKrc kernel.
# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
ROCM_ENV_OVERRIDES = (
    {"VLLM_ROCM_USE_SKINNY_GEMM": "0"} if current_platform.is_rocm() else {}
)
# ROCm: disable prefix caching and eliminate batch variance to reduce
# test flakiness.
ROCM_EXTRA_ARGS = (
    ["--no-enable-prefix-caching", "--max-num-seqs", "1"]
    if current_platform.is_rocm()
    else []
)
125
126
127
128
129
130
# Python-API equivalent of ROCM_EXTRA_ARGS for use with EngineArgs kwargs.
ROCM_ENGINE_KWARGS: dict = (
    {"enable_prefix_caching": False, "max_num_seqs": 1}
    if current_platform.is_rocm()
    else {}
)
131

132

133
134
135
136
137
138
139
class RemoteVLLMServer:
    """Base class for launching vLLM server subprocesses for testing.

    Subclasses must override ``_create_cli_subcommand`` and
    ``_start_server``.
    """

140
    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
141
142
143
144
145
    proc: subprocess.Popen

    def _create_cli_subcommand(self):
        """Return a CLISubcommand instance used to parse CLI args."""
        raise NotImplementedError
146

147
    def _start_server(
148
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
149
150
    ) -> None:
        """Subclasses override this method to customize server process launch"""
151
        raise NotImplementedError
152

153
154
155
156
157
158
159
160
161
162
163
    def _pre_download_model(self, model: str, args) -> None:
        """Download model weights before starting the server to avoid timeout."""
        is_local = os.path.isdir(model)
        if not is_local:
            engine_args = AsyncEngineArgs.from_cli_args(args)
            model_config = engine_args.create_model_config()
            load_config = engine_args.create_load_config()

            model_loader = get_model_loader(load_config)
            model_loader.download_model(model_config)

164
165
166
167
168
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
        *,
169
        env_dict: dict[str, str] | None = None,
170
        seed: int = 0,
171
        auto_port: bool = True,
172
173
        max_wait_seconds: float | None = None,
        override_hf_configs: dict[str, Any] | None = None,
174
    ) -> None:
175
        if auto_port:
176
            if "-p" in vllm_serve_args or "--port" in vllm_serve_args:
177
178
179
                raise ValueError(
                    "You have manually specified the port when `auto_port=True`."
                )
180

181
182
183
            # No need for a port if using unix sockets
            if "--uds" not in vllm_serve_args:
                # Don't mutate the input args
184
                vllm_serve_args = vllm_serve_args + ["--port", str(get_open_port())]
185
186
        if seed is not None:
            if "--seed" in vllm_serve_args:
187
188
189
                raise ValueError(
                    f"You have manually specified the seed when `seed={seed}`."
                )
190
191

            vllm_serve_args = vllm_serve_args + ["--seed", str(seed)]
192

193
194
195
        if override_hf_configs is not None:
            vllm_serve_args = vllm_serve_args + [
                "--hf-overrides",
196
                json.dumps(override_hf_configs),
197
198
            ]

199
        parser = FlexibleArgumentParser(description="vLLM's remote server.")
200
        subparsers = parser.add_subparsers(required=False, dest="subparser")
201
        parser = self._create_cli_subcommand().subparser_init(subparsers)
202
        args = parser.parse_args(["--model", model, *vllm_serve_args])
203
204
205
206
207
        self.uds = args.uds
        if args.uds:
            self.host = None
            self.port = None
        else:
208
            self.host = str(args.host or "127.0.0.1")
209
            self.port = int(args.port)
210

211
212
213
        self.show_hidden_metrics = (
            getattr(args, "show_hidden_metrics_for_version", None) is not None
        )
214

215
        self._pre_download_model(model, args)
216

217
218
219
220
221
222
        # Record GPU memory before server start so we know what
        # "released" looks like.
        self._pre_server_gpu_memory = self._get_gpu_memory_used()
        if self._pre_server_gpu_memory is not None:
            pre_gb = self._pre_server_gpu_memory / 1e9
            print(
223
224
                f"[{type(self).__name__}] GPU memory before server start: "
                f"{pre_gb:.2f} GB"
225
226
            )

227
        self._start_server(model, vllm_serve_args, env_dict)
228
229
230
231
232
233
234
235
236
237
238
        max_wait_seconds = max_wait_seconds or 480
        try:
            self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
        except Exception:
            # If the server never became healthy, we must still clean up
            # the subprocess tree. Without this, a timeout in __init__
            # leaks the server + EngineCore processes (and their GPU
            # memory), because __exit__ is never called when __init__
            # raises inside a ``with`` statement.
            self._shutdown()
            raise
239
240
241
242
243

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
244
245
246
247
248
249
250
251
252
        self._shutdown()

    def _shutdown(self) -> None:
        """Kill the server process tree and wait for GPU memory release.

        Called from both ``__exit__`` (normal path) and ``__init__``
        (when the server fails to start). Must be safe to call even if
        the process is already dead.
        """
253
        pid = self.proc.pid
254
255
256
257
258
259
260
261

        # Get the process group ID. Because we used
        # start_new_session=True the pgid equals the server's pid.
        try:
            pgid = os.getpgid(pid)
        except (ProcessLookupError, OSError):
            pgid = None

262
263
        # Phase 1: graceful SIGTERM to the root process
        with contextlib.suppress(ProcessLookupError, OSError):
264
            self.proc.terminate()
265
            print(f"[RemoteOpenAIServer] Sent SIGTERM to process {pid}")
266

267
        try:
268
269
            self.proc.wait(timeout=15)
            print(f"[RemoteOpenAIServer] Server {pid} terminated gracefully")
270
        except subprocess.TimeoutExpired:
271
            # Phase 2: SIGKILL the entire process group
272
273
            print(
                f"[RemoteOpenAIServer] Server {pid} did not respond "
274
                "to SIGTERM, sending SIGKILL to process group"
275
            )
276
277
278
279
280
281
            if pgid is not None:
                with contextlib.suppress(ProcessLookupError, OSError):
                    os.killpg(pgid, signal.SIGKILL)
            else:
                self.proc.kill()

282
            try:
283
                self.proc.wait(timeout=10)
284
                print(f"[RemoteOpenAIServer] Server {pid} killed")
285
            except subprocess.TimeoutExpired:
286
                pass
287

288
289
290
291
292
293
294
        # After killing the root process, ensure all children in the
        # process group (e.g. EngineCore workers) are also dead.
        # On ROCm especially, surviving children hold GPU contexts and
        # prevent VRAM from being reclaimed by the driver.
        self._kill_process_group_survivors(pgid)

        # Wait for GPU memory to actually be freed, not just
295
        # "stabilized at whatever level it's at".
296
297
        self._wait_for_gpu_memory_release()

298
299
300
301
302
    def _kill_process_group_survivors(
        self, pgid: int | None, timeout: float = 15.0
    ) -> None:
        """SIGKILL any processes still in the server's process group
        and wait for them to exit.
303

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
        Because the server is launched with ``start_new_session=True``,
        all its children (EngineCore, workers, etc.) share the same
        pgid. After the root process is killed, stragglers -- especially
        on ROCm where GPU contexts linger until the *process* exits --
        must be reaped explicitly.

        Uses ``/proc`` to scan for pgid members so this works even after
        the parent has been reaped (unlike ``psutil.Process.children``).
        """
        if pgid is None:
            return

        # Send SIGKILL to the entire process group one more time.
        # This is cheap and harmless if everyone is already dead.
        with contextlib.suppress(ProcessLookupError, OSError):
            os.killpg(pgid, signal.SIGKILL)

        # Collect surviving PIDs by scanning /proc for matching pgid.
        # This works on Linux even after the parent has been waited on
        # and is more reliable than psutil.Process(parent).children().
        survivor_pids = self._find_pgid_members(pgid)

        if not survivor_pids:
            return

        print(
            f"[RemoteOpenAIServer] {len(survivor_pids)} process(es) still "
            f"in pgid {pgid} after SIGKILL: {survivor_pids}"
        )

        # Wait for each survivor to actually exit so the GPU driver
        # releases its VRAM.
        deadline = time.time() + timeout
        while survivor_pids and time.time() < deadline:
            still_alive = []
            for spid in survivor_pids:
                try:
                    os.kill(spid, 0)  # Check if still alive
                    still_alive.append(spid)
                except (ProcessLookupError, OSError):
                    pass
            survivor_pids = still_alive
            if survivor_pids:
                time.sleep(0.5)

        if survivor_pids:
            print(
                f"[RemoteOpenAIServer] WARNING: processes {survivor_pids} "
                f"in pgid {pgid} could not be killed within {timeout}s"
            )

    @staticmethod
    def _find_pgid_members(pgid: int) -> list[int]:
        """Return PIDs of all living processes whose pgid matches."""
        members: list[int] = []
        proc_path = Path("/proc")
        if not proc_path.is_dir():
            return members
        for entry in proc_path.iterdir():
            if not entry.name.isdigit():
                continue
            pid = int(entry.name)
            try:
                if os.getpgid(pid) == pgid:
                    members.append(pid)
            except OSError:
                continue
        return members
372

373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
    def _get_gpu_memory_used(self) -> float | None:
        """Get total GPU memory used across all visible devices in bytes."""
        try:
            if current_platform.is_rocm():
                with _nvml():
                    handles = amdsmi_get_processor_handles()
                    total_used = 0
                    for handle in handles:
                        vram_info = amdsmi_get_gpu_vram_usage(handle)
                        total_used += vram_info["vram_used"]
                    return total_used
            elif current_platform.is_cuda():
                with _nvml():
                    total_used = 0
                    device_count = cuda_device_count_stateless()
                    for i in range(device_count):
                        handle = nvmlDeviceGetHandleByIndex(i)
                        mem_info = nvmlDeviceGetMemoryInfo(handle)
                        total_used += mem_info.used
                    return total_used
        except Exception as e:
            print(f"[RemoteOpenAIServer] Could not query GPU memory: {e}")
            return None
        return None

398
399
400
    def _wait_for_gpu_memory_release(
        self, timeout: float = 120.0, log_interval: float = 10.0
    ):
401
402
        """Wait for GPU memory to drop back toward pre-server levels.

403
404
405
406
407
        Waits the full timeout for memory to return close to the
        pre-server baseline. Does NOT fall back to a "stabilization"
        heuristic -- if memory is still held when the timeout expires,
        the test fails so the problem is surfaced immediately rather
        than causing cascading OOM failures in every subsequent test.
408
409
410
411
412
413
414
415
416
417
418
        """
        baseline = self._pre_server_gpu_memory
        if baseline is None:
            # Can't query GPU memory - nothing to do
            return

        # Allow up to 2 GiB overhead above baseline for driver/context state
        # that may persist between server instances.
        headroom_bytes = 2 * 1024 * 1024 * 1024
        target = baseline + headroom_bytes

419
        start = time.time()
420
        next_log_time = start + log_interval
421
422
423
424
425
426
427

        while time.time() - start < timeout:
            used = self._get_gpu_memory_used()

            if used is None:
                return  # Can't query, assume ok

428
429
430
            used_gb = used / 1e9
            target_gb = target / 1e9
            elapsed = time.time() - start
431

432
433
434
435
436
437
438
439
            if used <= target:
                print(
                    f"[RemoteOpenAIServer] GPU memory released to "
                    f"{used_gb:.2f} GB (target: {target_gb:.2f} GB) "
                    f"in {elapsed:.1f}s"
                )
                return

440
441
442
443
444
445
446
447
448
            now = time.time()
            if now >= next_log_time:
                print(
                    f"[RemoteOpenAIServer] Waiting for GPU memory release: "
                    f"{used_gb:.2f} GB (target: {target_gb:.2f} GB) "
                    f"[{elapsed:.0f}s/{timeout:.0f}s]"
                )
                next_log_time = now + log_interval

449
450
            time.sleep(1.0)

451
452
        # Timeout -- raise so the current test fails with a clear
        # message instead of silently poisoning subsequent tests.
453
454
        final_used = self._get_gpu_memory_used()
        final_gb = final_used / 1e9 if final_used else 0.0
455
        raise RuntimeError(
456
457
458
459
460
            f"[RemoteOpenAIServer] GPU memory did not release within "
            f"{timeout}s. Current: {final_gb:.2f} GB, "
            f"target: {target / 1e9:.2f} GB, "
            f"baseline: {baseline / 1e9:.2f} GB. "
            f"Child processes may still be holding GPU memory."
461
        )
462

463
    def _poll(self) -> int | None:
464
465
466
        """Subclasses override this method to customize process polling"""
        return self.proc.poll()

467
468
469
    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
470
471
472
473
474
        client = (
            httpx.Client(transport=httpx.HTTPTransport(uds=self.uds))
            if self.uds
            else requests
        )
475
476
        while True:
            try:
477
                if client.get(url).status_code == 200:
478
                    break
479
480
481
482
483
            except Exception:
                # this exception can only be raised by requests.get,
                # which means the server is not ready yet.
                # the stack trace is not useful, so we suppress it
                # by using `raise from None`.
484
                result = self._poll()
485
                if result is not None and result != 0:
486
                    raise RuntimeError("Server exited unexpectedly.") from None
487
488
489

                time.sleep(0.5)
                if time.time() - start > timeout:
490
                    raise RuntimeError("Server failed to start in time.") from None
491
492
493

    @property
    def url_root(self) -> str:
494
495
496
497
498
        return (
            f"http://{self.uds.split('/')[-1]}"
            if self.uds
            else f"http://{self.host}:{self.port}"
        )
499
500
501
502

    def url_for(self, *parts: str) -> str:
        return self.url_root + "/" + "/".join(parts)

503
504
505
    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
506
507
508
        return openai.OpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
509
510
            max_retries=0,
            **kwargs,
511
512
        )

513
    def get_async_client(self, **kwargs):
514
515
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
516
517
518
519
520
521
        return openai.AsyncOpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
            max_retries=0,
            **kwargs,
        )
522

523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
    def get_client_anthropic(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
        return anthropic.Anthropic(
            base_url=self.url_for(),
            api_key=self.DUMMY_API_KEY,
            max_retries=0,
            **kwargs,
        )

    def get_async_client_anthropic(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
        return anthropic.AsyncAnthropic(
            base_url=self.url_for(), api_key=self.DUMMY_API_KEY, max_retries=0, **kwargs
        )

540

541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
class RemoteOpenAIServer(RemoteVLLMServer):
    """Launches ``vllm serve`` for testing OpenAI-compatible endpoints."""

    def _create_cli_subcommand(self):
        return ServeSubcommand()

    def _start_server(
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
    ) -> None:
        env = os.environ.copy()
        # the current process might initialize cuda,
        # to be safe, we should use spawn method
        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
        if env_dict is not None:
            env.update(env_dict)
        serve_cmd = ["vllm", "serve", model, *vllm_serve_args]
        print(f"Launching RemoteOpenAIServer with: {' '.join(serve_cmd)}")
        print(f"Environment variables: {env}")
        self.proc: subprocess.Popen = subprocess.Popen(
            serve_cmd,
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
            # Create a dedicated process group so we can kill
            # the entire tree (parent + EngineCore + workers) at once.
            start_new_session=True,
        )


class RemoteLaunchRenderServer(RemoteVLLMServer):
    """Launches ``vllm launch render`` for GPU-less serving tests."""

    def _create_cli_subcommand(self):
        return ServeSubcommand()

    def _start_server(
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
    ) -> None:
        env = os.environ.copy()
        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
        if env_dict is not None:
            env.update(env_dict)
        serve_cmd = ["vllm", "launch", "render", model, *vllm_serve_args]
        print(f"Launching RemoteLaunchRenderServer with: {' '.join(serve_cmd)}")
        self.proc: subprocess.Popen = subprocess.Popen(
            serve_cmd,
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
            start_new_session=True,
        )

593
594
595
596
597
598
599
600
601
602
603
604
605
    def _pre_download_model(self, model: str, args) -> None:
        """Download only the tokenizer files (no model weights needed)."""
        is_local = os.path.isdir(model)
        if not is_local:
            engine_args = AsyncEngineArgs.from_cli_args(args)
            model_config = engine_args.create_model_config()
            get_tokenizer(
                model_config.tokenizer,
                tokenizer_mode=model_config.tokenizer_mode,
                trust_remote_code=model_config.trust_remote_code,
                revision=model_config.tokenizer_revision,
            )

606
607
608
    def _wait_for_gpu_memory_release(
        self, timeout: float = 30.0, log_interval: float = 10.0
    ):
609
610
611
        pass  # No GPU used


612
613
614
class RemoteOpenAIServerCustom(RemoteOpenAIServer):
    """Launch test server with custom child process"""

615
    def _start_server(
616
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
617
    ) -> None:
618
        self.proc: Process = Process(
619
620
            target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)
        )  # type: ignore[assignment]
621
622
        self.proc.start()

623
624
625
626
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
627
        child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],
628
        *,
629
        env_dict: dict[str, str] | None = None,
630
        seed: int = 0,
631
        auto_port: bool = True,
632
        max_wait_seconds: float | None = None,
633
    ) -> None:
634
635
636
        """Store custom child process function then invoke superclass
        constructor which will indirectly launch it."""
        self.child_process_fxn = child_process_fxn
637
638
639
640
641
642
643
644
        super().__init__(
            model=model,
            vllm_serve_args=vllm_serve_args,
            env_dict=env_dict,
            seed=seed,
            auto_port=auto_port,
            max_wait_seconds=max_wait_seconds,
        )
645

646
    def _poll(self) -> int | None:
647
648
649
650
651
652
653
654
655
656
        return self.proc.exitcode

    def __exit__(self, exc_type, exc_value, traceback):
        self.proc.terminate()
        self.proc.join(8)
        if self.proc.is_alive():
            # force kill if needed
            self.proc.kill()


657
658
659
660
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
661
    token_ids: list[int],
662
663
664
665
):
    results = []

    # test with text prompt
666
667
668
669
670
671
672
673
674
675
676
677
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=5, temperature=0.0
    )

    results.append(
        {
            "test": "single_completion",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
678
679
680
681
682
683
684
685
686

    # test using token IDs
    completion = client.completions.create(
        model=model,
        prompt=token_ids,
        max_tokens=5,
        temperature=0.0,
    )

687
688
689
690
691
692
693
694
    results.append(
        {
            "test": "token_ids",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
695
696

    # test seeded random sampling
697
698
699
700
701
702
703
704
705
706
707
708
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=5, seed=33, temperature=1.0
    )

    results.append(
        {
            "test": "seeded_sampling",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
709
710

    # test seeded random sampling with multiple prompts
711
712
713
714
715
716
717
718
719
720
721
722
    completion = client.completions.create(
        model=model, prompt=[prompt, prompt], max_tokens=5, seed=33, temperature=1.0
    )

    results.append(
        {
            "test": "seeded_sampling",
            "text": [choice.text for choice in completion.choices],
            "finish_reason": [choice.finish_reason for choice in completion.choices],
            "usage": completion.usage,
        }
    )
723
724
725
726
727
728
729
730
731

    # test simple list
    batch = client.completions.create(
        model=model,
        prompt=[prompt, prompt],
        max_tokens=5,
        temperature=0.0,
    )

732
733
734
735
736
737
738
    results.append(
        {
            "test": "simple_list",
            "text0": batch.choices[0].text,
            "text1": batch.choices[1].text,
        }
    )
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754

    # test streaming
    batch = client.completions.create(
        model=model,
        prompt=[prompt, prompt],
        max_tokens=5,
        temperature=0.0,
        stream=True,
    )

    texts = [""] * 2
    for chunk in batch:
        assert len(chunk.choices) == 1
        choice = chunk.choices[0]
        texts[choice.index] += choice.text

755
756
757
758
759
760
    results.append(
        {
            "test": "streaming",
            "texts": texts,
        }
    )
761
762
763
764

    return results


765
766
767
768
769
770
771
772
def _test_completion_close(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

    # test with text prompt
773
774
775
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=1, logprobs=5, temperature=0.0
    )
776

777
778
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
779

780
781
782
783
784
785
    results.append(
        {
            "test": "completion_close",
            "logprobs": logprobs,
        }
    )
786
787
788
789

    return results


790
791
792
793
794
795
796
def _test_chat(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

797
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
798
799

    # test with text prompt
800
801
802
803
804
805
806
807
808
809
810
811
    chat_response = client.chat.completions.create(
        model=model, messages=messages, max_tokens=5, temperature=0.0
    )

    results.append(
        {
            "test": "completion_close",
            "text": chat_response.choices[0].message.content,
            "finish_reason": chat_response.choices[0].finish_reason,
            "usage": chat_response.usage,
        }
    )
812
813
814
815

    return results


816
817
818
819
820
821
822
823
824
825
826
827
828
829
def _test_embeddings(
    client: openai.OpenAI,
    model: str,
    text: str,
):
    results = []

    # test with text input
    embeddings = client.embeddings.create(
        model=model,
        input=text,
        encoding_format="float",
    )

830
831
832
833
834
835
836
    results.append(
        {
            "test": "single_embedding",
            "embedding": embeddings.data[0].embedding,
            "usage": embeddings.usage,
        }
    )
837
838
839
840

    return results


841
842
843
844
845
846
847
848
def _test_image_text(
    client: openai.OpenAI,
    model_name: str,
    image_url: str,
):
    results = []

    # test pure text input
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "How do you feel today?"},
            ],
        }
    ]

    chat_completion = client.chat.completions.create(
        model=model_name,
        messages=messages,
        temperature=0.0,
        max_tokens=1,
        logprobs=True,
        top_logprobs=5,
    )
866
867
868
869
870
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

    for x in top_logprobs:
        x.logprob = round(x.logprob, 2)

871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
    results.append(
        {
            "test": "pure_text",
            "logprobs": top_logprobs,
        }
    )

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": image_url}},
                {"type": "text", "text": "What's in this image?"},
            ],
        }
    ]

    chat_completion = client.chat.completions.create(
        model=model_name,
        messages=messages,
        temperature=0.0,
        max_tokens=1,
        logprobs=True,
        top_logprobs=5,
    )
896
897
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

898
899
900
901
902
903
    results.append(
        {
            "test": "text_image",
            "logprobs": top_logprobs,
        }
    )
904
905
906
907

    return results


908
909
910
911
def compare_two_settings(
    model: str,
    arg1: list[str],
    arg2: list[str],
912
913
    env1: dict[str, str] | None = None,
    env2: dict[str, str] | None = None,
914
915
    *,
    method: str = "generate",
916
    max_wait_seconds: float | None = None,
917
) -> None:
918
    """
919
920
921
922
923
924
925
926
927
    Launch API server with two different sets of arguments/environments
    and compare the results of the API calls.

    Args:
        model: The model to test.
        arg1: The first set of arguments to pass to the API server.
        arg2: The second set of arguments to pass to the API server.
        env1: The first set of environment variables to pass to the API server.
        env2: The second set of environment variables to pass to the API server.
928
929
    """

930
931
932
933
934
935
936
937
938
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


939
940
941
def compare_all_settings(
    model: str,
    all_args: list[list[str]],
942
    all_envs: list[dict[str, str] | None],
943
944
    *,
    method: str = "generate",
945
    max_wait_seconds: float | None = None,
946
) -> None:
947
948
949
950
951
952
953
954
955
    """
    Launch API server with several different sets of arguments/environments
    and compare the results of the API calls with the first set of arguments.
    Args:
        model: The model to test.
        all_args: A list of argument lists to pass to the API server.
        all_envs: A list of environment dictionaries to pass to the API server.
    """

956
    trust_remote_code = False
957
    for args in all_args:
958
959
960
961
962
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
963
    for args in all_args:
964
965
966
967
968
969
970
971
972
        if "--tokenizer-mode" in args:
            tokenizer_mode = args[args.index("--tokenizer-mode") + 1]
            break

    tokenizer = get_tokenizer(
        model,
        trust_remote_code=trust_remote_code,
        tokenizer_mode=tokenizer_mode,
    )
973

974
975
976
977
978
979
980
    can_force_load_format = True

    for args in all_args:
        if "--load-format" in args:
            can_force_load_format = False
            break

981
    prompt = "Hello, my name is"
982
    token_ids = tokenizer(prompt).input_ids
983
    ref_results: list = []
984
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
985
986
987
988
989
990
991
992
993
        if can_force_load_format:
            # we are comparing the results and
            # usually we don't need real weights.
            # we force to use dummy weights by default,
            # and it should work for most of the cases.
            # if not, we can use VLLM_TEST_FORCE_LOAD_FORMAT
            # environment variable to force the load format,
            # e.g. in quantization tests.
            args = args + ["--load-format", envs.VLLM_TEST_FORCE_LOAD_FORMAT]
994
        compare_results: list = []
995
        results = ref_results if i == 0 else compare_results
996
997
998
        with RemoteOpenAIServer(
            model, args, env_dict=env, max_wait_seconds=max_wait_seconds
        ) as server:
999
1000
1001
1002
1003
1004
            client = server.get_client()

            # test models list
            models = client.models.list()
            models = models.data
            served_model = models[0]
1005
1006
1007
1008
1009
1010
1011
            results.append(
                {
                    "test": "models_list",
                    "id": served_model.id,
                    "root": served_model.root,
                }
            )
1012

1013
1014
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
1015
1016
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
1017
1018
            elif method == "generate_chat":
                results += _test_chat(client, model, prompt)
1019
1020
            elif method == "generate_with_image":
                results += _test_image_text(
1021
1022
                    client,
                    model,
1023
                    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
1024
                )
1025
1026
1027
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
1028
                raise ValueError(f"Unknown method: {method}")
1029

1030
1031
1032
1033
1034
1035
            if i > 0:
                # if any setting fails, raise an error early
                ref_args = all_args[0]
                ref_envs = all_envs[0]
                compare_args = all_args[i]
                compare_envs = all_envs[i]
1036
                for ref_result, compare_result in zip(ref_results, compare_results):
1037
1038
1039
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
1040
1041
1042
1043
1044
1045
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
1046
                            f"Embedding for {model=} are not the same.\n"
1047
1048
                            f"cosine_similarity={sim}\n"
                        )
1049
1050
                        del ref_result["embedding"]
                        del compare_result["embedding"]
1051
1052
1053
1054
1055
                    assert ref_result == compare_result, (
                        f"Results for {model=} are not the same.\n"
                        f"{ref_args=} {ref_envs=}\n"
                        f"{compare_args=} {compare_envs=}\n"
                        f"{ref_result=}\n"
1056
1057
                        f"{compare_result=}\n"
                    )
1058
1059


1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
@contextmanager
def ensure_current_vllm_config():
    """Ensures a vllm config is set for the duration of the context.

    If a config is already set, this is a no-op. Otherwise, it creates a default
    VllmConfig and sets it for the duration of the context.

    Used for tests that call functions which require a vllm config but don't
    need a specific config.

    Example:
        with ensure_current_vllm_config():
            init_distributed_environment(...)
            ensure_model_parallel_initialized(...)
    """
    from vllm.config import (
        VllmConfig,
        get_current_vllm_config_or_none,
        set_current_vllm_config,
    )

    if get_current_vllm_config_or_none() is not None:
        # Config already set, just yield
        yield
    else:
        # No config set, create a default one for the duration
        with set_current_vllm_config(VllmConfig()):
            yield


1090
1091
1092
1093
1094
1095
1096
def init_test_distributed_environment(
    tp_size: int,
    pp_size: int,
    rank: int,
    distributed_init_port: str,
    local_rank: int = -1,
) -> None:
1097
1098
1099
1100
1101
1102
1103
    # Note: This function is often called from Ray worker processes, so we
    # can't rely on pytest fixtures to set the config. We check if the config
    # is already set and only create a default one if needed.
    from vllm.config import (
        VllmConfig,
        get_current_vllm_config_or_none,
        set_current_vllm_config,
1104
    )
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115

    distributed_init_method = f"tcp://localhost:{distributed_init_port}"

    if get_current_vllm_config_or_none() is not None:
        # Config already set, use it directly
        init_distributed_environment(
            world_size=pp_size * tp_size,
            rank=rank,
            distributed_init_method=distributed_init_method,
            local_rank=local_rank,
        )
1116
        ensure_model_parallel_initialized(tp_size, pp_size)
1117
1118
1119
1120
1121
1122
1123
1124
1125
    else:
        # No config set, create a default one for the test
        with set_current_vllm_config(VllmConfig()):
            init_distributed_environment(
                world_size=pp_size * tp_size,
                rank=rank,
                distributed_init_method=distributed_init_method,
                local_rank=local_rank,
            )
1126
            ensure_model_parallel_initialized(tp_size, pp_size)
1127
1128


1129
def multi_process_parallel(
1130
    monkeypatch: pytest.MonkeyPatch,
1131
1132
    tp_size: int,
    pp_size: int,
1133
    test_target: Any,
1134
) -> None:
1135
1136
    import ray

1137
1138
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
1139
1140
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
1141
1142
1143
1144
1145
1146
    # NOTE: Force ray not to use gitignore file as excluding, otherwise
    # it will not move .so files to working dir.
    # So we have to manually add some of large directories
    os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1"
    ray.init(
        runtime_env={
1147
            "working_dir": VLLM_PATH,
1148
            "excludes": [
1149
1150
1151
1152
1153
1154
1155
1156
1157
                "build",
                ".git",
                "cmake-build-*",
                "shellcheck",
                "dist",
                "ep_kernels_workspace",
            ],
        }
    )
1158
1159
1160
1161
1162

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
1163
1164
1165
1166
1167
1168
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
1169
1170
            ),
        )
1171
1172
1173
    ray.get(refs)

    ray.shutdown()
1174
1175
1176


@contextmanager
1177
def error_on_warning(category: type[Warning] = Warning):
1178
1179
    """
    Within the scope of this context manager, tests will fail if any warning
1180
    of the given category is emitted.
1181
1182
    """
    with warnings.catch_warnings():
1183
        warnings.filterwarnings("error", category=category)
1184
1185

        yield
1186
1187


1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
def get_physical_device_indices(devices):
    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
    if visible_devices is None:
        return devices

    visible_indices = [int(x) for x in visible_devices.split(",")]
    index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
    return [index_mapping[i] for i in devices if i in index_mapping]


1198
@_nvml()
1199
1200
1201
def wait_for_gpu_memory_to_clear(
    *,
    devices: list[int],
1202
1203
    threshold_bytes: int | None = None,
    threshold_ratio: float | None = None,
1204
1205
    timeout_s: float = 120,
) -> None:
1206
    assert threshold_bytes is not None or threshold_ratio is not None
1207
1208
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
1209
    devices = get_physical_device_indices(devices)
1210
1211
    start_time = time.time()
    while True:
1212
        output: dict[int, str] = {}
1213
        output_raw: dict[int, tuple[float, float]] = {}
1214
        for device in devices:
1215
            if current_platform.is_rocm():
1216
1217
1218
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
1219
                gb_total = mem_info["vram_total"] / 2**10
1220
1221
1222
1223
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
1224
1225
                gb_total = mem_info.total / 2**30
            output_raw[device] = (gb_used, gb_total)
1226
            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
1227

1228
        print("gpu memory used/total (GiB): ", end="")
1229
        for k, v in output.items():
1230
1231
            print(f"{k}={v}; ", end="")
        print("")
1232

1233
1234
        if threshold_bytes is not None:
            is_free = lambda used, total: used <= threshold_bytes / 2**30
1235
            threshold = f"{threshold_bytes / 2**30} GiB"
1236
1237
1238
1239
        else:
            is_free = lambda used, total: used / total <= threshold_ratio
            threshold = f"{threshold_ratio:.2f}"

1240
        dur_s = time.time() - start_time
1241
        if all(is_free(used, total) for used, total in output_raw.values()):
1242
1243
1244
1245
            print(
                f"Done waiting for free GPU memory on devices {devices=} "
                f"({threshold=}) {dur_s=:.02f}"
            )
1246
1247
1248
            break

        if dur_s >= timeout_s:
1249
1250
1251
1252
            raise ValueError(
                f"Memory of devices {devices=} not free after "
                f"{dur_s=:.02f} ({threshold=})"
            )
1253
1254

        time.sleep(5)
1255
1256


1257
1258
1259
_P = ParamSpec("_P")


1260
def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
1261
1262
1263
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
1264

1265
    @functools.wraps(func)
1266
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
1267
1268
1269
1270
        # Make the process the leader of its own process group
        # to avoid sending SIGTERM to the parent process
        os.setpgrp()
        from _pytest.outcomes import Skipped
1271
1272
1273

        # Create a unique temporary file to store exception info from child
        # process. Use test function name and process ID to avoid collisions.
1274
1275
        with (
            tempfile.NamedTemporaryFile(
1276
                delete=False,
1277
                mode="w+b",
1278
                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
1279
1280
1281
1282
                suffix=".exc",
            ) as exc_file,
            ExitStack() as delete_after,
        ):
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
            exc_file_path = exc_file.name
            delete_after.callback(os.remove, exc_file_path)

            pid = os.fork()
            print(f"Fork a new process to run a test {pid}")
            if pid == 0:
                # Parent process responsible for deleting, don't delete
                # in child.
                delete_after.pop_all()
                try:
                    func(*args, **kwargs)
                except Skipped as e:
                    # convert Skipped to exit code 0
                    print(str(e))
                    os._exit(0)
                except Exception as e:
                    import traceback
1300

1301
1302
1303
1304
1305
1306
1307
                    tb_string = traceback.format_exc()

                    # Try to serialize the exception object first
                    exc_to_serialize: dict[str, Any]
                    try:
                        # First, try to pickle the actual exception with
                        # its traceback.
1308
                        exc_to_serialize = {"pickled_exception": e}
1309
1310
1311
1312
1313
                        # Test if it can be pickled
                        cloudpickle.dumps(exc_to_serialize)
                    except (Exception, KeyboardInterrupt):
                        # Fall back to string-based approach.
                        exc_to_serialize = {
1314
1315
1316
                            "exception_type": type(e).__name__,
                            "exception_msg": str(e),
                            "traceback": tb_string,
1317
1318
                        }
                    try:
1319
                        with open(exc_file_path, "wb") as f:
1320
1321
1322
1323
1324
1325
1326
                            cloudpickle.dump(exc_to_serialize, f)
                    except Exception:
                        # Fallback: just print the traceback.
                        print(tb_string)
                    os._exit(1)
                else:
                    os._exit(0)
1327
            else:
1328
1329
1330
                pgid = os.getpgid(pid)
                _pid, _exitcode = os.waitpid(pid, 0)
                # ignore SIGTERM signal itself
1331
                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
1332
1333
1334
1335
1336
1337
1338
1339
                # kill all child processes
                os.killpg(pgid, signal.SIGTERM)
                # restore the signal handler
                signal.signal(signal.SIGTERM, old_signal_handler)
                if _exitcode != 0:
                    # Try to read the exception from the child process
                    exc_info = {}
                    if os.path.exists(exc_file_path):
1340
1341
1342
1343
                        with (
                            contextlib.suppress(Exception),
                            open(exc_file_path, "rb") as f,
                        ):
1344
1345
                            exc_info = cloudpickle.load(f)

1346
1347
1348
                    if (
                        original_exception := exc_info.get("pickled_exception")
                    ) is not None:
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
                        # Re-raise the actual exception object if it was
                        # successfully pickled.
                        assert isinstance(original_exception, Exception)
                        raise original_exception

                    if (original_tb := exc_info.get("traceback")) is not None:
                        # Use string-based traceback for fallback case
                        raise AssertionError(
                            f"Test {func.__name__} failed when called with"
                            f" args {args} and kwargs {kwargs}"
                            f" (exit code: {_exitcode}):\n{original_tb}"
                        ) from None

                    # Fallback to the original generic error
                    raise AssertionError(
                        f"function {func.__name__} failed when called with"
                        f" args {args} and kwargs {kwargs}"
1366
1367
                        f" (exit code: {_exitcode})"
                    ) from None
1368
1369

    return wrapper
1370
1371


1372
1373
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
    """Decorator to spawn a new process for each test function."""
1374
1375
1376
1377

    @functools.wraps(f)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
        # Check if we're already in a subprocess
1378
        if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
1379
1380
1381
1382
            # If we are, just run the function directly
            return f(*args, **kwargs)

        import torch.multiprocessing as mp
1383

1384
        with suppress(RuntimeError):
1385
            mp.set_start_method("spawn")
1386
1387
1388
1389
1390
1391

        # Get the module
        module_name = f.__module__

        # Create a process with environment variable set
        env = os.environ.copy()
1392
        env["RUNNING_IN_SUBPROCESS"] = "1"
1393
1394
1395
1396
1397
1398
1399

        with tempfile.TemporaryDirectory() as tempdir:
            output_filepath = os.path.join(tempdir, "new_process.tmp")

            # `cloudpickle` allows pickling complex functions directly
            input_bytes = cloudpickle.dumps((f, output_filepath))

1400
1401
1402
1403
1404
            repo_root = str(VLLM_PATH.resolve())

            env = dict(env or os.environ)
            env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")

1405
1406
            cmd = [sys.executable, "-m", f"{module_name}"]

1407
1408
1409
            returned = subprocess.run(
                cmd, input=input_bytes, capture_output=True, env=env
            )
1410
1411
1412
1413
1414
1415

            # check if the subprocess is successful
            try:
                returned.check_returncode()
            except Exception as e:
                # wrap raised exception to provide more information
1416
1417
1418
                raise RuntimeError(
                    f"Error raised in subprocess:\n{returned.stderr.decode()}"
                ) from e
1419
1420
1421
1422
1423

    return wrapper


def create_new_process_for_each_test(
1424
    method: Literal["spawn", "fork"] | None = None,
1425
1426
1427
1428
) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
    """Creates a decorator that runs each test function in a new process.

    Args:
1429
        method: The process creation method. Can be either "spawn" or "fork".
1430
1431
               If not specified, it defaults to "spawn" on ROCm and XPU
               platforms and "fork" otherwise.
1432
1433
1434
1435
1436

    Returns:
        A decorator to run test functions in separate processes.
    """
    if method is None:
1437
1438
        use_spawn = current_platform.is_rocm() or current_platform.is_xpu()
        method = "spawn" if use_spawn else "fork"
1439

1440
    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
1441
1442
1443
1444
1445
1446
1447

    if method == "fork":
        return fork_new_process_for_each_test

    return spawn_new_process_for_each_test


1448
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
1449
1450
1451
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
1452

1453
1454
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
1455
1456
    """
    try:
1457
        if current_platform.is_cpu():
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
            memory_gb = 0
        else:
            memory_gb = current_platform.get_device_total_memory() / GB_bytes
    except Exception as e:
        warnings.warn(
            f"An error occurred when finding the available memory: {e}",
            stacklevel=2,
        )
        memory_gb = 0

1468
    return pytest.mark.skipif(
1469
        memory_gb < min_gb,
1470
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
1471
1472
    )

1473

1474
1475
1476
1477
1478
1479
1480
requires_fp8 = pytest.mark.skipif(
    not current_platform.supports_fp8(),
    reason="FP8 is not supported on this GPU (requires Hopper or "
    "Ada architecture, compute capability 8.9+)",
)


1481
1482
1483
1484
1485
1486
1487
def large_gpu_test(*, min_gb: int):
    """
    Decorate a test to be skipped if no GPU is available or it does not have
    sufficient memory.

    Currently, the CI machine uses L4 GPU which has 24 GB VRAM.
    """
1488
    mark = large_gpu_mark(min_gb)
1489

1490
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1491
        return mark(f)
1492
1493
1494
1495

    return wrapper


1496
1497
1498
def multi_gpu_marks(*, num_gpus: int):
    """Get a collection of pytest marks to apply for `@multi_gpu_test`."""
    test_selector = pytest.mark.distributed(num_gpus=num_gpus)
1499
1500
1501
1502
1503
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

1504
1505
1506
1507
1508
1509
1510
1511
1512
    return [test_selector, test_skipif]


def multi_gpu_test(*, num_gpus: int):
    """
    Decorate a test to be run only when multiple GPUs are available.
    """
    marks = multi_gpu_marks(num_gpus=num_gpus)

1513
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1514
        func = create_new_process_for_each_test()(f)
1515
1516
1517
1518
        for mark in reversed(marks):
            func = mark(func)

        return func
1519
1520
1521
1522

    return wrapper


1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None):
    """
    Mark a test to only run when the GPU count falls within [min_gpus, max_gpus].

    Examples:
        @gpu_tier_mark(min_gpus=2)          # only on multi-GPU
        @gpu_tier_mark(max_gpus=1)          # only on single-GPU
        @gpu_tier_mark(min_gpus=2, max_gpus=4)  # 2-4 GPUs only
    """
    gpu_count = cuda_device_count_stateless()
    marks = []

    if min_gpus > 1:
        marks.append(pytest.mark.distributed(num_gpus=min_gpus))

    reasons = []
    if gpu_count < min_gpus:
        reasons.append(f"Need at least {min_gpus} GPUs (have {gpu_count})")
    if max_gpus is not None and gpu_count > max_gpus:
        reasons.append(f"Need at most {max_gpus} GPUs (have {gpu_count})")

    if reasons:
        marks.append(pytest.mark.skipif(True, reason="; ".join(reasons)))

    return marks


def single_gpu_only(f=None):
    """Skip this test when running in a multi-GPU environment."""
    marks = gpu_tier_mark(max_gpus=1)

    def wrapper(func):
        for mark in reversed(marks):
            func = mark(func)
        return func

    return wrapper(f) if f is not None else wrapper


def multi_gpu_only(*, num_gpus: int = 2):
    """Skip this test when running on fewer than num_gpus GPUs."""
    marks = gpu_tier_mark(min_gpus=num_gpus)

    def wrapper(f):
        for mark in reversed(marks):
            f = mark(f)
        return f

    return wrapper


1574
async def completions_with_server_args(
1575
    prompts: list[str],
1576
    model_name: str,
1577
    server_cli_args: list[str],
1578
    num_logprobs: int | None,
1579
    max_wait_seconds: int = 240,
1580
    max_tokens: int | list = 5,
1581
) -> list[Completion]:
1582
    """Construct a remote OpenAI server, obtain an async client to the
1583
1584
1585
1586
1587
1588
1589
1590
1591
    server & invoke the completions API to obtain completions.

    Args:
      prompts: test prompts
      model_name: model to spin up on the vLLM server
      server_cli_args: CLI args for starting the server
      num_logprobs: Number of logprobs to report (or `None`)
      max_wait_seconds: timeout interval for bringing up server.
                        Default: 240sec
1592
1593
1594
      max_tokens: max_tokens value for each of the given input prompts.
        if only one max_token value is given, the same value is used
        for all the prompts.
1595
1596
1597

    Returns:
      OpenAI Completion instance
1598
    """
1599

1600
1601
1602
1603
1604
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

1605
    outputs = None
1606
1607
1608
    with RemoteOpenAIServer(
        model_name, server_cli_args, max_wait_seconds=max_wait_seconds
    ) as server:
1609
        client = server.get_async_client()
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
        outputs = [
            client.completions.create(
                model=model_name,
                prompt=[p],
                temperature=0,
                stream=False,
                max_tokens=max_tok,
                logprobs=num_logprobs,
            )
            for p, max_tok in zip(prompts, max_tokens)
        ]
1621
1622
        outputs = await asyncio.gather(*outputs)

1623
    assert outputs is not None, "Completion API call failed."
1624
1625
1626
1627

    return outputs


1628
def get_client_text_generations(completions: list[Completion]) -> list[str]:
1629
    """Extract generated tokens from the output of a
1630
    request made to an Open-AI-protocol completions endpoint.
1631
    """
1632
1633
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
1634
1635
1636


def get_client_text_logprob_generations(
1637
1638
1639
    completions: list[Completion],
) -> list[TextTextLogprobs]:
    """Operates on the output of a request made to an Open-AI-protocol
1640
    completions endpoint; obtains top-rank logprobs for each token in
1641
    each {class}`SequenceGroup`
1642
    """
1643
    text_generations = get_client_text_generations(completions)
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
    text = "".join(text_generations)
    return [
        (
            text_generations,
            text,
            (None if x.logprobs is None else x.logprobs.top_logprobs),
        )
        for completion in completions
        for x in completion.choices
    ]
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664


def has_module_attribute(module_name, attribute_name):
    """
    Helper function to check if a module has a specific attribute.
    """
    try:
        module = importlib.import_module(module_name)
        return hasattr(module, attribute_name)
    except ImportError:
        return False
1665
1666
1667
1668


def get_attn_backend_list_based_on_platform() -> list[str]:
    if current_platform.is_cuda():
1669
        return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"]
1670
    elif current_platform.is_rocm():
1671
        attn_backend_list = ["TRITON_ATTN"]
1672
1673
        try:
            import aiter  # noqa: F401
1674

1675
            attn_backend_list.append("ROCM_AITER_FA")
1676
        except Exception:
1677
            print("Skip ROCM_AITER_FA on ROCm as aiter is not installed")
1678
1679

        return attn_backend_list
1680
1681
    elif current_platform.is_xpu():
        return ["FLASH_ATTN", "TRITON_ATTN"]
1682
1683
    else:
        raise ValueError("Unsupported platform")
1684
1685
1686
1687
1688


@contextmanager
def override_cutlass_fp8_supported(value: bool):
    with patch(
1689
1690
1691
        "vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
        return_value=value,
    ):
1692
        yield
1693
1694


1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
def disable_aiter_plain_rmsnorm(monkeypatch) -> None:
    """Patch dispatch_rocm_rmsnorm_func so the plain (non-fused) rms_norm path
    always uses the native float32 kernel for the duration of a test.

    The fused path (rms_norm2d_with_add, selected when with_fused_add=True) is
    left on AITER -- only the plain path is redirected to native.

    AITER's plain rms_norm accumulates variance in bfloat16 (~1 ULP/call),
    which drifts the KV cache over many decode steps. This drift is irrelevant
    for a trained model (rank-1/rank-2 gap ~1-3 nats >> 1 ULP), but breaks
    logprob comparison tests with randomly-initialised models like
    TitanML/tiny-mixtral whose rank-1/rank-2 gap is only O(1/sqrt(V)) ~0.006
    nats -- smaller than the accumulated per-step error.
    """
    import torch

    import vllm.model_executor.layers.layernorm as _ln_mod
    from vllm.model_executor.layers.layernorm import rms_norm as _native

    _orig = _ln_mod.dispatch_rocm_rmsnorm_func

    def _native_plain(
        with_fused_add: bool, dtype: torch.dtype, use_aiter: bool = False
    ):
        if (
            use_aiter
            and not with_fused_add
            and dtype in (torch.float16, torch.bfloat16)
        ):
            return _native
        return _orig(with_fused_add, dtype, use_aiter)

    monkeypatch.setattr(_ln_mod, "dispatch_rocm_rmsnorm_func", _native_plain)


1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
    """
    Generate prompts which a bunch of assignments,
    then asking for the value of one of them.
    The prompt is just under 10k tokens; sliding window is 4k
    so the answer is outside sliding window, but should still be correct.
    Args:
        batch_size: number of prompts to generate
        ln_range: an argument to control the length of the prompt
    """
    prompts: list[str] = []
    answer: list[int] = []
    indices: list[int] = []
    random.seed(1)
    for _ in range(batch_size):
        idx = random.randint(30, 90)
        indices.append(idx)
1747
1748
        prompt = (
            "```python\n# We set a number of variables, "
1749
            f"x{idx} will be important later\n"
1750
        )
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
        ln = random.randint(*ln_range)
        for k in range(30, ln):
            v = random.randint(10, 99)
            if k == idx:
                answer.append(v)
            prompt += f"x{k} = {v}\n"
        prompt += f"# Now, we check the value of x{idx}:\n"
        prompt += f"assert x{idx} == "
        prompts.append(prompt)
    return prompts, answer, indices


1763
1764
1765
def check_answers(
    indices: list[int], answer: list[int], outputs: list[str], accept_rate: float = 0.7
):
1766
1767
1768
1769
1770
1771
1772
1773
1774
    answer2 = [int(text[0:2].strip()) for text in outputs]
    print(list(zip(indices, zip(answer, answer2))))
    numok = 0
    for a1, a2 in zip(answer, answer2):
        if a1 == a2:
            numok += 1
    frac_ok = numok / len(answer)
    print(f"Num OK: {numok}/{len(answer)} {frac_ok}")
    assert frac_ok >= accept_rate
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794


def flat_product(*iterables: Iterable[Any]):
    """
    Flatten lists of tuples of the cartesian product.
    Useful when we want to avoid nested tuples to allow
    test params to be unpacked directly from the decorator.

    Example:
    flat_product([(1, 2), (3, 4)], ["a", "b"]) ->
    [
      (1, 2, "a"),
      (1, 2, "b"),
      (3, 4, "a"),
      (3, 4, "b"),
    ]
    """
    for element in itertools.product(*iterables):
        normalized = (e if isinstance(e, tuple) else (e,) for e in element)
        yield tuple(itertools.chain(*normalized))
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908


class TestFP8Layer(torch.nn.Module):
    """
    Test helper for FP8 linear operations. Creates random weights and scales
    based on quantization configuration.

    Args:
        weight_shape: Shape of the weight tensor (out_features, in_features).
        activation_quant_key: Activation quantization configuration.
        weight_quant_key: Weight quantization configuration.
        out_dtype: Output dtype. Defaults to current default dtype.
        force_kernel: Optional kernel to force use of specific implementation.
    """

    def __init__(
        self,
        weight_shape: tuple[int, int],
        activation_quant_key: QuantKey,
        weight_quant_key: QuantKey,
        out_dtype: torch.dtype | None = None,
        device: torch.device | None = None,
        force_kernel: FP8ScaledMMLinearKernel | None = None,
    ):
        super().__init__()
        per_tensor_weights = weight_quant_key.scale.group_shape.is_per_tensor()
        is_static_activation_scale = activation_quant_key.scale.static
        weight_scale_shape = (1,) if per_tensor_weights else (weight_shape[0], 1)

        self.weight_scale = torch.rand(
            weight_scale_shape, dtype=torch.float32, device=device
        )
        self.input_scale = (
            torch.rand(1, dtype=torch.float32, device=device)
            if is_static_activation_scale
            else None
        )
        self.weight = torch.rand(weight_shape, device=device).to(dtype=FP8_DTYPE).t()
        self.input_scale_ub = None

        out_dtype = torch.get_default_dtype() if out_dtype is None else out_dtype

        self.kernel = init_fp8_linear_kernel(
            activation_quant_key=activation_quant_key,
            weight_quant_key=weight_quant_key,
            out_dtype=out_dtype,
            force_kernel=force_kernel,
        )

    def is_quant_fp8_enabled(self) -> bool:
        return self.kernel.quant_fp8.enabled()

    def forward(
        self, y: torch.Tensor, bias: torch.Tensor | None = None
    ) -> torch.Tensor:
        return self.kernel.apply_weights(self, y, bias)


# TODO: Drop TestBlockFP8Layer in favour of a unified TestFP8Layer
# after refactoring W8A8BlockFp8LinearOp.
# https://github.com/vllm-project/vllm/issues/31818
class TestBlockFP8Layer:
    """
    Test helper for blockwise FP8 linear operations. Creates random weights
    and scales for W8A8BlockFp8LinearOp.

    This is a workaround until W8A8BlockFp8LinearOp implements the kernel
    abstraction (ScaledMMLinearKernel) for blockwise quantization.

    Args:
        weight_shape: Shape of the weight tensor (out_features, in_features).
        group_shape: Blockwise quantization group shape.
        cutlass_block_fp8_supported: Whether CUTLASS blockwise FP8 is available.
        use_aiter_and_is_supported: Whether to use aiter quantization ops.
        transpose_weights: Whether to transpose weights after creation.
    """

    def __init__(
        self,
        weight_shape: tuple[int, int],
        group_shape: GroupShape,
        cutlass_block_fp8_supported: bool = False,
        use_aiter_and_is_supported: bool = False,
        transpose_weights: bool = False,
    ):
        weight_scale_shape = weight_shape[0] // group_shape[1]
        self.weight_scale = torch.rand(
            (weight_scale_shape, weight_scale_shape), dtype=torch.float32
        )
        self.weight = torch.rand(weight_shape).to(dtype=FP8_DTYPE)
        self.input_scale = None
        if transpose_weights:
            self.weight = self.weight.t()

        self.linear_op = W8A8BlockFp8LinearOp(
            weight_group_shape=GroupShape(group_shape[1], group_shape[1]),
            act_quant_group_shape=group_shape,
            cutlass_block_fp8_supported=cutlass_block_fp8_supported,
            use_aiter_and_is_supported=use_aiter_and_is_supported,
        )

    def __call__(
        self, y: torch.Tensor, bias: torch.Tensor | None = None
    ) -> torch.Tensor:
        return self.linear_op.apply(
            input=y,
            weight=self.weight,
            weight_scale=self.weight_scale,
            input_scale=self.input_scale,
            bias=bias,
        )

    def is_quant_fp8_enabled(self) -> bool:
        return self.linear_op.input_quant_op.enabled()