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

Change gqa to use repeat instead of concatenate (#443)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-02-14 17:40:11 -08:00 committed by GitHub
parent 06ddb8414d
commit f71e965d57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 16 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
llms/gguf_llm/models.py
View File

@ -107,12 +107,9 @@ class Attention(nn.Module):
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2)
return a.reshape([B, self.n_heads, L, -1])
if self.repeats > 1: if self.repeats > 1:
keys, values = map(repeat, (keys, values)) keys = mx.repeat(keys, self.repeats, axis=1)
values = mx.repeat(values, self.repeats, axis=1)
if cache is not None: if cache is not None:
key_cache, value_cache = cache key_cache, value_cache = cache

7
llms/mistral/mistral.py
View File

@ -73,11 +73,8 @@ class Attention(nn.Module):
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
def repeat(a): keys = mx.repeat(keys, self.repeats, axis=1)
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2) values = mx.repeat(values, self.repeats, axis=1)
return a.reshape([B, self.n_heads, L, -1])
keys, values = map(repeat, (keys, values))
if cache is not None: if cache is not None:
key_cache, value_cache = cache key_cache, value_cache = cache

7
llms/mixtral/mixtral.py
View File

@ -93,11 +93,8 @@ class Attention(nn.Module):
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
def repeat(a): keys = mx.repeat(keys, self.repeats, axis=1)
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2) values = mx.repeat(values, self.repeats, axis=1)
return a.reshape([B, self.n_heads, L, -1])
keys, values = map(repeat, (keys, values))
if cache is not None: if cache is not None:
key_cache, value_cache = cache key_cache, value_cache = cache

7
llms/mlx_lm/models/llama.py
View File

@ -93,12 +93,9 @@ class Attention(nn.Module):
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2)
return a.reshape([B, self.n_heads, L, -1])
if self.repeats > 1: if self.repeats > 1:
keys, values = map(repeat, (keys, values)) keys = mx.repeat(keys, self.repeats, axis=1)
values = mx.repeat(values, self.repeats, axis=1)
if cache is not None: if cache is not None:
key_cache, value_cache = cache key_cache, value_cache = cache

7
llms/mlx_lm/models/mixtral.py
View File

@ -95,12 +95,9 @@ class MixtralAttention(nn.Module):
0, 2, 1, 3 0, 2, 1, 3
) )
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2)
return a.reshape([B, self.num_heads, L, -1])
if self.repeats > 1: if self.repeats > 1:
keys, values = map(repeat, (keys, values)) keys = mx.repeat(keys, self.repeats, axis=1)
values = mx.repeat(values, self.repeats, axis=1)
if cache is not None: if cache is not None:
key_cache, value_cache = cache key_cache, value_cache = cache

7
llms/mlx_lm/models/phi.py
View File

@ -86,12 +86,9 @@ class PhiAttention(nn.Module):
B, L, self.num_key_value_heads, self.head_dim B, L, self.num_key_value_heads, self.head_dim
).transpose(0, 2, 1, 3) ).transpose(0, 2, 1, 3)
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2)
return a.reshape([B, self.num_heads, L, -1])
if self.repeats > 1: if self.repeats > 1:
keys, values = map(repeat, (keys, values)) keys = mx.repeat(keys, self.repeats, axis=1)
values = mx.repeat(values, self.repeats, axis=1)
# Add RoPE to the queries and keys and combine them with the cache # Add RoPE to the queries and keys and combine them with the cache
if cache is not None: if cache is not None:

7
llms/mlx_lm/models/qwen2.py
View File

@ -93,12 +93,9 @@ class Attention(nn.Module):
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2)
return a.reshape([B, self.n_heads, L, -1])
if self.repeats > 1: if self.repeats > 1:
keys, values = map(repeat, (keys, values)) keys = mx.repeat(keys, self.repeats, axis=1)
values = mx.repeat(values, self.repeats, axis=1)
if cache is not None: if cache is not None:
key_cache, value_cache = cache key_cache, value_cache = cache

7
llms/mlx_lm/models/stablelm_epoch.py
View File

@ -87,12 +87,9 @@ class Attention(nn.Module):
B, L, self.num_key_value_heads, self.head_dim B, L, self.num_key_value_heads, self.head_dim
).transpose(0, 2, 1, 3) ).transpose(0, 2, 1, 3)
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * self.repeats, axis=2)
return a.reshape([B, self.num_heads, L, -1])
if self.repeats > 1: if self.repeats > 1:
keys, values = map(repeat, (keys, values)) keys = mx.repeat(keys, self.repeats, axis=1)
values = mx.repeat(values, self.repeats, axis=1)
# Add RoPE to the queries and keys and combine them with the cache # Add RoPE to the queries and keys and combine them with the cache
if cache is not None: if cache is not None: