A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Attention] add DCP support for FLASH_ATTN_MLA backend (#24453)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Lucas Wilkinson committed
Sep 10, 2025 at 05:19 UTC
0ae43dbf8cb28a299ae724fc742b0c5bcddea868
2 files changed
+19
-2
vllm/v1/attention/backends/mla/flashattn_mla.py
+16
-2
@@ -11,6 +11,7 @@ from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,
11
from vllm.attention.utils.fa_utils import (flash_attn_supports_mla,
12
get_flash_attn_version)
13
from vllm.config import VllmConfig
14
+from vllm.distributed.parallel_state import get_dcp_group
15
from vllm.logger import init_logger
16
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
17
MLACommonDecodeMetadata,
@@ -98,6 +99,11 @@ class FlashAttnMLAMetadataBuilder(
99
# pre-allocated during capture.
100
self.max_num_splits = _DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
101
102
+ # TODO(lucas): Until we add support for the DCP custom masking we need
103
+ # to restrict decodes to q_len == 1 when DCP is enabled.
104
+ self.__class__.reorder_batch_threshold = 1 \
105
+ if get_dcp_group().world_size > 1 else self.reorder_batch_threshold
106
+
107
def _schedule_decode(self, num_reqs, cu_query_lens, max_query_len, seqlens,
108
max_seq_len, causal):
109
if self.fa_aot_schedule:
@@ -172,6 +178,7 @@ class FlashAttnMLAMetadataBuilder(
178
179
180
class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
181
+ can_return_lse_for_decode: bool = True
182
183
def __init__(
184
self,
@@ -239,7 +246,7 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
246
# to prevent invalid grid configuration during graph capture.
247
max_seqlen_q = max(attn_metadata.decode.max_query_len, 1)
248
242
- o = flash_attn_varlen_func(
249
+ attn_out = flash_attn_varlen_func(
250
q=q_pe,
251
k=k_pe_cache.unsqueeze(-2), # Add head dim of 1
252
v=kv_c_cache.unsqueeze(-2), # Add head dim of 1
@@ -251,9 +258,16 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
258
block_table=attn_metadata.decode.block_table,
259
softmax_scale=self.scale,
260
causal=True,
261
+ return_softmax_lse=self.need_to_return_lse_for_decode,
262
fa_version=3, # only version 3 is supported
263
scheduler_metadata=attn_metadata.decode.scheduler_metadata,
264
num_splits=attn_metadata.decode.max_num_splits,
265
)
266
259
- return self._v_up_proj(o)
267
+ if self.need_to_return_lse_for_decode:
268
+ o, lse = attn_out
269
+ # FA returns LSE in shape [ H, B ] but DCP wants [ B, H ]
270
+ return o, lse.transpose(0, 1) # [ H, B ] -> [ B, H ]
271
+ else:
272
+ o = attn_out
273
+ return o, None
vllm/v1/worker/gpu_model_runner.py
+3
@@ -440,6 +440,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
440
return
441
442
if self.reorder_batch_threshold is not None:
443
+ # NOTE(lucas): currently no backend supports the custom masking
444
+ # required for DCP with q_len > 1, so we assert here. Remove this
445
+ # assert once the custom mask is support is added to FA3.
446
if self.dcp_world_size > 1:
447
assert self.reorder_batch_threshold == 1, \
448
"DCP not support reorder_batch_threshold > 1 now."