zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-02 09:18:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

clean

Browse Source
This commit is contained in:
Alex Barron
2024-10-22 20:13:32 -07:00
parent 6649244686
commit 852336b8a2
4 changed files with 37 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
benchmarks/python/sdpa_vector_bench.py
View File

@@ -1,18 +1,15 @@
import mlx.core as mx
import numpy as np
from mlx.utils import tree_map
from time_utils import time_fn
L = 16
L = 65536
H = 32
H_k = 32 // 4
D = 128
def attention(q, k, v):
k = mx.quantize(k)
v = mx.quantize(v)
k = mx.dequantize(*k)
v = mx.dequantize(*v)
B, Hq, L, D = q.shape
_, Hk, S, _ = k.shape
q = q.reshape(B, Hk, Hq // Hk, L, D)
@@ -25,21 +22,31 @@ def attention(q, k, v):
def sdpa(q, k, v):
k = mx.quantize(k, bits=8)
v = mx.quantize(v, bits=8)
k = mx.dequantize(*k, bits=8)
v = mx.dequantize(*v, bits=8)
return mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=None)
def quant_sdpa(q, k, v):
k = mx.quantize(k, bits=8)
v = mx.quantize(v, bits=8)
return mx.fast.quantized_scaled_dot_product_attention(
q, *k, *v, scale=1.0, mask=None, bits=8
)
def quant_attention(q, k, v):
B, Hq, L, D = q.shape
Hk = k[0].shape[1]
q = q.reshape((B, Hk, Hq // Hk, L, D))
k = tree_map(lambda x: mx.expand_dims(x, axis=2), k)
v = tree_map(lambda x: mx.expand_dims(x, axis=2), v)
scores = mx.quantized_matmul(q, *k, transpose=True)
scores = mx.softmax(scores, axis=-1)
out = mx.quantized_matmul(scores, *v, transpose=False)
out = out.reshape((B, Hq, L, D))
return out
def time_self_attention_primitives(q, k, v):
time_fn(attention, q, k, v)
@@ -52,34 +59,22 @@ def time_self_attention_quant_sdpa(q, k, v):
time_fn(quant_sdpa, q, k, v)
def time_self_attention_quant_primitives(q, k, v):
time_fn(quant_attention, q, k, v)
if __name__ == "__main__":
mx.random.seed(3)
# q = mx.random.uniform(shape=(1, H, 1, D))
# k = mx.random.uniform(shape=(1, H_k, L, D))
# v = mx.random.uniform(shape=(1, H_k, L, D))
q = mx.array(np.load("/Users/alexbarron/mlx-examples/llms/queries.npy"))
k = mx.array(np.load("/Users/alexbarron/mlx-examples/llms/keys.npy"))
v = mx.array(np.load("/Users/alexbarron/mlx-examples/llms/values.npy"))
print(q.dtype)
print(q.shape, k.shape, v.shape)
q = mx.random.uniform(shape=(1, H, 1, D))
k = mx.random.uniform(shape=(1, H_k, L, D))
v = mx.random.uniform(shape=(1, H_k, L, D))
mx.eval(q, k, v)
k_quant = mx.quantize(k)
v_quant = mx.quantize(v)
mx.eval(k_quant, v_quant)
# time_self_attention_sdpa(q, k, v)
# time_self_attention_quant_sdpa(q, k_quant, v_quant)
# time_self_attention_primitives(q, k, v)
q_sdpa = quant_sdpa(q, k, v)
print(q_sdpa)
# o_attention = attention(q, k, v)
# print(o_attention)
# np.testing.assert_allclose(q_sdpa, o_attention, atol=1e-5)
o_sdpa = sdpa(q, k, v)
print(o_sdpa)
np.testing.assert_allclose(q_sdpa, o_sdpa, atol=1e-5)
# print(o_sdpa[..., :64])
# print()
# print(o_attention[..., :64])
# np.testing.assert_allclose(o_sdpa, o_attention)
time_self_attention_sdpa(q, k, v)
time_self_attention_quant_sdpa(q, k_quant, v_quant)
time_self_attention_primitives(q, k, v)
time_self_attention_quant_primitives(q, k_quant, v_quant)