ONNX Attention CUDA: Performance Optimization Opportunities vs Contrib GQA

[titaiwangms]

ONNX Attention Op — Performance Gap vs Contrib Op GQA

Parent issue: #27516
Related: #28351 (spec compliance gaps)

After #27992 merged, ONNX Attention CUDA uses the same Flash/MEA kernels as Contrib GQA. However, decode latency is ~34-111μs/token higher due to spec-level constraints and missing optimizations.

Baseline benchmark recipe needed. The ~34-111μs/decode figure comes from informal profiling during #27992 development and lacks reproducible context (GPU model, dtype, batch×heads×head_size×seq_len, KV cache size, single-stream vs concurrent). PR #27992 description does not preserve the raw numbers. TODO: capture a reproducible benchmark recipe (script + machine spec + model config) and link it here before relying on these numbers for prioritization. See "Baseline Benchmark Recipe" checklist below.

Performance Gap Summary

#	Area	Impact	Root Cause	Code Location	Mitigation	Success Criteria
1	KV cache concat (decode w/ past)	~34-111μs/decode	ONNX spec: immutable past tensors → full copy each step	Flash decode concat: core/providers/cuda/llm/attention.cc:398 (LaunchConcatNewToPastKV) · MEA decode concat: attention.cc:680 · Source: contrib_ops/cuda/bert/attention_kv_cache.cu (LaunchConcatNewToPastKV impl)	TensorScatter mode (external cache) bypasses this	TensorScatter decode latency within ±5% of Contrib GQA decode at matched config
2	No XQA / DecoderMaskedMHA kernel	2-5x slower decode	Requires shared-buffer KV cache (blocked by spec for past/present path)	Dispatch site (TensorScatter Flash path): attention.cc:297-310 (Path 1: nonpad_kv_seqlen → mha_fwd_kvcache) · XQA reference impl: contrib_ops/cuda/bert/xqa/xqa_loader.h and xqa/mha.h · GQA dispatch precedent: contrib_ops/cuda/bert/group_query_attention.cc:318-335 (XQA eligibility predicate)	Add XQA dispatch tier for TensorScatter mode (S_q==1, sm≥80, no softcap, no smooth-softmax, no local window — same predicates as GQA)	TensorScatter+XQA decode within ±10% of GQA+XQA decode at matched config
3	Extra concat kernel launch	~5-10μs	Explicit concat before Flash call (separate from prep)	Same call sites as item 1 (attention.cc:398, 680)	Fuse into prep kernel (see item 4)	Eliminate 1+ kernel launch per decode step (verify via Nsight Systems)
4	No fused prep kernel	2-3 extra launches	Prep is QKV reshape + (optional) RoPE + KV concat as separate launches	TODO marker: core/providers/cuda/llm/attention.cc:345-348 · GQA fused-prep precedent (LaunchUnpackRoPEAppend): contrib_ops/cuda/bert/group_query_attention_impl.cu:126, 589, 818, 903 · Description in group_query_attention_qdq.cuh:200	Fuse QKV reshape + KV concat (and RoPE if added later) mirroring GQA's UnpackRoPEAppend	Reduce decode prep from 4-6 kernel launches → ~2 (per the existing TODO comment)

Reference implementations to study

• GQA contrib op (perf baseline): onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc (op skeleton, XQA eligibility) and onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu (LaunchUnpackRoPEAppend fused prep at lines 126, 589, 818, 903 — direct precedent for items 3 and 4).
• XQA kernel: onnxruntime/contrib_ops/cuda/bert/xqa/ — xqa_loader.h is the public entry, mha.h declares the launch API. Loader specializations (e.g., xqa_loader_fp16.cu, xqa_loader_bf16_int8_256.cu) cover the supported dtype/quantization matrix.
• DecoderMaskedMHA (legacy decode kernel): onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu — earlier optimized decode kernel; useful for comparing assumptions XQA can/cannot relax.

When ONNX Attention already matches GQA:

• TensorScatter mode (external cache via nonpad_kv_seqlen) — no concat, near-parity. Dispatch: attention.cc:297-310 (Flash) and attention.cc:540-560 comment block (MEA).
• Prefill (large S_q) — Flash kernel dominates, overhead amortized.

Proposed Optimizations (priority order)

1. XQA kernel for TensorScatter decode — The external-cache path already has shared-buffer semantics. Integrating XQA via the same eligibility predicate used by group_query_attention.cc:318-335 would close the 2-5x gap for the most performance-critical path. Reuse the GQA-side scratch allocation pattern (group_query_attention.cc:306-307).

2. Fused KV cache preparation — Combine past_key/past_value concat + (future) RoPE into a single kernel modeled on LaunchUnpackRoPEAppend. Replaces the explicit LaunchConcatNewToPastKV calls at attention.cc:398, 680.

3. In-place KV buffer reuse — If the graph optimizer can prove past_key/past_value have no other consumers downstream, reuse the buffer for present_key/present_value (avoid O(seq_len) copy). Requires a graph-transform pass; not a kernel change.

4. Fused prep kernel — Consolidate QKV reshape, type conversion, and layout transforms into one kernel (TODO already noted at attention.cc:345-348). Subsumes item 2 once layout transforms are folded in.

Baseline Benchmark Recipe (TODO before relying on the table's μs numbers)

• Pick a representative decode config (e.g., Llama-3 8B: B=1, H=32, H_kv=8, head_size=128, S_kv ramps 128→4096)
• Capture GPU model, driver, CUDA version, ORT build flags
• Script: standalone harness invoking ONNX Attention vs Contrib GQA at matched config, single-stream
• Profile with Nsight Systems for per-kernel breakdown (attribute the μs to KV concat vs Flash kernel vs prep)
• Repeat at fp16 and bf16
• Commit script under onnxruntime/test/python/transformers/benchmark_*.py and link from this issue

Notes

• The immutable-past constraint is fundamental to the ONNX spec design (functional semantics). Items 1-2 work within this constraint via the TensorScatter mode escape hatch. Item 3 requires runtime graph analysis.
• For users who need maximum decode performance TODAY, recommend using TensorScatter mode (external cache management) — this is the only path that can ever reach Contrib GQA parity without spec changes.

[namgyu-youn]

Reproducible benchmark for the decode-latency gap

@titaiwangms this addresses the Baseline Benchmark Recipe TODO in the issue description:

The ~34-111μs/decode figure comes from informal profiling during #27992 development and lacks reproducible context (GPU model, dtype, batch×heads×head_size×seq_len, KV cache size, single-stream vs concurrent). [...] TODO: capture a reproducible benchmark recipe (script + machine spec + model config) and link it here before relying on these numbers for prioritization.

Recipe + measured data below; the measurements also surface one cost missing from the gap table (buffer-length-dependent dispatch, last section). Script: #29684. Raw artifacts (CSVs, nsys captures, logs): https://github.com/namgyu-youn/onnxruntime/releases/tag/bench-28352-artifacts.

Setup: five single-stream decode arms at matched config — S_q=1, causal, no RoPE, no mask, no softcap, io-binding, triton.testing.do_bench:

arm	op	kernel tier (pinned via env)
gqa_xqa	contrib GQA, shared KV buffer	XQA
gqa_flash	contrib GQA	Flash (ORT_ENABLE_XQA=0, ORT_ENABLE_CUDNN_FLASH_ATTENTION=0)
gqa_cudnn	contrib GQA	cuDNN SDPA (SM90+ non-XQA default)
attn_past	ONNX Attention op23, past/present	Flash
attn_scatter	ONNX Attention op24, external cache (in-place TensorScatter + nonpad_kv_seqlen, no present_*)	Flash

Machine: RTX PRO 6000 Blackwell (SM120, 96 GB), driver 580.95.05, CUDA 13.0, onnxruntime-gpu 1.29.0.dev20260708001 (2cd570577f). Llama-3-8B decode shape: B=1, H_q=32, H_kv=8, head 128, KV buffer 8192.

ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1 python benchmark_onnx_attention_vs_gqa.py --sanity
python benchmark_onnx_attention_vs_gqa.py --dtype float16 --csv sweep.csv   # + bfloat16
nsys profile --export=sqlite python benchmark_onnx_attention_vs_gqa.py --profile --arms <arm> --past-seq-len 2048

Sanity: all arms produce identical outputs (max |diff| ≤ 8e-6 fp16 / 3.2e-5 bf16); backends verified per arm (debug-info print for GQA, nsys kernel names for ONNX Attention).

Results

e2e µs/token, fp16 (bf16 equivalent; repeatable ±2 µs):

S_past	gqa_xqa	gqa_flash	gqa_cudnn	attn_past	attn_scatter
128	31.0	41.3	31.8	39.0	40.0
512	34.4	59.0	34.7	42.3	46.5
1024	36.1	58.6	37.8	46.1	57.6
2048	40.8	60.8	42.3	60.1	67.8
4096	48.9	75.9	48.6	80.4	78.9

Measured gap vs GQA(XQA): ~8–31 µs/token — same order as, but below, the informal 34–111 µs.

Per-kernel GPU time (nsys, S_past=2048 fp16, per step):

arm	GPU µs	breakdown (µs)
gqa_xqa	11.0	XQA 9.1 + UnpackRoPEAppend 1.2 + seqlens 0.7
gqa_cudnn	10.1	cuDNN sdpa 5.4 + reduce 2.4 + prep 2.4
gqa_flash	22.7	flash_splitkv 17.4 + combine 5.3 (append fused in)
attn_past	21.4	flash_splitkv 13.2 + concat 3.8 + combine 3.2 + 2×fill 1.2
attn_scatter (buffer 8192)	40.3	flash_splitkv 34.1 + combine 3.2 + 2×scatter 2.4 + convert 0.7
attn_scatter (buffer 2560)	19.6	flash_splitkv 13.4 + same prep

Against the gap table

1. KV concat: measured 3.8 µs at S=2048 (B=1) — not tens of µs at this config.
2. No XQA tier (claimed 2–5×): measured ~1.5× e2e, ~1.9× GPU-side. It's a kernel-tier gap, not op-specific: GQA pinned to the same Flash tier is equally slow.
3. Extra concat launch: confirmed — 5 launches/step vs GQA-flash's 2.
4. Unfused prep: confirmed — Fill + Concat separate launches vs GQA's fused append.

Finding not in the table (largest term)

TensorScatter-mode Flash launch is sized by the cache buffer length, not the valid length: shrinking the buffer 8192→2560 at fixed S=2048 drops the identical flash kernel 34.1→13.4 µs. The true length exists only as a device tensor, so host-side split sizing and seqlen_k use the K/V shape (attention.cc, nonpad path). GQA sizes splits from its CPU-input total_sequence_length (group_query_attention.cc:595-602). Cache layout ruled out (4D BNSH cache: same 34 µs).

Consequently the "TensorScatter → near-parity" note and item 1's ±5% criterion don't hold today: +66% e2e, worse with larger buffers — i.e. worst exactly for long max-context deployments. Possible mitigation: accept the valid length (or a max hint) as a CPU input, like GQA.

Caveats

• On SM120, GQA's non-XQA default (cuDNN SDPA) matches XQA — a cuDNN tier would be an alternative item-2 mitigation on new arches.
• B=1 decode is launch-bound (~20–30 µs host floor per step); e2e ratios understate GPU-side ratios. Kernel tables above are the ground truth.
• Single GPU model; kernel-tier ratios will differ on A100/H100. The buffer-length effect is arch-independent logic.

Developed with AI assistance (Claude Code); measurements on real hardware, reviewed by the author.