From 588c129559d12657f549283ad67034206eb20821 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sat, 17 Oct 2026 16:56:20 -0700 Subject: [PATCH] [executorch][cuda] Reuse PyTorch's access analysis of a repeated user Triton kernel call Pull Request resolved: https://github.com/pytorch/executorch/pull/23739 Tracing a call of a user-defined Triton kernel runs `torch._higher_order_ops.triton_kernel_wrap.identify_accessed_tensors`, which generates the kernel's TTIR and walks it to find the tensors the kernel reads and writes. PyTorch runs it in functionalization, in Inductor's mutation propagation and in Inductor's lowering, and caches nothing, so a model whose every layer launches the same kernels repeats it for every call site of every exported method. In the Muse Glimmer batching export it was the largest single item on the main thread: ~21 minutes. The analysis depends only on the kernel and on what `generate_ttir` makes of the arguments: a fake or Inductor tensor becomes a fresh tensor of its dtype, a symbolic value becomes 2, and any other value is used as is (a real tensor also by its pointer's 16-byte alignment). `backends/cuda/triton/kernels/triton_mutation_cache.py` keys calls on exactly that. The first call with a given key goes through PyTorch; repeats return a copy of its result. Calls it cannot key exactly (TMA descriptors, unhashable values) always go through PyTorch. It is installed when `executorch.backends.cuda.triton.kernels` is imported, by replacing the function on `triton_kernel_wrap`, which both of PyTorch's callers look it up on. Full Muse Glimmer 30B batching export (`export_solo.py`, 10 methods; A100; empty Triton and Inductor caches), stacked on D124736179 (min): | | before | this diff | |---|---|---| | main thread in the access analysis | ~21.1 | ~8.7 | | lowering start -> forward_1 compiled | 23.3 | 19.1 | | AOTI forward_2 / 4 / 8 / 16 / 32 / 64 | 13.3 / 12.0 / 8.5 / 9.8 / 11.6 / 15.5 | 9.8 / 8.3 / 7.2 / 9.5 / 10.8 / 13.8 | | total | 123.3 | 108.6 | (The rest of the host was busier during this run: ~15 other cores vs ~11.) What is left of the analysis is mostly its first call per key, since each method has its own shapes. The weight file is byte-identical, and the export generates the same tokens (greedy; 512 and 4 x 2k prompts) at the same speed (63.0 vs 63.2 tok/s; 110.6 vs 110.4). ghstack-source-id: 451021042 @exported-using-ghexport Differential Revision: [D124772365](https://our.internmc.facebook.com/intern/diff/D124772365/) --- backends/cuda/BUCK | 1 + backends/cuda/cuda_backend.py | 2 + backends/cuda/tests/targets.bzl | 19 + .../cuda/tests/test_triton_mutation_cache.py | 328 ++++++++++++++++++ .../triton/kernels/triton_mutation_cache.py | 131 +++++++ 5 files changed, 481 insertions(+) create mode 100644 backends/cuda/tests/test_triton_mutation_cache.py create mode 100644 backends/cuda/triton/kernels/triton_mutation_cache.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index bd2373e3dd8..ca8f9eee753 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -176,6 +176,7 @@ fbcode_target( "triton/kernels/sdpa.py", "triton/kernels/topk.py", "triton/kernels/tq4_sdpa.py", + "triton/kernels/triton_mutation_cache.py", ], visibility = [ "//executorch/backends/cuda/...", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index f28faabb624..aad054fa288 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -42,6 +42,7 @@ ReplaceInt64FloorDivWithFloatPass, ) from executorch.backends.cuda.target_smem import target_smem_context +from executorch.backends.cuda.triton.kernels import triton_mutation_cache from executorch.backends.cuda.triton.replacement_pass import ( ReplaceEdgeOpWithTritonOpPass, ) @@ -958,6 +959,7 @@ def _combined(): # them. SDPA ops already replaced by Triton kernels via # `ReplaceEdgeOpWithTritonOpPass` are unaffected. stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH])) + stack.enter_context(triton_mutation_cache.enabled()) stack.enter_context(target_smem_context()) # On ROCm, is_big_gpu gates on architecture, not SM count. if torch.version.hip is None: diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index 8630a2e2c31..576943f147b 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -306,6 +306,25 @@ def define_common_targets(is_fbcode = False): ], ) + python_unittest_remote_gpu( + name = "test_triton_mutation_cache", + srcs = [ + "test_triton_mutation_cache.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:triton_kernels", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100-exclusive", + ), + ) + python_unittest_remote_gpu( name = "test_cuda_weight_reuse", srcs = [ diff --git a/backends/cuda/tests/test_triton_mutation_cache.py b/backends/cuda/tests/test_triton_mutation_cache.py new file mode 100644 index 00000000000..61fb675e3b5 --- /dev/null +++ b/backends/cuda/tests/test_triton_mutation_cache.py @@ -0,0 +1,328 @@ +# 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. + +"""triton_mutation_cache: repeats of PyTorch's Triton access analysis are reused.""" + +import re +import threading +import time +import unittest +from unittest import mock + +import torch +import triton +import triton.language as tl +from executorch.backends.cuda.triton.kernels import triton_mutation_cache +from torch._higher_order_ops import triton_kernel_wrap +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.library import wrap_triton + + +@triton.jit +def _scale_kernel( + x_ptr, y_ptr, z_ptr, n, scale, WRITE_Z: tl.constexpr, BLOCK: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < n + x = tl.load(x_ptr + offsets, mask=mask) + tl.store(y_ptr + offsets, x * scale, mask=mask) + if WRITE_Z: + tl.store(z_ptr + offsets, x + scale, mask=mask) + + +@triton.jit +def _store_z_kernel( + x_ptr, y_ptr, z_ptr, n, scale, WRITE_Z: tl.constexpr, BLOCK: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < n + x = tl.load(x_ptr + offsets, mask=mask) + tl.store(z_ptr + offsets, x * scale, mask=mask) + + +def _scale(x: torch.Tensor, scale: float) -> torch.Tensor: + y = torch.empty_like(x) + z = torch.empty_like(x) + n = x.numel() + wrap_triton(_scale_kernel)[(triton.cdiv(n, 256),)]( + x, y, z, n, scale, WRITE_Z=False, BLOCK=256 + ) + return y + + +class _Layers(torch.nn.Module): + """Calls the kernel directly, so export traces (and analyzes) each call.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Like a model's layers: the same kernel, dtypes and scalars each time. + for _ in range(4): + x = _scale(x, 0.5) + return x + + +def _written(analysis) -> set: + return {dep.name for dep in analysis.read_writes.writes} + + +def _scale_with_num_warps(x: torch.Tensor, scale: float) -> torch.Tensor: + y = torch.empty_like(x) + z = torch.empty_like(x) + n = x.numel() + wrap_triton(_scale_kernel)[(triton.cdiv(n, 256),)]( + x, y, z, n, scale, WRITE_Z=False, BLOCK=256, num_warps=4 + ) + return y + + +class _NumWarpsLayers(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + for _ in range(4): + x = _scale_with_num_warps(x, 0.5) + return x + + +class TritonMutationCacheTest(unittest.TestCase): + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA required") + self.original = triton_kernel_wrap.identify_accessed_tensors + + def _kwargs(self, mode, **overrides): + with mode: + x = torch.empty(1000, device="cuda") + kwargs = { + "x_ptr": x, + "y_ptr": torch.empty_like(x), + "z_ptr": torch.empty_like(x), + "n": 1000, + "scale": 2.0, + "WRITE_Z": False, + "BLOCK": 256, + } + kwargs.update(overrides) + return kwargs + + def _counting_stand_in(self, delay: float = 0.0): + """A stand-in analysis that counts its calls, to patch in for PyTorch's.""" + calls = [] + + def analysis(kernel, kwargs, tma_descriptor_metadata): + calls.append(1) + time.sleep(delay) + return {"writes": sorted(kwargs)} + + return ( + mock.patch.object( + triton_kernel_wrap, "identify_accessed_tensors", analysis + ), + calls, + ) + + def test_the_context_patches_the_function_pytorch_calls_and_restores_it( + self, + ) -> None: + with triton_mutation_cache.enabled(): + self.assertIsNot( + triton_kernel_wrap.identify_accessed_tensors, self.original + ) + self.assertIs(triton_kernel_wrap.identify_accessed_tensors, self.original) + + def test_first_call_goes_through_pytorch_and_repeats_reuse_it(self) -> None: + mode = FakeTensorMode() + with triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + first = analyze(_scale_kernel, self._kwargs(mode), {}) + # Other fake tensors of the same dtypes and the same scalars. + for _ in range(3): + again = analyze(_scale_kernel, self._kwargs(FakeTensorMode()), {}) + self.assertEqual(again, first) + self.assertIsNot(again, first) + self.assertEqual(counts, {"hits": 3, "misses": 1, "uncached": 0}) + self.assertEqual(first, self.original(_scale_kernel, self._kwargs(mode), {})) + + def test_what_changes_the_kernel_is_not_reused(self) -> None: + mode = FakeTensorMode() + variants = [ + {}, + {"WRITE_Z": True}, + {"BLOCK": 128}, + {"n": 1}, + {"scale": 3.0}, + ] + with mode: + half = torch.empty(1000, device="cuda", dtype=torch.float16) + variants.append({"x_ptr": half}) + with triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + for overrides in variants: + kwargs = self._kwargs(mode, **overrides) + self.assertEqual( + analyze(_scale_kernel, kwargs, {}), + self.original(_scale_kernel, kwargs, {}), + overrides, + ) + self.assertEqual(counts["misses"], len(variants)) + self.assertEqual( + _written( + self.original(_scale_kernel, self._kwargs(mode, WRITE_Z=True), {}) + ), + {"y_ptr", "z_ptr"}, + ) + + def test_kernels_with_the_same_signature_are_kept_apart(self) -> None: + mode = FakeTensorMode() + kwargs = self._kwargs(mode, WRITE_Z=True) + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + scale = analyze(_scale_kernel, kwargs, {}) + store = analyze(_store_z_kernel, kwargs, {}) + self.assertEqual(_written(scale), {"y_ptr", "z_ptr"}) + self.assertEqual(_written(store), {"z_ptr"}) + + def test_an_autotuners_first_config_is_part_of_the_key(self) -> None: + tuned = triton.autotune( + configs=[triton.Config({"WRITE_Z": False, "BLOCK": 256})], key=["n"] + )(_scale_kernel) + kwargs = self._kwargs(FakeTensorMode()) + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + before = analyze(tuned, kwargs, {}) + tuned.configs = [triton.Config({"WRITE_Z": True, "BLOCK": 256})] + after = analyze(tuned, kwargs, {}) + self.assertEqual(_written(before), {"y_ptr"}) + self.assertEqual(_written(after), {"y_ptr", "z_ptr"}) + + def test_fresh_autotuners_of_one_kernel_share_an_entry(self) -> None: + mode = FakeTensorMode() + with triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + for _ in range(2): + tuned = triton.autotune( + configs=[ + triton.Config({"WRITE_Z": False, "BLOCK": 256}, num_warps=4) + ], + key=["n"], + )(_scale_kernel) + analyze(tuned, self._kwargs(mode), {}) + self.assertEqual(counts["hits"], 1) + self.assertEqual(counts["misses"], 1) + + def test_tma_descriptors_and_unhashable_values_always_go_through_pytorch( + self, + ) -> None: + patch, calls = self._counting_stand_in() + with patch, triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + for _ in range(2): + analyze(_scale_kernel, {"x_ptr": 1}, {"x_ptr": ("stable", ([16],))}) + analyze(_scale_kernel, {"x_ptr": [1, 2]}, {}) + self.assertEqual(len(calls), 4) + self.assertEqual(counts, {"hits": 0, "misses": 0, "uncached": 4}) + + def test_concurrent_calls_of_a_key_analyze_once(self) -> None: + patch, calls = self._counting_stand_in(delay=0.2) + results = [] + with patch, triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + threads = [ + threading.Thread( + target=lambda: results.append(analyze(_scale_kernel, {"n": 7}, {})) + ) + for _ in range(4) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + self.assertEqual(len(calls), 1) + self.assertEqual(results, [{"writes": ["n"]}] * 4) + self.assertEqual(counts, {"hits": 3, "misses": 1, "uncached": 0}) + + def test_real_tensors_are_keyed_by_dtype_and_pointer_alignment(self) -> None: + patch, calls = self._counting_stand_in() + base = torch.empty(64, dtype=torch.float32, device="cuda") + with patch, triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + # base[4:] starts 16 bytes in, so its pointer has base's alignment; + # base[1:] starts 4 bytes in and does not. + for x in (base, base[4:], base[1:], base[1:]): + analyze(_scale_kernel, {"x_ptr": x}, {}) + self.assertEqual(len(calls), 2) + self.assertEqual(counts, {"hits": 2, "misses": 2, "uncached": 0}) + + def test_entries_do_not_outlive_the_context(self) -> None: + patch, calls = self._counting_stand_in() + with patch: + for _ in range(2): + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + analyze(_scale_kernel, {"n": 1}, {}) + analyze(_scale_kernel, {"n": 1}, {}) + self.assertEqual(len(calls), 2) + + def test_tensor_shapes_and_strides_do_not_change_the_analysis(self) -> None: + # The key leaves them out; PyTorch's own analysis must not depend on them. + mode = FakeTensorMode() + with mode: + tensors = [ + torch.empty(1000, device="cuda"), + torch.empty(64, 32, device="cuda").t(), + torch.empty(2048, device="cuda")[1:], + ] + first = self.original(_scale_kernel, self._kwargs(mode), {}) + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + for x in tensors: + kwargs = self._kwargs(mode, x_ptr=x, y_ptr=x, z_ptr=x) + self.assertEqual(self.original(_scale_kernel, kwargs, {}), first) + self.assertEqual(analyze(_scale_kernel, kwargs, {}), first) + + def test_export_traces_the_same_graph_and_analyzes_once(self) -> None: + x = torch.randn(1000, device="cuda") + + def exported_code(): + program = torch.export.export(_Layers(), (x,)).run_decompositions({}) + # Side-table indices grow with every export; the rest must match. + return re.sub( + r"(kernel_idx|constant_args_idx) = \d+", + r"\1 = N", + program.graph_module.code, + ) + + with triton_mutation_cache.enabled() as counts: + cached = exported_code() + self.assertEqual(counts["misses"], 1) + self.assertGreaterEqual(counts["hits"], 3) + self.assertEqual(cached, exported_code()) + + def test_a_launch_with_num_warps_is_analyzed_once(self) -> None: + x = torch.randn(1000, device="cuda") + with triton_mutation_cache.enabled() as counts: + torch.export.export(_NumWarpsLayers(), (x,)).run_decompositions({}) + self.assertEqual(counts["misses"], 1) + self.assertGreaterEqual(counts["hits"], 3) + + def test_inductor_compile_gives_eager_results_and_reaches_the_analysis( + self, + ) -> None: + x = torch.randn(1000, device="cuda") + with triton_mutation_cache.enabled() as counts: + out = torch.compile(_Layers(), fullgraph=True)(x) + torch.testing.assert_close(out, _Layers()(x)) + self.assertGreaterEqual(counts["misses"], 1) + self.assertEqual(counts["uncached"], 0) + + def test_a_dynamic_shape_export_analyzes_the_symbolic_size_once(self) -> None: + x = torch.randn(1000, device="cuda") + n = torch.export.Dim("n", min=2, max=1 << 20) + with triton_mutation_cache.enabled() as counts: + program = torch.export.export( + _Layers(), (x,), dynamic_shapes=({0: n},) + ).run_decompositions({}) + self.assertEqual(counts["misses"], 1) + self.assertGreaterEqual(counts["hits"], 3) + other = torch.randn(4096, device="cuda") + torch.testing.assert_close(program.module()(other), _Layers()(other)) diff --git a/backends/cuda/triton/kernels/triton_mutation_cache.py b/backends/cuda/triton/kernels/triton_mutation_cache.py new file mode 100644 index 00000000000..096d2b3ee26 --- /dev/null +++ b/backends/cuda/triton/kernels/triton_mutation_cache.py @@ -0,0 +1,131 @@ +# 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. + +"""Remember PyTorch's analysis of which tensors a user Triton kernel accesses. + +Tracing a call of a user-defined Triton kernel (functionalization in +``torch.export``, and Inductor's mutation propagation and lowering) runs +``torch._higher_order_ops.triton_kernel_wrap.identify_accessed_tensors``: it +generates the kernel's TTIR and walks it to find the tensors the kernel reads +and writes. Nothing caches it, so a model whose every layer launches the same +kernels repeats it for every call site of every exported method. + +The analysis depends only on the traced kernel and on what ``generate_ttir`` +turns the arguments into: a fake or Inductor tensor becomes a fresh tensor of +its dtype, a symbolic value becomes 2, and every other value is used as is. +``enabled()`` caches it for the duration of a compile. Calls this cannot key +exactly (TMA descriptors, unhashable values) always go through PyTorch. +""" + +from __future__ import annotations + +import contextlib +import copy +import threading +from typing import Any, Dict, Iterator, Optional, Tuple + +import torch +import triton + +_UNKEYABLE = object() + +_lazy_types: Optional[Tuple[Tuple[type, ...], Tuple[type, ...]]] = None + + +def _symbolic_and_traced_tensor_types() -> Tuple[Tuple[type, ...], Tuple[type, ...]]: + """The symbolic-value and fake/Inductor tensor types, imported on first use + so that tracing does not load Inductor before it is needed.""" + global _lazy_types + if _lazy_types is None: + import sympy + from torch._inductor.ir import TensorBox + from torch._subclasses.fake_tensor import FakeTensor + + _lazy_types = ( + (torch.SymInt, torch.SymFloat, torch.SymBool, sympy.Expr), + (FakeTensor, TensorBox), + ) + return _lazy_types + + +def _argument_key(value: Any) -> Any: + """What ``generate_ttir`` makes of ``value``, as a hashable key.""" + symbolic, traced_tensor = _symbolic_and_traced_tensor_types() + if isinstance(value, symbolic): + return ("symbolic",) + if isinstance(value, traced_tensor): + return ("tensor", value.dtype) + if isinstance(value, torch.Tensor): + # Triton specializes on the pointer's alignment, so that is what matters. + return ("real tensor", value.dtype, value.data_ptr() % 16) + try: + hash(value) + except TypeError: + return _UNKEYABLE + return (type(value).__module__, type(value).__qualname__, value) + + +def _traced_kernel_and_kwargs(kernel: Any, kwargs: Dict[str, Any]) -> Tuple[Any, Dict]: + """The JIT function ``generate_ttir`` traces, and the arguments it traces it + with: an autotuner's first config is merged into the arguments.""" + if isinstance(kernel, triton.runtime.Autotuner) and kernel.configs: + return kernel.fn, {**kwargs, **kernel.configs[0].kwargs} + return kernel, kwargs + + +def _call_key( + kernel: Any, kwargs: Dict[str, Any], tma_descriptor_metadata +) -> Optional[Tuple[Tuple, Any]]: + """The cache key of a call and the traced kernel it belongs to, or None.""" + if tma_descriptor_metadata: + return None + traced, merged = _traced_kernel_and_kwargs(kernel, kwargs) + arguments = [] + for name in sorted(merged): + key = _argument_key(merged[name]) + if key is _UNKEYABLE: + return None + arguments.append((name, key)) + return (id(traced), tuple(arguments)), traced + + +@contextlib.contextmanager +def enabled() -> Iterator[Dict[str, int]]: + """Caches PyTorch's access analysis while the context is active. Yields the + hit, miss and uncached call counts. Contexts must not overlap: each one + patches the function PyTorch calls.""" + from torch._higher_order_ops import triton_kernel_wrap + + original = triton_kernel_wrap.identify_accessed_tensors + # key -> (traced kernel, result). The kernel is held so its id stays unique. + entries: Dict[Tuple, Tuple[Any, Any]] = {} + counts = {"hits": 0, "misses": 0, "uncached": 0} + lock = threading.RLock() + + def cached(kernel, kwargs, tma_descriptor_metadata): + keyed = _call_key(kernel, kwargs, tma_descriptor_metadata) + if keyed is None: + with lock: + counts["uncached"] += 1 + return original(kernel, kwargs, tma_descriptor_metadata) + key, traced = keyed + with lock: + entry = entries.get(key) + if entry is not None: + counts["hits"] += 1 + return copy.deepcopy(entry[1]) + counts["misses"] += 1 + result = original(kernel, kwargs, tma_descriptor_metadata) + entries[key] = (traced, copy.deepcopy(result)) + return result + + # Both callers (get_mutated_tensors and Inductor's UserDefinedTritonKernel) + # look the function up on this module. + triton_kernel_wrap.identify_accessed_tensors = cached + try: + yield counts + finally: + triton_kernel_wrap.identify_accessed_tensors = original