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."""