KV Cache-Centric Analysis of Long-Context Methods

В последние годы большие языковые модели (LLMs) значительно улучшили свои возможности обработки длинных контекстов, что открывает новые горизонты для приложений, таких как понимание кода на уровне репозитория, ответ на вопросы по длинным документам и многослойное обучение в контексте. Однако, с увеличением длины контекста возникают серьезные проблемы, связанные с вычислительной сложностью и эффективностью использования памяти. В этой статье мы рассмотрим концепции, представленные в недавней работе, посвященной анализу методов длинного контекста с фокусом на кэш KV (key-value).

Методы обработки длинных контекстов становятся стандартом для LLM, многие из которых поддерживают окна контекста от 128K до 10M токенов. Эти расширенные окна контекста открывают широкий спектр приложений, однако они также создают уникальные вызовы для инференса LLM из-за высоких вычислительных затрат и требований к памяти. Это привело к разработке эффективных решений для длинного контекста, использующих разреженность на различных этапах работы с кэшем KV.

Проблема кэша KV

Кэш KV играет ключевую роль в снижении вычислительных затрат во время декодирования, что делает его важным элементом в архитектуре LLM. Однако существующие бенчмарки часто оценивают методы обработки длинного контекста только на основе одиночных запросов, игнорируя полный жизненный цикл кэша KV в реальных приложениях. Это упущение критично, поскольку повторное использование кэша KV стало важным компонентом в популярных фреймворках инференса, таких как vLLM и SGLang.

SCBench: Новый бенчмарк для оценки методов длинного контекста

Чтобы заполнить этот пробел, был предложен SCBench (Shared Context Bench), который служит комплексным бенчмарком для оценки методов длинного контекста с фокусом на кэш KV. SCBench включает четыре ключевых аспекта:

1. Генерация кэша KV: Этот этап включает в себя обработку входного запроса и генерацию кэша KV для последующего декодирования.
2. Сжатие кэша KV: На этом этапе применяются техники сжатия для уменьшения объема памяти, необходимого для хранения кэша.
3. Извлечение кэша KV: Этот этап направлен на использование ранее сгенерированного кэша KV для обработки новых запросов, что позволяет пропускать этап генерации.
4. Загрузка кэша KV: На последнем этапе осуществляется динамическая загрузка только части кэша KV для каждого шага декодирования, что экономит память и вычислительные ресурсы.

Задачи и сценарии

SCBench включает в себя 12 задач с двумя режимами общего контекста, охватывающими четыре категории возможностей обработки длинного контекста: извлечение строк, семантическое извлечение, глобальная информация и многозадачность. Каждая задача включает общий контекст и несколько последовательных запросов.

Анализ методов длинного контекста

В рамках SCBench были проанализированы 13 методов длинного контекста, включая модели на основе Gated Linear RNN, гибридные модели Mamba-Attention и методы разреженного внимания. Результаты показали, что методы с под-O(n) памятью испытывают трудности в многослойных сценариях, в то время как разреженное кодирование с O(n) памятью и под-O(n²) вычислительная сложность в процессе предзаполнения демонстрируют устойчивую производительность.

Влияние разреженности

Динамическая разреженность обеспечивает более выразительные кэши KV по сравнению со статическими паттернами, а разреженность на уровне слоев в гибридных архитектурах снижает использование памяти при сохранении высокой производительности. Кроме того, были выявлены проблемы смещения распределения внимания в сценариях длинной генерации.

Выводы и рекомендации

В ходе исследования было установлено, что методы, поддерживающие кэш KV с O(n), превосходят в многозадачных сценариях, в то время как методы с под-O(n) показывают хорошие результаты только в одиночных запросах. Эти результаты подчеркивают важность многослойных сценариев с общим контекстом при разработке и оценке методов длинного контекста.

SCBench предоставляет более реалистичную оценку методов длинного контекста, предлагая ключевые идеи для улучшения архитектуры LLM и будущих исследований в этой области.