Hugging Face KV 캐시 양자화

Hugging Face는 Transformers 라이브러리에 KV 캐시 양자화를 도입하여, 키-값(KV) 캐시의 메모리 사용량을 줄임으로써 대형 언어 모델(LLM)이 더 긴 텍스트 시퀀스를 생성할 수 있게 했습니다. 이 기능을 사용하면 소비자용 GPU에서 약간의 생성 속도 감소를 감수하고도 메모리 효율성을 크게 향상시켜 컨텍스트 길이를 확장할 수 있습니다.

자동회귀 생성에서 KV 캐시의 역할

키-값(KV) 캐시는 텍스트 토큰을 하나씩 예측하는 자동회귀 모델을 최적화하는 데 필수적입니다. 새로운 토큰을 예측하려면 모델은 이전 모든 토큰의 정보를 필요로 합니다. 캐시가 없으면 모델은 매 단계마다 모든 이전 토큰에 대해 행렬 곱셈을 다시 계산해야 합니다. KV 캐시는 메모리 뱅크 역할을 하여, 이전에 처리된 토큰에 대한 self‑attention 레이어의 키‑값 쌍을 저장하고, 이를 재계산하는 대신 검색할 수 있게 함으로써 텍스트 생성 속도를 크게 가속합니다.

하지만 컨텍스트 길이 또는 배치 크기가 증가하면 KV 캐시는 메모리 병목 현상이 됩니다. 예를 들어, 컨텍스트 길이가 10,000 토큰인 7B Llama‑2 모델의 경우, KV 캐시는 float16 정밀도로 약 5GB의 메모리를 필요로 하며, 이는 모델 파라미터 자체가 차지하는 메모리의 거의 1/3에 해당합니다.

KV 캐시 양자화의 기술 구현

KIVI 논문에서 영감을 받아, Hugging Face의 구현은 어파인 양자화를 사용해 KV 캐시를 낮은 정밀도 형식으로 압축합니다.

양자화 방법

이 구현은 키와 값 모두에 대해 토큰 단위 양자화를 수행합니다. 매 생성 단계마다 양자화와 역양자화로 인한 잠재적 속도 저하를 완화하기 위해, Hugging Face는 고정 크기의 잔여 캐시를 사용합니다. 이 잔여 캐시는 가장 최근의 키와 값을 원래 정밀도로 저장하며, 캐시가 최대 용량에 도달하면 저장된 값들을 양자화하고 캐시를 비웁니다. 정확성을 유지하기 위해 기본값으로 잔여 길이 128을 사용합니다.

지원되는 백엔드 및 정밀도

  • Quanto: int2 및 int4 정밀도를 지원합니다.
  • HQQ: int2, int4, 및 int8 정밀도를 지원합니다.

성능 및 품질 트레이드오프

모델 품질 및 정확도

PG‑19 데이터셋에서 Llama2‑7b‑chat 모델을 사용한 테스트 결과, int4 캐시 정밀도는 fp16 정밀도와 거의 동일하게 동작함을 보여줍니다. 그러나 int2를 사용할 경우 품질이 저하됩니다. LongBench 벤치마크에서는 Quanto 백엔드를 통한 int4 정밀도가 TREC, SAMSum, TriviaQA 등 여러 데이터셋에서 fp16과 비슷하거나 약간 더 우수한 성능을 보였습니다.

메모리 효율성 vs. 지연 시간

int4로 캐시를 양자화하면 약 2.5배의 메모리 절감 효과를 얻을 수 있습니다. 메모리 부담이 줄어들지만, 특히 배치 크기가 커질수록 생성 속도가 감소할 수 있습니다. 또한 KV 캐시 양자화와 가중치 양자화를 동시에 적용하면 생성 속도가 세 배 정도 감소할 수 있습니다.

컨텍스트 길이 용량

80GB A100 GPU에서 Flash Attention과 결합하면, KV 캐시 양자화를 통해 모델이 최대 128k 토큰을 지원할 수 있게 되며, 반정밀도 캐시를 사용할 경우 최대 40k 토큰에 불과합니다.

Transformers와의 통합

KV 캐시 양자화는 장치에 구애받지 않으며 CPU, GPU, MPS(Apple Silicon)에서 모두 동작합니다. 🤗 Transformers에서 이를 구현하려면 사용자가 quanto 라이브러리를 설치하고 generate 호출 시 cache_implementation="quantized" 인자와 함께 cache_config 딕셔너리를 지정해야 합니다.

# Example usage
out = model.generate(
    **inputs, 
    do_sample=False, 
    max_new_tokens=20, 
    cache_implementation="quantized", 
    cache_config={"backend": "quanto", "nbits": 4}
)

Sources