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 15, 19:45 UTC
[Perf]:Optimize qwen2-vl to reduce cudaMemcpyAsync (#14377)
Signed-off-by: cynthieye <987073381@qq.com>
yexin(叶鑫) committed
Mar 11, 2025 at 15:39 UTC
70b808fe1a63322bc6bf5f46a91981a8f6b8af00
2 files changed
+70
-24
vllm/model_executor/models/qwen2_5_vl.py
+33
-12
@@ -255,10 +255,12 @@ class Qwen2_5_VisionAttention(nn.Module):
255
return q, k, v
256
257
def forward(
258
- self,
259
- x: torch.Tensor,
260
- cu_seqlens: torch.Tensor,
261
- rotary_pos_emb: torch.Tensor,
258
+ self,
259
+ x: torch.Tensor,
260
+ cu_seqlens: torch.Tensor,
261
+ rotary_pos_emb: torch.Tensor,
262
+ max_seqlen: Optional[int] = None, # Only used for Flash Attention
263
+ seqlens: Optional[list[int]] = None, # Only used for xFormers
264
) -> torch.Tensor:
265
# [s, b, c] --> [s, b, head * 3 * head_dim]
266
x, _ = self.qkv(x)
@@ -285,7 +287,6 @@ class Qwen2_5_VisionAttention(nn.Module):
287
288
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
289
288
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
290
output = flash_attn_varlen_func(q,
291
k,
292
v,
@@ -321,7 +322,6 @@ class Qwen2_5_VisionAttention(nn.Module):
322
from xformers import ops as xops
323
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
324
324
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
325
attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
326
kv_seqlen=None,
327
device=q.device)
@@ -364,11 +364,20 @@ class Qwen2_5_VisionBlock(nn.Module):
364
quant_config=quant_config,
365
prefix=f"{prefix}.mlp")
366
367
- def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
368
- rotary_pos_emb: torch.Tensor) -> torch.Tensor:
367
+ def forward(
368
+ self,
369
+ x: torch.Tensor,
370
+ cu_seqlens: torch.Tensor,
371
+ rotary_pos_emb: torch.Tensor,
372
+ max_seqlen: Optional[int] = None, # Only used for Flash Attention
373
+ seqlens: Optional[list[int]] = None, # Only used for xFormers
374
+ ) -> torch.Tensor:
375
x = x + self.attn(self.norm1(x),
376
cu_seqlens=cu_seqlens,
371
- rotary_pos_emb=rotary_pos_emb)
377
+ rotary_pos_emb=rotary_pos_emb,
378
+ max_seqlen=max_seqlen,
379
+ seqlens=seqlens)
380
+
381
x = x + self.mlp(self.norm2(x))
382
return x
383
@@ -528,6 +537,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
537
quant_config=quant_config,
538
prefix=f"{prefix}.merger",
539
)
540
+ self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
541
542
@property
543
def dtype(self) -> torch.dtype:
@@ -633,14 +643,25 @@ class Qwen2_5_VisionTransformer(nn.Module):
643
644
# transformers
645
hidden_states = hidden_states.unsqueeze(1)
646
+
647
+ max_seqlen = None
648
+ seqlens = None
649
+ if self.attn_backend == _Backend.FLASH_ATTN:
650
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
651
+ elif self.attn_backend == _Backend.XFORMERS:
652
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
653
for layer_num, blk in enumerate(self.blocks):
654
if layer_num in self.fullatt_block_indexes:
655
cu_seqlens_now = cu_seqlens
656
else:
657
cu_seqlens_now = cu_window_seqlens
641
- hidden_states = blk(hidden_states,
642
- cu_seqlens=cu_seqlens_now,
643
- rotary_pos_emb=rotary_pos_emb)
658
+ hidden_states = blk(
659
+ hidden_states,
660
+ cu_seqlens=cu_seqlens_now,
661
+ rotary_pos_emb=rotary_pos_emb,
662
+ max_seqlen=max_seqlen,
663
+ seqlens=seqlens,
664
+ )
665
666
# For Qwen2.5-VL-3B, float16 will overflow at last block
667
# for long visual tokens sequences.
vllm/model_executor/models/qwen2_vl.py
+37
-12
@@ -303,10 +303,12 @@ class Qwen2VisionAttention(nn.Module):
303
return q, k, v
304
305
def forward(
306
- self,
307
- x: torch.Tensor,
308
- cu_seqlens: torch.Tensor,
309
- rotary_pos_emb: torch.Tensor,
306
+ self,
307
+ x: torch.Tensor,
308
+ cu_seqlens: torch.Tensor,
309
+ rotary_pos_emb: torch.Tensor,
310
+ max_seqlen: Optional[int] = None, # Only used for Flash Attention
311
+ seqlens: Optional[list[int]] = None, # Only used for xFormers
312
) -> torch.Tensor:
313
314
# [s, b, c] --> [s, b, 3 * head * head_dim]
@@ -329,7 +331,6 @@ class Qwen2VisionAttention(nn.Module):
331
332
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
333
332
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
334
output = flash_attn_varlen_func(q,
335
k,
336
v,
@@ -365,7 +366,6 @@ class Qwen2VisionAttention(nn.Module):
366
from xformers import ops as xops
367
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
368
368
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
369
attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
370
kv_seqlen=None,
371
device=q.device)
@@ -409,11 +409,22 @@ class Qwen2VisionBlock(nn.Module):
409
quant_config=quant_config,
410
prefix=f"{prefix}.mlp")
411
412
- def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
413
- rotary_pos_emb: torch.Tensor) -> torch.Tensor:
414
- x = x + self.attn(self.norm1(x),
415
- cu_seqlens=cu_seqlens,
416
- rotary_pos_emb=rotary_pos_emb)
412
+ def forward(
413
+ self,
414
+ x: torch.Tensor,
415
+ cu_seqlens: torch.Tensor,
416
+ rotary_pos_emb: torch.Tensor,
417
+ max_seqlen: Optional[int] = None, # Only used for Flash Attention
418
+ seqlens: Optional[list[int]] = None, # Only used for xFormers
419
+ ) -> torch.Tensor:
420
+ x = x + self.attn(
421
+ self.norm1(x),
422
+ cu_seqlens=cu_seqlens,
423
+ rotary_pos_emb=rotary_pos_emb,
424
+ max_seqlen=max_seqlen,
425
+ seqlens=seqlens,
426
+ )
427
+
428
x = x + self.mlp(self.norm2(x))
429
return x
430
@@ -570,6 +581,7 @@ class Qwen2VisionTransformer(nn.Module):
581
quant_config=quant_config,
582
prefix=f"{prefix}.merger",
583
)
584
+ self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
585
586
@property
587
def dtype(self) -> torch.dtype:
@@ -624,8 +636,21 @@ class Qwen2VisionTransformer(nn.Module):
636
637
# transformers
638
x = x.unsqueeze(1)
639
+
640
+ max_seqlen = None
641
+ seqlens = None
642
+ if self.attn_backend == _Backend.FLASH_ATTN:
643
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
644
+ elif self.attn_backend == _Backend.XFORMERS:
645
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
646
for blk in self.blocks:
628
- x = blk(x, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)
647
+ x = blk(
648
+ x,
649
+ cu_seqlens=cu_seqlens,
650
+ rotary_pos_emb=rotary_pos_emb,
651
+ max_seqlen=max_seqlen,
652
+ seqlens=seqlens,
653
+ )
654
655
# adapter
656
x = self.merger(x)