[onnxruntime] ONNX Runtime: INT4 Paged KV Cache 도입을 통한 메모리 최적화

PR 링크: microsoft/onnxruntime#32515 상태: Merged | 변경: +1698 / -151

들어가며

대규모 언어 모델(LLM)의 긴 컨텍스트(Long-context) 처리 시 가장 큰 병목 중 하나는 KV Cache가 차지하는 메모리 공간입니다. 기존의 INT8 기반 Paged KV Cache는 메모리 효율성을 높였으나, 여전히 더 큰 모델을 배포하기에는 메모리 점유율이 높았습니다. 본 PR은 microsoft/onnxruntime의 CUDA PagedAttention 커널에 INT4 Paged KV Cache를 도입하여, 기존 INT8 대비 메모리 풋프린트를 절반으로 줄이고 추론 성능을 크게 향상시켰습니다.

코드 분석

1. CMake 빌드 옵션 변경 (cmake/CMakeLists.txt)

INT4 KV Cache 기능을 선택적으로 활성화할 수 있도록 빌드 플래그를 수정했습니다.

-option(onnxruntime_USE_INT4_KV_CACHE "Build cuda kernels for int4 kv cache" OFF)
+option(onnxruntime_USE_INT4_KV_CACHE "Build cuda kernels for int4 kv cache" ON)

2. PagedAttention 연산자 정의 업데이트 (docs/ContribOperators.md)

INT4는 4비트 값을 2개씩 1바이트(uint8)에 패킹하여 저장합니다. 이에 따라 캐시 텐서의 형태가 변경되었습니다.

-<dd>Block-based key cache with shape (num_blocks, block_size, kv_num_heads, head_size). ... </dd>
+<dd>Block-based key cache with shape (num_blocks, block_size, kv_num_heads, cache_head_size), where cache_head_size is (head_size + 1) / 2 for packed INT4 and head_size otherwise. ... </dd>

3. XQA(Cross-Query Attention) 커널 최적화

INT4 XQA 디코딩 시, PER_CHANNEL 스케일을 쿼리(Q)에 폴딩(Folding)하여 연산 효율을 높였습니다. 특히 FP16의 정밀도 한계를 극복하기 위해 max|k_scale|의 2의 거듭제곱(power of two) 값을 정규화 인자로 사용하여 오버플로우를 방지했습니다.

왜 이게 좋은가

1. 메모리 절감: INT4 포맷은 INT8 대비 동일한 컨텍스트 길이에 대해 메모리 사용량을 정확히 50% 줄입니다. 이는 H200과 같은 고성능 GPU에서 더 긴 시퀀스를 처리하거나 더 큰 배치 사이즈를 운용할 수 있게 합니다.
2. 성능 향상: Qwen3.8-27B 모델 기준, XQA 커널을 통해 포터블 커널 대비 최대 6배 이상의 속도 향상을 확인했습니다.
3. 정밀도 유지: PER_CHANNEL 스케일링을 정교하게 폴딩하여, INT8과 비교했을 때 MMLU-Pro 벤치마크에서 오차 범위 내의 정확도를 유지했습니다.

일반적 교훈

• 정밀도 관리: FP16 연산 시 오버플로우를 방지하기 위해 스케일링 인자를 2의 거듭제곱으로 정규화하는 기법은 수치적 안정성을 확보하는 데 매우 효과적입니다.
• 테스트 전략: 새로운 커널 도입 시, 기존 포터블 커널과의 결과 비교(Parity Check)뿐만 아니라, 극단적인 스케일 값(FLT_MAX 등)에 대한 수치적 안정성 테스트가 필수적임을 확인했습니다.

리뷰어 피드백 반영

리뷰 과정에서 FP16 폴딩 시 발생할 수 있는 NaN 문제와 오버플로우 위험이 제기되었습니다. 이를 해결하기 위해 max|k_scale| 기반의 정규화 로직을 도입하고, 특정 조건에서는 안전한 포터블 FP32 커널로 자동 폴백(Fallback)되도록 설계하여 정밀도와 성능의 균형을 맞췄습니다.

참고 자료

• torch.compile
• ONNX Runtime Contrib Ops

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