BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholding — лучшая статья MLSys '26.

Классная идея ускорения расчёта attention в трансформерах, заявляют большие цифры, около 50% ускорения, но на замере всей модели, а не отдельного блока, выходит ~10%. В основном потому, что в целом всё сильно заоптимизировано, и на безрыбье и 10% — много.

В чём идея: обычный FlashAttention при обработке последовательности распиливает её на блоки (по KV) и начинает обработку слева направо. Поскольку в Attention есть softmax, а значит и экспонента, то существует риск переполнения, когда степенью становится очень большое число — поэтому применяется классический трюк вычитания максимума (Softmax инвариантен к сдвигу). Но если FlashAttention работает по блокам, то как узнать максимум заранее?

Можно предварительно пройтись по всем токенам и посчитать max(), но это медленно. Ещё с 2018-го года известен трюк онлайн-подсчёта, который использовал FlashAttention 1. Его смысл в том, что нам не нужно знать максимум сразу — мы помним текущий максимум и для каждого нового блока сравниваем его и локальный максимум блока. Если они отличаются — нормализацию надо пересчитать, в том числе задним числом для предыдущих блоков. Таким образом обработав последний блок в последовательности все предыдущие блоки уже скорректированы, и расчёт Attention становится точным.

FlashAttention 4 предложил новую идею: делать коррекцию предыдущих блоков не каждый раз, когда встречается новый максимум, а только тогда, когда есть угроза потери точности вычислений из-за переполнения. Но при этом все вычисления всё равно так или иначе производятся по честному, никакие расчёты не выкидываются, и Attention получается математически точным.

BLASST предлагает идти дальше: если в блоке максимум сильно меньше, чем посчитанный по предыдущим блокам, то его полностью выбрасывают. Value-векторы токенов этого блока вообще не участвуют в вычислениях. Так можно сделать потому, что если максимум (то есть самые «важные» токены) маленький, то получается, что ничего важного в блоке и не было.

Это изображено на первой картинке — есть 6 блоков, первый всегда обрабатывается, а некоторые последующие пропускаются.

Авторы показывают, что можно выкидывать чуть ли не 50-60% блоков (то есть примерно столько же токенов в последовательности) и при этом почти не терять в качестве — до 1% на бенчмарках на длинный контекст. А в некоторых бенчмарках на рассуждения (AIME2024, GPQA) даже наблюдается маленький прирост качества, я бы сказал в рамках погрешности. Авторы объясняют это тем, что выкидывают токены, которые являются шумом и не важны для текущего токена. Логика в этом есть, в целом все sparse attention на это опираются.

Все изменения можно делать без дообучения моделей, просто поменяв кернелы для инференса. Но также показывают, что если модель немного поучить с новыми кернелами, чтобы она привыкла, что некоторые блоки выкидываются, то качество подрастает.

Один из плюсов подхода — он совместим со множеством других оптимизаций, включая DeepSeek MLA (правда статья вышла до v4, поэтому MLA уже не так актуален).

А в серьезные минусы статьи запишу, что как-то поскупились на оценки больших моделей на реально тяжелых бенчмарках с длинным контекстом — работу написали исследователи из Nvidia, уж у кого, а у них мощности точно были. В аппендиксе тестируют Llama 3.1 70b на одном датасете и DeepSeek R1 на трёх, но не long context — и для обеих моделей не пишут про ускорение 👨‍🦳