[Pierluigi Urru (stwykd)]

Description

CPU GroupQueryAttention with fp32 or fp16 KV caches is far slower than its cache size allows, mostly in decode. On main:

• The non-quantized path runs one task per query head (gqa_attention_base.h:1475, :1729), so each KV head is read once per query head, and with fp16 caches each task copies its KV head's whole K and V to fp32 (:1589, :1812).
• The flash path is fp32 only (group_query_attention.cc:811). Its decode is one GEMV task per query head, splits the cache only when batch_size * num_heads < threads (gqa_attention_base.h:1291), and its GEMVs (flashattn_gqa.cpp:51-101) are built with baseline flags (onnxruntime_mlas.cmake:73), so the x64 object has no AVX or FMA.

This PR adds grouped_flash_attention.{h,cc}, a portable kernel on MLAS SGEMM, fp16 conversion and the thread pool. Per KV head it stacks the (token, head) rows of the group as the GEMM M dimension against each cache tile, so a tile is read, converted and packed once. Blocks of a single row (decode with num_heads == kv_num_heads) skip packing and use MLAS's single-row kernels instead. The softmax runs online in fp32 with causal and window tile skipping. With fewer tasks than threads (decode) the cache is split and merged with a log-sum-exp reduction. Planning arithmetic is in ptrdiff_t.

GQA uses it for fp32 and fp16 caches, covering causal and bidirectional attention, local window, softcap, attention bias, smooth softmax and head sink, rotary, packed QKV, ragged and query_seqlens batches, shared, separate and read-only caches, and the sliding window cache. Quantized caches, qk_output and ORT_GQA_DISABLE_FLASH_ATTENTION=1 keep the existing paths, as do fp16 caches where MLAS has fp16 GEMMs (ARM64 with fp16 arithmetic), where the per-head path is still faster. The fp32 flash dispatch stays behind the new path: off x64, ungrouped fp32 decode still uses its GEMV kernel, which is faster there because MLAS has no single-row kernel for the transposed Q*K' product. docs/contrib_ops/cpu/gqa.md describes the dispatch.

The kernel takes separate K and V caches with their own widths and strides, so other CPU attention ops can share it, for example latent attention (#33914) with V as a column window of K.

Tests

New tests check the kernel against fp64 (fp32/fp16, group sizes 1, 3, 8, split decode, padded prefill, window, softcap, bias, sinks, bidirectional, NaN-poisoned unread rows) and the op against the per-head path; the Python decode parity test now covers fp16. GQA C++ and Python tests pass on x64; on ARM64 the C++ tests pass and Python fails only where main does (below). Over 1000 random configurations, outputs match main within 2.6e-6 (fp32) and one fp16 ulp.

Performance

One layer, 4 threads, Xeon w5-2565X, cache flushed between decode calls, median of 3 runs, ms:

	2K	8K	32K
decode fp16, 24 q / 8 kv heads (Llama 3.2 3B)	0.87 → 0.49	5.43 → 1.72	32.4 → 5.86
decode fp16, 32/8 (Llama 3 8B)	1.08 → 0.49	6.88 → 1.76	48.0 → 6.00
decode fp16, 28/4 (Qwen2.5 7B)	1.10 → 0.30	6.52 → 1.05	42.6 → 3.59
decode fp32, 24/8	0.99 → 0.59	5.72 → 2.17	21.3 → 7.75
prefill 128, fp16, 24/8	6.28 → 5.77	25.1 → 20.8	113.6 → 81.4
prefill 128, fp32, 24/8	6.46 → 5.88	24.0 → 22.1	93.0 → 83.1

Ungrouped decode (num_heads == kv_num_heads, as Phi-3-mini exports) doesn't regress: 8/8 heads fp32 at 2K and 8K is 0.83x and 0.81x the time of main, 32/32 heads fp16 at 2K 0.91x, on a busy host. On Neoverse N1 it stays on the GEMV path and matches main.

fp16 decode at 32K reads the cache at 22 GB/s, a third of 4-core bandwidth. On Neoverse N1, fp32 decode is 1.4 to 2.4x faster and prefill 1.1x.

Follow-ups

• Share the kernel with latent attention, or move it into MLAS beside MlasFlashAttentionGQA, with an ARM64 single-row Q*K' kernel so ungrouped decode can use the grouped path there too.
• Small-group decode is bound by MLAS packing each tile; a hand-written AVX2 kernel is 1.3x faster at 32K, so an MLAS multi-row fp16 GEMV would help.
• An fp16-GEMM variant for ARM64, whose per-head fp16 path is less accurate (test_gqa_interactive_one_batch fails on main on Neoverse N1).
• Quantized caches through the same tiling.

[azure-pipelines]

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

[Pierluigi Urru (stwykd)]

the failing checks aren't from this change, which only touches the CPU GQA kernel: the T5 beam search tests in test_generation.py fail on main too (tiktoken missing and a protobuf version mismatch in the CI image), the kernel docs check diff is in Gather and Relu rows this PR doesn't touch, and the CUDA plugin EP job failed one FpA-IntB tuning cache test by a tiny margin (it passes on main and on #33914, so it looks flaky, could someone rerun it?). test_gqa_cpu.py and the GQA C++ tests pass in every job that runs them

[Pierluigi Urru (stwykd)]

Tianlei Wu (@tianleiwu) you added the CPU flash attention path, would you mind taking a look? this applies the same tiling to grouped heads and fp16 caches, reading each cache tile once per KV head instead of once per query head

[Copilot]

🟡 Changes recommended

Deferring the sink can produce infinite results when the correct output is finite, and incorrect ARM64EC guards bypass the intended GEMV fallback.

2 open issues

• Include the sink in the running maximum before row writes · New
• Use MLAS_TARGET_AMD64 for ARM64EC GEMV and packing guards · New

What changed in this PR

This PR accelerates CPU GroupQueryAttention, especially decode, by sharing each KV-cache tile across query heads that use it.

Changes:

• Adds a portable FP32/FP16 grouped attention kernel with tiled online softmax and split-cache processing.
• Integrates the kernel with GQA dispatch and cache handling while preserving existing fallback paths.
• Extends reference and parity tests and documents dispatch behavior and benchmarks.

File	Description
onnxruntime/​test/​python/​transformers/​test_gqa_cpu.py	Extends decode parity tests to FP16.
onnxruntime/​test/​contrib_ops/​group_query_attention_op_test.cc	Adds FP64-reference kernel tests and per-head-path comparisons.
onnxruntime/​contrib_ops/​cpu/​bert/​grouped_flash_attention.h	Defines the reusable kernel interface and cache strides.
onnxruntime/​contrib_ops/​cpu/​bert/​grouped_flash_attention.cc	Implements grouped tiled attention and split-state reduction.
onnxruntime/​contrib_ops/​cpu/​bert/​group_query_attention.cc	Dispatches eligible GQA calls to the grouped kernel.
onnxruntime/​contrib_ops/​cpu/​bert/​gqa_attention_base.h	Adds eligibility checks, argument setup, and shared cache appending.
docs/​contrib_ops/​cpu/​gqa.md	Documents the algorithm, activation conditions, tests, and performance.

_{Review effort: Max}