[CUDA] Use Triton TMA for long-context causal prefill (pytorch#22414)

## Stack dependency

Depends on pytorch#22190. This PR supersedes pytorch#22193 after reordering the stack
so TMA sits directly on the split-K baseline and the optional FP8 linear
path sits above TMA.

## Summary

Use Triton tensor descriptors/TMA for long-context global causal
attention on SM90+ GPUs.

- Dispatch from semantic and hardware properties rather than a model
identity: global causal attention, a device-resident KV bound,
sufficiently long KV context, and prefill-sized queries.
- Select the optimized path from the requested CUDA target, honoring
explicit `TORCH_CUDA_ARCH_LIST`, excluding ROCm, and avoiding CUDA
context creation at module import.
- Support common transformer head dimensions 64 and 128 with separately
profiled tile, stage, and warp configurations.
- For D=128, select a smaller tile for 512-1023 query rows and a larger
tile at 1024+ rows.
- Use real output strides and keep non-unit inner-stride inputs on the
existing portable Triton path.
- Preserve finite outputs for fully masked rows with guarded
online-softmax normalization.
- Keep unsupported or non-beneficial shapes on the existing portable
Triton path; head dimension 256 currently falls back because tested TMA
configurations were slower.
- Preserve online-softmax numerical stability; this stack contains
neither PDL nor the former `phi=5` approximation.

The path is selected automatically inside the SDPA dispatcher. Eligible
prefill shapes use TMA when the AOTInductor compile target is SM90+;
decode, ROCm, pre-SM90 targets, unsupported layouts/head dimensions,
explicit masks, and non-beneficial shapes retain the portable Triton
path. The implementation remains pure Triton and is not tied to
Muse-Glimmer.

## Operator-level performance

RTX 5090, BF16 global causal GQA (`B=1, Hq=16, Hkv=2`), TMA versus the
existing Triton kernel:

| Head dim | Lq | Lkv | Existing | TMA | Speedup |
|---:|---:|---:|---:|---:|---:|
| 64 | 512 | 16,384 | 0.439 ms | 0.268 ms | 1.64x |
| 64 | 4,096 | 131,072 | 13.153 ms | 11.794 ms | 1.12x |
| 128 | 512 | 16,384 | 0.632 ms | 0.534 ms | 1.18x |
| 128 | 1,024 | 131,072 | 7.847 ms | 6.750 ms | 1.16x |
| 128 | 4,096 | 131,072 | 29.212 ms | 27.039 ms | 1.08x |

The latest post-rebase exact-stack A/B attributes the long-context gain
directly to this PR: relative to pytorch#22190, prefill improves from 2,971.25
to 3,152.94 tok/s at 32K (+6.11%) and from 2,226.87 to 2,630.43 tok/s at
130K (+18.12%). Cumulatively versus the measured main baseline, the
stack through this PR improves prefill by 8.59% at 32K and 25.24% at
130K. Decode changes by at most 0.53% relative to pytorch#22190. KV buffers
below 16K remain on the portable attention path.

## Post-rebase automatic-TMA end-to-end performance

This branch is `pytorch#22190 + TMA`, without FP8/Q8 activation or the later
linear optimizations. These figures supersede the earlier PR-summary
measurements. Stable reruns used a single otherwise-idle RTX 5090,
single-GPU `llama-cli`, the P2409802035 protocol, CUDA graphs, 512
decoded tokens, seven runs, and the mean of runs 3-7. The measurements
used the same TMA kernel/configuration as this revision; the subsequent
cleanup only removed the manual feature gate and made eligible-shape
selection automatic.

| Prompt length | ET prefill tok/s | llama.cpp prefill tok/s | ET vs
llama.cpp prefill | ET decode tok/s | llama.cpp decode tok/s | ET vs
llama.cpp decode |
|---:|---:|---:|---:|---:|---:|---:|
| 512 | 2302.19 | 2391.04 | -3.72% | 76.51 | 79.89 | -4.23% |
| 2,048 | 3311.78 | 2607.19 | +27.02% | 76.05 | 77.40 | -1.74% |
| 8,192 | 3304.35 | 2574.57 | +28.35% | 74.37 | 76.96 | -3.37% |
| 32,768 | 3152.94 | 2465.93 | +27.86% | 74.56 | 75.31 | -1.00% |
| 130,048 | 2630.43 | 2228.04 | +18.06% | 69.45 | 70.77 | -1.87% |

Across the five lengths, ET is 18.83% faster geometric mean than
llama.cpp in prefill and 2.45% slower in decode. The long-context
incremental gain in this PR comes from TMA attention; the decode
improvement in the stack comes from pytorch#22190's split-K scheduling.

## NLL evaluation

`notes_v5`, 100 documents / 154,846 next-token labels. Relative to the
A100 BF16 reference, the same/close bucket counts move from `56/44` for
pytorch#22190 to `59/41` for TMA-only. Compared directly with pytorch#22190, TMA-only
has mean token absolute NLL gap 0.014022 (max 1.158961) and mean
document gap 0.000579 (max 0.002262), with document buckets `84 same /
16 close / 0 match / 0 mismatch`.

Buckets use document mean-NLL gap: `<1e-3`, `[1e-3,1e-2)`,
`[1e-2,1e-1)`, and `>=1e-1`. The same/close bucket distribution changes
by 3 documents, within the agreed maximum difference of 5 for enabling
the optimization automatically on eligible targets.

## Test plan

- `lintrunner init` completed; linting this PR's changed files with
`lintrunner -a` is clean. The full-repository run reports only the two
existing ARM `materialize_symints` mypy failures from current `main`.
- Focused TMA routing/correctness tests: passed.
- TMA correctness compares the forced TMA path, forced portable
fallback, and dense-mask reference for head dimensions 64 and 128.
- Covered automatic SM90+ dispatch, pre-SM90/ROCm exclusion,
non-contiguous fallback, fully masked rows, and D=256 portable fallback.
- Fresh 128K TMA-only PTE/PTD export and smoke run on RTX 5090.
- 35-pair final E2E comparison on an otherwise idle dual-RTX-5090 host,
using only GPU0 for execution.
- `git diff --check`.

Author: Gasoonjia

backends/cuda/tests/test_triton_sdpa.py

@@ -13,8 +13,10 @@
 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
@@ -646,6 +648,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 mock.patch.object(
+                    sdpa_module,
+                    "_cuda_compile_target_is_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_compile_target_is_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")
+
+        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")
+
+        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

backends/cuda/triton/kernels/sdpa.py

@@ -52,6 +52,38 @@ def _is_power_of_2(n: int) -> bool:
 _SPLITK_LKV_THRESHOLD = 256
 
 
+_TMA_PREFILL_LKV_THRESHOLD = 16384
+
+
+def _cuda_compile_target_is_sm90_or_newer() -> bool:
+    """Return whether the CUDA code being compiled targets SM90 or newer."""
+    if torch.version.hip is not None:
+        return False
+
+    # Match the target selected by AOTInductor rather than assuming that the
+    # export host is also the inference target. Multi-architecture PTE builds
+    # compile each native variant independently, so this value corresponds to
+    # the variant currently being lowered.
+    from torch._inductor.codegen.cuda.compile_utils import _nvcc_arch_as_compile_option
+
+    target = _nvcc_arch_as_compile_option().removesuffix("a")
+    return target.isdigit() and int(target) >= 90
+
+
+def _tma_prefill_config(
+    head_dim: int, query_len: int
+) -> Optional[tuple[int, int, int, int]]:
+    """Return a validated TMA config for common transformer head dimensions."""
+    if head_dim == 64:
+        return 64, 64, 3, 4
+    if head_dim == 128:
+        if query_len >= 1024:
+            return 128, 64, 3, 8
+        if query_len >= 512:
+            return 64, 64, 3, 4
+    return None
+
+
 # Decode split-K occupancy target. A sweep across both production attention
 # families showed that targeting 16/9 waves gives a good balance between split
 # kernel occupancy and reduction/empty-CTA overhead. Keep the ratio integral so
@@ -780,6 +812,102 @@ def _sdpa_fwd_kernel(
     )
 
 
+@triton.jit
+def _sdpa_prefill_tma_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    O_ptr,
+    KV_LEN_ptr,
+    H_grid,
+    Lq,
+    Lk,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    sm_scale,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    NUM_STAGES: tl.constexpr,
+):
+    """TMA global-causal prefill for SM90+ and a device-resident KV bound."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H_grid
+    h_q = pid_bh % H_grid
+    h_kv = h_q // NUM_GROUPS
+
+    q_desc = tl.make_tensor_descriptor(
+        Q_ptr + b * stride_qb + h_q * stride_qh,
+        shape=[Lq, HEAD_DIM],
+        strides=[stride_qm, 1],
+        block_shape=[BLOCK_M, HEAD_DIM],
+    )
+    k_desc = tl.make_tensor_descriptor(
+        K_ptr + b * stride_kb + h_kv * stride_kh,
+        shape=[Lk, HEAD_DIM],
+        strides=[stride_kn, 1],
+        block_shape=[BLOCK_N, HEAD_DIM],
+    )
+    v_desc = tl.make_tensor_descriptor(
+        V_ptr + b * stride_vb + h_kv * stride_vh,
+        shape=[Lk, HEAD_DIM],
+        strides=[stride_vn, 1],
+        block_shape=[BLOCK_N, HEAD_DIM],
+    )
+
+    q_start = pid_m * BLOCK_M
+    offs_m = q_start + tl.arange(0, BLOCK_M)
+    offs_n_base = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, HEAD_DIM)
+    q = tl.load_tensor_descriptor(q_desc, [q_start, 0])
+    kv_len = tl.minimum(tl.load(KV_LEN_ptr), Lk)
+    absolute_q = (kv_len - Lq) + offs_m
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], tl.float32)
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
+    scale = sm_scale.to(tl.float32)
+
+    for start_n in tl.range(0, kv_len, BLOCK_N, num_stages=NUM_STAGES):
+        offs_n = start_n + offs_n_base
+        k = tl.load_tensor_descriptor(k_desc, [start_n, 0])
+        qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale
+        valid = (offs_n[None, :] < kv_len) & (offs_n[None, :] <= absolute_q[:, None])
+        qk = tl.where(valid, qk, -float("inf"))
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        safe_m = tl.where(m_ij == -float("inf"), 0.0, m_ij)
+        alpha = tl.exp(m_i - safe_m)
+        p = tl.exp(qk - safe_m[:, None])
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None]
+        v = tl.load_tensor_descriptor(v_desc, [start_n, 0])
+        acc = tl.dot(p.to(tl.bfloat16), v, acc).to(tl.float32)
+        m_i = m_ij
+
+    inv_l = tl.where(l_i > 0, 1.0 / l_i, 0.0)
+    out = acc * inv_l[:, None]
+    o_ptrs = (
+        O_ptr
+        + b * stride_ob
+        + h_q * stride_oh
+        + offs_m[:, None] * stride_om
+        + offs_d[None, :] * stride_od
+    )
+    tl.store(o_ptrs, out.to(tl.bfloat16), mask=offs_m[:, None] < Lq)
+
+
 def _validate_sdpa_inputs(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -875,6 +1003,57 @@ def _launch_pow2_kernel(
     stride_vb, stride_vh, stride_vn, stride_vd = value.stride()
     stride_ob, stride_oh, stride_om, stride_od = out.stride()
 
+    # Hopper/Blackwell: TMA materially improves the long-context global
+    # causal path. Sliding-window attention keeps the existing kernel, as do
+    # pre-SM90 GPUs. The same TMA kernel is safe for short chunks and avoids a
+    # host sync on the device-resident kv_len scalar.
+    tma_config = _tma_prefill_config(D, L_q)
+    if (
+        HAS_KV_LEN
+        and mask_is_causal
+        and not HAS_MASK
+        and query.stride(-1) == 1
+        and key.stride(-1) == 1
+        and value.stride(-1) == 1
+        and tma_config is not None
+        and L_kv >= _TMA_PREFILL_LKV_THRESHOLD
+        and L_q > 4
+        and _cuda_compile_target_is_sm90_or_newer()
+    ):
+        block_m, block_n, tma_num_stages, tma_num_warps = tma_config
+        wrap_triton(_sdpa_prefill_tma_kernel)[(triton.cdiv(L_q, block_m), B * H_q)](
+            query,
+            key,
+            value,
+            out,
+            kv_len_ptr,
+            H_q,
+            L_q,
+            L_kv,
+            stride_qb,
+            stride_qh,
+            stride_qm,
+            stride_kb,
+            stride_kh,
+            stride_kn,
+            stride_vb,
+            stride_vh,
+            stride_vn,
+            stride_ob,
+            stride_oh,
+            stride_om,
+            stride_od,
+            sm_scale,
+            BLOCK_M=block_m,
+            BLOCK_N=block_n,
+            HEAD_DIM=D,
+            NUM_GROUPS=num_groups,
+            NUM_STAGES=tma_num_stages,
+            num_warps=tma_num_warps,
+            num_stages=tma_num_stages,
+        )
+        return
+
     if pack_gqa:
         H_grid = H_kv
         Lq_packed = L_q * num_groups