# Анатомия GPU и FlashAttention: как оптимизация памяти ускоряет LLM

Источник: https://www.youtube.com/watch?v=6OBtO9niT00
Канал: Stanford Online
Опубликовано: 01.05.2025

---

Эффективность современных больших языковых моделей напрямую зависит от аппаратного обеспечения, на котором они обучаются и запускаются. В рамках лекции курса CS336 Стэнфордского университета подробно разбираются внутреннее устройство графических процессоров (GPU) и ключевые методики оптимизации вычислений. Понимание этих низкоуровневых механизмов позволяет преодолеть ограничения пропускной способности памяти и создавать прорывные алгоритмы, такие как FlashAttention.

## 🛠️ Демистификация CUDA и цели лекции
[[JUMP:0:05]]

Начало лекции посвящено организационным моментам, связанным со сдачей первого домашнего задания и скорым выходом второго, в котором студентам предстоит работать с компилятором Triton и самостоятельно реализовать элементы алгоритма FlashAttention 2. Преподаватель ставит перед аудиторией ключевую цель — сделать работу графических процессоров (GPU) и архитектуры CUDA менее «магической» и более предсказуемой для разработчиков моделей.

Главные учебные задачи лекции включают в себя:

* Освоение принципов работы GPU как одиночного ускорителя вычислений.
* Понимание причин, по которым графические чипы внезапно замедляют работу при определенных конфигурациях данных.
* Изучение примитивов и компонентов, необходимых для ускорения собственных нейросетевых архитектур.

Лектор выражает признательность за материалы блогу Хоруса Хита (Horus He), сообществу CUDA mode и книге Google о тензорных процессорах (TPU). Программа занятия разделена на три блока: глубокий анализ аппаратного обеспечения одного GPU и краткое сравнение с TPU, разбор факторов производительности матричного умножения и детальный разбор механики FlashAttention.

## 📈 Эволюция вычислений: от одиночных потоков к параллелизму
[[JUMP:3:39]]

В современной индустрии обработки естественного языка (NLP) законы масштабирования (scaling laws) стали общепринятым стандартом. Рост вычислительных мощностей напрямую обеспечивает повышение качества работы моделей: увеличение объема обучающих данных и масштаба самих сетей стабильно улучшает их точность. По словам лектора, ключевым драйвером прогресса глубокого обучения стало именно развитие аппаратного обеспечения, эффективная утилизация чипов и параллелизация вычислений.

Исторический контекст кремниевого масштабирования демонстрирует фундаментальный сдвиг в проектировании микросхем:

1.  Эра масштабирования Деннарда (Dennard scaling) и закона Мура позволяла удваивать количество транзисторов на чипе каждые два года, увеличивая тактовую частоту и снижая энергопотребление.
2.  В период с 1980-х по 2000-е годы масштабирование Деннарда фактически исчерпало себя, что привело к стагнации однопоточной производительности процессоров (согласно известным графикам Хеннесси и Паттерсона).
3.  Плотность транзисторов продолжала расти, но это больше не приводило к автоматическому ускорению выполнения одиночных программных потоков.

В результате индустрия перешла от абсолютного ускорения вычислений к параллельному масштабированию, когда множество задач обрабатывается одновременно. Лектор демонстрирует знаменитый график Билла Далли (Bill Dally), иллюстрирующий суперэкспоненциальный рост числа целочисленных операций в секунду от ранних архитектур NVIDIA K20 до современных чипов H100. Понимание этой кривой необходимо для максимального извлечения эффективности при обучении языковых моделей.

## ⚖️ Сравнение архитектур CPU и GPU: задержка против пропускной способности
[[JUMP:6:15]]

Традиционная модель выполнения программ на центральном процессоре (CPU) ориентирована на последовательное пошаговое исполнение инструкций в рамках одиночных потоков. Чтобы поддерживать такую логику, CPU требует огромных блоков управления (control units), систем предсказания переходов и развитой условной логики. В силу этого архитектура CPU жертвует количеством вычислительных ядер ради минимизации задержки выполнения конкретной задачи.

В отличие от CPU, графический процессор (GPU) устроен принципиально иначе:

* Большая часть площади кристалла выделена под массивные вычислительные блоки (ALU), обозначенные на схемах зелеными квадратами.
* На логику управления (control) и кэш-память отводится значительно меньшая доля чипа.
* Вместо оптимизации времени выполнения отдельной задачи (latency), GPU создавался для максимизации совокупного объема обрабатываемых данных в единицу времени (throughput).

Лектор приводит концептуальную аналогию с обработкой пакета задач от T1 до T4. CPU стремится завершить задачу T1 как можно быстрее. GPU же запускает огромное количество легковесных потоков, которые могут мгновенно «засыпать» и «просыпаться», ожидая данные. В итоге вся пачка задач на GPU завершается быстрее, хотя латентность выполнения каждой индивидуальной операции оказывается выше, чем на CPU.

## 🧠 Анатомия видеокарты: вычислительные блоки и иерархия памяти
[[JUMP:8:23]]

С точки зрения внутренней структуры, GPU состоит из множества потоковых мультипроцессоров (Streaming Multiprocessors, SM). Мультипроцессор SM является неделимым атомарным блоком управления при программировании на таких языках, как Triton. Каждый SM содержит внутри себя множество потоковых процессоров (Streaming Processors, SP), которые непосредственно выполняют вычисления, применяя одну и ту же инструкцию к разным элементам данных. Архитектура NVIDIA A100 содержит 128 блоков SM, что кардинально превосходит количество ядер в стандартных CPU.

Однако вычисления — это лишь половина уравнения. Пропускная способность памяти играет решающую, если не главную, роль в итоговой производительности системы. Физическое расположение памяти на подложке чипа жестко диктует скорость доступа к ней:

* **Память на уровне SM**: регистры и кэш L1 (совмещенный с разделяемой памятью / shared memory) находятся непосредственно внутри мультипроцессора и требуют всего около 20 тактов для доступа.
* **Кэш-память L2**: расположена на самом чипе рядом с блоками SM, обеспечивая высокую скорость, но доступ к ней в среднем в 10 раз медленнее, чем к L1.
* **Глобальная память (DRAM/HBM)**: вынесена за пределы вычислительного кристалла и соединяется с ним через специальные контроллеры. Время доступа составляет от 200 до 300 тактов.

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

## 🔄 Логические модели выполнения задач и памяти
[[JUMP:13:15]]

Для эффективного написания CUDA-ядер разработчику необходимо понимать трехуровневую иерархию параллелизма: блоки (blocks), варпы (warps) и потоки (threads). Блоки представляют собой крупные группы потоков, каждая из которых целиком назначается на определенный мультипроцессор SM. Внутри блока потоки группируются в варпы — связки из 32 последовательно пронумерованных потоков, выполняющих одну инструкцию одновременно над разными данными.

Логическая модель памяти GPU строго соответствует этой иерархии выполнения задач:

* Каждый индивидуальный поток обладает собственными сверхбыстрыми регистрами и локальной памятью.
* Потоки в рамках одного блока имеют совместный доступ к разделяемой памятью (shared memory), что позволяет им эффективно обмениваться данными.
* Если данным необходимо пересечь границы между разными блоками потоков, их требуется принудительно записать в медленную глобальную память.

Идеальный паттерн программирования подразумевает загрузку небольшого объема данных в быструю разделяемую память, выполнение над ними интенсивных вычислений всеми потоками блока и завершение работы без промежуточных обращений к глобальной памяти. Произвольный доступ к адресам по всему пространству DRAM сводит производительность к нулю.

## 🔲 Альтернативный подход: архитектура тензорных процессоров Google TPU
[[JUMP:16:38]]

Опираясь на недавние публикации и информацию от инженеров Google, лектор проводит параллель между GPU и тензорными процессорами (TPU), отмечая их концептуальное сходство при кардинальных различиях в архитектурных деталях. В структуре TPU аналогом мультипроцессора SM выступает тензорное ядро (Tensor Core). Оно включает в себя скалярный блок для управления и произвольной логики, векторный блок для поэлементных операций и массивный специализированный блок матричного умножения (Matrix Multiply Unit, MXU).

TPU использует аналогичную двухъярусную схему памяти:

* Внутри тензорного ядра размещается сверхбыстрая векторная память и SMEM.
* За пределами кристалла находится глобальная высокополосная память (HBM).

Главное отличие TPU заключается в максимальном упрощении логики ради одной задачи. В отличие от GPU, архитектура TPU не содержит усложнений вроде концепции варпов; она оптимизирована под пакетное перемножение матриц, выполняемое над тензорами. Подобная узкая специализация делает архитектуру чипа проще, сохраняя при этом фундаментальный принцип разделения на быструю внутреннюю и медленную внешнюю память.

## 🧱 Проклятие памяти и благословенные матричные операции
[[JUMP:19:13]]

Успех графических процессоров во многом предопределен легкостью их масштабирования: для наращивания мощности достаточно добавить больше блоков SM, не сталкиваясь с проблемами критического тепловыделения при разгоне тактовой частоты. Программирование под CUDA, несмотря на кажущуюся сложность, логично структурировано благодаря модели SIMT (одна инструкция — много потоков). Легковесность потоков позволяет аппаратно планировать их запуск и остановку, маскируя задержки памяти и удерживая высокий уровень утилизации SM.

Исторически GPU создавались для обработки графики, однако ученые быстро научились адаптировать текстурные буферы для научных расчетов и быстрого перемножения матриц. В современную эпоху глубокого обучения производители чипов осознали, что матричные вычисления составляют львиную долю нагрузок нейросетей, и сделали их привилегированными операциями. Появление специализированных тензорных ядер в архитектуре NVIDIA V100 создало колоссальный разрыв в производительности между матричными и обычными вычислительными операциями. Лектор подчеркивает: любая новая нейросетевая архитектура обязана опираться на матричные умножения, иначе она окажется неконкурентоспособной на современном оборудовании.

Анализ долгосрочных трендов масштабирования различных компонентов тренировочного стека выявляет драматический дисбаланс:

* Пропускная способность подключения GPU к хост-серверу (PCIe, NVLink) растет крайне медленно.
* Скорость глобальной памяти (переход от GDDR к HBM) увеличивается в логарифмическом масштабе, но все равно отстает от вычислительных ядер.
* Пиковая вычислительная мощность (flops) выросла лавинообразно — в диапазоне от 1 до 100 000 раз относительно базовых показателей.

Если на заре глубокого обучения бутылочным горлышком был дефицит вычислительных операций, то в эпоху чипов H100 главным ограничителем стала память. По этой причине разработка современных эффективных алгоритмов требует смещения фокуса с экономии операций на оптимизацию работы с памятью.

## 🧩 Загадка волнового графика утилизации
[[JUMP:24:51]]

Для демонстрации реального поведения железа лектор предлагает разобрать экспериментальный график утилизации GPU при перемножении квадратных матриц. На оси X отложен размер матриц, на оси Y — количество операций в секунду (эффективность использования оборудования). С ростом матриц утилизация в среднем растет, поскольку накладные расходы на запуск CUDA-ядер нивелируются огромным объемом работы. Однако график покрыт странными, непредсказуемыми волнами и резкими провалами производительности.

Для объяснения этого феномена используется классическая модель Roofline. Она разделяет работу ускорителя на два режима:

* **Режим ограничения памятью (Memory-bound)**: левая диагональная часть графика, где производительность жестко лимитирована скоростью доставки байт данных на единицу вычислительной операции (арифметической интенсивностью).
* **Режим ограничения вычислениями (Compute-bound)**: правое плато, где вычислительные блоки загружены полностью, и память успевает снабжать их данными.

Наша цель — находиться в зоне Compute-bound. Но достижение стабильного максимума требует применения целого арсенала инженерных подходов, защищающих алгоритм от специфических аппаратных ловушек.

## 🛠️ Шесть главных инструментов оптимизации вычислений на GPU
[[JUMP:28:07]]

### ### ❌ 1. Устранение условных ветвлений (Conditionals)
Модель выполнения SIMT обязывает все потоки в рамках одного варпа выполнять абсолютно одинаковую инструкцию. Если в коде ядра появляется условный оператор `if-else` (например, разделение потоков по индексу меньше или больше четырех), GPU физически не может выполнить обе ветки параллельно. Чип вынужден поставить на паузу потоки, попадающие в условие `else`, выполнить ветку `if`, а затем поменять их местами. Такое поочередное переключение уничтожает параллелизм, поэтому использование ветвлений внутри критических участков CUDA-кода категорически не рекомендуется.

### ### 📉 2. Снижение точности представления чисел (Lower Precision)
Переход на пониженную разрядность данных — один из самых мощных методов борьбы с дефицитом пропускной способности памяти. Использование меньшего числа бит (переход от FP32 к FP16, BF16 или INT8) пропорционально сокращает физический объем перемещаемых данных по шине из глобальной памяти. Лектор приводит математический пример поэлементной операции ReLU на векторе размера $N$. В формате FP32 для выполнения одной операции сравнения требуется прочитать и записать 8 байт данных (арифметическая интенсивность — 8 байт на один flop). Перевод вычислений в FP16 снижает этот показатель до 4 байт на flop, фактически удваивая доступную пропускную способность шины памяти без изменения аппаратной части. При этом mixed-precision обучение требует осторожности: в матричных умножениях входы подаются в 16-битном формате, но промежуточное накопление сумм (accumulation) обязательно происходит в FP32-регистрах тензорных ядер для обеспечения численной стабильности, после чего результат может быть снова приведен к 16 битам. Функции вроде экспоненты требуют широкого динамического диапазона, поэтому для них выбирают формат BF16.

### ### 🧪 3. Слияние операторов (Operator Fusion)
Обычный паттерн выполнения последовательных операций в PyTorch создает колоссальные накладные расходы. Согласно аналогии Хоруса Хита, GPU работает как фабрика с конвейерной лентой ограниченной ширины. В наивном сценарии вычисление функции $\sin^2(x) + \cos^2(x)$ заставит систему последовательно запускать пять независимых CUDA-ядер. Сначала вектор $x$ копируется из глобальной памяти в вычислительный блок для расчета синуса, результат отправляется обратно в DRAM; затем этот результат снова извлекается для возведения в квадрат и так далее. Бесконечные пересылки промежуточных данных полностью забивают шину. Суть слияния операторов (kernel fusion) состоит в том, чтобы удержать данные внутри регистров вычислительного блока как можно дольше, выполнив всю цепочку преобразований за один проход, и выгрузить в глобальную память только финальный результат. Инструмент `torch.compile` позволяет осуществлять подобное слияние автоматически.

### ### 🔄 4. Повторные вычисления (Recomputation)
Методика рекомпутации (известная также как градиентный чекпоинтинг) предлагает осознанно тратить избыточные вычислительные такты процессора ради экономии пропускной способности памяти. В процессе стандартного обратного распространения ошибки (backpropagation) активации, полученные на прямом проходе, должны сохраняться в глобальной памяти для последующего вычисления градиентов по цепному правилу. На примере цепочки из трех последовательных функций сигмоиды наивный подход требует выполнения 8 операций чтения/записи в DRAM при крайне низкой арифметической интенсивности. При включении рекомпутации промежуточные активации слоев вообще не сохраняются в памяти. На обратном проходе алгоритм заново рассчитывает эти значения прямо внутри локальной памяти SM на лету. В итоге количество обращений к медленной глобальной памяти сокращается до 5 из 8. Подобный размен избыточных flops на разгрузку шины памяти дает колоссальный выигрыш в общей скорости работы.

### ### ⚡ 5. Объединение запросов к памяти (Memory Coalescing)
Глобальная память DRAM на физическом уровне оптимизирована под так называемый «пакетный режим» (burst mode). При запросе конкретного адреса аппаратура всегда возвращает целый непрерывный блок памяти, поскольку самым медленным этапом является перемещение заряда к усилителю чтения. Если потоки внутри одного варпа одновременно обращаются к последовательным адресам памяти, эти запросы объединяются (коалесцируются) в одну транзакцию burst mode, повышая пропускную способность шины в 4 раза. Лектор детально разбирает обход матриц по строкам и столбцам. Если каждый поток варпа сдвигается по строке (меняя столбцы на каждом шаге), то в момент времени $T_1$ поток 1 читает элемент (0,0), поток 2 — элемент (1,0) и так далее. Они запрашивают разные burst-секции, что приводит к катастрофическому падению скорости. Правильный паттерн — движение потоков по столбцам (изменение индексов строк), когда на одном шаге весь варп считывает данные из одной burst-секции, полностью утилизируя аппаратный потенциал шины.

### ### 📦 6. Тайлинг (Tiling)
Тайлинг представляет собой разбиение крупных матриц на небольшие подматрицы (тайлы), размеры которых позволяют целиком уместить их в быстрой разделяемую память (shared memory) мультипроцессора SM. В наивном алгоритме умножения матриц каждый элемент считывается из глобальной памяти $N$ раз. При использовании тайлинга размером $T \times T$ данные один раз загружаются в кэш SM, где потоки выполняют над ними серию локальных перемножений, после чего загружается следующий тайл. С математической точки зрения количество обращений к глобальной памяти сокращается в $T$ раз. Тайлинг значительно снижает нагрузку на DRAM, однако его реализация сопряжена с серьезными трудностями дискретизации. Если размер матрицы не делится нацело на размер тайла, алгоритм вынужден создавать дополнительные пустые тайлы. Потоки в этих краевых блоках будут простаивать, приводя к недоутилизации мультипроцессоров. Попытка обойти это через предварительную выборку (prefetching) упирается в то, что разделяемая память обычно забита до предела, и свободного места для маневра не остается.

## 🌊 Разгадка аномалий: выравнивание burst-секций и квантование волн
[[JUMP:55:34]]

Взаимодействие тайлинга с физической разметкой памяти DRAM способно преподнести неприятные сюрпризы. Если размер строк матрицы или конфигурация тайла не кратны размеру burst-секции оборудования, строки данных начинают пересекать физические границы секторов. Для чтения одной строки тайла GPU приходится выполнять две транзакции вместо одной, что мгновенно удваивает объем обращений к памяти и режет скорость выполнения операции в два раза. Решением этой проблемы является принудительное выравнивание (padding) — искусственное дополнение размеров матриц нулями до красивых чисел. Лектор цитирует знаменитый твит Андрея Карпати (Andrej Karpathy) об оптимизации модели nanoGPT: увеличение размера словаря с 5257 до 5304 (ближайшее число, кратное 64) привело к росту заполняемости (occupancy) чипа и дало мгновенное ускорение всей модели на 25% исключительно за счет выравнивания данных.

Это подводит к полной разгадке волнообразного графика утилизации. Цветовая маркировка линий на графике в зависимости от делимости размера матрицы (параметр $K$) наглядно подтверждает теорию:

* Матрицы с размерами, кратными 32 или 16, демонстрируют идеальную утилизацию и находятся на верхней границе производительности.
* При кратности 8, 2 или 1 производительность падает, так как burst-секции перестают совпадать с шагами вычислений.

Но как объяснить гигантский вертикальный провал производительности при переходе от размера матрицы 1792 к 1794? Причина кроется в феномене квантования волн (wave quantization).

Рассмотрим математику этого процесса:

1.  При использовании оптимального размера тайла 256x128 матрица размером 1792 разбивается ровно на $7 \times 14 = 98$ тайлов (блоков).
2.  Поскольку графический процессор NVIDIA A100 содержит 108 мультипроцессоров SM, все 98 блоков запускаются одновременно в рамках одной «волны» (wave), обеспечивая стопроцентную параллельную загрузку чипа.
3.  Увеличение размерности матрицы всего на 2 единицы (до 1794) из-за неделимости заставляет алгоритм округлить сетку вверх, порождая уже 120 тайлов.
4.  В первую волну GPU выполняет 108 блоков. Оставшиеся 12 блоков переносятся во вторую волну. В этот момент всего 12 мультипроцессоров выполняют полезную работу, а остальные 96 ядер SM полностью простаивают, дожидаясь окончания их работы.

В результате микроскопическое изменение размера матрицы рушит утилизацию оборудования, заставляя его работать вхолостую.

## 🚀 FlashAttention: триумф инженерной мысли в архитектуре Transformer
[[JUMP:1:04:37]]

Изученные инструменты находят свое идеальное практическое воплощение в алгоритме FlashAttention, ставшем основой современных эффективных трансформеров. Стандартный расчет механизма внимания (Attention) в PyTorch включает перемножение матриц запросов ($Q$), ключей ($K$) и значений ($V$), разделенных операцией Softmax. Главная проблема кроется в Softmax: это глобальная операция, требующая суммирования всей строки для вычисления нормировочного знаменателя. Это вынуждает классический алгоритм материализовать в глобальной памяти (HBM) промежуточную матрицу весов внимания огромного размера $N \times N$, что создает квадратичную зависимость от длины контекста по памяти.

Создатели FlashAttention объединили две классические техники — тайлинг и повторные вычисления, чтобы сократить число обращений к HBM до субквадратичного уровня. Для решения проблемы Softmax используется концепция онлайн-вычисления (Online Softmax), предложенная Милаковым и Гимельштейном в 2018 году. Алгоритм позволяет вычислять точный Softmax последовательно, порция за порцией, без наличия всех элементов строки целиком. Путем математических преобразований и введения корректирующих коэффициентов при обновлении текущего максимума строки ($m_j$) и промежуточного знаменателя ($d_j$), Softmax рассчитывается локально внутри тайлов разделяемой памяти.

В итоге прямой проход FlashAttention-2 выглядит следующим образом:

* Матрицы $Q$, $K$ и $V$ нарезаются на тайлы и загружаются в разделяемую память.
* Вычисляются локальные произведения тайлов, и параллельно обновляются бегущие значения максимумов и нормировочных сумм онлайн-софтмакса.
* В конце прохода по строке, когда знаменатель полностью сформирован в shared memory последнего тайла, происходит финальное деление и мгновенное умножение на тайлы матрицы $V$. Матрица весов внимания ни разу не записывается во внешнюю память HBM.

В обратном проходе (backward pass) разработчики применяют трюк с рекомпутацией. Чтобы не хранить гигантские промежуточные значения для расчета градиентов, алгоритм на лету заново пересчитывает элементы Softmax внутри быстрой памяти SM по тайлам. Это требует дополнительных операций перемножения, но полностью избавляет систему от необходимости гонять огромные массивы данных через медленную шину глобальной памяти.

Лектор подводит итог: современный прогресс искусственного интеллекта держится на предельном внимании к деталям аппаратного уровня. Ключ к созданию высокопроизводительных систем будущего — это не бездумное сокращение количества математических операций (flops), а жесткая оптимизация перемещения данных между иерархическими слоями памяти.