Leave SDPA calls the Triton kernel cannot take to the regular CUDA lowering (pytorch#23349)

## What is wrong today

When the CUDA backend lowers a model, `ReplaceEdgeOpWithTritonOpPass`
swaps every `scaled_dot_product_attention` call for a Triton SDPA
kernel: the standard one, or split-K for long-cache decode. But those
kernels reject many calls that PyTorch accepts, and a rejected call
fails the whole export. For example:

```
RuntimeError: attn_mask must have dtype torch.bool
RuntimeError: Expected bfloat16 inputs
RuntimeError: attn_mask shape mismatch: expected [B=1, 1, L_q=8, L_kv=8], got torch.Size([1, 1, 1, 8])
RuntimeError: Head dimension must match across Q, K, V; got D_q=16, D_k=16, D_v=32.
RuntimeError: Causal masking requires L_q == L_kv; got L_q=4, L_kv=8.
```

The same happens with dropout, or a key/value batch or head count that
only broadcasts. The only way around it today is
`triton_kernel_mode="OFF"`, which also gives up the Triton kernels for
the calls they do support.

## What this change does

The pass still picks the kernel as before. It now runs that kernel's own
input checks on the fake tensors first (`_validate_sdpa_inputs`,
`_validate_qkv_shapes` and `_prepare_mask_params`, plus the standard
kernel's causal length rule), and replaces the call only if they pass.
Any other call stays in the graph, and the regular AOTInductor lowering
compiles it, as it does with `triton_kernel_mode="OFF"`. The pass logs
the kernel's own reason for each skipped call.

Because the pass asks the kernels instead of copying their rules, the
two cannot drift apart. Calls that were replaced before are still
replaced, with the same kernel. Masked decode and causal decode still
use split-K.

## What was tested

- New export tests in `test_sdpa_splitk_replacement.py` for a float
mask, float32 inputs, a broadcast mask, a batch-1 mask under batch 2,
dropout, a different value head size, a key/value batch of 1 under batch
2, one shared key/value head without `enable_gqa`, and causal attention
with different lengths. Each export now finishes and replaces no node;
each fails before this change. Exports with a full boolean mask still
replace the call, and masked decode, causal decode and decode with one
shared key/value head still use split-K on CUDA. Hand-built graphs cover
2D inputs and query, key and value passed by name.
- Removing any one of the kernel checks from the pass makes the matching
tests fail.
- The existing split-K replacement tests give the same result before and
after.
- Ran on a Linux aarch64 machine with a CUDA GPU.

Author: shoumikhin

backends/cuda/cuda_backend.py

@@ -769,7 +769,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
         Return CUDA-specific passes: ReplaceEdgeOpWithTritonOpPass.
 
         The Triton kernel replacement behavior can be controlled via compile_specs:
-        - triton_kernel_mode="ON": Always use Triton kernels
+        - triton_kernel_mode="ON": Use Triton kernels for the calls they support; other calls use the regular lowering
         - triton_kernel_mode="OFF": Never use Triton kernels and fallback to other implementations like cuda or decomposed operator.
         """
         # Parse compile_specs for triton_kernel_mode
@@ -929,8 +929,7 @@ def _combined():
                 # Force any remaining PyTorch SDPA ops to use the MATH
                 # backend during compilation so AOTI can lower / decompose
                 # them. SDPA ops already replaced by Triton kernels via
-                # `ReplaceEdgeOpWithTritonOpPass` are unaffected; this is
-                # only the fallback for the `triton_kernel_mode="OFF"` path.
+                # `ReplaceEdgeOpWithTritonOpPass` are unaffected.
                 stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH]))
                 stack.enter_context(target_smem_context())
                 if low_memory_mode == "ON":

backends/cuda/tests/test_sdpa_splitk_replacement.py

@@ -4,13 +4,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Test ReplaceEdgeOpWithTritonOpPass split-K SDPA kernel selection.
+"""Test which SDPA calls ReplaceEdgeOpWithTritonOpPass replaces, and with which kernel.
 
-Exports a minimal model containing F.scaled_dot_product_attention through the
-CUDA backend and verifies that CUDA routes eligible decode shapes to split-K,
-while ROCm and other shapes use standard SDPA.
+Exports minimal models containing F.scaled_dot_product_attention through the
+CUDA backend. CUDA routes eligible decode shapes to split-K, ROCm and other
+shapes use standard SDPA, and calls neither kernel accepts stay with the
+regular lowering.
 """
 
+import importlib.util
 import logging
 import unittest
 
@@ -24,6 +26,11 @@ def _require_cuda(tc: unittest.TestCase) -> None:
         tc.skipTest("CUDA required")
 
 
+def _require_splitk(tc: unittest.TestCase) -> None:
+    if torch.version.hip is not None:
+        tc.skipTest("split-K is off on ROCm")
+
+
 class SDPAModule(nn.Module):
     """Single-layer model with SDPA and a static KV cache buffer."""
 
@@ -199,5 +206,154 @@ def test_non_pow2_head_dim_uses_standard(self):
         self.assertEqual(len(splitk), 0, f"Expected no split-K for D=96. Got: {splitk}")
 
 
+class MaskedSDPAModule(nn.Module):
+    def __init__(self, is_causal=False, dropout_p=0.0):
+        super().__init__()
+        self.is_causal = is_causal
+        self.dropout_p = dropout_p
+
+    def forward(self, q, k, v, mask=None):
+        return F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=mask,
+            dropout_p=self.dropout_p,
+            is_causal=self.is_causal,
+        )
+
+
+def _bf16(*shape):
+    return torch.randn(*shape, dtype=torch.bfloat16)
+
+
+class TestUnsupportedSDPAStaysWithRegularLowering(unittest.TestCase):
+    """SDPA calls the Triton kernels do not accept are left to the regular lowering."""
+
+    def setUp(self):
+        _require_cuda(self)
+
+    def _logs(self, q, k, v, mask=None, **kwargs):
+        model = MaskedSDPAModule(**kwargs)
+        args = (q, k, v) if mask is None else (q, k, v, mask)
+        _, msgs = _capture_pass_logs(lambda: _export_through_cuda_backend(model, args))
+        return msgs
+
+    def assertStays(self, msgs, reason):
+        self.assertTrue(any(reason in m for m in msgs), msgs)
+        self.assertTrue(any("Replaced 0 nodes" in m for m in msgs), msgs)
+
+    def test_float_mask_is_not_replaced(self):
+        q = _bf16(1, 2, 8, 16)
+        mask = torch.zeros(1, 1, 8, 8, dtype=torch.bfloat16)
+        self.assertStays(self._logs(q, q, q, mask), "attn_mask must have dtype")
+
+    def test_float32_inputs_are_not_replaced(self):
+        q = torch.randn(1, 2, 8, 16)
+        self.assertStays(self._logs(q, q, q), "Expected bfloat16 inputs")
+
+    def test_broadcast_bool_mask_is_not_replaced(self):
+        q = _bf16(1, 2, 8, 16)
+        mask = torch.ones(1, 1, 1, 8, dtype=torch.bool)
+        self.assertStays(self._logs(q, q, q, mask), "attn_mask shape mismatch")
+
+    def test_batch_1_mask_under_batch_2_is_not_replaced(self):
+        q = _bf16(2, 2, 8, 16)
+        mask = torch.ones(1, 1, 8, 8, dtype=torch.bool)
+        self.assertStays(self._logs(q, q, q, mask), "attn_mask shape mismatch")
+
+    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_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")
+
+    def test_shared_key_value_head_without_gqa_is_not_replaced(self):
+        q, kv = _bf16(1, 4, 8, 16), _bf16(1, 1, 8, 16)
+        self.assertStays(self._logs(q, kv, kv), "Head counts must match")
+
+    def test_causal_with_other_lengths_is_not_replaced(self):
+        q, kv = _bf16(1, 2, 4, 16), _bf16(1, 2, 8, 16)
+        self.assertStays(self._logs(q, kv, kv, is_causal=True), "Causal masking")
+
+    def test_dropout_is_not_replaced(self):
+        q = _bf16(1, 2, 8, 16)
+        self.assertStays(self._logs(q, q, q, dropout_p=0.1), "dropout_p must be 0.0")
+
+    def test_bool_mask_is_still_replaced(self):
+        q = _bf16(1, 2, 8, 16)
+        mask = torch.ones(1, 1, 8, 8, dtype=torch.bool)
+        msgs = self._logs(q, q, q, mask)
+        self.assertTrue(any("Replaced 1 nodes" in m for m in msgs), msgs)
+
+    def assertSplitK(self, msgs):
+        self.assertTrue(any("Using split-K decode SDPA" in m for m in msgs), msgs)
+        self.assertTrue(any("Replaced 1 nodes" in m for m in msgs), msgs)
+
+    def test_masked_decode_still_uses_splitk(self):
+        _require_splitk(self)
+        q, kv = _bf16(1, 4, 1, 64), _bf16(1, 4, 512, 64)
+        mask = torch.ones(1, 1, 1, 512, dtype=torch.bool)
+        self.assertSplitK(self._logs(q, kv, kv, mask))
+
+    def test_causal_decode_still_uses_splitk(self):
+        _require_splitk(self)
+        q, kv = _bf16(1, 4, 1, 64), _bf16(1, 4, 512, 64)
+        self.assertSplitK(self._logs(q, kv, kv, is_causal=True))
+
+    def test_decode_with_a_shared_key_value_head_still_uses_splitk(self):
+        _require_splitk(self)
+        q, kv = _bf16(1, 4, 1, 64), _bf16(1, 1, 512, 64)
+        self.assertSplitK(self._logs(q, kv, kv))
+
+
+class TestSDPAKernelSupportCheck(unittest.TestCase):
+    """Calls that cannot be exported end to end, on hand-built graphs."""
+
+    def setUp(self):
+        if importlib.util.find_spec("triton") is None:
+            self.skipTest("Triton required")
+
+    def _sdpa_node(self, shape, dropout_p=0.0):
+        from executorch.exir.dialects._ops import ops as exir_ops
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        graph = torch.fx.Graph()
+        with FakeTensorMode():
+            q = torch.empty(*shape, dtype=torch.bfloat16, device="cuda")
+        inputs = []
+        for name in ("q", "k", "v"):
+            node = graph.placeholder(name)
+            node.meta["val"] = q
+            inputs.append(node)
+        return graph.call_function(
+            exir_ops.edge.aten.scaled_dot_product_attention.default,
+            (*inputs, None, dropout_p),
+        )
+
+    def test_2d_inputs_are_not_supported(self):
+        # The regular lowering of a 2D SDPA fails later in Inductor, so this
+        # case cannot be exported end to end.
+        from executorch.backends.cuda.triton.replacement_pass import (
+            ReplaceEdgeOpWithTritonOpPass,
+        )
+
+        node = self._sdpa_node((8, 16))
+        self.assertFalse(ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node))
+
+    def test_named_inputs_are_not_supported(self):
+        from executorch.backends.cuda.triton.replacement_pass import (
+            ReplaceEdgeOpWithTritonOpPass,
+        )
+
+        node = self._sdpa_node((1, 2, 8, 16))
+        query, key, value = node.args[:3]
+        node.args = ()
+        node.kwargs = {"query": query, "key": key, "value": value}
+        self.assertFalse(ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node))
+
+
 if __name__ == "__main__":
     unittest.main()

backends/cuda/triton/replacement_pass.py

@@ -13,6 +13,11 @@
 import logging
 
 import torch
+from executorch.backends.cuda.triton.kernels.sdpa import (
+    _prepare_mask_params,
+    _validate_qkv_shapes,
+    _validate_sdpa_inputs,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 
 from torch.fx import GraphModule, Node
@@ -114,7 +119,6 @@ def _pick_sdpa_kernel(node: Node):
             and D > 0
             and (D & (D - 1)) == 0  # power of 2
         ):
-            logger.info(f"Using split-K decode SDPA (L_kv={L_kv}, D={D})")
             return triton.sdpa_decode_splitk
 
         return triton.sdpa
@@ -145,6 +149,44 @@ def _should_replace_node(self, node: Node) -> bool:
                 logger.info(f"Skipping topk replacement: N={N} > {self._TOPK_MAX_N}")
                 return False
 
+        if node.target == exir_ops.edge.aten.scaled_dot_product_attention.default:
+            return self._sdpa_kernel_supports(node)
+
+        return True
+
+    @classmethod
+    def _sdpa_kernel_supports(cls, node: Node) -> bool:
+        """Whether the Triton SDPA kernel the pass would pick accepts this call.
+
+        Runs the kernels' own input checks on the fake tensors, so calls they
+        reject stay with the regular lowering instead of failing the export.
+        """
+        # The kernel picker reads query and key by position.
+        if len(node.args) < 3:
+            return False
+        names = ("attn_mask", "dropout_p", "is_causal", "scale", "enable_gqa")
+        defaults = (None, 0.0, False, None, False)
+        attn_mask, dropout_p, is_causal, _, enable_gqa = (
+            node.args[3 + i] if 3 + i < len(node.args) else node.kwargs.get(name, d)
+            for i, (name, d) in enumerate(zip(names, defaults))
+        )
+        query, key, value = (arg.meta["val"] for arg in node.args[:3])
+        try:
+            _validate_sdpa_inputs(query, key, value, dropout_p, enable_gqa)
+            splitk = cls._pick_sdpa_kernel(node) is triton.sdpa_decode_splitk
+            # Split-K groups query heads over key/value heads on its own.
+            B, _, _, L_q, L_kv, _, _ = _validate_qkv_shapes(
+                query, key, value, enable_gqa or splitk
+            )
+            if is_causal and L_q != L_kv and not splitk:
+                raise RuntimeError(
+                    f"Causal masking requires L_q == L_kv; got L_q={L_q}, L_kv={L_kv}."
+                )
+            mask = None if attn_mask is None else attn_mask.meta["val"]
+            _prepare_mask_params(mask, B, L_q, L_kv)
+        except RuntimeError as e:
+            logger.info(f"Skipping SDPA replacement: {e}")
+            return False
         return True
 
     def _replace_node_with_triton(self, graph_module: GraphModule, node: Node) -> None:
@@ -166,6 +208,9 @@ def _replace_node_with_triton(self, graph_module: GraphModule, node: Node) -> No
 
         if target == exir_ops.edge.aten.scaled_dot_product_attention.default:
             triton_kernel_fn = self._pick_sdpa_kernel(node)
+            if triton_kernel_fn is triton.sdpa_decode_splitk:
+                L_kv, D = node.args[1].meta["val"].shape[2:]
+                logger.info(f"Using split-K decode SDPA (L_kv={L_kv}, D={D})")
 
         # Create a new node with the Triton kernel
         with graph_module.graph.inserting_before(node):