[Attention][DCP] Enable FlashInfer MLA for DSpark drafting (#54277)
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Summer Yang committed
Aug 28, 2026 at 21:11 UTC
6d4562c59b97b4e35d459ff9389e71b6fe4995de
8 files changed
+98
-28
tests/v1/attention/test_flashinfer_mla_dcp.py
+33
-8
@@ -84,7 +84,8 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
84
85
86
@requires_flashinfer_mla
87
-def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
87
+@pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"])
88
+def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
89
import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla
90
91
impl = MagicMock()
@@ -106,7 +107,7 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
107
seq_lens = torch.tensor([4, 5], dtype=torch.int32)
108
global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32)
109
attn_metadata = MagicMock()
109
- attn_metadata.causal = True
110
+ attn_metadata.causal = causal
111
attn_metadata.num_decode_tokens = num_tokens
112
attn_metadata.num_decodes = num_reqs
113
attn_metadata.max_seq_len = 9
@@ -116,10 +117,23 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
117
118
query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16)
119
kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16)
120
+ kernel_batch = num_reqs if causal else num_tokens
121
+ kernel_query_len = query_len if causal else 1
122
kernel = MagicMock(
123
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),
124
+ torch.ones(
125
+ kernel_batch,
126
+ kernel_query_len,
127
+ 24,
128
+ 512,
129
+ dtype=torch.bfloat16,
130
+ ),
131
+ torch.ones(
132
+ kernel_batch,
133
+ kernel_query_len,
134
+ 24,
135
+ dtype=torch.float32,
136
+ ),
137
)
138
)
139
monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock())
@@ -130,14 +144,25 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch):
144
)
145
146
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
147
+ assert call["query"].shape == (kernel_batch, kernel_query_len, 24, 576)
148
+ expected_block_table = (
149
+ block_table if causal else block_table.repeat_interleave(query_len, dim=0)
150
+ )
151
+ expected_seq_lens = seq_lens if causal else seq_lens.repeat_interleave(query_len)
152
+ expected_global_seq_lens = (
153
+ global_causal_seq_lens
154
+ if causal
155
+ else global_causal_seq_lens.repeat_interleave(query_len)
156
+ )
157
+ torch.testing.assert_close(call["block_tables"], expected_block_table)
158
+ torch.testing.assert_close(call["seq_lens"], expected_seq_lens)
159
assert call["backend"] == "cute-dsl"
160
assert call["enable_dcp"] is True
161
assert call["cp_world"] == 8
162
assert call["cp_rank"] == 3
140
- assert call["causal_seqlens_kv_global"] is global_causal_seq_lens
163
+ torch.testing.assert_close(
164
+ call["causal_seqlens_kv_global"], expected_global_seq_lens
165
+ )
166
assert "multi_ctas_kv_counter_buffer" not in call
167
assert flashinfer_mla.FlashInferMLAImpl.lse_base_on_e
168
assert output.shape == (num_tokens, 24, 512)
tests/v1/attention/test_mla_backends.py
+2
-2
@@ -1021,7 +1021,7 @@ def test_tokenspeed_mla_noncausal_capability():
1021
assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()
1022
1023
1024
-def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch):
1024
+def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch):
1025
flashinfer_mla_module = pytest.importorskip(
1026
"vllm.v1.attention.backends.mla.flashinfer_mla"
1027
)
@@ -1053,7 +1053,7 @@ def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch):
1053
assert reason is None
1054
assert backend.supports_non_causal()
1055
assert builder.supports_non_causal_multi_token_decode
1056
- assert not backend.supports_non_causal_dcp()
1056
+ assert backend.supports_non_causal_dcp()
1057
1058
1059
@pytest.mark.parametrize(
tests/v1/attention/test_mla_noncausal.py
+3
-2
@@ -109,7 +109,7 @@ def test_noncausal_decode_metadata_keeps_live_request_buffers():
109
)
110
111
112
-def test_mla_cache_marker_is_promoted_to_group_capability():
112
+def test_mla_cache_marker_is_a_group_property():
113
kwargs = {
114
"block_size": 64,
115
"num_kv_heads": 1,
@@ -120,7 +120,8 @@ def test_mla_cache_marker_is_promoted_to_group_capability():
120
unmarked = MLAAttentionSpec(**kwargs)
121
122
assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode
123
- assert MLAAttentionSpec.merge([marked, unmarked]).non_causal_multi_token_decode
123
assert not MLAAttentionSpec.merge(
124
[unmarked, unmarked]
125
).non_causal_multi_token_decode
126
+ with pytest.raises(AssertionError, match="non_causal_multi_token_decode"):
127
+ MLAAttentionSpec.merge([marked, unmarked])
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+33
@@ -37,10 +37,14 @@ def _make_fake_speculator(
37
fake_input_buffers = SimpleNamespace(
38
query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32),
39
seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
40
+ dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
41
)
42
fake_block_tables = SimpleNamespace(
43
input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)],
44
slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64),
45
+ cp_size=1,
46
+ cp_rank=0,
47
+ cp_interleave=1,
48
)
49
return SimpleNamespace(
50
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():
130
bound = captured["seq_lens_cpu_upper_bound"]
131
# 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
132
assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
133
+
134
+
135
+def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
136
+ fake = _make_fake_speculator()
137
+ fake.block_tables.cp_size = 2
138
+ fake.block_tables.cp_rank = 1
139
+ fake.block_tables.cp_interleave = 4
140
+ fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16])
141
+
142
+ def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave):
143
+ assert seq_lens is fake.input_buffers.seq_lens
144
+ assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4)
145
+ out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32))
146
+ out[num_reqs:].zero_()
147
+
148
+ with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
149
+ captured = _run_build(
150
+ fake,
151
+ num_reqs=3,
152
+ num_reqs_padded=4,
153
+ num_tokens_padded=4,
154
+ base=torch.tensor([5, 9, 16]),
155
+ step=0,
156
+ )
157
+
158
+ local = captured["dcp_local_seq_lens"]
159
+ assert isinstance(local, torch.Tensor)
160
+ assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr()
161
+ assert local.tolist() == [1, 4, 8, 0]
vllm/v1/attention/backends/mla/flashinfer_mla.py
+7
-2
@@ -120,6 +120,7 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
120
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
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
+ supports_non_causal_multi_token_dcp: ClassVar[bool] = True
124
125
def __init__(
126
self,
@@ -288,8 +289,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
289
seq_lens = attn_metadata.decode.seq_lens
290
291
if not attn_metadata.causal:
291
- # FlashInfer decode has no causal flag. For TP-only DSpark, flatten
292
- # each non-causal query block into independent single-token rows.
292
+ # FlashInfer decode has no causal flag. Flatten each non-causal
293
+ # query block into independent single-token rows.
294
query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
295
q = q.unsqueeze(1)
296
if query_len > 1:
@@ -326,6 +327,10 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
327
assert self.cp_kv_cache_interleave_size == 1
328
causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens
329
assert causal_seqlens_kv_global is not None
330
+ if not attn_metadata.causal and query_len > 1:
331
+ causal_seqlens_kv_global = causal_seqlens_kv_global.repeat_interleave(
332
+ query_len
333
+ )
334
extra_kwargs.update(
335
enable_dcp=True,
336
cp_world=self.dcp_world_size,
vllm/v1/kv_cache_interface.py
+6
-3
@@ -570,6 +570,11 @@ class MLAAttentionSpec(FullAttentionSpec):
570
"All attention layers in the same KV cache group must use the same "
571
"quantization method, tokens per state, and model version."
572
)
573
+ non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
574
+ assert len(non_causal_mtd_set) == 1, (
575
+ "All attention layers in the same KV cache group must agree on "
576
+ "non_causal_multi_token_decode."
577
+ )
578
merged_spec = cls(
579
block_size=specs[0].block_size,
580
num_kv_heads=specs[0].num_kv_heads,
@@ -582,9 +587,7 @@ class MLAAttentionSpec(FullAttentionSpec):
587
cache_dtype_str=cache_dtype_str_set.pop(),
588
tokens_per_state=tokens_per_state_set.pop(),
589
model_version=model_version_set.pop(),
585
- non_causal_multi_token_decode=any(
586
- spec.non_causal_multi_token_decode for spec in specs
587
- ),
590
+ non_causal_multi_token_decode=non_causal_mtd_set.pop(),
591
)
592
for spec in specs:
593
for f in fields(AttentionSpec):
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+1
-11
@@ -18,7 +18,7 @@ from vllm.v1.attention.backends.utils import PAD_SLOT_ID
18
from vllm.v1.kv_cache_interface import KVCacheConfig
19
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
20
from vllm.v1.worker.gpu.block_table import BlockTables
21
-from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens
21
+from vllm.v1.worker.gpu.cp_utils import cp_local_slot
22
from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
23
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
24
from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -291,16 +291,6 @@ class DFlashSpeculator(DraftModelSpeculator):
291
if not self.draft_attn_layer_names:
292
return None
293
assert num_query_per_req is None # Omitted for DFlash, read from self instead
294
- if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
295
- prepare_dcp_local_seq_lens(
296
- self.input_buffers.dcp_local_seq_lens,
297
- self.input_buffers.seq_lens,
298
- num_reqs,
299
- self.block_tables.cp_size,
300
- self.block_tables.cp_rank,
301
- self.block_tables.cp_interleave,
302
- )
303
- dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
294
return super()._build_draft_attn_metadata(
295
num_reqs,
296
num_reqs_padded,
vllm/v1/worker/gpu/spec_decode/speculator.py
+13
@@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.attn_utils import (
21
init_attn_backend,
22
)
23
from vllm.v1.worker.gpu.block_table import BlockTables
24
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
25
from vllm.v1.worker.gpu.dp_utils import DPSyncState
26
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
27
from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -292,6 +293,18 @@ class DraftModelSpeculator(BaseSpeculator):
293
out=draft_seq_lens_cpu_upper_bound[:num_reqs],
294
)
295
draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
296
+ if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
297
+ # Draft steps advance and rewind their own global sequence lengths,
298
+ # so the target model's DCP-local lengths may already be stale.
299
+ prepare_dcp_local_seq_lens(
300
+ self.input_buffers.dcp_local_seq_lens,
301
+ self.input_buffers.seq_lens,
302
+ num_reqs,
303
+ self.block_tables.cp_size,
304
+ self.block_tables.cp_rank,
305
+ self.block_tables.cp_interleave,
306
+ )
307
+ dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
308
attn_metadata = build_attn_metadata(
309
attn_groups=self.attn_groups,
310
num_reqs=num_reqs_padded,