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

[Bugfix][Model] Fix Mllama SDPA illegal memory access for batched multi-image (#9626)

Browse Source 
Signed-off-by: mgoin <michael@neuralmagic.com>
This commit is contained in:
Michael Goin 2024-10-23 22:03:44 -04:00 committed by GitHub
parent b548d7a5f4
commit bb01f2915e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/model_executor/models/mllama.py
View File

@ -795,17 +795,19 @@ class MllamaTextCrossAttention(nn.Module):
kv_len = k.shape[0] kv_len = k.shape[0]
q = q.transpose(0, 1).view(self.num_local_key_value_heads, q = q.transpose(0, 1).view(self.num_local_key_value_heads,
self.num_key_value_groups, q_len, self.num_key_value_groups, q_len,
self.head_dim) self.head_dim).contiguous()
k = k.transpose(0, k = k.transpose(0,
1)[:, 1)[:,
None, :, :].expand(self.num_local_key_value_heads, None, :, :].expand(self.num_local_key_value_heads,
self.num_key_value_groups, self.num_key_value_groups,
kv_len, self.head_dim) kv_len,
self.head_dim).contiguous()
v = v.transpose(0, v = v.transpose(0,
1)[:, 1)[:,
None, :, :].expand(self.num_local_key_value_heads, None, :, :].expand(self.num_local_key_value_heads,
self.num_key_value_groups, self.num_key_value_groups,
kv_len, self.head_dim) kv_len,
self.head_dim).contiguous()
attention_mask = attention_mask.view(1, 1, q_len, kv_len) attention_mask = attention_mask.view(1, 1, q_len, kv_len)
output = F.scaled_dot_product_attention(q, output = F.scaled_dot_product_attention(q,
k, k,