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

from vllm.logger import init_logger
5
from vllm.platforms import current_platform
6
import torch
7
8
9

logger = init_logger(__name__)

10
11
if current_platform.is_cuda():
    from vllm import _custom_ops as ops
12
    
13
    reshape_and_cache_flash = ops.reshape_and_cache_flash
14
    from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
15
16
elif current_platform.is_xpu():
    from vllm._ipex_ops import ipex_ops as ops
17

18
19
20
    reshape_and_cache_flash = ops.reshape_and_cache_flash
    flash_attn_varlen_func = ops.flash_attn_varlen_func
    get_scheduler_metadata = ops.get_scheduler_metadata
21
22
elif current_platform.is_rocm():
    try:
23
24
25
26
        # from flash_attn import flash_attn_varlen_func  # noqa: F401
        from vllm import _custom_ops as ops
        from flash_attn import vllm_flash_attn_varlen_func
        reshape_and_cache_cuda = ops.reshape_and_cache_cuda
27
28
29
30
31
    except ImportError as e:
        raise ImportError(
            "Rocm platform requires upstream flash-attn "
            "to be installed. Please install flash-attn first."
        ) from e
32
    
33

34
def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
35
36
    # import here to avoid circular dependencies
    from vllm.platforms import current_platform
37

38
39
    if current_platform.is_xpu():
        return 2
40
41
    try:
        from vllm.vllm_flash_attn.flash_attn_interface import (
42
43
44
45
            fa_version_unsupported_reason,
            is_fa_version_supported,
        )

46
47
48
49
50
        device_capability = current_platform.get_device_capability()

        assert device_capability is not None

        # 1. default version depending on platform
51
52
53
        fa_version = (
            3 if (device_capability.major == 9 and is_fa_version_supported(3)) else 2
        )
54

55
56
57
58
59
60
        # 2. override if passed by environment or config
        from vllm.config import get_current_vllm_config

        vllm_config = get_current_vllm_config()
        if vllm_config.attention_config.flash_attn_version is not None:
            fa_version = vllm_config.attention_config.flash_attn_version
61
62
63

        # 3. fallback for unsupported combinations
        if device_capability.major == 10 and fa_version == 3:
64
65
            logger.warning_once(
                "Cannot use FA version 3 on Blackwell platform "
66
67
                "defaulting to FA version 2."
            )
68
69
70
            fa_version = 2

        if requires_alibi and fa_version == 3:
71
72
73
            logger.warning_once(
                "Cannot use FA version 3 with ALiBi, defaulting to FA version 2."
            )
74
75
76
            fa_version = 2

        if not is_fa_version_supported(fa_version):
77
78
79
80
81
            logger.error(
                "Cannot use FA version %d is not supported due to %s",
                fa_version,
                fa_version_unsupported_reason(fa_version),
            )
82
83
84
85
86

        assert is_fa_version_supported(fa_version)
        return fa_version
    except (ImportError, AssertionError):
        return None
87
88
89


def flash_attn_supports_fp8() -> bool:
90
91
    if torch.cuda.get_device_properties("cuda").gcnArchName.split(':')[0] == "gfx938":
        return True
92
93
94
95
    return (
        get_flash_attn_version() == 3
        and current_platform.get_device_capability().major == 9
    )
96
97


98
99
100
101
102
def flash_attn_supports_sinks() -> bool:
    if current_platform.is_xpu():
        return True
    else:
        return get_flash_attn_version() == 3
103
104


105
106
def flash_attn_supports_mla():
    from vllm.platforms import current_platform
107

108
109
110
    if current_platform.is_cuda():
        try:
            from vllm.vllm_flash_attn.flash_attn_interface import (
111
112
113
114
115
                is_fa_version_supported,
            )

            return (
                is_fa_version_supported(3)
116
                and current_platform.get_device_capability()[0] == 9
117
            )
118
119
120
121
122
        except (ImportError, AssertionError):
            pass
    return False


123
def is_flash_attn_varlen_func_available() -> bool:
zhuwenwen's avatar
zhuwenwen committed
124
    return current_platform.is_cuda() or current_platform.is_rocm() or current_platform.is_xpu()