[Cleanup][MLA] Remove FlashInfer DSpark DCP support (#53139)
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Summer Yang committed
Aug 21, 2026 at 11:42 UTC
d6c2fec9fd72eeac44f42c884e19a1c6bd1142a7
2 files changed
+55
-183
tests/v1/attention/test_mla_backends.py
+37
-62
@@ -886,79 +886,54 @@ def test_tokenspeed_mla_noncausal_capability():
886
assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()
887
888
889
-def test_flashinfer_mla_dcp_multi_token_decode_uses_per_query_bounds(monkeypatch):
889
+def test_flashinfer_mla_dspark_support_is_tp_only(monkeypatch):
890
flashinfer_mla_module = pytest.importorskip(
891
"vllm.v1.attention.backends.mla.flashinfer_mla"
892
)
893
-
894
- decode_call = None
895
-
896
- def fake_decode(**kwargs):
897
- nonlocal decode_call
898
- decode_call = kwargs
899
- query = kwargs["query"]
900
- output = torch.empty(*query.shape[:-1], 512, dtype=torch.bfloat16)
901
- lse = torch.empty(query.shape[0], query.shape[-2], dtype=torch.float32)
902
- return output, lse
903
-
904
- monkeypatch.setattr(
905
- flashinfer_mla_module,
906
- "trtllm_batch_decode_with_kv_cache_mla",
907
- fake_decode,
893
+ vllm_config = SimpleNamespace(
894
+ speculative_config=SimpleNamespace(method="dspark"),
895
+ parallel_config=SimpleNamespace(decode_context_parallel_size=2),
896
+ model_config=None,
897
)
898
monkeypatch.setattr(
899
flashinfer_mla_module,
911
- "_get_workspace_buffer",
912
- lambda return_lse: torch.empty(1, dtype=torch.int8),
900
+ "get_current_vllm_config",
901
+ lambda: vllm_config,
902
)
903
915
- impl = object.__new__(flashinfer_mla_module.FlashInferMLAImpl)
916
- impl.dcp_world_size = 2
917
- impl.dcp_rank = 1
918
- impl.cp_kv_cache_interleave_size = 1
919
- impl.need_to_return_lse_for_decode = True
920
- impl.kv_lora_rank = 512
921
- impl.qk_nope_head_dim = 128
922
- impl.qk_rope_head_dim = 64
923
- impl.bmm1_scale = 1.0
924
- impl.bmm2_scale = 1.0
925
-
926
- block_table = torch.tensor([[1], [2]], dtype=torch.int32)
927
- metadata = SimpleNamespace(
928
- num_decodes=2,
929
- num_decode_tokens=6,
930
- max_seq_len=7,
931
- causal=True,
932
- decode=SimpleNamespace(
933
- block_table=block_table,
934
- seq_lens=torch.tensor([5, 6], dtype=torch.int32),
935
- dcp_tot_seq_lens=torch.tensor([10, 13], dtype=torch.int32),
936
- flattened_block_table=None,
937
- flattened_seq_lens=None,
938
- query_len=0,
939
- ),
904
+ backend = flashinfer_mla_module.FlashInferMLABackend
905
+ builder = flashinfer_mla_module.FlashInferMLAMetadataBuilder
906
+ reason = backend.supports_combination(
907
+ head_size=576,
908
+ dtype=torch.bfloat16,
909
+ kv_cache_dtype="fp8",
910
+ block_size=64,
911
+ use_mla=True,
912
+ has_sink=False,
913
+ use_sparse=False,
914
+ use_mm_prefix=False,
915
+ device_capability=SimpleNamespace(),
916
)
941
- query = torch.empty(6, 2, 576, dtype=torch.bfloat16)
942
- kv_cache = torch.empty(3, 16, 576, dtype=torch.bfloat16)
917
944
- output, lse = impl.forward_mqa(
945
- query,
946
- kv_cache,
947
- metadata,
948
- SimpleNamespace(),
949
- )
918
+ assert reason == "FlashInfer MLA does not support DSpark with DCP"
919
+ assert backend.supports_non_causal()
920
+ assert builder.supports_non_causal_multi_token_decode
921
+ assert not backend.supports_non_causal_dcp()
922
951
- assert output.shape == (6, 2, 512)
952
- assert lse is not None
953
- assert lse.shape == (6, 2)
954
- assert decode_call is not None
955
- assert decode_call["query"].shape == (6, 1, 2, 576)
956
- torch.testing.assert_close(
957
- decode_call["seq_lens"],
958
- torch.tensor([4, 4, 5, 5, 6, 6], dtype=torch.int32),
959
- )
960
- torch.testing.assert_close(
961
- decode_call["block_tables"], block_table.repeat_interleave(3, dim=0)
923
+ vllm_config.parallel_config.decode_context_parallel_size = 1
924
+ assert (
925
+ backend.supports_combination(
926
+ head_size=576,
927
+ dtype=torch.bfloat16,
928
+ kv_cache_dtype="fp8",
929
+ block_size=64,
930
+ use_mla=True,
931
+ has_sink=False,
932
+ use_sparse=False,
933
+ use_mm_prefix=False,
934
+ device_capability=SimpleNamespace(),
935
+ )
936
+ is None
937
)
938
939
vllm/v1/attention/backends/mla/flashinfer_mla.py
+18
-121
@@ -1,8 +1,7 @@
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 TYPE_CHECKING, ClassVar
4
+from typing import ClassVar
5
6
import torch
7
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
@@ -16,7 +15,6 @@ from vllm.config.cache import CacheDType
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.attention.mla_attention import (
17
MLACommonBackend,
19
- MLACommonDecodeMetadata,
18
MLACommonImpl,
19
MLACommonMetadata,
20
MLACommonMetadataBuilder,
@@ -32,10 +30,6 @@ from vllm.v1.attention.backend import (
30
)
31
from vllm.v1.attention.backends.utils import KVCacheLayoutType
32
35
-if TYPE_CHECKING:
36
- from vllm.config import VllmConfig
37
- from vllm.v1.kv_cache_interface import AttentionSpec
38
-
33
logger = init_logger(__name__)
34
35
@@ -118,56 +112,12 @@ def _get_multi_ctas_kv_counter_buffer(
112
return _fi_multi_ctas_kv_counter
113
114
121
-@dataclass
122
-class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata):
123
- flattened_block_table: torch.Tensor | None = None
124
- flattened_seq_lens: torch.Tensor | None = None
125
- query_len: int = 0
126
-
127
-
128
-@dataclass
129
-class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]):
130
- pass
131
-
132
-
133
-class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]):
115
+class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
116
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
117
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
118
# Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
119
supports_non_causal_multi_token_decode: ClassVar[bool] = True
120
139
- def __init__(
140
- self,
141
- kv_cache_spec: "AttentionSpec",
142
- layer_names: list[str],
143
- vllm_config: "VllmConfig",
144
- device: torch.device,
145
- ) -> None:
146
- super().__init__(
147
- kv_cache_spec,
148
- layer_names,
149
- vllm_config,
150
- device,
151
- FlashInferMLAMetadata,
152
- supports_dcp_with_varlen=True,
153
- )
154
-
155
- def _build_decode(
156
- self,
157
- block_table_tensor: torch.Tensor,
158
- seq_lens_device: torch.Tensor,
159
- max_seq_len: int,
160
- query_start_loc_cpu: torch.Tensor,
161
- query_start_loc_device: torch.Tensor,
162
- num_decode_tokens: int,
163
- dcp_tot_seq_lens_device: torch.Tensor | None,
164
- ) -> FlashInferMLADecodeMetadata:
165
- return FlashInferMLADecodeMetadata(
166
- block_table=block_table_tensor,
167
- seq_lens=seq_lens_device,
168
- dcp_tot_seq_lens=dcp_tot_seq_lens_device,
169
- )
170
-
121
122
class FlashInferMLABackend(MLACommonBackend):
123
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -224,10 +174,16 @@ class FlashInferMLABackend(MLACommonBackend):
174
use_mm_prefix: bool,
175
device_capability: DeviceCapability,
176
) -> str | None:
227
- # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192]
228
- from vllm.config import get_current_vllm_config
229
-
177
vllm_config = get_current_vllm_config()
178
+ speculative_config = vllm_config.speculative_config
179
+ if (
180
+ speculative_config is not None
181
+ and speculative_config.method == "dspark"
182
+ and vllm_config.parallel_config.decode_context_parallel_size > 1
183
+ ):
184
+ return "FlashInfer MLA does not support DSpark with DCP"
185
+
186
+ # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192]
187
if vllm_config.model_config is not None:
188
hf_text_config = vllm_config.model_config.hf_text_config
189
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
@@ -243,7 +199,7 @@ class FlashInferMLABackend(MLACommonBackend):
199
return "HND"
200
201
246
-class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
202
+class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
203
can_return_lse_for_decode: bool = True
204
# trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own
205
# reference at flashinfer/trace/templates/attention.py:81:
@@ -314,7 +270,7 @@ class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
270
self,
271
q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
272
kv_c_and_k_pe_cache: torch.Tensor,
317
- attn_metadata: FlashInferMLAMetadata,
273
+ attn_metadata: MLACommonMetadata,
274
layer: AttentionLayer,
275
) -> tuple[torch.Tensor, torch.Tensor | None]:
276
assert kv_c_and_k_pe_cache.numel() > 0
@@ -326,28 +282,15 @@ class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
282
283
block_table = attn_metadata.decode.block_table
284
seq_lens = attn_metadata.decode.seq_lens
329
- query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
285
286
if not attn_metadata.causal:
332
- # Non-causal DSpark block: flatten to single-token decode rows with
333
- # per-row context seq_lens (trtllm-gen has no causal flag and would
334
- # otherwise mask the block causally).
287
+ # FlashInfer decode has no causal flag. For TP-only DSpark, flatten
288
+ # each non-causal query block into independent single-token rows.
289
+ query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
290
q = q.unsqueeze(1)
291
if query_len > 1:
337
- block_table, seq_lens = self._prepare_flattened_decode_metadata(
338
- attn_metadata,
339
- query_len,
340
- causal=False,
341
- )
342
- elif self.dcp_world_size > 1 and query_len > 1:
343
- # Causal DCP block: flatten to single-token decode rows with
344
- # per-row rank-local seq_lens for each query's visible prefix.
345
- block_table, seq_lens = self._prepare_flattened_decode_metadata(
346
- attn_metadata,
347
- query_len,
348
- causal=True,
349
- )
350
- q = q.unsqueeze(1)
292
+ block_table = block_table.repeat_interleave(query_len, dim=0)
293
+ seq_lens = seq_lens.repeat_interleave(query_len)
294
# trtllm API requires extra dimension q_len_per_request for MTP
295
elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
296
logger.warning_once(
@@ -420,49 +363,3 @@ class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
363
o = o.view(-1, o.shape[-2], o.shape[-1])
364
365
return o, lse
423
-
424
- def _prepare_flattened_decode_metadata(
425
- self,
426
- attn_metadata: FlashInferMLAMetadata,
427
- query_len: int,
428
- *,
429
- causal: bool,
430
- ) -> tuple[torch.Tensor, torch.Tensor]:
431
- """Prepare flattened decode metadata once for all layers in the group."""
432
- decode = attn_metadata.decode
433
- assert decode is not None
434
- if decode.query_len:
435
- assert decode.query_len == query_len
436
- assert decode.flattened_block_table is not None
437
- assert decode.flattened_seq_lens is not None
438
- return decode.flattened_block_table, decode.flattened_seq_lens
439
-
440
- block_table = decode.block_table.repeat_interleave(query_len, dim=0)
441
- if causal:
442
- global_seq_lens = decode.dcp_tot_seq_lens
443
- assert global_seq_lens is not None
444
- offsets = torch.arange(
445
- query_len - 1,
446
- -1,
447
- -1,
448
- device=global_seq_lens.device,
449
- dtype=global_seq_lens.dtype,
450
- )
451
- per_query_global_lens = torch.clamp(
452
- (global_seq_lens.unsqueeze(1) - offsets).reshape(-1), min=0
453
- )
454
- interleave = self.cp_kv_cache_interleave_size
455
- dcp_span = self.dcp_world_size * interleave
456
- remainder = torch.clamp(
457
- per_query_global_lens % dcp_span - self.dcp_rank * interleave,
458
- min=0,
459
- max=interleave,
460
- )
461
- seq_lens = per_query_global_lens // dcp_span * interleave + remainder
462
- else:
463
- seq_lens = decode.seq_lens.repeat_interleave(query_len)
464
-
465
- decode.flattened_block_table = block_table
466
- decode.flattened_seq_lens = seq_lens
467
- decode.query_len = query_len
468
- return block_table, seq_lens