Use lower-right causal mask alignment consistently (#2967)

Co-authored-by: Awni Hannun <[email protected]>

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));

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h

@@ -461,7 +461,7 @@ template <
   metal::vec<AccumType, kRowsPT> 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<MulOp>(rcp);

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);

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``.
 

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(