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

8 bit working

Browse Source
This commit is contained in:
Alex Barron
2024-10-22 20:09:27 -07:00
parent ef14b1e9c3
commit 047a584e3d
6 changed files with 127 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
benchmarks/python/sdpa_vector_bench.py
View File

@@ -25,18 +25,18 @@ def attention(q, k, v):
def sdpa(q, k, v):
k = mx.quantize(k)
v = mx.quantize(v)
k = mx.dequantize(*k)
v = mx.dequantize(*v)
return mx.fast.scaled_dot_product_attention(q, k, v, scale=0.08, mask=None)
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)
v = mx.quantize(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=0.08, mask=None
q, *k, *v, scale=1.0, mask=None, bits=8
)