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, 22:15 UTC
[Bugfix] Initialize attention bias on the same device as Query/Key/Value (#13468)
Junlin Zhou committed
Feb 25, 2025 at 18:13 UTC
75e9d4979658f19c8a507a241acd6be8b83a8f55
1 file changed
+6
-4
vllm/attention/backends/xformers.py
+6
-4
@@ -673,7 +673,9 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
673
674
# Cross-attention mask is non-causal
675
attn_bias = BlockDiagonalMask.from_seqlens(
676
- attn_metadata.seq_lens, attn_metadata.encoder_seq_lens)
676
+ attn_metadata.seq_lens,
677
+ attn_metadata.encoder_seq_lens,
678
+ device=query.device)
679
680
# Encoder branch of encoder-decoder model uses
681
# attn_metadata.encoder_seq_lens
@@ -683,7 +685,7 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
685
686
# Encoder self-attention mask is non-causal
687
attn_bias = BlockDiagonalMask.from_seqlens(
686
- attn_metadata.encoder_seq_lens)
688
+ attn_metadata.encoder_seq_lens, device=query.device)
689
690
# Self-attention block of encoder-only model just
691
# uses the seq_lens directly.
@@ -692,7 +694,7 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
694
695
# Encoder self-attention mask is non-causal
696
attn_bias = BlockDiagonalMask.from_seqlens(
695
- attn_metadata.seq_lens)
697
+ attn_metadata.seq_lens, device=query.device)
698
699
# Self-attention block of decoder branch just
700
# uses the seq_lens directly
@@ -701,7 +703,7 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
703
704
# Decoder self-attention mask is causal
705
attn_bias = BlockDiagonalCausalMask.from_seqlens(
704
- attn_metadata.seq_lens)
706
+ attn_metadata.seq_lens, device=query.device)
707
else:
708
raise ValueError("Unknown AttentionType: %s", attn_type)
709