onnxruntime
a203dfaf - [CPU] Add FP32 GEMV decode kernel for GroupQueryAttention (#29216)

Commit
30 days ago
[CPU] Add FP32 GEMV decode kernel for GroupQueryAttention (#29216) ### Description PR1 https://github.com/microsoft/onnxruntime/pull/28962 adds flash attention for **prefill**, and removed flash decoding. This PR will add optimized kernel for **single-token decode**, which will be faster than other kernels including flash decoding. This PR builds on the prefill-only flash attention change and additionally introduces a dedicated decode kernel. #### What's included - **Decode (GEMV) kernel** — A dedicated single-token decode kernel (`MlasGQADecodeGQAThreaded`) for `sequence_length == 1`, parallelized over (batch, head) with a two-pass softmax, using GEMV (`acc[8]`-lane dot product / AXPY) helpers instead of per-block M=1 SGEMM calls. This fixes the per-block SGEMM decode regression. - The FP32 flash gate (`group_query_attention.cc`) is enabled for `total_sequence_length > 1`, routing prefill to the tiled kernel and decode to the GEMV kernel. - The quantized KV-cache path is unchanged (FP32-only scope). #### Results (AMD EPYC 7763, AVX2, 8 threads) - **Decode:** correctness ~1e-8 vs naive; long-context decode ~1.0–1.5x (T = 4097 ~1.3–1.5x). ### Motivation and Context The naive GQA path materializes the full score matrix, which is memory-bound for long sequences. Flash attention reduces memory traffic for prefill, and the GEMV decode kernel avoids SGEMM overhead for the M=1 decode case. ### Testing - Built with `--compile_no_warning_as_error`. - Correctness verified against the naive path for both prefill and decode (max abs diff ~1e-8). - Benchmarked via `benchmark_gqa_cpu_flash.py`.
Author
Parents
Loading