[onnxruntime] [ONNX Runtime] CPU GQA 성능의 한계를 넘다: FP16 입력과 양자화된 KV 캐시 최적화 분석

PR 링크: microsoft/onnxruntime#29825 상태: Merged | 변경: +1501 / -285

들어가며\n\n최근 대규모 언어 모델(LLM)의 추론 성능을 최적화하기 위한 핵심 기술 중 하나는 GroupQueryAttention(GQA)과 KV 캐시 양자화입니다. 하지만 CPU 환경에서 FP16 모델을 실행할 때, 양자화된 KV 캐시를 사용하려면 데이터 타입 간의 빈번한 변환 오버헤드가 발생하여 성능이 저하되는 고질적인 문제가 있었습니다.\n\n이번에 분석할 microsoft/onnxruntime의 PR은 CPU GQA에서 FP16 입력/출력을 지원하면서도 INT8 및 INT4로 양자화된 KV 캐시를 효율적으로 처리할 수 있도록 아키텍처별 SIMD 커널을 최적화한 사례입니다. 특히 FP16의 성능을 FP32 수준 혹은 그 이상으로 끌어올린 MLAS(Microsoft Linear Algebra Subprograms) 계층의 개선 사항은 시니어 엔지니어들에게 매우 흥미로운 기술적 통찰을 제공합니다.\n\n## 코드 분석: 핵심 변경 사항\n\n### 1. SIMD 가속을 위한 빌드 및 커널 확장\n\n가장 먼저 눈에 띄는 변화는 FP16 연산을 하드웨어 수준에서 가속하기 위해 컴파일 플래그를 조정한 것입니다. 특히 AVX-512 VNNI 환경에서 FP16 변환을 지원하기 위해 -mf16c 플래그가 추가되었습니다.\n\ncmake\n// Before\nset_source_files_properties(\n ${mlas_platform_srcs_avx512vnni} PROPERTIES \n COMPILE_FLAGS \"-mfma -mavx512vnni -mavx512bw -mavx512dq -mavx512vl -mavx512f\")\n\n// After\nset_source_files_properties(\n ${mlas_platform_srcs_avx512vnni} PROPERTIES \n COMPILE_FLAGS \"-mfma -mf16c -mavx512vnni -mavx512bw -mavx512dq -mavx512vl -mavx512f\")\n\n\n이 변경을 통해 F16C 인스트럭션 셋을 활용하여 FP16과 FP32 간의 변환을 고속으로 수행할 수 있게 되었으며, 이는 후술할 'Query Tile 변환 재사용' 전략의 기반이 됩니다.\n\n### 2. FP16 Query Tile 변환 및 재사용 전략\n\n기존 FP16 커널은 QK(Query-Key) 내적 연산 시 매번 FP16 쿼리 값을 변환하는 비효율이 있었습니다. 이번 PR에서는 **'한 번 변환하고 여러 번 재사용(Convert once, reuse many)'**하는 전략을 도입했습니다.\n\ncpp\n// 신규 도입된 MLAS FP16 전용 API\nMlasQKGemmFp16(..., query_fp16, ..., quantized_kv_cache, ..., scores);\nMlasSVGemmFp16(..., scores, ..., quantized_kv_cache, ..., output_fp16);\n\n\n- AVX2: 단일 행(Single-row) INT8 디코딩 시에는 직접 FP16 벡터 로드를 유지하지만, INT4 디코딩이나 멀티 행(Multi-row) 프리필(Prefill) 시에는 변환된 FP32 타일을 재사용합니다.\n- AVX-512 VNNI: INT8의 경우 모든 행 수에 대해 직접 FP16 로드를 사용하며, 오직 INT4에서만 변환된 타일을 재사용하도록 최적화되었습니다.\n- NEON: 한 번 변환 후 기존 FP32 커널을 재사용하는 방식을 택했습니다.\n\n### 3. SV Gemm의 FP32 Accumulation 최적화\n\nSV(Score-Value) 연산에서는 성능과 정밀도의 균형을 맞추기 위해 독특한 방식을 사용합니다. FP16으로 직접 계산하는 대신, 캐시 친화적인 FP32 축적 커널을 재사용합니다.\n\n1. FP32 스크래치 행(Scratch row)에 중간 결과를 축적합니다.\n2. 최종 단계에서 벡터화된 변환을 통해 결과를 FP16으로 한꺼번에 씁니다.\n\n이 방식은 특히 INT8 캐시 환경에서 기존 FP16 커널 대비 약 50% 이상의 성능 향상을 가져왔습니다.\n\n## 왜 이게 좋은 최적화인가?\n\n### 성능 수치로 증명된 효율성 (AMD EPYC 7763 기준)\n\n마이크로벤치마크 결과는 이 최적화의 가치를 명확히 보여줍니다.\n\n- QK (Query-Key): FP16 커널이 기존 대비 4.7~9.0% 개선되었으며, 모든 디코딩 케이스에서 FP32보다 빠르게 동작합니다.\n- SV (Score-Value): INT8 양자화 시 성능 향상이 **50.7%**에 달하며, FP32 성능의 91% 수준까지 추격했습니다.\n- End-to-End 스레드 스케일링: 8개 스레드를 사용할 때 FP16은 FP32 대비 23~26% 더 빠른 성능을 보였습니다. 이는 FP16이 메모리 대역폭을 적게 사용하면서도 연산 효율을 극대화했음을 의미합니다.\n\n### 일반적 교훈: 캐시 친화성이 정밀도보다 중요하다\n\n단순히 데이터 타입을 FP16으로 낮춘다고 성능이 오르지 않습니다. 오히려 데이터 타입 변환 오버헤드가 더 클 수 있습니다. 이번 PR은 **'캐시 계층 구조에 맞춘 데이터 흐름 설계'**와 **'SIMD 인스트럭션의 특성을 고려한 레지스터 재사용'**이 실제 성능에 얼마나 큰 영향을 미치는지 잘 보여줍니다.\n\n## 리뷰어 피드백 및 논의 사항\n\n리뷰 과정에서 시니어 엔지니어들이 지적한 몇 가지 중요한 기술적 포인트가 있었습니다.\n\n1. 데이터 레이스(Data Race) 노출: gqa_attention_base.h에서 여러 쿼리 헤드가 동일한 KV 헤드를 공유할 때 발생하는 잠재적인 데이터 레이스 문제가 지적되었습니다. 이는 기존 FP32에도 존재하던 문제였으나, FP16 경로가 추가되면서 노출 빈도가 높아졌습니다. 향후 KV 캐시 결합(Concat)을 병렬 루프 외부로 빼내는 후속 작업이 제안되었습니다.\n2. 테스트 정밀도 조정: FP16의 수치적 특성을 고려하여 테스트 코드의 허용 오차(Tolerance)를 조정했습니다.\n cpp\n // Before (test_qkv_quant.cpp)\n 1e-5f\n // After\n 1e-3f // FP16의 반올림 오차와 FMA 연산 순서 차이를 고려하여 완화\n \n3. Flash Attention 지수 함수 최적화: std::exp 계산이 루프 내부에서 과도하게 호출되는 문제를 발견하고, 이를 루프 외부로 빼내어 계산 횟수를 $O(head\_size \times kv\_chunk)$에서 $O(kv\_chunk)$로 줄이는 최적화가 논의되었습니다.\n\n## 마무리하며\n\n이번 PR은 CPU 환경에서도 LLM의 FP16 추론이 충분히 고성능으로 동작할 수 있음을 입증했습니다. 특히 MLAS 계층에서의 정교한 SIMD 최적화와 캐시 전략은 고성능 컴퓨팅(HPC) 분야의 정수를 보여줍니다. CPU 기반 AI 서비스를 운영하는 엔지니어라면, 이러한 저수준 최적화 기법을 통해 인프라 비용을 절감하고 사용자 경험을 개선할 수 있을 것입니다.

참고 자료

• https://intel.github.io/intrinsic-guide/#text=_mm256_cvtps_ph
• https://developer.arm.com/architectures/instruction-sets/intrinsics/
• https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/mlas/lib/mlasi.h

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.