CS336, лекции 5–6: GPU, ядра и Triton

Перевод лекций весны 2025 года: как устроена видеокарта и почему она «тормозит» (иерархия памяти, roofline, слияние операций, тайлинг, квантование волн, FlashAttention), затем бенчмаркинг, профилирование и пять способов написать GeLU и softmax — от PyTorch до CUDA и Triton.

Полный переводс языка: английский

Оригинал: Stanford CS336 (Spring 2025) — Lecture 5: GPUs; Lecture 6: Kernels, Triton · слайды (PDF, 51 слайд), исполняемая лекция lecture_06.py и видео на YouTube · весна 2025
Автор: Тацунори Хашимото (Tatsunori Hashimoto), Stanford University
Лицензия оригинала: MIT (© 2025 Stanford University)

Текст оригинала переведён целиком, его лицензия это разрешает.

К урокам курса

Лекция 5 читалась по слайдам — PDF «2025 Lecture 5 - GPUs.pdf» в папке nonexecutable. Номер слайда — в скобках у заголовка; рисунки взяты лектором из статей и блогов, мы их не копируем, а описываем в квадратных скобках. Лекция 6 — исполняемая программа lecture_06.py: её текст переведён по порядку, код приведён с русскими комментариями, а результаты замеров (сделаны авторами на NVIDIA H100 80 ГБ) взяты из трассы выполнения var/traces/lecture_06.json и даны в комментариях # → …. Пояснения, которых нет в оригинале, помечены «Примечание переводчика». Видео — в плейлисте 2025 года, лекции 5 и 6.

Как читать вместе с курсом

Лекция 5 — «почему GPU быстрый и когда он медленный». Она почти один в один ложится на урок 5.7.1 «Как устроена видеокарта и куда уходит время», а последняя треть — на урок 5.7.2 о FlashAttention. Раздел о низкой точности дополняет урок 5.4.1.
Лекция 6 — практика: замер времени, профайлер PyTorch, свои ядра на CUDA и Triton, torch.compile. Читайте вместе с уроком 5.7.3 «Пишем ядра на Triton»: там те же ядра (поэлементное, слитый softmax), но с замерами на вашей RTX 4080.
Эти две лекции — теоретическая база задания 2 CS336 (профилирование и FlashAttention-2 на Triton), см. путеводитель по заданию 2.

Лекция 5. GPU

План и цели (2)

Сделать CUDA и GPU менее «волшебными».
Понять, когда GPU работают медленно. [График из блога Хораса Хе «What shapes do matrix multiplications like?»: достигнутые TFLOPS для квадратных матмулов разного размера — рваная пила с подписями «арифметическая интенсивность», «тайлинг», «квантование волн».]
Понять, как делать быстрые алгоритмы. [Рисунок 1 из статьи FlashAttention (Dao et al.): тайлы Q, K, V копируются в SRAM и обрабатываются во внутреннем и внешнем циклах; справа — время внимания GPT-2 в PyTorch и во FlashAttention.]

Прежде чем начать (3)

Большая часть материала опирается на несколько источников, которые хочется отметить: блог Хораса Хе (Thonk From First Principles) и сообщество CUDA MODE (сейчас GPU MODE), а также https://nichijou.co/ и https://jonathan-hui.medium.com/.

Что сегодня (4)

Часть 1: GPU изнутри — как они работают и из каких важных частей состоят.
Часть 2: как понимать производительность GPU.
Часть 3: собираем всё вместе — разбираем FlashAttention.

Отправная точка: вычисления дают предсказуемый выигрыш (5)

Часто больше вычислений даёт языковым моделям предсказуемый прирост качества. [График из «Neural Scaling Laws» Kaplan et al.: потери на валидации падают как степенная функция от вычислений, $L = 2.57 \cdot C^{-0.048}$ .]

Более быстрое железо, лучшая загрузка и лучшее распараллеливание сами по себе могут двигать прогресс (пока что…).

Откуда берётся рост вычислений? Сначала — масштабирование Деннарда (6)

[График «42 года данных о процессорах»: число транзисторов продолжает расти по закону Мура, а частота и однопоточная производительность с середины 2000-х выходят на плато.]

Но традиционная форма роста (масштабирование Деннарда), работавшая с 1980-х по 2000-е, исчерпалась. …Как же накормить ненасытный аппетит LLM к вычислениям?

Параллельное масштабирование продолжается (7)

[Слайд из доклада Билла Далли на HotChips: производительность инференса на одном чипе выросла в 1000 раз за 10 лет. Вклады: представление чисел (FP32 → FP16 → Int8) — около 16×; сложные инструкции (DP4, HMMA, IMMA) — около 12,5×; техпроцесс (28 → 5 нм) — около 2,5×; разреженность — около 2×.]

Параллельные вычисления на GPU выросли более чем в 1000 раз за 10 лет. Без роста GPU не было бы роста LLM.

Чем GPU отличается от CPU? (8)

CPU оптимизированы под несколько быстрых потоков, GPU — под очень-очень много потоков.

[Схема: у CPU большой блок управления, большой кэш и несколько АЛУ; у GPU — сетка из множества крошечных АЛУ.] Много маленьких вычислительных блоков (АЛУ); гораздо меньше поддержки ветвлений (управления, кэша).
CPU оптимизированы на задержку (каждый поток завершается быстро), GPU — на пропускную способность (общий объём обработанных данных). [Диаграмма: на GPU много потоков, каждый подолгу ждёт данные, но вместе они загружают железо; на CPU один поток выполняется почти без ожиданий.]

Анатомия GPU: исполнительные блоки (9)

[Схема чипа GA100: 128 SM, выстроенных рядами.] У GPU много SM (streaming multiprocessors, потоковые мультипроцессоры), которые независимо выполняют блоки (задания). Каждый SM, в свою очередь, содержит много SP (streaming processors, потоковых процессоров), которые параллельно выполняют «потоки».

Анатомия GPU: память (10)

Чем ближе память к SM, тем она быстрее. L1 и разделяемая память (shared memory) находятся внутри SM. Кэш L2 расположен на кристалле, а глобальная память — это микросхемы памяти рядом с GPU.

[Таблица задержек доступа (в тактах): глобальная память — 290, кэш L2 — 200, кэш L1 — 33, разделяемая память (загрузка/запись) — 23/19. Рядом — фото платы A100 и снимок кристалла.]

SRAM (разделяемая память/кэш) намного дороже (в 100 раз), но примерно в 8 раз быстрее DRAM (глобальной памяти).

Модель исполнения GPU (11)

В модели исполнения три важных участника:

Потоки (threads) «делают работу» параллельно: все потоки выполняют одни и те же инструкции, но над разными входными данными (SIMT — single instruction, multiple threads).
Блоки (blocks) — группы потоков. Каждый блок выполняется на одном SM со своей разделяемой памятью.
Варп (warp): потоки всегда выполняются варпами по 32 потока с последовательными номерами.

Модель памяти GPU (12)

[Схема из документации CUDA: у каждого потока свои регистры и локальная память; у блока — разделяемая память; у всей сетки (grid) — глобальная и константная память, которую хост (CPU) может читать и писать.]

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

Отступление: а как же TPU? (13)

GPU, TPU и многие другие ускорители на высоком уровне похожи. [Абстрактная схема ядра TPU (TensorCore): скалярный блок раздаёт инструкции векторному блоку (VPU) и матричному блоку (MXU); рядом — память HBM.] У GPU больше SM, у TPU меньше TensorCore (но производительность на матмулах похожая).

Общее устройство: лёгкое управление, быстрый (большой) блок умножения матриц, быстрая память.

Различия:

как ускорители соединены в сеть (об этом — в лекции о параллелизме);
нет варпов (только блоки — компромиссы между матмулами и остальными операциями).

Сильные стороны модели GPU (14)

Тяжёлые задачи легко масштабируются (добавлением SM).
Программировать легко (?) благодаря модели SIMT.
Потоки «лёгкие»: их можно останавливать и запускать.

GPU как быстрые умножители матриц (15)

Ранние GPU NVIDIA — это программируемые шейдеры. Исследователи «взламывали» их, чтобы умножать матрицы. [Снимок статьи «Fast Matrix Multiplies using Graphics Hardware» (Larsen, McAllister, 2001).]

Новое железо для матмулов: матмулы быстрые и особенные (16)

Тензорные ядра (появились в сериях V и T) — специализированные схемы умножения матриц. [График TFLOPS по поколениям K80 → M80 → P100 → V100 → A100 → H100: у матмулов кривая уходит вверх гораздо круче, чем у остальных операций.] Матмулы более чем в 10 раз быстрее прочих операций с плавающей точкой!

Вычисления растут быстрее памяти (17)

[График из статьи «AI and Memory Wall»: за 20 лет пиковые FLOPS железа выросли в 60 000 раз (3,0× за 2 года), пропускная способность DRAM — в 100 раз (1,6× за 2 года), межсоединений — в 30 раз (1,4× за 2 года).]

FLOPs растут быстрее памяти — всё труднее «кормить» вычислительные блоки данными!

Итоги части 1: что такое GPU и как он работает (18)

GPU массово параллельны: одни и те же инструкции выполняются множеством исполнителей.
Вычисления (особенно матмулы) росли быстрее памяти.
Чтобы всё работало быстро, нужно уважать иерархию памяти.

Часть 2. Как сделать ML-нагрузки быстрыми на GPU (19)

Производительность GPU может быть сложной даже для такой простой вещи, как умножение квадратных матриц. [Тот же график TFLOPS от размера матрицы: в целом растёт, но с пилообразными провалами.]

Что делает ML-нагрузки быстрыми? (20)

Модель roofline («линия крыши»). [График: по горизонтали — операционная интенсивность (FLOP на байт), по вертикали — производительность. Слева наклонные линии — ограничение пропускной способностью памяти (регистры, разделяемая память, основная память GPU, память CPU), справа горизонтальные «крыши» — пиковая производительность АЛУ GPU и CPU. Плотное умножение матриц лежит справа, разреженное — слева.]

Главный вопрос этой части: как не упереться в память?

Как заставить GPU работать быстро? (21)

Расхождение потоков управления (это не про память…).
Вычисления в низкой точности.
Слияние операторов.
Пересчёт.
Объединение обращений к памяти (coalescing).
Тайлинг.

Расхождение управления — не про память (22)

GPU работают в модели SIMT: все потоки варпа выполняют одну и ту же инструкцию. Условные операторы допустимы, но из-за модели исполнения дают заметные накладные расходы.

if (threadIdx.x < 4) {
    A;
    B;
} else {
    X;
    Y;
}
Z;

[Диаграмма: потоки варпа расходятся; сначала часть потоков выполняет A и B, пока остальные простаивают, затем другие выполняют X и Y, и только потом все вместе — Z. Время растёт.]

Приём 1: вычисления в низкой точности (23)

Меньше битов — меньше битов перемещать. [Тот же слайд Билла Далли: больше всего (около 16×) дала смена представления чисел.]

Низкая точность повышает арифметическую интенсивность (24)

Пример: поэлементный ReLU ( $x = \max(0, x)$ ) над вектором длины $n$ .

Случай float32. Обращения к памяти: 1 чтение ( $x$ ), 1 запись (если $x$ меньше 0); float32 — это 4 байта на каждое, итого 8 байт. Операции: 1 сравнение, 1 FLOP. Интенсивность: 8 байт на FLOP.

Случай float16. Обращения к памяти: те же 1 чтение и 1 запись, но float16 — это 2 байта на каждое, итого 4 байта. Операции: 1 сравнение, 1 FLOP. Интенсивность: 4 байта на FLOP.

Здесь интенсивность записана «наоборот» — байты на FLOP; меньше значит лучше. В уроке 5.7.1 используется привычная форма арифметической интенсивности — FLOP на байт: 1/8 для float32 и 1/4 для float16. Переход на половинную точность вдвое сдвигает операцию вправо на графике roofline.

Низкая точность ускоряет умножение матриц (25)

Многие операции в современных GPU ускорены за счёт низкой/смешанной точности. [Схема тензорного ядра: два 16-битных входа перемножаются с полной точностью, произведения суммируются в аккумуляторе FP32.]

Операции, которые могут использовать 16-битное хранение (FP16/BF16): умножение матриц; большинство поэлементных операций (relu, tanh, add, sub, mul).
Операции, которым нужна большая точность (FP32/FP16): прибавление малых значений к большим суммам (ошибки округления); редукции (сумма, softmax, нормализация).
Операции, которым нужен больший диапазон (FP32/BF16): поэлементные операции, где $|f(x)| \gg |x|$ (exp, log, pow); функции потерь.

Приём 2: слияние операторов (26)

Представьте GPU как фабрику: сырьё приходит со склада (из памяти) и обрабатывается на фабрике. [Рисунок Хораса Хе: склад «Память» и фабрика «Вычисления», соединённые конвейером; фабрику расширили вдвое, а конвейер — нет.] Вычисления растут, память — нет.

Слияние операторов, чтобы меньше ходить в память (27)

А если операций много? Возить туда-обратно довольно глупо. [Слева — «наивный (не слитый)» вариант: после каждой операции результат уезжает в память и возвращается обратно. Справа — «слитое ядро»: данные один раз приезжают на фабрику, проходят все операции и один раз уезжают.]

Пример — синусы и косинусы (28)

Наивное вычисление $\sin^2 x + \cos^2 x$ запускает 5 ядер CUDA (данные ходят туда и обратно). [Граф FX из PyTorch: sin, cos, два возведения в квадрат и сложение — пять отдельных узлов.]

Пример слияния (29)

Все 5 поточечных операций можно слить в один вызов ядра CUDA. [Слева — граф до слияния, справа — один узел после слияния в TorchInductor.] «Простые» слияния вроде этого компиляторы делают автоматически (torch.compile).

Приём 3: пересчёт (30)

[Из курса CS221.] При обратном распространении мы храним активации (жёлтые) и вычисляем якобианы (зелёные). [Пример графа вычислений для $\text{Loss}(x, y, w) = (w \cdot \phi(x) - y)^2$ с прямыми значениями $f_i$ и обратными $g_i = \partial \text{loss} / \partial f_i$ .]

Хранить (и доставать) активации бывает дорого! (31)

Допустим, мы поставили друг на друга три сигмоиды. [Схема: старый прямой проход — 1 чтение из памяти, 3 записи (сохраняем промежуточные s1, s2 и выход); старый обратный проход — 3 чтения, 1 запись.] Для производительности это ужасно: 8 операций чтения/записи и очень низкая арифметическая интенсивность.

Выбросьте активации, пересчитайте их! (32)

[Новый прямой проход — 1 чтение, 1 запись; новый обратный проход — 2 чтения, 1 запись: промежуточные сигмоиды вычисляются заново прямо в ядре обратного прохода.] Выбросить вычисленное может быть даже оптимально: обращений к памяти — 5/8 от исходного!

Приём (?) 4: объединение обращений и DRAM (33)

DRAM (глобальная память) читается в «пакетном режиме» (burst mode) — каждое чтение даёт сразу много байтов!

Всё адресное пространство разбито на пакетные секции.
При обращении к любому адресу все остальные адреса той же секции тоже доставляются процессору.
Простой пример: 16-байтное адресное пространство, пакетные секции по 4 байта. На практике у нас как минимум 4 ГБ адресного пространства и секции по 128 байт или больше.

Пакетный режим возникает из-за медленного построчного копирования в усилитель считывания. [Схема матрицы ячеек DRAM.]

Объединение обращений к памяти (34)

Обращения к памяти объединяются (coalesced), если все потоки варпа попадают в одну и ту же пакетную секцию. [Схема: потоки T0–T3 читают адреса 0–3 — одна секция, один запрос.] Когда все потоки варпа выполняют инструкцию загрузки и все адреса попадают в одну секцию, делается только один запрос к DRAM, и обращение полностью объединено.

Напоминание: варп — это 32 потока с последовательными номерами, которые выполняются вместе в одном блоке. Их обращения к памяти происходят одновременно.

Объединение при умножении матриц (35)

Для матриц, хранящихся построчно (row-major), обращения потоков, которые идут вдоль строк, не объединяются. [Схема (A): поток 1 и поток 2 идут каждый по своей строке матрицы — в каждый момент они читают из разных строк, то есть из разных секций. Схема (B): потоки идут вниз по столбцам — в каждый момент соседние потоки читают соседние элементы одной строки.] Обратите внимание: на второй схеме на каждом шаге читается весь вектор целиком!

Приём 5 (главный): тайлинг (36)

Тайлинг — идея группировать и упорядочивать потоки так, чтобы свести к минимуму обращения к глобальной памяти. Вернёмся к умножению матриц… [Схема: поток $(0,0)$ вычисляет $P_{0,0}$ , читая строку $M$ и столбец $N$ ; поток $(0,1)$ снова читает ту же строку $M$ .] Обращения не объединены и повторяются ( $M_{0,0}$ и $N_{1,0}$ читаются несколькими потоками).

Тайлинг: храним и переиспользуем данные в разделяемой памяти (37)

Режем матрицу на маленькие «плитки» (тайлы) и загружаем их в разделяемую память. Считаем произведение «фазами»:

Загружаем тайлы $M_{0,0}$ и $N_{0,0}$ в разделяемую память.
Считаем частичные суммы для $P$ . (С одним тайлом закончили.)
Загружаем следующие тайлы $M_{0,1}$ и $N_{1,0}$ в разделяемую память.
…

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

Математика тайлинга (38)

[Схема: матрицы A, B, C размера $N$ с тайлами размера $T$ ; внешний цикл — по тайлам, внутренний — по элементам.]

Без тайлинга: каждый входной элемент читается из глобальной памяти $N$ раз.
С тайлингом: каждый входной элемент читается из глобальной памяти $N/T$ раз и $T$ раз внутри тайла. Обращений к глобальной памяти становится в $T$ раз меньше.

Сложности тайлинга (39)

Размер тайла может не делить размер матрицы, и тогда загрузка падает. [Рисунок из руководства NVIDIA: матрица 256×256 покрывается тайлами 128×128 ровно четырьмя тайлами; матрица 257×257 требует уже девяти, причём пять из них почти пустые.]

От чего зависит размер тайла:

объединение обращений к памяти;
объём разделяемой памяти;
делимость размеров матрицы.

Сложности тайлинга 2 — выравнивание памяти (40)

Память приходит пакетами. Тайлы загружаются быстро, если пакеты выровнены с матрицей. [Рисунок Хораса Хе: «выровненная раскладка» — тайл ложится точно на границы секций («один хороший тайл»); «невыровненная» — тот же тайл задевает две секции («два плохих тайла»).]

Объединённый доступ может оказаться невозможным — в зависимости от размерности матрицы (приходится добавлять паддинг).

Собираем всё вместе: загадка матриц (41)

[Твит Андрея Карпатого, февраль 2023: «Самая впечатляющая оптимизация nanoGPT на сегодня (~25% ускорения) — просто увеличить размер словаря с 50257 до 50304 (ближайшее кратное 64). Это добавляет бесполезные размерности, но вычисления идут по другому пути ядра с гораздо большей загрузкой. Осторожнее со степенями двойки».]

Почему бо́льшие матрицы считаются быстрее? (Раздел основан на https://www.thonking.ai/p/what-shapes-do-matrix-multiplications.)

Загадка матриц (42)

Часть картины мы уже понимаем (арифметическая интенсивность, тайлинг). Посмотрим внимательнее… [Снова график TFLOPS от размера матрицы.]

Часть 1: тайлинг (43)

[График TFLOPS для умножения $N \times N$ , раскрашенный по тому, на что делится размер: $K = 2, 8, 16, 32, 128$ . Чем на большую степень двойки делится $N$ , тем выше кривая.] Тайлинг сильно влияет на скорость через выравнивание.

Часть 2: квантование волн (44)

Откуда периодичность? [Увеличенный кусок графика между 1536 и 2048: резкие провалы.]

Это происходит при переходе от размера 1792 к 1793. Почему? При тайле 256 × 128 получается

\frac{1792}{256} \times \frac{1792}{128} = 7 \times 14 = 98

тайлов. Если увеличить размер до 1793, будет $8 \times 15 = 120$ тайлов. У A100 108 SM, поэтому все 120 тайлов она одновременно выполнить не может.

Итоги части 2: как ускорить ML-нагрузки (45)

Уменьшать число обращений к памяти:
- объединение обращений (coalescing);
- слияние операций (fusion).
Переносить данные в разделяемую память:
- тайлинг.
Обменивать память на вычисления/точность:
- квантование;
- пересчёт.

Часть 3. Применим знания, чтобы понять FlashAttention (46)

FlashAttention [Dao et al.] резко ускоряет внимание… Но как? [Слева — время внимания GPT-2: в PyTorch основная часть уходит на маскирование, softmax и dropout, а не на матмулы; во FlashAttention всё слито в одно ядро. Таблица: GFLOPs 66,6 против 75,2 (FlashAttention считает даже больше!), чтение/запись HBM 40,3 ГБ против 4,4 ГБ, время 41,7 мс против 7,3 мс.]

Метод из статьи: «Мы применяем две известные техники (тайлинг и пересчёт), чтобы преодолеть техническую трудность вычисления точного внимания с субквадратичным числом обращений к HBM».

Напоминание: как считается внимание (47)

Вычисление внимания — это 3 умножения матриц (K, Q, V) с softmax посередине. [Схема: $XQ$ умножается на $K^\top X^\top$ — получаем $XQK^\top X^\top \in \mathbb{R}^{3 \times n \times n}$ , «три набора оценок для всех пар»; к ним применяется softmax, результат умножается на $XV$ — получается выход $\in \mathbb{R}^{n \times d}$ .]

Тайлинг, часть 1: тайлинг для умножения KQV (48)

[Рисунок 1 из статьи: пирамида иерархии памяти — SRAM 19 ТБ/с (20 МБ), HBM 1,5 ТБ/с (40 ГБ), DRAM CPU 12,8 ГБ/с (больше 1 ТБ); справа — внешний цикл копирует блоки K и V в SRAM, внутренний — блоки Q, результат пишется в HBM.]

Рисунок 1 из статьи — это буквально тайлинг для умножения KQV… Но что делать с softmax?

Тайлинг, часть 2: инкрементальное вычисление softmax (49)

Из работы Milakov и Gimelshein (2018). Обычный softmax — это

y_i = \frac{e^{x_i - \max_k x_k}}{\sum_{j=1}^{V} e^{x_j - \max_k x_k}}.

[Алгоритм 2 «Безопасный softmax»: три прохода — найти максимум, посчитать сумму экспонент, нормировать. Алгоритм 3 «Безопасный softmax с онлайн-вычислением нормировки»: максимум $m_j$ и сумма $d_j$ обновляются за один проход: $m_j = \max(m_{j-1}, x_j)$ , $d_j = d_{j-1} \cdot e^{m_{j-1} - m_j} + e^{x_j - m_j}$ .]

Чтобы отслеживать максимум, обновляем его инкрементально и строим «телескопическую» сумму. Это позволяет считать softmax тайл за тайлом.

Прочитаем формулу $d_j$ словами: «старую сумму экспонент пересчитываем к новому максимуму (умножаем на $e^{m_{j-1} - m_j}$ ) и добавляем новый член». Если максимум не изменился, множитель равен 1. Пошаговый пример с числами — в уроке 5.7.2 (раздел об online softmax).

Собираем всё: прямой проход FlashAttention (50)

[Схема из статьи FlashAttention-2 (Dao, 2023): $S^{(1)} = Q (K^{(1)})^\top$ и $S^{(2)} = Q (K^{(2)})^\top$ считаются по тайлам в SRAM; $A^{(j)} = \exp(S^{(j)})$ ; выход $O^{(2)}$ получается из $O^{(1)}$ пересчётом знаменателя.] Здесь мы видим:

вычисление скалярных произведений ( $S$ ) по тайлам;
слияние операции взятия экспоненты;
вычисление softmax по тайлам приёмом онлайн-телескопической суммы.

(Обратный проход разбирать не будем — но там тоже пересчитывают тайл за тайлом…)

Итоги всей лекции (51)

Мощность железа растёт, и низкоуровневые детали определяют, что масштабируется, а что нет.
Современные вычисления на GPU сильно подталкивают думать о матмулах и перемещении данных.
Внимательное отношение к GPU (объединение обращений, тайлинг, слияние) приводит к хорошей производительности.

Лекция 6. Ядра и Triton

Объявления

Таблица лидеров задания 1: Leaderboard. Вышло задание 2: A2.

Прошлая лекция: обзор GPU и производительности на высоком уровне. Эта лекция: бенчмаркинг/профилирование и написание ядер.

Эту лекцию стоит запускать на GPU, чтобы получить всё.

Обзор: GPU

Железо

[Схема A100: SM, кэши, память.]

Вычисления: потоковые мультипроцессоры (SM) [у A100 — 108].
Память:
- DRAM [у A100 — 80 ГБ] — большая, медленная;
- кэш L2 [у A100 — 40 МБ];
- кэш L1 [у A100 — 192 КБ на SM] — маленький, быстрый.

Характеристики своей GPU можно посмотреть:

def print_gpu_specs():
    num_devices = torch.cuda.device_count()  # → 8
    for i in range(num_devices):
        properties = torch.cuda.get_device_properties(i)
        # → _CudaDeviceProperties(name='NVIDIA H100 80GB HBM3', major=9, minor=0,
        #      total_memory=81090MB, multi_processor_count=132, L2_cache_size=50MB, …)

Базовая структура: выполнить $f(i)$ для всех $i = 0, \dots, N-1$ .

Модель исполнения

[Рисунок из документации PTX: сетка (grid) из блоков, каждый блок — из потоков.]

Поток (thread): обрабатывает отдельный индекс (то есть $f(i)$ ).
Блок потоков (thread block, он же CTA — cooperative thread array): планируется на один SM.
Сетка (grid): набор блоков потоков.

Зачем нужны блоки потоков? Ради разделяемой памяти.

Интуиция: собрать вместе те $f(i)$ , которые читают похожие данные.
У потоков внутри блока есть разделяемая память (быстрая, как кэш L1) [у A100 — 164 КБ].
Потоки внутри блока можно синхронизировать (для чтения/записи), а между блоками — нельзя.

Железо и исполнение взаимодействуют

[Рисунок из блога NVIDIA: блоки потоков распределяются по SM волнами.]

Блоки потоков планируются на SM волнами. Проблема: в последней волне блоков меньше, часть SM простаивает (низкая загрузка, occupancy). Квантование волн: делайте число блоков потоков кратным числу SM. Эмпирическое правило: блоков должно быть не меньше чем 4 × (число SM). Трудность: некоторые стороны железа скрыты от модели исполнения (например, планирование, число SM).

Арифметическая интенсивность: число FLOP / число байтов

Если высокая — операция ограничена вычислениями (compute-bound, хорошо).
Если низкая — операция ограничена памятью (memory-bound, плохо).

Общее правило: умножение матриц ограничено вычислениями, всё остальное — памятью.

Бенчмаркинг и профилирование

ВАЖНО: делайте бенчмарки и профилируйте свой код!

Можно читать спецификации (маркетинговые материалы) и статьи… но производительность зависит от версии библиотек, от вашего железа, от вашей нагрузки… так что бенчмаркинг и профилирование своего кода ничем не заменить.

Пример вычисления: прямой и обратный проходы MLP.

class MLP(nn.Module):
    """Простой MLP: linear -> GeLU -> linear -> GeLU -> ... -> linear -> GeLU"""
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(dim, dim) for _ in range(num_layers)])

    def forward(self, x: torch.Tensor):
        for layer in self.layers:
            x = layer(x)
            x = torch.nn.functional.gelu(x)
        return x


def run_mlp(dim: int, num_layers: int, batch_size: int, num_steps: int) -> Callable:
    # Определяем модель (со случайными весами)
    model = MLP(dim, num_layers).to(get_device())
    # Определяем вход (случайный)
    x = torch.randn(batch_size, dim, device=get_device())

    def run():
        # Прогоняем модель num_steps раз (заметьте: без шагов оптимизатора)
        for step in range(num_steps):
            y = model(x).mean()  # прямой проход
            y.backward()         # обратный проход
    return run


def run_operation1(dim: int, operation: Callable) -> Callable:
    # Подготовка: одна случайная матрица dim x dim
    x = torch.randn(dim, dim, device=get_device())
    # Возвращаем функцию, выполняющую операцию
    return lambda : operation(x)


def run_operation2(dim: int, operation: Callable) -> Callable:
    # Подготовка: две случайные матрицы dim x dim
    x = torch.randn(dim, dim, device=get_device())
    y = torch.randn(dim, dim, device=get_device())
    return lambda : operation(x, y)

Бенчмаркинг

Бенчмаркинг измеряет реальное (wall-clock) время выполнения операции. Он даёт только время «от начала до конца», а не то, где это время тратится (это профилирование). Но он всё равно полезен:

чтобы сравнивать реализации (какая быстрее?);
чтобы понимать, как производительность масштабируется (например, с размерностью).

Определим удобную функцию для бенчмарка произвольной функции.

def benchmark(description: str, run: Callable, num_warmups: int = 1, num_trials: int = 3):
    """Замеряем run(), выполняя её num_trials раз, и возвращаем среднее время (мс)."""
    # Прогрев: первые запуски бывают медленнее из-за компиляции и холодных кэшей.
    # Ядро мы будем вызывать много раз, поэтому важно время в установившемся режиме.
    for _ in range(num_warmups):
        run()
    if torch.cuda.is_available():
        torch.cuda.synchronize()  # Ждём, пока завершатся потоки CUDA (важно!)

    # Теперь меряем по-настоящему!
    times: list[float] = []
    for trial in range(num_trials):  # Несколько раз, чтобы увидеть разброс
        start_time = time.time()
        run()  # Собственно вычисление
        if torch.cuda.is_available():
            torch.cuda.synchronize()  # Ждём, пока завершатся потоки CUDA (важно!)
        end_time = time.time()
        times.append((end_time - start_time) * 1000)

    mean_time = mean(times)
    return mean_time

benchmark("sleep", lambda : time.sleep(50 / 1000))  # → 50.6 мс

Бенчмарк умножения матриц. Сначала замерим умножение квадратных матриц.

dims = (1024, 2048, 4096, 8192, 16384)
for dim in dims:
    result = benchmark(f"matmul(dim={dim})", run_operation2(dim=dim, operation=lambda a, b: a @ b))
# → matmul_results = [(1024, 0.87), (2048, 0.84), (4096, 3.15), (8192, 21.69), (16384, 162.94)]  (мс, H100)

Теперь замерим наш MLP!

dim = 256; num_layers = 4; batch_size = 256; num_steps = 2
mlp_base = benchmark("run_mlp", run_mlp(dim=dim, num_layers=num_layers, batch_size=batch_size, num_steps=num_steps))
# → 6.29 мс

Масштабируем число шагов: step_results → (2×: 11,33 мс), (3×: 16,76), (4×: 22,03), (5×: 27,31).

Масштабируем число слоёв: layer_results → (2×: 9,70 мс), (3×: 13,27), (4×: 16,86), (5×: 20,43).

Масштабируем размер батча: batch_results → (2×: 6,09 мс), (3×: 5,94), (4×: 6,08), (5×: 5,96).

Масштабируем размерность: dim_results → (2×: 6,07 мс), (3×: 6,05), (4×: 6,01), (5×: 6,09).

Время не всегда предсказуемо из-за неоднородной природы ядер CUDA, железа и т. д.

Можно также пользоваться torch.utils.benchmark — там больше удобств (рецепт). Мы его не использовали, чтобы бенчмаркинг был прозрачнее.

Профилирование

Бенчмаркинг смотрит на время «от начала до конца», а профилирование — на то, куда уходит время. Очевидное: профилирование помогает понять, где тратится время. Глубже: профилирование помогает понять, что вообще вызывается.

В PyTorch есть удобный встроенный профайлер (рецепт). Попрофилируем код и посмотрим, что происходит «под капотом».

def profile(description: str, run: Callable, num_warmups: int = 1, with_stack: bool = False):
    # Прогрев
    for _ in range(num_warmups):
        run()
    if torch.cuda.is_available():
        torch.cuda.synchronize()  # Ждём, пока завершатся потоки CUDA (важно!)

    # Запускаем код под профайлером
    with torch.profiler.profile(
            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
            # Записывать стек вызовов для визуализации
            with_stack=with_stack,
            # Нужно, чтобы экспортировать стек для визуализации
            experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)) as prof:
        run()
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    # Печатаем таблицу
    table = prof.key_averages().table(sort_by="cuda_time_total",
                                      max_name_column_width=80,
                                      row_limit=10)
    # Записываем стек вызовов для визуализации
    if with_stack:
        text_path = f"var/stacks_{description}.txt"
        prof.export_stacks(text_path, "self_cuda_time_total")
    return table

Начнём с простых операций (матрицы 2048 × 2048). Ниже — сокращённые таблицы профайлера (столбец «Name»; полные таблицы с процентами и временем — в трассе лекции).

## add
aten::add
void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add...
cudaLaunchKernel
cudaDeviceSynchronize
Self CUDA time total: 17.119us

## matmul
aten::matmul
aten::mm
void cutlass::Kernel2<cutlass_80_simt_sgemm_256x128_8x4_nn_align1>(cutlass_80...
cuLaunchKernel
cudaDeviceSynchronize
Self CUDA time total: 342.620us

## matmul(dim=128)
aten::matmul
aten::mm
sm80_xmma_gemm_f32f32_f32f32_f32_nn_n_tilesize32x32x8_stage3_warpsize1x2x1_ff...
cudaLaunchKernelExC
Self CUDA time total: 4.992us

Наблюдения:

Видно, какие ядра CUDA реально вызываются.
В зависимости от размеров тензоров вызываются разные ядра CUDA.

Имя ядра CUDA кое-что говорит о реализации. Пример: cutlass_80_simt_sgemm_256x128_8x4_nn_align1:

cutlass — библиотека NVIDIA для линейной алгебры на CUDA;
256x128 — размер тайла.

Теперь посмотрим на составные операции.

## cdist  (попарные расстояния)
aten::cdist → aten::_euclidean_dist → aten::matmul → aten::mm → sm80_xmma_gemm_..._tilesize128x128x8_...
aten::cat → CatArrayBatchedCopy...
aten::pow → vectorized_elementwise_kernel...
aten::sum
Self CUDA time total: 440.121us

## gelu  (gelu(a + b))
aten::add  → vectorized_elementwise_kernel<4, CUDAFunctor_add...>
aten::gelu → vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>
Self CUDA time total: 27.360us

## softmax  (softmax(a + b))
aten::softmax → aten::_softmax → cunn_SoftMaxForwardSmem<4, float, ...>
aten::add → vectorized_elementwise_kernel<4, CUDAFunctor_add...>
Self CUDA time total: 38.719us

Теперь попрофилируем наш MLP (dim=2048, 64 слоя, батч 1024, 2 шага). Заодно визуализируем стек вызовов в виде «flame graph» (диаграммы-пламени), которая показывает, куда уходит время.

## mlp
autograd::engine::evaluate_function: AddmmBackward0
AddmmBackward0
aten::mm
aten::linear
aten::addmm
sm80_xmma_gemm_f32f32_f32f32_f32_tn_n_tilesize128x128x8_stage3_warpsize2x2x1_...
void cutlass::Kernel2<cutlass_80_simt_sgemm_256x128_8x4_nn_align1>(...)
void cutlass::Kernel2<cutlass_80_simt_sgemm_256x128_8x4_nt_align1>(...)
autograd::engine::evaluate_function: torch::autograd::AccumulateGrad
torch::autograd::AccumulateGrad
Self CPU time total: 39.129ms
Self CUDA time total: 73.598ms

Каждый раз, когда что-то меняете, — делайте бенчмарк и профилируйте!

Зачем сливать ядра

Пост Хораса Хе: Making Deep Learning Go Brrrr From First Principles.

Аналогия: склад : DRAM :: фабрика : SRAM. [Рисунок: склад и фабрика, связанные конвейером ограниченной ширины.]

Каждая операция должна прочитать, вычислить и записать: [рисунок: несколько операций подряд, каждая возит данные со склада и обратно].

Если слить операции (kernel fusion), читать и писать нужно лишь один раз: [рисунок: данные один раз приезжают, проходят все операции и уезжают].

Чтобы увидеть эффект слияния, возьмём функцию активации GeLU (документация). Есть два способа её вычислить.

def pytorch_gelu(x: torch.Tensor):
    # Используем приближение через tanh, чтобы совпасть с нашей реализацией
    return torch.nn.functional.gelu(x, approximate="tanh")

def manual_gelu(x: torch.Tensor):
    return 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x * x * x)))

x = torch.tensor([1.])
y1 = pytorch_gelu(x)  # 1. Стандартная реализация PyTorch (слитая)  → [0.8412]
y2 = manual_gelu(x)   # 2. Своя, написанная вручную (не слитая)     → [0.8412]
assert torch.allclose(y1, y2)

Замерим (матрица 16384 × 16384):

manual_time = benchmark("manual_gelu", run_operation1(dim=16384, operation=manual_gelu))    # → 8.15 мс
pytorch_time = benchmark("pytorch_gelu", run_operation1(dim=16384, operation=pytorch_gelu)) # → 1.13 мс

Слитая версия значительно быстрее: 8,15 мс против 1,13 мс.

Заглянем под капот:

## manual_gelu
aten::mul  → vectorized_elementwise_kernel<4, BinaryFunctor<...>>      (несколько раз)
           → vectorized_elementwise_kernel<4, AUnaryFunctor<...>>
aten::add  → vectorized_elementwise_kernel<4, CUDAFunctor_add...>
           → vectorized_elementwise_kernel<4, CUDAFunctorOnSelf...>
aten::tanh → vectorized_elementwise_kernel<4, tanh_kernel_cuda...>
Self CUDA time total: 7.669ms

## pytorch_gelu
aten::gelu → vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>
Self CUDA time total: 701.560us

PyTorch вызывает всего одно ядро, тогда как в ручной версии каждая операция атомарна (вспомните склад и фабрику).

Дальше в лекции — просмотр MLP в профайлере Nsight (это делалось вживую, в трассе только заголовок).

Ядра на CUDA

Теперь откроем «чёрный ящик» и поймём, что происходит внутри ядра CUDA, написав своё. Напишем функцию GeLU на CUDA.

CUDA — это расширение C/C++ с API для управления GPU. Упрощённая картина: вы пишете $f(i)$ , а ядро CUDA вычисляет $f(i)$ для всех $i$ .

[Рисунок из документации PTX: сетка блоков потоков.]

Сетка (grid) — набор блоков потоков: numBlocks = (2, 4), blockDim = (1, 8).
Блок потоков — набор потоков: blockIdx = (0, 1).
Поток — одна единица работы: threadIdx = (0, 3).

Вы пишете код, который выполняет один поток, и по (blockIdx, blockDim, threadIdx) он решает, что делать.

Установите CUDA_LAUNCH_BLOCKING, чтобы при ошибке CUDA сообщила, что пошло не так:

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

Функция load_inline позволяет удобно написать код на CUDA и сразу привязать его к модулю Python. Вот код ядра (gelu.cu):

#include <math.h>
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>

__global__ void gelu_kernel(float* in, float* out, int num_elements) {
    // Вычисляем индекс в тензоре
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (i < num_elements) {  // На случай, когда n меньше numBlocks * blockDim
        // Собственно вычисление
        out[i] = 0.5 * in[i] * (1.0 + tanh(0.79788456 * (in[i] + 0.044715 * in[i] * in[i] * in[i])));
    }
}

inline unsigned int cdiv(unsigned int a, unsigned int b) {
    // Вычисляем ceil(a / b)
    return (a + b - 1) / b;
}

torch::Tensor gelu(torch::Tensor x) {
    TORCH_CHECK(x.device().is_cuda());
    TORCH_CHECK(x.is_contiguous());

    // Выделяем пустой тензор
    torch::Tensor y = torch::empty_like(x);

    // Определяем сетку (элементы делятся на блоки)
    int num_elements = x.numel();
    int block_size = 1024;  // Число потоков
    int num_blocks = cdiv(num_elements, block_size);

    // Запускаем ядро
    gelu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), num_elements);
    C10_CUDA_KERNEL_LAUNCH_CHECK();  // Ловим ошибки сразу

    return y;
}

Компилируем код CUDA и привязываем к модулю Python:

cpp_gelu_src = "torch::Tensor gelu(torch::Tensor x);"   # C++: объявление функции gelu
module = load_inline(
    cuda_sources=[cuda_gelu_src],
    cpp_sources=[cpp_gelu_src],
    functions=["gelu"],
    extra_cflags=["-O2"],
    verbose=True,
    name="inline_gelu",
    build_directory="var/cuda_gelu",
)
cuda_gelu = getattr(module, "gelu")

Проверяем правильность (check_equal(cuda_gelu, manual_gelu)) и замеряем нашу версию на CUDA:

pytorch_time = benchmark("pytorch_gelu", ...)  # → 1.13 мс
manual_time = benchmark("manual_gelu", ...)    # → 8.10 мс
cuda_time = benchmark("cuda_gelu", ...)        # → 1.80 мс

## cuda_gelu
gelu_kernel(float*, float*, int)
aten::empty_like → aten::empty_strided
cudaLaunchKernel
cudaDeviceSynchronize
Self CUDA time total: 1.664ms

Наша реализация на CUDA быстрее ручной, но не так хороша, как PyTorch. Поэлементные операции на CUDA пишутся легко (хотя и тут можно быть умнее). Но большинство интересных операций (например, matmul, softmax, RMSNorm) требуют читать много значений. Для них приходится думать об управлении разделяемой памятью и т. д.

Ядра на Triton

Triton разработан OpenAI в 2021 году. Цель — сделать программирование GPU доступнее:

пишем на Python;
думаем о блоках потоков, а не об отдельных потоках.

Что даёт Triton?

                                             CUDA      Triton
- Объединение обращений (перенос из DRAM)    вручную   автоматически
- Управление разделяемой памятью             вручную   автоматически
- Планирование внутри SM                     вручную   автоматически
- Планирование между SM                      вручную   вручную

Компилятор делает больше работы и может даже обгонять реализации PyTorch! Большое преимущество Triton — код на Python можно проходить по шагам (в режиме интерпретатора TRITON_INTERPRET=1). Пройдём по шагам ядро Triton.

def triton_gelu(x: torch.Tensor):
    assert x.is_cuda
    assert x.is_contiguous()

    # Выделяем выходной тензор
    y = torch.empty_like(x)

    # Определяем сетку (элементы делятся на блоки)
    num_elements = x.numel()
    block_size = 1024  # Число потоков
    num_blocks = triton.cdiv(num_elements, block_size)

    triton_gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=block_size)
    return y


@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    # Вход лежит по адресу x_ptr, выход — по адресу y_ptr
    #     |        Блок 0             |          Блок 1           |      ...      |
    #                            BLOCK_SIZE                                 num_elements

    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE

    # Индексы, с которыми должен работать этот блок
    offsets = block_start + tl.arange(0, BLOCK_SIZE)

    # Обрабатываем границу
    mask = offsets < num_elements

    # Читаем
    x = tl.load(x_ptr + offsets, mask=mask)

    # Приближённый gelu: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
    # Вычисляем (tl.tanh нет, используем tanh(a) = (exp(2a) - 1) / (exp(2a) + 1))
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    exp = tl.exp(2 * a)
    tanh = (exp - 1) / (exp + 1)
    y = 0.5 * x * (1 + tanh)

    # Записываем
    tl.store(y_ptr + offsets, y, mask=mask)

Смотрим на PTX

PTX (parallel thread execution) — это что-то вроде ассемблера для GPU (документация). Можно посмотреть PTX, который сгенерировал Triton. Фрагмент:

mov.u32 %r1, %ctaid.x;
mov.u32 %r44, %tid.x;
@%p1 ld.global.v4.b32 { %r2, %r3, %r4, %r5 }, [ %rd1 + 0 ];
@%p2 ld.global.v4.b32 { %r6, %r7, %r8, %r9 }, [ %rd2 + 0 ];
ex2.approx.f32 %f1, %f2;
...                                   (ещё 7 таких инструкций)
@%p1 st.global.v4.b32 [ %rd3 + 0 ], { %r34, %r35, %r36, %r37 };
@%p2 st.global.v4.b32 [ %rd4 + 0 ], { %r38, %r39, %r40, %r41 };

Наблюдения:

ld.global.* и st.global.* читают из глобальной памяти и пишут в неё;
%ctaid.x — индекс блока, %tid.x — индекс потока;
%f* — регистры с плавающей точкой, %r* — целочисленные регистры;
один поток обрабатывает 8 элементов одновременно (thread coarsening, «укрупнение потоков»).

Проверяем правильность (check_equal(triton_gelu, manual_gelu)) и сравниваем с реализациями PyTorch и CUDA. Не забудьте выставить TRITON_INTERPRET=0 для хорошей производительности.

manual_time = benchmark("manual_gelu", ...)    # → 8.16 мс
pytorch_time = benchmark("pytorch_gelu", ...)  # → 1.17 мс
cuda_time = benchmark("cuda_gelu", ...)        # → 1.85 мс
triton_time = benchmark("triton_gelu", ...)    # → 1.85 мс

## triton_gelu
triton_gelu_kernel
aten::empty_like → aten::empty_strided
cuLaunchKernel
cudaDeviceSynchronize
Self CUDA time total: 705.240us

Наша реализация на Triton (triton_gelu):

почти так же хороша, как реализация PyTorch (pytorch_gelu);
на деле медленнее нашей наивной реализации на CUDA (cuda_gelu).

Triton работает с блоками, CUDA — с потоками. Блоки позволяют компилятору Triton делать другие оптимизации (например, укрупнение потоков). Всё гораздо быстрее ручной реализации (manual_gelu).

Компиляция в PyTorch

Итак, мы видели три способа написать GeLU (на самом деле четыре):

стандартная функция PyTorch;
на Python (manual_gelu);
на CUDA (create_cuda_gelu);
на Triton (triton_gelu);
и ещё один: написать на Python и скомпилировать в Triton.

compiled_gelu = torch.compile(manual_gelu)
check_equal(compiled_gelu, manual_gelu)  # проверка правильности

manual_time = benchmark("manual_gelu", ...)      # → 8.10 мс
pytorch_time = benchmark("pytorch_gelu", ...)    # → 1.11 мс
cuda_time = benchmark("cuda_gelu", ...)          # → 1.82 мс
triton_time = benchmark("triton_gelu", ...)      # → 1.73 мс
compiled_time = benchmark("compiled_gelu", ...)  # → 1.47 мс

Заглянем под капот:

## compiled_gelu
Torch-Compiled Region: 0/1
triton_poi_fused_add_mul_tanh_0
TorchDynamo Cache Lookup
cuLaunchKernel
cudaDeviceSynchronize
Self CUDA time total: 707.261us

Примечание переводчика

Имя ядра triton_poi_fused_add_mul_tanh_0 расшифровывается так: torch.compile сгенерировал ядро на Triton (triton_), поточечное (poi — pointwise), в котором слиты (fused) сложения, умножения и tanh. То есть компилятор автоматически сделал то, что мы писали руками.

Softmax на Triton

До сих пор мы смотрели на поэлементные операции в Triton (например, GeLU). Теперь посмотрим на операции, которые агрегируют несколько значений. Будем примерно следовать учебнику Triton о слитом softmax (fused softmax). Вспомним: softmax используется во внимании и для получения вероятностей. Нормируем каждую строку матрицы:

[A1 A2 A3]   =>   [A1/A A2/A A3/A]
[B1 B2 B3]   =>   [B1/B B2/B B3/B]

Начнём с наивной реализации и посчитаем чтения/записи.

def manual_softmax(x: torch.Tensor):
    # M: число строк, N: число столбцов
    M, N = x.shape

    # Максимум каждой строки (MN чтений, M записей)
    x_max = x.max(dim=1)[0]

    # Вычитаем максимум (MN + M чтений, MN записей)
    x = x - x_max[:, None]

    # Берём экспоненту (MN чтений, MN записей)
    numerator = torch.exp(x)

    # Нормировочная константа (MN чтений, M записей)
    denominator = numerator.sum(dim=1)

    # Нормируем (MN чтений, MN записей)
    y = numerator / denominator[:, None]

    # Итого: 5MN + M чтений, 3MN + 2M записей
    # В принципе достаточно MN чтений и MN записей (ускорение в 4 раза!)
    return y

x = torch.tensor([[5., 5, 5],
                  [0, 0, 100]], device=get_device())
y1 = manual_softmax(x)
# → [[0.3333, 0.3333, 0.3333],
#    [3.78e-44, 3.78e-44, 1.0]]

Теперь напишем ядро на Triton.

def triton_softmax(x: torch.Tensor):
    # Выделяем выходной тензор
    y = torch.empty_like(x)

    # Определяем сетку
    M, N = x.shape                          # Число строк x число столбцов
    block_size = triton.next_power_of_2(N)  # Каждый блок содержит все столбцы
    num_blocks = M                          # Каждый блок — это строка

    # Запускаем ядро
    triton_softmax_kernel[(M,)](
        x_ptr=x, y_ptr=y,
        x_row_stride=x.stride(0), y_row_stride=y.stride(0),
        num_cols=N, BLOCK_SIZE=block_size
    )
    return y


@triton.jit
def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride, num_cols, BLOCK_SIZE: tl.constexpr):
    assert num_cols <= BLOCK_SIZE

    # Каждую строку обрабатываем независимо
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)

    # Читаем из глобальной памяти
    x_start_ptr = x_ptr + row_idx * x_row_stride
    x_ptrs = x_start_ptr + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))

    # Вычисляем
    x_row = x_row - tl.max(x_row, axis=0)
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator

    # Записываем обратно в глобальную память
    y_start_ptr = y_ptr + row_idx * y_row_stride
    y_ptrs = y_start_ptr + col_offsets
    tl.store(y_ptrs, y_row, mask=col_offsets < num_cols)

Проверяем, что реализации правильные (check_equal2(pytorch_softmax, manual_softmax), check_equal2(pytorch_softmax, triton_softmax)), добавляем compiled_softmax = torch.compile(manual_softmax) и замеряем всё:

manual_time = benchmark("manual_softmax", ...)      # → 3.71 мс
compiled_time = benchmark("compiled_softmax", ...)  # → 1.36 мс
pytorch_time = benchmark("pytorch_softmax", ...)    # → 1.54 мс
triton_time = benchmark("triton_softmax", ...)      # → 1.92 мс

Затем лекция смотрит на всё это профайлером и заканчивается PTX-кодом ядра triton_softmax_kernel (полный текст в репозитории).

Примечание переводчика: матмул на Triton

В файле лекции есть ещё функция triton_matmul_main (в 2025 году её не вызывали в main, и в трассе её нет). В ней — краткий план: наивный матмул требует

MKN

чтений и

MN

записей; вычисление

C_4

C_5

требует одних и тех же

A_4, A_5, A_6

— их можно один раз прочитать из DRAM в разделяемую память; тайлинг «загрузить блок A и блок B → мини-матмул → записать частичную сумму»; переупорядочивание блоков для L2 («грузим 90 блоков против 54»); и вывод: свой матмул пишут ради слияния с другой операцией, например gelu(A @ B). Подробный разбор с кодом — в уроке 5.7.3, раздел «Ядро 3: тайловый матмул».

Итоги

Разрыв между моделью программирования (PyTorch, Triton, PTX) и железом ⇒ загадки производительности.
Бенчмаркинг — чтобы понимать масштабирование.
Профилирование — чтобы понимать внутренности функций PyTorch (на самом дне — ядра).
Чтение ассемблера PTX — чтобы понимать внутренности ядер CUDA.
5 способов написать функцию: вручную, PyTorch, компиляция, CUDA, Triton.
GeLU (поэлементная), softmax (по строкам), matmul (сложная агрегация).
Главный принцип: организовать вычисления так, чтобы минимизировать чтения/записи.
Главные идеи: слияние ядер (аналогия склада и фабрики), тайлинг (разделяемая память).
Автоматические компиляторы (Triton, torch.compile) со временем будут становиться лучше.

Что почитать дальше

Пост Хораса Хе: Making Deep Learning Go Brrrr From First Principles.
CUDA MODE, лекция 1: как профилировать ядра CUDA в PyTorch — видео.
CUDA MODE, лекция 2: главы 1–3 книги PPMP — видео.
CUDA MODE, лекция 3: начало работы с CUDA для программистов на Python — видео.
CUDA MODE, лекция 4: основы вычислений и памяти — видео.
CUDA MODE, лекция 8: чек-лист производительности CUDA — видео.
Курс HetSys, лекции 1–5: программирование гетерогенных систем с GPU, SIMD и GPU, программная и аппаратная иерархия GPU, иерархия памяти, производительность — лекция 1, 2, 3, 4, 5.
A100 GPU with NVIDIA Ampere Architecture.
NVIDIA Deep Learning Performance Guide.
GPU Puzzles Саши Раша.
Статья о Triton (Tillet, Kung, Cox, 2019).
How PyTorch 2.0 accelerates deep learning with operator fusion.

Что сделать руками

Бенчмарк матмула на своей карте. Скопируйте функцию benchmark и замерьте a @ b для размеров 1024…8192 в float32 и bfloat16. Пересчитайте в TFLOPS ( $2n^3$ / время) и сравните с паспортными числами RTX 4080 из урока 5.7.1. Затем пройдите размеры 1536…2048 с шагом 1 и найдите «зубья» квантования волн.
Слияние. Повторите сравнение manual_gelu, pytorch_gelu и torch.compile(manual_gelu) на матрице 8192 × 8192. Посмотрите имена ядер в torch.profiler — найдите triton_poi_fused_….
Ядро на Triton. Напишите triton_gelu сами, не подглядывая, и проверьте torch.allclose против pytorch_gelu. Затем сделайте triton_softmax для матрицы с числом столбцов больше BLOCK_SIZE (подсказка: цикл по блокам столбцов и online softmax из слайда 49 лекции 5) — это шаг к FlashAttention из задания 2.
Посчитайте трафик. Для manual_softmax на матрице 16384 × 16384 в float32 оцените байты чтения/записи по комментариям в коде ( $5MN + M$ чтений, $3MN + 2M$ записей) и минимальное время при 716,8 ГБ/с. Сравните с замером.

Словарик

EN	RU
streaming multiprocessor (SM), streaming processor (SP)	потоковый мультипроцессор, потоковый процессор
thread, thread block (CTA), grid, warp	поток, блок потоков, сетка, варп (32 потока)
SIMT	одна инструкция — много потоков
global memory (DRAM, HBM), L2/L1 cache, shared memory (SRAM), registers	глобальная память, кэши L2/L1, разделяемая память, регистры
control divergence	расхождение потоков управления
tensor cores	тензорные ядра
roofline model	модель roofline («линия крыши»)
arithmetic intensity	арифметическая интенсивность
compute-bound / memory-bound	ограничено вычислениями / ограничено памятью
operator (kernel) fusion	слияние операторов (ядер)
recomputation	пересчёт (повторное вычисление вместо хранения)
memory coalescing, burst mode	объединение обращений к памяти, пакетный режим DRAM
tiling, tile	тайлинг, тайл (плитка)
wave quantization, occupancy	квантование волн, загрузка (заполненность SM)
online softmax	онлайн-softmax (за один проход)
benchmarking / profiling	бенчмаркинг (замер времени) / профилирование (где тратится время)
kernel, PTX	ядро (программа для GPU), PTX (ассемблер GPU)
thread coarsening	укрупнение потоков (один поток — несколько элементов)

Перевод выполнен по лицензии оригинала (MIT (© 2025 Stanford University)); автор оригинала — Тацунори Хашимото (Tatsunori Hashimoto), Stanford University. Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04