FlashMLA изнутри: ядро декодирования MLA и FP8-ядро для V3.2

От переводчика: зачем это читать и что повторить

FlashMLA — библиотека ядер внимания DeepSeek, открытая в первый день Open Source Week (февраль 2025). Здесь переведены две заметки из папки docs/. Они учат главному навыку инженера по производительности: сначала посчитать, во что упирается ядро (в память или в вычисления), и только потом оптимизировать.

Перед чтением повторите:

• как устроено многоголовое латентное вниманиемногоголовое латентное внимание (MLA)EN multi-head latent attention (MLA)Предложено DeepSeek в DeepSeek-V2 (2024). и почему при декодировании оно работает как MQA с одной головой ключа размером 576 = 512 + 64 (урок 5.2.2);
• модель rooflineмодель roofline («крыша»)EN roofline modelДостижимая производительность = min(пик FLOP/с, интенсивность × пропускная способность памяти). Точка перегиба (ridge point) = пик / пропускная способность; у H100 около 295 FLOP/байт. и арифметическую интенсивность (урок 7.1.1);
• потоковый мультипроцессорпотоковый мультипроцессор (SM)EN streaming multiprocessor (SM)Основной «цех» GPU NVIDIA: свои CUDA-ядра, тензорные ядра, регистры и shared memory. У RTX 4080 их 76, у H100 и H800 — 132. Блок потоков целиком выполняется на одном SM., варпварпEN warpГруппа из 32 потоков GPU NVIDIA, которые выполняют одну и ту же инструкцию одновременно, каждый над своими данными. Если потоки варпа расходятся по разным веткам if, ветки выполняются по очереди., регистры и общую память (урок 5.7.1);
• online softmax из FlashAttentionFlashAttentionEN FlashAttentionТочный алгоритм внимания (Tri Dao и соавт., 2022), который не хранит матрицу n×n в видеопамяти: Q, K, V обрабатываются блоками в SRAM с online softmax, а в обратном проходе блоки пересчитываются. Память — O(n) вместо O(n²). Версии 2 и 3 улучшают параллелизм и используют возможности Hopper. (урок 5.7.2);
• DeepSeek Sparse AttentionDeepSeek Sparse Attention (DSA)EN DeepSeek Sparse Attention (DSA)Разреженное внимание DeepSeek-V3.2 — дешёвый индексатор выбирает 2048 прошлых токенов для каждого запроса, полное MLA считается только по ним. Позже взято в GLM-5. — разреженное внимание V3.2 с выбором top-k токенов (урок 5.5.2).

Состояние на октябрь 2026. Обе заметки описывают ядра для GPU Hopper (H800). В выпуске FlashMLA от 30.09.2026 поддержку Hopper и старых моделей (V3, V3.2, V4.0) из основной ветки убрали и сменили формат FP8/FP4 KV-кэша; код, о котором идёт речь ниже, остался в коммите ba89a34. На ваших RTX 4080 (поколение Ada) эти ядра не запускаются — материал только для чтения. Третья заметка, о ядре для Huawei Ascend, переведена на отдельной странице.

Часть 1. Новое ядро Flash MLA изнутри

Оригинал: docs/20250422-new-kernel-deep-dive.md, 22 апреля 2025 года.

В предыдущей версии ядра Flash MLA мы добились впечатляющей производительности: 3000 ГБ/с в режиме, где всё упирается в память, и 580 TFLOPS в режиме, где всё упирается в вычисления. Теперь мы продвинулись ещё дальше — до 660 TFLOPS.

В этой заметке мы подробно разбираем новое ядро и объясняем оптимизации и приёмы, которые стоят за ростом производительности. Сначала объясним, почему ядро MLA упирается в вычисления, хотя это ядро внимания для стадии декодирования, затем обсудим высокоуровневое расписание ядра и в конце — технические детали.

Теоретический анализ алгоритма MLA

Ядра GPU делятся на ограниченные вычислениями (compute-bound: их предел — число операций с плавающей точкой в секунду, FLOPS) и ограниченные памятью (memory-bound: их предел — пропускная способность памяти). Чтобы понять, во что упирается ядро, мы считаем отношение числа операций к объёму памяти (FLOPs на байт) и сравниваем его с возможностями GPU.

Пусть число голов q равно $h_q$, число q-токенов в запросе — $s_q$ (равно 1, если MTP и спекулятивное декодирование выключены), число kv-токенов в запросе — $s_k$ ($s_k \gg h_q s_q$), а размерности голов K и V — $d_k$ и $d_v$. Число операций примерно равно $2 (h_q s_q \cdot d_k \cdot s_k + h_q s_q \cdot s_k \cdot d_v) = 2 h_q s_q s_k (d_k+d_v)$, а объём обращений к памяти (в байтах) — $\operatorname{sizeof}(\text{bfloat16}) \times (h_q s_q d_k + s_k d_k + h_q s_q d_v) \approx 2 s_k d_k$. Значит, отношение вычислений к памяти равно $h_q s_q \cdot \frac{d_k+d_v}{d_k} \approx 2 h_q s_q$.

У NVIDIA H800 SXM5 пиковая пропускная способность памяти — 3,35 ТБ/с, пиковая производительность — 990 TFLOPS. Но из-за троттлинга (у нас частота снижается примерно до 1600 МГц) практический пик падает примерно до 865 TFLOPS. Поэтому при $h_q s_q \ge \frac{1}{2} \cdot \frac{865}{3.35} = 128$ ядро ограничено вычислениями, иначе — памятью.

Согласно обзору онлайн-системы инференса DeepSeek, для экземпляров декодирования мы не используем тензорный параллелизм, то есть $h_q$ равно 128, и ядро ограничено вычислениями. Значит, ядро нужно оптимизировать именно для этого режима.

Комментарий переводчика: прочитаем формулы словами

• В MLA при декодировании K и V — это один и тот же латентный вектор: $d_k = 576$ (512 латентных + 64 RoPE), $d_v = 512$. Поэтому $\frac{d_k + d_v}{d_k} = \frac{1088}{576} \approx 1{,}9 \approx 2$.
• Память почти целиком уходит на чтение кэша: $2 s_k d_k$ байт (bfloat16 — 2 байта на число). Запросы и выход крошечные по сравнению с кэшем длиной в тысячи токенов.
• Порог: $\frac{865}{3.35} \approx 258$ FLOPs на байт может «переварить» GPU; делим пополам, потому что интенсивность ядра равна $2 h_q s_q$, и получаем $h_q s_q \approx 129$, авторы округляют до 128.
• Вывод неожиданный: у MHA каждая голова читает свой кэш и интенсивность низкая (memory-bound), а у MLA все 128 голов читают один общий латентный кэш — интенсивность выше в 128 раз, и декодирование становится compute-bound. Это прямое следствие поглощения матриц из урока 5.2.2.

Высокоуровневое устройство нового ядра

Чтобы полностью загрузить вычислительные ресурсы GPU, нужно накладывать операции на CUDA-ядрах на операции на тензорных ядрах, а обращения к памяти — на вычисления, чтобы тензорные ядра всё время были заняты. Для этого пришлось перепроектировать «расписание» ядра.

Статья о FlashAttention-3 вводит «пинг-понговое» расписание (ping-pong scheduling) и конвейеризацию GEMM и softmax внутри группы варпов, чтобы накладывать поблочные матричные умножения на операции CUDA-ядер. Но здесь эти приёмы напрямую не применить из-за ограничений ресурсов. Выходную матрицу (которая масштабируется и накапливается на каждом проходе основного цикла, как в алгоритме FlashAttention) приходится хранить в регистрах — этого требует инструкция WGMMA. Каждая выходная матрица $64 \times 512$ занимает 32 768 32-битных регистров. А в одном SM всего 65 536 32-битных регистров, поэтому в SM помещается только одна выходная матрица. Это исключает вариант с двумя выходными матрицами, которые поочерёдно пользуются CUDA-ядрами и тензорными ядрами. Нужен другой хитрый способ наложить вычисления CUDA-ядер на вычисления тензорных ядер.

(Можете остановиться и подумать — вдруг вы найдёте решение лучше нашего!)

Наше решение добавляет ещё одно математическое преобразование к online softmax и накоплению из FlashAttention. На каждом шаге мы берём два блока KV (назовём их $K_0$, $K_1$, $V_0$ и $V_1$). Выходная матрица занимает 32 768 регистров — слишком много для одной группы варпов, поэтому мы делим её по вертикали на $O_L$ и $O_R$ (каждая $64 \times 256$). Так же делим $V_0$ и $V_1$ на $V_{0L}$, $V_{0R}$, $V_{1L}$ и $V_{1R}$ (каждая $64 \times 256$). Затем выходная матрица вычисляется так:

0. Храним текущий максимум $m$ (начальное значение $-\infty$, общий для двух групп варпов) и выходные матрицы $\vec o_L, \vec o_R$ (начальное значение 0).
1. [0] Вычисляем $\vec p_0 = \vec q K_0^\intercal / \mathit{qk\_scale}$.
2. [1] Вычисляем $\vec p_1 = \vec q K_1^\intercal / \mathit{qk\_scale}$.
3. [0] Вычисляем $mp_0 = \max(\vec p_0)$, $m_{\text{new},0} = \max(m, mp_0)$ и $\mathit{scale}_0 = \exp(m_{\text{new},0} - m)$. Обновляем $m \gets m_{\text{new},0}$.
4. [0] Применяем softmax к $\vec p_0$: $\vec p_0 \gets \exp(\vec p_0 - m_{\text{new},0})$.
5. [0] Обновляем $\vec o_L \gets \vec o_L \cdot \mathit{scale}_0 + \vec p_0 V_{0L}$.
6. [1] Вычисляем $mp_1 = \max(\vec p_1)$, $m_{\text{new},1} = \max(m, mp_1)$ и $\mathit{scale}_1 = \exp(m_{\text{new},1} - m)$. Обновляем $m \gets m_{\text{new},1}$.
7. [1] Применяем softmax к $\vec p_1$: $\vec p_1 \gets \exp(\vec p_1 - m_{\text{new},1})$.
8. [1] Обновляем $\vec o_R \gets \vec o_R \cdot (\mathit{scale}_0 \cdot \mathit{scale}_1) + \vec p_1 V_{1R}$.
9. [0] Обновляем $\vec p_0 \gets \vec p_0 \cdot \mathit{scale}_1$.
10. [1] Обновляем $\vec o_R \gets \vec o_R + \vec p_0 V_{0R}$.
11. [0] Обновляем $\vec o_L \gets \vec o_L \cdot \mathit{scale}_1 + \vec p_1 V_{1L}$.

Примечание: для простоты считаем, что голова q одна, поэтому $\vec q$ и $\vec o$ — векторы. Число в скобках — номер группы варпов, которая выполняет операцию. Считаем, что $\vec o_L$ лежит в регистрах группы варпов 0, а $\vec o_R$ — в регистрах группы варпов 1.

Такое расписание можно считать вариантом «пинг-понга» с одной выходной матрицей — мы называем его «качелями» (seesaw scheduling). Математически оно эквивалентно online softmax из FlashAttention. Оно позволяет накладывать операции CUDA-ядер на операции тензорных ядер, чередуя две группы варпов, а также накладывать обращения к памяти на вычисления: соответствующие инструкции Tensor Memory Accelerator (TMA) можно запускать сразу, как только данные больше не нужны.

Полное расписание показано на рисунке ниже (помните, что в MLA $K$ и $V$ — одно и то же под разными именами).

Комментарий переводчика: что на рисунке расписания

Рисунок docs/assets/MLA Kernel Sched.drawio.svg в репозитории показывает две дорожки — группу варпов 0 и группу варпов 1 — и для каждой операции подписывает, где она выполняется: Tensor (тензорные ядра) или CUDA. Группа 0 считает rP0 = sQ @ sK0, группа 1 — rP1 = sQ @ sK1; дальше каждая на CUDA-ядрах считает свой максимум и масштаб, пока другая занята умножением на тензорных ядрах, и затем накапливает свою половину выхода (rO0 += … @ sV0L, rO1 += … @ sV1R, потом перекрёстные rO0 += sP1 @ sV1L и rO1 += sP0 @ sV0R). Промежуточные вероятности обмениваются через общую память (sP0, sP1). Между операциями вставлены запуски TMA на загрузку следующих блоков V. Приставка s — данные в общей памяти, r — в регистрах.

Почему это работает. Обычный online softmax держит одну выходную матрицу и перемасштабирует её целиком. Здесь выход разрезан пополам между двумя группами варпов, и каждая половина получает вклад обоих блоков KV — но в разное время. Пока группа 0 делает экспоненты (CUDA-ядра), группа 1 умножает матрицы (тензорные ядра), и наоборот: «качели». Поправочные множители $\mathit{scale}_0$, $\mathit{scale}_1$ гарантируют, что после шага 11 обе половины равны тому, что дал бы обычный FlashAttention.

Технические детали

Этот раздел посвящён техническим деталям нового ядра.

Во-первых, хотя ядро рассчитано на режим, ограниченный вычислениями (где пропускная способность памяти — не узкое место), задержку памяти игнорировать нельзя. Если данные не готовы к моменту, когда они нужны, приходится ждать. Чтобы решить эту проблему, мы используем такие приёмы:

• Мелкозернистая конвейеризация копирования TMA и GEMM. Для блока K размером $64 \times 576$ мы запускаем 9 копирований TMA (каждое переносит блок $64 \times 64$). Операции GEMM начинаются, как только завершается очередное копирование (когда готово первое копирование TMA, можно начинать первую операцию GEMM, и так далее). Это повышает устойчивость к задержкам памяти.
• Подсказки кэшу. Подсказка cute::TMA::CacheHintSm90::EVICT_FIRST для копирований TMA, как показали эксперименты, повышает долю попаданий в кэш L2.

Эти оптимизации дают до 80% загрузки тензорных ядер (от теоретического пика с учётом троттлинга) и 3 ТБ/с пропускной способности памяти на H800 SXM5. В режиме, ограниченном памятью, новое ядро немного (примерно на 2%) медленнее старой версии с пинг-понговыми буферами, но это приемлемо.

Другие улучшения производительности:

• Программный зависимый запуск (Programmatic Dependent Launch). Мы используем programmatic dependent launch, чтобы наложить друг на друга ядра splitkv_mla и combine.
• Планировщик тайлов (Tile Scheduler). Мы реализовали планировщик тайлов, который распределяет задания (запросы и блоки) по SM. Это обеспечивает равномерную загрузку SM.

Благодарности

Алгоритм и расписание FlashMLA вдохновлены FlashAttention, Flash-Decoding и CUTLASS, а также многими проектами, на которых они основаны. Благодарим авторов за их замечательную работу.

Как цитировать

@misc{flashmla2025,
      title={FlashMLA: Efficient MLA decoding kernels},
      author={Jiashi Li and Shengyu Liu and Yuanhang Sun},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}

Часть 2. FP8-ядро декодирования Flash MLA для Hopper

Оригинал: docs/20250929-hopper-fp8-sparse-deep-dive.md, 29 сентября 2025 года.

С выпуском DeepSeek-V3.2 мы удвоили длину контекста наших моделей — с 64 тыс. до 128 тыс. токенов. Это сильно нагружает память GPU (одному запросу длиной 128 тыс. токенов нужен KV-кэш размером $576 \times 2 \times 62 \times 128 \times 1024 = 8{,}72\ \mathrm{ГиБ}$), что может приводить к ошибкам нехватки памяти (OOM) или к недогрузке GPU из-за маленьких батчей. Чтобы решить эту проблему, мы ввели для DeepSeek-V3.2 KV-кэш в FP8.

Однако написать высокопроизводительное ядро декодирования трудно: нужно деквантовать данные, а доступ к памяти разреженный. В этой заметке мы рассказываем историю нашего нового разреженного FP8-ядра декодирования для GPU Hopper. Сначала объясним формат FP8 KV-кэша, затем проведём теоретический анализ тактов и в конце подробно опишем приёмы, использованные в новом ядре.

Комментарий переводчика: откуда 8,72 ГиБ

576 чисел на токен (латентный вектор MLA плюс RoPE-ключ) × 2 байта (BF16) × 62 слоя (61 слой модели и слой MTP) × 131 072 токена = 9 361 686 528 байт = 8,72 ГиБ. Это на один запрос, а батч на декодировании — десятки запросов.

Формат FP8 KV-кэша

Вспомним, что фаза декодирования в многоголовом латентном внимании (MLA) работает аналогично Multi-Query Attention (MQA) со 128 головами запросов и 1 головой ключа, где head_dim_k = 576 и head_dim_v = 512. Чтобы уменьшить KV-кэш и сохранить точность, мы используем мелкозернистое квантование. А именно: к первым 512 элементам KV-кэша каждого токена применяем квантование по тайлам (размер тайла $1 \times 128$). Получается 512 значений float8_e4m3 и 4 масштабных множителя float32. Оставшиеся 64 элемента (часть с RoPE) не квантуем, потому что они чувствительны к потере точности. Поэтому в памяти GPU KV-кэш каждого токена занимает 656 байт: 512 значений float8_e4m3, 4 значения float32 и 64 значения bfloat16.

Внутри ядра мы сначала деквантуем 512 значений float8_e4m3 в 512 значений bfloat16. Затем склеиваем их с 64 исходными значениями bfloat16 из RoPE-части. Наконец, выполняем вычисление MQA с помощью операций умножения-сложения матриц (MMA) в точности bfloat16 (то есть входы MMA в bfloat16, выходы в float32; это верно и для GEMM QK, и для GEMM «оценки внимания × V»).

Комментарий переводчика

512 + 4 × 4 + 64 × 2 = 656 байт против 1152 байт в BF16 — экономия около 43%. Тот же приём «мелкие блоки со своим масштабом» вы видели в FP8-обучении V3 и в групповом квантовании из урока 7.2.1.

Теоретический анализ тактов

Главная трудность в том, что тензорные ядра (они выполняют MMA) очень быстрые, а деквантование, которое выполняется на CUDA-ядрах, за ними не успевает.

Базовая единица GPU NVIDIA — потоковый мультипроцессор (Stream Multiprocessor, SM). Каждый SM можно представлять как независимое ядро GPU. Для простоты рассмотрим один SM. Каждый SM выполняет 4096 операций MMA за такт (на H800 это 989 TFlops / 1830 MHz / 132 SMs). В нашем ядре каждый CTA (блок потоков CUDA) работает на одном SM, и каждому SM сопоставлен только один CTA. Если каждый CTA обрабатывает 64 головы запросов, то на операции MMA для одного K/V-токена нужно всего $64 \times (576+512) \times 2 / 4096 \approx 34$ такта.

Однако H800 не умеет напрямую переводить float8_e4m3 в bfloat16, поэтому деквантование KV-кэша одного токена требует таких шагов:

1. Перевести float8_e4m3 в half.
2. Перевести half в float32.
3. Перевести float32 в bfloat16.
4. Умножить полученное значение bfloat16 на масштабный множитель float32.

Согласно документации NVIDIA, на деквантование каждого токена нужно не меньше $(\frac{1}{64} + \frac{1}{64} + \frac{1}{16} + \frac{1}{256}) \times 512 \approx 50$ тактов! Это заметно больше 34 тактов на операции MMA, то есть ядро ограничено деквантованием. Если ничего не делать, деквантование станет узким местом, и мощные тензорные ядра будут простаивать.

Комментарий переводчика: как читать дроби

Каждая дробь — сколько тактов SM тратит на одну операцию данного типа: $\frac{1}{64}$ такта означает, что SM выполняет 64 таких преобразования за такт. Складываем стоимость четырёх шагов для одного числа и умножаем на 512 чисел токена: 8 + 8 + 32 + 2 = 50 тактов. Выходит, что «подготовить» данные дороже, чем на них посчитать, — типичная ловушка низкой точности.

Кроссовер

Прежде чем продолжить, важно отметить ключевой факт: все головы запроса одного и того же q-токена обращаются к одним и тем же головам ключа, потому что это Multi-Query Attention (MQA).

Вспомним, что каждый CTA обрабатывает 64 головы запросов, а всего в DeepSeek-V3.2 128 голов запросов. Если найти способ «поделиться» деквантованными K/V между двумя CTA, которые обрабатывают разные наборы голов запросов, то каждому CTA придётся деквантовать лишь половину KV-кэша — это было бы отлично! Мы называем этот метод «кроссовер» (crossover), потому что идею подсказал кроссинговер хромосом при мейозе.

Следующий вопрос — как реализовать это в CUDA? До архитектуры Hopper обмениваться данными между CTA можно было только через глобальную память или кэш L2, а это медленно. Но новое решение дала мощная распределённая общая память.

На помощь приходит распределённая общая память

Распределённая общая память (Distributed Shared Memory, DSM) — новая возможность архитектуры Hopper, появившаяся вместе с кластерами CTA (thread block cluster). CTA одного кластера могут напрямую обращаться к общей памяти друг друга. Подробности — в статье NVIDIA Hopper Architecture In-Depth.

Вот как мы её используем. Мы запускаем CTA кластерами по 2. Каждый CTA в кластере отвечает за 64 головы запросов одного и того же q-токена. Каждый CTA выполняет такие шаги:

1. Загружает половину квантованных K/V из глобальной памяти. Для производительности используем широкую загрузку __ldg шириной 128 бит.
2. Деквантует свою половину на CUDA-ядрах.
3. Сохраняет деквантованные K/V в свою общую память.
4. Одновременно с помощью st.async записывает деквантованные K/V в общую память другого CTA кластера.

Для синхронизации этих операций мы используем транзакционный барьер кластера (cluster transaction barrier) — ещё один мощный примитив программирования, доступный в кластерах CTA. После обмена у каждого CTA в собственной общей памяти лежит полный набор деквантованных K и V, и на нём можно выполнять операции MMA.

Производительность

С этими приёмами мы получили 410 TFLOPS в конфигурации, ограниченной вычислениями (batch_size=128, num_heads=128, s_q=2, topk=2048), на GPU H800 SXM5. Это заметный прогресс по сравнению с 250 TFLOPS у нашего предыдущего разреженного FP8-ядра декодирования без кроссовера.

Это всё ещё меньше пиковых 640 TFLOPS нашего прежнего плотного ядра декодирования в bfloat16. Одна из причин в том, что это разреженное ядро, и его topk равен всего 2048. При меньшем topk относительные накладные расходы на пролог и эпилог ядра больше, чем у плотного декодирования с длинным контекстом. Если задать topk больше, например 32 768, ядро достигает 460 TFLOPS.

С другой стороны, время работы этого ядра в упомянутой конфигурации сопоставимо со временем плотного ядра декодирования при длине последовательности около 3000. Когда длина последовательности превышает 3000, преимущество нового ядра становится ещё заметнее. Это также показывает эффективность нашего алгоритма DeepSeek Sparse Attention.

От переводчика: что из этого перенести на свои задачи

1. Прежде чем оптимизировать, посчитайте такты или байты на единицу работы — как здесь 34 такта MMA против 50 тактов деквантования. Тот же расчёт для своей задачи можно сделать в уроке 7.1.1.
2. Если узкое место — подготовка данных, ищите, кто ещё может ей воспользоваться: здесь два CTA делят работу деквантования, потому что в MQA все головы смотрят в один и тот же кэш.
3. Разреженное внимание выгодно не всегда: при topk = 2048 ядро сравнивается с плотным только на контекстах длиннее ~3000 токенов. На коротких промптах DSA ничего не экономит, поэтому в V3.2 для коротких prefill предусмотрен отдельный «замаскированный» плотный режим (см. конспект отчёта V3.2).

Словарик

EN	RU
compute-bound / memory-bound	ограничено вычислениями / ограничено памятью
throttling	троттлинг (снижение частоты)
warpgroup	группа варпов (4 варпа, 128 потоков)
WGMMA	асинхронное умножение матриц на уровне группы варпов (Hopper)
TMA, Tensor Memory Accelerator	ускоритель копирования тензоров между глобальной и общей памятью
seesaw scheduling	«качельное» расписание
tile scheduler	планировщик тайлов
dequantization	деквантование
CTA, thread block cluster	блок потоков, кластер блоков потоков
distributed shared memory (DSM)	распределённая общая память
prologue / epilogue	пролог / эпилог ядра (подготовка и запись результата)