utils.py 45.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
46
47
48
49
50
51
52
53
54
55
from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
    init_fp8_linear_kernel,
)
from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import (  # noqa: E501
    FP8ScaledMMLinearKernel,
)
from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    GroupShape,
    QuantKey,
)
56
from vllm.model_executor.model_loader import get_model_loader
57
from vllm.platforms import current_platform
58
from vllm.tokenizers import get_tokenizer
Cyrus Leung's avatar
Cyrus Leung committed
59
from vllm.utils.argparse_utils import FlexibleArgumentParser
60
from vllm.utils.mem_constants import GB_bytes
61
from vllm.utils.network_utils import get_open_port
62
from vllm.utils.torch_utils import cuda_device_count_stateless
63

64
65
FP8_DTYPE = current_platform.fp8_dtype()

66
if current_platform.is_rocm():
67
68
69
70
71
72
    from amdsmi import (
        amdsmi_get_gpu_vram_usage,
        amdsmi_get_processor_handles,
        amdsmi_init,
        amdsmi_shut_down,
    )
73
74
75
76
77
78
79
80

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

    @contextmanager
    def _nvml():
        try:
            nvmlInit()
            yield
        finally:
            nvmlShutdown()
96
97
98
99
100
else:

    @contextmanager
    def _nvml():
        yield
101

102

103
104
VLLM_PATH = Path(__file__).parent.parent
"""Path to root of the vLLM repository."""
105
106


107
108
class RemoteOpenAIServer:
    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
109

110
    def _start_server(
111
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
112
113
    ) -> None:
        """Subclasses override this method to customize server process launch"""
114
115
116
        env = os.environ.copy()
        # the current process might initialize cuda,
        # to be safe, we should use spawn method
117
        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
118
119
        if env_dict is not None:
            env.update(env_dict)
120
121
        serve_cmd = ["vllm", "serve", model, *vllm_serve_args]
        print(f"Launching RemoteOpenAIServer with: {' '.join(serve_cmd)}")
122
        print(f"Environment variables: {env}")
123
        self.proc: subprocess.Popen = subprocess.Popen(
124
            serve_cmd,
125
126
127
128
129
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
        )

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

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

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

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

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

177
        self.show_hidden_metrics = args.show_hidden_metrics_for_version is not None
178

179
180
181
182
        # 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)
183
184
185
186
187
            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)
188

189
        self._start_server(model, vllm_serve_args, env_dict)
190
        max_wait_seconds = max_wait_seconds or 240
191
        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
192
193
194
195
196
197

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.proc.terminate()
198
        try:
199
            self.proc.wait(8)
200
201
202
        except subprocess.TimeoutExpired:
            # force kill if needed
            self.proc.kill()
203

204
    def _poll(self) -> int | None:
205
206
207
        """Subclasses override this method to customize process polling"""
        return self.proc.poll()

208
209
210
    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
211
212
213
214
215
        client = (
            httpx.Client(transport=httpx.HTTPTransport(uds=self.uds))
            if self.uds
            else requests
        )
216
217
        while True:
            try:
218
                if client.get(url).status_code == 200:
219
                    break
220
221
222
223
224
            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`.
225
                result = self._poll()
226
                if result is not None and result != 0:
227
                    raise RuntimeError("Server exited unexpectedly.") from None
228
229
230

                time.sleep(0.5)
                if time.time() - start > timeout:
231
                    raise RuntimeError("Server failed to start in time.") from None
232
233
234

    @property
    def url_root(self) -> str:
235
236
237
238
239
        return (
            f"http://{self.uds.split('/')[-1]}"
            if self.uds
            else f"http://{self.host}:{self.port}"
        )
240
241
242
243

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

244
245
246
    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
247
248
249
        return openai.OpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
250
251
            max_retries=0,
            **kwargs,
252
253
        )

254
    def get_async_client(self, **kwargs):
255
256
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
257
258
259
260
261
262
        return openai.AsyncOpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
            max_retries=0,
            **kwargs,
        )
263

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
    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
        )

281

282
283
284
class RemoteOpenAIServerCustom(RemoteOpenAIServer):
    """Launch test server with custom child process"""

285
    def _start_server(
286
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
287
    ) -> None:
288
        self.proc: Process = Process(
289
290
            target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)
        )  # type: ignore[assignment]
291
292
        self.proc.start()

293
294
295
296
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
297
        child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],
298
        *,
299
        env_dict: dict[str, str] | None = None,
300
        seed: int = 0,
301
        auto_port: bool = True,
302
        max_wait_seconds: float | None = None,
303
    ) -> None:
304
305
306
        """Store custom child process function then invoke superclass
        constructor which will indirectly launch it."""
        self.child_process_fxn = child_process_fxn
307
308
309
310
311
312
313
314
        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,
        )
315

316
    def _poll(self) -> int | None:
317
318
319
320
321
322
323
324
325
326
        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()


327
328
329
330
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
331
    token_ids: list[int],
332
333
334
335
):
    results = []

    # test with text prompt
336
337
338
339
340
341
342
343
344
345
346
347
    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,
        }
    )
348
349
350
351
352
353
354
355
356

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

357
358
359
360
361
362
363
364
    results.append(
        {
            "test": "token_ids",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
365
366

    # test seeded random sampling
367
368
369
370
371
372
373
374
375
376
377
378
    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,
        }
    )
379
380

    # test seeded random sampling with multiple prompts
381
382
383
384
385
386
387
388
389
390
391
392
    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,
        }
    )
393
394
395
396
397
398
399
400
401

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

402
403
404
405
406
407
408
    results.append(
        {
            "test": "simple_list",
            "text0": batch.choices[0].text,
            "text1": batch.choices[1].text,
        }
    )
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424

    # 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

425
426
427
428
429
430
    results.append(
        {
            "test": "streaming",
            "texts": texts,
        }
    )
431
432
433
434

    return results


435
436
437
438
439
440
441
442
def _test_completion_close(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

    # test with text prompt
443
444
445
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=1, logprobs=5, temperature=0.0
    )
446

447
448
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
449

450
451
452
453
454
455
    results.append(
        {
            "test": "completion_close",
            "logprobs": logprobs,
        }
    )
456
457
458
459

    return results


460
461
462
463
464
465
466
def _test_chat(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

467
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
468
469

    # test with text prompt
470
471
472
473
474
475
476
477
478
479
480
481
    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,
        }
    )
482
483
484
485

    return results


486
487
488
489
490
491
492
493
494
495
496
497
498
499
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",
    )

500
501
502
503
504
505
506
    results.append(
        {
            "test": "single_embedding",
            "embedding": embeddings.data[0].embedding,
            "usage": embeddings.usage,
        }
    )
507
508
509
510

    return results


511
512
513
514
515
516
517
518
def _test_image_text(
    client: openai.OpenAI,
    model_name: str,
    image_url: str,
):
    results = []

    # test pure text input
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
    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,
    )
536
537
538
539
540
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

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

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
    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,
    )
566
567
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

568
569
570
571
572
573
    results.append(
        {
            "test": "text_image",
            "logprobs": top_logprobs,
        }
    )
574
575
576
577

    return results


578
579
580
581
def compare_two_settings(
    model: str,
    arg1: list[str],
    arg2: list[str],
582
583
    env1: dict[str, str] | None = None,
    env2: dict[str, str] | None = None,
584
585
    *,
    method: str = "generate",
586
    max_wait_seconds: float | None = None,
587
) -> None:
588
    """
589
590
591
592
593
594
595
596
597
    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.
598
599
    """

600
601
602
603
604
605
606
607
608
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


609
610
611
def compare_all_settings(
    model: str,
    all_args: list[list[str]],
612
    all_envs: list[dict[str, str] | None],
613
614
    *,
    method: str = "generate",
615
    max_wait_seconds: float | None = None,
616
) -> None:
617
618
619
620
621
622
623
624
625
    """
    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.
    """

626
    trust_remote_code = False
627
    for args in all_args:
628
629
630
631
632
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
633
    for args in all_args:
634
635
636
637
638
639
640
641
642
        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,
    )
643

644
645
646
647
648
649
650
    can_force_load_format = True

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

651
    prompt = "Hello, my name is"
652
    token_ids = tokenizer(prompt).input_ids
653
    ref_results: list = []
654
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
655
656
657
658
659
660
661
662
663
        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]
664
        compare_results: list = []
665
        results = ref_results if i == 0 else compare_results
666
667
668
        with RemoteOpenAIServer(
            model, args, env_dict=env, max_wait_seconds=max_wait_seconds
        ) as server:
669
670
671
672
673
674
            client = server.get_client()

            # test models list
            models = client.models.list()
            models = models.data
            served_model = models[0]
675
676
677
678
679
680
681
            results.append(
                {
                    "test": "models_list",
                    "id": served_model.id,
                    "root": served_model.root,
                }
            )
682

683
684
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
685
686
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
687
688
            elif method == "generate_chat":
                results += _test_chat(client, model, prompt)
689
690
            elif method == "generate_with_image":
                results += _test_image_text(
691
692
                    client,
                    model,
693
                    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
694
                )
695
696
697
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
698
                raise ValueError(f"Unknown method: {method}")
699

700
701
702
703
704
705
            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]
706
                for ref_result, compare_result in zip(ref_results, compare_results):
707
708
709
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
710
711
712
713
714
715
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
716
                            f"Embedding for {model=} are not the same.\n"
717
718
                            f"cosine_similarity={sim}\n"
                        )
719
720
                        del ref_result["embedding"]
                        del compare_result["embedding"]
721
722
723
724
725
                    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"
726
727
                        f"{compare_result=}\n"
                    )
728
729


730
731
732
733
734
735
736
def init_test_distributed_environment(
    tp_size: int,
    pp_size: int,
    rank: int,
    distributed_init_port: str,
    local_rank: int = -1,
) -> None:
737
738
739
740
741
742
743
    # 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,
744
    )
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764

    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,
        )
    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,
            )
765
766
767
    ensure_model_parallel_initialized(tp_size, pp_size)


768
def multi_process_parallel(
769
    monkeypatch: pytest.MonkeyPatch,
770
771
    tp_size: int,
    pp_size: int,
772
    test_target: Any,
773
) -> None:
774
775
    import ray

776
777
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
778
779
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
780
781
782
783
784
785
    # 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={
786
            "working_dir": VLLM_PATH,
787
            "excludes": [
788
789
790
791
792
793
794
795
796
                "build",
                ".git",
                "cmake-build-*",
                "shellcheck",
                "dist",
                "ep_kernels_workspace",
            ],
        }
    )
797
798
799
800
801

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
802
803
804
805
806
807
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
808
809
            ),
        )
810
811
812
    ray.get(refs)

    ray.shutdown()
813
814
815


@contextmanager
816
def error_on_warning(category: type[Warning] = Warning):
817
818
    """
    Within the scope of this context manager, tests will fail if any warning
819
    of the given category is emitted.
820
821
    """
    with warnings.catch_warnings():
822
        warnings.filterwarnings("error", category=category)
823
824

        yield
825
826


827
828
829
830
831
832
833
834
835
836
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]


837
@_nvml()
838
839
840
def wait_for_gpu_memory_to_clear(
    *,
    devices: list[int],
841
842
    threshold_bytes: int | None = None,
    threshold_ratio: float | None = None,
843
844
    timeout_s: float = 120,
) -> None:
845
    assert threshold_bytes is not None or threshold_ratio is not None
846
847
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
848
    devices = get_physical_device_indices(devices)
849
850
    start_time = time.time()
    while True:
851
        output: dict[int, str] = {}
852
        output_raw: dict[int, tuple[float, float]] = {}
853
        for device in devices:
854
            if current_platform.is_rocm():
855
856
857
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
858
                gb_total = mem_info["vram_total"] / 2**10
859
860
861
862
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
863
864
                gb_total = mem_info.total / 2**30
            output_raw[device] = (gb_used, gb_total)
865
            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
866

867
        print("gpu memory used/total (GiB): ", end="")
868
        for k, v in output.items():
869
870
            print(f"{k}={v}; ", end="")
        print("")
871

872
873
        if threshold_bytes is not None:
            is_free = lambda used, total: used <= threshold_bytes / 2**30
874
            threshold = f"{threshold_bytes / 2**30} GiB"
875
876
877
878
        else:
            is_free = lambda used, total: used / total <= threshold_ratio
            threshold = f"{threshold_ratio:.2f}"

879
        dur_s = time.time() - start_time
880
        if all(is_free(used, total) for used, total in output_raw.values()):
881
882
883
884
            print(
                f"Done waiting for free GPU memory on devices {devices=} "
                f"({threshold=}) {dur_s=:.02f}"
            )
885
886
887
            break

        if dur_s >= timeout_s:
888
889
890
891
            raise ValueError(
                f"Memory of devices {devices=} not free after "
                f"{dur_s=:.02f} ({threshold=})"
            )
892
893

        time.sleep(5)
894
895


896
897
898
_P = ParamSpec("_P")


899
def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
900
901
902
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
903

904
    @functools.wraps(func)
905
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
906
907
908
909
        # 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
910
911
912

        # Create a unique temporary file to store exception info from child
        # process. Use test function name and process ID to avoid collisions.
913
914
        with (
            tempfile.NamedTemporaryFile(
915
                delete=False,
916
                mode="w+b",
917
                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
918
919
920
921
                suffix=".exc",
            ) as exc_file,
            ExitStack() as delete_after,
        ):
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
            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
939

940
941
942
943
944
945
946
                    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.
947
                        exc_to_serialize = {"pickled_exception": e}
948
949
950
951
952
                        # Test if it can be pickled
                        cloudpickle.dumps(exc_to_serialize)
                    except (Exception, KeyboardInterrupt):
                        # Fall back to string-based approach.
                        exc_to_serialize = {
953
954
955
                            "exception_type": type(e).__name__,
                            "exception_msg": str(e),
                            "traceback": tb_string,
956
957
                        }
                    try:
958
                        with open(exc_file_path, "wb") as f:
959
960
961
962
963
964
965
                            cloudpickle.dump(exc_to_serialize, f)
                    except Exception:
                        # Fallback: just print the traceback.
                        print(tb_string)
                    os._exit(1)
                else:
                    os._exit(0)
966
            else:
967
968
969
                pgid = os.getpgid(pid)
                _pid, _exitcode = os.waitpid(pid, 0)
                # ignore SIGTERM signal itself
970
                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
971
972
973
974
975
976
977
978
                # 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):
979
980
981
982
                        with (
                            contextlib.suppress(Exception),
                            open(exc_file_path, "rb") as f,
                        ):
983
984
                            exc_info = cloudpickle.load(f)

985
986
987
                    if (
                        original_exception := exc_info.get("pickled_exception")
                    ) is not None:
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
                        # 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}"
1005
1006
                        f" (exit code: {_exitcode})"
                    ) from None
1007
1008

    return wrapper
1009
1010


1011
1012
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
    """Decorator to spawn a new process for each test function."""
1013
1014
1015
1016

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

        import torch.multiprocessing as mp
1022

1023
        with suppress(RuntimeError):
1024
            mp.set_start_method("spawn")
1025
1026
1027
1028
1029
1030

        # Get the module
        module_name = f.__module__

        # Create a process with environment variable set
        env = os.environ.copy()
1031
        env["RUNNING_IN_SUBPROCESS"] = "1"
1032
1033
1034
1035
1036
1037
1038

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

1039
1040
1041
1042
1043
            repo_root = str(VLLM_PATH.resolve())

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

1044
1045
            cmd = [sys.executable, "-m", f"{module_name}"]

1046
1047
1048
            returned = subprocess.run(
                cmd, input=input_bytes, capture_output=True, env=env
            )
1049
1050
1051
1052
1053
1054

            # check if the subprocess is successful
            try:
                returned.check_returncode()
            except Exception as e:
                # wrap raised exception to provide more information
1055
1056
1057
                raise RuntimeError(
                    f"Error raised in subprocess:\n{returned.stderr.decode()}"
                ) from e
1058
1059
1060
1061
1062

    return wrapper


def create_new_process_for_each_test(
1063
    method: Literal["spawn", "fork"] | None = None,
1064
1065
1066
1067
) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
    """Creates a decorator that runs each test function in a new process.

    Args:
1068
        method: The process creation method. Can be either "spawn" or "fork".
1069
1070
               If not specified, it defaults to "spawn" on ROCm and XPU
               platforms and "fork" otherwise.
1071
1072
1073
1074
1075

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

1079
    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
1080
1081
1082
1083
1084
1085
1086

    if method == "fork":
        return fork_new_process_for_each_test

    return spawn_new_process_for_each_test


1087
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
1088
1089
1090
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
1091

1092
1093
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
1094
1095
    """
    try:
1096
        if current_platform.is_cpu():
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
            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

1107
    return pytest.mark.skipif(
1108
        memory_gb < min_gb,
1109
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
1110
1111
    )

1112

1113
1114
1115
1116
1117
1118
1119
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+)",
)


1120
1121
1122
1123
1124
1125
1126
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.
    """
1127
    mark = large_gpu_mark(min_gb)
1128

1129
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1130
        return mark(f)
1131
1132
1133
1134

    return wrapper


1135
1136
1137
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)
1138
1139
1140
1141
1142
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

1143
1144
1145
1146
1147
1148
1149
1150
1151
    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)

1152
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1153
        func = create_new_process_for_each_test()(f)
1154
1155
1156
1157
        for mark in reversed(marks):
            func = mark(func)

        return func
1158
1159
1160
1161

    return wrapper


1162
async def completions_with_server_args(
1163
    prompts: list[str],
1164
    model_name: str,
1165
    server_cli_args: list[str],
1166
    num_logprobs: int | None,
1167
    max_wait_seconds: int = 240,
1168
    max_tokens: int | list = 5,
1169
) -> list[Completion]:
1170
    """Construct a remote OpenAI server, obtain an async client to the
1171
1172
1173
1174
1175
1176
1177
1178
1179
    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
1180
1181
1182
      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.
1183
1184
1185

    Returns:
      OpenAI Completion instance
1186
    """
1187

1188
1189
1190
1191
1192
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

1193
    outputs = None
1194
1195
1196
    with RemoteOpenAIServer(
        model_name, server_cli_args, max_wait_seconds=max_wait_seconds
    ) as server:
1197
        client = server.get_async_client()
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
        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)
        ]
1209
1210
        outputs = await asyncio.gather(*outputs)

1211
    assert outputs is not None, "Completion API call failed."
1212
1213
1214
1215

    return outputs


1216
def get_client_text_generations(completions: list[Completion]) -> list[str]:
1217
    """Extract generated tokens from the output of a
1218
    request made to an Open-AI-protocol completions endpoint.
1219
    """
1220
1221
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
1222
1223
1224


def get_client_text_logprob_generations(
1225
1226
1227
    completions: list[Completion],
) -> list[TextTextLogprobs]:
    """Operates on the output of a request made to an Open-AI-protocol
1228
    completions endpoint; obtains top-rank logprobs for each token in
1229
    each {class}`SequenceGroup`
1230
    """
1231
    text_generations = get_client_text_generations(completions)
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
    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
    ]
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252


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
1253
1254
1255
1256


def get_attn_backend_list_based_on_platform() -> list[str]:
    if current_platform.is_cuda():
1257
        return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"]
1258
    elif current_platform.is_rocm():
1259
        attn_backend_list = ["TRITON_ATTN"]
1260
1261
        try:
            import aiter  # noqa: F401
1262

1263
            attn_backend_list.append("ROCM_AITER_FA")
1264
        except Exception:
1265
            print("Skip ROCM_AITER_FA on ROCm as aiter is not installed")
1266
1267

        return attn_backend_list
1268
1269
    elif current_platform.is_xpu():
        return ["FLASH_ATTN", "TRITON_ATTN"]
1270
1271
    else:
        raise ValueError("Unsupported platform")
1272
1273
1274
1275
1276


@contextmanager
def override_cutlass_fp8_supported(value: bool):
    with patch(
1277
1278
1279
        "vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
        return_value=value,
    ):
1280
        yield
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299


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)
1300
1301
1302
1303
        prompt = (
            "```python\n# We set a number of variables, "
            + f"x{idx} will be important later\n"
        )
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
        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


1316
1317
1318
def check_answers(
    indices: list[int], answer: list[int], outputs: list[str], accept_rate: float = 0.7
):
1319
1320
1321
1322
1323
1324
1325
1326
1327
    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
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347


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))
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
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


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