разборllmтрансформерыgpuоптимизация

FlashAttention: как ускорить внимание, ничего не приближая

· 2 мин · BadRabbit

Механизм внимания в трансформерах квадратичен по длине последовательности. Годами это пытались лечить приближёнными версиями внимания, которые на бумаге быстрее, но на практике часто не дают выигрыша по реальному времени. Три Дао с соавторами зашли с другой стороны.

Где на самом деле теряется время

У GPU два уровня памяти:

  • HBM: большая (гигабайты), но относительно медленная;
  • SRAM на кристалле: крошечная, но на порядок быстрее.

Обычное внимание считает матрицу QKᵀ размером N×N, записывает её в HBM, читает обратно для softmax, записывает результат, читает для умножения на V. Сама арифметика быстрая. Медленно гонять огромную матрицу туда-обратно по памяти.

Ключевая идея статьи — IO-awareness: алгоритм должен считать не только операции, но и обращения к памяти.

Как работает FlashAttention

  1. Тайлинг. Q, K и V режутся на блоки, которые помещаются в SRAM. Внимание считается по блокам, и полная матрица N×N никогда не записывается в HBM.
  2. Онлайн-softmax. Softmax вроде бы требует видеть всю строку сразу, но его можно считать по частям, поддерживая текущий максимум и сумму и пересчитывая частичный результат на лету.
  3. Пересчёт вместо хранения. Для обратного прохода промежуточная матрица не сохраняется, а пересчитывается из Q, K, V. Лишние вычисления дешевле лишних чтений.

Результат — точное внимание, без приближений, с памятью, линейной по длине.

Цифры

  • на 15% быстрее обучение BERT-large (последовательности 512) по реальному времени;
  • в 3 раза быстрее GPT-2 на контексте 1K;
  • в 2,4 раза быстрее на Long Range Arena (1K–4K);
  • блочно-разреженный вариант позволил впервые решить Path-X (16K) с точностью 61,4% и Path-256 (64K) с 63,1%.

Мой комментарий

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

На практике это даёт длинный контекст на маленькой видеокарте. Без FlashAttention-подобных ядер контекст в 16 тысяч токенов на 8 ГБ не поместился бы даже теоретически.