Как KV Sparsity обеспечивает ускорение в 1,5 раза для vLLM

Zachary, эксперт по алгоритмам в Novita AI

Направление исследований: Ускорение инференса

Введение

За последний год, с момента появления H2O, количество исследовательских работ по KV sparsity значительно возросло. Однако одной из главных проблем реального применения является существенный разрыв между академическими исследованиями и практической реализацией. Например, такие фреймворки, как vLLM, используют PagedAttention — технику управления памятью на основе страниц, которая несовместима с большинством алгоритмов разреживания или работает хуже, чем PagedAttention. Из-за этого алгоритмы разреживания не могут эффективно применяться в производственных средах.

Проанализировав недавние академические статьи за последний год в области KV sparsity, мы внесли изменения в фреймворк vLLM на основе KV sparsity. Эти изменения были согласованы с встроенными оптимизациями vLLM, такими как Continuous Batching, FlashAttention и PagedAttention. В итоге, сравнив модифицированный фреймворк с современными инференс-фреймворками на наиболее часто используемых моделях, параметрах и оборудовании в производстве, мы добились ускорения инференса в 1,5 раза.

Прежде чем обсуждать KV sparsity, необходимо рассмотреть характеристику Massive Activations в LLM. В LLM небольшое число значений активации значительно активнее остальных, иногда превосходя их в 100 000 раз и более. Иными словами, небольшая подгруппа токенов играет критическую роль. Это наблюдение лежит в основе применения техник KV sparsity, где сохранение только важных токенов может повысить производительность инференса.

Конкретные данные анализа Llama2 и Llama3 дополнительно подтверждают это наблюдение, как показано на рис. 1 и рис. 2.

Рис. 1. Величины активаций (ось 𝐳) в LLaMA2-7B. Оси 𝐱 и 𝐲 — размерности последовательности и признаков. Для данной модели мы наблюдаем, что активации с огромными величинами появляются в двух фиксированных размерностях признаков (1415, 2533) и двух типах токенов — начальном токене и первом символе точки (.) или новой строки (\n).

В отличие от Llama2, Llama3 не проявляет характеристику Massive Activations во всех слоях. В нижних слоях активации распределены более равномерно, в то время как средние слои демонстрируют локализованные паттерны внимания. Только в верхних слоях феномен Massive Activations становится заметным. Это наблюдение объясняет, почему такие модели, как Llama3, используют иерархические техники разреживания, которые адаптируют разреживание для разных слоёв в соответствии с их специфическими паттернами активации.

Рис. 2. Паттерны внимания при генерации с дополнением по retrieval по слоям в LLaMa (Touvron et al., 2023a;b) показывают, что в нижних слоях модель демонстрирует широкоспектральный режим внимания, равномерно распределяя веса внимания по всему контенту. В средних слоях внимание становится более локализованным внутри каждого документа, что указывает на уточнённую агрегацию информации (пунктирные красные треугольные формы на слоях 6 и 10). Это завершается в верхних слоях, где «массовое внимание» фокусируется на нескольких ключевых токенах (концентрированные столбцы внимания после 18-го слоя), эффективно извлекая важную информацию для ответов.

Системная модель

После введения характеристики Massive Activations в LLM мы получили базовое понимание принципов KV sparsity. Далее обсудим, как реализуется KV sparsity. В процессе инференса больших моделей обычно встречаются три основных узких места: объём памяти GPU, вычислительная мощность и пропускная способность ввода-вывода. Чтобы избежать избыточных вычислений, KV часто кэшируется во время инференса, что приводит к значительному потреблению памяти GPU. В то же время вычислительная нагрузка и пропускная способность ввода-вывода становятся критическими узкими местами на этапах Prefill и Decode соответственно.

Наша модификация фреймворка vLLM основана на характеристике Massive Activations LLM и реализует послойное разреживание (т.е. разный уровень разреживания для разных слоёв). Этот подход значительно снижает накладные расходы, связанные с кэшированием KV, тем самым ускоряя процесс инференса моделей LLM.

Как показано на рис. 3, в процессе инференса мы применяем KV sparsity к разным слоям модели. Используя стратегию прореживания, мы удаляем пары KV с более низкими оценками, сохраняя те, у которых оценки выше и которые находятся ближе. Это снижает использование памяти, а также вычислительные и вводно-выводные накладные расходы, что в конечном итоге приводит к ускорению инференса.

Рис. 3. Влияние KV sparsity на модель

Результаты ускорения инференса

Основной вопрос для большинства исследователей — каково фактическое влияние KV sparsity на фреймворк vLLM. Прежде чем углубляться в детали, сначала представим результаты оценки производительности.

1. Производительность инференса

В реальных приложениях LLM наиболее часто встречающейся конфигурацией является длина ввода/вывода 4000/500. Кроме того, GPU RTX 4090 — один из наиболее широко используемых GPU в индустрии. На основе этих условий мы провели сравнительные тесты производительности пакетной обработки. Контрольная группа — vLLM 0.6.1.p2, экспериментальная группа — Novita AI Sparse 0.5.1 (версия vLLM 0.5.1, модифицированная для KV sparsity). Между ними было проведено несколько раундов сравнительных тестов производительности, ключевые метрики включали TTFT (время до первого токена, критично для пользовательского опыта) и Throughput (фактическая скорость инференса).

Итоговые результаты тестов показаны в таблице 1. Применяя KV sparsity, мы смогли увеличить пропускную способность vLLM примерно в 1,58 раза, сохранив приемлемый TTFT (P50 менее 1 секунды). Кроме того, в сценариях с большим размером пакета vLLM 0.6.1.p2 достиг предела производительности при уровне параллелизма 10, тогда как Novita AI Sparse 0.5.1 смог поддерживать стабильный TTFT даже при уровне параллелизма 20. Это демонстрирует устойчивость KV sparsity в обеспечении стабильности производительности в производственных средах.

Таблица 1: Сравнение производительности vLLM 0.6.1.p2 и Novita AI Sparse 0.5.1

2. Точность модели

Поскольку KV sparsity является алгоритмом сжатия с потерями, оценка производительности модели также важна. Мы провели оценку производительности на основе Llama3–8B, включая оценку с использованием MMLU и других бенчмарков. Результаты показывают, что потеря точности остаётся в пределах примерно 3%.

Как показано в таблице 2, мы протестировали производительность модели с помощью бенчмарков MMLU и humanities:

Таблица 2: Результаты оценки производительности Llama3–8B с использованием MMLU и Humanities

Во-вторых, мы провели оценку задач QA в длинных текстовых сценариях, часто встречающихся в практических приложениях, с длиной ввода от 7k до 30k. Степень сжатия поддерживалась более чем в 10 раз, как показано в таблице 3. Результаты показали общую потерю точности примерно на 10%.

Таблица 3: Результаты оценки длинных текстовых задач QA с длиной ввода от 7k до 30k

Ключевые техники

1. Послойное разреживание и тензорный параллелизм

Фреймворк vLLM использует стратегию планирования на уровне итераций, обычно называемую Continuous Batching. Этот подход к планированию характеризуется тем, что не ожидает завершения генерации всех последовательностей в пакете; вместо этого он обеспечивает планирование на уровне итераций, что приводит к более высокому использованию GPU по сравнению со статическим пакетированием. Как упоминалось в начале этой статьи, такие модели, как Llama3–8B, изначально обладают свойствами иерархического разреживания. Следовательно, наши модификации на основе Continuous Batching столкнулись с рядом проблем, как показано на рис. 4.

Рис. 4. Завершение семи последовательностей с использованием непрерывного пакетирования. Слева вы видите пакет после одной итерации, справа — после нескольких итераций. Как только последовательность генерирует токен конца последовательности, мы вводим новую последовательность на её место (например, последовательности S5, S6 и S7). Этот метод повышает использование GPU, избегая необходимости ожидания завершения всех последовательностей перед запуском новой.

Основная проблема, которую мы решили, — управление памятью для разных слоёв. Поскольку vLLM использует единую модель управления памятью на основе страниц, а стратегия очереди поддерживает только Full Cache или Sliding Window, мы внесли изменения в базовую структуру vLLM, чтобы обеспечить поддержку обоих вариантов.

Эта корректировка даёт каждому слою возможность выбирать между Full Cache или Sliding Window, позволяя устанавливать разные уровни разреживания для разных слоёв, как показано на диаграмме ниже.

Первая проблема, которую необходимо решить, — это управление памятью для разных слоёв. Поскольку vLLM использует единую модель управления памятью на основе страниц для всех слоёв, а стратегия очереди позволяет использовать только Full Cache или Sliding Window, нам нужно скорректировать базовую структуру vLLM. В частности, это включает поддержку как Full Cache, так и Sliding Window, чтобы разные слои могли независимо выбирать предпочтительную стратегию управления памятью. Эта корректировка в конечном итоге позволяет реализовать разную степень разреживания в разных слоях. Соответствующая структурная диаграмма показана на рисунке 5.

Рис. 5. Послойная структура управления памятью в vLLM

В целом корректировки можно разделить на три отдельные фазы:

Фаза инициализации сервиса: Во время запуска сервиса выделенная память для трёх слоёв инициализируется в основном через вычисление оставшегося пространства, а также иерархическую инициализацию структуры Metadata. В отличие от этого, оригинальный vLLM работает как полный Full Cache, не учитывая логику выделения памяти для разных слоёв.

Фаза подготовки шага: В процессе инициализации шагов в vLLM, а именно в фазах, включающих методы block manager и model runner, для разных слоёв используются разные стратегии управления таблицами блоков для контроля механизма обновления KV. Как показано на предыдущем рисунке, мы используем механизм Full Cache для первых трёх слоёв, сохраняя все KV, в то время как слои после третьего используют стратегию Sliding Window для автоматической замены KV с самыми низкими оценками. Этот подход позволяет гибко распределять пространство памяти в соответствии с потребностями каждого слоя, экономя пространство и максимизируя эффективность алгоритма разреживания.

Фаза вычисления разреживания: Эта фаза обычно выполняется после выполнения модели, где операции разреживания выполняются на основе вычисленных оценок разреживания. Ключевым моментом на этом этапе является совместимость с такими технологиями, как Tensor Parallelism, CUDA graphs и GQA. В практической инженерной реализации Tensor Parallelism требует особого внимания. Как показано на рисунке ниже, в сценарии с одной машиной и несколькими GPU необходимо использовать соответствующий механизм разделения, чтобы ограничить операции разреживания KV в пределах каждого GPU, избегая связи между GPU. К счастью, блок Attention использует подход разделения на основе ColumnLinear, что делает его реализацию относительно простой.

Рис. 6. Трёхэтапный процесс управления памятью и разреживания в vLLM

Вышеизложенное в основном касается модификаций на уровне фреймворка для поддержки иерархического разреживания в vLLM. Далее мы переключим внимание на оптимизации, реализованные на уровне CUDA.

2. Модификация Attention

Модификации на уровне CUDA в основном сосредоточены на улучшениях блока вычисления Attention, а именно FlashAttention и PagedAttention. В то время как PagedAttention относительно прост, мы уделим особое внимание модификациям, связанным с FlashAttention.

Сначала представим формулу вычисления Attention:

Как показано в формуле, процесс вычисления softmax включает вычисление произведения QK и является двухпроходным алгоритмом (требуется два прохода). Однако конечный результат O может быть достигнут с помощью алгоритма FlashAttention за один проход (требуется один проход). Логика реализации FlashAttention может быть проиллюстрирована на рис. 7 из статьи FlashAttention2. Вкратце, она заключается в поблочном обходе Q для выполнения блочных вычислений над KV с постепенным обновлением O, P и данных rowmax до завершения цикла. Наконец, деление O на ℓ позволяет выполнить однопроходное вычисление FlashAttention.

Рис. 7. Логика реализации FlashAttention2

Для нашей реализации KV sparsity важно отметить, что в процессе вычисления FlashAttention значения, необходимые для softmax, а именно rowmax/ℓ (глобальный максимум и кумулятивные значения), уже вычислены. Однако FlashAttention не возвращает результаты этих двух значений.

Имея базовое понимание FlashAttention, мы можем обсудить формулы оценки для разреживания, представленные в таких работах, как PYRAMIDKV:

где [n − α, n] — диапазон токенов инструкции. В каждом слое l и для каждой головы h выбираются top k l токенов с наивысшими оценками, и их соответствующие KV-кэши сохраняются. Все остальные KV-кэши отбрасываются и не будут использоваться в последующих вычислениях в процессе генерации.

В частности, мы вычисляем оценку выбора i-го токена для сохранения в KV-кэше как S в каждой голове внимания h по формуле:

В приведенной выше формуле A представляет собой оценку softmax, а S — сумму оценок a ближайших значений softmax. Как указано, softmax сам по себе является двухпроходным алгоритмом, который не может быть напрямую интегрирован в FlashAttention. Это также означает, что оценка разреживания не может быть реализована за один проход.

Как упоминалось ранее, хотя FlashAttention не выводит оценки softmax, он вычисляет промежуточные значения, такие как rowmax/ℓ. Поэтому, ссылаясь на Online Softmax: двухпроходный алгоритм, мы можем использовать выходные данные rowmax/ℓ для получения конкретной оценки S для KV sparsity. На практике нам нужно лишь внести улучшения в FlashAttention, показанные на рис. 7, чтобы удовлетворить это требование.

Рис. 7. Модификация FlashAttention для оценки KV sparsity

• Вывод rowmax/ℓ в FlashAttention: FlashAttention включает вывод rowmax/ℓ. Поскольку эти значения представляют собой максимум и кумулятивные суммы, их размер данных минимален, что оказывает незначительное влияние на общую производительность FlashAttention. Повторное использование rowmax/ℓ позволяет избежать необходимости во втором проходе, тем самым повышая общую вычислительную эффективность.
• Добавление шага Online Softmax: Мы вводим шаг Online Softmax, который строится поверх FlashAttention, вычисляя оценки softmax по формуле p=e((kq)−row_max)/ℓ (логика одного прохода). Затем, выполняя суммирование по строкам, мы можем вычислить оценку разреживания S.

В реализации Online Softmax необходимо уделить особое внимание двум критическим деталям:

• Корректировка порядка циклов: Поскольку S является суммой оценок по размерности k, необходимо изменить вычисление в Online Softmax с q*k на *k**q по сравнению с механизмом внимания. Эта корректировка позволяет избежать ненужной связи между разными блоками, преобразуя суммирование по столбцам в суммирование по строкам, что позволяет выполнять все вычисления внутри регистров.
• Начальная и конечная позиции q: Из-за наличия каузального кодирования и логики суммирования для ближайших a нет необходимости обходить все q со всеми k. Ссылаясь на индикаторы маски, обозначенные как 0/1 на рис. 7, мы можем определить, требуется ли вычисление для соответствующего блока, на основе начальной и конечной позиций q, что значительно экономит вычислительные ресурсы. Например, если длина входного промпта составляет 5000, а длина ближайших a — 256, обычное вычисление потребовало бы 5000*5000 операций. Однако, записывая начальные и конечные позиции, требуется только *5000**256 операций, что повышает вычислительную эффективность.

PagedAttention выполняет работу, аналогичную FlashAttention, с общей целью вычисления итоговой оценки S. Детальное объяснение PagedAttention здесь не приводится.

Заключение

Наша модифицированная версия, Novita AI Sparse 0.5.1, основанная на vLLM 0.5.1, в настоящее время в основном поддерживает такие модели, как Llama3–8B и Llama3–70B. Она работает в режиме инференса с использованием CUDA graphs и в основном развертывается на потребительских GPU, таких как RTX 4090. Остаётся значительный потенциал для оптимизации как в инженерном, так и в алгоритмическом аспектах.

Предупреждение: Следующие функции пока не поддерживаются в vLLM-0.6.2:

• Chunked-prefill
• Prefix caching
• FlashInfer и другие бэкенды, отличные от FlashAttention

Для тех, кто заинтересован в Novita AI Sparse 0.5.1, вы можете посетить Novita AI, чтобы опробовать его раньше других.

Ссылки:

[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

[3] SnapKV: LLM Knows What You are Looking for Before Generation

[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling

[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference

[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models

[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models

[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding

[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion

[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation