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

backends/cuda/tests/test_cuda_export.py

@@ -683,6 +683,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

@@ -573,6 +573,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(
@@ -1273,7 +1280,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)