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

4
from abc import ABC, abstractmethod
5
from typing import Generic, List, Optional, Protocol, Tuple, Type, TypeVar
6
7
8

import torch

9
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
10

11

12
13
14
15
16
17
class AttentionType:
    """
    Attention type.
    Use string to be compatible with `torch.compile`.
    """
    DECODER = "decoder"
18
    """Decoder attention between previous layer Q/K/V."""
19
    ENCODER = "encoder"
20
    """Encoder attention between previous layer Q/K/V for encoder-decoder."""
21
    ENCODER_ONLY = "encoder_only"
22
    """Encoder attention between previous layer Q/K/V."""
23
    ENCODER_DECODER = "encoder_decoder"
24
    """Attention between dec. Q and enc. K/V for encoder-decoder."""
25
26


27
28
class AttentionBackend(ABC):
    """Abstract class for attention backends."""
29
30
31
32
    # For some attention backends, we allocate an output tensor before
    # calling the custom op. When piecewise cudagraph is enabled, this
    # makes sure the output tensor is allocated inside the cudagraph.
    accept_output_buffer: bool = False
33

34
35
36
37
38
39
40
41
    # Whether this backend supports receiving pre-quantized query input.
    # If True, the attention layer will handle query quantization instead
    # of the backend, allowing torch.compile to fuse quantization with
    # previous operations.
    # Needs to be worked through for all backends
    # https://github.com/vllm-project/vllm/issues/25584
    supports_quant_query_input: bool = False

42
43
44
45
46
    @staticmethod
    @abstractmethod
    def get_name() -> str:
        raise NotImplementedError

47
48
49
50
51
52
53
    @staticmethod
    @abstractmethod
    def get_impl_cls() -> Type["AttentionImpl"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
54
    def get_metadata_cls() -> Type["AttentionMetadata"]:
55
56
        raise NotImplementedError

57
58
59
60
    @classmethod
    def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
        return cls.get_metadata_cls()(*args, **kwargs)

61
62
    @staticmethod
    @abstractmethod
63
    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
64
65
        raise NotImplementedError

66
67
68
69
70
71
72
73
74
75
    @staticmethod
    @abstractmethod
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
    ) -> Tuple[int, ...]:
        raise NotImplementedError

76
77
78
79
    @staticmethod
    def get_kv_cache_stride_order() -> Tuple[int, ...]:
        raise NotImplementedError

80
81
82
83
    @classmethod
    def full_cls_name(cls) -> tuple[str, str]:
        return (cls.__module__, cls.__qualname__)

84

85
class AttentionMetadata:
86
    pass
87
88


89
T = TypeVar("T", bound=AttentionMetadata)
90
91


92
93
class AttentionLayer(Protocol):

94
    _q_scale: torch.Tensor
95
96
    _k_scale: torch.Tensor
    _v_scale: torch.Tensor
97
    _q_scale_float: float
98
99
    _k_scale_float: float
    _v_scale_float: float
100
    _prob_scale: torch.Tensor
101
102
103
104
105
106
107
108
109
110
111
112

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
        ...


113
class AttentionImpl(ABC, Generic[T]):
114

115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
    # Whether the attention impl can return the softmax lse for decode.
    # Some features like decode context parallelism require the softmax lse.
    can_return_lse_for_decode: bool = False

    # some attention backends might not always want to return lse
    # even if they can return lse (for efficiency reasons)
    need_to_return_lse_for_decode: bool = False

    dcp_world_size: int
    dcp_rank: int

    def __new__(cls, *args, **kwargs):
        # use __new__ so that all subclasses will call this
        self = super().__new__(cls)
        try:
            from vllm.distributed.parallel_state import get_dcp_group
            self.dcp_world_size = get_dcp_group().world_size
            self.dcp_rank = get_dcp_group().rank_in_group
        except AssertionError:
            # DCP might not be initialized in testing
            self.dcp_world_size = 1
            self.dcp_rank = 0
        self.need_to_return_lse_for_decode = self.dcp_world_size > 1 \
            and self.can_return_lse_for_decode
        return self

141
142
143
144
145
146
147
148
149
    @abstractmethod
    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: Optional[int] = None,
        alibi_slopes: Optional[List[float]] = None,
        sliding_window: Optional[int] = None,
150
        kv_cache_dtype: str = "auto",
151
        logits_soft_cap: Optional[float] = None,
152
        attn_type: str = AttentionType.DECODER,
153
        kv_sharing_target_layer_name: Optional[str] = None,
154
155
156
157
158
159
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def forward(
        self,
160
        layer: AttentionLayer,
161
162
163
164
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
165
        attn_metadata: T,
166
        output: Optional[torch.Tensor] = None,
167
        output_scale: Optional[torch.Tensor] = None,
168
        output_block_scale: Optional[torch.Tensor] = None,
169
170
    ) -> torch.Tensor:
        raise NotImplementedError
171

172
    def fused_output_quant_supported(self, quant_key: QuantKey):
173
174
175
176
177
        """
        Does this attention implementation support fused output quantization.
        This is used by the AttnFusionPass to only fuse output quantization
        onto implementations that support it.

178
        :param quant_key: QuantKey object that describes the quantization op
179
180
181
182
        :return: is fusion supported for this type of quantization
        """
        return False

183
184
185
186
187
188
189
190
191
192
193
194
195

class MLAAttentionImpl(AttentionImpl[T], Generic[T]):

    @abstractmethod
    def forward(
        self,
        layer: AttentionLayer,
        hidden_states_or_cq: torch.Tensor,
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: T,
        output: Optional[torch.Tensor] = None,
196
        output_scale: Optional[torch.Tensor] = None,
197
        output_block_scale: Optional[torch.Tensor] = None,
198
199
    ) -> torch.Tensor:
        raise NotImplementedError
200
201
202
203


def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
    return kv_cache_dtype != "auto"