diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index fa6d0398..401e53d7 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -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)