iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention][DCP] Use FlashInfer native CP for MLA decode (#54012)

Browse Source 
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
This commit is contained in:
Summer Yang
2026-08-27 10:32:45 -07:00
committed by GitHub
parent 32ad1400d7
commit 07242faa46
3 changed files with 118 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_flashinfer_mla_dcp.py
+64 -1
View File
@@ -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)
tests/v1/attention/test_mla_backends.py
+2 -18
View File
@@ -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"),
vllm/v1/attention/backends/mla/flashinfer_mla.py
+52 -21
View File
@@ -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