Квантование KV-кэша для локальных LLM
Запуск Llama 3.2 3B на Android-устройстве с 2 ГБ оперативной памяти — задача непростая. Большинство разработчиков сосредотачиваются на весах модели. Это ошибка. Настоящий «убийца» памяти — это KV-кэш.
KV-кэш растет по мере общения. Если использовать стандартную точность FP16, кэш поглощает сотни мегабайт. Это приводит к вылету приложения всего через несколько реплик.
Это можно исправить с помощью трех конкретных шагов.
- Используйте смешанное квантование (Mixed-Precision Quantization) Ключам (Keys) и значениям (Values) не требуется одинаковая точность. Кэши ключей хорошо справляются с низкой точностью. Кэши значений — нет.
- Используйте INT4 для ключей.
- Используйте INT8 для значений.
Такой подход уменьшает размер кэша на 62%. Для контекста в 2048 токенов объем памяти снижается с 224 МБ до 84 МБ. Это происходит без изменения весов модели.
Внедрите вытеснение скользящим окном (Sliding Window Eviction) Невозможно держать каждый токен в активной памяти. Используйте скользящее окно, чтобы хранить только последние 1536 токенов. Оставьте первые 64 токена в качестве «якорей» для сохранения системного промпта.
Используйте сброс в Flash-память (Flash Spilling) Когда токены выходят за пределы скользящего окна, перемещайте их во Flash-память. Используйте memory-mapped файлы в Android. Современная память UFS 4.0 достаточно быстрая, чтобы подгружать эти данные обратно в память без задержек.
Результаты значительны. На Snapdragon 8 Gen 3:
- Пиковое потребление памяти падает ниже лимита в 2 ГБ.
- Максимальное количество реплик в диалоге увеличивается с 4 до более чем 12.
- Скорость генерации токенов растет, так как меньшие кэши эффективнее используют пропускную способность памяти.
- Качество модели остается практически неизменным.
Избегайте этих ошибок:
- Не квантуйте ключи и значения до одного и того же уровня. Вы потеряете в качестве.
- Не игнорируйте троттлинг из-за перегрева. Длительный инференс вызывает нагрев. Используйте Android Thermal HAL для управления производительностью.
- Не забывайте о жизненном цикле кэша. Всегда привязывайте отображаемые (mapped) буферы к соответствующей области видимости, чтобы избежать утечек памяти.
Сначала планируйте бюджет памяти, а затем — функционал.
Источник: https://dev.to/software_mvp-factory/kv-cache-quantization-for-on-device-llms-kf