Flash-MSA: открытые ядра для обучения моделей с разреженным вниманием на миллионах токенов

Разработчик под ником rawsh опубликовал то, что называет первой производительной открытой реализацией ядер для ОБУЧЕНИЯ разреженного внимания MiniMax (MSA, MiniMax Sparse Attention), написанных на CuTeDSL для GPU архитектур Hopper и Blackwell (H100, B200). Несколько передовых моделей уже используют разреженное внимание, чтобы ускорить инференс, но код для эффективного ОБУЧЕНИЯ такой схемы внимания раньше никто не публиковал. Автор явно оговаривает: это неофициальная реализация, к MiniMax он отношения не имеет; вся разработка велась на арендованных у Spheron GPU H100 и B200, с опорой на ядра FA4 и MSA-инференс, а также с помощью ИИ-агента Codex.

MSA похож на DeepSeek Sparse Attention (DSA), но с тремя ключевыми отличиями. Во-первых, блочная разреженность: вместо выбора отдельных ключей и значений (KV) «прокси»-механизм внимания отбирает целые блоки по 128 токенов через max-pooling по прокси-оценкам, это даёт удобные свойства для кеширования. Во-вторых, для основного внимания используется GQA (group-query attention, группированное внимание по запросам) вместо MLA (multi-head latent attention). По словам автора, ни одна западная лаборатория, насколько ему известно, не обучает модели с MLA, поэтому доминирующая сейчас в передовых моделях (например, GLM-5.2, DeepSeek v4) формулировка разреженного внимания, заточенная под MLA, для остальных недоступна, переход на GQA решает эту проблему. В-третьих, добавлена групповая специализация прокси-голов: замена MLA на GQA создаёт независимые группы запросов в каждом слое, и теперь каждая прокси-голова может выбирать свой набор KV-блоков, а не единый набор на весь слой, как в DSA, это должно повышать выразительность основного внимания.

В технической части поста автор разбирает устройство ядер: в прямом проходе прокси-внимание вместо накопления выходов ведёт потоковый top-k отбор по блокам (со вставочной сортировкой по регистрам GPU), а локальный блок каждого токена всегда остаётся видимым, как в скользящем окне. Основное внимание, блочно-разреженная версия flash attention, использующая приём из архитектуры MoBA для переформулировки в variable-length flash. Ключевое свойство: поскольку индексы блоков кешируются с прямого прохода, квадратичная по длине контекста стоимость остаётся только у прокси-прохода, а весь остальной проход обучения работает уже с разреженными кешированными блоками, именно это и даёт ускорение при обучении на очень длинных (миллион+ токенов) контекстах.

Отдельно разобран обратный проход: поскольку KL-дивергенция между распределениями прокси- и основного внимания требует доступа к обеим сразу, оба прохода приходится считать слитно. Автор приводит вывод формулы и показывает, что градиент KL-потерь по прокси-оценке равен просто разности между прокси-вероятностью и основной вероятностью внимания, это позволяет считать градиент, не материализуя полную KL-дивергенцию в памяти, что критично для экономии регистров и разделяемой памяти GPU. Отдельно есть "разогревочный" режим (warmup), где основное внимание считается плотно (без блочных индексов), а прокси-часть обучается только в обратном проходе.

Корректность автор проверял, реализовав MSA также на eager PyTorch и сравнив выходы и градиенты обеих реализаций по косинусному сходству в разных конфигурациях (в точности bf16, допуск 0.01). Из ограничений: обратный проход сейчас упирается в нехватку регистров и разделяемой памяти, слитный проход требует 138 регистров на поток и 105 КБ разделяемой памяти на блок (CTA) при лимите GPU в 255 регистров и 228 КБ, из-за чего на потоковый мультипроцессор помещается только один CTA. Теоретическая загрузка (occupancy) обратного прохода Flash-MSA, 12.5% против 18.75% у обычного Flash-Attention.

В разделе "что дальше" автор перечисляет: ускорение маршрутизации через переиспользование прокси-голов между слоями (по примеру IndexShare у GLM), обучение индексатора в низкой точности, и параллелизм по контексту (context parallelism) для реально длинных последовательностей, либо через простое разбиение по прокси-головам (headwise all-gather, уже возможно без изменения ядер), либо через более сложный, но эффективный ring-style параллелизм, для которого автору нужна консультация специалиста. Также в планах, предобучение существующей GQA-модели вроде Qwen3 с использованием MSA для проверки конверсии архитектуры.

Ключевые факты

Выложены первые открытые и производительные ядра именно для ОБУЧЕНИЯ (не инференса) MiniMax Sparse Attention, на CuTeDSL для GPU Hopper и Blackwell (H100, B200)
MSA отличается от DeepSeek Sparse Attention тремя вещами: блочный отбор KV по 128 токенов, GQA вместо MLA для основного внимания, независимый выбор блоков каждой прокси-головой
За счёт кеширования индексов блоков квадратичной по контексту остаётся только прокси-часть прохода, остальное обучение идёт по разреженным кешированным блокам, что и даёт ускорение на контекстах в миллионы токенов
Обратный проход для KL-потери между прокси- и основным вниманием считается слитно; градиент выводится в замкнутом виде (разность вероятностей) без материализации полной KL-дивергенции
Корректность проверена сравнением с eager-реализацией на PyTorch по косинусному сходству; открытые ограничения, низкая загрузка GPU (occupancy 12.5%) из-за нехватки регистров и разделяемой памяти, и пока не реализованный параллелизм по контексту

Почему это важно

Разреженное внимание уже используют несколько передовых моделей для ускорения инференса, но код для эффективного ОБУЧЕНИЯ такой схемы внимания раньше никто не публиковал в открытом виде. Этот пробел и закрывает Flash-MSA: открытые ядра позволяют не просто запускать уже обученные модели с разреженным вниманием быстрее, а обучать их с нуля на очень длинных контекстах (миллион+ токенов), сохраняя почти линейную, а не квадратичную стоимость по длине последовательности.

Кому это важно

В первую очередь, инженерам и исследователям, которые обучают LLM с длинным контекстом на GPU Hopper/Blackwell (H100, B200) и используют архитектуру GQA, а не MLA. Автор прямо указывает, что западные лаборатории (в отличие от, например, GLM или DeepSeek) обычно не обучают модели с MLA, и поэтому доминирующая сейчас формулировка разреженного внимания, заточенная под MLA, для них была недоступна. Замена MLA на GQA в MSA делает технику применимой к куда более широкому кругу существующих моделей и инфраструктур обучения.

Как это применить

Ядра выложены в открытом доступе (ссылки на GitHub, статью MiniMax, на которую он ссылается, и тренер даны в самом посте) и рассчитаны на GPU Hopper и Blackwell. Автор упоминает, что интеграция headwise-параллелизма по контексту возможна уже сейчас без изменения самих ядер, в том числе в форке Megatron с поддержкой MSA. Есть архитектурное ограничение: число прокси-голов должно быть не меньше числа групп GQA, по мнению автора, для стабильного обучения нужно не меньше четырёх MSA-групп, поэтому на практике это ограничение редко становится проблемой.

Можно ли доверять

Автор явно указывает, что это неофициальная и независимая реализация, не связанная с MiniMax; вся разработка велась на арендованных GPU. Корректность ядер проверена сравнением с эталонной реализацией на eager PyTorch по косинусному сходству выходов и градиентов в точности bf16 (допуск 0.01), это стандартный и достаточно строгий способ верификации кастомных GPU-ядер. При этом публикация свежая (обсуждение на Hacker News пока набрало всего один комментарий), независимого стороннего аудита или бенчмарков от третьих лиц ещё не было.

Риски и подводные камни

Технология пока сырая по производительности: обратный проход упирается в нехватку регистров и разделяемой памяти GPU (используется 138 из 255 регистров на поток и 105 из 228 КБ разделяемой памяти на блок), из-за чего теоретическая загрузка GPU составляет только 12.5% против 18.75% у обычного Flash-Attention, то есть запас для ускорения ещё большой, но пока не реализован. Полноценный параллелизм по контексту (context parallelism), необходимый для действительно длинных последовательностей, реализован только частично, сложный и эффективный ring-style вариант автор пока не знает, как встроить в ядра, и откладывает на будущее. Также не реализован альтернативный вариант обучения MSA с совместным градиентом от целевой функции в индексатор, по словам автора, это замедлило бы обучение, а собственные эксперименты MiniMax показывают, что корректный разогрев (warmup) может это компенсировать.

«Насколько мне известно, ни одна западная лаборатория не обучает модели с MLA, из-за чего доминирующая сейчас в передовых моделях (таких как GLM-5.2, DeepSeek v4) формулировка разреженного внимания, заточенная под MLA, для остальных недоступна.»

— rawsh, автор Flash-MSA