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

[Gasoonjia]

Stack dependency

Depends on #22190. This PR supersedes #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 controlled by the enable_tma_causal_prefill CUDA compile spec and is disabled by default. Users must explicitly opt in; unsupported targets warn and 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

Holding the FP8 linear path fixed in the earlier isolation run, TMA improved end-to-end prefill from 5,019 to 5,388 tok/s at 32K (+7.4%) and from 3,341 to 4,172 tok/s at 130K (+24.8%). Inputs below 16K remain on the portable attention path.

TMA-only end-to-end performance

This branch is #22190 + TMA, without the upper FP8 linear optimization. Single otherwise-idle RTX 5090, single-GPU llama-cli, P2409802035 protocol, ET and llama.cpp interleaved per repetition, runs 3-7 mean:

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	2178.74	2844.52	-23.41%	76.52	78.20	-2.14%
2,048	3528.60	3605.18	-2.12%	75.65	77.02	-1.78%
8,192	3785.23	4091.16	-7.48%	73.07	76.46	-4.43%
32,768	3628.39	4064.62	-10.73%	73.30	74.90	-2.13%
130,048	3034.08	3509.10	-13.54%	70.36	69.42	+1.35%

Across the five lengths, TMA-only ET is 11.75% slower geometric mean than llama.cpp in prefill and 1.84% slower in decode. The result also shows the division of labor in the stack: TMA improves long-context attention, while the upper FP8 linear PR supplies the large model-wide prefill uplift.

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 #22190 to 59/41 for TMA-only. Compared directly with #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. Under the agreed conservative bucket-count rule, the optimization remains explicit opt-in.

Test plan

• lintrunner init && lintrunner -a: clean on this branch.
• Focused CUDA export and 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 compile-spec default/off/on behavior, invalid values, unsupported targets, 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.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22414

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 2 Unrelated Failures

As of commit 8da3f29 with merge base 5428092 ():

NEW FAILURE - The following job has failed:

• pull / android / run-emulator (gh)
  WARNING | Device 'cache' does not have the requested snapshot 'default_boot'

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-moshi-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  E: Failed to fetch http://security.ubuntu.com/ubuntu/pool/main/libs/libssh/libssh-gcrypt-4_0.9.6-2ubuntu0.22.04.7_amd64.deb 404 Not Found [IP: 91.189.91.81 80]

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (trunk failure)

This comment was automatically generated by Dr. CI and updates every 15 minutes.