diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 5ab15c33016..8ea87c56a79 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -192,6 +192,20 @@ fbcode_target( ], ) +fbcode_target( + _kind = runtime.python_library, + name = "triton_routing", + srcs = [ + "triton/routing.py", + ], + visibility = [ + "//executorch/backends/cuda/...", + ], + deps = [ + "//caffe2:torch", + ], +) + fbcode_target( _kind = runtime.python_library, name = "triton_replacement_pass", @@ -204,6 +218,7 @@ fbcode_target( ], deps = [ ":triton_kernels", + ":triton_routing", "//caffe2:torch", "//executorch/exir/dialects:lib", ], diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index f3a889eb39e..b07feeafd87 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -32,6 +32,26 @@ def define_common_targets(is_fbcode = False): ), ) + python_unittest_remote_gpu( + name = "test_triton_replacement", + srcs = [ + "test_triton_replacement.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:triton_replacement_pass", + "//executorch/backends/cuda:triton_routing", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100", + ), + ) + runtime.python_library( name = "autotune_test_utils", srcs = [ diff --git a/backends/cuda/tests/test_triton_replacement.py b/backends/cuda/tests/test_triton_replacement.py new file mode 100644 index 00000000000..31110bd6376 --- /dev/null +++ b/backends/cuda/tests/test_triton_replacement.py @@ -0,0 +1,184 @@ +# 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. + +"""The generic parts of the Triton replacement pass: kernel legality on the +export device, planned replacement and its transactional rewrite.""" + +import importlib.util +import unittest +from unittest import mock + +import torch + + +class DeviceSupportTest(unittest.TestCase): + def test_capability_of_the_export_device(self): + from executorch.backends.cuda.triton.routing import device_supports + + cuda = torch.device("cuda", 0) + with mock.patch.object(torch.version, "hip", None), mock.patch.object( + torch.cuda, "get_device_capability", return_value=(8, 6) + ) as capability: + self.assertTrue(device_supports(0, cuda)) + self.assertTrue(device_supports(80, cuda)) + self.assertTrue(device_supports(86, cuda)) + self.assertFalse(device_supports(90, cuda)) + self.assertFalse(device_supports(0, torch.device("cpu"))) + capability.assert_called_with(cuda) + + def test_rocm_runs_no_cuda_kernel(self): + from executorch.backends.cuda.triton.routing import device_supports + + with mock.patch.object(torch.version, "hip", "6.0"): + self.assertFalse(device_supports(0, torch.device("cuda", 0))) + + +_PASS_LOGGER = "executorch.backends.cuda.triton.replacement_pass" +_RELU = torch.ops.aten.relu.default +_SIGMOID = torch.ops.aten.sigmoid.default + + +def _graph(): + """relu(x) with two users.""" + graph = torch.fx.Graph() + x = graph.placeholder("x") + y = graph.call_function(_RELU, (x,)) + graph.output((graph.call_function(torch.ops.aten.add.Tensor, (y, 1.0)), y)) + return torch.fx.GraphModule(torch.nn.Module(), graph) + + +def _nodes(graph_module): + return [ + (n.op, n.target, n.name, n.args, n.kwargs, sorted(u.name for u in n.users)) + for n in graph_module.graph.nodes + ] + + +def _targets(graph_module): + return [n.target for n in graph_module.graph.nodes if n.op == "call_function"] + + +class PlannedReplacementTest(unittest.TestCase): + """A planned op through the pass, with a planner that returns a fixed route.""" + + def setUp(self): + if importlib.util.find_spec("triton") is None: + self.skipTest("Triton required") + from executorch.backends.cuda.triton import replacement_pass + + self.route = None + for patch in ( + mock.patch.dict(replacement_pass.EDGE_TO_TRITON_KERNELS, {_RELU: _SIGMOID}), + mock.patch.dict(replacement_pass.PLANNERS, {_RELU: self.plan}), + ): + patch.start() + self.addCleanup(patch.stop) + + def plan(self, node): + from executorch.backends.cuda.triton import routing + + return routing.ReplacementPlan( + self.route, + "only a kernel" if self.route is routing.Route.REQUIRED else "why", + "fake", + "test kernel", + _SIGMOID, + node.args, + "relu(x)", + ) + + def run_pass(self, graph_module): + from executorch.backends.cuda.triton.replacement_pass import ( + ReplaceEdgeOpWithTritonOpPass, + ) + + ReplaceEdgeOpWithTritonOpPass()(graph_module) + graph_module.graph.lint() + + def test_routes(self): + from executorch.backends.cuda.triton import routing + + for route in routing.Route: + self.route = route + graph_module = _graph() + with self.subTest(route=route), self.assertLogs( + _PASS_LOGGER, "INFO" + ) as logs: + self.run_pass(graph_module) + output = "\n".join(logs.output) + if route is routing.Route.UNSUPPORTED: + self.assertEqual(_targets(graph_module).count(_RELU), 1) + self.assertIn("Skipping fake replacement of relu_default: why", output) + self.assertIn("Replaced 0 nodes", output) + else: + self.assertNotIn(_RELU, _targets(graph_module)) + self.assertIn(_SIGMOID, _targets(graph_module)) + self.assertIn( + f"Replacing relu_default ({route.value}, test kernel)", output + ) + self.assertIn("Replaced 1 nodes", output) + + def faults(self): + def fail(*args, **kwargs): + raise RuntimeError("injected fault") + + real_erase = torch.fx.Graph.erase_node + + def fail_erasing_the_original(graph, node): + if node.target == _RELU: + fail() + real_erase(graph, node) + + real_replace_input = torch.fx.Node.replace_input_with + rewired = [] + + def fail_after_the_first_user(user, old, new): + # The first user is rewired, the second fails. + if rewired: + fail() + real_replace_input(user, old, new) + rewired.append(user) + + return { + "node creation": mock.patch.object(torch.fx.Graph, "call_function", fail), + "rewiring the second user": mock.patch.object( + torch.fx.Node, "replace_input_with", fail_after_the_first_user + ), + "erasing": mock.patch.object( + torch.fx.Graph, "erase_node", fail_erasing_the_original + ), + } + + def assertUnchanged(self, graph_module, before): + graph_module.graph.lint() + self.assertEqual(_nodes(graph_module), before) + self.assertNotIn(_SIGMOID, _targets(graph_module)) + + def test_optional_failures_keep_the_node(self): + from executorch.backends.cuda.triton import routing + + self.route = routing.Route.OPTIONAL + for step, fault in self.faults().items(): + graph_module = _graph() + before = _nodes(graph_module) + with self.subTest(step=step): + with fault: + self.run_pass(graph_module) + self.assertUnchanged(graph_module, before) + + def test_required_failures_raise(self): + from executorch.backends.cuda.triton import routing + + self.route = routing.Route.REQUIRED + for step, fault in self.faults().items(): + graph_module = _graph() + before = _nodes(graph_module) + with self.subTest(step=step): + with fault, self.assertRaisesRegex( + routing.ReplacementError, r"relu\(x\) \(only a kernel\)" + ): + self.run_pass(graph_module) + self.assertUnchanged(graph_module, before) diff --git a/backends/cuda/triton/replacement_pass.py b/backends/cuda/triton/replacement_pass.py index abfd4e92941..81625fa6a97 100644 --- a/backends/cuda/triton/replacement_pass.py +++ b/backends/cuda/triton/replacement_pass.py @@ -11,8 +11,10 @@ """ import logging +from typing import Callable import torch +from executorch.backends.cuda.triton import routing from executorch.backends.cuda.triton.kernels.sdpa import ( _prepare_mask_params, _validate_qkv_shapes, @@ -32,6 +34,10 @@ exir_ops.edge.aten.topk.default: triton.topk, } +# Ops whose calls are classified before they are replaced (see routing.py): +# edge op -> planner(node) returning a ``routing.ReplacementPlan``. +PLANNERS: dict[object, Callable[[Node], routing.ReplacementPlan]] = {} + _SPLITK_LKV_THRESHOLD = 256 @@ -49,6 +55,7 @@ def __init__(self): """Initialize the pass.""" super().__init__() self._replacement_count = 0 + self._plans: dict[Node, routing.ReplacementPlan] = {} def call(self, graph_module: GraphModule) -> PassResult: """ @@ -61,21 +68,29 @@ def call(self, graph_module: GraphModule) -> PassResult: PassResult indicating success/failure and the modified graph module """ self._replacement_count = 0 + self._plans = {} modified = False if not EDGE_TO_TRITON_KERNELS: return PassResult(graph_module, False) - # Iterate through all nodes in the graph - for node in graph_module.graph.nodes: - if self._should_replace_node(node): + for node in list(graph_module.graph.nodes): + if not self._should_replace_node(node): + continue + if node.target in PLANNERS: + # Rolls back on failure and raises only for calls that must be + # replaced. + replaced = self._replace_node_with_triton(graph_module, node) + else: try: - self._replace_node_with_triton(graph_module, node) - modified = True - self._replacement_count += 1 + replaced = self._replace_node_with_triton(graph_module, node) except Exception as e: logger.warning(f"Failed to replace node {node.name}: {e}") # Continue with other replacements even if one fails + replaced = False + if replaced: + modified = True + self._replacement_count += 1 if modified: # Recompile the graph module after modifications @@ -152,6 +167,16 @@ def _should_replace_node(self, node: Node) -> bool: if node.target == exir_ops.edge.aten.scaled_dot_product_attention.default: return self._sdpa_kernel_supports(node) + if node.target in PLANNERS: + plan = PLANNERS[node.target](node) + if plan.route is routing.Route.UNSUPPORTED: + logger.info( + f"Skipping {plan.kind} replacement of {node.name}: {plan.reason}" + ) + return False + self._plans[node] = plan + logger.info(f"Replacing {node.name} ({plan.route.value}, {plan.label})") + return True @classmethod @@ -189,13 +214,16 @@ def _sdpa_kernel_supports(cls, node: Node) -> bool: return False return True - def _replace_node_with_triton(self, graph_module: GraphModule, node: Node) -> None: + def _replace_node_with_triton(self, graph_module: GraphModule, node: Node) -> bool: """ Replace an edge dialect node with a Triton kernel call. Args: graph_module: The graph module containing the node node: The node to replace + + Returns: + Whether the node was replaced """ # Get the target operator (should be an exir_ops edge dialect op) target = node.target @@ -204,6 +232,9 @@ def _replace_node_with_triton(self, graph_module: GraphModule, node: Node) -> No if target not in EDGE_TO_TRITON_KERNELS: raise ValueError(f"No replacement kernel found for {target}") + if target in PLANNERS: + return routing.replace(graph_module.graph, node, self._plans.pop(node)) + triton_kernel_fn = EDGE_TO_TRITON_KERNELS[target] if target == exir_ops.edge.aten.scaled_dot_product_attention.default: @@ -230,3 +261,4 @@ def _replace_node_with_triton(self, graph_module: GraphModule, node: Node) -> No # Remove the old node graph_module.graph.erase_node(node) + return True diff --git a/backends/cuda/triton/routing.py b/backends/cuda/triton/routing.py new file mode 100644 index 00000000000..96b5f11bba0 --- /dev/null +++ b/backends/cuda/triton/routing.py @@ -0,0 +1,100 @@ +# 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. + +"""Preflight plans and the transactional rewrite of the Triton replacement pass. + +An op whose calls a Triton kernel takes only sometimes registers a planner in +``replacement_pass.PLANNERS``. The planner is a pure preflight: it inspects a +node, never changes the graph, and classifies it: + +- ``OPTIONAL``: a Triton kernel runs it and the regular lowering could too. A + failed rewrite is rolled back and the node keeps the regular lowering. +- ``REQUIRED``: a Triton kernel runs it and the regular lowering cannot. A + failure raises ``ReplacementError`` with the call's signature. +- ``UNSUPPORTED``: the node stays as it is; ``reason`` says why. +""" + +from __future__ import annotations + +import enum +import logging +from dataclasses import dataclass +from typing import Any, Optional + +import torch +from torch.fx import Graph, Node + +logger: logging.Logger = logging.getLogger(__name__) + + +class Route(enum.Enum): + OPTIONAL = "optional" + REQUIRED = "required" + UNSUPPORTED = "unsupported" + + +class ReplacementError(RuntimeError): + """A call only a Triton kernel can compile could not be replaced.""" + + +@dataclass(frozen=True) +class ReplacementPlan: + route: Route + reason: str + # The kind of op ("conv") and the kernel chosen, for the pass's logs. + kind: str = "" + label: str = "" + kernel: Any = None + args: tuple[Any, ...] = () + # The normalized call, for errors. + signature: str = "" + + +def device_supports(min_capability: int, device: torch.device) -> bool: + """Whether a kernel that needs compute capability ``min_capability`` (e.g. + 80 for sm_80; 0 for every NVIDIA GPU) runs on ``device``, the GPU the model + is lowered and autotuned on. Newer GPUs run it through the package's PTX; + older ones load a variant exported for them.""" + if torch.version.hip is not None or device.type != "cuda": + return False + if min_capability == 0: + return True + major, minor = torch.cuda.get_device_capability(device) + return 10 * major + minor >= min_capability + + +def unsupported(kind: str, reason: str) -> ReplacementPlan: + return ReplacementPlan(Route.UNSUPPORTED, reason, kind) + + +def replace(graph: Graph, node: Node, plan: ReplacementPlan) -> bool: + """Replaces ``node`` with ``plan.kernel(*plan.args)`` as one transaction. + + Every user's arguments are saved before they are rewired, so a failure at + any step restores them all and removes the new node; a REQUIRED plan then + raises, an OPTIONAL one logs and returns False. + """ + saved = {user: (user.args, user.kwargs) for user in node.users} + new_node: Optional[Node] = None + try: + with graph.inserting_before(node): + new_node = graph.call_function(plan.kernel, plan.args) + new_node.meta = node.meta.copy() + for user in saved: + user.replace_input_with(node, new_node) + graph.erase_node(node) + except Exception as e: + for user, (args, kwargs) in saved.items(): + user.args, user.kwargs = args, kwargs + if new_node is not None: + graph.erase_node(new_node) + if plan.route is Route.REQUIRED: + raise ReplacementError( + f"{plan.signature} ({plan.reason}): replacing {node.name} failed: {e!r}" + ) from e + logger.warning("Kept %s: replacing it failed: %r", node.name, e) + return False + return True