При авторегрессионной генерации каждый новый токен «смотрит» на все предыдущие токены. Без кеширования на каждом шаге пришлось бы заново считать key/value-представления для прошлых токенов и внимание по всей последовательности — это квадратичная стоимость. KV-кешпамять, в которой хранятся проекции key и value из каждого слоя для всех токенов, сгенерированных до текущего момента, что избавляет от их пересчёта хранит проекции key и value из каждого слоя для каждого токена, сгенерированного до текущего момента, и избавляет от пересчёта. Но сам кеш стоит памяти, и на длинных контекстах и больших батчах он перерастает веса модели. Разберём, как устроен разреженный инференс, который снижает эту цену, и что при этом теряется.
Сколько памяти съедает KV-кеш
Размер KV-кеша на один токен считается как per_token = 2 × layers × kv_heads × head_dim × dtype, где множитель 2 учитывает отдельные тензоры key и value. Для Llama 3 70B (80 слоёв, 8 KV-голов с GQA, head_dim 128, float16) это даёт per_token = 2 × 80 × 8 × 128 × 2 = 327 680 байт ≈ 320 КБ. Для последовательности длиной 4096 токенов — около 1,3 ГБ, а при batch size 32 — около 41 ГБ. Этот объём часто больше веса самой модели, потому что кеш растёт с длиной последовательности и размером батча, тогда как веса фиксированы. GQAсхема, в которой несколько query-голов разделяют одни key- и value-головы уменьшает размер кеша именно за счёт этого разделения.
Почему полное внимание не масштабируется
Переход к механизму внимания (Bahdanau et al., 2014) был мотивирован ограничениями RNNрекуррентных сетей, обрабатывающих последовательность по шагам. У RNN три проблемы: потеря контекста на больших последовательностях, потому что всю историю приходится сжимать в вектор фиксированного размера (information bottleneckузкое место, при котором вся история сжимается в вектор фиксированной длины); трудности параллелизма, поскольку вычисления на каждом временном шаге начинаются только после завершения предыдущего, что ограничивает загрузку GPU; нестабильность обучения — при обратном распространении во времени граф разворачивается с одними и теми же параметрами, и если собственные значения матрицы весов скрытого состояния больше или меньше 1, это приводит к взрыву или затуханию градиентов. Изначально внимание применялось в машинном переводе: декодер динамически обращался к разным частям исходного предложения через веса внимания как взвешенную сумму энкодерных состояний.
Как устроено самоиндексирующееся внимание
Сжатые ключи
Входом служит исходная последовательность векторов-ключей до текущего момента. Она разбивается на блоки длины l с шагом d: длина блока l — это сколько последовательных ключей попадает в один блок, а шаг d — на сколько токенов вперёд начинается следующий блок; если шаг меньше длины блока, соседние блоки частично перекрываются, что сохраняет непрерывность информации. Каждый блок из l последовательных ключей преобразуется обучаемым MLPмногослойным перцептроном в один сжатый вектор размерности d_model: токены блока конкатенируются, к ним добавляется позиционное кодирование, чтобы сохранить порядок внутри блока. Результат — множество сжатых ключей, состоящее из меньшего числа элементов, чем исходное число токенов t. Пример: блок 32 токена × 512 = 16384 параметров сжимается до 1 × 512, сокращение в 32 раза по числу токенов.
Параметр сжатия выбирается для уменьшения фрагментации информации. Для длинных последовательностей эффективнее большие размеры блоков (256 вместо 32), а оптимальное соотношение stride/block_size близко к 0.5 — это даёт баланс между сжатием и сохранением информации.
Оценки важности и выбор блоков
Для выборочного внимания оценки важности вычисляются как скалярное произведение текущего запроса на каждый сжатый ключ с последующей Softmaxнормализацией в распределение вероятностей. Затем выбираются блоки с наивысшими оценками: берётся множество индексов по рангу элементов вектора вероятностей, где ранг 1 соответствует наивысшей вероятности, а число выбираемых блоков задаётся гиперпараметром.
Три ветви и gate-оценка
Выход внимания формируется как комбинация трёх стратегий: сжатия, выбора и скользящего окна. Они работают параллельно: сжатое внимание собирает глобальную информацию путём сжатия блоков токенов, выборочное обрабатывает важные блоки и сохраняет детали, скользящее окно обрабатывает локальный контекст. Вкладом каждой стратегии и общим количеством переназначенных ключей и значений управляет gate-оценкавеличина, вычисляемая на основе входных признаков с помощью MLP и сигмоидной активации.
Скользящее окно
Для каждого текущего запроса q_t скользящее окно шириной w включает только w последних токенов; токены за пределами окна игнорируются. Это даёт постоянную вычислительную сложность O(w) на токен вместо линейного роста O(t), фиксированное потребление памяти и стабильное время обработки. В практическом примере при генерации слова в позиции 10000 окно w=512 учитывает токены с позиций 9488–10000, сжатое внимание даёт общую информацию обо всём предыдущем контексте, а выборочное выделяет ключевые токены. Скользящее окно предотвращает «короткое замыкание», когда сжатие и выбор могут пропустить важную локальную информацию, и повышает арифметическую интенсивность.
Что даёт каждая ветвь
Сжатие фраз в обобщающие блоки сокращает объём данных, повышает эффективность обработки и сохраняет основную информацию текста — так обеспечивается глобальное понимание контекста. Выборочное внимание сохраняет детальную информацию на уровне отдельных токенов, но только для наиболее важных частей последовательности. Скользящее окно фиксирует важную информацию, находящуюся поблизости. Вместе они дают и глобальную осведомлённость, и локальную точность.
Совместимость с FlashAttention и обучение
Для ускорения на уровне FlashAttention при обучении и prefillэтапе обработки входного промпта до начала генерации реализовано аппаратно-оптимизированное ядро разреженного внимания на Triton, оптимизированное под Tensor Core и доступ к памяти. Применены блочный режим доступа к памяти, циклическое планирование, групповая загрузка данных, совместное получение KV и внешний цикл с планированием сетки.
Совместимость обеспечивается через kernel fusionобъединение всех основных операций вычисления attention в единый super-kernel — это избавляет от промежуточных обращений к HBM и многократных запусков ядер. Block-sparse Flash Attention вводит структурную разреженность на уровне блоков и выполняет шаги только для заранее заданных важных блоков, отмеченных в маске.
Настройки обучения: архитектура GQA и MoEсмесь экспертов, 27 млрд параметров, 3 млрд активных, 30 слоёв, скрытая размерность 2560, 4 группы GQA, 64 головы внимания, 72 маршрутизируемых эксперта и 2 общих, 6 выбираемых экспертов. Предварительное обучение — на 270 млрд токенов с длиной 8k, затем дообучение на 32k с YaRNметодом расширения контекста.
Оптимизации FlashAttention 2 и 3
FlashAttention-2 использует атомарные операциинизкоуровневые инструкции, выполняющие операцию над ячейкой памяти как единую неделимую транзакцию — они исключают перезапись результатов другими потоками, но не дают детерминизма. При необходимости детерминизма режим делают опциональным с последовательной обработкой блоков столбцов K/V, теряя 10–20% производительности. FA2 эффективно разделяет работу между варпамигруппами потоков внутри блока для минимизации чтения/записи в shared memory: в forward pass вместо схемы «split-K» работа распределяется иначе, в backward pass применяется аналогичная схема с некоторой синхронизацией.
FlashAttention-3 добавляет три нововведения. Асинхронность Tensor Cores и TMATensor Memory Accelerator, обеспечивающего перемещение данных через warp-специализацию: одни варпы грузят данные через TMA, другие выполняют матричные умножения через WGMMAинструкцию матричного умножения на тензорных ядрах архитектуры Hopper, выполняемую группой варпов. Перекрытие GEMM и Softmax через ping-pong scheduling на основе синхронизации барьеров — плюс 10–15%. Эффективная работа с FP8 через блочное квантование и Incoherent Processing: умножение Q и K на случайную ортогональную матрицу (преобразование Адамара) снижает ошибку квантования в 2.6 раза.
Для новых архитектур (Hopper и выше) FA2 даёт низкую загрузку GPU — лишь 35% от теоретического максимума на H100 против 80–90% для GEMM. FA3 как адаптация под новые архитектуры даёт прирост скорости в 1.5–2 раза и загрузку до 75% на H100.
Пограничные случаи и режимы отказа
Сжатие сокращает число обрабатываемых токенов, фиксируя семантическую информацию более высокого уровня, — это даёт глобальную осведомлённость. Скользящее окно шириной w учитывает только w последних токенов и гарантирует, что непосредственный локальный контекст не будет потерян из-за сжатия или селективного выбора, которые могут упустить важные локальные связи. Компромисс: сжатое внимание жертвует гранулярностью на уровне отдельных токенов ради обработки длинных контекстов, и для полной архитектуры это компенсируется комбинированием с другими механизмами внимания.
Что видит клиент: качество и задержки
На бенчмарках MMLU, MMLU-PRO, CMMLU, BBH, GSM8K, MATH, DROP, MBPP и HumanEval NSA превосходит полное внимание по большинству показателей, несмотря на более высокую разреженность. На LongBench средний балл NSA — 0,469, что на 0,032 выше полного внимания и на 0,046 выше Exact-Top. На AIME NSA-R после контролируемой тонкой настройки превосходит Full Attention-R при длине контекста 8 тыс. и 16 тыс. Для последовательностей 64k NSA даёт значительный выигрыш в скорости на декодировании, прямом проходе и обратном распространении. Реализация ядра на Triton значительно снижает задержку по сравнению с FlashAttention-2 при всех длинах контекста, и улучшение становится заметнее с ростом длины входных данных.
На задачах длинного контекста сжатое внимание даёт существенное преимущество для последовательностей >10K токенов, тогда как для <1000 токенов стандартное внимание эффективнее, и потенциально позволяет обрабатывать контексты 100K+ токенов на стандартном оборудовании.
Как сравнивали и что получилось
Стенд и методология
Популярные реализации Flash Attention 2 работают исключительно с архитектурами Ampere и новее, исключая Turing (T4, RTX 20-й серии), а в версиях Triton 3.3.0 и выше есть баг с увеличенным потреблением shared memory, заметный на старых архитектурах Turing и Ampere. Для кастомной реализации на Turing приведён пример конфигурации с num_warps=4 и num_stages=1, а также ограничения min_warps=4, max_warps=8.
Участники сравнения — модели с полным вниманием и с NSA: 27 млрд параметров, 3 млрд активных, 30 слоёв, скрытая размерность 2560, GQA с 4 группами и 64 головами внимания, MoE DeepSeekMoE с 72 маршрутизируемыми и 2 общими экспертами, 6 выбираемых экспертов. Обе модели предварительно обучаются на 270 млрд токенов с длиной 8k, затем дообучаются на 32k с YaRN. Для NSA заданы размер сжатого блока, шаг скользящего окна, размер выбранного блока, количество выбранных блоков (1 фиксированный начальный и 2 локальных) и размер скользящего окна. Для примеров Flash Attention 2 входные тензоры имеют форму (batch, n_heads, seq_len, head_dim) = (16, 8, 512, 64).
Оценка длинного контекста — на LongBench с бюджетом активируемых токенов 2560, что соответствует среднему числу токенов, активируемых NSA при длине последовательности 32 тыс. Оценка рассуждений — на AIME при длине контекста 8 тыс. и 16 тыс. BABILong доступен на Hugging Face и может быть сгенерирован с собственными задачами и шумом с помощью кода на GitHub.
Результаты
| Сценарий | Результат | ||
|---|---|---|---|
| MMLU, MMLU-PRO, CMMLU, BBH, GSM8K, MATH, DROP, MBPP, HumanEval | NSA превосходит полное внимание по большинству показателей | ||
| LongBench (средний балл) | NSA — 0,469; на 0,032 выше полного внимания и на 0,046 выше Exact-Top | ||
| AIME, контекст 8k и 16k | [[term:NSA-R | вариант NSA, дообученный с контролируемой тонкой настройкой для задач рассуждения по инструкциям]] превосходит [[term:Full Attention-R | вариант модели с полным вниманием, дообученный с контролируемой тонкой настройкой для задач рассуждения по инструкциям]] |
| Последовательности 64k | Значительный выигрыш в скорости на декодировании, прямом и обратном проходе | ||
| Задержка (ядро на Triton) | Снижается при всех длинах контекста, улучшение заметнее с ростом длины |
Оговорки
Для единообразия разреженности число активируемых токенов на запрос во всех базовых уровнях разреженного внимания установлено равным 2560. Некоторые подмножества LongBench имеют в целом низкие оценки у всех моделей и не могут служить ценным источником для сравнения — авторы исключили их из области оценки. Needle in a Haystack из-за включения схожих задач в обучающую выборку быстро привёл к переобучению на них и вырождению бенчмарка. Большинство несинтетических бенчмарков вроде LongBench не работает на контексте более 16 тысяч токенов. BABILong показывает, что с увеличением длины контекста качество у всех моделей падает, большинство моделей эффективно работают лишь с 10–20% от заявленной длины контекста, а стандартные бенчмарки не могут полноценно оценить растущий контекст.
Почему сделано именно так: компромиссы
Три механизма работают совместно: сжатие даёт глобальный обзор, выборка сохраняет важные детали, скользящее окно удерживает локальный контекст. Сжатие резко сокращает объём вычислений, но может потерять детали — поэтому добавляется точный выбор наиболее релевантных блоков. Скользящее окно ограничивает внимание фиксированным числом последних токенов, обеспечивая постоянную сложность O(w) и предотвращая потерю локальных связей. Один механизм не смог бы одновременно покрыть все эти аспекты.
От чистого RAGподхода с извлечением релевантных фрагментов и подстановкой их в промпт авторы отказываются, потому что простой семантический поиск не сохраняет порядок фактов, что даёт низкую точность даже на коротких контекстах, а за одну итерацию трудно извлечь всю необходимую информацию. Длинный контекст оправдан для пошагового рассуждения: модель может «проговорить» промежуточные шаги, и чем больше длина контекста, тем сложнее размышления и точнее результат. Загрузка всего репозитория в одно окно контекста тоже оправдывает разреженный инференс: агент SubQ Code загружает весь репозиторий в одно окно, и вопрос «где используется эта библиотека?» обрабатывается с учётом всей кодовой базы. При этом есть гипотеза, что высокая скорость достигается агрессивным прореживанием внимания, что может «срезать» важные детали в длинных зависимостях.
Практический воркфлоу оптимизации инференса
Воркфлоу начинается с профилирования: с помощью nvidia-smi, PyTorch Profiler или NSight Systems находят узкое место. Затем включается continuous batchingвставка новых запросов в батч сразу после завершения предыдущих, что часто удваивает производительность при практически нулевом риске для качества и поднимает загрузку GPU с 30–40% до 80–90%. Третьим шагом идёт квантование: начинают с FP8 или INT8, измеряя качество на своей eval-выборке; INT4 применяют только при нехватке памяти. Квантование снижает точность представления с FP16/BF16 до INT8, INT4 и ниже; для INT4 (GPTQ) потеря качества — 1–3% на бенчмарках, для INT4 (AWQ) — 1–2%, для GGUF зависит от уровня квантования. Далее включается prefix caching для переиспользования KV-кеша общего префикса, настраиваются batch size и concurrency, рассматривается спекулятивный декодинггенерация K кандидатов draft-моделью с проверкой их основной моделью одним проходом — принимаются 70–90% кандидатов, что даёт 3–6 токенов за один прямой проход основной модели, — и масштабирование через TP/PP. Разреженное внимание в воркфлоу не упоминается как отдельный шаг; вместо него описан Flash Attention, используемый по умолчанию во всех основных фреймворках и снижающий память под вычисление внимания с O(N²) до O(N).
Что из этого следует на практике
- KV-кеш растёт линейно с длиной последовательности и размером батча, поэтому на длинных контекстах он может превысить веса модели — это и есть основная причина переходить к разреженному вниманию.
- Разреженное внимание не заменяет полное, а комбинирует три ветви: глобальный обзор через сжатие, детали через выбор важных блоков, локальный контекст через скользящее окно. Отказ от любой из них оставляет непокрытый аспект.
- Сжатие выгодно только на длинных последовательностях: для >10K токенов оно даёт существенное преимущество, а для <1000 токенов стандартное внимание эффективнее.
- Разреженность даёт ускорение, но её цена — потеря гранулярности на уровне отдельных токенов; компенсируется она только остальными механизмами внимания.
- Качество на бенчмарках не деградирует: NSA превосходит полное внимание по большинству показателей, а на LongBench — на 0,032 балла. Но цифры не показывают поведение на контекстах, которые бенчмарки не покрывают: большинство несинтетических тестов не работает за пределами 16 тысяч токенов, а модели эффективно используют лишь 10–20% заявленной длины.
- В практическом воркфлоу разреженное внимание не фигурирует как отдельный шаг — порядок оптимизации идёт от профилирования к батчингу, квантованию и кешированию префиксов, а Flash Attention используется по умолчанию.
