실시간 KV 캐시 양자화를 위한 레이어별 민감도 및 그룹화 기반 효율적인 오프라인 혼합 정밀도 할당 (Efficient Offline Mixed-Precision Allocation Based on Layer-wise Sensitivity and Grouping for Real-Time KV Cache Quantization)

작성자 신승민
학과 또는 소속(회사명) 컴퓨터공학과
조회수 8
평가(좋아요)수 0
댓글수 0
요 약
대규모 언어 모델의 장문 컨텍스트 추론에서 캐시는 배치 및 입력 길이에 비례해 커지며 병목의 주요 원인이 된다. 캐시로 인한 병목을 해결하기 위한 기존 균일 정밀도 양자화와 같은 방법은 레이어별 민감도를 무시해 저정밀도 양자화에서 성능 붕괴를 초래하고 온라인 적응형 방식은 추가 계산으로 실시간성을 저해한다. 본 논문은 레이어별 K4 attention score error를 민감도로 삼아 상위 민감 레이어를 브랜치 및 기타 레이어를 구간으로 그룹화하고 평균 정밀도 예산을 만족하도록 오프라인 혼합 정밀도 할당을 수행하는 방법을 제안한다. 브랜치에는 높은 정밀도를 우선 배정하고 구간은 상대적으로 낮은 정밀도를 배정하며 이후 예산을 맞추기 위해 정밀도를 보정한다. 이를 통해 레이어별 정밀도 쌍 탐색 공간을 획기적으로 축소하여 효율적인 오프라인 탐색을 가능하게 한다. 온라인 추론 시에는 프리셋을 그대로 양자화에 적용해 온라인 오버헤드가 사실상 0이 된다. Qwen2.5-3B-Instruct 모델을 사용한 GSM8K 데이터셋(4, 8, 16-shot) 평가에서 평균 정밀도 예산 4bit 조건에서도 원본 모델(BF16) 대비 약 95% 수준의 정확도를 유지하며 균일 정밀도의 성능 붕괴를 회피하였다. 이는 민감 레이어의 정밀도를 우선 배정하고 비민감 레이어의 정밀도를 낮추는 설계가 저정밀도 양자화에서도 안정적 성능과 예산 균형을 제공함을 보여준다.

배경지식
트랜스포머와 KV 캐시
대규모 언어 모델은 다수의 트랜스포머 레이어가 누적되어 동작한다. 각 레이어에서 계산된 Key와 Value는 이후 스텝에서도 재사용되므로 캐시에 저장되어 재계산 비용을 줄인다. 이에 따라 캐시 크기는 배치, 입력 길이, 레이어 수에 비례해 커지며 장문 컨텍스트나 큰 배치 환경에서 병목으로 작용한다. 따라서 캐시를 양자화하여 크기를 줄이고 병목을 해결하는 연구가 필요하다.

캐시 양자화
캐시의 메모리 및 대역폭 부담을 낮추기 위해 B-bit 양자화를 적용한다. 입력에 대해 토큰별, 채널별로 비대칭 양자화 및 역양자화를 수행하여 복원한다.

방법
문제 공식화
본 연구의 목표는 동일 평균 정밀도 예산 이하에서 정확도를 극대화하도록 각 레이어의 Key와 Value 정밀도 쌍을 선택하는 것이다. 평균 정밀도 제약 조건을 만족하도록 정밀도를 선택하며, 최종 출력은 각 레이어별 프리셋으로 구성된다. 추론 시에는 이 프리셋을 그대로 적용해 캐시를 양자화한다.

레이어별 측정
원본 모델 대비 각 레이어에서 Key 또는 Value를 저정밀도로 양자화했을 때의 Attention score error를 측정한다. 이는 원본 모델의 Attention score와 저정밀도 양자화된 모델의 Attention score 간의 차이로 정의되며, 민감도를 평가하는 기준으로 사용된다.

민감도 정의
여러 연구에서 Key 분포가 Value 분포보다 이상치(outlier)가 많아 균일 양자화에 더 취약함이 밝혀졌다. 이러한 이상치 분포를 별도로 처리하지 않는다면 낮은 정밀도에서 성능 저하가 발생한다. 본 연구에서는 이상치 분포 처리를 하지 않는 토큰별 양자화를 가정한다. 따라서 GSM8K 0-shot의 첫 20개 프롬프트에 대한 K4 attention score error의 평균을 민감도로 사용한다. K4를 선택한 이유는 K2는 변별력이 약하고 K8은 안정적으로 사용 가능하지만 민감도 구분이 약하기 때문이다.

레이어 그룹화를 통한 탐색 공간 축소
모든 레이어에 대해 Key와 Value 정밀도 쌍을 탐색하려면 매우 큰 탐색 공간을 가진다. 최근 대규모 언어 모델은 레이어 수가 많기 때문에 그룹화를 통해 탐색 공간을 줄인다. 비슷한 특성이나 민감도를 가지는 레이어를 하나의 그룹으로 묶어 동일한 정밀도를 공유하게 함으로써 탐색 공간을 축소한다.

민감도 기반의 그룹화 및 그룹별 정밀도 탐색
민감도가 큰 다섯 개의 레이어를 브랜치로 지정하고, 브랜치 사이의 레이어들을 각각의 그룹으로 묶는다. 브랜치는 구간보다 더 높은 민감도를 가지므로 브랜치 우선 탐색 방식을 사용한다. 브랜치는 높은 정밀도 후보에서, 구간은 낮은 정밀도 후보에서 정밀도를 선택한다. 같은 구간에 속한 레이어는 동일한 정밀도를 가지며, 브랜치는 구간보다 같거나 높은 정밀도를 갖는다. 각 정밀도 조합은 GSM8K 0-shot의 첫 20개 프롬프트를 통해 정확도를 측정한다. 탐색 후 평균 정밀도 예산을 만족하도록 Value의 정밀도를 낮추며 조정한다. 이러한 탐색을 통해 민감한 레이어에 우선적으로 높은 정밀도를 할당하면서도 전체 탐색 공간을 크게 줄일 수 있다.

실험
실험 설정
실험은 Qwen2.5-3B-Instruct 모델을 대상으로 NVIDIA RTX 3090 24GB 환경에서 수행되었다. 평가에는 GSM8K 데이터셋을 사용하였고 4, 8, 16-shot 설정에서 Exact Match 정확도를 측정하였다.

실험 결과
평균 정밀도 예산을 4bit로 설정하고 오프라인 탐색은 GSM8K 0-shot의 20개 프롬프트를 사용하여 진행되었다. 민감도를 기준으로 상위 다섯 개 레이어가 브랜치로 그룹화되었고, 브랜치 사이의 레이어들은 다섯 개 구간으로 그룹화되었다. 브랜치 우선 탐색을 통해 가장 높은 정확도를 보인 조합을 선택하였다. 브랜치는 높은 정밀도 값으로, 구간은 동일한 낮은 정밀도로 할당되었다. 이후 Value 정밀도를 조정하여 평균 4bit 예산을 만족하도록 하였다. 결과적으로 균일 정밀도 양자화(KV4, KV2)에서는 성능 붕괴가 발생했으나 제안한 방법은 동일한 평균 정밀도에서도 원본 모델 대비 약 95% 수준의 정확도를 유지하였다. 이는 민감 레이어의 Key 정밀도를 보호하고 비민감 레이어의 Value 정밀도를 낮추는 설계가 성능과 예산의 균형 면에서 효과적임을 보여준다.

결론
본 연구는 장문 컨텍스트에서 병목의 주 원인으로 작용하는 KV 캐시에 대해 온라인 단계의 추가 연산 없이 실시간으로 적용 가능한 민감도 및 그룹화 기반 오프라인 혼합 정밀도 할당 방법을 제안하였다. 핵심 절차는 레이어별 K4 attention score error를 민감도로 정의하고, 민감도가 큰 레이어를 브랜치와 구간으로 그룹화하여 탐색 공간을 구조적으로 축소하며, 브랜치 우선 정밀도 할당 규칙을 기반으로 탐색하고, 평균 정밀도 예산을 맞추기 위해 Value 정밀도 보정을 수행하는 순서로 구성된다. 최종 프리셋은 추론 시 그대로 양자화되어 실시간 오버헤드가 사실상 0이 된다.

평가 결과, 평균 정밀도 4bit의 저정밀도 환경에서 원본(BF16) 대비 약 95% 수준의 정확도를 유지하였다. 반면 균일 정밀도 KV4에서는 성능 붕괴가 관측되었다. 동일 평균 정밀도 예산 대비 제안법의 성능 우위를 확인할 수 있었으며, 민감 레이어의 Key 정밀도를 우선적으로 보호하고 비민감 레이어의 Value 정밀도를 낮추는 설계가 저정밀도 환경에서도 안정적인 운용을 가능하게 함을 시사한다.

Github

실시간 KV 캐시 양자화를 위한 레이어별 민감도 및 그룹화 기반 효율적인 오프라인 혼합 정밀도 할당 (Efficient Offline Mixed-Precision Allocation Based on Layer-wise Sensitivity and Grouping for Real-Time KV Cache Quantization)

조회수 8
평가(좋아요)수 0
댓글수 0
게시 : 2025년 11월 07일
신승민 컴퓨터공학과

실시간 KV 캐시 양자화를 위한 레이어별 민감도 및 그룹화 기반 효율적인 오프라인 혼합 정밀도 할당 (Efficient Offline Mixed-Precision Allocation Based on Layer-wise Sensitivity and Grouping for Real-Time KV Cache Quantization)

조회수 8
평가(좋아요)수 0
댓글수 0
게시 : 2025-11-07

Member

신승민

Keyword

KV Cahce, Quantization

수상여부

  • 인공지능