test_attention_selector.py 4.39 KB
Newer Older
1
from unittest.mock import Mock, patch
2
3
4
5

import pytest
import torch

6
from tests.kernels.utils import override_backend_env_variable
7
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
8
9
10
11
from vllm.platforms.cpu import CpuPlatform
from vllm.platforms.cuda import CudaPlatform
from vllm.platforms.openvino import OpenVinoPlatform
from vllm.platforms.rocm import RocmPlatform
12
from vllm.utils import STR_FLASH_ATTN_VAL, STR_INVALID_VAL
zhuwenwen's avatar
zhuwenwen committed
13
from vllm.platforms import current_platform
14
15


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


23
@pytest.mark.parametrize(
zhuwenwen's avatar
zhuwenwen committed
24
    "name", ["TORCH_SDPA", "ROCM_FLASH", "XFORMERS", "FLASHINFER", "OPENVINO"] if not current_platform() else ["ROCM_FLASH"])
25
@pytest.mark.parametrize("device", ["cpu", "openvino", "hip", "cuda"])
26
def test_env(name: str, device: str, monkeypatch):
27
28
29
    """Test that the attention selector can be set via environment variable.
    Note that we do not test FlashAttn because it is the default backend.
    """
30
31

    override_backend_env_variable(monkeypatch, name)
32
33

    if device == "cpu":
34
        with patch("vllm.attention.selector.current_platform", CpuPlatform()):
35
36
37
            backend = get_attn_backend(16, torch.float16, torch.float16, 16,
                                       False)
        assert backend.get_name() == "TORCH_SDPA"
38
    elif device == "hip":
39
        with patch("vllm.attention.selector.current_platform", RocmPlatform()):
40
41
42
            backend = get_attn_backend(16, torch.float16, torch.float16, 16,
                                       False)
        assert backend.get_name() == "ROCM_FLASH"
43
    elif device == "openvino":
44
        with patch("vllm.attention.selector.current_platform",
45
46
47
48
49
                   OpenVinoPlatform()), patch.dict('sys.modules',
                                                   {'openvino': Mock()}):
            backend = get_attn_backend(16, torch.float16, torch.float16, 16,
                                       False)
        assert backend.get_name() == "OPENVINO"
50
    else:
51
52
53
54
55
56
        if name in ["XFORMERS", "FLASHINFER"]:
            with patch("vllm.attention.selector.current_platform",
                       CudaPlatform()):
                backend = get_attn_backend(16, torch.float16, torch.float16,
                                           16, False)
            assert backend.get_name() == name
57
58


59
def test_flash_attn(monkeypatch):
60
    """Test FlashAttn validation."""
Joe Runde's avatar
Joe Runde committed
61
    # TODO: When testing for v1, pipe in `use_v1` as an argument to
62
    # get_attn_backend
63
64

    override_backend_env_variable(monkeypatch, STR_FLASH_ATTN_VAL)
65
66

    # Unsupported CUDA arch
67
    with patch("torch.cuda.get_device_capability", return_value=(7, 5)):
68
69
        backend = get_attn_backend(16, torch.float16, None, 16, False)
        assert backend.get_name() != STR_FLASH_ATTN_VAL
70
71

    # Unsupported data type
72
73
    backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16, False)
    assert backend.get_name() != STR_FLASH_ATTN_VAL
74
75

    # Unsupported kv cache data type
76
77
    backend = get_attn_backend(16, torch.float16, "fp8", 16, False)
    assert backend.get_name() != STR_FLASH_ATTN_VAL
78
79

    # Unsupported block size
80
81
    backend = get_attn_backend(16, torch.float16, None, 8, False)
    assert backend.get_name() != STR_FLASH_ATTN_VAL
82
83
84

    # flash-attn is not installed
    with patch.dict('sys.modules', {'vllm_flash_attn': None}):
85
86
        backend = get_attn_backend(16, torch.float16, None, 16, False)
        assert backend.get_name() != STR_FLASH_ATTN_VAL
87
88

    # Unsupported head size
89
90
    backend = get_attn_backend(17, torch.float16, None, 16, False)
    assert backend.get_name() != STR_FLASH_ATTN_VAL
91
92

    # Attention-free models should bypass env and use PlaceholderAttention
93
94
    backend = get_attn_backend(16, torch.float16, torch.float16, 16, True)
    assert backend.get_name() != STR_FLASH_ATTN_VAL
95
96


97
def test_invalid_env(monkeypatch):
98
    """Ignore the invalid env variable if it is set."""
99
    override_backend_env_variable(monkeypatch, STR_INVALID_VAL)
100
101
102
103
104
105
106
    with patch("vllm.attention.selector.current_platform", CudaPlatform()):
        backend = get_attn_backend(32, torch.float16, None, 16, False)
        assert backend.get_name() == "FLASH_ATTN"

        # when block size == 16, backend will fall back to XFORMERS
        backend = get_attn_backend(16, torch.float16, None, 16, False)
        assert backend.get_name() == "XFORMERS"