From 4a201ca280674d055e85d54deadb9a2ef89621e1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 1 Nov 2026 22:29:53 -0800 Subject: [PATCH] Lower fused attention variants in the CUDA backend Decomposing scaled_dot_product_attention on a CUDA device turns it into one of its fused variants: _scaled_dot_product_flash_attention without a mask, or _scaled_dot_product_efficient_attention with one. torch.export's default decompositions do this, and so does Torch-TensorRT for an attention it leaves to the CUDA backend, for example one whose keys are a constant. The CUDA backend had no kernel for these ops: RuntimeError: Method forward missing fallback kernels (1 total): - at::_ops::_scaled_dot_product_flash_attention::call A new pass, ReplaceFusedSdpaWithSdpaPass, turns the flash, efficient and cuDNN variants back into scaled_dot_product_attention with the same query, key, value, bias, dropout, causal flag and scale. It sets enable_gqa unless the query and key head counts are known to be equal, since the fused variants group fewer key heads the same way. It leaves out a scale of None, because the Triton SDPA kernels take a float scale. The Triton SDPA kernel or Inductor's math decomposition then lowers the attention, as it does any other. The pass also looks inside control flow branches. The pass rewrites a node only when scaled_dot_product_attention computes the same result. Only its attention output may be used, which is the case for inference. A causal call must have no bias, since scaled_dot_product_attention takes no mask together with is_causal. Its query and key lengths must also be known to be equal: flash attention aligns a causal mask to the bottom right, the other variants and scaled_dot_product_attention to the top left, and the Triton decode kernel for one query ignores it. Otherwise the node stays, the pass logs why, and the error above still names it. test_fused_sdpa decomposes attention on a CUDA device into a causal flash attention with fewer key heads and a custom scale, an efficient attention with a mask and a flash attention of one query over 256 keys. It lowers the program, checks that the first and last calls get the Triton kernels, runs it, and compares the outputs with eager. The tests in passes/tests/test_replace_fused_sdpa.py check on the CPU which calls the pass rewrites and with which keyword arguments (cuDNN, causal with equal static or dynamic lengths, fewer key heads, inside a torch.cond branch), which it keeps (causal with lengths that differ or may differ, causal with a bias, a used log-sum-exp), that the program stays valid, and that the pass runs before the Triton kernel replacement. They fail without this change. --- backends/cuda/BUCK | 1 + backends/cuda/cuda_backend.py | 9 +- backends/cuda/passes/replace_fused_sdpa.py | 136 ++++++++++++++++++ .../passes/tests/test_replace_fused_sdpa.py | 124 ++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 67 +++++++++ 5 files changed, 336 insertions(+), 1 deletion(-) create mode 100644 backends/cuda/passes/replace_fused_sdpa.py create mode 100644 backends/cuda/passes/tests/test_replace_fused_sdpa.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index ca8f9eee753..7f88126cd5f 100644 --- a/backends/cuda/BUCK +++ b/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 = [ diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 891fb81bd06..f2bd1a202f1 100644 --- a/backends/cuda/cuda_backend.py +++ b/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" diff --git a/backends/cuda/passes/replace_fused_sdpa.py b/backends/cuda/passes/replace_fused_sdpa.py new file mode 100644 index 00000000000..1da395e8e09 --- /dev/null +++ b/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) diff --git a/backends/cuda/passes/tests/test_replace_fused_sdpa.py b/backends/cuda/passes/tests/test_replace_fused_sdpa.py new file mode 100644 index 00000000000..96be05e84f1 --- /dev/null +++ b/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) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 637ff4036b1..725a95a60c7 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/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."""