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, 21:15 UTC
[Perf] Remove sync point in vit torch sdpa attn backend (#30232)
Signed-off-by: Dazhi Jiang <dazhi_jiang@163.com>
Dazhi Jiang committed
Dec 8, 2025 at 15:12 UTC
bcb6f5947f8acac22f5d1d5fc92a91b06fd57e77
4 files changed
+24
-24
vllm/attention/ops/vit_attn_wrappers.py
+6
-6
@@ -93,12 +93,12 @@ def torch_sdpa_wrapper(
93
cu_seqlens: torch.Tensor,
94
) -> torch.Tensor:
95
outputs = []
96
- for i in range(1, len(cu_seqlens)):
97
- start_idx = cu_seqlens[i - 1]
98
- end_idx = cu_seqlens[i]
99
- q_i = q[:, start_idx:end_idx]
100
- k_i = k[:, start_idx:end_idx]
101
- v_i = v[:, start_idx:end_idx]
96
+
97
+ lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
98
+ q_chunks = torch.split(q, lens, dim=1)
99
+ k_chunks = torch.split(k, lens, dim=1)
100
+ v_chunks = torch.split(v, lens, dim=1)
101
+ for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
102
q_i, k_i, v_i = (
103
einops.rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
104
)
vllm/model_executor/models/ernie45_vl.py
+6
-6
@@ -289,12 +289,12 @@ class Ernie4_5_VisionAttention(nn.Module):
289
elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
290
# Execute attention entry by entry for speed & less VRAM.
291
outputs = []
292
- for i in range(1, len(cu_seqlens)):
293
- start_idx = cu_seqlens[i - 1]
294
- end_idx = cu_seqlens[i]
295
- q_i = q[:, start_idx:end_idx]
296
- k_i = k[:, start_idx:end_idx]
297
- v_i = v[:, start_idx:end_idx]
292
+
293
+ lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
294
+ q_chunks = torch.split(q, lens, dim=1)
295
+ k_chunks = torch.split(k, lens, dim=1)
296
+ v_chunks = torch.split(v, lens, dim=1)
297
+ for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
298
q_i, k_i, v_i = (
299
rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
300
)
vllm/model_executor/models/glm4_1v.py
+6
-6
@@ -377,12 +377,12 @@ class Glm4vVisionAttention(nn.Module):
377
elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
378
# Execute attention entry by entry for speed & less VRAM.
379
outputs = []
380
- for i in range(1, len(cu_seqlens)):
381
- start_idx = cu_seqlens[i - 1]
382
- end_idx = cu_seqlens[i]
383
- q_i = q[:, start_idx:end_idx]
384
- k_i = k[:, start_idx:end_idx]
385
- v_i = v[:, start_idx:end_idx]
380
+
381
+ lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
382
+ q_chunks = torch.split(q, lens, dim=1)
383
+ k_chunks = torch.split(k, lens, dim=1)
384
+ v_chunks = torch.split(v, lens, dim=1)
385
+ for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
386
q_i, k_i, v_i = (
387
rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
388
)
vllm/model_executor/models/qwen2_vl.py
+6
-6
@@ -424,12 +424,12 @@ class Qwen2VisionAttention(nn.Module):
424
k = k.contiguous()
425
v = v.contiguous()
426
outputs = []
427
- for i in range(1, len(cu_seqlens)):
428
- start_idx = cu_seqlens[i - 1]
429
- end_idx = cu_seqlens[i]
430
- q_i = q[:, start_idx:end_idx]
431
- k_i = k[:, start_idx:end_idx]
432
- v_i = v[:, start_idx:end_idx]
427
+
428
+ lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
429
+ q_chunks = torch.split(q, lens, dim=1)
430
+ k_chunks = torch.split(k, lens, dim=1)
431
+ v_chunks = torch.split(v, lens, dim=1)
432
+ for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
433
q_i, k_i, v_i = (
434
rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
435
)