Ускорение E2E-инференса через оптимизацию KV-кэша. Часть I

Существует много способов ускорить инференс LLM: менять архитектуру, использовать speculative decoding или просто добавлять вычислительные ресурсы. Но есть и более практичный путь — оптимизация KV-кэша.

Её можно разделить на pre-train и post-train. Первые требуют изменений до обучения модели: это архитектурные решения вроде GQA/MQA/MLA, смешивание глобального и локального атеншена, а также другие модификации, которые обычно стоят дорого из-за переобучения.

Post-train-методы можно применять к уже готовой модели: это различные sparse-стратегии, pruning, удаление повторов токенов и другие техники, которые уменьшают объём KV или сокращают число обращений к нему во время инференса.

KV-бюджеты удобно делить на dense и sparse, отдельно для prefill и отдельно для decode. В варианте dense prefill + dense decode (обычный KV-кэш) каждый новый Q взаимодействует со всеми K и V до него: ко всем токенам промпта и всем ранее сгенерированным токенам. Тогда KV-бюджет равен сумме длины промпта и длины генерации.

Если сделать sparse только на prefill, а decode оставить плотным, то Q перестаёт смотреть на весь промпт, но общий выигрыш заметен в основном в сценариях «длинный промпт — короткий ответ». Если же оставить dense prefill и сделать sparse decode, это часто релевантно reasoning/CoT-сценариям. Sparse и на prefill, и на decode даёт максимальную экономию бюджета, но обычно сильнее всего ухудшает качество.

Sparse можно строить по-разному. Если пересчитывать важные токены на каждом шаге decode, то качество станет выше, но скорость падает. Если пересчитывать раз в несколько токенов, то получается быстрее, но нужно удерживать локальный контекст между пересчётами, иначе модель начинает терять связность.

Один из сильных post-train-методов оптимизации KV-кэша — ShadowKV, который позволяет получать минимальные просадки на бенчмарках без дообучения и увеличивает throughput до трёх раз. О нём мы подробно поговорим в следующей части.

Разбор подготовил❣ Владислав Кругликов

Душный NLP