Сколько памяти требует 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 8B32 / 8 / 1288,1921 GiB
Llama 3.1 8B32 / 8 / 128131,07216 GiB
Llama 3.1 70B80 / 8 / 1284,0961.25 GiB
Llama 3.1 70B80 / 8 / 128131,07240 GiB

При бюджете кэша 40 GiB этот расчёт даёт максимум 32 полные последовательности модели 70B по 4,096 токенов. Здесь учтены только данные кэша; задавайте предел параллельности после резервирования памяти для рантайма.

Проверьте фактическое выделение памяти на сервере

Веса, временные буферы, активации, частично заполненные блоки и зарезервированная ёмкость требуют дополнительной памяти. Распределённые серверы могут разделять или реплицировать KV-головы. Поддерживаемый кэш с одним байтом на значение вдвое уменьшает объём самих данных, но коэффициенты масштабирования, поддержка в реализации и качество могут изменить доступную ёмкость.

Измерьте пиковое выделение памяти и вытеснение последовательностей при типичных длинах промптов и ответов. PagedAttention улучшает распределение блоков; он не устраняет сохранённые данные ключей/значений и не делает длину последовательности несущественной.

Инженерное руководство: память KV cache содержит исполняемый расчёт ёмкости.