CUDA backend: lower fused attention variants

backends/cuda/BUCK

@@ -79,6 +79,7 @@ fbcode_target(
         "passes/__init__.py",
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
+        "passes/replace_fused_sdpa.py",
         "passes/replace_int64_floordiv.py",
     ],
     visibility = [

backends/cuda/cuda_backend.py

@@ -39,6 +39,9 @@
 from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
     MoveCondPredicateToCpuPass,
 )
+from executorch.backends.cuda.passes.replace_fused_sdpa import (
+    ReplaceFusedSdpaWithSdpaPass,
+)
 from executorch.backends.cuda.passes.replace_int64_floordiv import (
     ReplaceInt64FloorDivWithFloatPass,
 )
@@ -1013,7 +1016,11 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                         f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
                 triton_kernel_mode = mode
-        passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
+        passes = [
+            MoveCondPredicateToCpuPass(),
+            ReplaceInt64FloorDivWithFloatPass(),
+            ReplaceFusedSdpaWithSdpaPass(),
+        ]
         cell_layout = any(
             spec.key == OFFGRAPH_KV_COMPILE_SPEC
             and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"

backends/cuda/passes/replace_fused_sdpa.py

@@ -0,0 +1,136 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Graph Transformation Pass for Fused Attention Replacement.
+
+Decomposing ``scaled_dot_product_attention`` on CUDA, as ``torch.export``'s
+default decompositions and Torch-TensorRT do, turns it into one of its fused
+CUDA variants. The backend has no kernel for those, so this pass turns them
+back into ``scaled_dot_product_attention``, which the Triton SDPA kernel or
+Inductor's math decomposition lowers.
+"""
+
+import logging
+import operator
+from typing import Any, Dict, Optional
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import ExportedProgram
+from torch.export.graph_signature import ExportGraphSignature
+from torch.fx import GraphModule, Node
+from torch.fx.experimental.symbolic_shapes import statically_known_true
+
+logger = logging.getLogger(__name__)
+
+_FUSED_SDPA_OPS = (
+    exir_ops.edge.aten._scaled_dot_product_flash_attention.default,
+    exir_ops.edge.aten._scaled_dot_product_efficient_attention.default,
+    exir_ops.edge.aten._scaled_dot_product_cudnn_attention.default,
+)
+
+
+class ReplaceFusedSdpaWithSdpaPass:
+    """
+    Replace each fused SDPA node with ``scaled_dot_product_attention`` when
+    only its attention output is used and ``scaled_dot_product_attention``
+    computes the same result. Any other node stays as it is.
+    """
+
+    requires_exported_program = True
+
+    def __call__(self, exported_program: ExportedProgram) -> None:
+        program_module = exported_program.graph_module
+        for module in program_module.modules():
+            if not isinstance(module, GraphModule):
+                continue
+            # The signature names the outputs of the program's own graph only.
+            signature = (
+                exported_program.graph_signature if module is program_module else None
+            )
+            modified = False
+            for node in list(module.graph.nodes):
+                if node.op != "call_function" or node.target not in _FUSED_SDPA_OPS:
+                    continue
+                values = _arguments(node)
+                reason = _reason_to_keep(node, values)
+                if reason is None:
+                    _replace_with_sdpa(module, node, values, signature)
+                    modified = True
+                else:
+                    logger.warning(
+                        "The CUDA backend cannot lower %s, since %s",
+                        node.format_node(),
+                        reason,
+                    )
+            if modified:
+                module.recompile()
+
+
+def _arguments(node: Node) -> Dict[str, Any]:
+    arguments = node.target._schema.arguments
+    values = {a.name: a.default_value for a in arguments}
+    values.update(zip((a.name for a in arguments), node.args))
+    values.update(node.kwargs)
+    return values
+
+
+def _reason_to_keep(node: Node, values: Dict[str, Any]) -> Optional[str]:
+    if any(
+        user.target is not operator.getitem or user.args[1] != 0 for user in node.users
+    ):
+        return "an output other than the attention output is used"
+    if not values["is_causal"]:
+        return None
+    if values.get("attn_bias") is not None:
+        return "scaled_dot_product_attention takes no mask together with is_causal"
+    # Flash attention aligns the causal mask to the bottom right, the other
+    # variants and scaled_dot_product_attention to the top left, and the Triton
+    # decode kernel ignores it.
+    query, key = values["query"].meta["val"], values["key"].meta["val"]
+    if not statically_known_true(query.shape[-2] == key.shape[-2]):
+        return "its causal mask may have different query and key lengths"
+    return None
+
+
+def _replace_with_sdpa(
+    module: GraphModule,
+    node: Node,
+    values: Dict[str, Any],
+    signature: Optional[ExportGraphSignature],
+) -> None:
+    query, key = values["query"], values["key"]
+    kwargs = {}
+    # The Triton kernels take a float scale, so leave out the default.
+    if values["scale"] is not None:
+        kwargs["scale"] = values["scale"]
+    # The fused variants that accept fewer key heads than query heads group
+    # them as enable_gqa does.
+    if not statically_known_true(
+        query.meta["val"].shape[-3] == key.meta["val"].shape[-3]
+    ):
+        kwargs["enable_gqa"] = True
+    with module.graph.inserting_before(node):
+        sdpa = module.graph.call_function(
+            exir_ops.edge.aten.scaled_dot_product_attention.default,
+            args=(
+                query,
+                key,
+                values["value"],
+                # Only the flash variant has no bias.
+                values.get("attn_bias"),
+                values["dropout_p"],
+                values["is_causal"],
+            ),
+            kwargs=kwargs,
+        )
+    sdpa.meta = {**node.meta, "val": node.meta["val"][0]}
+    for user in list(node.users):
+        user.replace_all_uses_with(sdpa)
+        if signature is not None:
+            signature.replace_all_uses(user.name, sdpa.name)
+        module.graph.erase_node(user)
+    module.graph.erase_node(node)

backends/cuda/passes/tests/test_replace_fused_sdpa.py

@@ -0,0 +1,124 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.aoti.passes.replace_view_copy_with_view import (
+    ReplaceViewCopyWithViewPass,
+)
+from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+from executorch.backends.cuda.passes.replace_fused_sdpa import (
+    ReplaceFusedSdpaWithSdpaPass,
+)
+from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import export
+
+
+def _program_for_cuda_passes(model, inputs, dynamic_shapes=None):
+    """The Edge program of ``model`` as the CUDA backend's passes get it: no
+    op decomposed, as the CUDA partitioner asks, and views for copies."""
+    program = export(model, inputs, dynamic_shapes=dynamic_shapes, strict=True)
+    # to_edge_transform_and_lower functionalizes the program before it asks
+    # the partitioner which ops to keep.
+    program = program.run_decompositions({})
+    ops, _ = CudaPartitioner([]).ops_to_not_decompose(program)
+    program = to_edge(
+        program,
+        compile_config=EdgeCompileConfig(_check_ir_validity=False, preserve_ops=ops),
+    ).exported_program()
+    ReplaceViewCopyWithViewPass()(program.graph_module)
+    return program
+
+
+class _Fn(torch.nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, *args):
+        return self.fn(*args)
+
+
+class TestReplaceFusedSdpaWithSdpaPass(unittest.TestCase):
+    def test_runs_before_triton_replacement(self):
+        """The calls it rewrites must still reach the Triton SDPA kernels."""
+        names = [type(p).__name__ for p in CudaBackend.get_custom_passes([])]
+        self.assertLess(
+            names.index("ReplaceFusedSdpaWithSdpaPass"),
+            names.index("ReplaceEdgeOpWithTritonOpPass"),
+        )
+
+    def test_rewrite(self):
+        """Fused attention becomes scaled_dot_product_attention only when that
+        computes the same result."""
+        aten = torch.ops.aten
+
+        def cudnn(q, k):
+            return aten._scaled_dot_product_cudnn_attention(q, k, k, None, False)[0]
+
+        def causal(q, k, bias=None):
+            return aten._scaled_dot_product_efficient_attention(
+                q, k, k, bias, False, is_causal=True
+            )[0]
+
+        def grouped(q, k):
+            return aten._scaled_dot_product_flash_attention(q, k, k, scale=0.2)[0]
+
+        def with_log_sum_exp(q, k):
+            return aten._scaled_dot_product_efficient_attention(q, k, k, None, True)[:2]
+
+        def flash(q, k):
+            return aten._scaled_dot_product_flash_attention(q, k, k)[0]
+
+        def in_branch(q, k):
+            return torch.cond(q.sum() > 0, flash, lambda q, k: q * 2, (q, k))
+
+        q = torch.randn(1, 2, 16, 64)
+        one_head = torch.randn(1, 1, 16, 64)
+        decode = q[:, :, :1].clone()
+        length = torch.export.Dim("length", min=2, max=64)
+        keys = torch.export.Dim("keys", min=2, max=64)
+        # The last item is the keyword arguments of the new call, or None when
+        # the fused call has to stay.
+        for fn, inputs, lengths, kwargs in (
+            (cudnn, (q, q), None, {}),
+            (causal, (q, q), None, {}),
+            (causal, (q, q.clone()), (length, length), {}),
+            (grouped, (q, one_head), None, {"scale": 0.2, "enable_gqa": True}),
+            (in_branch, (q, q), None, {}),
+            # With different query and key lengths, the variants and the Triton
+            # decode kernel place the causal mask differently.
+            (causal, (decode, q), None, None),
+            (causal, (q, q.clone()), (length, keys), None),
+            (causal, (q, q, torch.randn(1, 2, 16, 16)), None, None),
+            (with_log_sum_exp, (q, q), None, None),
+        ):
+            shapes = [t.shape for t in inputs]
+            with self.subTest(fn=fn.__name__, shapes=shapes, lengths=lengths):
+                dynamic_shapes = None
+                if lengths:
+                    dynamic_shapes = {"args": tuple({2: n} for n in lengths)}
+                program = _program_for_cuda_passes(_Fn(fn), inputs, dynamic_shapes)
+                ReplaceFusedSdpaWithSdpaPass()(program)
+                program.validate()
+                nodes = [
+                    node
+                    for module in program.graph_module.modules()
+                    if isinstance(module, torch.fx.GraphModule)
+                    for node in module.graph.nodes
+                ]
+                sdpa = exir_ops.edge.aten.scaled_dot_product_attention.default
+                calls = [node.kwargs for node in nodes if node.target == sdpa]
+                self.assertEqual(calls, [] if kwargs is None else [kwargs])
+                fused = [
+                    node
+                    for node in nodes
+                    if "aten._scaled_dot_product_" in str(node.target)
+                ]
+                self.assertEqual(not fused, kwargs is not None)

backends/cuda/tests/test_cuda_export.py

@@ -708,6 +708,73 @@ def test_triton_kernel_mode_off(self):
             "SDPA kernel export with triton_kernel_mode=OFF failed",
         )
 
+    @unittest.skipIf(
+        torch.version.hip is not None,
+        "loads through the Python runtime, which the ROCm build does not include",
+    )
+    def test_fused_sdpa(self):
+        """Decomposing attention on CUDA, as Torch-TensorRT does, gives fused
+        variants that the backend has no kernel for."""
+        from executorch.runtime import Runtime
+
+        class Attention(torch.nn.Module):
+            def forward(self, q, k, v, mask):
+                attention = torch.nn.functional.scaled_dot_product_attention
+                # Flash attention with fewer key heads and a custom scale.
+                grouped = attention(
+                    q,
+                    k[:, :2, :16],
+                    v[:, :2, :16],
+                    is_causal=True,
+                    scale=0.2,
+                    enable_gqa=True,
+                )
+                # Efficient attention, with the mask turned into a bias.
+                masked = attention(q, k[:, :, :16], v[:, :, :16], attn_mask=mask)
+                # Flash attention of one query over many keys, which runs in
+                # the Triton decode kernel.
+                decode = attention(q[:, :, :1], k, v)
+                return grouped, masked, decode
+
+        model = Attention()
+        inputs = (
+            torch.randn(1, 4, 16, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 4, 256, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 4, 256, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.ones(16, 16, dtype=torch.bool, device="cuda").tril(),
+        )
+        program = export(model, inputs, strict=True).run_decompositions()
+        fused = [
+            node
+            for node in program.graph.nodes
+            if "aten._scaled_dot_product_" in str(node.target)
+        ]
+        self.assertEqual(len(fused), 3)
+        with self.assertLogs(
+            "executorch.backends.cuda.triton.replacement_pass", "INFO"
+        ) as logs:
+            edge = to_edge_transform_and_lower(
+                program,
+                partitioner=[
+                    CudaPartitioner(
+                        [CudaBackend.generate_method_name_compile_spec("forward")]
+                    )
+                ],
+                compile_config=EdgeCompileConfig(_check_ir_validity=False),
+            )
+        # The masked call keeps a float bias, which the Triton kernels do not take.
+        output = "\n".join(logs.output)
+        self.assertIn("Replaced 2 nodes with Triton kernels", output)
+        self.assertIn("Using split-K decode SDPA", output)
+        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([t.cpu() for t in inputs])
+        for output, reference in zip(outputs, model(*inputs), strict=True):
+            torch.testing.assert_close(output, reference.cpu(), atol=2e-2, rtol=2e-2)
+
     def test_missing_scatter_kernel_fails_at_export(self):
         """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks."""