Сколько памяти требует KV cache для одного запроса к LLM?
Автоматический перевод
Эта статья была автоматически переведена с оригинальной английской версии.
У обычного трансформера с полным аттеншном объём памяти KV cache растёт с числом сохранённых токенов, слоёв, голов ключей/значений и байтов на значение. Умножьте объём на одну последовательность на число параллельных последовательностей, чтобы оценить объём самих данных. Затем зарезервируйте дополнительную память для весов и работы рантайма.
Эта оценка помогает инженерам сервинга подобрать ресурсы под нагрузку. Латентный многоголовый аттеншн (MLA), кэши со скользящим окном, гибридные и рекуррентные кэши требуют расчётов с учётом архитектуры.
Используйте KV-головы, а не головы запросов
Кэш хранит проекции ключей и значений предыдущих токенов, чтобы модель не вычисляла их заново на каждом шаге генерации. Плотный аттеншн по-прежнему читает растущую историю токенов. Для кэшей полного аттеншна с отдельными головами ключей/значений, как в MHA, MQA или GQA, и последовательностей одинаковой длины:
cache bytes = 2 × layers × KV_heads × head_dimension
× cached_tokens × concurrent_sequences × bytes_per_element
Множитель два учитывает ключи и значения. Для последовательностей разной длины используйте сумму количеств сохранённых токенов. Если в конфигурации модели явно задан head_dim, используйте его; hidden_size / num_attention_heads — лишь запасной вариант для архитектур, в которых действует это соотношение.
Таблица архитектуры Llama 3 от Meta задаёт размерности для этих примеров с кэшем FP16/BF16 по два байта на значение:
| Модель | Слои / KV-головы / размерность головы | Токены на последовательность | Объём самих данных кэша на последовательность |
|---|---|---|---|
| Llama 3.1 8B | 32 / 8 / 128 | 8,192 | 1 GiB |
| Llama 3.1 8B | 32 / 8 / 128 | 131,072 | 16 GiB |
| Llama 3.1 70B | 80 / 8 / 128 | 4,096 | 1.25 GiB |
| Llama 3.1 70B | 80 / 8 / 128 | 131,072 | 40 GiB |
При бюджете кэша 40 GiB этот расчёт даёт максимум 32 полные последовательности модели 70B по 4,096 токенов. Здесь учтены только данные кэша; задавайте предел параллельности после резервирования памяти для рантайма.
Проверьте фактическое выделение памяти на сервере
Веса, временные буферы, активации, частично заполненные блоки и зарезервированная ёмкость требуют дополнительной памяти. Распределённые серверы могут разделять или реплицировать KV-головы. Поддерживаемый кэш с одним байтом на значение вдвое уменьшает объём самих данных, но коэффициенты масштабирования, поддержка в реализации и качество могут изменить доступную ёмкость.
Измерьте пиковое выделение памяти и вытеснение последовательностей при типичных длинах промптов и ответов. PagedAttention улучшает распределение блоков; он не устраняет сохранённые данные ключей/значений и не делает длину последовательности несущественной.
Инженерное руководство: память KV cache содержит исполняемый расчёт ёмкости.