CUDA backend: keep wide attention heads that are not a power of 2 off the Triton SDPA kernel (#24427)

### What is wrong today

The CUDA backend replaces `scaled_dot_product_attention` with its own
Triton kernel when the kernel accepts the call. For head sizes that are
not a power of 2, that kernel holds the whole head in one tile of at
most 256 values. Wider heads of that kind, such as 288, 320, 384, 640 or
768, are still accepted. The kernel then uses only the first 256 values
of each query and key, and it never writes the rest of the output. The
program runs and returns wrong numbers, with no error. For example, the
mid-block attention of the LingBot-VA video VAE encoder has heads of
640, and its lowered encoder drifts far from eager.

```python
q, k, v = (torch.randn(2, 1, 64, 640, device="cuda", dtype=torch.bfloat16) for _ in range(3))
sdpa(q, k, v)  # max error against PyTorch: about 4
```

### What this change does

The shape check that the kernel and the replacement pass share now
rejects head sizes that are not a power of 2 and are wider than that one
tile. The pass already leaves calls the kernel rejects to the regular
CUDA lowering, so such attention in an exported model now gives the
right result. Code that calls the Triton kernel directly, such as the
off-graph KV lowering, now gets an error instead of wrong numbers. Heads
that are a power of 2, including 512 and 1024, use the other Triton
kernel, which covers the whole head, so they still use Triton.

Because this fixes wrong results, it should also go to `release/1.6`.

### What was tested

On an A100 with torch 2.14:

- `test_sdpa_head_sizes` lowers attention with heads of 96, 192, 256,
512, 288, 320, 384, 640 and 768, runs it, and compares each result with
eager. Without this change the last five fail.
- `test_heads_the_kernel_takes_are_supported` checks that the pass still
gives heads of 192, 256, 512 and 1024 to the Triton kernel.
`test_wide_head_that_is_not_a_power_of_2_is_not_replaced` checks that a
head of 640 exports through the regular lowering and logs why.
`test_wide_non_pow2_head_rejected` checks that a direct call with a head
of 640 raises.
- The other Triton SDPA kernel tests and
`test_sdpa_splitk_replacement.py` still pass.
- At 64 tokens and head 640, one lowered attention call takes about as
long as before: 767 µs now, 784 µs with the wrong result. This was
measured at 64 tokens only. The regular lowering builds the full score
matrix, so for long sequences it can need much more time and memory than
the Triton kernel did.
- With torch 2.15, the LingBot-VA VAE encoder, lowered with default
settings, now matches eager within bfloat16 precision.

Author: shoumikhin

backends/cuda/tests/test_cuda_export.py

@@ -815,6 +815,43 @@ def test_sdpa_single_kernel(self):
             "SDPA single kernel operation export failed",
         )
 
+    @unittest.skipIf(
+        torch.version.hip is not None,
+        "loads through the Python runtime, which the ROCm build does not include",
+    )
+    def test_sdpa_head_sizes(self):
+        from executorch.runtime import Runtime
+
+        class Attention(torch.nn.Module):
+            def forward(self, *qkv):
+                return tuple(
+                    torch.nn.functional.scaled_dot_product_attention(*qkv[i : i + 3])
+                    for i in range(0, len(qkv), 3)
+                )
+
+        # Heads that are not a power of 2 and are above 256 are wider than one
+        # tile of the Triton kernel for such heads.
+        head_sizes = (96, 192, 256, 512, 288, 320, 384, 640, 768)
+        inputs = tuple(
+            torch.randn(1, 2, 64, size, dtype=torch.bfloat16)
+            for size in head_sizes
+            for _ in range(3)
+        )
+        model = Attention()
+        edge = self._export_to_cuda_with_lower(model, inputs)
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "model.pte")
+            with open(path, "wb") as f:
+                edge.to_executorch().write_to_file(f)
+            method = Runtime.get().load_program(path).load_method("forward")
+            outputs = method.execute(list(inputs))
+        expected = model(*(x.cuda() for x in inputs))
+        for size, output, reference in zip(head_sizes, outputs, expected, strict=True):
+            with self.subTest(head_size=size):
+                torch.testing.assert_close(
+                    output.cpu(), reference.cpu(), atol=2e-2, rtol=2e-2
+                )
+
     def test_triton_kernel_mode_off(self):
         """
         Test CUDA export with triton_kernel_mode set to OFF for SDPA kernel.

backends/cuda/tests/test_sdpa_splitk_replacement.py

@@ -275,6 +275,10 @@ def test_value_head_size_other_than_query_is_not_replaced(self):
         q, v = _bf16(1, 2, 1, 16), _bf16(1, 2, 256, 32)
         self.assertStays(self._logs(q, _bf16(1, 2, 256, 16), v), "Head dimension")
 
+    def test_wide_head_that_is_not_a_power_of_2_is_not_replaced(self):
+        q = _bf16(1, 2, 64, 640)
+        self.assertStays(self._logs(q, q, q), "not a power of 2")
+
     def test_key_value_batch_1_under_batch_2_is_not_replaced(self):
         q, kv = _bf16(2, 2, 1, 16), _bf16(1, 2, 256, 16)
         self.assertStays(self._logs(q, kv, kv), "Batch dimension must match")
@@ -320,7 +324,7 @@ def test_decode_with_a_shared_key_value_head_still_uses_splitk(self):
 
 
 class TestSDPAKernelSupportCheck(unittest.TestCase):
-    """Calls that cannot be exported end to end, on hand-built graphs."""
+    """Support checks on hand-built graphs, which compile nothing."""
 
     def setUp(self):
         if importlib.util.find_spec("triton") is None:
@@ -364,6 +368,19 @@ def test_named_inputs_are_not_supported(self):
         node.kwargs = {"query": query, "key": key, "value": value}
         self.assertFalse(ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node))
 
+    def test_heads_the_kernel_takes_are_supported(self):
+        # Some CI GPUs lack the shared memory to compile heads of 1024.
+        from executorch.backends.cuda.triton.replacement_pass import (
+            ReplaceEdgeOpWithTritonOpPass,
+        )
+
+        for size in (192, 256, 512, 1024):
+            with self.subTest(head_size=size):
+                node = self._sdpa_node((1, 2, 64, size))
+                self.assertTrue(
+                    ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node)
+                )
+
 
 if __name__ == "__main__":
     unittest.main()

backends/cuda/tests/test_triton_sdpa.py

@@ -603,6 +603,12 @@ def test_causal_lq_ne_lkv_rejected(self):
         with self.assertRaises(RuntimeError):
             self.sdpa(q, k, v, is_causal=True)
 
+    def test_wide_non_pow2_head_rejected(self):
+        """A non-pow2 head wider than one tile should raise, not be cut short."""
+        q = torch.randn(1, 2, 64, 640, dtype=torch.bfloat16, device="cuda")
+        with self.assertRaises(RuntimeError):
+            self.sdpa(q, q, q)
+
     def test_non_pow2_no_mask(self):
         """Non-pow2 head dim without mask should work (mask_ptr=0 path)."""
         B, H, Lq, Lk, D = 1, 4, 4, 64, 40  # D=40 is not pow2

backends/cuda/triton/kernels/sdpa.py

@@ -270,6 +270,13 @@ def _validate_qkv_shapes(
         raise RuntimeError(
             f"Head dimension must match across Q, K, V; got D_q={D_q}, D_k={D_k}, D_v={D_v}."
         )
+    # Heads that are not a power of 2 must fit in one tile, and
+    # _next_power_of_2 caps the tile at 256.
+    if not _is_power_of_2(D_q) and _next_power_of_2(D_q) < D_q:
+        raise RuntimeError(
+            "Head dimensions that are not a power of 2 must be at most "
+            f"{_next_power_of_2(D_q)}; got D={D_q}."
+        )
     # Key and Value sequence lengths must match
     if L_kv_k != L_kv_v:
         raise RuntimeError(
@@ -1301,7 +1308,8 @@ def sdpa(
     share K/V loads) based on a tile-utilization heuristic from FlashAttention.
 
     Args:
-        query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16
+        query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16. A D that
+            is not a power of 2 must be at most 256.
         key: Key tensor [B, H_kv, L_kv, D], dtype torch.bfloat16
         value: Value tensor [B, H_kv, L_kv, D], dtype torch.bfloat16
         attn_mask: Optional bool mask [B, 1, L_q, L_kv] (broadcast over heads)