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

[Gasoonjia]

Stack dependency

Depends on #22192.

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 the original 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 dispatch is not tied to Muse-Glimmer and can benefit any model whose attention semantics and shape satisfy the validated conditions.

This is the top of the following stack:

1. [CUDA] Elide dense causal masks in long-context prefill #22189: D111 causal-mask elision.
2. [CUDA] Tune hardware-aware split-K decode scheduling #22190: hardware-aware online-softmax split-K scheduling.
3. [CUDA] Add Triton Q4_K-to-FP8 prefill linear #22192: opt-in Q4_K-to-FP8 prefill linear.
4. [CUDA] Use Triton TMA for long-context causal prefill #22193: opt-in TMA long-context causal prefill.

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 rest of the stack fixed, TMA alone improves 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 path.

NLL evaluation

notes_v5, 100 documents / 154,846 next-token labels. The supplied .pt stores one token chunk per document; labels are chunk[1:] for inputs chunk[:-1] (all chunks begin with BOS 200000).

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. Relative to the A100 BF16 reference, the same/close bucket counts move from 56/44 for #22190 to 59/41 for TMA-only. Under the agreed conservative bucket-count rule, the optimization therefore remains explicit opt-in.

End-to-end comparison with llama.cpp

Both optional prefill paths enabled. 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	3216.25	2877.14	+11.79%	77.69	78.34	-0.83%
2,048	5378.22	3615.14	+48.77%	77.21	77.12	+0.11%
8,192	5740.71	4060.96	+41.36%	74.77	76.62	-2.42%
32,768	5378.54	4012.94	+34.03%	74.70	74.94	-0.32%
130,048	4167.30	3455.06	+20.61%	70.63	69.54	+1.56%

Across the five lengths, ET prefill is 30.61% faster geometric mean than single-GPU llama-cli; decode is effectively tied overall at -0.39% geometric mean.

The isolated 8K miss is primarily launch overlap rather than SDPA kernel latency. In a matched Nsight trace, ET's stable decode graph is serialized (12.421 ms summed kernel time, 12.742 ms wall time), while llama.cpp overlaps up to seven kernels (16.114 ms summed kernel time, 12.460 ms wall time). llama.cpp enables Programmatic Dependent Launch by default: an idle-GPU 7-run A/B measured 76.58 tok/s with PDL and 75.24 tok/s with GGML_CUDA_PDL=0 (+1.78%). ET reaches 74.77 tok/s, only 0.63% below the no-PDL comparison. PDL remains intentionally outside this stack.

Test plan

• lintrunner init && lintrunner -a on [CUDA] Add Triton Q4_K-to-FP8 prefill linear #22192 and on the complete top stack: clean.
• Focused final top-stack FP8/TMA routing tests: 13 passed.
• Full int4 dispatch suite: 26 passed.
• Standard Triton SDPA and dedicated TMA correctness coverage passed before the test-only ROCm follow-up.
• 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, unsupported-target fallback, ROCm exclusion, non-contiguous fallback, fully masked rows, and restored D=256 portable coverage.
• 35-pair final E2E comparison on an otherwise idle dual-RTX-5090 host, using only GPU1 for execution.
• git diff --check.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[digantdesai]

curious how much you are getting just from using TMA?

[Gasoonjia]

We can get like 15%-ish operator-level perf improvement. Detailed can be found in PR summary