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 5, 18:30 UTC
[Bugfix][DCP] Fix CUDA graph capture for Decode Context Parallelism (#36070)
Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
sungsoo ha committed
Mar 30, 2026 at 17:20 UTC
4ac227222fc2e1ec093425ce2334a1b11b6d282f
1 file changed
+34
-8
vllm/v1/attention/backends/flash_attn.py
+34
-8
@@ -22,9 +22,11 @@ from vllm.v1.attention.backends.fa_utils import (
22
get_flash_attn_version,
23
is_flash_attn_varlen_func_available,
24
)
25
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
26
from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
27
from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
28
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
29
+from vllm.v1.worker.workspace import current_workspace_manager
30
31
if is_flash_attn_varlen_func_available():
32
from vllm.v1.attention.backends.fa_utils import (
@@ -52,7 +54,6 @@ from vllm.v1.attention.backend import (
54
CommonAttentionMetadata,
55
)
56
from vllm.v1.attention.backends.utils import (
55
- get_dcp_local_seq_lens,
57
get_kv_cache_layout,
58
)
59
from vllm.v1.kv_cache_interface import AttentionSpec
@@ -356,6 +357,14 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
357
self.attention_config.flash_attn_max_num_splits_for_cuda_graph
358
)
359
360
+ if self.dcp_world_size > 1:
361
+ max_num_reqs = vllm_config.scheduler_config.max_num_seqs
362
+ self._dcp_context_kv_lens = torch.zeros(
363
+ max_num_reqs,
364
+ dtype=torch.int32,
365
+ device=self.device,
366
+ )
367
+
368
# Sliding window size to be used with the AOT scheduler will be
369
# populated on first build() call.
370
self.aot_sliding_window: tuple[int, int] | None = None
@@ -452,15 +461,18 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
461
prefix_scheduler_metadata = None
462
463
if self.dcp_world_size > 1:
455
- query_kv_lens = query_start_loc[1:] - query_start_loc[:-1]
456
- dcp_context_kv_lens = seq_lens - query_kv_lens
457
-
458
- dcp_context_kv_lens = get_dcp_local_seq_lens(
459
- dcp_context_kv_lens,
464
+ query_lens = query_start_loc[1:] - query_start_loc[:-1]
465
+ context_kv_lens = seq_lens - query_lens
466
+ local_context_kv_lens = get_dcp_local_seq_lens(
467
+ context_kv_lens,
468
self.dcp_world_size,
469
self.dcp_rank,
470
self.cp_kv_cache_interleave_size,
471
)
472
+ self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens
473
+ self._dcp_context_kv_lens[num_reqs:] = 0
474
+ dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs]
475
+
476
# After DCP distribution, the maximum number of tokens for any rank is
477
# ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
478
# and I is cp_kv_cache_interleave_size.
@@ -637,6 +649,10 @@ class FlashAttentionImpl(AttentionImpl):
649
)
650
self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs
651
652
+ self._dcp_dtype: torch.dtype | None = None
653
+ if vllm_config is not None and self.dcp_world_size > 1:
654
+ self._dcp_dtype = vllm_config.model_config.dtype
655
+
656
def forward(
657
self,
658
layer: torch.nn.Module,
@@ -862,11 +878,18 @@ class FlashAttentionImpl(AttentionImpl):
878
sliding_window_size = (
879
list(self.sliding_window) if self.sliding_window is not None else None
880
)
881
+ n = query_across_dcp.shape[0]
882
+ (dcp_context_out,) = current_workspace_manager().get_simultaneous(
883
+ (
884
+ (n, self.num_heads * self.dcp_world_size, self.head_size),
885
+ self._dcp_dtype,
886
+ ),
887
+ )
888
context_attn_out, context_lse = flash_attn_varlen_func(
889
q=query_across_dcp,
890
k=key_cache,
891
v=value_cache,
869
- out=None,
892
+ out=dcp_context_out,
893
cu_seqlens_q=cu_seqlens_q,
894
max_seqlen_q=max_seqlen_q,
895
seqused_k=attn_metadata.dcp_context_kv_lens,
@@ -894,11 +917,14 @@ class FlashAttentionImpl(AttentionImpl):
917
)
918
context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()
919
920
+ (dcp_query_out,) = current_workspace_manager().get_simultaneous(
921
+ ((query.shape[0], self.num_heads, self.head_size), self._dcp_dtype),
922
+ )
923
query_attn_out, query_lse = flash_attn_varlen_func(
924
q=query,
925
k=key,
926
v=value,
901
- out=None,
927
+ out=dcp_query_out,
928
cu_seqlens_q=cu_seqlens_q,
929
max_seqlen_q=max_seqlen_q,
930
cu_seqlens_k=cu_seqlens_q,