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

fix mask in sdpa (#1980)

Browse Source 
* fix mask in sdpa

* fix attention mask

* Re-enable routing for array mask

---------

Co-authored-by: Jagrit Digani <digani@apple.com>
This commit is contained in:
Awni Hannun
2025-03-20 14:53:12 -07:00
committed by GitHub
parent b42d13ec84
commit 005e7efa64
3 changed files with 34 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/tests/test_fast_sdpa.py
View File

@@ -527,6 +527,22 @@ class TestSDPA(mlx_tests.MLXTestCase):
diff = mx.abs(out_fst - out_ref) - atol * mx.abs(out_ref)
self.assertLessEqual(mx.max(diff).item(), atol)
def test_sdpa_broadcast_mask(self):
mask = mx.array(True)
D = 64
Nq = 4
Nkv = 1
scale = 1.0
L = 256
mx.random.seed(0)
q = 5e-1 * mx.random.normal(shape=(1, Nq, L, D))
k = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D))
v = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D))
ref = mlx_primitives_sdpa(q, k, v, scale, mask=mask)
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
if __name__ == "__main__":
unittest.main(failfast=True)