L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Cleanup test_fast_sdpa.py (#3112)

Browse Source
This commit is contained in:
Cheng
2026-02-08 19:04:24 +09:00
committed by GitHub
parent 8fe1d09207
commit 566bc16b7c
1 changed files with 71 additions and 265 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_fast_sdpa.py
+71 -265
View File
@@ -1,5 +1,6 @@
import math
import unittest
from itertools import product
import mlx.core as mx
import mlx_tests
@@ -74,31 +75,23 @@ def do_attention(f, q, k, v, scale, mask=None, transpose=False):
def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype):
np.random.seed(0)
np_dtype = getattr(np, dtype)
mx.random.seed(0)
scale = 1.0 / math.sqrt(D)
shape_q = (B, qL, qH, D) if transpose else (B, qH, qL, D)
shape_kv = (B, kL, kH, D) if transpose else (B, kH, kL, D)
scale = 1.0 / math.sqrt(D)
q_np = np.random.normal(0.0, 0.5, shape_q).astype(np_dtype)
k_np = np.random.normal(0.0, 0.5, shape_kv).astype(np_dtype)
v_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
q_mx = mx.array(q_np)
k_mx = mx.array(k_np)
v_mx = mx.array(v_np)
q = mx.random.uniform(0.0, 0.5, shape_q, dtype)
k = mx.random.uniform(0.0, 0.5, shape_kv, dtype)
v = mx.random.uniform(0.0, scale, shape_kv, dtype)
if mask is not None:
if mask == "additive":
mask_np = np.random.normal(0.0, 0.5, (B, qH, qL, kL)).astype(np_dtype)
mask = mx.array(mask_np)
mask = mx.random.uniform(0.0, 0.5, (B, qH, qL, kL), dtype)
elif mask == "bool":
mask_np = np.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5
mask = mx.array(mask_np)
mask = mx.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5
return q_mx, k_mx, v_mx, scale, mask
return q, k, v, scale, mask
# SDPA for MHA (n_heads == n_kv_heads)
@@ -122,192 +115,8 @@ def mlx_primitives_sdpa(q, k, v, scale, mask=None):
return scores @ v
# SDPA for GQA (n_heads > n_kv_heads, n_kv_heads > 1, n_heads % n_kv_heads == 0)
def mlx_primitives_sdpa_with_gqa(q, k, v, scale, mask=None):
n_repeats = q.shape[1] // k.shape[1]
# borrowing kv cache tiling from mlx-examples/llms/mistral/mistral.py
n_heads = q.shape[1]
B = q.shape[0]
L = k.shape[2]
def repeat(a):
a = mx.concatenate([mx.expand_dims(a, 2)] * n_repeats, axis=2)
return a.reshape([B, n_heads, L, -1])
k, v = map(repeat, (k, v))
return mlx_primitives_sdpa(q, k, v, scale, mask=mask)
class TestFastSelfAttentionSDPA(mlx_tests.MLXTestCase):
def test_fast_sdpa(self):
# Not yet supported:
# * K pre-transposed in kernel, V pre-transposed in kernel
np.random.seed(0)
R = 20
L = R
Dk = 64
H = 3
scale = float(1.0 / np.sqrt(Dk))
q_npy = np.random.normal(0.0, 1.0, (1, H, R, Dk)).astype(np.float32)
k_npy = np.random.normal(0.0, 1.0, (1, H, L, Dk)).astype(np.float32)
v_npy = np.random.normal(0.0, 1.0, (1, H, L, Dk)).astype(np.float32)
q_mlx = mx.array(q_npy)
k_mlx = mx.array(k_npy)
v_mlx = mx.array(v_npy)
reference = mlx_primitives_sdpa(q_mlx, k_mlx, v_mlx, scale)
o_mlx = mx.fast.scaled_dot_product_attention(
q_mlx, k_mlx, v_mlx, scale=scale, mask=None
)
self.assertListEqual(list(reference.shape), list(o_mlx.shape))
self.assertTrue(mx.allclose(o_mlx, reference, atol=1e-4))
dtypes = [np.float32]
Dk = 64
if mx.is_available(mx.gpu):
dtypes.append(np.half)
for SEQUENCE_LENGTH in [63, 129, 400]:
for DTYPE in dtypes:
B = 2
H = 24
n_kv_heads = H
q_npy = np.random.normal(0.0, 1.0, (B, H, SEQUENCE_LENGTH, Dk)).astype(
DTYPE
)
k_npy = np.random.normal(
0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk)
).astype(DTYPE)
v_npy = np.random.normal(
0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk)
).astype(DTYPE)
q_mlx = mx.array(q_npy)
k_mlx = mx.array(k_npy)
v_mlx = mx.array(v_npy)
reference = mlx_primitives_sdpa_with_gqa(q_mlx, k_mlx, v_mlx, scale)
o_mlx = mx.fast.scaled_dot_product_attention(
q_mlx,
k_mlx,
v_mlx,
scale=scale,
)
self.assertListEqual(list(reference.shape), list(o_mlx.shape))
rtol = 1e-3
atol = 1e-2
if SEQUENCE_LENGTH > 500:
rtol = 1e-2
if DTYPE == np.half:
rtol = 1e-2
self.assertTrue(mx.allclose(o_mlx, reference, rtol=rtol, atol=atol))
class TestFastSDPA(mlx_tests.MLXTestCase):
def test_fast_sdpa(self):
# Not yet supported:
# * K pre-transposed in kernel, V pre-transposed in kernel
np.random.seed(0)
L = 43
R = 1
Dk = 128
scale = float(1.0 / np.sqrt(128.0))
q_npy = np.random.normal(0.0, 1.0, (1, 32, R, Dk)).astype(np.float32)
k_npy = np.random.normal(0.0, 1.0, (1, 32, L, Dk)).astype(np.float32)
v_npy = np.random.normal(0.0, 1.0, (1, 32, L, Dk)).astype(np.float32)
q_mlx = mx.array(q_npy)
k_mlx = mx.array(k_npy)
v_mlx = mx.array(v_npy)
reference = mlx_primitives_sdpa(q_mlx, k_mlx, v_mlx, scale)
o_mlx = mx.fast.scaled_dot_product_attention(
q_mlx, k_mlx, v_mlx, scale=scale, mask=None
)
self.assertListEqual(list(reference.shape), list(o_mlx.shape))
self.assertTrue(mx.allclose(o_mlx, reference, atol=1e-4))
B = 1
H = 32
dtypes = [np.float32]
if mx.is_available(mx.gpu):
dtypes.append(np.half)
for SEQUENCE_LENGTH in [1, 7, 9, 32, 63, 67, 129, 400, 2000]:
for DO_GQA in [0, 1]:
for DTYPE in dtypes:
n_kv_heads = 8 if DO_GQA else 32
q_npy = np.random.normal(0.0, 1.0, (B, H, R, Dk)).astype(DTYPE)
k_npy = np.random.normal(
0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk)
).astype(DTYPE)
v_npy = np.random.normal(
0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk)
).astype(DTYPE)
q_mlx = mx.array(q_npy)
k_mlx = mx.array(k_npy)
v_mlx = mx.array(v_npy)
reference = mlx_primitives_sdpa_with_gqa(q_mlx, k_mlx, v_mlx, scale)
o_mlx = mx.fast.scaled_dot_product_attention(
q_mlx, k_mlx, v_mlx, scale=scale
)
self.assertListEqual(list(reference.shape), list(o_mlx.shape))
rtol = 1e-5
atol = 1e-1
if SEQUENCE_LENGTH > 500:
rtol = 1e-2
if DTYPE == np.half:
rtol = 1e-2
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))
# Test with per-example mask
q = mx.random.normal(shape=(2, 8, 4, 32))
k = mx.random.normal(shape=(2, 2, 8, 32))
v = mx.random.normal(shape=(2, 2, 8, 32))
mask = 10 * mx.random.normal(shape=(2, 1, 4, 8))
y = mlx_primitives_sdpa_with_gqa(q, k, v, scale, mask=mask)
y_hat = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
self.assertTrue(mx.allclose(y, y_hat, atol=atol))
# Test with boolean causal mask
indices = mx.arange(8)
bool_mask = indices[:, None] >= indices[None]
additive_mask = (~bool_mask).astype(mx.float32) * mx.finfo(mx.float32).min
x = mx.random.normal(shape=(1, 2, 8, 32))
y = mlx_primitives_sdpa_with_gqa(x, x, x, scale, mask=additive_mask)
y_hat = mx.fast.scaled_dot_product_attention(
x, x, x, scale=scale, mask=bool_mask
)
self.assertTrue(mx.allclose(y, y_hat, atol=atol))
def test_fast_sdpa_vector_kv_transposed_head_seq(self):
def test_sdpa_vector_kv_transposed_head_seq(self):
D = 64
Nq = 4
Nkv = 1
@@ -339,7 +148,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
def test_fast_sdpa_vector(self):
def test_sdpa_vector(self):
D = 64
L = 43
Nq = 4
@@ -411,7 +220,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
def test_fully_masked(self):
def test_sdpa_fully_masked(self):
Lkv = 8
mask = mx.array(False)
for D in [128]:
@@ -423,7 +232,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
out = mx.fast.scaled_dot_product_attention(q, k, v, mask=mask, scale=1)
self.assertFalse(mx.any(mx.isnan(out)))
def test_inf_score(self):
def test_sdpa_inf_score(self):
Lkv = 8
for D in [4, 128]:
for Lq in [1, 8]:
@@ -435,7 +244,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
out = mx.fast.scaled_dot_product_attention(q, k, v, mask=None, scale=1)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
def test_fast_sdpa_few_query(self):
def test_sdpa_few_query(self):
D = 64
L = 43
Lq = 8
@@ -494,7 +303,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
@unittest.skip("Different head and value dims is not enabled")
def test_fast_sdpa_vector_value_dims(self):
def test_sdpa_vector_value_dims(self):
D = 192
V = 128
Nq = 4
@@ -550,91 +359,88 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
ref = mlx_ref_attn(q, k, v, mask=mask)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
class TestSDPA(mlx_tests.MLXTestCase):
@unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU")
def test_sdpa(self):
if not mx.is_available(mx.gpu):
return
# fmt: off
shapes_64 = (
shapes_64 = [
# ( B, qsl, ksl, head_dim, n_qh, n_kvh)
( 1, 20, 20, 64, 3, 3),
( 1, 63, 63, 64, 24, 24),
( 1, 129, 129, 64, 24, 24),
( 1, 400, 400, 64, 24, 24),
( 1, 128, 128, 64, 32, 32),
( 1, 64, 128, 64, 32, 32),
( 1, 65, 128, 64, 32, 8),
( 1, 64, 127, 64, 32, 8),
( 1, 65, 127, 64, 32, 8),
( 1, 127, 65, 64, 32, 8),
)
shapes_128 = (
]
shapes_128 = [
# ( B, qsl, ksl, head_dim, n_qh, n_kvh)
( 1, 128, 128, 128, 32, 8),
( 1, 64, 128, 128, 32, 8),
( 1, 65, 127, 128, 32, 8),
( 1, 127, 65, 128, 32, 8),
)
]
for ksl in [7, 9, 32, 63, 67, 129, 400, 2000]:
shapes_128.append((1, 1, ksl, 128, 32, 32))
shapes_128.append((1, 1, ksl, 128, 32, 8))
# fmt: on
shapes = shapes_64 + shapes_128
dtypes = ["float32", "float16"]
dtypes = [mx.float16]
if mx.metal.is_available():
dtypes.append(mx.float32)
masks = [None, "additive", "bool", "causal"]
transposes = (False, True)
for dtype in dtypes:
for t in transposes:
for mask_str in masks:
for B, qL, kL, D, qH, kH in shapes:
with self.subTest(
B=B,
qsl=qL,
ksl=kL,
head_dim=D,
n_q_heads=qH,
n_kv_heads=kH,
mask=mask_str,
transpose=t,
dtype=dtype,
):
for dtype, t, mask_str, (B, qL, kL, D, qH, kH) in product(
dtypes, transposes, masks, shapes
):
with self.subTest(
B=B,
qsl=qL,
ksl=kL,
head_dim=D,
n_q_heads=qH,
n_kv_heads=kH,
mask=mask_str,
transpose=t,
dtype=dtype,
):
q, k, v, scale, mask = prepare_inputs(
B, qL, kL, D, qH, kH, mask_str, t, dtype
)
np.random.seed(0)
q_mx, k_mx, v_mx, scale, mask = prepare_inputs(
B, qL, kL, D, qH, kH, mask_str, t, dtype
)
out_ref = do_attention(mlx_ref_attn, q, k, v, scale, mask, t)
out_ref = do_attention(
mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, t
)
out_fst = do_attention(
mx.fast.scaled_dot_product_attention,
q,
k,
v,
scale,
mask,
t,
)
out_fst = do_attention(
mx.fast.scaled_dot_product_attention,
q_mx,
k_mx,
v_mx,
scale,
mask,
t,
)
# For causal mask when qL > kL, first qL-kL rows are undefined
# Compare only the valid portion
if mask_str == "causal" and qL > kL:
offset = qL - kL
if t: # transpose=True: shape is (B, qL, qH, D)
out_ref = out_ref[:, offset:, :, :]
out_fst = out_fst[:, offset:, :, :]
else: # transpose=False: shape is (B, qH, qL, D)
out_ref = out_ref[:, :, offset:, :]
out_fst = out_fst[:, :, offset:, :]
# For causal mask when qL > kL, first qL-kL rows are undefined
# Compare only the valid portion
if mask_str == "causal" and qL > kL:
offset = qL - kL
if t: # transpose=True: shape is (B, qL, qH, D)
out_ref = out_ref[:, offset:, :, :]
out_fst = out_fst[:, offset:, :, :]
else: # transpose=False: shape is (B, qH, qL, D)
out_ref = out_ref[:, :, offset:, :]
out_fst = out_fst[:, :, offset:, :]
atol = 2e-5 if dtype == mx.float32 else 3e-4
atol = 2e-5 if dtype == "float32" else 3e-4
self.assertListEqual(list(out_ref.shape), list(out_fst.shape))
self.assertListEqual(
list(out_ref.shape), list(out_fst.shape)
)
diff = mx.abs(out_fst - out_ref) - atol * mx.abs(out_ref)
self.assertLessEqual(mx.max(diff).item(), atol)
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)