[Perf][Attention] Avoid CPU-GPU sync in DCP sequence lengths (#58169)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Jiangyun Zhu committed
Sep 23, 2026 at 16:28 UTC
6632ed559e2a02e2a71a4464ddff52bdd0809874
1 file changed
+3
-1
vllm/v1/attention/backends/utils.py
+3
-1
@@ -1096,6 +1096,7 @@ def get_dcp_local_seq_lens(
1096
Only consider dcp now, we can extend the case of cp based on this.
1097
"""
1098
seq_lens_i32 = seq_lens.to(torch.int32)
1099
+ rank_offsets: int | torch.Tensor
1100
if dcp_rank is None:
1101
rank_offsets = torch.arange(
1102
dcp_size,
@@ -1113,7 +1114,8 @@ def get_dcp_local_seq_lens(
1114
rank_offsets = dcp_rank
1115
seq_lens_tiled = seq_lens_i32
1116
else:
1116
- rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device)
1117
+ # Use the Python scalar directly to avoid a synchronizing H2D copy.
1118
+ rank_offsets = dcp_rank
1119
seq_lens_tiled = seq_lens_i32
1120
base = (
1121
seq_lens_tiled