onnxruntime

56640661 - CPU GroupQueryAttention: Quantized KV Cache with SIMD-optimized MLAS kernels (#28578)

CPU GroupQueryAttention: Quantized KV Cache with SIMD-optimized MLAS kernels (#28578) ## Description Add INT8/INT4 symmetric quantized KV-cache support for the CPU `GroupQueryAttention` contrib operator, with progressive SIMD acceleration (AVX2, AVX512-VNNI, NEON). This builds on the scalar baseline in https://github.com/microsoft/onnxruntime/pull/28576 and adds hardware-specialized MLAS kernels for the quantized GEMM operations (`QKGemm`: query × K-cache^T, `SVGemm`: attn_weights × V-cache). ### Motivation LLM inference on CPU is memory-bandwidth-bound during decoding. Quantizing the KV cache to INT8 or INT4 reduces memory traffic by 4-8x, directly translating to throughput improvement. The CUDA EP already supports quantized KV cache; this PR extends it to the CPU EP. ### Key Changes | Area | Files | Description | |------|-------|-------------| | **MLAS Public API** | `mlas_qkv_quant.h`, `qkv_quant.cpp` | `MlasKVQuantize`, `MlasKVDequantize`, `MlasQKGemm`, `MlasSVGemm` — thread-safe quantize/dequantize and fused dequant-GEMM | | **MLAS Dispatch** | `qkv_quant_kernel.h`, `mlasi.h`, `platform.cpp` | Runtime CPUID detection selects best kernel (scalar → AVX2 → AVX512-VNNI → NEON) | | **AVX2 kernel** | `qkv_quant_kernel_avx2.cpp` | Fused dequant-dot: 256-bit FMA, 8 elements/iteration for INT8, INT4 nibble-unpack | | **AVX512-VNNI kernel** | `qkv_quant_kernel_avx512vnni.cpp` | `_mm512_dpbusd_epi32` for INT8 per-tensor (4x throughput vs FMA), MultiDot4 ILP optimization, 512-bit FMA for per-channel/INT4 | | **NEON kernel** | `qkv_quant_kernel_neon.cpp` | ARM NEON fused dequant-dot for INT8/INT4 | | **GQA Operator** | `group_query_attention.cc`, `gqa_attention_base.h` | Accept quantized `past_key`/`past_value` (uint8/int8), quantize-on-write, dequant-on-read via MLAS | | **CMake** | `onnxruntime_mlas.cmake` | Build rules for all platforms (Linux x64, Windows x64, ARM64) | ### Performance (Intel Xeon Platinum 8480C, single core, Release -O3) **QKGemm (query × K-cache^T), M=1, N=512, K=128, INT8 per-tensor:** | Implementation | Latency (ns) | vs Scalar | |---|---|---| | Scalar fallback | 30,179 | 1.0x | | AVX2 (fused dequant-dot) | 4,219 | **7.2x** | | AVX512 FP32 fused dequant-dot | 3,736 | **8.1x** | | AVX512-VNNI (dpbusd + MultiDot4) | 1,938 | **15.6x** | **QKGemm, M=1, N=512, K=128, INT4 per-tensor:** | Implementation | Latency (ns) | vs Scalar | |---|---|---| | Scalar fallback | 141,467 | 1.0x | | AVX2 (fused dequant-dot) | 24,946 | **5.7x** | | AVX512-VNNI (512-bit FMA) | 5,817 | **24.3x** | ### Quantization Modes - `S8_PerTensor` — INT8, single scale for entire cache slice - `S8_PerChannel` — INT8, per-head-dimension scale vector - `S4_PerTensor` — INT4 (biased nibble packing), single scale - `S4_PerChannel` — INT4, per-head-dimension scale vector ### Testing - **MLAS unit tests** (`test_qkv_quant.cpp`): Exhaustive sweep of all M/N/K/QuantType combinations, validates correctness against scalar reference - **MLAS benchmarks** (`bench_qkv_quant.cpp`): Scalar vs AVX2 vs AVX512-VNNI comparison - **C++ op test** (`group_query_attention_op_test.cc`): End-to-end GQA with quantized KV cache - **Python integration test** (`test_gqa_cpu_quantized.py`): Multi-step decoding accuracy validation against FP32 reference ### How to Test ```bash # MLAS unit test ./build/Release/onnxruntime_mlas_test --gtest_filter="KVQuant.*" # MLAS benchmark (Release build with --build_micro_benchmarks) ./build/Release/onnxruntime_mlas_benchmark --benchmark_filter="BM_QKGemm|BM_SVGemm" # C++ GQA op test ./build/Release/onnxruntime_test_all --gtest_filter="*GroupQueryAttentionQuantized*" # Python integration test python onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py ```