[CUDA] Add Triton Q4_K-to-FP8 prefill linear

[Gasoonjia]

Stack dependency

Depends on #22414.

Summary

Add an SM90+ pure-Triton Q4_K-to-FP8 prefill linear path on top of the TMA attention PR.

• Dequantize Q4_K weights directly to FP8 and feed an FP8 matrix multiplication with BF16 output.
• Register the custom-op schema unconditionally so exported programs remain loadable on every host.
• Select the optimized path from the requested CUDA target: explicit TORCH_CUDA_ARCH_LIST takes precedence, while native export falls back to the local device capability.
• Avoid CUDA context creation at module import and explicitly exclude ROCm.
• Retain the existing low-token decode kernel and BF16 prefill fallback for pre-SM90 targets, non-BF16 activations, non-contiguous activations, and unsupported shapes.
• Apply the same input validation to the real and fake implementations.
• Add float8 dtype support to the slim AOT runtime shims.

The path is controlled by the enable_q4k_fp8_prefill CUDA compile spec and is disabled by default. Users must explicitly opt into the lossy performance path. Unsupported targets warn and retain the BF16 fallback. This change remains Triton-only and does not introduce CUDA or CuteDSL kernels.

This is the top of the reordered 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] Use Triton TMA for long-context causal prefill #22414: opt-in TMA long-context causal prefill.
4. [CUDA] Add Triton Q4_K-to-FP8 prefill linear #22192: opt-in Q4_K-to-FP8 prefill linear.

Operator-level performance

RTX 5090, representative Muse MLP projection (N=39,936, K=6,656), Q4_K dequantization plus linear:

Prefill rows M	BF16 dequant + linear	Q4_K-to-FP8 path	Speedup
128	5.033 ms	0.507 ms	9.93x
512	5.818 ms	0.945 ms	6.16x
2,048	9.182 ms	2.688 ms	3.42x

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).

Comparison	Mean token abs NLL gap	Max token gap	Mean document gap	Max document gap	Document buckets (same/close/match/mismatch)
A100 BF16 reference vs #22190	0.023180	1.462394	0.001020	0.003171	56 / 44 / 0 / 0
A100 BF16 reference vs final TMA+FP8 top	0.159700	-	0.022701	0.067222	1 / 12 / 87 / 0

Buckets use document mean-NLL gap: <1e-3, [1e-3,1e-2), [1e-2,1e-1), and >=1e-1. The FP8 path materially changes NLL, so it is intentionally opt-in rather than the default.

End-to-end top-stack performance

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.

For comparison, the lower TMA-only PR is 11.75% slower geometric mean than llama.cpp in prefill. Adding this FP8 linear path above TMA supplies the large model-wide prefill uplift; TMA supplies the additional long-context attention gain.

The isolated 8K decode 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%). PDL remains intentionally outside this stack.

Test plan

• lintrunner init && lintrunner -a: clean on both reordered branches.
• Focused final FP8/TMA routing tests: passed.
• Full int4 dispatch suite: 26 passed before the reorder; the reordered focused dispatch suite also passed.
• Accuracy is checked against _dequant_matmul with normalized max error below 5%.
• Covered compile-spec default/off/on behavior, invalid values, unsupported targets, explicit target-architecture selection, ROCm exclusion, non-BF16 fallback, and non-contiguous fallback.
• Verified that an opted-in SM90+ export captures triton.q4k_fp8_linear, while default export retains the BF16 path.
• Final opted-in PTE profile contains 208 FP8 matmul and 208 Q4_K-to-FP8 dequant kernel calls.
• git diff --check.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ 13 Pending, 2 Unrelated Failures

As of commit fc45c0d with merge base 5428092 ():

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: 185.125.190.83 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.

[digantdesai]

why does this feels like this can be lossy or clamp, let's do some evals? NLL might be a good one just comparing these two..

[Gasoonjia]

I've add a accuracy test comparing with _dequant_matmul to guard the correctness. Let me do nll verfication on top of the stack