온디바이스 LLM을 위한 KV 캐시 양자화
2GB RAM을 탑재한 안드로이드 기기에서 Llama 3.2 3B를 실행하는 것은 어렵습니다. 대부분의 개발자는 모델 가중치(weights)에만 집중하지만, 이는 실수입니다. 진짜 메모리 도둑은 바로 KV 캐시입니다.
KV 캐시는 대화가 진행됨에 따라 커집니다. 표준 FP16 정밀도를 사용하면 캐시가 수백 메가바이트를 차지하게 됩니다. 이로 인해 단 몇 번의 대화만으로도 앱이 충돌할 수 있습니다.
다음 세 가지 단계를 통해 이 문제를 해결할 수 있습니다.
- 혼합 정밀도 양자화(Mixed-Precision Quantization) 사용 Key와 Value는 동일한 정밀도가 필요하지 않습니다. Key 캐시는 낮은 정밀도에서도 잘 작동하지만, Value 캐시는 그렇지 않습니다.
- Key에는 INT4를 사용하세요.
- Value에는 INT8을 사용하세요.
이 방식을 사용하면 캐시 크기를 62% 줄일 수 있습니다. 2048 토큰 컨텍스트의 경우, 224MB에서 84MB로 줄어듭니다. 이는 모델 가중치를 변경하지 않고도 가능합니다.
슬라이딩 윈도우 제거(Sliding Window Eviction) 구현 모든 토큰을 활성 메모리에 유지할 수는 없습니다. 슬라이딩 윈도우를 사용하여 가장 최근의 1536개 토큰만 유지하세요. 시스템 프롬프트를 보존하기 위해 처음 64개 토큰은 앵커(anchor)로 남겨둡니다.
플래시 스필링(Flash Spilling) 활용 토큰이 슬라이딩 윈도우를 벗어나면 플래시 저장소로 이동시킵니다. 안드로이드에서는 메모리 맵 파일(memory-mapped files)을 사용하세요. 최신 UFS 4.0 스토리지는 지연 없이 이 데이터를 다시 메모리로 페이징할 수 있을 만큼 충분히 빠릅니다.
결과는 놀랍습니다. Snapdragon 8 Gen 3 기준:
- 피크 메모리가 2GB 제한 미만으로 떨어집니다.
- 최대 대화 횟수가 4회에서 12회 이상으로 증가합니다.
- 캐시 크기가 작아지면 메모리 대역폭을 더 효율적으로 사용하므로 토큰 생성 속도가 향상됩니다.
- 모델 품질은 거의 동일하게 유지됩니다.
다음 실수를 피하세요:
- Key와 Value를 동일한 수준으로 양자화하지 마세요. 품질 저하가 발생합니다.
- 서멀 스로틀링(thermal throttling)을 무시하지 마세요. 지속적인 추론은 발열을 유발합니다. 성능 관리를 위해 Android Thermal HAL을 확인하세요.
- 캐시 수명 주기를 잊지 마세요. 메모리 누수를 방지하기 위해 항상 매핑된 버퍼를 적절한 스코프(scope)에 연결하세요.
기능을 구현하기 전에 메모리 예산을 먼저 세우세요.
출처: https://dev.to/software_mvp-factory/kv-cache-quantization-for-on-device-llms-kf