[CUDA] Use Triton TMA for long-context causal prefill

backends/cuda/BUCK

@@ -3,6 +3,31 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 oncall("executorch")
 
+fbcode_target(
+    _kind = runtime.python_library,
+    name = "target_arch",
+    srcs = [
+        "target_arch.py",
+    ],
+    visibility = [
+        "//executorch/backends/cuda/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
+fbcode_target(
+    _kind = runtime.python_library,
+    name = "optimization_config",
+    srcs = [
+        "optimization_config.py",
+    ],
+    visibility = [
+        "//executorch/backends/cuda/...",
+    ],
+)
+
 fbcode_target(
     _kind = runtime.python_library,
     name = "coalesced_int4_tensor",
@@ -99,6 +124,8 @@ fbcode_target(
     visibility = ["PUBLIC"],
     deps = [
         ":cuda_passes",
+        ":optimization_config",
+        ":target_arch",
         ":triton_replacement_pass",
         "//caffe2:torch",
         "//executorch/backends/aoti/passes:passes",
@@ -137,6 +164,8 @@ fbcode_target(
         "//executorch/backends/cuda/...",
     ],
     deps = [
+        ":optimization_config",
+        ":target_arch",
         "//caffe2:torch",
     ],
 )

backends/cuda/cuda_backend.py

@@ -24,12 +24,14 @@
     CudaWeightCollector,
     trim_host_memory,
 )
+from executorch.backends.cuda.optimization_config import cuda_optimization_context
 from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
     MoveCondPredicateToCpuPass,
 )
 from executorch.backends.cuda.passes.replace_int64_floordiv import (
     ReplaceInt64FloorDivWithFloatPass,
 )
+from executorch.backends.cuda.target_arch import cuda_targets_are_sm90_or_newer
 from executorch.backends.cuda.target_smem import target_smem_context
 from executorch.backends.cuda.triton.replacement_pass import (
     ReplaceEdgeOpWithTritonOpPass,
@@ -866,6 +868,7 @@ def get_extra_aoti_compile_context_manager(
         # Parse compile_specs for low_memory_mode (default OFF). compile_specs
         # may be None when called without specs (parity with base default).
         low_memory_mode = "OFF"
+        tma_causal_prefill = False
         for spec in compile_specs or []:
             if spec.key == "low_memory_mode":
                 mode = spec.value.decode("utf-8").upper()
@@ -874,6 +877,14 @@ def get_extra_aoti_compile_context_manager(
                         f"Invalid low_memory_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
                 low_memory_mode = mode
+            elif spec.key == "enable_tma_causal_prefill":
+                tma_causal_prefill = _on_off_compile_spec_value(spec)
+
+        if tma_causal_prefill and not cuda_targets_are_sm90_or_newer():
+            logging.warning(
+                "enable_tma_causal_prefill requires an SM90+ CUDA target; disabling it"
+            )
+            tma_causal_prefill = False
 
         @contextlib.contextmanager
         def _combined():
@@ -885,6 +896,9 @@ def _combined():
                 # only the fallback for the `triton_kernel_mode="OFF"` path.
                 stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH]))
                 stack.enter_context(target_smem_context())
+                stack.enter_context(
+                    cuda_optimization_context(tma_causal_prefill=tma_causal_prefill)
+                )
                 if low_memory_mode == "ON":
                     # Force AOTI's mutated-buffer clones onto CPU during
                     # compile so we stay under tight GPU memory caps (e.g.

backends/cuda/optimization_config.py

@@ -0,0 +1,29 @@
+# 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.
+
+"""Scoped CUDA optimization choices used while AOTInductor traces kernels."""
+
+import contextlib
+import contextvars
+from typing import Iterator
+
+
+_TMA_CAUSAL_PREFILL_ENABLED = contextvars.ContextVar(
+    "tma_causal_prefill_enabled", default=False
+)
+
+
+def tma_causal_prefill_enabled() -> bool:
+    return _TMA_CAUSAL_PREFILL_ENABLED.get()
+
+
+@contextlib.contextmanager
+def cuda_optimization_context(*, tma_causal_prefill: bool) -> Iterator[None]:
+    tma_token = _TMA_CAUSAL_PREFILL_ENABLED.set(tma_causal_prefill)
+    try:
+        yield
+    finally:
+        _TMA_CAUSAL_PREFILL_ENABLED.reset(tma_token)

backends/cuda/target_arch.py

@@ -0,0 +1,42 @@
+# 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.
+
+"""Helpers for selecting CUDA export paths from the requested target."""
+
+import os
+import re
+
+import torch
+
+
+def cuda_targets_are_sm90_or_newer() -> bool:
+    """Return whether every requested NVIDIA CUDA target is SM90 or newer.
+
+    Explicit ``TORCH_CUDA_ARCH_LIST`` targets take precedence over the local
+    export device. This keeps per-architecture AOT exports deterministic while
+    retaining local-device detection for the usual native-export workflow.
+    """
+    if torch.version.hip is not None:
+        return False
+
+    arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST")
+    if arch_list:
+        target_majors = []
+        for target in re.split(r"[;,\s]+", arch_list):
+            target = target.strip().lower().removeprefix("sm_").removeprefix("compute_")
+            target = target.removesuffix("+ptx").removesuffix("a")
+            if not target:
+                continue
+            match = re.fullmatch(r"(\d+)(?:\.(\d+))?", target)
+            if match is None:
+                return False
+            major = int(match.group(1))
+            if match.group(2) is None and major >= 10:
+                major //= 10
+            target_majors.append(major)
+        return bool(target_majors) and min(target_majors) >= 9
+
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9

backends/cuda/tests/test_cuda_export.py

@@ -104,6 +104,49 @@ def test_invalid_autotune_at_compile_time_compile_spec(self):
                 [CompileSpec(key="autotune_at_compile_time", value=b"MAYBE")]
             )
 
+    def test_tma_causal_prefill_defaults_off(self):
+        from executorch.backends.cuda.optimization_config import (
+            tma_causal_prefill_enabled,
+        )
+
+        with CudaBackend.get_extra_aoti_compile_context_manager([]):
+            self.assertFalse(tma_causal_prefill_enabled())
+
+    def test_tma_causal_prefill_compile_spec(self):
+        from executorch.backends.cuda.optimization_config import (
+            tma_causal_prefill_enabled,
+        )
+
+        with patch(
+            "executorch.backends.cuda.cuda_backend.cuda_targets_are_sm90_or_newer",
+            return_value=True,
+        ), CudaBackend.get_extra_aoti_compile_context_manager(
+            [CompileSpec(key="enable_tma_causal_prefill", value=b"ON")]
+        ):
+            self.assertTrue(tma_causal_prefill_enabled())
+
+    def test_invalid_tma_causal_prefill_compile_spec(self):
+        with self.assertRaisesRegex(ValueError, "Invalid enable_tma_causal_prefill"):
+            CudaBackend.get_extra_aoti_compile_context_manager(
+                [CompileSpec(key="enable_tma_causal_prefill", value=b"MAYBE")]
+            )
+
+    def test_tma_causal_prefill_unsupported_target_is_disabled(self):
+        from executorch.backends.cuda.optimization_config import (
+            tma_causal_prefill_enabled,
+        )
+
+        with patch(
+            "executorch.backends.cuda.cuda_backend.cuda_targets_are_sm90_or_newer",
+            return_value=False,
+        ), self.assertLogs(
+            level="WARNING"
+        ) as logs, CudaBackend.get_extra_aoti_compile_context_manager(
+            [CompileSpec(key="enable_tma_causal_prefill", value=b"ON")]
+        ):
+            self.assertFalse(tma_causal_prefill_enabled())
+        self.assertIn("requires an SM90+ CUDA target", "\n".join(logs.output))
+
     def test_target_smem_context_is_applied(self):
         with patch(
             "executorch.backends.cuda.cuda_backend.target_smem_context",

backends/cuda/tests/test_triton_sdpa.py

@@ -13,11 +13,14 @@
 Test parametrization adapted from FlashAttention (tests/cute/test_flash_attn.py).
 """
 
+import importlib
 import itertools
 import unittest
+from unittest import mock
 
 import torch
 import torch.nn.functional as F
+from executorch.backends.cuda.optimization_config import cuda_optimization_context
 
 
 def _skip_if_no_cuda():
@@ -646,6 +649,136 @@ def test_explicit_mask_composes_with_causal_kv_len(self):
         self.assertFalse(torch.isnan(out).any())
         self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)
 
+    @unittest.skipIf(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9,
+        "TMA requires SM90+",
+    )
+    def test_tma_causal_prefill_common_head_dims(self):
+        """TMA causal prefill supports common power-of-two head dimensions."""
+        import triton
+        from triton.runtime._allocation import _allocator
+
+        sdpa_module = importlib.import_module(
+            "executorch.backends.cuda.triton.kernels.sdpa"
+        )
+
+        # Tensor descriptors require a small runtime descriptor workspace in
+        # eager mode. Inductor supplies this allocator in the production path.
+        self.addCleanup(triton.set_allocator, _allocator.get())
+        triton.set_allocator(
+            lambda size, alignment, stream: torch.empty(
+                size, dtype=torch.int8, device="cuda"
+            )
+        )
+        B, H_q, H_kv = 1, 4, 2
+        Lq, kv_len, Lk = 512, 4096, 16384
+
+        for D in (64, 128):
+            with self.subTest(D=D):
+                self.assertIsNotNone(sdpa_module._tma_prefill_config(D, Lq))
+                torch.manual_seed(D)
+                q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
+                k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
+                v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
+                kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda")
+                dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda")
+
+                with cuda_optimization_context(
+                    tma_causal_prefill=True
+                ), mock.patch.object(
+                    sdpa_module, "cuda_targets_are_sm90_or_newer", return_value=True
+                ):
+                    out_tma = self.sdpa(
+                        q,
+                        k,
+                        v,
+                        attn_mask=None,
+                        enable_gqa=True,
+                        kv_len=kv_len_t,
+                        is_causal=True,
+                    )
+                with mock.patch.object(
+                    sdpa_module, "cuda_targets_are_sm90_or_newer", return_value=False
+                ):
+                    out_existing = self.sdpa(
+                        q,
+                        k,
+                        v,
+                        attn_mask=None,
+                        enable_gqa=True,
+                        kv_len=kv_len_t,
+                        is_causal=True,
+                    )
+                out_dense = self.sdpa(
+                    q,
+                    k,
+                    v,
+                    attn_mask=dense,
+                    enable_gqa=True,
+                    kv_len=kv_len_t,
+                )
+
+                self.assertFalse(torch.isnan(out_tma).any())
+                self.assertLess(_max_abs_error(out_tma, out_dense), MAX_ABS_TOL)
+                self.assertLess(_max_abs_error(out_tma, out_existing), MAX_ABS_TOL)
+
+    @unittest.skipIf(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9,
+        "TMA requires SM90+",
+    )
+    def test_tma_falls_back_for_noncontiguous_head_dim(self):
+        """TMA descriptors are used only when Q/K/V have unit inner stride."""
+        B, H_q, H_kv, Lq, kv_len, Lk, D = 1, 4, 2, 512, 4096, 16384, 128
+        torch.manual_seed(6)
+        q = torch.randn(B, H_q, Lq, D * 2, dtype=torch.bfloat16, device="cuda")[
+            ..., ::2
+        ]
+        k = torch.randn(B, H_kv, Lk, D * 2, dtype=torch.bfloat16, device="cuda")[
+            ..., ::2
+        ]
+        v = torch.randn(B, H_kv, Lk, D * 2, dtype=torch.bfloat16, device="cuda")[
+            ..., ::2
+        ]
+        kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda")
+        dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda")
+
+        with cuda_optimization_context(tma_causal_prefill=True):
+            out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True)
+        ref = _reference_sdpa(q, k, v, attn_mask=dense)
+
+        self.assertFalse(torch.isnan(out).any())
+        self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)
+
+    @unittest.skipIf(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9,
+        "TMA requires SM90+",
+    )
+    def test_tma_fully_masked_rows_are_finite(self):
+        """Rows before the beginning of a short KV prefix return zero, not NaN."""
+        import triton
+        from triton.runtime._allocation import _allocator
+
+        self.addCleanup(triton.set_allocator, _allocator.get())
+        triton.set_allocator(
+            lambda size, alignment, stream: torch.empty(
+                size, dtype=torch.int8, device="cuda"
+            )
+        )
+        B, H_q, H_kv, Lq, kv_len, Lk, D = 1, 4, 2, 512, 128, 16384, 128
+        torch.manual_seed(7)
+        q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
+        v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
+        kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda")
+        dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda")
+
+        with cuda_optimization_context(tma_causal_prefill=True):
+            out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True)
+        ref = _reference_sdpa(q, k, v, attn_mask=dense)
+
+        self.assertFalse(torch.isnan(out).any())
+        self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)
+
     def test_mask_is_causal_matches_dense_decode(self):
         """is_causal + kv_len is a no-op vs dense for L_q==1 decode over a KV cache."""
         D, B, H_q, H_kv = 128, 1, 16, 2