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 14, 05:30 UTC
[Multi Modal][Performance] Fused Q,K's apply_rope into one (#24511)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Wenlong Wang committed
Sep 14, 2025 at 01:10 UTC
cc3173ae982ce4e589b24b70263cc25c3035d172
1 file changed
+10
-6
vllm/model_executor/models/qwen2_5_vl.py
+10
-6
@@ -363,8 +363,10 @@ class Qwen2_5_VisionAttention(nn.Module):
363
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
364
for x in (q, k, v))
365
if rotary_pos_emb is not None:
366
- q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
367
- k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
366
+ # [2 * b, s, heads, head_dim]
367
+ qk_concat = torch.cat([q, k], dim=0)
368
+ qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
369
+ q, k = torch.chunk(qk_rotated, 2, dim=0)
370
371
if self.is_flash_attn_backend:
372
if self.attn_backend == _Backend.ROCM_AITER_FA:
@@ -388,8 +390,8 @@ class Qwen2_5_VisionAttention(nn.Module):
390
causal=False)
391
392
context_layer = rearrange(output,
391
- "(b s) ... -> b s ...",
392
- b=batch_size)
393
+ "(b s) h d -> s b (h d)",
394
+ b=batch_size).contiguous()
395
elif self.attn_backend == _Backend.TORCH_SDPA:
396
# Execute attention entry by entry for speed & less VRAM.
397
outputs = []
@@ -408,6 +410,8 @@ class Qwen2_5_VisionAttention(nn.Module):
410
output_i = rearrange(output_i, "b h s d -> b s h d ")
411
outputs.append(output_i)
412
context_layer = torch.cat(outputs, dim=1)
413
+ context_layer = rearrange(context_layer,
414
+ "b s h d -> s b (h d)").contiguous()
415
elif self.attn_backend == _Backend.XFORMERS:
416
from xformers import ops as xops
417
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
@@ -418,8 +422,8 @@ class Qwen2_5_VisionAttention(nn.Module):
422
423
context_layer = xops.memory_efficient_attention_forward(
424
q, k, v, attn_bias=attn_bias, p=0, scale=None)
421
- context_layer = rearrange(context_layer,
422
- "b s h d -> s b (h d)").contiguous()
425
+ context_layer = rearrange(context_layer,
426
+ "b s h d -> s b (h d)").contiguous()
427
428
output, _ = self.proj(context_layer)
429
return output