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

from dataclasses import dataclass
5
from typing import ClassVar
6
7
8

import torch

9
from vllm.attention.backends.abstract import AttentionLayer, AttentionType, MultipleOf
10
11
12
from vllm.attention.ops.flashmla import (
    flash_mla_with_kvcache,
    get_mla_metadata,
13
    is_flashmla_dense_supported,
14
)
15
from vllm.config import VllmConfig
16
from vllm.config.cache import CacheDType
17
from vllm.logger import init_logger
18
from vllm.model_executor.layers.batch_invariant import (
19
    vllm_is_batch_invariant,
20
)
21
from vllm.platforms.interface import DeviceCapability
22
23
24
25
26
27
from vllm.v1.attention.backends.mla.common import (
    MLACommonBackend,
    MLACommonDecodeMetadata,
    MLACommonImpl,
    MLACommonMetadata,
    MLACommonMetadataBuilder,
28
29
30
31
32
33
    QueryLenSupport,
)
from vllm.v1.attention.backends.utils import (
    AttentionCGSupport,
    reshape_attn_output_for_spec_decode,
    reshape_query_for_spec_decode,
34
)
35
from vllm.v1.kv_cache_interface import AttentionSpec
36
37
38
39
40

logger = init_logger(__name__)


class FlashMLABackend(MLACommonBackend):
41
42
43
44
45
46
47
    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "fp8",
        "fp8_e4m3",
    ]

48
49
50
51
    @staticmethod
    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
        return [64]

52
53
    @staticmethod
    def get_name() -> str:
54
        return "FLASHMLA"
55
56

    @staticmethod
57
    def get_builder_cls() -> type["FlashMLAMetadataBuilder"]:
58
59
60
        return FlashMLAMetadataBuilder

    @staticmethod
61
    def get_impl_cls() -> type["FlashMLAImpl"]:
62
63
        return FlashMLAImpl

64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability.major in [9, 10]

    @classmethod
    def supports_combination(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: CacheDType | None,
        block_size: int,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        device_capability: DeviceCapability,
    ) -> str | None:
        if use_sparse:
            from vllm.attention.ops.flashmla import is_flashmla_sparse_supported

            return is_flashmla_sparse_supported()[1]
        else:
            from vllm.attention.ops.flashmla import is_flashmla_dense_supported

            return is_flashmla_dense_supported()[1]
88

89
90

@dataclass
91
class FlashMLADecodeMetadata(MLACommonDecodeMetadata):
92
    tile_scheduler_metadata: torch.Tensor
93
94
95
96
97
98
    num_splits: torch.Tensor


@dataclass
class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
    pass
99
100
101


class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
102
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
103
    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
104
    reorder_batch_threshold: int = 128  # process small prefills with decode pathway
105
    # ^ TODO(matt): tune this
106

107
108
109
110
111
112
113
114
115
116
    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ):
        super().__init__(
            kv_cache_spec, layer_names, vllm_config, device, FlashMLAMetadata
        )
117

118
        self.num_q_heads = vllm_config.model_config.get_num_attention_heads(
119
120
            vllm_config.parallel_config
        )
121

122
123
        self.cg_buf_tile_scheduler_metadata = None
        self.cg_buf_num_splits = None
124
        self.is_fp8_kvcache = vllm_config.cache_config.cache_dtype.startswith("fp8")
125

126
127
128
        device_properties = torch.cuda.get_device_properties(self.device)
        num_sms = device_properties.multi_processor_count

129
        if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
130
131
132
133
134
135
136
137
138
139
            self.cg_buf_tile_scheduler_metadata = torch.zeros(
                # Upper bound on size (<= #SMs, TileSchedulerMetaDataSize)
                # TileSchedulerMetaDataSize = 8
                (num_sms, 8),
                device=self.device,
                dtype=torch.int32,
            )
            self.cg_buf_num_splits = torch.empty(
                (vllm_config.scheduler_config.max_num_seqs + 1),
                device=self.device,
140
141
142
143
144
145
146
                dtype=torch.int32,
            )

    def _build_decode(
        self,
        block_table_tensor: torch.Tensor,
        seq_lens_device: torch.Tensor,
147
        max_seq_len: int,
148
149
150
        query_start_loc_cpu: torch.Tensor,
        query_start_loc_device: torch.Tensor,
        num_decode_tokens: int,
151
        dcp_tot_seq_lens_device: torch.Tensor | None,
152
    ) -> FlashMLADecodeMetadata:
153
154
155
156
        query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
        # we use the max but all should be the same due to uniform length requirement
        max_query_len = query_lens_cpu.max().item()
        num_q_tokens_per_head_k = max_query_len * self.num_q_heads // 1
157
        tile_scheduler_metadata, num_splits = get_mla_metadata(
158
            seq_lens_device,
159
            num_q_tokens_per_head_k,
160
            1,  # MQA for the decode path
161
            is_fp8_kvcache=self.is_fp8_kvcache,
162
        )
163

164
165
166
167
        # TODO: we can disambiguate between decode and mixed-prefill decode here
        # so we can only use the persistent buffer if a cudagraph is actually
        # being used.
        if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
168
169
170
171
172
173
            assert self.cg_buf_tile_scheduler_metadata is not None
            assert self.cg_buf_num_splits is not None

            sm_parts = tile_scheduler_metadata.size(0)
            # Metadata per-SM, upper bound on size (<= #SMs, TileMetadataSize)
            assert sm_parts <= self.cg_buf_tile_scheduler_metadata.size(0)
174
175
176
            tile_scheduler_metadata_view = self.cg_buf_tile_scheduler_metadata[
                :sm_parts
            ]
177
178
179
180
181
182
183
184
185
186
187
188
189
190
            tile_scheduler_metadata_view.copy_(tile_scheduler_metadata)
            tile_scheduler_metadata = tile_scheduler_metadata_view

            # Num splits is per-batch, varying size (batch_size,)
            n = num_splits.size(0)
            # make sure static buffer is large enough
            assert n <= self.cg_buf_num_splits.size(0)
            num_splits_view = self.cg_buf_num_splits[:n]
            num_splits_view.copy_(num_splits)
            # Num splits needs to monotonically increasing
            # (with: https://github.com/vllm-project/FlashMLA/pull/3, otherwise
            #  it needs to monotonically increasing by 1)
            self.cg_buf_num_splits[n:].fill_(num_splits[-1])
            num_splits = num_splits_view
191

192
        return FlashMLADecodeMetadata(
193
            block_table=block_table_tensor,
194
            seq_lens=seq_lens_device,
195
196
            tile_scheduler_metadata=tile_scheduler_metadata,
            num_splits=num_splits,
197
            dcp_tot_seq_lens=dcp_tot_seq_lens_device,
198
        )
199
200
201


class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
202
203
    can_return_lse_for_decode: bool = True

204
    def __init__(
205
206
207
208
209
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
210
211
        alibi_slopes: list[float] | None,
        sliding_window: int | None,
212
        kv_cache_dtype: str,
213
        logits_soft_cap: float | None,
214
        attn_type: str,
215
        kv_sharing_target_layer_name: str | None,
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
        # MLA Specific Arguments
        **mla_args,
    ) -> None:
        super().__init__(
            num_heads,
            head_size,
            scale,
            num_kv_heads,
            alibi_slopes,
            sliding_window,
            kv_cache_dtype,
            logits_soft_cap,
            attn_type,
            kv_sharing_target_layer_name,
            **mla_args,
        )
232

233
        is_supported, reason = is_flashmla_dense_supported()
234
        assert is_supported, reason
235

236
        unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
237
238
239
        if any(unsupported_features):
            raise NotImplementedError(
                "FlashMLAImpl does not support one of the following: "
240
241
                "alibi_slopes, sliding_window, logits_soft_cap"
            )
242
243

        if attn_type != AttentionType.DECODER:
244
245
246
247
248
249
            raise NotImplementedError(
                "Encoder self-attention and "
                "encoder/decoder cross-attention "
                "are not implemented for "
                "FlashMLAImpl"
            )
250
251
252

    def _forward_decode(
        self,
253
        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
254
255
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: FlashMLAMetadata,
256
        layer: AttentionLayer,
257
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
258
        # TODO: (zyongye) decode function for mla here
259
        assert kv_c_and_k_pe_cache.numel() > 0
260
261
        assert attn_metadata.decode is not None

262
263
        if type(q) is tuple:
            q = torch.cat(q, dim=-1)
264

265
        # mypy assertion: q is now always a tensor
266
        assert isinstance(q, torch.Tensor)
267
268
269
270

        num_decodes = attn_metadata.num_decodes
        q = reshape_query_for_spec_decode(q, num_decodes)

271
272
        tile_scheduler_metadata = attn_metadata.decode.tile_scheduler_metadata
        num_splits = attn_metadata.decode.num_splits
273
        if vllm_is_batch_invariant():
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
            device = q.device
            dtype = torch.int32

            B = q.shape[0]
            # block_table shape: [batch_size, max_num_blocks_per_seq]
            # The number of blocks per sequence is in the second dimension
            topk = attn_metadata.decode.block_table.shape[-1]
            B_TOPK = 64
            assert topk % B_TOPK == 0, f"topk ({topk}) must be divisible by {B_TOPK}"
            end_block_idx = topk // B_TOPK

            # Single partition => num_sm_parts = 1
            # TileSchedulerMetaDataSize = 8, layout:
            # [begin_idx, begin_block_idx, end_idx, end_block_idx,
            #  begin_n_split_idx, _, _, _]
            tile_scheduler_metadata = torch.zeros((1, 8), dtype=dtype, device=device)
            tile_scheduler_metadata[0, 0] = 0  # begin_idx
            tile_scheduler_metadata[0, 1] = 0  # sched_begin_block_idx
            tile_scheduler_metadata[0, 2] = B - 1  # end_idx
            tile_scheduler_metadata[0, 3] = end_block_idx
            tile_scheduler_metadata[0, 4] = 0  # begin_n_split_idx
            # fields [5..7] stay 0

            # Non-split path ignores num_splits, but the API requires it:
            # zeros of length B+1
            num_splits = torch.zeros((B + 1,), dtype=dtype, device=device)

301
        o, lse = flash_mla_with_kvcache(
302
            q=q,
303
            k_cache=kv_c_and_k_pe_cache.unsqueeze(-2),  # Add head dim of 1
304
305
            block_table=attn_metadata.decode.block_table,
            cache_seqlens=attn_metadata.decode.seq_lens,
306
            head_dim_v=self.kv_lora_rank,
307
308
            tile_scheduler_metadata=tile_scheduler_metadata,
            num_splits=num_splits,
309
310
            softmax_scale=self.scale,
            causal=True,
311
312
            descale_q=layer._q_scale.reshape(1),
            descale_k=layer._k_scale.reshape(1),
313
314
        )

315
316
        o = reshape_attn_output_for_spec_decode(o)

317
        return o, lse