@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3b736e1c38

Fork 0

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][DCP] Support DCP with query length > 1 (MTP) with FA3 (#25049)

Signed-off-by: Ming Yang <minos.future@gmail.com>

Ming Yang committed Oct 9, 2025 at 08:06 UTC 3b736e1c3869a593fed5b2eb58abe0186081d7cd

8 files changed +45 -13

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -38,7 +38,7 @@ else()

38 FetchContent_Declare(

39 vllm-flash-attn

40 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

41 - GIT_TAG 4695e6bed5366c41e28c06cd86170166e4f43d00

41 + GIT_TAG 8f468e7da54a8e2f98abfa7c38636aac91c0cba1

42 GIT_PROGRESS TRUE

43 # Don't share the vllm-flash-attn build between build types

44 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

vllm/v1/attention/backends/mla/common.py

+14 -2

@@ -370,6 +370,7 @@ class CudnnPrefillMetadata(MLACommonPrefillMetadata):

370 class MLACommonDecodeMetadata:

371 block_table: torch.Tensor

372 seq_lens: torch.Tensor

373 + dcp_tot_seq_lens: Optional[torch.Tensor]

374

375

376 D = TypeVar("D", bound=MLACommonDecodeMetadata)

@@ -682,10 +683,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

683 query_start_loc_cpu: torch.Tensor,

684 query_start_loc_device: torch.Tensor,

685 num_decode_tokens: int,

686 + dcp_tot_seq_lens_device: Optional[torch.Tensor],

687 ) -> MLACommonDecodeMetadata:

688 return MLACommonDecodeMetadata(

689 block_table=block_table_tensor,

690 seq_lens=seq_lens_device,

691 + dcp_tot_seq_lens=dcp_tot_seq_lens_device,

692 )

693

694 def build_for_cudagraph_capture(

@@ -727,6 +730,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

730 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

731 seq_lens = common_attn_metadata.seq_lens

732 seq_lens_cpu = common_attn_metadata.seq_lens_cpu

733 + dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens

734

735 query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

736

@@ -742,7 +746,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

746

747 # Note(hc): update seq_lens of decode reqs under DCP.

748 if self.dcp_world_size > 1:

745 - seq_lens[:num_decodes] = seq_lens[:num_decodes] // self.dcp_world_size + (

749 + assert dcp_local_seq_lens is not None

750 + dcp_local_seq_lens[:num_decodes] = seq_lens[

751 + :num_decodes

752 + ] // self.dcp_world_size + (

753 self.dcp_rank <= (seq_lens[:num_decodes] - 1) % self.dcp_world_size

754 )

755

@@ -899,10 +906,15 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

906 decode_metadata = self._build_decode(

907 block_table_tensor=block_table_tensor[:num_decodes, ...],

908 seq_lens_cpu=seq_lens_cpu[:num_decodes],

902 - seq_lens_device=seq_lens[:num_decodes],

909 + seq_lens_device=dcp_local_seq_lens[:num_decodes]

910 + if self.dcp_world_size > 1 and dcp_local_seq_lens is not None

911 + else seq_lens[:num_decodes],

912 query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],

913 query_start_loc_device=query_start_loc[: num_decodes + 1],

914 num_decode_tokens=num_decode_tokens,

915 + dcp_tot_seq_lens_device=seq_lens[:num_decodes]

916 + if self.dcp_world_size > 1

917 + else None,

918 )

919

920 attn_metadata = self.metadata_cls(

vllm/v1/attention/backends/mla/flashattn_mla.py

+7 -9

@@ -17,7 +17,6 @@ from vllm.attention.utils.fa_utils import (

17 get_flash_attn_version,

18 )

19 from vllm.config import VllmConfig

20 -from vllm.distributed.parallel_state import get_dcp_group

20 from vllm.logger import init_logger

21 from vllm.v1.attention.backends.mla.common import (

22 MLACommonBackend,

@@ -107,12 +106,6 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

106 # pre-allocated during capture.

107 self.max_num_splits = envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH

108

110 - # TODO(lucas): Until we add support for the DCP custom masking we need

111 - # to restrict decodes to q_len == 1 when DCP is enabled.

112 - self.reorder_batch_threshold = (

113 - 1 if get_dcp_group().world_size > 1 else self.reorder_batch_threshold

114 - )

115 -

109 def _schedule_decode(

110 self, num_reqs, cu_query_lens, max_query_len, seqlens, max_seq_len, causal

111 ):

@@ -121,7 +114,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

114 batch_size=num_reqs,

115 max_seqlen_q=max_query_len,

116 max_seqlen_k=max_seq_len,

124 - num_heads_q=self.num_heads,

117 + num_heads_q=self.num_heads * self.dcp_world_size,

118 num_heads_kv=1,

119 headdim=self.mla_dims.qk_rope_head_dim,

120 cache_seqlens=seqlens,

@@ -142,10 +135,11 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

135 query_start_loc_cpu: torch.Tensor,

136 query_start_loc_device: torch.Tensor,

137 num_decode_tokens: int,

138 + dcp_tot_seq_lens_device: Optional[torch.Tensor],

139 ) -> FlashAttnMLADecodeMetadata:

140 query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

141 max_query_len = query_lens_cpu.max().item()

148 - max_seq_len = seq_lens_cpu.max().item()

142 + max_seq_len = seq_lens_device.max().item()

143

144 scheduler_metadata = self._schedule_decode(

145 num_reqs=seq_lens_cpu.numel(),

@@ -188,6 +182,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

182 max_seq_len=max_seq_len,

183 scheduler_metadata=scheduler_metadata,

184 max_num_splits=max_num_splits,

185 + dcp_tot_seq_lens=dcp_tot_seq_lens_device,

186 )

187

188

@@ -289,6 +284,9 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

284 fa_version=3, # only version 3 is supported

285 scheduler_metadata=attn_metadata.decode.scheduler_metadata,

286 num_splits=attn_metadata.decode.max_num_splits,

287 + cp_world_size=self.dcp_world_size,

288 + cp_rank=self.dcp_rank,

289 + cp_tot_seqused_k=attn_metadata.decode.dcp_tot_seq_lens,

290 )

291

292 if self.need_to_return_lse_for_decode:

vllm/v1/attention/backends/mla/flashmla.py

@@ -106,6 +106,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

106 query_start_loc_cpu: torch.Tensor,

107 query_start_loc_device: torch.Tensor,

108 num_decode_tokens: int,

109 + dcp_tot_seq_lens_device: Optional[torch.Tensor],

110 ) -> FlashMLADecodeMetadata:

111 tile_scheduler_metadata, num_splits = get_mla_metadata(

112 seq_lens_device,

@@ -146,6 +147,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

147 seq_lens=seq_lens_device,

148 tile_scheduler_metadata=tile_scheduler_metadata,

149 num_splits=num_splits,

150 + dcp_tot_seq_lens=dcp_tot_seq_lens_device,

151 )

152

153

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -116,6 +116,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

116 query_start_loc_cpu: torch.Tensor,

117 query_start_loc_device: torch.Tensor,

118 num_decode_tokens: int,

119 + dcp_tot_seq_lens_device: Optional[torch.Tensor],

120 ) -> AiterMLADecodeMetadata:

121 page_size = self.kv_cache_spec.block_size

122 block_table_bounds = (seq_lens_device + page_size - 1) // page_size

@@ -174,6 +175,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

175 paged_kv_indices=paged_kv_indices,

176 paged_kv_last_page_len=paged_kv_last_page_len,

177 qo_indptr=qo_indptr,

178 + dcp_tot_seq_lens=dcp_tot_seq_lens_device,

179 )

180

181 return attn_metadata

vllm/v1/attention/backends/utils.py

@@ -93,6 +93,9 @@ class CommonAttentionMetadata:

93 # Needed by CrossAttentionBuilder

94 encoder_seq_lens: Optional[np.ndarray] = None

96 + dcp_local_seq_lens: Optional[torch.Tensor] = None

97 + """Sequence lengths of the local rank in decode context parallelism world"""

98 +

100 def slice_query_start_locs(

101 query_start_loc: torch.Tensor,

vllm/v1/spec_decode/eagle.py

@@ -597,6 +597,7 @@ class EagleProposer:

597 block_table_tensor=common_attn_metadata.block_table_tensor,

598 slot_mapping=common_attn_metadata.slot_mapping[token_indices],

599 causal=True,

600 + dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,

601 )

602

603 token_indices_to_sample = (

@@ -868,6 +869,7 @@ class EagleProposer:

869 block_table_tensor=common_attn_metadata.block_table_tensor,

870 slot_mapping=common_attn_metadata.slot_mapping[token_indices],

871 causal=True,

872 + dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,

873 )

874

875 return spec_common_attn_metadata, token_indices

vllm/v1/worker/gpu_model_runner.py

+14 -1

@@ -398,6 +398,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

398 self.max_num_reqs + 1, dtype=torch.int32

399 )

400 self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

401 + if self.dcp_world_size > 1:

402 + self.dcp_local_seq_lens = self._make_buffer(

403 + self.max_num_reqs, dtype=torch.int32

404 + )

405 # Because inputs_embeds may be bfloat16 and we don't need a numpy

406 # version of this tensor, avoid a RuntimeError by not creating a

407 # numpy buffer.

@@ -581,7 +585,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

585 # NOTE(lucas): currently no backend supports the custom masking

586 # required for DCP with q_len > 1, so we assert here. Remove this

587 # assert once the custom mask is support is added to FA3.

584 - if self.dcp_world_size > 1:

588 + if (

589 + self.dcp_world_size > 1

590 + and envs.VLLM_ATTENTION_BACKEND != "FLASH_ATTN_MLA"

591 + ):

592 assert self.reorder_batch_threshold == 1, (

593 "DCP not support reorder_batch_threshold > 1 now."

594 )

@@ -1335,6 +1342,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1342 num_logits_indices=logits_indices.size(0),

1343 causal=True,

1344 encoder_seq_lens=encoder_seq_lens,

1345 + dcp_local_seq_lens=self.dcp_local_seq_lens.gpu[:num_reqs]

1346 + if self.dcp_world_size > 1

1347 + else None,

1348 )

1349

1350 if self.speculative_config and spec_decode_common_attn_metadata is None:

@@ -3310,6 +3320,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3320 kv_cache_group_id

3321 ].slot_mapping.gpu[:num_tokens],

3322 causal=True,

3323 + dcp_local_seq_lens=self.dcp_local_seq_lens.gpu[:num_reqs]

3324 + if self.dcp_world_size > 1

3325 + else None,

3326 )

3327 for attn_group in self.attn_groups[kv_cache_group_id]:

3328 if ubatch_slices is not None: