@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 67475a6e81
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 11, 18:45 UTC

[DCP][Bugfix][CI] Fix accuracy issue of DCP when using FLASH_ATTN_MLA (#30309)

Signed-off-by: FENP <yuanyongjie.yyj@antgroup.com>

Jaya Yuan committed Dec 9, 2025 at 16:22 UTC 67475a6e81abea915857f82e6f10d80b03b842c9
2 files changed +6 -2
tests/distributed/test_context_parallel.py
+4 -1
@@ -123,8 +123,11 @@ class CPTestSettings:
123
124 CP_TEXT_GENERATION_MODELS = {
125 "deepseek-ai/DeepSeek-V2-Lite-Chat": [
126 + CPTestSettings.detailed(dcp_multipliers=[1]),
127 CPTestSettings.detailed(
127 - dcp_multipliers=[0.5, 1], cp_kv_cache_interleave_size=64
128 + dcp_multipliers=[0.5],
129 + cp_kv_cache_interleave_size=64,
130 + attn_backend="FLASHMLA",
131 ),
132 ],
133 "Qwen/Qwen2.5-1.5B-Instruct": [
vllm/v1/attention/backends/mla/flashattn_mla.py
+2 -1
@@ -105,13 +105,14 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
105 vllm_config: VllmConfig,
106 device: torch.device,
107 ):
108 + interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
109 super().__init__(
110 kv_cache_spec,
111 layer_names,
112 vllm_config,
113 device,
114 FlashAttnMLAMetadata,
114 - supports_dcp_with_varlen=True,
115 + supports_dcp_with_varlen=(interleave_size == 1),
116 )
117 self.max_num_splits = 0 # No upper bound on the number of splits.
118 self.fa_aot_schedule = get_flash_attn_version() == 3