From 07242faa4667adcec25e40cb5f580cec60f2faac Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Thu, 27 Aug 2026 10:32:45 -0700 Subject: [PATCH] [Attention][DCP] Use FlashInfer native CP for MLA decode (#54012) Signed-off-by: Summer Yang --- tests/v1/attention/test_flashinfer_mla_dcp.py | 65 ++++++++++++++++- tests/v1/attention/test_mla_backends.py | 20 +---- .../attention/backends/mla/flashinfer_mla.py | 73 +++++++++++++------ 3 files changed, 118 insertions(+), 40 deletions(-) diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py index 1841b25d49..4fe545407f 100644 --- a/tests/v1/attention/test_flashinfer_mla_dcp.py +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -45,7 +45,9 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): impl.bmm1_scale = 1.0 impl.bmm2_scale = 1.0 impl.need_to_return_lse_for_decode = True - impl.dcp_world_size = 2 + impl.dcp_world_size = 1 + impl.dcp_rank = 0 + impl.cp_kv_cache_interleave_size = 1 impl.num_heads = 6 impl.qk_nope_head_dim = 128 impl.kv_lora_rank = 512 @@ -79,3 +81,64 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): assert kernel.call_args.kwargs["backend"] == "cute-dsl" assert output.shape == (2, 24, 512) assert lse is not None and lse.shape == (2, 24) + + +@requires_flashinfer_mla +def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch): + import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla + + impl = MagicMock() + impl.bmm1_scale = 1.0 + impl.bmm2_scale = 1.0 + impl.need_to_return_lse_for_decode = True + impl.num_heads = 6 + impl.qk_nope_head_dim = 128 + impl.kv_lora_rank = 512 + impl.qk_rope_head_dim = 64 + impl.kv_cache_dtype = "auto" + impl.dcp_world_size = 8 + impl.dcp_rank = 3 + impl.cp_kv_cache_interleave_size = 1 + + num_reqs, query_len = 2, 3 + num_tokens = num_reqs * query_len + block_table = torch.tensor([[10], [20]], dtype=torch.int32) + seq_lens = torch.tensor([4, 5], dtype=torch.int32) + global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32) + attn_metadata = MagicMock() + attn_metadata.causal = True + attn_metadata.num_decode_tokens = num_tokens + attn_metadata.num_decodes = num_reqs + attn_metadata.max_seq_len = 9 + attn_metadata.decode.block_table = block_table + attn_metadata.decode.seq_lens = seq_lens + attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens + + query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16) + kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16) + kernel = MagicMock( + return_value=( + torch.ones(num_reqs, query_len, 24, 512, dtype=torch.bfloat16), + torch.ones(num_reqs, query_len, 24, dtype=torch.float32), + ) + ) + monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock()) + monkeypatch.setattr(flashinfer_mla, "trtllm_batch_decode_with_kv_cache_mla", kernel) + + output, lse = flashinfer_mla.FlashInferMLAImpl.forward_mqa( + impl, query, kv_cache, attn_metadata, MagicMock() + ) + + call = kernel.call_args.kwargs + assert call["query"].shape == (num_reqs, query_len, 24, 576) + assert call["block_tables"] is block_table + assert call["seq_lens"] is seq_lens + assert call["backend"] == "cute-dsl" + assert call["enable_dcp"] is True + assert call["cp_world"] == 8 + assert call["cp_rank"] == 3 + assert call["causal_seqlens_kv_global"] is global_causal_seq_lens + assert "multi_ctas_kv_counter_buffer" not in call + assert flashinfer_mla.FlashInferMLAImpl.lse_base_on_e + assert output.shape == (num_tokens, 24, 512) + assert lse is not None and lse.shape == (num_tokens, 24) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cbccbfc9a2..f054d1a7e0 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -902,7 +902,7 @@ def test_tokenspeed_mla_noncausal_capability(): assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal() -def test_flashinfer_mla_dspark_support_is_tp_only(monkeypatch): +def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch): flashinfer_mla_module = pytest.importorskip( "vllm.v1.attention.backends.mla.flashinfer_mla" ) @@ -931,27 +931,11 @@ def test_flashinfer_mla_dspark_support_is_tp_only(monkeypatch): device_capability=SimpleNamespace(), ) - assert reason == "FlashInfer MLA does not support DSpark with DCP" + assert reason is None assert backend.supports_non_causal() assert builder.supports_non_causal_multi_token_decode assert not backend.supports_non_causal_dcp() - vllm_config.parallel_config.decode_context_parallel_size = 1 - assert ( - backend.supports_combination( - head_size=576, - dtype=torch.bfloat16, - kv_cache_dtype="fp8", - block_size=64, - use_mla=True, - has_sink=False, - use_sparse=False, - use_mm_prefix=False, - device_capability=SimpleNamespace(), - ) - is None - ) - @pytest.mark.parametrize( ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"), diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index 99dd7f1625..fbe5f6844f 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import torch from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla @@ -29,6 +29,10 @@ from vllm.v1.attention.backend import ( MultipleOf, ) +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import AttentionSpec + logger = init_logger(__name__) @@ -117,6 +121,31 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa. supports_non_causal_multi_token_decode: ClassVar[bool] = True + def __init__( + self, + kv_cache_spec: "AttentionSpec", + layer_names: list[str], + vllm_config: "VllmConfig", + device: torch.device, + ) -> None: + parallel_config = vllm_config.parallel_config + dcp_size = parallel_config.decode_context_parallel_size + interleave_size = parallel_config.cp_kv_cache_interleave_size + if dcp_size > 1 and interleave_size != 1: + raise ValueError( + "FlashInfer MLA native DCP requires " + "cp_kv_cache_interleave_size=1; got " + f"{interleave_size}." + ) + super().__init__( + kv_cache_spec, + layer_names, + vllm_config, + device, + MLACommonMetadata, + supports_dcp_with_varlen=True, + ) + class FlashInferMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -165,16 +194,8 @@ class FlashInferMLABackend(MLACommonBackend): use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: - vllm_config = get_current_vllm_config() - speculative_config = vllm_config.speculative_config - if ( - speculative_config is not None - and speculative_config.method == "dspark" - and vllm_config.parallel_config.decode_context_parallel_size > 1 - ): - return "FlashInfer MLA does not support DSpark with DCP" - # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192] + vllm_config = get_current_vllm_config() if vllm_config.model_config is not None: hf_text_config = vllm_config.model_config.hf_text_config qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) @@ -188,13 +209,9 @@ class FlashInferMLABackend(MLACommonBackend): class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): can_return_lse_for_decode: bool = True - # trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own - # reference at flashinfer/trace/templates/attention.py:81: - # `logsumexp / log(2.0)`). Override the AttentionImplBase default - # so MLAAttention's DCP combine branches on the correct base - # (IS_BASE_E=False uses tl.exp2/tl.log2 natively, avoiding an FP - # multiply per decode step). - lse_base_on_e: bool = False + # DCP is the only path that consumes LSE. It uses monolithic CuTeDSL, + # whose public LSE contract is natural-log. + lse_base_on_e: bool = True def __init__( self, @@ -303,10 +320,23 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): workspace_buffer = _get_workspace_buffer(return_lse) # Parallel gathers can change the runtime Q heads from TP-local num_heads. runtime_num_heads = q.shape[-2] - # trtllm-gen rejects MLA head counts it can't tile (e.g. 96); - # fall back to cute-dsl for those. - decode_backend = _select_mla_decode_backend(runtime_num_heads) - extra_kwargs = {} + extra_kwargs: dict[str, Any] = {} + decode_backend: str | None + if self.dcp_world_size > 1: + assert self.cp_kv_cache_interleave_size == 1 + causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens + assert causal_seqlens_kv_global is not None + extra_kwargs.update( + enable_dcp=True, + cp_world=self.dcp_world_size, + cp_rank=self.dcp_rank, + causal_seqlens_kv_global=causal_seqlens_kv_global, + ) + decode_backend = "cute-dsl" + else: + # trtllm-gen rejects MLA head counts it can't tile (e.g. 96); + # fall back to cute-dsl for those. + decode_backend = _select_mla_decode_backend(runtime_num_heads) if decode_backend: extra_kwargs["backend"] = decode_backend elif kv_c_and_k_pe_cache.shape[-2] in (32, 64): @@ -343,6 +373,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): ) if return_lse: o, lse = kernel_out + lse = lse.view(-1, lse.shape[-1]) else: o, lse = kernel_out, None