mirrors/vllm
0
0 Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-10-28 09:45:26 +08:00
Code Activity

[Bugfix][Attention] Fix FlashInfer MLA DCP decode on ragged spec-decode batches (#59392)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: Pavani Majety <pmajety@nvidia.com>
This commit is contained in:
Jiangyun ZhuClaude Opus 5.5Kevin H. LuuPavani Majety
 authored and GitHub committed 2026-10-27 16:59:59 -07:00
1 parent 36f5ff3172
commit 8697bbae4c
2 files changed
+41 -23
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/v1/attention/test_flashinfer_mla_dcp.py
+22 -16
View File
@@ -87,7 +87,8 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
@requires_flashinfer_mla
@pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"])
def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
@pytest.mark.parametrize("query_lens", [(3, 3), (3, 1)], ids=["uniform", "ragged"])
def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal, query_lens):
import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla
impl = MagicMock()
@@ -106,8 +107,11 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
flashinfer_mla.FlashInferMLAImpl._flattened_decode_metadata, impl
)
num_reqs, query_len = 2, 3
num_tokens = num_reqs * query_len
num_reqs = len(query_lens)
num_tokens = sum(query_lens)
query_len = max(query_lens)
split_rows = not causal or len(set(query_lens)) > 1
lens = torch.tensor(query_lens)
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)
@@ -120,15 +124,15 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
attn_metadata.decode.seq_lens = seq_lens
attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens
attn_metadata.decode.max_query_len = query_len
attn_metadata.decode.query_start_loc = torch.arange(
0, num_tokens + 1, query_len, dtype=torch.int32
attn_metadata.decode.query_start_loc = torch.tensor(
[0, *lens.cumsum(0).tolist()], dtype=torch.int32
)
attn_metadata.decode.query_len = 0
query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16)
kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16)
kernel_batch = num_reqs if causal else num_tokens
kernel_query_len = query_len if causal else 1
kernel_batch = num_tokens if split_rows else num_reqs
kernel_query_len = 1 if split_rows else query_len
kernel = MagicMock(
return_value=(
torch.ones(
@@ -155,15 +159,17 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
call = kernel.call_args.kwargs
assert call["query"].shape == (kernel_batch, kernel_query_len, 24, 576)
expected_block_table = (
block_table if causal else block_table.repeat_interleave(query_len, dim=0)
)
expected_seq_lens = seq_lens if causal else seq_lens.repeat_interleave(query_len)
expected_global_seq_lens = (
global_causal_seq_lens
if causal
else global_causal_seq_lens.repeat_interleave(query_len)
)
expected_block_table = block_table
expected_seq_lens = seq_lens
expected_global_seq_lens = global_causal_seq_lens
if split_rows:
expected_block_table = block_table.repeat_interleave(lens, dim=0)
expected_seq_lens = seq_lens.repeat_interleave(lens)
expected_global_seq_lens = global_causal_seq_lens.repeat_interleave(lens)
if causal and split_rows:
# Each split row sees KV only up to its own position in its request.
rows_after = torch.cat([torch.arange(n - 1, -1, -1) for n in query_lens])
expected_global_seq_lens = expected_global_seq_lens - rows_after.int()
torch.testing.assert_close(call["block_tables"], expected_block_table)
torch.testing.assert_close(call["seq_lens"], expected_seq_lens)
assert call["backend"] == "cute-dsl"
vllm/v1/attention/backends/mla/flashinfer_mla.py
+19 -7
View File
@@ -341,6 +341,8 @@ class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
cum_seq_lens_q: torch.Tensor | None = None
max_q_len: int | None = None
row_req: torch.Tensor | None = None
# Per split row: query tokens of its request after it.
causal_rows_after: torch.Tensor | None = None
if not attn_metadata.causal:
# FlashInfer decode has no causal flag. Flatten each non-causal
@@ -357,15 +359,23 @@ class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
# do not return on the ragged path (flashinfer #3238).
cum_seq_lens_q = attn_metadata.decode.query_start_loc
max_q_len = attn_metadata.decode.max_query_len
# trtllm API requires extra dimension q_len_per_request for MTP
elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
logger.warning_once(
"""FlashInferMLAImpl got a query of uneven length.
This usually indicates an issue in batch reordering
or incorrect setup in dummy_run."""
)
elif (
attn_metadata.num_decode_tokens
!= attn_metadata.num_decodes * attn_metadata.decode.max_query_len
):
# Ragged causal DCP decode, e.g. requests with fewer drafts: split
# into single-token rows, each bounded at its own causal position.
q = q.unsqueeze(1)
block_table, seq_lens, row_req = self._flattened_decode_metadata(
attn_metadata, q.shape[0]
)
query_start_loc = attn_metadata.decode.query_start_loc
assert query_start_loc is not None
query_end = query_start_loc[1:]
rows = torch.arange(q.shape[0], device=q.device, dtype=query_end.dtype)
causal_rows_after = query_end[row_req] - 1 - rows
else:
# trtllm API requires extra dimension q_len_per_request for MTP
q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
if self.bmm1_scale is None:
@@ -389,6 +399,8 @@ class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
assert causal_seqlens_kv_global is not None
if row_req is not None:
causal_seqlens_kv_global = causal_seqlens_kv_global[row_req]
if causal_rows_after is not None:
causal_seqlens_kv_global = causal_seqlens_kv_global - causal_rows_after
extra_kwargs.update(
enable_dcp=True,
cp_world=self.dcp_world_size,