[Attention] Add FLASH_ATTN_MLA_SPARSE backend for Hopper sparse MLA (#46189)
Signed-off-by: Dakai An <dakaian108@gmail.com>
Dakai An committed
Jun 24, 2026 at 17:33 UTC
dc0d318177e1086d483bff12ad71f1d3b5b980e7
4 files changed
+292
docs/design/attention_backends.md
+1
@@ -226,6 +226,7 @@ MLA decode backends are selected using the standard
226
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
227
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
228
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
229
+| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
230
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
231
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
232
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
vllm/platforms/cuda.py
+1
@@ -138,6 +138,7 @@ def _get_backend_priorities(
138
AttentionBackendEnum.FLASHMLA,
139
AttentionBackendEnum.FLASHINFER_MLA,
140
AttentionBackendEnum.TRITON_MLA,
141
+ AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
142
AttentionBackendEnum.FLASHMLA_SPARSE,
143
]
144
else:
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
new
+287
@@ -0,0 +1,287 @@
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 numpy as np
8
+import torch
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.config.cache import CacheDType
12
+from vllm.platforms.interface import DeviceCapability
13
+from vllm.utils.torch_utils import np_to_pinned_tensor
14
+from vllm.v1.attention.backend import (
15
+ AttentionBackend,
16
+ AttentionCGSupport,
17
+ AttentionLayer,
18
+ AttentionMetadata,
19
+ AttentionMetadataBuilder,
20
+ CommonAttentionMetadata,
21
+ MultipleOf,
22
+ SparseMLAAttentionImpl,
23
+)
24
+from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
25
+from vllm.v1.attention.backends.mla.sparse_utils import (
26
+ triton_convert_req_index_to_global_index,
27
+)
28
+from vllm.v1.kv_cache_interface import AttentionSpec
29
+from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
30
+
31
+
32
+class FlashAttnMLASparseBackend(AttentionBackend):
33
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
34
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
35
+ "auto",
36
+ "float16",
37
+ "bfloat16",
38
+ ]
39
+
40
+ @staticmethod
41
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
42
+ return [64]
43
+
44
+ @staticmethod
45
+ def get_name() -> str:
46
+ return "FLASH_ATTN_MLA_SPARSE"
47
+
48
+ @staticmethod
49
+ def get_builder_cls() -> type["FlashAttnMLASparseMetadataBuilder"]:
50
+ return FlashAttnMLASparseMetadataBuilder
51
+
52
+ @staticmethod
53
+ def get_impl_cls() -> type[SparseMLAAttentionImpl[Any]]:
54
+ return FlashAttnMLASparseImpl
55
+
56
+ @classmethod
57
+ def get_supported_head_sizes(cls) -> list[int]:
58
+ return []
59
+
60
+ @classmethod
61
+ def is_mla(cls) -> bool:
62
+ return True
63
+
64
+ @classmethod
65
+ def is_sparse(cls) -> bool:
66
+ return True
67
+
68
+ @classmethod
69
+ def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
70
+ return capability.major == 9
71
+
72
+ @classmethod
73
+ def supports_combination(
74
+ cls,
75
+ head_size: int,
76
+ dtype: torch.dtype,
77
+ kv_cache_dtype: CacheDType | None,
78
+ block_size: int | None,
79
+ use_mla: bool,
80
+ has_sink: bool,
81
+ use_sparse: bool,
82
+ use_mm_prefix: bool,
83
+ device_capability: DeviceCapability,
84
+ ) -> str | None:
85
+ if kv_cache_dtype not in (None, "auto", "float16", "bfloat16"):
86
+ return (
87
+ "FlashAttention MLA Sparse currently supports only FP16/BF16 KV cache"
88
+ )
89
+
90
+ if not flash_attn_supports_mla():
91
+ return "FlashAttention MLA not supported on this device"
92
+
93
+ from vllm.config import get_current_vllm_config_or_none
94
+
95
+ vllm_config = get_current_vllm_config_or_none()
96
+ if vllm_config is not None and vllm_config.model_config is not None:
97
+ if vllm_config.parallel_config.decode_context_parallel_size > 1:
98
+ return "FlashAttention MLA Sparse does not support DCP for now"
99
+
100
+ hf_config = vllm_config.model_config.hf_config
101
+ if not hasattr(hf_config, "index_topk"):
102
+ return "FlashAttention MLA Sparse requires model with index_topk"
103
+ return None
104
+
105
+ @staticmethod
106
+ def get_kv_cache_shape(
107
+ num_blocks: int,
108
+ block_size: int,
109
+ num_kv_heads: int,
110
+ head_size: int,
111
+ cache_dtype_str: str = "auto",
112
+ ) -> tuple[int, ...]:
113
+ return (num_blocks, block_size, head_size)
114
+
115
+
116
+@dataclass
117
+class FlashAttnMLASparseMetadata(AttentionMetadata):
118
+ num_reqs: int
119
+ max_query_len: int
120
+ max_seq_len: int
121
+
122
+ num_actual_tokens: int
123
+ query_start_loc: torch.Tensor
124
+ slot_mapping: torch.Tensor
125
+
126
+ block_table: torch.Tensor
127
+ req_id_per_token: torch.Tensor
128
+ block_size: int = 64
129
+ topk_tokens: int = 2048
130
+
131
+
132
+class FlashAttnMLASparseMetadataBuilder(
133
+ AttentionMetadataBuilder[FlashAttnMLASparseMetadata]
134
+):
135
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
136
+
137
+ def __init__(
138
+ self,
139
+ kv_cache_spec: AttentionSpec,
140
+ layer_names: list[str],
141
+ vllm_config: VllmConfig,
142
+ device: torch.device,
143
+ ) -> None:
144
+ self.vllm_config = vllm_config
145
+ self.layer_names = layer_names
146
+ self.kv_cache_spec = kv_cache_spec
147
+ self.model_config = vllm_config.model_config
148
+ self.device = device
149
+
150
+ self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
151
+
152
+ self.topk_tokens = vllm_config.model_config.hf_config.index_topk
153
+ self.req_id_per_token_buffer = torch.empty(
154
+ (vllm_config.scheduler_config.max_num_batched_tokens,),
155
+ dtype=torch.int32,
156
+ device=device,
157
+ )
158
+
159
+ def build(
160
+ self,
161
+ common_prefix_len: int,
162
+ common_attn_metadata: CommonAttentionMetadata,
163
+ fast_build: bool = False,
164
+ ) -> FlashAttnMLASparseMetadata:
165
+ cm = common_attn_metadata
166
+ num_tokens = cm.num_actual_tokens
167
+ starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
168
+ seg_lengths = np.diff(starts)
169
+ req_id_per_token = np.repeat(
170
+ np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
171
+ )
172
+
173
+ self.req_id_per_token_buffer.fill_(0)
174
+ self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
175
+ np_to_pinned_tensor(req_id_per_token), non_blocking=True
176
+ )
177
+
178
+ return FlashAttnMLASparseMetadata(
179
+ num_reqs=cm.num_reqs,
180
+ max_query_len=cm.max_query_len,
181
+ max_seq_len=cm.max_seq_len,
182
+ num_actual_tokens=cm.num_actual_tokens,
183
+ query_start_loc=cm.query_start_loc,
184
+ slot_mapping=cm.slot_mapping,
185
+ block_table=cm.block_table_tensor,
186
+ req_id_per_token=self.req_id_per_token_buffer[:num_tokens],
187
+ block_size=self.kv_cache_spec.block_size,
188
+ topk_tokens=self.topk_tokens,
189
+ )
190
+
191
+
192
+class FlashAttnMLASparseImpl(SparseMLAAttentionImpl[FlashAttnMLASparseMetadata]):
193
+ def __init__(
194
+ self,
195
+ num_heads: int,
196
+ head_size: int,
197
+ scale: float,
198
+ num_kv_heads: int,
199
+ alibi_slopes: list[float] | None,
200
+ sliding_window: int | None,
201
+ kv_cache_dtype: str,
202
+ logits_soft_cap: float | None,
203
+ attn_type: str,
204
+ kv_sharing_target_layer_name: str | None,
205
+ topk_indices_buffer: torch.Tensor | None = None,
206
+ indexer: Any | None = None,
207
+ **mla_args: Any,
208
+ ) -> None:
209
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
210
+ if any(unsupported_features):
211
+ raise NotImplementedError(
212
+ "FlashAttnMLASparseImpl does not support alibi, sliding window, "
213
+ "or logits soft cap."
214
+ )
215
+ if kv_cache_dtype not in ("auto", "float16", "bfloat16"):
216
+ raise NotImplementedError(
217
+ "FlashAttnMLASparseImpl currently supports only FP16/BF16 KV cache."
218
+ )
219
+
220
+ self.num_heads = num_heads
221
+ self.head_size = head_size
222
+ self.scale = float(scale)
223
+ self.num_kv_heads = num_kv_heads
224
+ self.kv_cache_dtype = kv_cache_dtype
225
+ self.kv_lora_rank: int = mla_args["kv_lora_rank"]
226
+ self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
227
+ self.topk_indices_buffer: torch.Tensor | None = (
228
+ indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
229
+ )
230
+ assert self.topk_indices_buffer is not None, (
231
+ "Indexer or topk_indices_buffer required for sparse MLA"
232
+ )
233
+ self.supports_quant_query_input = False
234
+ self.dcp_world_size = -1
235
+ self.q_pad_num_heads = None
236
+
237
+ def forward_mqa(
238
+ self,
239
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
240
+ kv_c_and_k_pe_cache: torch.Tensor,
241
+ attn_metadata: FlashAttnMLASparseMetadata,
242
+ layer: AttentionLayer,
243
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
244
+ if not isinstance(q, tuple):
245
+ raise NotImplementedError(
246
+ "FlashAttnMLASparseImpl expects split (q_nope, q_rope) input."
247
+ )
248
+ q_nope, q_rope = q
249
+ num_actual_toks = q_rope.shape[0]
250
+
251
+ assert self.topk_indices_buffer is not None
252
+ topk_indices = self.topk_indices_buffer[:num_actual_toks]
253
+ topk_indices, valid_counts = triton_convert_req_index_to_global_index(
254
+ attn_metadata.req_id_per_token[:num_actual_toks],
255
+ attn_metadata.block_table,
256
+ topk_indices,
257
+ BLOCK_SIZE=attn_metadata.block_size,
258
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
259
+ return_valid_counts=True,
260
+ )
261
+
262
+ cu_seqlens_q = torch.arange(
263
+ 0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device
264
+ )
265
+ kv_cache = kv_c_and_k_pe_cache.view(
266
+ -1, attn_metadata.block_size, self.head_size
267
+ )
268
+ k_cache = kv_cache[:, :, self.kv_lora_rank :].view(
269
+ -1, 1, 1, self.qk_rope_head_dim
270
+ )
271
+ v_cache = kv_cache[:, :, : self.kv_lora_rank].view(-1, 1, 1, self.kv_lora_rank)
272
+
273
+ out = flash_attn_varlen_func(
274
+ q=q_rope,
275
+ k=k_cache,
276
+ v=v_cache,
277
+ q_v=q_nope,
278
+ max_seqlen_q=1,
279
+ cu_seqlens_q=cu_seqlens_q,
280
+ max_seqlen_k=topk_indices.shape[1],
281
+ seqused_k=valid_counts,
282
+ block_table=topk_indices,
283
+ softmax_scale=self.scale,
284
+ causal=True,
285
+ fa_version=3,
286
+ )
287
+ return out, None
vllm/v1/attention/backends/registry.py
+3
@@ -95,6 +95,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
95
"vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend"
96
)
97
FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
98
+ FLASH_ATTN_MLA_SPARSE = (
99
+ "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend"
100
+ )
101
MINIMAX_M3_SPARSE = (
102
"vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend"
103
)