⚡ Как ускорить LLM на длинном контексте
На
AI RnD Day Никита Арсенин разобрал альтернативы классическому механизму внимания и показал результаты экспериментов команды. Он рассказал, как современные LLM избавляются от главного ограничения обычного attention — квадратичной зависимости стоимости от длины контекста, и почему в новых моделях всё чаще появляются гибрные подходы.
📌 TLDR
Ускорять внимание можно двумя способами: выбирать из контекста только нужные токены или сжимать историю в состояние фиксированного размера. Первый подход сокращает дорогую агрегацию, второй убирает растущий KV-кеш. Но у обоих есть ограничения по доступу к деталям, поэтому их комбинируют с полным вниманием. В экспериментах команды
sparse attention дал ускорение генерации в
4,8 раза на длинном контексте, а гибрид
GDN + Attention Residuals — отдельно
+6,2 п. п. на
MMLU.
🔎 Sparse attention: сначала найти, потом прочитать
В
классическом трансформере query сравнивается со всеми
keys, затем модель агрегирует соответствующие
values. А в показанной схеме разреженного внимания дешёвый индексатор просматривает контекст и выбирает top-k позиций. Дорогая агрегация выполняется только по ним. Для последовательности длины L её стоимость снижается с квадратичной O(L²) до O(Lk). Но поиск не бесплатный: в этой схеме индексатор всё ещё даёт квадратичный член. Выигрыш в том, что он заметно дешевле полного внимания.
🧠 Linear attention
Отдельные
key/value всех прошлых токенов не хранятся. Они последовательно обновляют матрицу состояния S. В простейшей схеме со слайда: Sₜ = Sₜ₋₁ + vₜkₜᵀ
При фиксированных размерностях состояния память и стоимость одного рекуррентного шага не растут с длиной контекста: O(1) вместо O(L). Суммарная работа механизма внимания для L шагов становится O(L), а не O(L²).
Обратная сторона медали — история сжата. Нельзя просто обратиться к отдельному старому KV, как в full attention. Поэтому гибрид может чередовать, например, три линейных блока с одним полным: большую часть слоёв сделать дешевле, но оставить доступ ко всем позициям контекста.
🛠 Инфраструктура
Наивные реализации
DSA и линейного внимания поверх
LLaMA-3.1-8B-Instruct показали ускорение в
2,6–3,75 раза на длинном контексте. А при сравнении PyTorch-реализации
Gated DeltaNet с FlashAttention-3 линейное внимание стало быстрее лишь примерно на 200K токенов.
Наша команда не ограничилась заменой слоя: адаптировала cuDNN-ядра для DSA/NSA, написала Triton-ядра для DSA и DSA + sliding window attention, оптимизировала линейные механизмы и
Attention Residuals. Последние меняют смешивание выходов по глубине сети, а не отбор токенов контекста.
📊 Что получилось
Модель с разреженным вниманием MSA справилась с генерацией примерно за 1,9 секунды, а базовый вариант с обычным вниманием MQA — за 9,2 секунды. Ускорение — в 4,8 раза.
В отдельном эксперименте гибрид GDN + Attention Residuals обошёл full-attention baseline на
6,2 п. п. по усреднённому MMLU после примерно 110В обучающих токенов. Важно отметить, что это результаты конкретных конфигураций, а не гарантия для любой LLM.
🔗 Более подробно — в презентации Никиты (в комментариях файлом) и в
записи выступления (тайминг:
1:07:00).
#llm #attention #conference #paperwatch
