onnxruntime

7e82d19d - Flash Attention style tiled computation for CPU GQA quantized KV cache (#28695)

Flash Attention style tiled computation for CPU GQA quantized KV cache (#28695) ## Description Add a flash attention-style tiled computation path to the CPU GroupQueryAttention operator for quantized KV cache (INT8/INT4). Instead of materializing the full `[B, N, S, T]` attention probability matrix, this processes K/V in L2-cache-sized blocks with online softmax — reducing peak memory from O(S×T) to O(S×Bc) per head where Bc is the KV block size. Additionally, implements **flash decoding** for the decode phase (S=1): when `batch×heads < threads`, idle threads are repurposed to partition the KV sequence across parallel workers. Each worker computes partial softmax statistics on its KV chunk, then a lightweight reduce phase merges the partials — achieving 2–5x decode speedup for long sequences. ### Motivation For long-sequence LLM inference with quantized KV cache on CPU: - **Prefill**: The full attention matrix allocation becomes a significant memory bottleneck. With 16 heads and S=4096, the naive path allocates ~1 GB for attention scores alone. The tiled approach reduces peak memory by 13–24x and latency by 1.2–2.7x. - **Decode**: When batch size is small relative to available threads, many threads sit idle. Flash decoding partitions the KV sequence across these idle threads, achieving 2–5x speedup for long KV lengths. ## Key Changes | File | Change | |------|--------| | `onnxruntime/core/mlas/lib/flashattn_qkv.cpp` | MLAS kernel: tiled prefill with online softmax, flash decoding (two-phase KV partitioning), and reduce | | `onnxruntime/core/mlas/inc/mlas_qkv_quant.h` | `MlasFlashAttentionQuantizedKVArgs` struct with `flash_decoding_partials` and `kv_chunk_count` fields | | `onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h` | `ApplyAttentionQuantizedFlash()` with L2-cache-aware block sizing, KV concat, flash decoding setup | | `onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc` | Dispatch logic: activates flash path when no softcap/smooth softmax/output_qk | | `cmake/onnxruntime_mlas.cmake` | Added `flashattn_qkv.cpp` to the MLAS build | | `docs/contrib_ops/cpu/gqa.md` | Documentation with algorithm details, benchmark results, and reproduction steps | | `onnxruntime/test/mlas/bench/bench_qkv_quant.cpp` | MLAS-level C++ benchmark (`BM_GQA_Naive` vs `BM_GQA_Flash`) | | `onnxruntime/test/python/transformers/benchmark_gqa_cpu_flash.py` | Operator-level Python benchmark | ## Algorithm ### Prefill (S > 1): Tiled Flash Attention Per (batch, head, q_block) tile: 1. **QK GEMM** — `MlasQKGemm` on a block slice of quantized K cache 2. **Causal + local window masking** — Set masked positions to -inf before softmax 3. **Online softmax** — Track running max `m` and sum `l`, rescale accumulated output with `exp(m_old - m_new)` 4. **SV accumulation** — Dequantize V block to FP32, then accumulate weighted V into output ### Decode (S = 1): Flash Decoding When `sequence_length == 1 && batch_size * num_heads < thread_count && kv_chunk_count > 1`: **Phase 1 — Parallel KV scan**: Each idle thread processes a disjoint KV chunk for a (batch, head) pair. For each chunk: compute QK dot products, find local max, compute local softmax sum, and accumulate partial weighted V output. Store per-chunk `(max_score, sum_exp, partial_output[head_size])` into a partials buffer. **Phase 2 — Reduce**: One thread per (batch, head) merges all chunk partials using the log-sum-exp trick: find global max, rescale each chunk's sum and partial output, then normalize by global sum. This is analogous to GPU flash decoding (Dao et al.) but adapted for CPU threading. ### Activation Conditions Flash path activates when ALL of: - `ORT_GQA_DISABLE_FLASH_ATTENTION` env var is not set - `total_sequence_length > 1` - No softcap, no smooth softmax, no output_qk (attention bias IS supported) Flash decoding additionally requires: - `sequence_length == 1` (decode phase) - `batch_size * num_heads < thread_count` (idle threads available) - `kv_chunk_count > 1` (enough KV to partition) ## Benchmark Results Measured on Intel Xeon Platinum 8480C, 96 CPUs, threads=8. MLAS-level C++ benchmark. ### Latency — Prefill (S = T) Shape: B=1, num_heads=16, kv_num_heads=8, head_size=128. | Seq Length | Naive (ms) | Flash (ms) | Speedup | Quant | |---:|---:|---:|---:|:---| | 512 | 9.9 | 8.1 | 1.2x | per-tensor | | 1024 | 44.4 | 27.0 | 1.6x | per-tensor | | 2048 | 190.9 | 116.9 | 1.6x | per-tensor | | 4096 | 1257.8 | 461.6 | 2.7x | per-tensor | | 512 | 10.7 | 10.8 | 1.0x | per-channel | | 1024 | 49.5 | 41.7 | 1.2x | per-channel | | 2048 | 212.1 | 164.1 | 1.3x | per-channel | | 4096 | 1223.9 | 607.8 | 2.0x | per-channel | ### Latency — Decode (S = 1, no flash decoding) Shape: B=1, num_heads=16, kv_num_heads=8, head_size=128. Flash decoding NOT active (batch×heads=16 > threads=8). | Total Seqlen | Naive (us) | Flash (us) | Speedup | Quant | |---:|---:|---:|---:|:---| | 512 | 32 | 22 | 1.4x | per-tensor | | 1024 | 71 | 47 | 1.5x | per-tensor | | 2048 | 120 | 87 | 1.4x | per-tensor | | 4096 | 210 | 174 | 1.2x | per-tensor | | 512 | 53 | 31 | 1.7x | per-channel | | 1024 | 86 | 52 | 1.7x | per-channel | | 2048 | 172 | 97 | 1.8x | per-channel | | 4096 | 299 | 191 | 1.6x | per-channel | ### Latency — Flash Decoding (S = 1, KV partitioned across threads) Shape: B=1, num_heads=4, kv_num_heads=4 (MHA), head_size=128. Flash decoding IS active (batch×heads=4 < threads=8). | Total Seqlen | Naive (us) | Flash (us) | Speedup | Quant | |---:|---:|---:|---:|:---| | 512 | 31 | 25 | 1.2x | per-tensor | | 1024 | 41 | 25 | 1.6x | per-tensor | | 2048 | 67 | 34 | 2.0x | per-tensor | | 4096 | 197 | 54 | 3.7x | per-tensor | | 512 | 25 | 28 | 0.9x | per-channel | | 1024 | 72 | 27 | 2.7x | per-channel | | 2048 | 144 | 37 | 3.9x | per-channel | | 4096 | 304 | 60 | 5.1x | per-channel | ### Peak Memory — Prefill | Seq Length | Naive Peak | Flash Peak | Memory Reduction | |---:|---:|---:|---:| | 2048 (N=16) | +294 MB | +44 MB | 6.7x | | 4096 (N=16) | +1107 MB | +82 MB | 13.5x | | 4096 (N=32) | +2131 MB | +87 MB | 24.5x | **Summary**: Prefill gains 1.2–2.7x latency + 7–24x memory reduction from tiled online softmax. Decode gains 1.2–1.8x from fused dequant+dot alone. Flash decoding adds 2–5x for long sequences when idle threads are available to partition the KV scan. ### How to Reproduce ```bash # Build ORT python tools/ci_build/build.py --build_dir build/cpu --config Release \ --parallel --build_wheel --skip_tests # MLAS-level C++ benchmark: cd build/cpu/Release ./onnxruntime_mlas_benchmark \ --benchmark_filter='BM_GQA_(Naive|Flash)' \ --benchmark_min_time=0.5s \ --benchmark_repetitions=3 \ --benchmark_report_aggregates_only=true ``` ## Testing - All 35 CPU `GroupQueryAttentionTest.*` tests pass (INT8/INT4, per-tensor/per-channel, multi-batch, large head, GQA ratio variants) - Set `ORT_GQA_DISABLE_FLASH_ATTENTION=1` to verify fallback path still works - End-to-end verified with `quantized_kv_cache_cpu_demo.py` - Numerical agreement between flash and naive paths: max diff < 1e-7