xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-16 11:26:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix/CI] Fix broken kernels/test_mha.py (#12450)

Browse Source
This commit is contained in:
Tyler Michael Smith 2025-01-26 13:39:03 -05:00 committed by GitHub
parent aa2cd2c43d
commit 72f4880425
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/kernels/test_mha_attn.py
View File

@ -26,7 +26,7 @@ def clear_cache():
@pytest.mark.parametrize("device", ["cpu", "hip", "cuda"])
def test_mha_attn_platform(device: str):
"""
Test that the attention selector between different platform and device.
Test the attention selector between different platform and device.
"""
torch.set_default_dtype(torch.float16)
@ -41,7 +41,7 @@ def test_mha_attn_platform(device: str):
else:
with patch("vllm.attention.selector.current_platform", CudaPlatform()):
attn = MultiHeadAttention(16, 64, scale=1)
assert attn.attn_backend == _Backend.FLASH_ATTN
assert attn.attn_backend == _Backend.XFORMERS
with patch("vllm.attention.selector.current_platform", CudaPlatform()):
attn = MultiHeadAttention(16, 72, scale=1)

8
vllm/attention/layer.py
View File

@ -210,6 +210,9 @@ class MultiHeadAttention(nn.Module):
self.scale = scale
self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
assert self.num_heads % self.num_kv_heads == 0
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
dtype = torch.get_default_dtype()
attn_backend = get_attn_backend(head_size,
dtype,
@ -240,6 +243,11 @@ class MultiHeadAttention(nn.Module):
key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)
if (num_repeat := self.num_queries_per_kv) > 1:
# Handle MQA and GQA
key = torch.repeat_interleave(key, num_repeat, dim=2)
value = torch.repeat_interleave(value, num_repeat, dim=2)
if self.attn_backend == _Backend.XFORMERS:
from xformers import ops as xops