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, 18:30 UTC
[Misc] Add attention mask pre-computation optimization back to Qwen2.5-VL (#15273)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Mar 21, 2025 at 18:32 UTC
47c7126213163e454019efc7d913125c22df9d6e
2 files changed
+35
-16
vllm/model_executor/models/qwen2_5_vl.py
+23
-10
@@ -608,6 +608,17 @@ class Qwen2_5_VisionTransformer(nn.Module):
608
window_index = torch.cat(window_index, dim=0)
609
return window_index, cu_window_seqlens
610
611
+ def compute_attn_mask_seqlen(
612
+ self,
613
+ cu_seqlens: torch.Tensor,
614
+ ) -> tuple[Optional[int], Optional[list[int]]]:
615
+ max_seqlen, seqlens = None, None
616
+ if self.attn_backend == _Backend.FLASH_ATTN:
617
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
618
+ elif self.attn_backend == _Backend.XFORMERS:
619
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
620
+ return max_seqlen, seqlens
621
+
622
def forward(
623
self,
624
x: torch.Tensor,
@@ -645,25 +656,27 @@ class Qwen2_5_VisionTransformer(nn.Module):
656
# transformers
657
hidden_states = hidden_states.unsqueeze(1)
658
648
- max_seqlen = None
649
- seqlens = None
659
+ # pre-compute seqlens for window/full attn to reduce cuMemcpy operations
660
+ max_seqlen_full, seqlens_full = self.compute_attn_mask_seqlen(
661
+ cu_seqlens)
662
+ max_seqlen_window, seqlens_window = self.compute_attn_mask_seqlen(
663
+ cu_window_seqlens)
664
for layer_num, blk in enumerate(self.blocks):
665
if layer_num in self.fullatt_block_indexes:
666
cu_seqlens_now = cu_seqlens
667
+ max_seqlen_now = max_seqlen_full
668
+ seqlens_now = seqlens_full
669
else:
670
cu_seqlens_now = cu_window_seqlens
655
- # pre-compute cu_seqlens for window attn
656
- if self.attn_backend == _Backend.FLASH_ATTN:
657
- max_seqlen = (cu_seqlens_now[1:] -
658
- cu_seqlens_now[:-1]).max().item()
659
- elif self.attn_backend == _Backend.XFORMERS:
660
- seqlens = (cu_seqlens_now[1:] - cu_seqlens_now[:-1]).tolist()
671
+ max_seqlen_now = max_seqlen_window
672
+ seqlens_now = seqlens_window
673
+
674
hidden_states = blk(
675
hidden_states,
676
cu_seqlens=cu_seqlens_now,
677
rotary_pos_emb=rotary_pos_emb,
665
- max_seqlen=max_seqlen,
666
- seqlens=seqlens,
678
+ max_seqlen=max_seqlen_now,
679
+ seqlens=seqlens_now,
680
)
681
682
# For Qwen2.5-VL-3B, float16 will overflow at last block
vllm/model_executor/models/qwen2_vl.py
+12
-6
@@ -617,6 +617,16 @@ class Qwen2VisionTransformer(nn.Module):
617
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
618
return rotary_pos_emb
619
620
+ def compute_attn_mask_seqlen(
621
+ self, cu_seqlens: torch.Tensor
622
+ ) -> tuple[Optional[int], Optional[list[int]]]:
623
+ max_seqlen, seqlens = None, None
624
+ if self.attn_backend == _Backend.FLASH_ATTN:
625
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
626
+ elif self.attn_backend == _Backend.XFORMERS:
627
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
628
+ return max_seqlen, seqlens
629
+
630
def forward(
631
self,
632
x: torch.Tensor,
@@ -638,12 +648,8 @@ class Qwen2VisionTransformer(nn.Module):
648
# transformers
649
x = x.unsqueeze(1)
650
641
- max_seqlen = None
642
- seqlens = None
643
- if self.attn_backend == _Backend.FLASH_ATTN:
644
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
645
- elif self.attn_backend == _Backend.XFORMERS:
646
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
651
+ # pre-compute seqlens for attn mask to reduce cuMemcpy operations
652
+ max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
653
for blk in self.blocks:
654
x = blk(
655
x,