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

[Attention][DCP] Enable FlashInfer MLA for DSpark drafting (#54277)

Browse Source 
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
This commit is contained in:
Summer Yang
2026-08-28 21:11:28 -07:00
committed by GitHub
parent b2f685834a
commit 6d4562c59b
8 changed files with 98 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_flashinfer_mla_dcp.py
+33 -8
View File
@@ -84,7 +84,8 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
@requires_flashinfer_mla
def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
@pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"])
def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla
impl = MagicMock()
@@ -106,7 +107,7 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
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.causal = causal
attn_metadata.num_decode_tokens = num_tokens
attn_metadata.num_decodes = num_reqs
attn_metadata.max_seq_len = 9
@@ -116,10 +117,23 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
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 = MagicMock(
return_value=(
torch.ones(num_reqs, query_len, 24, 512, dtype=torch.bfloat16),
torch.ones(num_reqs, query_len, 24, dtype=torch.float32),
torch.ones(
kernel_batch,
kernel_query_len,
24,
512,
dtype=torch.bfloat16,
),
torch.ones(
kernel_batch,
kernel_query_len,
24,
dtype=torch.float32,
),
)
)
monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock())
@@ -130,14 +144,25 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
)
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["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)
)
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"
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
torch.testing.assert_close(
call["causal_seqlens_kv_global"], expected_global_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)
tests/v1/attention/test_mla_backends.py
+2 -2
View File
@@ -1021,7 +1021,7 @@ def test_tokenspeed_mla_noncausal_capability():
assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()
def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch):
def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch):
flashinfer_mla_module = pytest.importorskip(
"vllm.v1.attention.backends.mla.flashinfer_mla"
)
@@ -1053,7 +1053,7 @@ def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch):
assert reason is None
assert backend.supports_non_causal()
assert builder.supports_non_causal_multi_token_decode
assert not backend.supports_non_causal_dcp()
assert backend.supports_non_causal_dcp()
@pytest.mark.parametrize(
tests/v1/attention/test_mla_noncausal.py
+3 -2
View File
@@ -109,7 +109,7 @@ def test_noncausal_decode_metadata_keeps_live_request_buffers():
)
def test_mla_cache_marker_is_promoted_to_group_capability():
def test_mla_cache_marker_is_a_group_property():
kwargs = {
"block_size": 64,
"num_kv_heads": 1,
@@ -120,7 +120,8 @@ def test_mla_cache_marker_is_promoted_to_group_capability():
unmarked = MLAAttentionSpec(**kwargs)
assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode
assert MLAAttentionSpec.merge([marked, unmarked]).non_causal_multi_token_decode
assert not MLAAttentionSpec.merge(
[unmarked, unmarked]
).non_causal_multi_token_decode
with pytest.raises(AssertionError, match="non_causal_multi_token_decode"):
MLAAttentionSpec.merge([marked, unmarked])
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+33
View File
@@ -37,10 +37,14 @@ def _make_fake_speculator(
fake_input_buffers = SimpleNamespace(
query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32),
seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
)
fake_block_tables = SimpleNamespace(
input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)],
slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64),
cp_size=1,
cp_rank=0,
cp_interleave=1,
)
return SimpleNamespace(
arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),
@@ -126,3 +130,32 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len():
bound = captured["seq_lens_cpu_upper_bound"]
# 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
fake = _make_fake_speculator()
fake.block_tables.cp_size = 2
fake.block_tables.cp_rank = 1
fake.block_tables.cp_interleave = 4
fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16])
def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave):
assert seq_lens is fake.input_buffers.seq_lens
assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4)
out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32))
out[num_reqs:].zero_()
with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
captured = _run_build(
fake,
num_reqs=3,
num_reqs_padded=4,
num_tokens_padded=4,
base=torch.tensor([5, 9, 16]),
step=0,
)
local = captured["dcp_local_seq_lens"]
assert isinstance(local, torch.Tensor)
assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr()
assert local.tolist() == [1, 4, 8, 0]
vllm/v1/attention/backends/mla/flashinfer_mla.py
+7 -2
View File
@@ -120,6 +120,7 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
# Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
supports_non_causal_multi_token_decode: ClassVar[bool] = True
supports_non_causal_multi_token_dcp: ClassVar[bool] = True
def __init__(
self,
@@ -288,8 +289,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
seq_lens = attn_metadata.decode.seq_lens
if not attn_metadata.causal:
# FlashInfer decode has no causal flag. For TP-only DSpark, flatten
# each non-causal query block into independent single-token rows.
# FlashInfer decode has no causal flag. Flatten each non-causal
# query block into independent single-token rows.
query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
q = q.unsqueeze(1)
if query_len > 1:
@@ -326,6 +327,10 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
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
if not attn_metadata.causal and query_len > 1:
causal_seqlens_kv_global = causal_seqlens_kv_global.repeat_interleave(
query_len
)
extra_kwargs.update(
enable_dcp=True,
cp_world=self.dcp_world_size,
vllm/v1/kv_cache_interface.py
+6 -3
View File
@@ -570,6 +570,11 @@ class MLAAttentionSpec(FullAttentionSpec):
"All attention layers in the same KV cache group must use the same "
"quantization method, tokens per state, and model version."
)
non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
assert len(non_causal_mtd_set) == 1, (
"All attention layers in the same KV cache group must agree on "
"non_causal_multi_token_decode."
)
merged_spec = cls(
block_size=specs[0].block_size,
num_kv_heads=specs[0].num_kv_heads,
@@ -582,9 +587,7 @@ class MLAAttentionSpec(FullAttentionSpec):
cache_dtype_str=cache_dtype_str_set.pop(),
tokens_per_state=tokens_per_state_set.pop(),
model_version=model_version_set.pop(),
non_causal_multi_token_decode=any(
spec.non_causal_multi_token_decode for spec in specs
),
non_causal_multi_token_decode=non_causal_mtd_set.pop(),
)
for spec in specs:
for f in fields(AttentionSpec):
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+1 -11
View File
@@ -18,7 +18,7 @@ from vllm.v1.attention.backends.utils import PAD_SLOT_ID
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cp_utils import cp_local_slot
from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -291,16 +291,6 @@ class DFlashSpeculator(DraftModelSpeculator):
if not self.draft_attn_layer_names:
return None
assert num_query_per_req is None # Omitted for DFlash, read from self instead
if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
prepare_dcp_local_seq_lens(
self.input_buffers.dcp_local_seq_lens,
self.input_buffers.seq_lens,
num_reqs,
self.block_tables.cp_size,
self.block_tables.cp_rank,
self.block_tables.cp_interleave,
)
dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
return super()._build_draft_attn_metadata(
num_reqs,
num_reqs_padded,
vllm/v1/worker/gpu/spec_decode/speculator.py
+13
View File
@@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.attn_utils import (
init_attn_backend,
)
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.dp_utils import DPSyncState
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -292,6 +293,18 @@ class DraftModelSpeculator(BaseSpeculator):
out=draft_seq_lens_cpu_upper_bound[:num_reqs],
)
draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
# Draft steps advance and rewind their own global sequence lengths,
# so the target model's DCP-local lengths may already be stale.
prepare_dcp_local_seq_lens(
self.input_buffers.dcp_local_seq_lens,
self.input_buffers.seq_lens,
num_reqs,
self.block_tables.cp_size,
self.block_tables.cp_rank,
self.block_tables.cp_interleave,
)
dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
attn_metadata = build_attn_metadata(
attn_groups=self.attn_groups,
num_reqs=num_reqs_padded,