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 in more models (#25005)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Wenlong Wang committed
Sep 20, 2025 at 20:55 UTC
035fd2bd2cd2fb70f5834f5ca6c2ea30cdae9187
3 files changed
+29
-19
vllm/model_executor/models/ernie45_vl.py
+9
-6
@@ -234,8 +234,9 @@ class Ernie4_5_VisionAttention(nn.Module):
234
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
235
for x in (q, k, v))
236
if rotary_pos_emb is not None:
237
- q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
238
- k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
237
+ qk_concat = torch.cat([q, k], dim=0)
238
+ qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
239
+ q, k = torch.chunk(qk_rotated, 2, dim=0)
240
241
if self.is_flash_attn_backend:
242
# from vllm_flash_attn.flash_attn_interface import (
@@ -261,8 +262,8 @@ class Ernie4_5_VisionAttention(nn.Module):
262
causal=False)
263
264
context_layer = rearrange(output,
264
- "(b s) ... -> b s ...",
265
- b=batch_size)
265
+ "(b s) h d -> s b (h d)",
266
+ b=batch_size).contiguous()
267
elif self.attn_backend == _Backend.TORCH_SDPA:
268
# Execute attention entry by entry for speed & less VRAM.
269
outputs = []
@@ -281,6 +282,8 @@ class Ernie4_5_VisionAttention(nn.Module):
282
output_i = rearrange(output_i, "b h s d -> b s h d ")
283
outputs.append(output_i)
284
context_layer = torch.cat(outputs, dim=1)
285
+ context_layer = rearrange(context_layer,
286
+ "b s h d -> s b (h d)").contiguous()
287
elif self.attn_backend == _Backend.XFORMERS:
288
from xformers import ops as xops
289
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
@@ -291,8 +294,8 @@ class Ernie4_5_VisionAttention(nn.Module):
294
295
context_layer = xops.memory_efficient_attention_forward(
296
q, k, v, attn_bias=attn_bias, p=0, scale=None)
294
- context_layer = rearrange(context_layer,
295
- "b s h d -> s b (h d)").contiguous()
297
+ context_layer = rearrange(context_layer,
298
+ "b s h d -> s b (h d)").contiguous()
299
300
output, _ = self.proj(context_layer)
301
return output
vllm/model_executor/models/glm4_1v.py
+10
-7
@@ -315,8 +315,10 @@ class Glm4vVisionAttention(nn.Module):
315
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
316
for x in (q, k, v))
317
if rotary_pos_emb is not None:
318
- q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
319
- k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
318
+ # [2 * b, s, heads, head_dim]
319
+ qk_concat = torch.cat([q, k], dim=0)
320
+ qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
321
+ q, k = torch.chunk(qk_rotated, 2, dim=0)
322
323
if self.attn_backend == _Backend.FLASH_ATTN:
324
# from vllm_flash_attn.flash_attn_interface import (
@@ -341,8 +343,8 @@ class Glm4vVisionAttention(nn.Module):
343
)
344
345
context_layer = rearrange(output,
344
- "(b s) ... -> b s ...",
345
- b=batch_size)
346
+ "(b s) h d -> s b (h d)",
347
+ b=batch_size).contiguous()
348
elif self.attn_backend == _Backend.TORCH_SDPA:
349
# Execute attention entry by entry for speed & less VRAM.
350
outputs = []
@@ -361,6 +363,8 @@ class Glm4vVisionAttention(nn.Module):
363
output_i = rearrange(output_i, "b h s d -> b s h d ")
364
outputs.append(output_i)
365
context_layer = torch.cat(outputs, dim=1)
366
+ context_layer = rearrange(context_layer,
367
+ "b s h d -> s b (h d)").contiguous()
368
elif self.attn_backend == _Backend.XFORMERS:
369
from xformers import ops as xops
370
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
@@ -371,9 +375,8 @@ class Glm4vVisionAttention(nn.Module):
375
376
context_layer = xops.memory_efficient_attention_forward(
377
q, k, v, attn_bias=attn_bias, p=0, scale=None)
374
-
375
- context_layer = rearrange(context_layer,
376
- "b s h d -> s b (h d)").contiguous()
378
+ context_layer = rearrange(context_layer,
379
+ "b s h d -> s b (h d)").contiguous()
380
381
output, _ = self.proj(context_layer)
382
return output
vllm/model_executor/models/qwen2_vl.py
+10
-6
@@ -377,8 +377,10 @@ class Qwen2VisionAttention(nn.Module):
377
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
378
for x in (q, k, v))
379
if rotary_pos_emb is not None:
380
- q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
381
- k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
380
+ # [2 * b, s, heads, head_dim]
381
+ qk_concat = torch.cat([q, k], dim=0)
382
+ qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
383
+ q, k = torch.chunk(qk_rotated, 2, dim=0)
384
385
if self.is_flash_attn_backend:
386
if self.attn_backend == _Backend.ROCM_AITER_FA:
@@ -402,8 +404,8 @@ class Qwen2VisionAttention(nn.Module):
404
causal=False)
405
406
context_layer = rearrange(output,
405
- "(b s) ... -> b s ...",
406
- b=batch_size)
407
+ "(b s) h d -> s b (h d)",
408
+ b=batch_size).contiguous()
409
elif self.attn_backend == _Backend.TORCH_SDPA:
410
# Execute attention entry by entry for speed & less VRAM.
411
outputs = []
@@ -422,6 +424,8 @@ class Qwen2VisionAttention(nn.Module):
424
output_i = rearrange(output_i, "b h s d -> b s h d ")
425
outputs.append(output_i)
426
context_layer = torch.cat(outputs, dim=1)
427
+ context_layer = rearrange(context_layer,
428
+ "b s h d -> s b (h d)").contiguous()
429
elif self.attn_backend == _Backend.XFORMERS:
430
from xformers import ops as xops
431
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
@@ -432,8 +436,8 @@ class Qwen2VisionAttention(nn.Module):
436
437
context_layer = xops.memory_efficient_attention_forward(
438
q, k, v, attn_bias=attn_bias, p=0, scale=None)
435
- context_layer = rearrange(context_layer,
436
- "b s h d -> s b (h d)").contiguous()
439
+ context_layer = rearrange(context_layer,
440
+ "b s h d -> s b (h d)").contiguous()
441
442
output, _ = self.proj(context_layer)
443
return output