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

fix copies in sdpa (#2563)

Browse Source
This commit is contained in:
Awni Hannun
2025-09-02 11:00:36 -07:00
committed by GitHub
parent 04cbb4191c
commit b61a65e313
2 changed files with 12 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/tests/test_fast_sdpa.py
View File

@@ -619,6 +619,17 @@ class TestSDPA(mlx_tests.MLXTestCase):
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))
def test_sdpa_noncontiguous_inputs(self):
mask = mx.ones(shape=(4, 1, 7, 7), dtype=mx.bool_)
mx.random.seed(0)
q = mx.random.normal(shape=(4, 7, 32, 64)).swapaxes(1, 2)
k = mx.random.normal(shape=(4, 7, 8, 64)).swapaxes(1, 2)
v = mx.random.normal(shape=(4, 7, 8, 64)).swapaxes(1, 2)
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
ref = mlx_ref_attn(q, k, v, scale=1.0, mask=mask)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
def test_sdpa_promote_mask(self):
mask = mx.array(2.0, mx.bfloat16)
D = 64