FlashAttention и новые ядра: почему память, а не вычисления, тормозят LLM

Слои оптимизации и роль Kernel

В стеке инференса больших языковых моделей (LLM) ключевым слоем является Kernel. Это не абстракция, а конкретная функция, запускаемая на GPU, где тысячи потоков выполняются параллельно. Именно здесь рождаются имена, ставшие стандартом индустрии: FlashAttention, FlashInfer, Triton, CUTLASS. Все они существуют по одной причине: арифметические операции на современных GPU дешевы, а главная проблема — пропускная способность памяти.

Иерархия памяти: разрыв в 1600 раз

Большинство описаний иерархии памяти GPU неверны. На самом деле это не плавная лестница, а структура с одним критическим «обрывом». Разница между кэшем L2 и основной памятью HBM3 составляет 1600 раз.

Самый важный ресурс для оптимизации — Shared Memory (scratchpad на чипе). Он доступен потокам одного блока и работает на порядок быстрее основной памяти. Однако его крайне мало: один Streaming Multiprocessor (SM) имеет 228 KiB, а один блок — всего 48 KiB (если ядро явно не запросит больше). Задача хорошего kernel — не считать быстрее, а эффективно загружать данные в этот «скретчпад» и выжимать из него максимум, прежде чем обращаться к медленной HBM.

Проблема Attention: квадратичная сложность

Классический механизм Attention требует вычисления матрицы оценок размером S × S (где S — длина последовательности). Для каждого токена нужно:

• Вычислить оценки и записать их в память.
• Прочитать обратно для применения Softmax.
• Записать результат.
• Прочитать снова для умножения на значения.

Это четыре прохода через самое медленное звено системы. При длине контекста 64K токенов матрица оценок для одной головы запроса в одном слое занимает 8.6 ГБ. Учитывая, что модель gpt-oss-120b имеет 64 головы запроса и 18 слоев, объем данных колоссален. Даже каузальное маскирование (запрет смотреть в будущее) не уменьшает объем выделяемой памяти, так как наивная реализация все равно заполняет половину матрицы значениями «минус бесконечность».

Решение: FlashAttention и работа без записи

FlashAttention решает проблему, отказываясь от материализации полной матрицы оценок. Алгоритм никогда не записывает матрицу S × S в память. Вместо этого он обрабатывает последовательность тайлами (блоками), пропуская полностью замаскированные участки. Это устраняет лишние операции чтения/записи и обеспечивает значительный прирост скорости на декодерах.

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

Понимание того, что ядро (kernel) — это инструмент управления памятью, а не вычислениями, объясняет, почему индустрия массово переписывает базовые операции. Оптимизация доступа к Shared Memory и минимизация обращений к HBM — единственный путь к эффективному инференсу длинных контекстов.

Источник: Towards AI pub ↗