Stratified Sparse Attention ускоряет декодинг в 2,96 раза на контексте 1M

Команда исследователей предложила Stratified Sparse Attention (SSA) — способ обращаться к малой доле KV-кэша и при этом терять заметно меньше качества, чем у прежних разреженных схем. На RULER-HARD с Ministral-3-8B при 50-кратной разреженности внимания SSA поднимает точность с 81,8% до 84,2% на контексте 32K и с 63,8% до 69,3% на 128K. Это 94% и 85% от результата полного dense-внимания. На контексте в миллион токенов оптимизированные CUDA-ядра дают ускорение декодирования до 2,96 раза относительно FlashInfer.

Где разреженное внимание теряет качество

При длинном контексте узким местом инференса становится не столько вычисление, сколько чтение растущего KV-кэша на каждом шаге декодирования. Sparse attention снижает цену, обращаясь только к части кэша. Здесь есть развилка.

Первый путь — детерминированный отбор. Top-k или top-p берут токены с наибольшими оценками. Это надёжно ловит главное, но всё остальное внимание, которое по отдельности мало, а в сумме может быть значимым, просто выбрасывается. Второй путь — семплирование. Оно лучше покрывает хвост распределения, но платит дисперсией: результат от запуска к запуску шумит сильнее.

Гибриды склеивают top-k с равномерной выборкой. Авторы показывают, что при агрессивной разреженности разрыв с dense-вниманием у них всё равно остаётся большим. Равномерная выборка слепа: она тратит бюджет на токены, которые почти не влияют на ответ.

Как SSA делит кэш на страты

SSA строит единую схему из двух частей. Токены с высокими оценками модель читает детерминированно. Остаток кэша не отбрасывается и не выбирается наугад: токены из него семплируются с вероятностями, пропорциональными прокси-оценкам внимания. Чтобы вклад выбранных токенов не искажал картину, SSA вводит поправку на вероятность их включения в выборку. Так восстанавливается масса внимания за пределами отобранного набора.

Идея близка к классической стратифицированной выборке: важную область покрываем целиком, а остальную оцениваем по подвыборке с весами. Авторы подчёркивают, что SSA — обобщение. Top-k, top-p и гибрид top-k с семплированием оказываются его частными случаями, так что метод можно рассматривать как общую рамку для всего семейства.

Цифры и ускорение

Оценка шла на четырёх бенчмарках — RULER, LongBench, LOFT и AIME — и пяти моделях от 3B до 27B параметров. При одинаковой разреженности SSA стабильно обходит базовые методы. Главные показатели:

• RULER-HARD, Ministral-3-8B, 32K: 81,8% → 84,2% против сильнейшего практического бейзлайна, 94% от dense-внимания;
• RULER-HARD, Ministral-3-8B, 128K: 63,8% → 69,3%, 85% от dense-внимания;
• режим: 50-кратная разреженность внимания;
• скорость: до 2,96x по декодированию относительно FlashInfer на контексте 1M;
• код открыт: github.com/skylight-org/SSA-Paper-Code.

Важна и вторая часть результата. Метод с лучшей точностью обычно требует более сложных операций, и ускорение может исчезнуть. Здесь авторы показывают на оптимизированных CUDA-ядрах, что выигрыш в качестве не съедается накладными расходами. Рост точности на 128K заметнее, чем на 32K (5,5 п.п. против 2,4 п.п.), то есть метод помогает сильнее там, где контекст длиннее, а разреженность болезненнее.

Что получают практики

Для команд, которые обслуживают длинный контекст, KV-кэш — прямая статья расходов на GPU-память и пропускную способность. Если 50-кратное сокращение обращений к кэшу сохраняет 85–94% качества dense-внимания на сложных задачах, у инженеров появляется ещё один рычаг между стоимостью и точностью. Особенно это интересно для агентных сценариев и работы с большими кодовыми базами, где контекст растёт быстро.

Ограничения тоже видны. Часть цифр относится к одной модели, Ministral-3-8B, и 128K остаётся с заметным отставанием от dense-режима: 85% — хороший результат, но не паритет. Для продакшена придётся проверять метод на собственных задачах и совместимость ядер с используемым стеком. Открытый код при этом снижает порог: воспроизвести эксперименты и сравнить SSA с текущими top-k-схемами можно без ожидания реализации в популярных фреймворках.