Give causal attention the mask PyTorch asks for (pytorch#22443)

Fixes pytorch#22426.

### The problem

Attention with `is_causal=True` returns wrong values, with nothing
raised, whenever the query is shorter than the keys:

| shape | error against eager |
| --- | --- |
| query 1, keys 16 | **2.7 to 3.4** |
| query 16, keys 16 | 3.6e-07 |

The two libraries anchor a causal mask at opposite corners. PyTorch puts
it at the top left, so the query at row `i` attends to keys `0..i`. MLX
puts it at the bottom right, aligning the last query with the last key.
They agree only when the lengths are equal, and nothing checked that.

This is the decode step of any model that generates a token at a time,
so it is not a corner case.

### The fix

Slice the keys and values to the query length when the lengths are not
provably equal, then let the kernel apply its own mask to what is now a
square problem.

That is what PyTorch computes: a row past the last query never reads a
later key, so dropping those keys changes nothing. Measured over query
lengths 1 to 63 against key lengths 2 to 512, the sliced form and the
original agree **exactly** in float32. A slice before attention is also
what the cache-aware path in this file already does, so no mask tensor
is built and the fused kernel is kept.

A causal query **longer** than its keys is left to decompose instead.
PyTorch clamps each row to the keys that exist, and neither the kernel's
own mask nor a slice reproduces that.

### Why the speech example changes too

This backend already knew about the conflict. `mlx::custom_sdpa` exists
for cached attention and takes a start position so the kernel gets a
square problem, and seven models here use it.

The speech example does not. It slices the cache itself and then asks
for a causal mask over the result, which in PyTorch means **the new
token sees only the first cached key**. It has been getting attention
over the whole cache because the kernel read its request the other way
round, so it depends on the behaviour this change corrects. Moving it
onto the same path as the other models makes it ask for what it wants.

Worth stating plainly: that example's own eager math was already 3.3
away from what it intended at a decode step, so this is not a regression
being introduced, it is one being removed.

### Test plan

Two cases added to the operator suite, a query of 1 and of 6 against 32
keys, which run the compiled runtime and compare against eager. **Both
fail before this change on numerics and pass after.**

Ran the shapes on an Apple Silicon Mac against eager:

- query 1, 2 and 6 against longer keys, float32 and bfloat16: now exact
or within 4.2e-07, previously off by around 3
- equal lengths and non-causal calls untouched, fused kernel kept in
every case, no mask tensor built
- a causal query longer than its keys decomposes and matches eager to
3.6e-07, previously off by 2.7
- equal-length causal attention serializes to a **byte-identical
program** before and after, so that path is bit-for-bit unchanged

Exported the speech model and compared its logits against HuggingFace's
own decoder at a cached decode step: unchanged from before this change,
and the top token still matches.

Ran the neighbouring backend test files, all passing.

Co-authored-by: PyTorch Bot <pytorchbot@users.noreply.github.com>

Author: shoumikhin

backends/mlx/examples/whisper/export_whisper.py

@@ -122,14 +122,17 @@ def forward(
         # Update KV cache
         k_cache, v_cache = self.kv_cache.update(pos_int, k, v)
 
-        # Explicit windowing: slice cache to valid positions
-        end_pos = pos_int + T
-        k_win = k_cache[:, :, :end_pos, :]
-        v_win = v_cache[:, :, :end_pos, :]
-
-        # SDPA with causal mask
-        attn_out = F.scaled_dot_product_attention(
-            q, k_win, v_win, attn_mask=None, is_causal=True, scale=self.scale
+        # The cache-aware op slices the cache to start_pos + query length itself, and
+        # applies the causal mask the kernel wants for a cached decode step. Passing
+        # a hand-sliced window with is_causal instead means something different: torch
+        # would let the new token see only the first cached key.
+        attn_out = torch.ops.mlx.custom_sdpa(
+            q,
+            k_cache,
+            v_cache,
+            start_pos=pos_int,
+            is_causal=True,
+            scale=self.scale,
         )
 
         # Reshape back

backends/mlx/patterns.py

@@ -583,6 +583,14 @@ def _kernel_can_compute(cls, sdpa_node: Node) -> bool:
             if mask_val is None or mask_val.dim() > 4:
                 return False
 
+        # A causal query longer than its keys is left to decompose. Torch clamps each
+        # row to the keys that exist, and neither the kernel's own mask nor slicing the
+        # keys to the query length reproduces that.
+        if is_causal and not statically_known_true(
+            operand_vals[0].shape[-2] <= operand_vals[1].shape[-2]
+        ):
+            return False
+
         return True
 
     @classmethod
@@ -704,6 +712,31 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
         if output_rank < 4:
             _, sdpa_out = P.make_tmp_slot()
 
+        # MLX anchors its causal mask at the bottom right and torch at the top left, so
+        # the flag alone only means the same thing when the lengths are equal. Slicing
+        # the keys and values to the query length makes the problem square, where the
+        # two conventions agree, and it is what torch computes: a query at row i attends
+        # to keys 0..i, so keys past the last query row are never read.
+        q_len = self.q_node.meta["val"].shape[-2]
+        k_len = self.k_node.meta["val"].shape[-2]
+        if is_causal and not statically_known_true(q_len == k_len):
+            _, rows = P.slot_manager.make_tmp_value_slot()
+            P.emit(
+                SymSizeNode(a=P.slot_to_tid(inputs[0]), dim=2, out=P.slot_to_vid(rows))
+            )
+            for i in (1, 2):
+                _, sliced = P.make_tmp_slot()
+                P.emit(
+                    SliceNode(
+                        x=P.slot_to_tid(inputs[i]),
+                        out=P.slot_to_tid(sliced),
+                        axis=IntOrVid.from_literal(2),
+                        start=IntOrVid.from_literal(0),
+                        stop=P.to_int_or_vid(rows),
+                    )
+                )
+                inputs[i] = sliced
+
         P.emit(
             SdpaNode(
                 q=P.slot_to_tid(inputs[0]),

backends/mlx/test/test_ops.py

@@ -6212,6 +6212,7 @@ def __init__(
         is_causal: bool = False,
         use_mask: bool = False,
         use_bool_mask: bool = False,
+        kv_seq_len: Optional[int] = None,
     ):
         self.batch_size = batch_size
         self.num_heads = num_heads
@@ -6221,12 +6222,15 @@ def __init__(
         self.is_causal = is_causal
         self.use_mask = use_mask
         self.use_bool_mask = use_bool_mask
+        self.kv_seq_len = kv_seq_len if kv_seq_len is not None else seq_len
 
         parts = ["sdpa"]
         if num_kv_heads is not None:
             parts.append(f"gqa{num_kv_heads}")
         if is_causal:
             parts.append("causal")
+        if self.kv_seq_len != seq_len:
+            parts.append(f"q{seq_len}kv{self.kv_seq_len}")
         if use_mask:
             parts.append("mask")
         if use_bool_mask:
@@ -6241,6 +6245,11 @@ def get_test_configs(cls) -> List["SDPATest"]:
             cls(num_kv_heads=4),
             cls(use_mask=True),
             cls(use_bool_mask=True),  # Test boolean mask conversion
+            # A decode step against a longer key cache. MLX anchors its causal mask at
+            # the bottom right and torch at the top left, so they only agree when the
+            # lengths match.
+            cls(is_causal=True, seq_len=1, kv_seq_len=32),
+            cls(is_causal=True, seq_len=6, kv_seq_len=32),
         ]
 
     def create_model(self) -> nn.Module:
@@ -6256,8 +6265,8 @@ def create_model(self) -> nn.Module:
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         q = torch.randn(self.batch_size, self.num_heads, self.seq_len, self.head_dim)
         kv_heads = self.num_kv_heads if self.num_kv_heads else self.num_heads
-        k = torch.randn(self.batch_size, kv_heads, self.seq_len, self.head_dim)
-        v = torch.randn(self.batch_size, kv_heads, self.seq_len, self.head_dim)
+        k = torch.randn(self.batch_size, kv_heads, self.kv_seq_len, self.head_dim)
+        v = torch.randn(self.batch_size, kv_heads, self.kv_seq_len, self.head_dim)
 
         if self.use_mask:
             # Additive float mask: 0 = attend, -inf = masked