utils.py 56.8 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
class RemoteVLLMServer:
    """Base class for launching vLLM server subprocesses for testing.

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

120
    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
121
122
123
124
125
    proc: subprocess.Popen

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

127
    def _start_server(
128
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
129
130
    ) -> None:
        """Subclasses override this method to customize server process launch"""
131
        raise NotImplementedError
132

133
134
135
136
137
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
        *,
138
        env_dict: dict[str, str] | None = None,
139
        seed: int = 0,
140
        auto_port: bool = True,
141
142
        max_wait_seconds: float | None = None,
        override_hf_configs: dict[str, Any] | None = None,
143
    ) -> None:
144
        if auto_port:
145
            if "-p" in vllm_serve_args or "--port" in vllm_serve_args:
146
147
148
                raise ValueError(
                    "You have manually specified the port when `auto_port=True`."
                )
149

150
151
152
            # No need for a port if using unix sockets
            if "--uds" not in vllm_serve_args:
                # Don't mutate the input args
153
                vllm_serve_args = vllm_serve_args + ["--port", str(get_open_port())]
154
155
        if seed is not None:
            if "--seed" in vllm_serve_args:
156
157
158
                raise ValueError(
                    f"You have manually specified the seed when `seed={seed}`."
                )
159
160

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

162
163
164
        if override_hf_configs is not None:
            vllm_serve_args = vllm_serve_args + [
                "--hf-overrides",
165
                json.dumps(override_hf_configs),
166
167
            ]

168
        parser = FlexibleArgumentParser(description="vLLM's remote server.")
169
        subparsers = parser.add_subparsers(required=False, dest="subparser")
170
        parser = self._create_cli_subcommand().subparser_init(subparsers)
171
        args = parser.parse_args(["--model", model, *vllm_serve_args])
172
173
174
175
176
        self.uds = args.uds
        if args.uds:
            self.host = None
            self.port = None
        else:
177
            self.host = str(args.host or "127.0.0.1")
178
            self.port = int(args.port)
179

180
181
182
        self.show_hidden_metrics = (
            getattr(args, "show_hidden_metrics_for_version", None) is not None
        )
183

184
185
186
187
        # download the model before starting the server to avoid timeout
        is_local = os.path.isdir(model)
        if not is_local:
            engine_args = AsyncEngineArgs.from_cli_args(args)
188
189
190
191
192
            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)
193

194
195
196
197
198
199
        # 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(
200
201
                f"[{type(self).__name__}] GPU memory before server start: "
                f"{pre_gb:.2f} GB"
202
203
            )

204
        self._start_server(model, vllm_serve_args, env_dict)
205
        max_wait_seconds = max_wait_seconds or 360
206
        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
207
208
209
210
211

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
212
        pid = self.proc.pid
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228

        # 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

        # Phase 1: graceful SIGTERM to the entire process group
        if pgid is not None:
            with contextlib.suppress(ProcessLookupError, OSError):
                os.killpg(pgid, signal.SIGTERM)
                print(f"[RemoteOpenAIServer] Sent SIGTERM to process group {pgid}")
        else:
            self.proc.terminate()

229
        try:
230
231
            self.proc.wait(timeout=15)
            print(f"[RemoteOpenAIServer] Server {pid} terminated gracefully")
232
        except subprocess.TimeoutExpired:
233
            # Phase 2: SIGKILL the entire process group
234
235
            print(
                f"[RemoteOpenAIServer] Server {pid} did not respond "
236
                "to SIGTERM, sending SIGKILL to process group"
237
            )
238
239
240
241
242
243
            if pgid is not None:
                with contextlib.suppress(ProcessLookupError, OSError):
                    os.killpg(pgid, signal.SIGKILL)
            else:
                self.proc.kill()

244
            try:
245
                self.proc.wait(timeout=10)
246
                print(f"[RemoteOpenAIServer] Server {pid} killed")
247
248
249
250
251
252
            except subprocess.TimeoutExpired:
                # Phase 3: last resort - find and kill any orphaned children
                self._kill_orphaned_children(pid)

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

255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
    def _kill_orphaned_children(self, parent_pid: int) -> None:
        """Best-effort cleanup of any lingering child processes."""
        try:
            import psutil

            parent = psutil.Process(parent_pid)
            children = parent.children(recursive=True)
            for child in children:
                print(
                    f"[RemoteOpenAIServer] Killing orphaned child "
                    f"pid={child.pid} name={child.name()}"
                )
                child.kill()
            psutil.wait_procs(children, timeout=5)
        except Exception as e:
            # psutil may not be installed, or processes already gone
            print(f"[RemoteOpenAIServer] Orphan cleanup failed: {e}")
            # Fallback: try to kill by pgid one more time
            with contextlib.suppress(ProcessLookupError, OSError):
                os.killpg(parent_pid, signal.SIGKILL)

276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
    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

301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
    def _wait_for_gpu_memory_release(self, timeout: float = 60.0):
        """Wait for GPU memory to drop back toward pre-server levels.

        Two-phase strategy:
          1. Try to wait for memory to return close to pre-server baseline.
          2. If that doesn't happen, fall back to waiting for stabilization
             and log a warning (the next server might still OOM).
        """
        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

319
        start = time.time()
320
        last_used: float | None = None
321
322
323
324
325
326
327
328
        stable_count = 0

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

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

329
330
331
            used_gb = used / 1e9
            target_gb = target / 1e9
            elapsed = time.time() - start
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
            # Phase 1: memory dropped to near baseline - we're done.
            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

            # Phase 2 (after 40s): fall back to stabilization check.
            # This handles cases where another process is using GPU memory
            # and we'll never reach baseline.
            if elapsed > 40.0 and last_used is not None:
                delta = abs(used - last_used)
                if delta < 200 * 1024 * 1024:  # 200 MB
                    stable_count += 1
                    if stable_count >= 3:
                        print(
                            f"[RemoteOpenAIServer] WARNING: GPU memory "
                            f"stabilized at {used_gb:.2f} GB "
                            f"(target was {target_gb:.2f} GB). "
                            f"Proceeding - next server may OOM."
                        )
                        return
                else:
                    stable_count = 0

            last_used = used
            time.sleep(1.0)

        # Timeout - log clearly so CI failures are diagnosable
        final_used = self._get_gpu_memory_used()
        final_gb = final_used / 1e9 if final_used else 0.0
366
        raise RuntimeError(
367
368
369
370
371
            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."
372
        )
373

374
    def _poll(self) -> int | None:
375
376
377
        """Subclasses override this method to customize process polling"""
        return self.proc.poll()

378
379
380
    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
381
382
383
384
385
        client = (
            httpx.Client(transport=httpx.HTTPTransport(uds=self.uds))
            if self.uds
            else requests
        )
386
387
        while True:
            try:
388
                if client.get(url).status_code == 200:
389
                    break
390
391
392
393
394
            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`.
395
                result = self._poll()
396
                if result is not None and result != 0:
397
                    raise RuntimeError("Server exited unexpectedly.") from None
398
399
400

                time.sleep(0.5)
                if time.time() - start > timeout:
401
                    raise RuntimeError("Server failed to start in time.") from None
402
403
404

    @property
    def url_root(self) -> str:
405
406
407
408
409
        return (
            f"http://{self.uds.split('/')[-1]}"
            if self.uds
            else f"http://{self.host}:{self.port}"
        )
410
411
412
413

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

414
415
416
    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
417
418
419
        return openai.OpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
420
421
            max_retries=0,
            **kwargs,
422
423
        )

424
    def get_async_client(self, **kwargs):
425
426
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
427
428
429
430
431
432
        return openai.AsyncOpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
            max_retries=0,
            **kwargs,
        )
433

434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
    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
        )

451

452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
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,
        )

    def _wait_for_gpu_memory_release(self, timeout: float = 30.0):
        pass  # No GPU used


508
509
510
class RemoteOpenAIServerCustom(RemoteOpenAIServer):
    """Launch test server with custom child process"""

511
    def _start_server(
512
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
513
    ) -> None:
514
        self.proc: Process = Process(
515
516
            target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)
        )  # type: ignore[assignment]
517
518
        self.proc.start()

519
520
521
522
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
523
        child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],
524
        *,
525
        env_dict: dict[str, str] | None = None,
526
        seed: int = 0,
527
        auto_port: bool = True,
528
        max_wait_seconds: float | None = None,
529
    ) -> None:
530
531
532
        """Store custom child process function then invoke superclass
        constructor which will indirectly launch it."""
        self.child_process_fxn = child_process_fxn
533
534
535
536
537
538
539
540
        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,
        )
541

542
    def _poll(self) -> int | None:
543
544
545
546
547
548
549
550
551
552
        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()


553
554
555
556
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
557
    token_ids: list[int],
558
559
560
561
):
    results = []

    # test with text prompt
562
563
564
565
566
567
568
569
570
571
572
573
    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,
        }
    )
574
575
576
577
578
579
580
581
582

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

583
584
585
586
587
588
589
590
    results.append(
        {
            "test": "token_ids",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
591
592

    # test seeded random sampling
593
594
595
596
597
598
599
600
601
602
603
604
    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,
        }
    )
605
606

    # test seeded random sampling with multiple prompts
607
608
609
610
611
612
613
614
615
616
617
618
    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,
        }
    )
619
620
621
622
623
624
625
626
627

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

628
629
630
631
632
633
634
    results.append(
        {
            "test": "simple_list",
            "text0": batch.choices[0].text,
            "text1": batch.choices[1].text,
        }
    )
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650

    # 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

651
652
653
654
655
656
    results.append(
        {
            "test": "streaming",
            "texts": texts,
        }
    )
657
658
659
660

    return results


661
662
663
664
665
666
667
668
def _test_completion_close(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

    # test with text prompt
669
670
671
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=1, logprobs=5, temperature=0.0
    )
672

673
674
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
675

676
677
678
679
680
681
    results.append(
        {
            "test": "completion_close",
            "logprobs": logprobs,
        }
    )
682
683
684
685

    return results


686
687
688
689
690
691
692
def _test_chat(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

693
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
694
695

    # test with text prompt
696
697
698
699
700
701
702
703
704
705
706
707
    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,
        }
    )
708
709
710
711

    return results


712
713
714
715
716
717
718
719
720
721
722
723
724
725
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",
    )

726
727
728
729
730
731
732
    results.append(
        {
            "test": "single_embedding",
            "embedding": embeddings.data[0].embedding,
            "usage": embeddings.usage,
        }
    )
733
734
735
736

    return results


737
738
739
740
741
742
743
744
def _test_image_text(
    client: openai.OpenAI,
    model_name: str,
    image_url: str,
):
    results = []

    # test pure text input
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
    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,
    )
762
763
764
765
766
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

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

767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
    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,
    )
792
793
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

794
795
796
797
798
799
    results.append(
        {
            "test": "text_image",
            "logprobs": top_logprobs,
        }
    )
800
801
802
803

    return results


804
805
806
807
def compare_two_settings(
    model: str,
    arg1: list[str],
    arg2: list[str],
808
809
    env1: dict[str, str] | None = None,
    env2: dict[str, str] | None = None,
810
811
    *,
    method: str = "generate",
812
    max_wait_seconds: float | None = None,
813
) -> None:
814
    """
815
816
817
818
819
820
821
822
823
    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.
824
825
    """

826
827
828
829
830
831
832
833
834
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


835
836
837
def compare_all_settings(
    model: str,
    all_args: list[list[str]],
838
    all_envs: list[dict[str, str] | None],
839
840
    *,
    method: str = "generate",
841
    max_wait_seconds: float | None = None,
842
) -> None:
843
844
845
846
847
848
849
850
851
    """
    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.
    """

852
    trust_remote_code = False
853
    for args in all_args:
854
855
856
857
858
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
859
    for args in all_args:
860
861
862
863
864
865
866
867
868
        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,
    )
869

870
871
872
873
874
875
876
    can_force_load_format = True

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

877
    prompt = "Hello, my name is"
878
    token_ids = tokenizer(prompt).input_ids
879
    ref_results: list = []
880
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
881
882
883
884
885
886
887
888
889
        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]
890
        compare_results: list = []
891
        results = ref_results if i == 0 else compare_results
892
893
894
        with RemoteOpenAIServer(
            model, args, env_dict=env, max_wait_seconds=max_wait_seconds
        ) as server:
895
896
897
898
899
900
            client = server.get_client()

            # test models list
            models = client.models.list()
            models = models.data
            served_model = models[0]
901
902
903
904
905
906
907
            results.append(
                {
                    "test": "models_list",
                    "id": served_model.id,
                    "root": served_model.root,
                }
            )
908

909
910
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
911
912
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
913
914
            elif method == "generate_chat":
                results += _test_chat(client, model, prompt)
915
916
            elif method == "generate_with_image":
                results += _test_image_text(
917
918
                    client,
                    model,
919
                    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
920
                )
921
922
923
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
924
                raise ValueError(f"Unknown method: {method}")
925

926
927
928
929
930
931
            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]
932
                for ref_result, compare_result in zip(ref_results, compare_results):
933
934
935
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
936
937
938
939
940
941
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
942
                            f"Embedding for {model=} are not the same.\n"
943
944
                            f"cosine_similarity={sim}\n"
                        )
945
946
                        del ref_result["embedding"]
                        del compare_result["embedding"]
947
948
949
950
951
                    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"
952
953
                        f"{compare_result=}\n"
                    )
954
955


956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
@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


986
987
988
989
990
991
992
def init_test_distributed_environment(
    tp_size: int,
    pp_size: int,
    rank: int,
    distributed_init_port: str,
    local_rank: int = -1,
) -> None:
993
994
995
996
997
998
999
    # 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,
1000
    )
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011

    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,
        )
1012
        ensure_model_parallel_initialized(tp_size, pp_size)
1013
1014
1015
1016
1017
1018
1019
1020
1021
    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,
            )
1022
            ensure_model_parallel_initialized(tp_size, pp_size)
1023
1024


1025
def multi_process_parallel(
1026
    monkeypatch: pytest.MonkeyPatch,
1027
1028
    tp_size: int,
    pp_size: int,
1029
    test_target: Any,
1030
) -> None:
1031
1032
    import ray

1033
1034
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
1035
1036
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
1037
1038
1039
1040
1041
1042
    # 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={
1043
            "working_dir": VLLM_PATH,
1044
            "excludes": [
1045
1046
1047
1048
1049
1050
1051
1052
1053
                "build",
                ".git",
                "cmake-build-*",
                "shellcheck",
                "dist",
                "ep_kernels_workspace",
            ],
        }
    )
1054
1055
1056
1057
1058

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
1059
1060
1061
1062
1063
1064
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
1065
1066
            ),
        )
1067
1068
1069
    ray.get(refs)

    ray.shutdown()
1070
1071
1072


@contextmanager
1073
def error_on_warning(category: type[Warning] = Warning):
1074
1075
    """
    Within the scope of this context manager, tests will fail if any warning
1076
    of the given category is emitted.
1077
1078
    """
    with warnings.catch_warnings():
1079
        warnings.filterwarnings("error", category=category)
1080
1081

        yield
1082
1083


1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
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]


1094
@_nvml()
1095
1096
1097
def wait_for_gpu_memory_to_clear(
    *,
    devices: list[int],
1098
1099
    threshold_bytes: int | None = None,
    threshold_ratio: float | None = None,
1100
1101
    timeout_s: float = 120,
) -> None:
1102
    assert threshold_bytes is not None or threshold_ratio is not None
1103
1104
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
1105
    devices = get_physical_device_indices(devices)
1106
1107
    start_time = time.time()
    while True:
1108
        output: dict[int, str] = {}
1109
        output_raw: dict[int, tuple[float, float]] = {}
1110
        for device in devices:
1111
            if current_platform.is_rocm():
1112
1113
1114
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
1115
                gb_total = mem_info["vram_total"] / 2**10
1116
1117
1118
1119
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
1120
1121
                gb_total = mem_info.total / 2**30
            output_raw[device] = (gb_used, gb_total)
1122
            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
1123

1124
        print("gpu memory used/total (GiB): ", end="")
1125
        for k, v in output.items():
1126
1127
            print(f"{k}={v}; ", end="")
        print("")
1128

1129
1130
        if threshold_bytes is not None:
            is_free = lambda used, total: used <= threshold_bytes / 2**30
1131
            threshold = f"{threshold_bytes / 2**30} GiB"
1132
1133
1134
1135
        else:
            is_free = lambda used, total: used / total <= threshold_ratio
            threshold = f"{threshold_ratio:.2f}"

1136
        dur_s = time.time() - start_time
1137
        if all(is_free(used, total) for used, total in output_raw.values()):
1138
1139
1140
1141
            print(
                f"Done waiting for free GPU memory on devices {devices=} "
                f"({threshold=}) {dur_s=:.02f}"
            )
1142
1143
1144
            break

        if dur_s >= timeout_s:
1145
1146
1147
1148
            raise ValueError(
                f"Memory of devices {devices=} not free after "
                f"{dur_s=:.02f} ({threshold=})"
            )
1149
1150

        time.sleep(5)
1151
1152


1153
1154
1155
_P = ParamSpec("_P")


1156
def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
1157
1158
1159
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
1160

1161
    @functools.wraps(func)
1162
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
1163
1164
1165
1166
        # 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
1167
1168
1169

        # Create a unique temporary file to store exception info from child
        # process. Use test function name and process ID to avoid collisions.
1170
1171
        with (
            tempfile.NamedTemporaryFile(
1172
                delete=False,
1173
                mode="w+b",
1174
                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
1175
1176
1177
1178
                suffix=".exc",
            ) as exc_file,
            ExitStack() as delete_after,
        ):
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
            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
1196

1197
1198
1199
1200
1201
1202
1203
                    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.
1204
                        exc_to_serialize = {"pickled_exception": e}
1205
1206
1207
1208
1209
                        # Test if it can be pickled
                        cloudpickle.dumps(exc_to_serialize)
                    except (Exception, KeyboardInterrupt):
                        # Fall back to string-based approach.
                        exc_to_serialize = {
1210
1211
1212
                            "exception_type": type(e).__name__,
                            "exception_msg": str(e),
                            "traceback": tb_string,
1213
1214
                        }
                    try:
1215
                        with open(exc_file_path, "wb") as f:
1216
1217
1218
1219
1220
1221
1222
                            cloudpickle.dump(exc_to_serialize, f)
                    except Exception:
                        # Fallback: just print the traceback.
                        print(tb_string)
                    os._exit(1)
                else:
                    os._exit(0)
1223
            else:
1224
1225
1226
                pgid = os.getpgid(pid)
                _pid, _exitcode = os.waitpid(pid, 0)
                # ignore SIGTERM signal itself
1227
                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
1228
1229
1230
1231
1232
1233
1234
1235
                # 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):
1236
1237
1238
1239
                        with (
                            contextlib.suppress(Exception),
                            open(exc_file_path, "rb") as f,
                        ):
1240
1241
                            exc_info = cloudpickle.load(f)

1242
1243
1244
                    if (
                        original_exception := exc_info.get("pickled_exception")
                    ) is not None:
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
                        # 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}"
1262
1263
                        f" (exit code: {_exitcode})"
                    ) from None
1264
1265

    return wrapper
1266
1267


1268
1269
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
    """Decorator to spawn a new process for each test function."""
1270
1271
1272
1273

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

        import torch.multiprocessing as mp
1279

1280
        with suppress(RuntimeError):
1281
            mp.set_start_method("spawn")
1282
1283
1284
1285
1286
1287

        # Get the module
        module_name = f.__module__

        # Create a process with environment variable set
        env = os.environ.copy()
1288
        env["RUNNING_IN_SUBPROCESS"] = "1"
1289
1290
1291
1292
1293
1294
1295

        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))

1296
1297
1298
1299
1300
            repo_root = str(VLLM_PATH.resolve())

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

1301
1302
            cmd = [sys.executable, "-m", f"{module_name}"]

1303
1304
1305
            returned = subprocess.run(
                cmd, input=input_bytes, capture_output=True, env=env
            )
1306
1307
1308
1309
1310
1311

            # check if the subprocess is successful
            try:
                returned.check_returncode()
            except Exception as e:
                # wrap raised exception to provide more information
1312
1313
1314
                raise RuntimeError(
                    f"Error raised in subprocess:\n{returned.stderr.decode()}"
                ) from e
1315
1316
1317
1318
1319

    return wrapper


def create_new_process_for_each_test(
1320
    method: Literal["spawn", "fork"] | None = None,
1321
1322
1323
1324
) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
    """Creates a decorator that runs each test function in a new process.

    Args:
1325
        method: The process creation method. Can be either "spawn" or "fork".
1326
1327
               If not specified, it defaults to "spawn" on ROCm and XPU
               platforms and "fork" otherwise.
1328
1329
1330
1331
1332

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

1336
    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
1337
1338
1339
1340
1341
1342
1343

    if method == "fork":
        return fork_new_process_for_each_test

    return spawn_new_process_for_each_test


1344
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
1345
1346
1347
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
1348

1349
1350
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
1351
1352
    """
    try:
1353
        if current_platform.is_cpu():
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
            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

1364
    return pytest.mark.skipif(
1365
        memory_gb < min_gb,
1366
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
1367
1368
    )

1369

1370
1371
1372
1373
1374
1375
1376
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+)",
)


1377
1378
1379
1380
1381
1382
1383
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.
    """
1384
    mark = large_gpu_mark(min_gb)
1385

1386
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1387
        return mark(f)
1388
1389
1390
1391

    return wrapper


1392
1393
1394
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)
1395
1396
1397
1398
1399
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

1400
1401
1402
1403
1404
1405
1406
1407
1408
    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)

1409
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1410
        func = create_new_process_for_each_test()(f)
1411
1412
1413
1414
        for mark in reversed(marks):
            func = mark(func)

        return func
1415
1416
1417
1418

    return wrapper


1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
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


1470
async def completions_with_server_args(
1471
    prompts: list[str],
1472
    model_name: str,
1473
    server_cli_args: list[str],
1474
    num_logprobs: int | None,
1475
    max_wait_seconds: int = 240,
1476
    max_tokens: int | list = 5,
1477
) -> list[Completion]:
1478
    """Construct a remote OpenAI server, obtain an async client to the
1479
1480
1481
1482
1483
1484
1485
1486
1487
    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
1488
1489
1490
      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.
1491
1492
1493

    Returns:
      OpenAI Completion instance
1494
    """
1495

1496
1497
1498
1499
1500
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

1501
    outputs = None
1502
1503
1504
    with RemoteOpenAIServer(
        model_name, server_cli_args, max_wait_seconds=max_wait_seconds
    ) as server:
1505
        client = server.get_async_client()
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
        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)
        ]
1517
1518
        outputs = await asyncio.gather(*outputs)

1519
    assert outputs is not None, "Completion API call failed."
1520
1521
1522
1523

    return outputs


1524
def get_client_text_generations(completions: list[Completion]) -> list[str]:
1525
    """Extract generated tokens from the output of a
1526
    request made to an Open-AI-protocol completions endpoint.
1527
    """
1528
1529
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
1530
1531
1532


def get_client_text_logprob_generations(
1533
1534
1535
    completions: list[Completion],
) -> list[TextTextLogprobs]:
    """Operates on the output of a request made to an Open-AI-protocol
1536
    completions endpoint; obtains top-rank logprobs for each token in
1537
    each {class}`SequenceGroup`
1538
    """
1539
    text_generations = get_client_text_generations(completions)
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
    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
    ]
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560


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
1561
1562
1563
1564


def get_attn_backend_list_based_on_platform() -> list[str]:
    if current_platform.is_cuda():
1565
        return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"]
1566
    elif current_platform.is_rocm():
1567
        attn_backend_list = ["TRITON_ATTN"]
1568
1569
        try:
            import aiter  # noqa: F401
1570

1571
            attn_backend_list.append("ROCM_AITER_FA")
1572
        except Exception:
1573
            print("Skip ROCM_AITER_FA on ROCm as aiter is not installed")
1574
1575

        return attn_backend_list
1576
1577
    elif current_platform.is_xpu():
        return ["FLASH_ATTN", "TRITON_ATTN"]
1578
1579
    else:
        raise ValueError("Unsupported platform")
1580
1581
1582
1583
1584


@contextmanager
def override_cutlass_fp8_supported(value: bool):
    with patch(
1585
1586
1587
        "vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
        return_value=value,
    ):
1588
        yield
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607


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)
1608
1609
        prompt = (
            "```python\n# We set a number of variables, "
1610
            f"x{idx} will be important later\n"
1611
        )
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
        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


1624
1625
1626
def check_answers(
    indices: list[int], answer: list[int], outputs: list[str], accept_rate: float = 0.7
):
1627
1628
1629
1630
1631
1632
1633
1634
1635
    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
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655


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))
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
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
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769


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()