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 13, 19:30 UTC
[BugFix] Fix DCP Assert (AssertionError: DCP not support reorder_batch_threshold > 1 now.) (#28100)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Nov 5, 2025 at 01:54 UTC
d43ad5a75790e4d97394940187bbf37402c4fa97
3 files changed
+18
-3
vllm/v1/attention/backends/mla/common.py
+2
-1
@@ -545,6 +545,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
545
vllm_config: VllmConfig,
546
device: torch.device,
547
metadata_cls: type[M] | None = None,
548
+ supports_dcp_with_varlen: bool = False,
549
):
550
self.metadata_cls = (
551
metadata_cls if metadata_cls is not None else MLACommonMetadata
@@ -638,7 +639,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
639
640
supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
641
self._init_reorder_batch_threshold(
641
- self.reorder_batch_threshold, supports_spec_decode
642
+ self.reorder_batch_threshold, supports_spec_decode, supports_dcp_with_varlen
643
)
644
645
# Validate consistency between query_len_support and reorder_batch_threshold
vllm/v1/attention/backends/mla/flashattn_mla.py
+6
-1
@@ -81,7 +81,12 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
81
device: torch.device,
82
):
83
super().__init__(
84
- kv_cache_spec, layer_names, vllm_config, device, FlashAttnMLAMetadata
84
+ kv_cache_spec,
85
+ layer_names,
86
+ vllm_config,
87
+ device,
88
+ FlashAttnMLAMetadata,
89
+ supports_dcp_with_varlen=True,
90
)
91
self.max_num_splits = 0 # No upper bound on the number of splits.
92
self.fa_aot_schedule = get_flash_attn_version() == 3
vllm/v1/attention/backends/utils.py
+10
-1
@@ -264,7 +264,10 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
264
self.device = device
265
266
def _init_reorder_batch_threshold(
267
- self, reorder_batch_threshold: int = 1, supports_spec_as_decode: bool = False
267
+ self,
268
+ reorder_batch_threshold: int = 1,
269
+ supports_spec_as_decode: bool = False,
270
+ supports_dcp_with_varlen: bool = False,
271
) -> None:
272
self.reorder_batch_threshold = reorder_batch_threshold
273
if self.reorder_batch_threshold is not None and supports_spec_as_decode:
@@ -281,6 +284,12 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
284
1 + speculative_config.num_speculative_tokens,
285
)
286
287
+ if (
288
+ self.vllm_config.parallel_config.decode_context_parallel_size > 1
289
+ and not supports_dcp_with_varlen
290
+ ):
291
+ self.reorder_batch_threshold = 1
292
+
293
@abstractmethod
294
def build(
295
self,