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

import pytest
import torch

6
from tests.kernels.utils import override_backend_env_variable
7
from vllm.attention.selector import which_attn_to_use
8
from vllm.utils import STR_FLASH_ATTN_VAL, STR_INVALID_VAL
9
10
11


@pytest.mark.parametrize(
12
13
    "name", ["TORCH_SDPA", "ROCM_FLASH", "XFORMERS", "FLASHINFER", "OPENVINO"])
@pytest.mark.parametrize("device", ["cpu", "openvino", "hip", "cuda"])
14
def test_env(name: str, device: str, monkeypatch):
15
16
17
    """Test that the attention selector can be set via environment variable.
    Note that we do not test FlashAttn because it is the default backend.
    """
18
19

    override_backend_env_variable(monkeypatch, name)
20
21
22

    if device == "cpu":
        with patch("vllm.attention.selector.is_cpu", return_value=True):
23
24
            backend = which_attn_to_use(16, torch.float16, torch.float16, 16,
                                        False)
25
26
27
        assert backend.name == "TORCH_SDPA"
    elif device == "hip":
        with patch("vllm.attention.selector.is_hip", return_value=True):
28
29
            backend = which_attn_to_use(16, torch.float16, torch.float16, 16,
                                        False)
30
        assert backend.name == "ROCM_FLASH"
31
32
    elif device == "openvino":
        with patch("vllm.attention.selector.is_openvino", return_value=True):
33
34
            backend = which_attn_to_use(16, torch.float16, torch.float16, 16,
                                        False)
35
        assert backend.name == "OPENVINO"
36
    else:
37
        backend = which_attn_to_use(16, torch.float16, torch.float16, 16,
38
                                    False)
39
40
41
        assert backend.name == name


42
def test_flash_attn(monkeypatch):
43
    """Test FlashAttn validation."""
44
45

    override_backend_env_variable(monkeypatch, STR_FLASH_ATTN_VAL)
46
47

    # Unsupported CUDA arch
48
    with patch("torch.cuda.get_device_capability", return_value=(7, 5)):
49
        backend = which_attn_to_use(16, torch.float16, None, 16, False)
50
        assert backend.name != STR_FLASH_ATTN_VAL
51
52

    # Unsupported data type
53
    backend = which_attn_to_use(16, torch.float8_e4m3fn, None, 16, False)
54
    assert backend.name != STR_FLASH_ATTN_VAL
55
56

    # Unsupported kv cache data type
57
    backend = which_attn_to_use(16, torch.float16, "fp8", 16, False)
58
    assert backend.name != STR_FLASH_ATTN_VAL
59
60

    # Unsupported block size
61
    backend = which_attn_to_use(16, torch.float16, None, 8, False)
62
    assert backend.name != STR_FLASH_ATTN_VAL
63
64
65

    # flash-attn is not installed
    with patch.dict('sys.modules', {'vllm_flash_attn': None}):
66
        backend = which_attn_to_use(16, torch.float16, None, 16, False)
67
        assert backend.name != STR_FLASH_ATTN_VAL
68
69

    # Unsupported head size
70
    backend = which_attn_to_use(17, torch.float16, None, 16, False)
71
72
73
    assert backend.name != STR_FLASH_ATTN_VAL

    # Attention-free models should bypass env and use PlaceholderAttention
74
    backend = which_attn_to_use(16, torch.float16, torch.float16, 16, True)
75
    assert backend.name != STR_FLASH_ATTN_VAL
76
77


78
def test_invalid_env(monkeypatch):
79
    """Throw an exception if the backend name is invalid."""
80
    override_backend_env_variable(monkeypatch, STR_INVALID_VAL)
81
    with pytest.raises(ValueError):
82
        which_attn_to_use(16, torch.float16, None, 16, False)