Апдейт: идеи, изложенные в этой статье, позволили сформулировать оптимальные стратегии warp-специализации, описанные в научной публикации, которую можно посмотреть здесь.
Недавно я глубоко задумался о специализации варпов в контексте высокопроизводительных ядер для современных графических процессоров (GPU) на тензорных ядрах. Примеры таких процессоров — H100 и B200 от NVIDIA. Я стал полнее понимать, чего можно добиться при помощи специализации варпов, а также задался интересным вопросом: а нужна ли нам вообще специализация варпов (и вся та сложность, которую она с собой влечёт)? В итоге я пришёл к выводу, что, да, нуждаемся, но она не столь обязательна, как может показаться. В этом посте обсудим, в каких случаях без специализации варпов действительно не обойтись, а также я опишу, на каком пространстве компромиссов она зиждется, и какие границы этого пространства я вижу. Притом, что я обрисую некоторый контекст, касающийся графических процессоров, необходимый для обсуждения тем, которые мы взялись здесь рассмотреть, эту статью нельзя считать туториалом. Предполагается, что читатель имеет некоторый опыт работы с GPU и имеет опыт параллельного программирования.
Контекст
GPU — это набор процессоров, именуемых «потоковыми мультипроцессорами» (SM). В рамках данной дискуссии мы сосредоточимся на программировании отдельно взятого SM. Потоковый мультипроцессор программируется при помощи иерархии нитей, так называемого «блока нитей» (thread block). Нити в блоке далее группируются в варпы, в каждом варпе — по 32 нити. Каждый варп выполняет вычисления по модели «одна инструкция, много потоков» (SIMT). Таким образом, у каждой нити в варпе — собственный поток инструкций, и каждый варп выдаёт одну инструкцию по инициативе своих потоков в каждый слот выдачи. Как будет рассмотрено выше, производительность доводится до максимума, когда все нити в пределах варпа хотят одновременно выдать одну и ту же инструкцию. Так, у потокового мультипроцессора Hopper (изображённого ниже) четыре контекста выполнения, в каждом из которых может размещаться один активный варп. Варпы изображены в виде 4 квадрантов.
На каждом такте максимум 4 варпа могут выдавать потоковому мультипроцессору инструкции для выполнения. Если в некотором блоке содержится нитей более чем на 4 варпа (128), то включается аппаратный компонент под названием «планировщик варпов» (warp scheduler), выбирающий для выполнения инструкций 4 доступных варпа.
Таким образом, потоковый мультипроцессор можно трактовать как набор функциональных единиц (арифметико-логических устройств, блоков загрузки/сохранения, тензорных ядер), на каждом такте процессора выдающих инструкции из 4 контекстов выполнения. Свойства этих функциональных единиц варьируются. Арифметические логические устройства (АЛУ) выполняют отдельные математические операции с соблюдением коротких и фиксированных тактовых задержек. Тензорные ядра в рамках единственной инструкции выполняют тысячи операций с плавающей точкой (FLOP), но допускают длинные тактовые задержки. При этом блоки загрузки/сохранения (LSU) допускают при работе (непредсказуемо) долгие задержки, поскольку им приходится взаимодействовать с системой памяти. Высокопроизводительные программы для GPU эффективно используют имеющиеся в распоряжении функциональные единицы. Те программы, которые располагают ограниченными вычислительными ресурсами, должны на каждом такте процессора использовать тензорное ядро и арифметическое логическое устройство, а программы, у которых ограничена полоса передачи данных, должны постоянно задействовать блоки загрузки/сохранения, чтобы довести до максимума пропускную способность. Чтобы добиться высокой степени задействования ресурсов, всегда должна быть в наличии работа, которой могут заниматься функциональные единицы (например, в приложении с ограниченными вычислительными ресурсами операции с плавающей точкой не должны накапливаться, дожидаясь, пока завершатся загрузки). Кроме того, такая доступная работа должна выдаваться функциональным единицам по мере их доступности. Именно во втором аспекте нам пригодится специализация варпов.
Суть специализации варпов
Специализация варпов — это технология, популяризованная в ходе работы над CUDA-DMA и компилятором языка Singe, а сегодня превратившаяся в необходимый минимум для достижения высокой производительности тензорных ядер на графических процессорах Hopper и Blackwell. Специализация варпов опирается на иерархическое группирование нитей в рамках блока. Когда начинают расходиться нити, относящиеся к одному и тому же варпу (т.e., по-разному ветвятся в потоке управления), производительность снижается именно по той причине, что каждый варп устроен по принципу SIMT. Допустим, варп доходит до такой точки, в которой путь у половины нитей ветвится, а у половины — нет. Теперь варп будет выполнять инструкции, оказавшиеся по обе стороны от точки ветвления. Когда варп выберет инструкцию с конкретной стороны от ветвления, нити, выполняющие задачи по другую сторону, ничего делать не будут. В результате на выполнение всех задач может уйти вдвое больше времени, чем если бы все нити в результате ветвления свернули на один и тот же путь. В наихудшем случае, если по разным траекториям в потоке управления пойдут все 32 нити, входящие в варп, то код будет выполняться в 32 раза медленнее, чем в идеале. Разные варпы в пределах блока нитей (в отличие от разных нитей в пределах варпа) выполняются независимо друг от друга в разных контекстах выполнения. Таким образом, не возникает издержек при расхождении между варпами. Именно это свойство, присущее расхождению варпов, используется при варп-специализации для реструктурирования программ для GPU. Стандартная программа для GPU выполняет в каждом варпе одну и ту же логику, тогда как программа, использующая специализацию варпов, задействует разные варпы для выполнения различных компонентов всей программы. Ниже рассмотрим некоторые из таких стратегий специализации варпов в вышеупомянутых контекстах.
В рамках проекта CUDA-DMA предложили отделить загрузку данных из глобальной (медленной) памяти GPU в быструю (разделяемую) память от вычислений, производимых на тех данных, что обрабатываются в разделяемой быстрой памяти. Таким образом, в рамках проекта CUDA-DMA все варпы подразделили на те, что загружают данные а память, и те, что занимаются вычислениями. Варпы-загрузчики выдают загрузки и сигнализируют варпам-вычислителям, когда именно загруженные данные будут доступны для обработки.
Компилятор Singe предназначался для работы с целым поколением эффективных ядер, применяемых для вычислений в области химии горения. В контексте этого поста такие ядра можно рассматривать, в сущности, как крупные распараллеленные вычисления данных (например, с применением конкретной функции f к каждому элементу в массиве), но с оговоркой. Дело в том, что такие вычисления f требуют хранить большой объём промежуточного состояния (многочисленные временные переменные в составе химических формул). Чтобы напрямую реализовать такие ядра, потребовалось бы выделить слишком много регистров под хранение промежуточного состояния и сливать значения в стек – из-за этого сильно снижалась бы производительность. Наиболее досадно в данном случае то, что в файле регистров SM вполне достаточно места для хранения всех временных данных. Но архитектура устроена так, что каждой нити предоставляется фиксированное количество доступных регистров (например, в Hopper — по 255 регистров на нить). В Singe специализация варпов помогает обойти этот предел количества регистров на нить, подразделяя вычисление f на разные варпы. А именно, предположим, что:
Учитывая, что на каждую нить можно выделить ограниченное количество регистров, реализация f c применением специализации варпов предполагает, что вычисление 1 + x + 2 ·x можно вынести в первый варп, а вычисление x2 + 8 ·x3 — во второй варп. Далее два варпа должны обменяться информацией, чтобы суммировать промежуточные значения.
Наконец, специализация варпов применяется при работе с высокопроизводительными тензорными ядрами для процессоров Hopper и Blackwell и позволяет взаимодействовать с ускорителями, фигурирующими в потоковом мультипроцессоре. На этих графических процессорах SM содержит ускорители, выполняющие перемножение матриц (Tensor Core) и перемещающие данные в глобальную память и из неё (тензорный ускоритель памяти, он же TMA). Эти ускорители предлагают инструкции для перемножения тайлов данных или для копирования тайлов данных в глобальную память или из неё. Кроме того, эти ускорители являются асинхронными в таком смысле: всего одна инструкция запускает работу на ускорителе, после чего должна быть назначена блокирующая операция ожидания, перед тем, как можно будет воспользоваться результатами выполнения инструкции. Специализированные варпы применяются на процессорах Hopper и Blackwell для выдачи либо копий TMA, либо произведений матриц, полученных при помощи тензорных ядер. Варп TMA выдаёт копии и уведомляет варпы тензорных ядер, когда данные будут готовы к перемножению, а варпы тензорных ядер уведомляют варп TMA, когда данные будут потреблены и, соответственно, память освободится под новые копии. Этот код выглядит примерно так:
if warpid() == LOAD: for i, tile in enumerate(tiles): if i > 0: wait_for_tile_release() async_tma_load(tile) wait_for_tma_load() signal_tile_loaded() else: for tile in enumerate(tiles): wait_for_tile_loaded() tile_data = get_loaded_tile(tile) async_mma(tile_data) wait_for_async_mma() signal_tile_released()
В ядрах Tensor Core встречаются значительно более сложные стратегии специализации варпов, функциональность которых далеко не ограничивается перемножением матриц. Например, в высокопроизводительной реализации Flash-внимания на Blackwell используется не менее 5 видов различных специализированных варпов! В реализации Flash-внимания предусмотрены варпы для загрузки данных, перемножения матриц, вычисления softmax, шкалирования промежуточных результатов и хранения данных. Именно поэтому код получается сложным; стратегия как таковая тщательно выстраивается под достижение высокой производительности, а приходится иметь дело с массовым перемещением данных между варпами и синхронизацией. Представьте себе, если бы в вышеприведённом коде было бы описано 5 вариантов варпов, и каждый из этих вариантов сигнализировал бы остальным продолжать работу в разные временные окна!
Чем хороша специализация варпов?
Учитывая, как сложно было реализовать этот вариант Flash-внимания, я решил отойти немного назад и исследовать, какова роль специализации варпов в достижении высокой производительности при работе с тензорными ядрами. Такую необходимость специализации варпов в контексте тензорных ядер я воспринимал как данность. Более сведущие коллеги сказали мне, что она необходима, и я этого не оспаривал (что меня как академического учёного раздражает). Кроме того, в других объяснениях специализации варпов, которые мне удалось найти, попадались зыбкие формулировки из разряда «этого требует архитектура» или «это необходимо для создания конвейеров производитель-потребитель».
Давайте, отталкиваясь от базовых принципов, выведем, в каких случаях полезны специализации варпов. Потоковый мультипроцессор располагает ограниченным количеством вычислительных ресурсов (таких, как АЛУ, блоки загрузки/сохранения, тензорные ядра) и слотов для выдачи инструкций на каждый такт процессора — независимо от того, сколько варпов использует данный блок нитей. Следовательно, при использовании потокового мультипроцессора H100 ядро обладает одинаковой теоретической пиковой пропускной способностью при вычислениях и пиковым количеством инструкций на слот выдачи, независимо от того, сколько варпов оно использует — 4 или 64. Так в чём же польза? Рассмотрим две версии целевой программы: одна со специализированными варпами, а другая без них. Ядро со специализацией варпов использует более 4 варпов для выдачи потоковому мультипроцессору потенциально разных потоков инструкций, а сами варпы динамически перемежаются друг с другом — за это отвечает планировщик инструкций. Стандартная программа использует всего один поток инструкций, выдаваемых от 4 идентичных варпов. Очевидно, специализация варпов может влиять на производительность лишь при условии, что поток динамически перемежаемых инструкций от 4 различных варпов отличается от статически заданного потока инструкций, выдаваемых от 4 варпов в стандартной программе. Условия, при которых возникнут отличия между двумя такими потоками инструкций — как раз те самые, при которых специализация варпов может обеспечить выигрыш в производительности. Думаю, такие условия складываются в трёх случаях.
Первый случай легко выявляется, и именно на него ориентирован Singe: из-за ограниченности ресурсов неспециализированной версии просто не существует! Если бы неспециализированная версия использовала слишком много регистров, предикатов или других ресурсов потокового мультипроцессора, которыми варп располагает в ограниченном количестве, то специализация варпов помогала бы укладываться в эти ограничения. Специализация варпов, обусловленная ресурсными ограничениями — обычное дело в ядрах Hopper Tensor Core, где аккумуляторные тайлы распределены по разным группам варпов. Это нужно, чтобы оставаться в пределах количества регистров, выделенных на нить, и не сливать регистры в стек. Если сливать в стек регистры в высокопроизводительных ядрах для линейной алгебры при работе на графическом процессоре, это может приводить к сильному снижению производительности. Предпринимаются значительные усилия, чтобы оптимизировать код и настраивать параметры, укладывая их в лимиты регистров.
Второй случай чуть нетривиальнее, поскольку он обязательно предполагает существование неспециализированной версии целевой программы. В таком случае требуется обсудить планирование инструкций. В состав потокового мультипроцессора входит несколько независимых функциональных компонентов — например, FP для арифметики с плавающей точкой и INT для целочисленной арифметики. Они могут выполнять операции одновременно. Представьте себе программу, в которой за 2 операциями с плавающей точкой следуют 2 не зависящие от них целочисленные операции. Хороший планировщик инструкций назначил бы сначала одну операцию с плавающей точкой, а за ней – целочисленную операцию, чтобы блоки FP использовались одновременно с блоками INT. Когда компилятор (например, NVCC) обладает точной информацией о том, сколько тактов процессора затрачивается на каждую инструкцию, он может выдавать высококачественные статические планы и задействовать параллелизм на уровне инструкций (ILP), перемежая независимые инструкции, относящиеся к разным функциональным блокам. Но компилятору становится сложно планировать инструкции, когда такое количество тактов процессора на инструкцию определяется неточно. Если на инструкцию может уходить от 10 до 100 тактов процессора, а не строго 25, то статически собрать плотный план будет значительно сложнее. Поэтому именно во втором случае специализация варпов особенно полезна: мы имеем дело с динамическим планированием, при котором планировщик варпов может аккуратно справляться с переменной длительностью задержки при обработке инструкций, а такая изменчивость типична для операций, связанных с памятью. В данном случае статически спланированная и неспециализированная программа вынуждена угадывать, сколько времени потребуется на каждую операцию, задержка при которой может варьироваться, и на основании этого составить план, при котором операции с переменной задержкой перемежаются с операциями с фиксированной задержкой. Если операции с переменной задержкой будут выполняться быстрее, чем ожидалось, то функциональные блоки окажутся недоиспользованными, а если медленнее – то временами работа этих блоков будет стопориться. При реализации с применением специализированных варпов гадать не приходится, и инструкции перемежаются прямо во время выполнения — за это отвечает планировщик варпов.
Третья (и последняя) известная мне причина связана со сложностью обработки инструкций, сопряжённых с задержками переменной длительности, но требует ещё подробнее углубиться в детали архитектуры ЦП. Чтобы обрисовать контекст, сравним архитектуру графического процессора с архитектурой обычного центрального процессора. Современные центральные процессоры выполняют инструкции в относительно произвольном порядке (OOO). Притом, что процессор располагает последовательностью инструкций, он изыскивает способы их переупорядочивать, так, чтобы можно было задействовать параллелизм на уровне инструкций (ILP), в то же время, поддерживая иллюзию последовательного выполнения. Конкретный пример: выполняя инструкции addf r1 r2; addf r3 r4; addi r5 r6; addi r7 r8 (из вышеприведённого примера с планировщиком), ЦП может заранее автоматически подтягивать независимые инструкции addi и выполнять их в то же время, пока выполняются инструкции addf. Компилятор при этом может помочь аппаратному планировщику, частично переупорядочивая инструкции, но именно благодаря возможности OOO процессор выполняет часть тяжёлой работы. Основной недостаток OOO заключается в дороговизне оборудования, необходимого для реализации этой технологии. Действительно, в GPU площадь чипа и энергия, необходимая для его работы, экономится за счёт того, что инструкции выполняются строго по порядку. Когда поток инструкций на GPU содержит, например, addf r1 r2; addf r3 r4; addi r5 r6; addi r7 r8, инструкции выдаются потоковому мультипроцессору именно в таком порядке. Компилятор никак их не переставляет; в противном случае недоиспользовались бы функциональные блоки потокового мультипроцессора. Эта проблема усугубляется при применении инструкций синхронизации, которые используются для взаимодействия с асинхронными ускорителями (Tensor Core, TMA). Эти инструкции синхронизации действуют, в сущности, как семафоры, временно усыпляющие варпы до тех пор, пока вызванный ускоритель не справится с работой. Если разместить эти инструкции синхронизации неоптимальным образом в пределах потока, то варп может оказаться выключен из выполнения независимых инструкций до полного завершения той операции, относительно которой он синхронизировался. Притом, что компиляторы достаточно хорошо справляются с планированием, есть две сложности, провоцируемые такими синхронизационными действиями: 1) зачастую операции по синхронизации превращаются в серьёзную преграду на пути переноса кода, поскольку сложно доказать корректность операций переупорядочивания, связанных с синхронизацией (особенно тех, что касаются работы с памятью) и 2) компилятору бывает сложно точно отследить, какая именно операция приведёт к разрешению точки синхронизации (то есть, какая именно нить снимет блокировку). Благодаря специализации варпов, программист может игнорировать эффекты такой синхронизации: при разбиении вычислений на отдельные варпы, другие варпы могут выполнять работу прямо в тот период, пока другие блоки варпов дожидаются синхронизации.
Резюмируя сказанное в нескольких предыдущих абзацах, отметим, что специализация варпов может приносить пользу в следующих ситуациях:
Когда потребность приложения в ресурсах превышает объём ресурсов, доступных у одной нити или в пределах одного варпа.
В приложении выполняются операции, которым присуща переменная задержка, и их нужно перемежать, чтобы максимально эффективно использовать процессор.
В приложении есть акты блокирующей синхронизации, поэтому точки синхронизации нужно продуманно расставлять, чтобы избежать его пробуксовки.
Условия 2 и 3 по природе своей переплетены друг с другом, поскольку оба они обусловлены осложнениями, из-за которых так непросто оптимально статически спланировать последовательность выполнения инструкций для процессоров, обрабатывающих инструкции по порядку. Чтобы понять, почему специализация варпов помогает в таких обстоятельствах, нужно понять следующее: фактически, специализация варпов превращает потоковый мультипроцессор из упорядоченного в квази-неупорядоченный, особенно на тех границах, где стыкуются варпы с разными специализациями. Точки специализации, выбранные на этапе составления стратегии специализации варпов, позволяют судить, где может встретиться параллелизм на уровне инструкций (но в статике составить такую картину сложно). В то же время, внутри специализированного варпа вполне реально статически спланировать порядок следования инструкций.
Проверка гипотезы
В предыдущем разделе мы пришли к выводу, что специализация варпов полезна в следующих случаях: 1) специализация продиктована ограниченностью ресурсов, 2) сложно статически спланировать порядок выполнения инструкций, для которых характерны задержки с переменной длительностью и 3) блокирующая синхронизация нарушает порядок выдачи инструкций. Давайте исследуем, как возникают (и возникают ли вообще) такие сложности при разработке высокопроизводительной реализации H100 GEMM для задачи размером 8192x8192x8192. Для этих экспериментов я вручную изменил код, сгенерированный компилятором Cypress; то, что у него получилось, выглядит настолько уродливо, что не хочется вам это показывать. Поэтому рассмотрим результат в виде псевдокода, отдалённо напоминающего Python. Структура FP16 H100 GEMM примерно напоминает следующую.
Эффективная операция GEMM координирует работу программного конвейера так, что для большинства PIPE необслуженные загрузки тайлов (поступающие через TMA) из глобальной памяти в разделяемую остаются в состоянии ожидания, в то время как операции GEMM распараллеливаются для выполнения на тензорном ядре. В имеющихся сейчас реализациях две этих составляющих конвейера выполняются в отдельных специализированных варпах. Ниже показан более подробный код, описывающий умножение C = A @ B.
# Инициализация кольцевых буферов, в которых содержатся фрагменты PIPE от A и B. Abuf = [tile()] * PIPE Bbuf = [tile()] * PIPE if warpid() == LOAD: for k in K / KTILE: if k > PIPE: wait_for_compute_iter(k - PIPE) # Индекс кольцевого буфера, используемый на итерации k — это k % PIPE. async_tma_load(tile(A, k), Abuf[k % PIPE]) async_tma_load(tile(B, k), Bbuf[k % PIPE]) signal_when_load_complete_iter(k) wait_all_mmas_done() copy_C_shared_memory_to_global() else: C = init_accumulator() for k in K / KTILE: wait_for_load_complete(k) async_mma(C, Abuf[k % PIPE], Bbuf[k % PIPE]) wait_for_mma() signal_compute_iter_done(k) store_C_into_shared_memory() notify_all_mmas_done()
Давайте разберём этот код и выясним, какие здесь действуют условия, касающиеся специализации варпов. В результате настройки оказывается, что наиболее высокая производительность применительно к аккумулятору C достигается при размере тайла 256x256. Аккумулятор C для H100 нужно размещать в регистрах. Нехитрая математика показывает, что (256 * 256 элементов FP16 / 128 нитей (4 варпа, по 1 для каждого контекста выполнения) / 2 элемента FP16 на 32-разрядный регистр = 256 регистров) аккумулятор уже выходит за пределы регистров на нить, то есть, наступает условие 1. Это значит, что для хранения C нам потребуется, как минимум, две группы по 4 варпа. Что же насчёт двух прочих условий? Не вполне очевидно, что синхронизация или планирование инструкций окажутся настолько сложными, что нам придётся вынести загрузку данных в отдельный варп. Если использовать достаточно глубокий конвейер и синхронизироваться лишь в том случае, когда у нас уже накопились ожидающие задачи, то в обычных условиях можно обойтись и без специализации варпов. Ниже показана неспециализированная версия того GEMM, который мы рассматривали выше; для запуска конвейера потребуется некоторое дополнительное переупорядочивание операций и расщепление цикла, но в целом структура похожа.
Abuf = [tile()] * PIPE Bbuf = [tile()] * PIPE # Разделим аккумулятор на 2 части по одной для # каждой группы варпов, которые будут к нему обращаться. C = init_accumulator(warpid()) # Выдать загрузки, ожидающие выполнения в первом PIPE. for k in PIPE: async_tma_load(tile(A, k), Abuf[k]) async_tma_load(tile(B, k), Bbuf[k]) # Главный цикл выполняет MMA и загрузки на будущее. for k in [PIPE, K / KTILE): # Логично, что этот MMA — для итерации k - PIPE. wait_for_load_complete(k - PIPE) async_mma(C, Abuf[k % PIPE], Bbuf[k % PIPE]) # Эта операция ожидания MMA также захватывает ожидание # спаренного блока нитей wait_for_mma() # Начинаем следующую загрузку. async_tma_load(tile(A, k), Abuf[k % PIPE]) async_tma_load(tile(B, k), Bbuf[k % PIPE]) # Выполняем идущие в хвосте операции PIPE MMA. for k in [K / KTILE - PIPE, K / KTILE): wait_for_load_complete(k) async_mma(C, Abuf[k % PIPE], Bbuf[k % PIPE]) wait_for_mma() # Копируем аккумулятор в глобальную память (пропускается).
К сожалению, на практике работает совсем не так хорошо, как я ожидал:
> ./gemm_nows_v1 --m 8192 --n 8192 --k 8192 MY_GEMM: [675868.8]GFlop/s (1.6268)ms CUBLAS_GEMM: [805409.4]GFlop/s (1.3652)ms
Соответственно, применяется либо условие 2, либо условие 3, в зависимости от того, что именно будет недоиспользоваться — TMA или тензорное ядро. Но надежда ещё есть; меня в этом коде особенно впечатлила операция wait_for_mma() внутри главного цикла. Из-за неё варп блокируется и не может выдавать новые загрузки, пока не завершатся ожидающие выполнения задачи MMA. В свою очередь, это может застопорить выдачу дальнейших MMA. Эта ситуация решается, если вновь конвейеризовать цикл, где у нас ещё остались ожидающие обработки MMA, выданные до синхронизации. Тогда можно надеяться, что вопросы синхронизации к тому моменту будут закрыты за счёт сделанной работы. Новый код приобретает такой вид:
Abuf = [tile()] * PIPE Bbuf = [tile()] * PIPE # Разделим аккумулятор на 2 части по одной для # каждой группы варпов, которые будут к нему обращаться. C = init_accumulator(warpid()) # Выдаём загрузки из первого PIPE, ожидающие выполнения. for k in PIPE: async_tma_load(tile(A, k), Abuf[k]) async_tma_load(tile(B, k), Bbuf[k]) # Запускаем k=0 MMA, но не дожидаемся его. wait_for_load_complete(0) async_mma(C, Abuf[0], BBuf[0]) for k in [PIPE, K / KTILE): # Логично, что этот MMA — для итерации k - PIPE + 1. wait_for_load_complete(k - PIPE + 1) async_mma(C, Abuf[(k - PIPE + 1) % PIPE], Bbuf[(k - PIPE + 1) % PIPE]) # Дожидаемся MMA от итерации k - PIPE, а не # того MMA, который мы только что запустили (то есть, не k - PIPE + 1). wait_for_mma(k - PIPE) # Запускаем следующую загрузку. async_tma_load(tile(A, k), Abuf[k % PIPE]) async_tma_load(tile(B, k), Bbuf[k % PIPE]) # Выполняем идущие в хвосте операции PIPE MMA. for k in [K / KTILE - PIPE + 1, K / KTILE): wait_for_load_complete(k) async_mma(C, Abuf[k % PIPE], Bbuf[k % PIPE]) wait_for_mma()
Полюбуйтесь!
> ./gemm_nows_v2 --m 8192 --n 8192 --k 8192 MY_GEMM: [815881.7]GFlop/s (1.3476)ms CUBLAS_GEMM: [807708.0]GFlop/s (1.3613)ms
В данном конкретном случае можно добиться производительности, сопоставимой с CuBLAS, даже не вынося рабочие нагрузки TMA в отдельный варп. Здесь удалось разобраться с циклами вручную, чтобы избежать эффектов, характерных для условий 2 и 3. Даже берусь утверждать, что, хотя здесь и присутствует условие 1, в получившейся программе мы всё равно обошлись без полноценной специализации варпов! Аккумулятор пришлось разнести на несколько варпов, чтобы уложиться в ограничения по количеству регистров, но сами варпы выполняли одну и ту же программу. Наше упражнение показало, что, как минимум, при решении задач такого масштаба на H100, пиковая производительность достижима и без специализации варпов.
Специализация варпов - это компромисс
В предыдущих разделах были изложены некоторые принципы специализации варпов, после чего продемонстрировано, что бывает, как минимум, один полезный случай, в котором специализация варпов казалась необходимым условием (как минимум, мне), но удалось добиться высокой производительности и без неё. Но я полагаю, что это упражнение подводит нас к следующему выводу: не увлекайтесь написанием программ со специализацией варпов. Напротив, следует рассматривать специализацию варпов как пространство для компромисса, а не обязательный вариант при реализации ядер, к которым предъявляются очень высокие требования по производительности (например, к Tensor Core). Берусь утверждать, что при реализациях высокопроизводительных ядер Tensor Core требуется сориентироваться, сколько усилий уйдёт на написание ядра, а сколько — на разработку такого компилятора, который будет в состоянии проанализировать и правильно интерпретировать, чего хотел автор ядра. В данном случае анализ полезности и рисков заключается в том, как распределить силы, чтобы программист графических процессоров мог сэкономить как можно больше времени.
Если перерабатывать не хочется — то в компиляторе и так предусмотрено несколько видов анализа, на основе которых можно добиться высокой производительности. Во-первых, можно поручить планирование инструкций такому низкоуровневому компилятору как NVCC, но ему присущи проблемы, рассмотренные выше. Ещё одно направление — поручить специализацию варпов компилятору, написанному на более высокоуровневом языке, например, Triton или Cypress. Поддержка специализации варпов на уровне компиляторов кажется многообещающей, но компилятор пока уступает человеку в выработке хороших стратегий специализации варпов. В Triton даже добавили низкоуровневый язык Gluon, программируя на котором эксперт может обойти настройки компилятора, связанные со специализацией варпов, и всё сделать самостоятельно!
Если специалист готов ценой серьёзных усилий добиться высокой производительности, то вполне может вручную писать программы, в которых заложена специализация варпов и самостоятельно разбираться со сложностями, связанными с планированием и синхронизацией. Естественно, в таком случае придётся больше потрудиться над разработкой ядра, зато у вас получится программа, не столь зависимая от планирования инструкций на уровне компилятора или алгоритмов специализации варпов. Также можно попробовать написать реализацию без специализации варпов, например, показанное выше ядро GEMM, которое отличается высокой производительностью и без специализации варпов. Такая реализация также требует от разработчика ядра некоторых (но совершенно иных) усилий, в частности, переупорядочивать циклы для того, чтобы операции запускались и синхронизировались в правильном порядке. Такое переупорядочивание, в сущности «подсказывает» компилятору, как действовать. В сравнении с реализацией, написанной человеком и учитывающей специализацию варпов, реализация без специализации варпов может получиться более хрупкой. Другая проблема в том, что из-за неправильно подобранного размера тайлов может нарушиться тщательно спланированное статическое расписание. В свою очередь, реализация с динамическим планированием и специализацией варпов легко адаптируется к таким условиям.
Издержки, от которых зависит поле для манёвра в этом пространстве компромиссов, меняются вместе с возможностями архитектуры и также зависят от того, какое именно приложение вы разрабатываете. Например, для реализации GEMM под архитектуру Ampere характерны примерно такие же осложнения, как и для H100, связанные с асинхронностью и переменной длительностью задержек у команд загрузки. Но инженеры NVIDIA выяснили, что высокая производительность достижима ценой приемлемой сложности и без специализации варпов. Другой пример — реализация Flash-внимания для графического процессора H100. В таком приложении потоковый мультипроцессор должен выполнять массу работы, не связанной с ускорением (сокращение знаков после запятой и вычисление степеней), в то же время, выдавая работу ускорителю и блокируя его при асинхронных операциях. В то же время, как человеку, так и компилятору очень сложно подобрать статическое расписание, в рамках которого эти операции идеально перемежаются. Вот почему в большинстве реализаций Flash-внимания для H100 и более поздних применяется специализация варпов. Мне сложно сказать, насколько далеко можно зайти, развивая специализацию варпов в таком направлении, и как скоро эта работа окажется настолько трудоёмкой, что перестанет оправдывать затрачиваемые усилия. Как я упоминал выше, формулировать стратегии специализации варпов для эффективной работы ядер Blackwell по-настоящему сложно. Причём, становится запредельно сложно не только разрабатывать такие стратегии, но и выявлять их в корректном коде. В перспективе я усматриваю, как минимум, следующие возможности для развития специализации варпов (список не является исчерпывающим). 1) Возможно, программировать железо GPU станет проще, поэтому сократится и количество ситуаций, в которых нужна специализация варпов (маловероятно), 2) будут и далее совершенствоваться алгоритмы, применяемые для специализации варпов на уровне компилятора, в результате компилятор сравняется по этому показателю с человеком или 3) будет разработан системный софт, который устранит большинство мин, поджидающих программиста, рискующего писать код со специализацией варпов.
Благодарности
Этот текст вырос из бесед с коллегами и наставниками, среди которых: Рупаншу Сой, Фредрик Кьёльстад, Алекс Айкен, Майкл Гарланд, Майкл Бауэр и Дуэйн Меррилл. Майкл Бауэр и Рупаншу Сой сделали ценные замечания, позволившие улучшить эту статью. Спасибо Манье Бансал, вдохновившей меня собрать и изложить все эти мысли. Также благодарю Нэша Брауна, отловившего ряд ошибок в псевдокоде.