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

Sdpa fix (#1558)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-11-02 21:25:46 -07:00
committed by GitHub
parent 09bc32f62f
commit 62f297b51d
4 changed files with 17 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/tests/test_fast_sdpa.py
View File

@@ -167,6 +167,15 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(o_mlx, reference, rtol=rtol, atol=atol))
q = mx.random.normal(shape=(1, 32, 1, Dk))
k = mx.random.normal(shape=(1, 32, 32, Dk))
v = mx.random.normal(shape=(1, 32, 128, Dk))
atol = 1e-6
y = mlx_primitives_sdpa(q, k, v[:, :, :32], scale)
y_hat = mx.fast.scaled_dot_product_attention(q, k, v[:, :, :32], scale=scale)
self.assertTrue(mx.allclose(y, y_hat, atol=atol))
if __name__ == "__main__":
unittest.main(failfast=True)