test_full_cudagraph.py 5.47 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import contextlib
import os
5
import weakref
6
7
8

import pytest

9
from tests.utils import wait_for_gpu_memory_to_clear
10
from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
11
12
from vllm import LLM, SamplingParams
from vllm.config import CompilationConfig
13
from vllm.platforms import current_platform
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34


@contextlib.contextmanager
def temporary_environ(env_vars):
    """
    Temporarily set environment variables and restore them afterward.
    We have to do this vs monkeypatch because monkeypatch doesn't work
    with "module" scoped fixtures.
    """
    original_env = {k: os.environ.get(k) for k in env_vars}
    try:
        os.environ.update(env_vars)
        yield
    finally:
        for k, v in original_env.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v


35
36
37
test_params_full_cudagraph = []

# deepseek-ai/DeepSeek-V2-Lite with MLA
38
MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"]
39
40
for mla_backend in MLA_backends:
    test_params_full_cudagraph.append(
41
42
        pytest.param(("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend]))
    )
43
44
45
46
47
48
49

# Qwen/Qwen2-1.5B-Instruct with other backends
other_backend_configs = [
    backend_configs[c] for c in backend_configs if c not in MLA_backends
]
for backend_config in other_backend_configs:
    test_params_full_cudagraph.append(
50
51
        pytest.param(("Qwen/Qwen2-1.5B-Instruct", backend_config))
    )
52
53
54
55
56
57
58


@pytest.fixture(scope="class")
def llm_pair(request):
    model, backend_config = request.param

    # Dynamically skip test if GPU capability is not met
59
60
61
62
    if (
        backend_config.specific_gpu_arch
        and backend_config.specific_gpu_arch != current_platform.get_device_capability()
    ):
63
64
65
66
67
68
69
70
71
72
73
74
        if backend_config.specific_gpu_arch == (9, 0):
            pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
        elif backend_config.specific_gpu_arch == (10, 0):
            pytest.skip("Only Blackwell GPUs support Cutlass MLA")

    env_vars = {
        # Force native sampler to avoid potential nondeterminism in FlashInfer
        # when per-request generators are not used in V1.
        "VLLM_USE_FLASHINFER_SAMPLER": "0",
        **backend_config.env_vars,
    }
    with temporary_environ(env_vars):
75
76
        full = LLM(
            model=model,
77
            gpu_memory_utilization=0.43,
78
79
            trust_remote_code=True,
            max_model_len=1024,
80
            max_num_seqs=128,
81
            compilation_config=CompilationConfig(**backend_config.comp_config),
82
83
            generation_config="vllm",
            seed=42,
84
85
86
        )
        piecewise = LLM(
            model=model,
87
            gpu_memory_utilization=0.43,
88
89
            trust_remote_code=True,
            max_model_len=1024,
90
91
92
93
            max_num_seqs=128,
            compilation_config=CompilationConfig(cudagraph_mode="PIECEWISE"),
            generation_config="vllm",
            seed=42,
94
95
        )

96
    # PyTest caches the fixture values so we use weakref.proxy to enable GC
97
98
99
100
101
102
103
104
105
106
    yield weakref.proxy(full), weakref.proxy(piecewise)
    del full
    del piecewise

    wait_for_gpu_memory_to_clear(
        devices=[0],
        threshold_ratio=0.1,
    )


107
@pytest.mark.parametrize("llm_pair", test_params_full_cudagraph, indirect=True)
108
class TestFullCUDAGraph:
109
    """
110
111
    Use a class such that an llm pair is constructed once for all
    batch_size/max_tokens combinations and released immediately after.
112

113
114
    Module-scope fixtures would stick around the whole time,
    meaning there would be multiple LLM instances hogging memory simultaneously.
115
116
    """

117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
    @pytest.mark.parametrize(
        ("batch_size", "max_tokens"),
        [
            (1, 10),
            (7, 10),
            (16, 10),
            (25, 10),
            (32, 10),
            (45, 10),
            (64, 10),
            (123, 10),
            (8, 5),
            (8, 30),
        ],
    )
    def test_full_cudagraph(self, batch_size, max_tokens, llm_pair: tuple[LLM, LLM]):
133
134
135
136
137
        """
        Test various batch sizes and max_tokens to ensure that the
        full cudagraph compilation works for padded cases too.
        """

138
        full_cudagraph_llm, piecewise_llm = llm_pair
139

140
141
142
        prompts = ["the quick brown fox"] * batch_size
        # Use purely greedy decoding to avoid top-p truncation sensitivity
        # that can amplify tiny numeric differences across runtimes.
143
144
145
        sampling_params = SamplingParams(
            temperature=0.0, max_tokens=max_tokens, top_p=1.0
        )
146
147
148
149
150

        piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
        full_responses = full_cudagraph_llm.generate(prompts, sampling_params)

        # Check that all responses are the same
151
152
153
154
155
        for piecewise_res, full_res in zip(piecewise_responses, full_responses):
            assert (
                piecewise_res.outputs[0].text.lower()
                == full_res.outputs[0].text.lower()
            )
156
157


158
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
159
def test_full_cudagraph_with_invalid_backend():
160
161
162
163
164
165
166
167
168
169
170
171
172
    with (
        temporary_environ(
            {
                "VLLM_ATTENTION_BACKEND": "FLEX_ATTENTION",
                # Flex_Attention is not supported with full cuda graph
            }
        ),
        pytest.raises(RuntimeError),
    ):
        LLM(
            model="Qwen/Qwen2-1.5B-Instruct",
            compilation_config=CompilationConfig(cudagraph_mode="FULL"),
        )