zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-24 01:17:28 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix PLaMo model to support Grouped Query Attention (#1037)

Browse Source
This commit is contained in:
Shunta Saito 2024-10-13 07:26:50 +09:00 committed by GitHub
parent d8611dd69f
commit 7612c646f3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
llms/mlx_lm/models/plamo.py
View File

@ -89,6 +89,9 @@ class Attention(nn.Module):
queries = self.rotary_emb(queries)
keys = self.rotary_emb(keys)
keys = mx.tile(keys, [1, self.config.n_shared_head, 1, 1])
values = mx.tile(values, [1, self.config.n_shared_head, 1, 1])
output = mx.fast.scaled_dot_product_attention(
queries,
keys,