import mlx.core as mx

B, T, N, D = (1, 128, 32, 64) 
scale = D ** -0.5

q = mx.random.normal((B, T, N, D), dtype=mx.float16).transpose(0, 2, 1, 3)
k = mx.random.normal((B, T, N, D), dtype=mx.float16).transpose(0, 2, 1, 3)
v = mx.random.normal((B, T, N, D), dtype=mx.float16).transpose(0, 2, 1, 3)

o = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale).transpose(0, 2, 1, 3)
check_row_contiguous = mx.flatten(o)
mx.eval(check_row_contiguous)