Add decode benchmark comparing ONNX Attention CUDA with contrib GQA

[namgyu-youn]

Description

Adds onnxruntime/test/python/transformers/benchmark_onnx_attention_vs_gqa.py: a single-stream decode benchmark comparing ONNX Attention (CUDA EP) with contrib GroupQueryAttention, as requested by the "Baseline Benchmark Recipe" TODO in #28352.

Five arms at matched config (S_q=1, causal, no RoPE/mask/softcap; Llama-3-8B shape by default): GQA on its XQA / Flash / cuDNN SDPA kernel tiers, ONNX Attention with past/present inputs, and ONNX Attention opset-24 external cache (in-place TensorScatter + nonpad_kv_seqlen). Includes a --sanity cross-arm output-parity check, a --profile NVTX mode for Nsight Systems captures, and GPU/driver/wheel provenance in every log.

Motivation and Context

The decode-latency figures in #28352 came from informal profiling with no reproducible context; this script is the recipe to reproduce them. Benchmark-only change — no ORT source modifications, not collected by CI.

Verified on RTX PRO 6000 (SM120): all arms produce identical outputs on identical inputs (max |diff| ≤ 8e-6 fp16); fp16/bf16 sweeps repeatable within ±2 µs. Measured results will be posted on #28352.

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

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[namgyu-youn]

Measured results and per-kernel attribution from this script (RTX PRO 6000, fp16/bf16 sweep + nsys captures): #28352 (comment)

[Copilot]

Pull request overview

Adds a standalone Python benchmark script under onnxruntime/test/python/transformers/ to reproduce and compare single-token decode latency across five matched “arms”: contrib GroupQueryAttention (XQA/Flash/cuDNN tiers) vs ONNX Attention (past/present path and opset-24 TensorScatter external-cache path). This supports the “Baseline Benchmark Recipe” TODO from issue #28352 by making the measurement setup reproducible and logging key provenance.

Changes:

• Introduces benchmark_onnx_attention_vs_gqa.py that programmatically builds ONNX Attention graphs (opset 23 past/present, opset 24 TensorScatter external cache) and benchmarks them against contrib GQA sessions.
• Adds --sanity parity mode and --profile NVTX-annotated mode, plus CSV logging and environment pinning for backend selection.
• Prints GPU/driver/ORT build provenance to make published results attributable and repeatable.

[titaiwangms]

Review synthesis (multi-model review team)

Benchmark-only file (650 lines, no ORT source changes). No Critical issues, no blockers. Integration is clean — the imported APIs (gqa_test_helper, CudaSession) and the hand-built ONNX Attention (opset 23/24) + TensorScatter graphs are spec-faithful and correctly wired for the ORT CUDA target. Findings below focus on fair-comparison integrity, since that's the point of the file.

Major (fair-comparison integrity — recommend addressing)

1. Named kernel arms aren't verified to use their advertised kernel (ARM_ENV, ~L71-76; build_arm ~L406). ORT_ENABLE_XQA=1 etc. only enable a backend — GQA still silently falls back on shape/device ineligibility, and ambient ORT_DISABLE_FLASH_ATTENTION isn't overridden or reported. A row labeled gqa_xqa could be measuring another kernel. Consider parsing the resolved backend (ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1) after session build and asserting/aborting on mismatch.

2. CPU EP fallback can silently invalidate a CUDA number (~L240). Every session appends CPUExecutionProvider; an unsupported dtype/shape can offload nodes to CPU while still emitting a latency row. Consider inspecting node assignment (or dropping the CPU EP) and aborting unless all compute nodes are on CUDA.

3. --device is not honored by the session (create_cuda_session ~L232-238). Tensors use args.device, but the EP device_id + stream come from torch.cuda.current_device()/current_stream(). --device cuda:1 puts tensors on GPU1 and the EP on GPU0. Default "cuda" is fine. Fix: torch.cuda.set_device() from a parsed torch.device and derive id/stream from it.

4. CSV conflates different experiments (~L497-530). Neither --attention-only nor --cache-4d is recorded, so appended rows for structurally different runs share an arm label. Add both as columns (or use distinct arm names).

5. Timer choice changes methodology, not just precision (~L56-62, L416-438). Triton's do_bench clears L2 each iteration; the event_bench fallback runs warm back-to-back. Numbers aren't comparable across environments depending on whether Triton is installed. Consider making both paths use one explicit cache policy.

Minor

• attn_scatter uses is_causal=1 (~L200). Executing the ONNX reference impl and comparing against ORT's own test_tensorscatter_attention.py: for q_len=1 this is a no-op under ORT's bottom-right causal convention (correct, matches GQA), but would compute V[0] under the spec-conformant ONNX reference / CPU EP. Harmless for this ORT-only benchmark. Recommend setting is_causal=0 (matches GQA, the ONNX reference, and ORT's default test config) or adding a one-line comment.
• Sanity check is weak (~L555-568): compares against another optimized arm (not an fp32 reference), abs-error only, bf16 tol 0.05 vs input std 0.1; a single-arm run passes vacuously. Consider an fp32 reference + rtol and requiring >= 2 arms.
• No argument validation (~L600-634): zero/negative dims, q_heads % kv_heads != 0, or sweep >= max_seq_len (TensorScatter linear mode) reach graph construction unchecked.

Nits

• do_bench(warmup=15, rep=100) (ms) vs event_bench(warmup=20, rep=100) (iters) — same kwarg names, different units; add a clarifying comment at the do_bench call.
• ok &= status == "PASS" -> ok = ok and ...; add a bench_arm docstring ("returns mean latency in ms"); add a scoped_env docstring; module docstring says the scatter arm has "no mask" but it passes nonpad_kv_seqlen (a padding mask) — say "no attn_mask".

Praise

Unusually thorough docstrings; clean ARM_ENV/ARM_BUILDERS dispatch (2-line to add an arm); correct TensorScatter in-place buffer-sharing; BSNH/BNSH layout handling and nonpad/write_indices consistency are all spec-correct.

---

GPU-dependent checks (kernel-dispatch verification, on-device parity) were not run — no CUDA-built onnxruntime in the review environment. The PR author reports all arms produce identical outputs (max |diff| <= 8e-6 fp16) on an RTX PRO 6000.