| 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 Any, ClassVar |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.config.cache import CacheDType |
| 11 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 12 | SparseMLACommonImpl, |
| 13 | SparseMLACommonMetadata, |
| 14 | SparseMLACommonMetadataBuilder, |
| 15 | ) |
| 16 | from vllm.platforms.interface import DeviceCapability |
| 17 | from vllm.v1.attention.backend import ( |
| 18 | AttentionBackend, |
| 19 | AttentionCGSupport, |
| 20 | AttentionLayer, |
| 21 | MLAAttentionImpl, |
| 22 | MultipleOf, |
| 23 | ) |
| 24 | from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla |
| 25 | from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup |
| 26 | from vllm.v1.attention.backends.mla.sparse_utils import ( |
| 27 | flat_kv_row_view, |
| 28 | triton_convert_req_index_to_global_index, |
| 29 | ) |
| 30 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 31 | from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func |
| 32 | |
| 33 | |
| 34 | class FlashAttnMLASparseBackend(AttentionBackend): |
| 35 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] |
| 36 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 37 | "auto", |
| 38 | "float16", |
| 39 | "bfloat16", |
| 40 | ] |
| 41 | |
| 42 | @staticmethod |
| 43 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 44 | return [64] |
| 45 | |
| 46 | @staticmethod |
| 47 | def get_name() -> str: |
| 48 | return "FLASH_ATTN_MLA_SPARSE" |
| 49 | |
| 50 | @staticmethod |
| 51 | def get_builder_cls() -> type["FlashAttnMLASparseMetadataBuilder"]: |
| 52 | return FlashAttnMLASparseMetadataBuilder |
| 53 | |
| 54 | @staticmethod |
| 55 | def get_impl_cls() -> type[MLAAttentionImpl[Any]]: |
| 56 | return FlashAttnMLASparseImpl |
| 57 | |
| 58 | @classmethod |
| 59 | def get_supported_head_sizes(cls) -> list[int]: |
| 60 | return [] |
| 61 | |
| 62 | @classmethod |
| 63 | def is_mla(cls) -> bool: |
| 64 | return True |
| 65 | |
| 66 | @classmethod |
| 67 | def is_sparse(cls) -> bool: |
| 68 | return True |
| 69 | |
| 70 | @classmethod |
| 71 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 72 | return capability.major == 9 |
| 73 | |
| 74 | @classmethod |
| 75 | def supports_combination( |
| 76 | cls, |
| 77 | head_size: int, |
| 78 | dtype: torch.dtype, |
| 79 | kv_cache_dtype: CacheDType | None, |
| 80 | block_size: int | None, |
| 81 | use_mla: bool, |
| 82 | has_sink: bool, |
| 83 | use_sparse: bool, |
| 84 | use_mm_prefix: bool, |
| 85 | device_capability: DeviceCapability, |
| 86 | ) -> str | None: |
| 87 | if kv_cache_dtype not in (None, "auto", "float16", "bfloat16"): |
| 88 | return ( |
| 89 | "FlashAttention MLA Sparse currently supports only FP16/BF16 KV cache" |
| 90 | ) |
| 91 | |
| 92 | if not flash_attn_supports_mla(): |
| 93 | return "FlashAttention MLA not supported on this device" |
| 94 | |
| 95 | from vllm.config import get_current_vllm_config_or_none |
| 96 | |
| 97 | vllm_config = get_current_vllm_config_or_none() |
| 98 | if vllm_config is not None and vllm_config.model_config is not None: |
| 99 | if vllm_config.parallel_config.decode_context_parallel_size > 1: |
| 100 | return "FlashAttention MLA Sparse does not support DCP for now" |
| 101 | |
| 102 | hf_text_config = vllm_config.model_config.hf_text_config |
| 103 | if not hasattr(hf_text_config, "index_topk"): |
| 104 | return "FlashAttention MLA Sparse requires model with index_topk" |
| 105 | return None |
| 106 | |
| 107 | |
| 108 | @dataclass |
| 109 | class FlashAttnMLASparseMetadata(SparseMLACommonMetadata): |
| 110 | pass |
| 111 | |
| 112 | |
| 113 | class FlashAttnMLASparseMetadataBuilder( |
| 114 | SparseMLACommonMetadataBuilder[FlashAttnMLASparseMetadata] |
| 115 | ): |
| 116 | metadata_cls = FlashAttnMLASparseMetadata |
| 117 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 118 | |
| 119 | def __init__( |
| 120 | self, |
| 121 | kv_cache_spec: AttentionSpec, |
| 122 | layer_names: list[str], |
| 123 | vllm_config: VllmConfig, |
| 124 | device: torch.device, |
| 125 | ) -> None: |
| 126 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 127 | |
| 128 | num_q_heads = self.model_config.get_num_attention_heads( |
| 129 | vllm_config.parallel_config |
| 130 | ) |
| 131 | threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) |
| 132 | self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True) |
| 133 | |
| 134 | |
| 135 | class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): |
| 136 | def __init__( |
| 137 | self, |
| 138 | num_heads: int, |
| 139 | head_size: int, |
| 140 | scale: float, |
| 141 | num_kv_heads: int, |
| 142 | alibi_slopes: list[float] | None, |
| 143 | sliding_window: int | None, |
| 144 | kv_cache_dtype: str, |
| 145 | logits_soft_cap: float | None, |
| 146 | attn_type: str, |
| 147 | kv_sharing_target_layer_name: str | None, |
| 148 | topk_indices_buffer: torch.Tensor | None = None, |
| 149 | indexer: Any | None = None, |
| 150 | **mla_args: Any, |
| 151 | ) -> None: |
| 152 | unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] |
| 153 | if any(unsupported_features): |
| 154 | raise NotImplementedError( |
| 155 | "FlashAttnMLASparseImpl does not support alibi, sliding window, " |
| 156 | "or logits soft cap." |
| 157 | ) |
| 158 | if kv_cache_dtype not in ("auto", "float16", "bfloat16"): |
| 159 | raise NotImplementedError( |
| 160 | "FlashAttnMLASparseImpl currently supports only FP16/BF16 KV cache." |
| 161 | ) |
| 162 | |
| 163 | super().__init__( |
| 164 | num_heads, |
| 165 | head_size, |
| 166 | scale, |
| 167 | num_kv_heads, |
| 168 | alibi_slopes, |
| 169 | sliding_window, |
| 170 | kv_cache_dtype, |
| 171 | logits_soft_cap, |
| 172 | attn_type, |
| 173 | kv_sharing_target_layer_name, |
| 174 | indexer=indexer, |
| 175 | topk_indices_buffer=topk_indices_buffer, |
| 176 | **mla_args, |
| 177 | ) |
| 178 | assert self.topk_indices_buffer is not None, ( |
| 179 | "Indexer or topk_indices_buffer required for sparse MLA" |
| 180 | ) |
| 181 | self.cu_seqlens_q_buffer = torch.arange( |
| 182 | self.topk_indices_buffer.shape[0] + 1, |
| 183 | dtype=torch.int32, |
| 184 | device=self.topk_indices_buffer.device, |
| 185 | ) |
| 186 | self.supports_quant_query_input = False |
| 187 | |
| 188 | def forward_mqa( |
| 189 | self, |
| 190 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 191 | kv_c_and_k_pe_cache: torch.Tensor, |
| 192 | attn_metadata: FlashAttnMLASparseMetadata, |
| 193 | layer: AttentionLayer, |
| 194 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 195 | if not isinstance(q, tuple): |
| 196 | raise NotImplementedError( |
| 197 | "FlashAttnMLASparseImpl expects split (q_nope, q_rope) input." |
| 198 | ) |
| 199 | q_nope, q_rope = q |
| 200 | num_actual_toks = q_rope.shape[0] |
| 201 | |
| 202 | assert self.topk_indices_buffer is not None |
| 203 | topk_indices = self.topk_indices_buffer[:num_actual_toks] |
| 204 | index_group = self.index_group |
| 205 | if isinstance(index_group, HiSparseMLAIndexGroup): |
| 206 | num_decode_tokens = attn_metadata.num_decode_tokens |
| 207 | outputs = [] |
| 208 | if num_decode_tokens: |
| 209 | physical_topk, valid_counts = ( |
| 210 | index_group.convert_logical_to_physical_topk( |
| 211 | self.index_group_index, |
| 212 | topk_indices[:num_decode_tokens], |
| 213 | attn_metadata, |
| 214 | block_stride_rows=None, |
| 215 | return_valid_counts=True, |
| 216 | ) |
| 217 | ) |
| 218 | outputs.append( |
| 219 | self._run_mqa_kernel( |
| 220 | q_nope[:num_decode_tokens], |
| 221 | q_rope[:num_decode_tokens], |
| 222 | index_group.physical_kv_cache(self.index_group_index).view( |
| 223 | kv_c_and_k_pe_cache.dtype |
| 224 | ), |
| 225 | physical_topk, |
| 226 | valid_counts, |
| 227 | attn_metadata.block_size, |
| 228 | ) |
| 229 | ) |
| 230 | if num_decode_tokens < num_actual_toks: |
| 231 | cache = index_group.cache(self.index_group_index) |
| 232 | if num_decode_tokens == 0 and cache.all_context_pages_resident: |
| 233 | physical_topk, valid_counts = ( |
| 234 | index_group.convert_logical_to_physical_topk( |
| 235 | self.index_group_index, |
| 236 | topk_indices, |
| 237 | attn_metadata, |
| 238 | block_stride_rows=None, |
| 239 | return_valid_counts=True, |
| 240 | ) |
| 241 | ) |
| 242 | prefill_cache = index_group.physical_kv_cache( |
| 243 | self.index_group_index |
| 244 | ).view(kv_c_and_k_pe_cache.dtype) |
| 245 | else: |
| 246 | prefill_cache, block_table, req_ids = ( |
| 247 | index_group.stage_prefill_rows( |
| 248 | self.index_group_index, |
| 249 | kv_c_and_k_pe_cache, |
| 250 | attn_metadata, |
| 251 | ) |
| 252 | ) |
| 253 | physical_topk, valid_counts = ( |
| 254 | triton_convert_req_index_to_global_index( |
| 255 | req_ids, |
| 256 | block_table, |
| 257 | topk_indices[num_decode_tokens:], |
| 258 | BLOCK_SIZE=attn_metadata.block_size, |
| 259 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 260 | return_valid_counts=True, |
| 261 | ) |
| 262 | ) |
| 263 | outputs.append( |
| 264 | self._run_mqa_kernel( |
| 265 | q_nope[num_decode_tokens:], |
| 266 | q_rope[num_decode_tokens:], |
| 267 | prefill_cache, |
| 268 | physical_topk, |
| 269 | valid_counts, |
| 270 | attn_metadata.block_size, |
| 271 | ) |
| 272 | ) |
| 273 | return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None |
| 274 | |
| 275 | kv_rows, block_stride_rows = flat_kv_row_view( |
| 276 | kv_c_and_k_pe_cache, attn_metadata.block_size |
| 277 | ) |
| 278 | topk_indices, valid_counts = triton_convert_req_index_to_global_index( |
| 279 | attn_metadata.req_id_per_token[:num_actual_toks], |
| 280 | attn_metadata.block_table, |
| 281 | topk_indices, |
| 282 | BLOCK_SIZE=attn_metadata.block_size, |
| 283 | BLOCK_STRIDE_ROWS=block_stride_rows, |
| 284 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 285 | return_valid_counts=True, |
| 286 | ) |
| 287 | return ( |
| 288 | self._run_mqa_kernel( |
| 289 | q_nope, |
| 290 | q_rope, |
| 291 | kv_rows, |
| 292 | topk_indices, |
| 293 | valid_counts, |
| 294 | attn_metadata.block_size, |
| 295 | cache_is_flat=True, |
| 296 | ), |
| 297 | None, |
| 298 | ) |
| 299 | |
| 300 | def _run_mqa_kernel( |
| 301 | self, |
| 302 | q_nope: torch.Tensor, |
| 303 | q_rope: torch.Tensor, |
| 304 | kv_cache: torch.Tensor, |
| 305 | topk_indices: torch.Tensor, |
| 306 | valid_counts: torch.Tensor, |
| 307 | block_size: int, |
| 308 | *, |
| 309 | cache_is_flat: bool = False, |
| 310 | ) -> torch.Tensor: |
| 311 | kv_rows = ( |
| 312 | kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0] |
| 313 | ) |
| 314 | |
| 315 | cu_seqlens_q = self.cu_seqlens_q_buffer[: q_rope.shape[0] + 1] |
| 316 | v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) |
| 317 | if self.qk_rope_head_dim == 0: |
| 318 | # FA3's QV path requires the 64-wide Q/K specialization. For NoPE |
| 319 | # MLA, zero Q preserves QV-only attention scores while providing a |
| 320 | # valid TMA shape. |
| 321 | _FA3_QV_HEAD_DIM = 64 |
| 322 | q_rope = q_nope.new_zeros(*q_nope.shape[:-1], _FA3_QV_HEAD_DIM) |
| 323 | k_cache = v_cache[..., :_FA3_QV_HEAD_DIM] |
| 324 | else: |
| 325 | k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1) |
| 326 | |
| 327 | out = flash_attn_varlen_func( |
| 328 | q=q_rope, |
| 329 | k=k_cache, |
| 330 | v=v_cache, |
| 331 | q_v=q_nope, |
| 332 | max_seqlen_q=1, |
| 333 | cu_seqlens_q=cu_seqlens_q, |
| 334 | max_seqlen_k=topk_indices.shape[1], |
| 335 | seqused_k=valid_counts, |
| 336 | block_table=topk_indices, |
| 337 | softmax_scale=self.scale, |
| 338 | causal=True, |
| 339 | fa_version=3, |
| 340 | ) |
| 341 | return out |