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 12, 23:30 UTC
[BugFix] Fix async-scheduling + FlashAttn MLA (#28990)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Nov 19, 2025 at 10:04 UTC
48fc8b1e595766af9c91edfc1de43f3a352575eb
4 files changed
+18
-10
vllm/v1/attention/backends/mla/common.py
+9
-6
@@ -755,6 +755,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
755
seq_lens = common_attn_metadata.seq_lens
756
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
757
dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
758
+ dcp_local_seq_lens_cpu = common_attn_metadata.dcp_local_seq_lens_cpu
759
760
query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
761
@@ -944,18 +945,20 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
945
946
decode_metadata = None
947
if num_decodes > 0:
948
+ dcp_tot_seq_lens_device = None
949
+ if self.dcp_world_size > 1:
950
+ dcp_tot_seq_lens_device = seq_lens[:num_decodes]
951
+ seq_lens_cpu = dcp_local_seq_lens_cpu
952
+ seq_lens = dcp_local_seq_lens
953
+
954
decode_metadata = self._build_decode(
955
block_table_tensor=block_table_tensor[:num_decodes, ...],
956
seq_lens_cpu=seq_lens_cpu[:num_decodes],
950
- seq_lens_device=dcp_local_seq_lens[:num_decodes]
951
- if self.dcp_world_size > 1 and dcp_local_seq_lens is not None
952
- else seq_lens[:num_decodes],
957
+ seq_lens_device=seq_lens[:num_decodes],
958
query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],
959
query_start_loc_device=query_start_loc[: num_decodes + 1],
960
num_decode_tokens=num_decode_tokens,
956
- dcp_tot_seq_lens_device=seq_lens[:num_decodes]
957
- if self.dcp_world_size > 1
958
- else None,
961
+ dcp_tot_seq_lens_device=dcp_tot_seq_lens_device,
962
)
963
964
attn_metadata = self.metadata_cls(
vllm/v1/attention/backends/mla/flashattn_mla.py
+1
-1
@@ -173,7 +173,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
173
) -> FlashAttnMLADecodeMetadata:
174
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
175
max_query_len = query_lens_cpu.max().item()
176
- max_seq_len = seq_lens_device.max().item()
176
+ max_seq_len = seq_lens_cpu.max().item()
177
178
# For Flash Attention MLA + full cudagraph
179
max_num_splits = 0
vllm/v1/attention/backends/utils.py
+1
@@ -92,6 +92,7 @@ class CommonAttentionMetadata:
92
encoder_seq_lens: np.ndarray | None = None
93
94
dcp_local_seq_lens: torch.Tensor | None = None
95
+ dcp_local_seq_lens_cpu: torch.Tensor | None = None
96
"""Sequence lengths of the local rank in decode context parallelism world"""
97
98
vllm/v1/worker/gpu_model_runner.py
+7
-3
@@ -1451,9 +1451,12 @@ class GPUModelRunner(
1451
num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[
1452
:num_reqs
1453
]
1454
- dcp_local_seq_lens = (
1455
- self.dcp_local_seq_lens.gpu[:num_reqs] if self.dcp_world_size > 1 else None
1456
- )
1454
+
1455
+ dcp_local_seq_lens, dcp_local_seq_lens_cpu = None, None
1456
+ if self.dcp_world_size > 1:
1457
+ dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs]
1458
+ dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[:num_reqs]
1459
+
1460
spec_decode_common_attn_metadata = None
1461
1462
if for_cudagraph_capture:
@@ -1521,6 +1524,7 @@ class GPUModelRunner(
1524
causal=True,
1525
encoder_seq_lens=encoder_seq_lens,
1526
dcp_local_seq_lens=dcp_local_seq_lens,
1527
+ dcp_local_seq_lens_cpu=dcp_local_seq_lens_cpu,
1528
)
1529
1530
if self.speculative_config and spec_decode_common_attn_metadata is None: