Read the query length at build time where it is known (pytorch#22494)

Follow-up polish on the causal slicing added in pytorch#22443. No behaviour
change for any shape that worked before.

### Read the query length at build time where it is known

The slice that shortens the keys for a causal call always emitted a node
to read the query length at run time. For a decode step that length is a
constant known when the program is built, so the node was pure overhead
on exactly the path this slicing exists to speed up.

There is already a helper for this, used by the cached attention handler
a few hundred lines down in the same file: it returns a literal for a
concrete dimension and emits a size node only for a symbolic one.

| Call | Before | After |
| --- | --- | --- |
| query 1 against 32 keys, static | size node, 2336 bytes | no size
node, **2288 bytes** |
| query 6 against 32 keys, static | size node, 2336 bytes | no size
node, **2288 bytes** |
| dynamic query length | size node | size node, unchanged |

### Two comments promised less than the code delivered

The guard that declines causal attention also declines a query length
that **cannot be compared** at build time, two unrelated dynamic
dimensions for instance, and the comment above it described only the
case where the query is genuinely longer. It now says both, because the
second is the case a reader is more likely to meet.

### A half-wired test knob

The key length knob added to the attention test reached the query, key
and value but not either mask branch, so a case combining a mask with an
unequal key length built a mask of the wrong width. No shipped case does
that today, so nothing was failing, but the next person to use it would
have hit it.

### Test plan

Causal attention at every shape the operator suite covers, run against
eager on an Apple Silicon Mac: query 1, 2 and 6 against longer keys,
equal lengths, non-causal, float32 and bfloat16. All match to 4.8e-07 or
exactly, 8 of 8 cases.

A dynamic query length bounded by the keys keeps its size node, exports,
and matches eager to 1.2e-07.

Building the test mask both ways in eager torch: with the key length it
runs, and with the query length it raises the size mismatch this
removes.

Exported the speech model: unchanged, 8 fused attention nodes and no
size nodes.

Ran the neighbouring backend test files, all passing.

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

Author: shoumikhin

backends/mlx/patterns.py

@@ -21,6 +21,7 @@
 from executorch.backends.mlx.builder.op_helpers import (
     emit_quantized_biases,
     emit_quantized_gather,
+    emit_shape,
     emit_stop_position,
     mlx_qparams_supported,
     parse_dequant_int4_node,
@@ -583,9 +584,11 @@ 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.
+        # Causal attention is only taken when the query is known to be no longer than
+        # the keys. Torch clamps a longer query to the keys that exist, and neither the
+        # kernel's own mask nor slicing the keys reproduces that. A query length that
+        # cannot be compared at build time, two unrelated dynamic dimensions for
+        # instance, is also declined, because the relation has to hold for every call.
         if is_causal and not statically_known_true(
             operand_vals[0].shape[-2] <= operand_vals[1].shape[-2]
         ):
@@ -720,10 +723,9 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
         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))
-            )
+            # A literal when the query length is known at build time, which is the
+            # decode case this exists for, and a size node only when it is symbolic.
+            rows = emit_shape(P, self.q_node, inputs[0])[-2]
             for i in (1, 2):
                 _, sliced = P.make_tmp_slot()
                 P.emit(
@@ -732,7 +734,7 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
                         out=P.slot_to_tid(sliced),
                         axis=IntOrVid.from_literal(2),
                         start=IntOrVid.from_literal(0),
-                        stop=P.to_int_or_vid(rows),
+                        stop=rows,
                     )
                 )
                 inputs[i] = sliced

backends/mlx/test/test_ops.py

@@ -6270,16 +6270,16 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
 
         if self.use_mask:
             # Additive float mask: 0 = attend, -inf = masked
-            mask = torch.zeros(self.batch_size, 1, self.seq_len, self.seq_len)
-            mask[:, :, :, : self.seq_len // 4] = float("-inf")
+            mask = torch.zeros(self.batch_size, 1, self.seq_len, self.kv_seq_len)
+            mask[:, :, :, : self.kv_seq_len // 4] = float("-inf")
             return (q, k, v, mask)
         elif self.use_bool_mask:
             # Boolean mask: True = attend, False = masked
             # This tests that the backend correctly converts bool -> additive format
             mask = torch.ones(
-                self.batch_size, 1, self.seq_len, self.seq_len, dtype=torch.bool
+                self.batch_size, 1, self.seq_len, self.kv_seq_len, dtype=torch.bool
             )
-            mask[:, :, :, : self.seq_len // 4] = False  # Mask out first quarter
+            mask[:, :, :, : self.kv_seq_len // 4] = False  # Mask out first quarter
             return (q, k, v, mask)
         return (q, k, v)