[Selena Yang (selenayang888)]

Draft. This branch applies the Round-2 iteration 5 + 7 changes onto main 1f53982. It was built, tested and measured at this head together with the other Round-2 PRs on A100 and H200 (see Reported Validation and Performance At This Head). Additional evidence comes from staged runtimes built on ORT main 2bf400a.

Summary

• Instantiates the non-quantized fp16/bf16 XQA decode kernels for group size 3 (3 query heads per KV head, e.g. 24 query / 8 KV heads in Phi-3-mini, Phi-3.5-mini and Phi-4-mini) and adds the launch, shared-memory and IsSupportedGQAXqaGroupSize cases. Before this change such models never used XQA.
• Default from the device: on for fp16 on SM 8.0 (A100/A30) and SM 9.0 (H100/H200), the two measured architectures; off elsewhere. ORT_ENABLE_XQA_GROUP3=0/1 overrides the device default in both directions; ORT_ENABLE_XQA=0 still disables XQA entirely. bf16 and SM 8.6/8.9/10.x keep the previous behavior because they were not measured.
• Phi-4-mini FL package (Phi-4-mini-instruct-cuda-gpu:5) at this head, group 3 off → on: A100 decode +1.5% to +26.3% as shipped and +3.2% to +28.8% with capture; H200 +7.2% to +39.7% as shipped and +9.2% to +44.4% with capture. TTFT unchanged. Logits move by ~1e-3 relL2 (same order as switching to the memory-efficient attention backend); WikiText-2 perplexity within 0.03%.

Previous PR: GenAI #2979 (GenAI: skip the legacy Phi-3 LongRoPE rewind under graph capture). Source files are disjoint from #34234 and from the Round-1 PRs, so this PR targets main; its accumulated numbers were measured with all of them applied.

Series (Round 2, Phi-4-mini): B1 baseline → #34234 constant-branch If + capture → GenAI #2979 GenAI capture rewind → this PR (XQA group 3).

Companions: #34234 lets the Phi-4-mini package run with CUDA graph capture (its LongRoPE If blocked capture). Wide GEMV (w48) = #34212 INT4 + #33728 INT8 (opt-in).

Scope: group_query_attention.cc (default gate), group_query_attention_workspace.h (geometry predicate), xqa_loader_{fp16,bf16}_impl.cuh (instantiation and dispatch), C++ geometry test, new TestXQAGroup3Parity in test_gqa.py.

Description structure follows ORT #33932; performance tables use the columns of Mobius #813.

Reported Validation

• Built at this head together with Stage constant If branches on device and allow CUDA graph capture #34234 (ORT main 1f53982 + both commits; Release, SM 80/90, CUDA 13.0, Python bindings, onnxruntime_provider_test) and GenAI main da29a76b + GenAI #2979.
• test_gqa.py -k "TestXQAGroup3Parity or TestXQAHeadSinkParity or TestXQASlidingWindowParity": 112/112 passed on A100 and H200 (24 new group-3 cases: fp16/bf16, head 64/128/256, past 4/1000, with/without rotary).
• onnxruntime_provider_test --gtest_filter=If.*:GroupQueryAttentionWorkspaceTest.*: 14/14 passed on A100 and H200.
• Model-level gate check at this head (Phi-4-mini, teacher-forced logits, prompts 128/2048, 64 forced steps): the default is bit-identical to ORT_ENABLE_XQA_GROUP3=1 on A100 and on H200; =0 differs by relL2 1.6e-3.
• git clang-format and ruff clean.
• Raw timing protocol: 3 interleaved rounds x (2 warmups + 5 measured requests) per configuration and prompt length, llama.cpp in the same rounds, identical input IDs.
• Current PR-head CI, task-level quality (MMLU etc.), CPU performance and padded-batch coverage are not claimed.

Performance At This Head: Phi-4-mini

Runtime rt-pr = ORT main 1f53982 + #34234 + this PR; GenAI main da29a76b + GenAI #2979. Baseline = the same runtime with ORT_ENABLE_XQA_GROUP3=0 (main's behavior for this model); optimized = the device default (group 3 on). The wide GEMV (#34212 / #33728) is not on main, so w48 is not available at this head. Change = Optimized decode relative to Baseline (first table) or llama.cpp (second table). One cohort, 3 interleaved rounds, llama.cpp in the same rounds; cores 32-39, A100 GPU 2, H200 GPU 1.

A100-SXM4-80GB

Package as shipped (no capture). Baseline: a100-phi4-pr0; optimized: a100-phi4-prdef.

Tokens	Baseline TTFT ms	Optimized TTFT ms	Baseline decode tok/s	Optimized decode tok/s	Change
128	18.22	18.23	168.33	170.84	+1.49%
512	32.55	32.67	167.84	171.32	+2.07%
2048	96.03	96.00	156.65	169.16	+7.99%
4096	193.28	193.34	109.94	130.58	+18.77%
7680	380.68	381.16	124.82	157.61	+26.27%

Decoder capture on. Baseline: a100-phi4-prcap0; optimized: a100-phi4-prcap.

Tokens	Baseline TTFT ms	Optimized TTFT ms	Baseline decode tok/s	Optimized decode tok/s	Change
128	19.01	18.94	187.25	193.25	+3.20%
512	32.82	32.82	187.68	194.12	+3.43%
2048	96.63	96.53	174.80	190.41	+8.93%
4096	193.96	193.86	144.59	182.60	+26.29%
7680	382.47	381.95	135.88	175.06	+28.83%

Same captured optimized series versus llama.cpp a100-phi4-llama (identical input IDs, same rounds):

Tokens	llama.cpp TTFT ms	Optimized TTFT ms	llama.cpp decode tok/s	Optimized decode tok/s	Change
128	21.41	18.94	219.10	193.25	-11.80%
512	55.37	32.82	218.27	194.12	-11.06%
2048	215.21	96.53	211.47	190.41	-9.96%
4096	442.00	193.86	200.43	182.60	-8.89%
7680	881.13	381.95	189.10	175.06	-7.42%

• Prefill ms (shipped): baseline 18.11 / 32.45 / 95.92 / 193.14 / 380.57; optimized 18.13 / 32.56 / 95.90 / 193.22 / 381.00; llama.cpp 21.07 / 54.98 / 214.74 / 441.59 / 880.67.

H200

Package as shipped. Baseline: h200-phi4-prdef (group 3 off at that time); optimized: h200-phi4-pr1 (=1, identical to the updated default, see the gate check).

Tokens	Baseline TTFT ms	Optimized TTFT ms	Baseline decode tok/s	Optimized decode tok/s	Change
128	8.92	8.92	255.39	273.87	+7.24%
512	13.56	13.51	252.16	274.36	+8.80%
2048	37.98	37.94	233.91	270.29	+15.56%
4096	75.99	76.19	180.55	224.73	+24.47%
7680	161.67	161.59	182.65	255.14	+39.69%

Decoder capture on. Baseline: h200-phi4-prcap; optimized: h200-phi4-prcap1.

Tokens	Baseline TTFT ms	Optimized TTFT ms	Baseline decode tok/s	Optimized decode tok/s	Change
128	9.02	8.93	297.70	325.09	+9.20%
512	13.62	13.57	297.41	327.99	+10.28%
2048	38.26	37.95	274.09	322.19	+17.55%
4096	76.68	76.60	241.83	310.52	+28.40%
7680	162.35	162.75	209.30	302.31	+44.44%

Same captured optimized series versus llama.cpp h200-phi4-llama:

Tokens	llama.cpp TTFT ms	Optimized TTFT ms	llama.cpp decode tok/s	Optimized decode tok/s	Change
128	11.64	8.93	358.74	325.09	-9.38%
512	29.37	13.57	355.46	327.99	-7.73%
2048	111.63	37.95	342.10	322.19	-5.82%
4096	229.99	76.60	327.10	310.52	-5.07%
7680	457.13	162.75	314.53	302.31	-3.89%

• Prefill ms (shipped): baseline 8.85 / 13.48 / 37.91 / 75.90 / 161.59; optimized 8.84 / 13.44 / 37.86 / 76.10 / 161.51; llama.cpp 11.35 / 29.07 / 111.34 / 229.63 / 456.63.
• Why H200 is on at this head but was off in the staged runs: after 2bf400a, [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops #34164 made GQA's auto mode prefer FlashAttention over cuDNN SDPA for short shared-buffer steps (decode). For this model on H200, FlashAttention split-KV costs 39 + 11 us per layer per token at 4096 tokens versus 9.8 + 2.2 us for cuDNN SDPA (nsys, uncaptured), so main's default decode is 15-27% slower than the staged runtimes. Group-3 XQA recovers most of it. Forcing cuDNN SDPA (ORT_ENABLE_XQA_GROUP3=0 ORT_ENABLE_CUDNN_FLASH_ATTENTION=1) is still 7.5% / 6.1% faster than group-3 XQA at 128 / 4096 tokens (captured smoke, 2 rounds), so revisiting [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops #34164's decode policy for group-3 models is worth a separate look.

Staged evidence below adds the Round-1 wide GEMV and shows where this model stands against llama.cpp with all Round-1 and Round-2 changes.

Staged Evidence: ORT 2bf400a + Round-1 PRs

Baseline = it6: B1 + #34234 + GenAI #2979 + this PR's kernels with the previous opt-in default (group 3 off). Optimized = it7: same plus a device default (SM 8.0 at the time). Only the group-3 default differs (one libonnxruntime_providers_cuda.so).

A100-SXM4-80GB (staged)

Package as shipped (no capture, no flags). Baseline: a100-phi4-it6def; optimized: a100-phi4-it7def.

Tokens	Baseline TTFT ms	Optimized TTFT ms	Baseline decode tok/s	Optimized decode tok/s	Change
128	18.13	18.15	168.94	175.03	+3.61%
512	32.49	32.48	169.01	176.22	+4.27%
2048	95.85	96.00	156.99	171.91	+9.50%
4096	193.20	193.49	110.15	131.51	+19.38%
7680	380.89	381.29	124.70	158.62	+27.20%

Decoder capture on (runtime option, package files unchanged). Baseline: a100-phi4-it6cap; optimized: a100-phi4-it7cap.

Tokens	Baseline TTFT ms	Optimized TTFT ms	Baseline decode tok/s	Optimized decode tok/s	Change
128	18.72	18.84	189.11	197.67	+4.52%
512	32.59	32.57	188.81	199.08	+5.44%
2048	96.32	96.24	175.88	193.11	+9.80%
4096	193.90	193.78	145.39	183.79	+26.42%
7680	382.14	382.06	136.56	176.12	+28.97%

Best A100 configuration (capture + w48) versus llama.cpp a100-phi4-llama (identical input IDs, same rounds):

Tokens	llama.cpp TTFT ms	Optimized TTFT ms	llama.cpp decode tok/s	Optimized decode tok/s	Change
128	21.47	18.66	216.73	238.31	+9.95%
512	55.84	32.66	216.40	238.39	+10.16%
2048	216.92	96.28	210.11	230.06	+9.49%
4096	445.12	193.95	200.45	213.55	+6.54%
7680	879.80	381.45	190.34	202.27	+6.27%

• it7 capture + w48 (device default) vs the previous opt-in it6 capture + w48 + ORT_ENABLE_XQA_GROUP3=1: same kernels, greedy outputs 128/128 identical at every length, decode -0.24% to +0.07% at 128-2048 and -1.6% at 4096/7680 (shared-host noise).
• Without capture the package is still behind llama.cpp (-16.7% to -34.4%); see Known Limitations.
• Prefill ms (package as shipped): baseline 18.03 / 32.40 / 95.75 / 193.08 / 380.75; optimized 18.05 / 32.39 / 95.91 / 193.38 / 381.19; llama.cpp 21.08 / 55.31 / 216.47 / 444.65 / 879.33.

H200 (staged)

On the staged base the H200 decode default was cuDNN SDPA, and forcing group 3 on (capture + w48, it6 cohort) cost -4.34 / -3.83 / -5.07 / -6.81 / -5.61% at 128 / 512 / 2048 / 4096 / 7680 tokens, which is why the staged default was SM 8.0 only (the same switch on A100: +4.81 / +5.56 / +10.24 / +30.18 / +33.05%). Best staged H200 configuration (capture + w48, group 3 off, cuDNN SDPA) versus llama.cpp h200-phi4-llama:

Tokens	llama.cpp TTFT ms	Optimized TTFT ms	llama.cpp decode tok/s	Optimized decode tok/s	Change
128	11.51	8.95	357.72	436.05	+21.90%
512	29.41	13.56	353.38	434.39	+22.92%
2048	112.31	38.08	341.24	425.04	+24.56%
4096	230.89	76.68	327.90	406.72	+24.04%
7680	458.83	162.14	313.36	390.69	+24.68%

Shared-host medians are descriptive, not statistical significance; negative rows are kept.

Controls And Configuration

Runtimes (all source-built):

• At this head (rt-pr): ORT main 1f53982 + Stage constant If branches on device and allow CUDA graph capture #34234 + this PR; GenAI main da29a76b + GenAI #2979. Release, SM 80/90, CUDA 13.0, cuDNN 9.8 runtime.
• Staged (rt-it6 / rt-it7, differ only in libonnxruntime_providers_cuda.so): ORT main 2bf400a + Round-1 iterations (Add opt-in wide INT4 M1 GEMV with uint8 zero points #34212-Add opt-in chunked LinearAttention prefill using GatedDeltaNet #34215, with [CUDA] MatMulNBits 8-bit: wide fp16 M=1 kernel for decode GEMVs #33728's INT8 wide kernel) + Harden CUDA graph ownership, retained-buffer lifetime, and failed-capture recovery #33932 + [CUDA] Add opt-in single-warp INT4 GEMV for FP16 decode #33952 (= B1) + Stage constant If branches on device and allow CUDA graph capture #34234 + XQA group 3 (rt-it6: opt-in; rt-it7: SM 8.0 default). GenAI main faa19c7a + Weird requested input shape after running bert_model_optimization.py #2889 + building ONNXruntime with DNNL cause an error #2899 + Add section for distro installations #2900 + GenAI #2979 (0.20.0.dev0).
• Model: Foundry Local Phi-4-mini-instruct-cuda-gpu:5 (CLI 0.11.0), unchanged: INT4 + INT8 block 32 with uint8 zero points, INT8 LM head, LongRoPE If, GQA 24 / 8 heads, head size 128, fp16 KV. --capture turns on the decoder's enable_cuda_graph at runtime.
• Workload: token prefixes 128/512/2048/4096/7680 of a fixed source text (sha256 3496ea55...), context 8192, exactly 128 generated tokens, batch 1, greedy, FPA1, 8 CPU threads pinned to cores 32-39; A100 physical GPU 2, H200 physical GPU 1. Fresh generator per request; tokenization and generator creation excluded. 3 rounds with rotating order x (2 warmups + 5 measured); medians of 15 measured requests.
• llama.cpp rev 5d4c8eeb, microsoft/Phi-4-mini-instruct @cfbefacb converted to Q4_K_M (2,493,840,224 B), CUDA graphs, flash attention, fp16 KV, the same prompt token IDs. Quantization is not quality-matched.

Correctness Evidence And Limits

Teacher-forced logits, group 3 on vs off (same build, staged, A100):

Prompt	Uncaptured max relL2 / argmax agreement	Captured max relL2 / argmax agreement
128	1.31e-3 / 98.4%	1.20e-3 / 95.3%
2048	1.66e-3 / 100%	2.79e-3 / 100%

This is the same magnitude as switching to the existing memory-efficient attention backend (max relL2 1.1e-3 / 3.3e-3 at prompts 128 / 2048, same harness).

WikiText-2 (raw, test split) decode-path perplexity on A100: prefill prefix tokens, then append one token per step and score it with the previous step's logits, so every scored token goes through decode attention.

Window prefix:forced	Uncaptured off → on	Captured off → on
512:1536	7.0022 → 7.0008 (-0.020%), argmax 99.5%	7.1595 → 7.1609 (+0.020%), argmax 99.4%
2000:2000 (uncaptured) / 3000:2000 (captured)	12.0579 → 12.0612 (+0.027%), 99.8%	10.8172 → 10.8184 (+0.012%), 99.6%
6000:2000	5.5595 → 5.5602 (+0.012%), 99.95%	6.8777 → 6.8776 (-0.002%), 99.75%

Greedy outputs (first measured request, 128 tokens), group 3 on vs off at this head: A100 shipped first differs at token 10 for the 128-token prompt, identical at the other lengths; A100 captured differs at tokens 39 (2048) and 90 (4096); H200 shipped differs at token 25 (512); H200 captured at token 90 (4096). Staged: A100 it7 capture + w48 vs it6 capture + w48 + switch 128/128 at all lengths (same kernels).

Known Limitations

• The default changes on SM 8.0 and 9.0 (fp16) for every model with 3 query heads per KV head, but the evidence covers Phi-4-mini only. Other group-3 models (Phi-3/3.5-mini) and task-level quality are not measured.
• Not measured: SM 8.6/8.9 (A10, L4, RTX 30/40), SM 10.x, bf16 models. They keep the previous default. H100 shares SM 9.0 with H200 but was not run.
• On H200 the gain is relative to main's current FlashAttention decode default ([CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops #34164). cuDNN SDPA, the default before [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops #34164, is still faster than group-3 XQA for this model; if [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops #34164's decode policy changes, the SM 9.0 default here should be re-measured.
• XQA under capture sizes its multi-block split from the KV capacity (8192) rather than the current length; adapting the split to the device-side length may add more at short context (not attempted here).
• The Phi-4-mini package does not ship enable_cuda_graph, so the capture rows require the runtime option.

Required Tests

• TestXQAGroup3Parity (new): fp16 + bf16, head 64/128/256, past 4/1000, with and without rotary, group 3 forced on, against the PyTorch reference.
• GroupQueryAttentionWorkspaceTest.XqaGeometryPredicatesCoverSupportedDomain updated for group 3.
• Still needed: a test that the device default selects XQA on SM 8.0/9.0 and not elsewhere (verified at model level via bit-identity on A100 and H200, not in CI).

Resource Observations

GPU	Config (staged runtimes)	Package bytes	Peak process RSS bytes	Sampled peak VRAM bytes
A100	B1 (shipped)	4147570154	1478553600	16580083712
A100	it7 default (shipped)	4147570154	1445363712	16580083712
A100	it7 capture + w48	4147570154	1422139392	15512633344
H200	B1 (shipped)	4147570154	1607434240	16819159040
H200	it7 default (shipped)	4147570154	1583943680	16819159040
H200	it7 capture + w48	4147570154	1525772288	15745417216

Separate one-round resource pass (2 warmups + 1 measured request per length, fresh process); own-process NVML peak sampled every 20 ms, RSS high-water mark; excluded from the timing tables. XQA group 3 adds no persistent allocation; peak VRAM is dominated by the arena. llama.cpp resources were not sampled in this round (GGUF 2,493,840,224 B). The new instantiations add to libonnxruntime_providers_cuda.so size and compile time (fp16 + bf16 x supported head sizes).

Existing Work And Duplication

GitHub search (open and merged) for XQA group sizes: #32108 (merged) adds group 6 to paged XQA; #33759 (open) covers group 8 in a different split-decode kernel; #32071 (open) integrates the CUDA workspace with activation memory patterns and touches the same workspace header. None adds non-paged group 3. Round-1 #34214 changes xqa/mha_impl.cuh only (disjoint).

WORK.md Acceptance Criteria

Criterion	Evidence/status and remaining work
Exporter	Published package unchanged; no graph surgery.
Quality	NOT QUALIFIED: MMLU (and other task benchmarks) versus llama.cpp Q4_K_M not measured. Perplexity within 0.03% and logits within ~1e-3 of the previous default are reported above; they do not replace the <1% per-benchmark gate.
Performance	At this head: gain at every length on A100 and H200, with and without capture (+1.5% to +44.4%). With the Round-1 wide GEMV (staged) the best A100 config beats llama.cpp by 6.3-10.2% and H200 by 21.9-24.7%; at this head without it Phi-4-mini still trails llama.cpp. CPU EP not affected or measured.
Metrics	TTFT, prefill ms and decode tok/s reported separately.
Hardware/kernel changes	Both A100 and H200 measured at this head; gains on both.
Fair comparison	Identical token IDs, context 8192, 128 outputs, batch 1, 8 threads; INT4/INT8 block 32 vs Q4_K_M, fp16 KV and capture settings disclosed.
Correctness	Short and long prompts, captured and uncaptured, logits, perplexity and greedy outputs above. Padded batches and multi-turn not covered for this change.
Resources	Package size, RSS and sampled VRAM above; no agreed budget supplied.
Compatibility/defaults	Changes the default only for fp16 group-3 GQA on SM 8.0 and 9.0; ORT_ENABLE_XQA_GROUP3=0 restores the previous behavior. Experimental per WORK.md until quality is validated.
Reproducibility	Runtimes, library hashes, prompt IDs, raw samples and warmups retained (author's review package).
PR description	Configurations, A100/H200 tables, quality status, correctness, resources, limitations included.
Release gate	BLOCKED until CI, task-level quality and broader model coverage pass.
Avoid duplication	See above.

Quality Results

Benchmark	llama.cpp Q4_K_M	This candidate	Relative loss	Gate
MMLU	Not measured	Not measured	Unknown	BLOCKED
WikiText-2 decode-path perplexity (proxy, vs previous ORT default)	n/a	within 0.03%	n/a	Informational

Remaining Before Ready For Review

1. Pass CI on SM80 and SM90; add a CI-level check of the device default.
2. MMLU for Phi-4-mini (and a Phi-3.5-mini spot check) with the default on A100 and H200 versus llama.cpp Q4_K_M.
3. Coordinate with the owners of [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops #34164 on the decode policy for group-3 models on SM 9.0 (cuDNN SDPA is fastest there for this model).
4. Decide with reviewers whether to keep the device default or revert to opt-in.

Evidence

Round-2 iteration log (ITERATIONS.md, iterations 5 and 7 and the PR-head section), rebased-head cohort prh-{a100,h200}.tar.gz, staged cohorts it6/it7-{a100,h200}.tar.gz (timing JSONL with token IDs, logs, runtime hashes), it7-{a100,h200}-checks.tar.gz (gate and perplexity logs), PR-head verification logs, resource pass res2. Kept in the author's review package; available on request.