diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index f20bd8c2..d8ff1218 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -152,11 +152,7 @@ DnnGraph build_sdpa_graph( .set_attn_scale(graph.scalar("Scale", SCALE, float32)) .set_generate_stats(output_logsumexp); if (do_causal) { - if (q.shape(2) > k.shape(2)) { - options.set_causal_mask(do_causal); - } else { - options.set_causal_mask_bottom_right(do_causal); - } + options.set_causal_mask_bottom_right(do_causal); } if (mask_arr) { options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr)); @@ -201,11 +197,7 @@ DnnGraph build_sdpa_backward_graph( .set_name("sdpa_backward_cudnn") .set_attn_scale(graph.scalar("Scale", SCALE, float32)); if (do_causal) { - if (q.shape(2) > k.shape(2)) { - options.set_causal_mask(do_causal); - } else { - options.set_causal_mask_bottom_right(do_causal); - } + options.set_causal_mask_bottom_right(do_causal); } if (mask_arr) { options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr)); diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index 1814f9b9..6abee21f 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -461,7 +461,7 @@ template < metal::vec rcp; STEEL_PRAGMA_UNROLL for (short i = 0; i < kRowsPT; ++i) { - rcp[i] = (1.f / sum_score[i]); + rcp[i] = 1.f / sum_score[i]; } Otile.template row_bin_op(rcp); diff --git a/mlx/fast.cpp b/mlx/fast.cpp index dbc38cb1..bf140b7b 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -720,8 +720,8 @@ array scaled_dot_product_attention( if (do_causal) { int kL = k.shape(-2); int qL = q.shape(-2); - int q_off = (kL - qL) < 0 ? 0 : (kL - qL); - auto q_idx = arange(q_off, q_off + qL, s); + int offset = kL - qL; + auto q_idx = arange(offset, qL + offset, s); auto k_idx = arange(0, kL, s); q_idx = expand_dims(q_idx, 1, s); k_idx = expand_dims(k_idx, 0, s); diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 97dd632c..1a43d89d 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -272,6 +272,8 @@ void init_fast(nb::module_& parent_module) { can have at most 4 dimensions and must be broadcast-compatible with the shape ``[B, N, T_q, T_kv]``. If an additive mask is given its type must promote to the promoted type of ``q``, ``k``, and ``v``. + The ``"causal"`` mask uses lower-right alignment where the + last query aligns with the last key. sinks (array, optional): An optional array of attention sinks. Default: ``None``. diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index d7ad6070..715e4881 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -23,11 +23,12 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None, sinks=None): v = mx.expand_dims(v, 2) scores = q @ mx.swapaxes(k, -1, -2) + is_causal = mask == "causal" if mask is not None: - if mask == "causal": - q_offset = max(0, kL - L) - q_indices = mx.arange(q_offset, q_offset + L) + if is_causal: + offset = kL - L + q_indices = mx.arange(L) + offset k_indices = mx.arange(kL) mask = q_indices[:, None] >= k_indices[None] @@ -58,7 +59,6 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None, sinks=None): out = scores @ v if n_repeats > 1: out = mx.reshape(out, [B, n_q_heads, L, -1]) - return out @@ -104,11 +104,14 @@ def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype): # SDPA for MHA (n_heads == n_kv_heads) def mlx_primitives_sdpa(q, k, v, scale, mask=None): p = (q * scale) @ k.transpose(0, 1, 3, 2) + qL = q.shape[2] + kL = k.shape[2] + is_causal = mask == "causal" if mask is not None: - if mask == "causal": - q_offset = max(0, k.shape[2] - q.shape[2]) - q_indices = mx.arange(q_offset, q_offset + q.shape[2]) - k_indices = mx.arange(k.shape[2]) + if is_causal: + offset = kL - qL + q_indices = mx.arange(qL) + offset + k_indices = mx.arange(kL) mask = q_indices[:, None] >= k_indices[None] p = mx.where(mask, p, mx.finfo(mx.float32).min) elif mask.dtype == mx.bool_: @@ -613,6 +616,17 @@ class TestSDPA(mlx_tests.MLXTestCase): 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:, :] + atol = 2e-5 if dtype == "float32" else 3e-4 self.assertListEqual(