📖 Как выжать максимум из Decoder Attention на GPU Привет, меня зовут Андрей Шукшов, в Яндекс R&D я пишу YNMT — движок инференса, на котором работают почти все наши LLM. Большую часть времени я пытаюсь понять, почему некоторые вещи работают медленно, а потом делаю их быстрее. Хочу поделиться с вами кейсом из своей практики: Я разобрал по частям, почему Attention медленно работает на GPU в режиме генерации (декодирования), и написал fused kernel, который объединяет все операции, — чтобы выжать из железа максимум. В статье на Хабре я рассказал, как решал эту задачу по шагам: 🟢 Нашёл узкое место генерации — decoder attention. Классическая схема (QKᵀ → softmax → V) использует три GPU-kernel, поэтому промежуточные результаты приходится читать и записывать в DRAM несколько раз 🟢 Учёл особенности декодирования. Новый токен даёт один вектор Q, а K и V заставляют расти кеш. Из-за этого вычисление превращается из GEMM в GEMV и входит в узкое горлышко пропускной способности памяти 🟢 Использовал Online Softmax. Этот алгоритм позволяет считать нормализацию по мере обработки данных и не требует сначала вычислять всю матрицу Attention 🟢 Собрал fused CUDA-kernel. Объединил все шаги attention в одном ядре, чтобы обрабатывать K и V по тайлам в shared memory 🟢 Оценил эффективность. Результат — ускорение примерно в 1,5–1,7 раза и 85–91% утилизации пропускной способности памяти 🔶 В статье на Хабре я рассказываю, как железо работает с кодом LLM на глубинном уровне и какие математические трюки помогут сделать модели быстрее. Подписывайтесь: 💬 @Yandex4Backend 📹 @YandexforBackend