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

4
from unittest.mock import patch
5
6
7
8

import pytest
import torch

9
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
10
11
12
from vllm.platforms.cpu import CpuPlatform
from vllm.platforms.cuda import CudaPlatform
from vllm.platforms.rocm import RocmPlatform
13
from vllm.utils import STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL, STR_INVALID_VAL
14
15


16
17
@pytest.fixture(autouse=True)
def clear_cache():
18
    """Clear lru cache to ensure each test case runs without caching."""
19
20
21
    _cached_get_attn_backend.cache_clear()


22
23
# Define MLA and non-MLA backends separately
DEVICE_MLA_BACKENDS = {
24
    "cuda": [
25
26
27
28
29
        "TRITON_MLA",
        "FLASHMLA",
        "FLASHINFER_MLA",
        "FLASH_ATTN_MLA",
        "CUTLASS_MLA",
30
    ],
31
32
33
34
35
    "hip": ["TRITON_MLA", "ROCM_AITER_MLA"],
    "cpu": [],
}

DEVICE_REGULAR_ATTN_BACKENDS = {
36
    "cuda": ["XFORMERS", "FLASHINFER", "FLASH_ATTN"],
37
38
39
40
41
42
43
    "hip": ["ROCM_FLASH"],
    "cpu": ["TORCH_SDPA"],
}

DEVICE_MLA_BLOCK_SIZES = {
    "cuda": [16, 64],  # CUDA supports both standard and extended block sizes
    "hip": [16, 1],  # HIP requires special handling for block_size=1
44
    # "cpu": [16]  # CPU uses fixed block size from test cases
45
    "cpu": [],  # FIXME(woosuk): Temporarily disable CPU tests
46
47
48
49
50
51
52
}


def generate_params():
    params = []
    for use_mla in [True, False]:
        for device in ["cuda", "hip", "cpu"]:
53
54
55
56
57
            backends = (
                DEVICE_MLA_BACKENDS[device]
                if use_mla
                else DEVICE_REGULAR_ATTN_BACKENDS[device]
            )
58
            for name in backends:
59
                block_sizes = DEVICE_MLA_BLOCK_SIZES[device] if use_mla else [16]
60
61
62
63
64
65
66
                for block_size in block_sizes:
                    params.append(
                        pytest.param(
                            device,
                            name,
                            use_mla,
                            block_size,
67
68
69
                            id=f"{device}_{name}_mla_{str(use_mla)[0]}_blks{block_size}",
                        )
                    )
70
71
72
    return params


73
@pytest.mark.parametrize("device, name, use_mla, block_size", generate_params())
74
def test_env(
75
    device: str,
76
    name: str,
77
78
    use_mla: bool,
    block_size: int,
79
80
    monkeypatch: pytest.MonkeyPatch,
):
81
    """Test attention backend selection with valid device-backend pairs."""
82
83
    with monkeypatch.context() as m:
        m.setenv(STR_BACKEND_ENV_VAR, name)
84
        m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")
85
86

        if device == "cpu":
87
            with patch("vllm.attention.selector.current_platform", CpuPlatform()):
88
                backend = get_attn_backend(16, torch.float16, None, block_size)
89
            assert backend.get_name() == "TORCH_SDPA"
90

91
        elif device == "hip":
92
            with patch("vllm.attention.selector.current_platform", RocmPlatform()):
93
                if use_mla:
94
95
96
97
98
99
100
101
                    # ROCm MLA backend logic:
                    # - TRITON_MLA: supported when block_size != 1
                    # - ROCM_AITER_MLA: supported when block_size == 1
                    # If backend is forced but doesn't match block_size,
                    # should raise ValueError

                    if name == "TRITON_MLA" and block_size == 1:
                        # TRITON_MLA doesn't support block_size == 1
102
                        with pytest.raises(ValueError) as exc_info:
103
104
105
106
                            get_attn_backend(
                                16, torch.float16, None, block_size, use_mla=use_mla
                            )
                        assert f"The selected backend, {name}" in str(exc_info.value)
107
108
109
                    elif name == "ROCM_AITER_MLA" and block_size != 1:
                        # ROCM_AITER_MLA only supports block_size == 1
                        with pytest.raises(ValueError) as exc_info:
110
111
112
113
                            get_attn_backend(
                                16, torch.float16, None, block_size, use_mla=use_mla
                            )
                        assert f"The selected backend, {name}" in str(exc_info.value)
114
115
                    else:
                        # Valid backend-block_size combination
116
117
118
                        backend = get_attn_backend(
                            16, torch.float16, None, block_size, use_mla=use_mla
                        )
119
                        expected = name
120
                        assert backend.get_name() == expected
121
                else:
122
123
124
                    backend = get_attn_backend(
                        16, torch.float16, None, block_size, use_mla=use_mla
                    )
125
                    expected = "TRITON_ATTN"
126
127
128
                    assert backend.get_name() == expected

        elif device == "cuda":
129
            with patch("vllm.attention.selector.current_platform", CudaPlatform()):
130
                if use_mla:
131
132
133
                    # CUDA MLA backend logic:
                    # - CUTLASS_MLA: only supported with block_size == 128
                    #   and Blackwell GPUs (SM 10.0), V1 only
134
135
                    # - FLASHINFER_MLA: only supported on Blackwell GPUs
                    #   (SM 10.0+), V1 only
136
137
138
139
140
                    # - FLASHMLA: only supported with block_size == 64
                    # - FLASH_ATTN_MLA: V1 only
                    # - TRITON_MLA: fallback for other cases

                    if name == "CUTLASS_MLA":
141
                        if block_size != 128:
142
                            # CUTLASS_MLA only supports block_size == 128
143
                            pytest.skip("CUTLASS_MLA only supports block_size 128")
144
                        else:
145
146
147
                            backend = get_attn_backend(
                                16, torch.float16, None, block_size, use_mla=use_mla
                            )
148
                            expected = "CUTLASS_MLA"
149
                            assert backend.get_name() == expected
150
                    elif name == "FLASHINFER_MLA":
151
                        if block_size not in [32, 64]:
152
153
                            # FlashInfer MLA only supports block_size 32 or 64
                            pytest.skip(
154
155
                                "FlashInfer MLA only supports block_size 32 or 64"
                            )
156
                        else:
157
158
159
                            backend = get_attn_backend(
                                16, torch.float16, None, block_size, use_mla=use_mla
                            )
160
161
                            expected = "FLASHINFER_MLA"
                            assert backend.get_name() == expected
162
163
164
165
166
                    elif name == "FLASHMLA":
                        if block_size != 64:
                            # FlashMLA only supports block_size == 64
                            pytest.skip("FlashMLA only supports block_size 64")
                        else:
167
                            from vllm.v1.attention.backends.mla.flashmla import (
168
169
170
                                is_flashmla_supported,
                            )

171
172
                            is_supported, _ = is_flashmla_supported()
                            if not is_supported:
173
                                pytest.skip("FlashMLA not supported on this platform")
174
                            else:
175
176
177
                                backend = get_attn_backend(
                                    16, torch.float16, None, block_size, use_mla=use_mla
                                )
178
                                expected = name
179
180
                                assert backend.get_name() == expected
                    elif name == "FLASH_ATTN_MLA":
181
182
183
                        backend = get_attn_backend(
                            16, torch.float16, None, block_size, use_mla=use_mla
                        )
184
185
                        expected = "FLASH_ATTN_MLA"
                        assert backend.get_name() == expected
186
                    else:
187
                        # TRITON_MLA or other fallback
188
189
190
                        backend = get_attn_backend(
                            16, torch.float16, None, block_size, use_mla=use_mla
                        )
191
                        expected = "TRITON_MLA"
192
                        assert backend.get_name() == expected
193
                elif name == "FLASHINFER":
194
195
196
                    backend = get_attn_backend(
                        16, torch.float16, None, block_size, use_mla=use_mla
                    )
197
                    expected = "FLASHINFER"
198
                    assert backend.get_name() == expected
199
                elif name == "XFORMERS":
200
201
202
                    backend = get_attn_backend(
                        32, torch.float16, None, block_size, use_mla=use_mla
                    )
203
                    expected = "XFORMERS"
204
                    assert backend.get_name() == expected
205
                elif name == "FLASH_ATTN":
206
207
208
                    backend = get_attn_backend(
                        32, torch.float16, None, block_size, use_mla=use_mla
                    )
209
210
                    expected = "FLASH_ATTN"
                    assert backend.get_name() == expected
211

212

213
@pytest.mark.parametrize("device", ["cpu", "cuda"])
214
def test_fp32_fallback(device: str):
215
    """Test attention backend selection with fp32."""
216
217
218
219
    if device == "cpu":
        with patch("vllm.attention.selector.current_platform", CpuPlatform()):
            backend = get_attn_backend(16, torch.float32, None, 16)
        assert backend.get_name() == "TORCH_SDPA"
220

221
222
223
224
    elif device == "cuda":
        with patch("vllm.attention.selector.current_platform", CudaPlatform()):
            backend = get_attn_backend(16, torch.float32, None, 16)
        assert backend.get_name() == "FLEX_ATTENTION"
225
226


227
def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
228
    """Test FlashAttn validation."""
229
230
231
232
    pytest.skip(
        "Skipping as current backend selector does not "
        "handle fallbacks when a backend is set via env var."
    )
233

234
235
    with monkeypatch.context() as m:
        m.setenv(STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL)
236

237
        # Unsupported CUDA arch
238
        monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))
239
        backend = get_attn_backend(16, torch.float16, None, 16)
240
        assert backend.get_name() != STR_FLASH_ATTN_VAL
241

242
243
        # Reset the monkeypatch for subsequent tests
        monkeypatch.undo()
244

245
        # Unsupported data type
246
        backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16)
247
        assert backend.get_name() != STR_FLASH_ATTN_VAL
248

249
        # Unsupported kv cache data type
250
        backend = get_attn_backend(16, torch.float16, "fp8", 16)
251
        assert backend.get_name() != STR_FLASH_ATTN_VAL
252

253
        # Unsupported block size
254
        backend = get_attn_backend(16, torch.float16, None, 8)
255
256
257
258
        assert backend.get_name() != STR_FLASH_ATTN_VAL

        # flash-attn is not installed
        import sys
259
260
261

        original_module = sys.modules.get("vllm_flash_attn")
        monkeypatch.setitem(sys.modules, "vllm_flash_attn", None)
262
        backend = get_attn_backend(16, torch.float16, None, 16)
263
        assert backend.get_name() != STR_FLASH_ATTN_VAL
264

265
266
        # Restore the original module if it existed
        if original_module is not None:
267
            monkeypatch.setitem(sys.modules, "vllm_flash_attn", original_module)
268
        else:
269
            monkeypatch.delitem(sys.modules, "vllm_flash_attn", raising=False)
270

271
        # Unsupported head size
272
        backend = get_attn_backend(17, torch.float16, None, 16)
273
        assert backend.get_name() != STR_FLASH_ATTN_VAL
274
275


276
def test_invalid_env(monkeypatch: pytest.MonkeyPatch):
277
    """Test that invalid attention backend names raise ValueError."""
278
279
280
281
    with (
        monkeypatch.context() as m,
        patch("vllm.attention.selector.current_platform", CudaPlatform()),
    ):
282
        m.setenv(STR_BACKEND_ENV_VAR, STR_INVALID_VAL)
283

284
285
        # Should raise ValueError for invalid backend
        with pytest.raises(ValueError) as exc_info:
286
            get_attn_backend(32, torch.float16, None, 16)
287
        assert "Invalid value 'INVALID'" in str(exc_info.value)