[CUDA] Enable XQA by default for FP16/BF16 GQA

[namgyu-youn]

Description

Enable XQA by default for non-quantized FP16/BF16 GQA. XQA is a TensorRT-LLM-derived fused decode kernel (seq_len=1) that also fuses RoPE + KV-append into a single pass. Requires SM80+, shared KV buffer, no softcap.

Also add group_size=5 XQA support (e.g. Qwen3-14B: 40 Q-heads / 8 KV-heads) for the Qwen series.

Affected models

Model	group_size	Before	After
Qwen3-8B	4	FlashDecode (opt-in XQA)	XQA default
Qwen3-14B	5	No XQA	XQA default
Qwen2-72B	8	FlashDecode (opt-in XQA)	XQA default
Qwen2-7B	7	No XQA	No XQA (unsupported tile width)

Performance Result

repro:

LD_PRELOAD=/usr/local/cuda-12.8/lib64/libcudart.so.12 \
  LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:/workspace/onnxruntime/.venv/lib/python3.12/site-packages/nvidia/cudnn/lib \
  PYTHONPATH=build/Release:onnxruntime/test/python/transformers \
  python onnxruntime/test/python/transformers/benchmark_gqa.py

Benchmark: Qwen3-14B (40Q/8KV heads, head_dim=128, fp16, batch=1, SM80, CUDA12.8)

past_seq_len	XQA off	XQA on	Speedup
256	0.124 ms	0.092 ms	1.30×
512	0.122 ms	0.089 ms	1.33×
1024	0.133 ms	0.094 ms	1.42×
2048	0.136 ms	0.095 ms	1.43×
4096	0.152 ms	0.103 ms	1.42×
8191	0.182 ms	0.116 ms	1.57×

Speedup grows with context length (1.30× → 1.57×). No regression in the prompt/prefill path.

Testing

• pytest onnxruntime/test/python/transformers/test_gqa.py -v
• ORT_ENABLE_XQA=0 pytest onnxruntime/test/python/transformers/test_gqa.py -v (opt-out smoke test)

[namgyu-youn]

@microsoft-github-policy-service agree

[namgyu-youn]

update via 1959cec (this PR): now XQA doesn't require "check if quantized", so dropped dead-code (is_quantized); CI would be green now.

[tianleiwu]

Thanks for the XQA default-on + group_size=5 work. Functionally this looks sound: the env override semantics are preserved (ORT_ENABLE_XQA=0 still wins), and ineligible shapes/group sizes fall back safely via data.use_xqa. Replacing 64 % group_size == 0 with the explicit {1,2,4,5,8,16,32} allowlist is also a nice correctness fix — it drops group_size=64, which previously passed the gate but has no case 64 in the launcher switch and would have returned a hard FAIL instead of falling back.

Two follow-ups (both non-blocking):

1. Missing dedicated test for group_size=5. This is the first non-power-of-2 group (the M_TILESIZE=8 tile leaves 3 unused rows), but it isn't covered by any XQA parity test — gqa_xqa_test_cases() loops [4,8,16,32] and gqa_xqa_head_sink_test_cases() loops [1,4,8]. Please add a group_size=5 case (e.g. 40 Q / 8 KV) to a parity test in test_gqa.py so the new kernel path is validated, not only benchmarked. (test_gqa.py isn't in this diff, hence the note here.)

2. Stale docstring: gqa_xqa_head_sink_test_cases() still documents 64 % group_size == 0, which no longer matches the new allowlist.

Inline comments below for the constructor comment drift and the allowlist.

[tianleiwu]

@namgyu-youn, there is build error:

 error: unused variable ‘is_quantized’ [-Werror=unused-variable]
  113 |   bool is_quantized = (k_quant_type_ != KVQuantizationType::NONE || v_quant_type_ != KVQuantizationType::NONE);
      |        ^~~~~~~~~~~~

[namgyu-youn]

@namgyu-youn, there is build error:

 error: unused variable ‘is_quantized’ [-Werror=unused-variable]
  113 |   bool is_quantized = (k_quant_type_ != KVQuantizationType::NONE || v_quant_type_ != KVQuantizationType::NONE);
      |        ^~~~~~~~~~~~

@tianleiwu sorry it's my bad; it should be green now. Could you trigger CI again?

[tianleiwu]

There is another error:

unused variable ‘kIsFp16OrBf16’ [-Werror=unused-variable]
  402 |   constexpr bool kIsFp16OrBf16 = std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>;

[namgyu-youn]

Update via bb16aa8 (this PR): Dropped dead checker for XPU enable, which has been breaking CI (lint).

@tianleiwu could you please trigger CI again? It should be green now I believe.

[tianleiwu]

@namgyu-youn, please run lintrunner to format the code:

pip install -r requirements-dev.txt
lintrunner init
lintrunner -a

[namgyu-youn]

@namgyu-youn, please run lintrunner to format the code:

pip install -r requirements-dev.txt
lintrunner init
lintrunner -a

done; @tianleiwu sorry for my bad. Could you please trigger CI again?