[Attention][DCP] Use FlashInfer native CP for MLA decode (#54012)
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Summer Yang committed
Aug 27, 2026 at 10:32 UTC
07242faa4667adcec25e40cb5f580cec60f2faac
3 files changed
+118
-40
tests/v1/attention/test_flashinfer_mla_dcp.py
+64
-1
@@ -45,7 +45,9 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
45
impl.bmm1_scale = 1.0
46
impl.bmm2_scale = 1.0
47
impl.need_to_return_lse_for_decode = True
48
- impl.dcp_world_size = 2
48
+ impl.dcp_world_size = 1
49
+ impl.dcp_rank = 0
50
+ impl.cp_kv_cache_interleave_size = 1
51
impl.num_heads = 6
52
impl.qk_nope_head_dim = 128
53
impl.kv_lora_rank = 512
@@ -79,3 +81,64 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
81
assert kernel.call_args.kwargs["backend"] == "cute-dsl"
82
assert output.shape == (2, 24, 512)
83
assert lse is not None and lse.shape == (2, 24)
84
+
85
+
86
+@requires_flashinfer_mla
87
+def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
88
+ import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla
89
+
90
+ impl = MagicMock()
91
+ impl.bmm1_scale = 1.0
92
+ impl.bmm2_scale = 1.0
93
+ impl.need_to_return_lse_for_decode = True
94
+ impl.num_heads = 6
95
+ impl.qk_nope_head_dim = 128
96
+ impl.kv_lora_rank = 512
97
+ impl.qk_rope_head_dim = 64
98
+ impl.kv_cache_dtype = "auto"
99
+ impl.dcp_world_size = 8
100
+ impl.dcp_rank = 3
101
+ impl.cp_kv_cache_interleave_size = 1
102
+
103
+ num_reqs, query_len = 2, 3
104
+ num_tokens = num_reqs * query_len
105
+ block_table = torch.tensor([[10], [20]], dtype=torch.int32)
106
+ seq_lens = torch.tensor([4, 5], dtype=torch.int32)
107
+ global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32)
108
+ attn_metadata = MagicMock()
109
+ attn_metadata.causal = True
110
+ attn_metadata.num_decode_tokens = num_tokens
111
+ attn_metadata.num_decodes = num_reqs
112
+ attn_metadata.max_seq_len = 9
113
+ attn_metadata.decode.block_table = block_table
114
+ attn_metadata.decode.seq_lens = seq_lens
115
+ attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens
116
+
117
+ query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16)
118
+ kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16)
119
+ kernel = MagicMock(
120
+ return_value=(
121
+ torch.ones(num_reqs, query_len, 24, 512, dtype=torch.bfloat16),
122
+ torch.ones(num_reqs, query_len, 24, dtype=torch.float32),
123
+ )
124
+ )
125
+ monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock())
126
+ monkeypatch.setattr(flashinfer_mla, "trtllm_batch_decode_with_kv_cache_mla", kernel)
127
+
128
+ output, lse = flashinfer_mla.FlashInferMLAImpl.forward_mqa(
129
+ impl, query, kv_cache, attn_metadata, MagicMock()
130
+ )
131
+
132
+ call = kernel.call_args.kwargs
133
+ assert call["query"].shape == (num_reqs, query_len, 24, 576)
134
+ assert call["block_tables"] is block_table
135
+ assert call["seq_lens"] is seq_lens
136
+ assert call["backend"] == "cute-dsl"
137
+ assert call["enable_dcp"] is True
138
+ assert call["cp_world"] == 8
139
+ assert call["cp_rank"] == 3
140
+ assert call["causal_seqlens_kv_global"] is global_causal_seq_lens
141
+ assert "multi_ctas_kv_counter_buffer" not in call
142
+ assert flashinfer_mla.FlashInferMLAImpl.lse_base_on_e
143
+ assert output.shape == (num_tokens, 24, 512)
144
+ assert lse is not None and lse.shape == (num_tokens, 24)
tests/v1/attention/test_mla_backends.py
+2
-18
@@ -902,7 +902,7 @@ def test_tokenspeed_mla_noncausal_capability():
902
assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()
903
904
905
-def test_flashinfer_mla_dspark_support_is_tp_only(monkeypatch):
905
+def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch):
906
flashinfer_mla_module = pytest.importorskip(
907
"vllm.v1.attention.backends.mla.flashinfer_mla"
908
)
@@ -931,27 +931,11 @@ def test_flashinfer_mla_dspark_support_is_tp_only(monkeypatch):
931
device_capability=SimpleNamespace(),
932
)
933
934
- assert reason == "FlashInfer MLA does not support DSpark with DCP"
934
+ assert reason is None
935
assert backend.supports_non_causal()
936
assert builder.supports_non_causal_multi_token_decode
937
assert not backend.supports_non_causal_dcp()
938
939
- vllm_config.parallel_config.decode_context_parallel_size = 1
940
- assert (
941
- backend.supports_combination(
942
- head_size=576,
943
- dtype=torch.bfloat16,
944
- kv_cache_dtype="fp8",
945
- block_size=64,
946
- use_mla=True,
947
- has_sink=False,
948
- use_sparse=False,
949
- use_mm_prefix=False,
950
- device_capability=SimpleNamespace(),
951
- )
952
- is None
953
- )
954
-
939
940
@pytest.mark.parametrize(
941
("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"),
vllm/v1/attention/backends/mla/flashinfer_mla.py
+52
-21
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import ClassVar
4
+from typing import TYPE_CHECKING, Any, ClassVar
5
6
import torch
7
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
@@ -29,6 +29,10 @@ from vllm.v1.attention.backend import (
29
MultipleOf,
30
)
31
32
+if TYPE_CHECKING:
33
+ from vllm.config import VllmConfig
34
+ from vllm.v1.kv_cache_interface import AttentionSpec
35
+
36
logger = init_logger(__name__)
37
38
@@ -117,6 +121,31 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
121
# Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
122
supports_non_causal_multi_token_decode: ClassVar[bool] = True
123
124
+ def __init__(
125
+ self,
126
+ kv_cache_spec: "AttentionSpec",
127
+ layer_names: list[str],
128
+ vllm_config: "VllmConfig",
129
+ device: torch.device,
130
+ ) -> None:
131
+ parallel_config = vllm_config.parallel_config
132
+ dcp_size = parallel_config.decode_context_parallel_size
133
+ interleave_size = parallel_config.cp_kv_cache_interleave_size
134
+ if dcp_size > 1 and interleave_size != 1:
135
+ raise ValueError(
136
+ "FlashInfer MLA native DCP requires "
137
+ "cp_kv_cache_interleave_size=1; got "
138
+ f"{interleave_size}."
139
+ )
140
+ super().__init__(
141
+ kv_cache_spec,
142
+ layer_names,
143
+ vllm_config,
144
+ device,
145
+ MLACommonMetadata,
146
+ supports_dcp_with_varlen=True,
147
+ )
148
+
149
150
class FlashInferMLABackend(MLACommonBackend):
151
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -165,16 +194,8 @@ class FlashInferMLABackend(MLACommonBackend):
194
use_mm_prefix: bool,
195
device_capability: DeviceCapability,
196
) -> str | None:
168
- vllm_config = get_current_vllm_config()
169
- speculative_config = vllm_config.speculative_config
170
- if (
171
- speculative_config is not None
172
- and speculative_config.method == "dspark"
173
- and vllm_config.parallel_config.decode_context_parallel_size > 1
174
- ):
175
- return "FlashInfer MLA does not support DSpark with DCP"
176
-
197
# FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192]
198
+ vllm_config = get_current_vllm_config()
199
if vllm_config.model_config is not None:
200
hf_text_config = vllm_config.model_config.hf_text_config
201
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
@@ -188,13 +209,9 @@ class FlashInferMLABackend(MLACommonBackend):
209
210
class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
211
can_return_lse_for_decode: bool = True
191
- # trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own
192
- # reference at flashinfer/trace/templates/attention.py:81:
193
- # `logsumexp / log(2.0)`). Override the AttentionImplBase default
194
- # so MLAAttention's DCP combine branches on the correct base
195
- # (IS_BASE_E=False uses tl.exp2/tl.log2 natively, avoiding an FP
196
- # multiply per decode step).
197
- lse_base_on_e: bool = False
212
+ # DCP is the only path that consumes LSE. It uses monolithic CuTeDSL,
213
+ # whose public LSE contract is natural-log.
214
+ lse_base_on_e: bool = True
215
216
def __init__(
217
self,
@@ -303,10 +320,23 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
320
workspace_buffer = _get_workspace_buffer(return_lse)
321
# Parallel gathers can change the runtime Q heads from TP-local num_heads.
322
runtime_num_heads = q.shape[-2]
306
- # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);
307
- # fall back to cute-dsl for those.
308
- decode_backend = _select_mla_decode_backend(runtime_num_heads)
309
- extra_kwargs = {}
323
+ extra_kwargs: dict[str, Any] = {}
324
+ decode_backend: str | None
325
+ if self.dcp_world_size > 1:
326
+ assert self.cp_kv_cache_interleave_size == 1
327
+ causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens
328
+ assert causal_seqlens_kv_global is not None
329
+ extra_kwargs.update(
330
+ enable_dcp=True,
331
+ cp_world=self.dcp_world_size,
332
+ cp_rank=self.dcp_rank,
333
+ causal_seqlens_kv_global=causal_seqlens_kv_global,
334
+ )
335
+ decode_backend = "cute-dsl"
336
+ else:
337
+ # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);
338
+ # fall back to cute-dsl for those.
339
+ decode_backend = _select_mla_decode_backend(runtime_num_heads)
340
if decode_backend:
341
extra_kwargs["backend"] = decode_backend
342
elif kv_c_and_k_pe_cache.shape[-2] in (32, 64):
@@ -343,6 +373,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
373
)
374
if return_lse:
375
o, lse = kernel_out
376
+ lse = lse.view(-1, lse.shape[-1])
377
else:
378
o, lse = kernel_out, None
379