Каждый раз, когда открывается лента Instagram, на тысячах GPU стартует один и тот же конвейер: модель вытаскивает строки из миллионов embedding-таблиц, собирает из них признаки и решает, какой пост показать. Эту работу годами выполняли рукописные CUDA-ядра, за которые отвечала горстка экспертов. Команда PyTorch вместе с инфраструктурной командой рекомендаций Meta переписала ключевые операторы на Triton, и результат получился редкий: новый код не догнал рукописный CUDA, а обогнал его.

Что такое TBE

Table Batched Embedding (TBE) это оператор, который за один запуск на GPU выполняет поиск строк и пулинг сразу по множеству embedding-таблиц. Он обслуживает обучение и инференс рекомендательных моделей, а сами таблицы разложены по шардам тысяч GPU.

Каждая таблица это матрица E×D: E строк-эмбеддингов по D чисел в каждой. Разреженный признак это список идентификаторов вида [id1, id2, id3, id4], и forward-проход суммирует строки, которые им соответствуют, выдавая одно D-мерное представление на каждый запрос. Backward делает зеркальную работу: для каждой уникальной пары «таблица, строка», которую затронул батч, он суммирует градиенты со всех позиций и применяет ровно одно обновление оптимизатора.

Масштаб виден в арифметике. В одной из конфигураций 4,2 миллиона таблиц и батч из 128 тысяч запросов дают 83 миллиона обращений к строкам, три миллиона уникальных строк и отдельную строку, к которой за один шаг обратились 125 тысяч раз. Длины серий в одном батче расходятся от единицы до миллионов, и это расхождение определяет всю инженерную сложность оператора.

Почему эти ядра так сложно писать

У forward-прохода предсказуемая структура: собрать строки, умножить на веса, сложить в FP32, записать результат. У backward тензорных ядер нет вообще. Вся работа там это движение данных и редукция: для каждой уникальной строки нужно просуммировать градиенты из всех мест батча, которые её затронули, и только после этого обновить строку. Пропускная способность памяти здесь важнее арифметической плотности.

Дальше начинается неравномерность. Одна строка может собрать миллионы обращений, а миллион других по одному. CUDA-код решает это эскалацией: при длине серии (segment length, SL) от 32 он переключается на cooperative-ядро с одним CTA на строку. Хитрость в том, что синхронизация целого CTA стоит дорого, и окупается она только когда на строку приходит достаточно работы. Ниже порога ядро тратит на метаданные строки больше, чем на саму работу, а выше начинает выигрывать.

Triton в этом месте пошёл другим путём: вместо кооперации он тянет простой стриминг до SL = 256. Именно поэтому полоса длин между 32 и 256 стала зоной, где старый и новый код расходятся сильнее всего.

Forward: два ядра на разные формы данных

Forward реализован двумя ядрами. Универсальный gather-путь разбивает работу по формуле ceil(B / BAGS_PER_PROGRAM): каждая программа обслуживает несколько bags и прогоняет по ним таблицы циклом, вместо того чтобы запускать вычислительную сетку B×T. Внутри цикла выдаются четыре независимые загрузки строки, а на настроенном двухбэговом пути восемь. Индексы и смещения переключаются на int32, когда линеаризованный диапазон помещается в 2^31: это вдвое сокращает хранение индексов и ширину ключа radix-sort, которую использует CUB. Накопление идёт в FP32, а при FP32-весах в FP64, чтобы большие D не теряли точность.

Второй путь узкоспециализирован. Он включается для таблицы с E ≤ 64, размерностью D между 64 и 128, пулингом L ≥ 64, FP16-весами, без пер-сэмпловых весов и без variable-batched режима. Одна программа обрабатывает 16 bags, строит гистограмму по первым 256 индексам и вычисляет counts × table через tl.dot, а хвост индексов добирает скалярным циклом. Все прочие формы остаются универсальному ядру.

Отдельная забота это валидация входов. Обновлённый шаг проверки границ выполняется до Triton-прохода и на B200 ускоряет эту часть до 1,24×. Опция fused_bounds_check идёт дальше: подходящее ядро само валидирует и чинит невалидные входные тензоры, тогда как проверка смещений остаётся отдельным маленьким ядром. Веса, variable-batched режимы и AMD-конфигурации падают в стандартный путь валидации. Через TorchRec опция выключена по умолчанию.

Есть и переиспользование состояния между проходами. Режим exact row-wise Adagrad умеет сохранить гистограмму, построенную в forward, а backward использует эти counts для компенсированного GEMM из FP16 в FP32 и затем применяет оптимизатор. Ещё один опциональный путь переносит транспонирование индексов, сортировку и run-length encoding в forward и возвращает метаданные через autograd. На большой B200-конфигурации это выглядит так: forward подрастает с 22,844 до 33,252 мс, зато backward падает с 56,693 до 32,931 мс, а суммарная латентность снижается с 79,537 до 66,183 мс, то есть на 16,8%. Подготовку данных просто убрали с критического пути обучения.

Backward: три маршрута по длине серии

Backward начинается с операции transpose_embedding_input, которая разворачивает батч в серии: уникальная строка плюс сэмплы, которые её тронули. Эту работу подняли в forward, чтобы снять нагрузку с обратного прохода. Дальше всё решает длина серии, и по ней серии раскладываются по трём маршрутам.

Короткие серии (SL < 256) целиком отдаются одной программе: она собирает строку, копит сумму в регистрах, применяет оптимизатор и делает обычный store. Эксклюзивное владение означает, что атомики не нужны вообще, и это самый дешёвый случай.

Длинные серии (SL ≥ 256) нарезаются на куски по 256 lookups в стиле split-K: частичные суммы ложатся в workspace, а второе ядро применяет оптимизатор. На Blackwell с очень большими батчами тот же сплит работает за один запуск: device-scope fence делает все частичные суммы видимыми, счётчик декрементируется, и последний сабпрограмм, увидевший remaining == 1, применяет обновление на месте.

Порядок здесь не деталь, а сам аргумент корректности. Без fence какой-нибудь сабпрограмм мог выиграть обратный отсчёт, пока чужая атомарная запись ещё в полёте, и тогда в сумму попал бы неполный градиент. Вместо ускорения получилась бы тихая численная ошибка, которую потом ищи неделями.

Слияние частичных сумм тоже было болью. Каждый сабпрограмм атомарно добавляет целую строку в слот workspace, а при восьми тысячах кусков это восемь тысяч претендентов на одну строку, и tl.atomic_add выполняет такой merge поэлементно. Blackwell даёт инструкцию лучше: cp.reduce.async.bulk.tensor, доступная из Triton как tlx.async_descriptor_store со store_reduce="add". Редукция на уровне TMA вместо соревнования атомиков.

tl.atomic_add(temp_grad_buffer_ptr + temp_grad_offset + col_offsets, grad, mask=mask)
tlx.fence("gpu")
remaining = tl.atomic_add(grad_accum_counter_ptr + grad_buffer_id, -1)
if remaining == 1:
    ...  # последний сабпрограмм применяет оптимизатор и сохраняет

Грабли: регистровый обрыв и константы

Самая полезная для практиков часть отчёта это список конкретных проблем, которые пришлось решать по ходу портирования.

Первая: неравномерность программ. Программа, которой достался run из двух миллионов lookups, держит GPU занятым в одиночку, а хвост из одиночных обращений платит полную цену за почти нулевую работу. Решение: длинные серии нарезаются на фиксированные куски по 256, и двухмиллионный run превращается примерно в восемь тысяч сабпрограмм, которыми машина загружается целиком. На Blackwell дополнительно включается CLC, аппаратный планировщик: ядро стартует как persistent и «ворует» свободные run_id после завершения текущего. Софтовые схемы вроде сортировки run_id по частоте обращения больше не нужны.

Вторая проблема это регистровый обрыв. Каждая буферизованная строка держит в живых BLOCK_SIZE 64-битных адресов, потому что выражение вида dout_row_start_ptr[:, None] + col_offsets[None, :] материализует целую плитку указателей. Цена растёт как ширина × BLOCK_SIZE, и ширина, подобранная под одну форму, ломается на другой. После подбора ширины под цель получились цифры, которые говорят сами за себя. На коротких невзвешенных сериях переход с ширины 8 на 2 снизил число регистров с 184 до 64, поднял occupancy с 12,5% до 49,9% и перевёл относительную производительность с 0,41 на 1,05. На длинных сериях с накоплением регистры снизились со 158 до 62, occupancy выросла с 17,6% до 44,7%, а fleet parity с 82% до 87%. На взвешенных сериях: регистры со 125 до 64, occupancy с 24,8% до 49,3%, а parity поднялась с 51% до 69%.

Третья: BLOCK_SIZE это константа времени компиляции. Один запуск обязан выбрать размер по самой широкой таблице в наборе, поэтому узкие таблицы гоняют маскированные впустую линии. Решение: бакетирование по next_pow2(D), каждый бакет запускается со своим BLOCK_SIZE, и всё это складывается в ядро классификации без дополнительного прохода. Профилировщик подтверждает: выигрыш идёт от сокращения пустых линий, а не от occupancy.

Четвёртая проблема прячется в том, какие серии считать длинными. Ответ зависит от данных и меняется от батча к батчу, а выясняется он только на GPU. Чтение счётчиков обратно через .item() это cudaStreamSynchronize на каждом шаге, то есть остановка конвейера. Решение: держать формы на устройстве. Workspace преаллоцируется по границе, вычислимой из числа индексов, а классификация идёт в ядре через атомарные счётчики и stream compaction: один атомик на блок вместо атомика на элемент, смещения внутри блока из prefix sum, а дальше ядра сами распределяют работу через while-циклы.

Две последние грабли относятся к split-сериям и закрылись device-scope fence с обратным отсчётом плюс TMA-редукцией, о которых шла речь выше. Ценность всего списка в том, что каждый из этих фиксов в Triton это значение конфига, а не переписывание архитектуры ядра.

Результаты: 307 конфигураций на GB200

Замер шёл на 307 конфигурациях шардов, из которых 283 уникальные формы, на GB200, с exact row-wise Adagrad и FP16-весами. Метрика это отношение времени Triton к времени CUDA TBE. Медианный forward-прирост составил 1,28×, то есть 28% быстрее.

Куда интереснее, где именно берётся выигрыш. В полосе длин серий между 32 и 256 Triton идёт быстрее в 4,3 раза, передвигая при этом примерно тот же объём DRAM-байтов (0,91× от CUDA) и выдавая на 29% больше load-запросов. Nsight Compute показывает разницу без прикрас: CUDA-ядро на этой форме выжимает 678 ГБ/с против 3948 ГБ/с у Triton при одинаковой occupancy, это 5,8×. Причина в том, что CUDA включает CTA-кооперацию уже на SL = 32, когда амортизировать синхронизацию ещё нечем, а Triton тянет простой стриминг до 256.

Есть и места, где Triton проигрывает, и команда их не скрывает. Формы, вся работа которых сидит в сериях короче четырёх, уходят ниже паритета из-за чистой latency зависимых загрузок: CUDA с warp-per-row лучше амортизирует метаданные на коротких сериях. Это 11 из 307 шардов, каждый меньше миллисекунды. Формы с длинами выше 256 держатся на паритете, без отрыва. Но полоса между 32 и 256, где живёт основная масса реальных нагрузок, остаётся за Triton. Причём все победы в этом диапазоне одержаны значениями конфигов, а не новыми алгоритмами.

Что это меняет

Самый долгоиграющий результат FBTriton даже не в процентах. Весь разреженный путь, forward и backward, теперь написан на обычном Python и занимает меньше кода, чем одни только CUDA-шаблоны, которые он заменил. Для инженеров это означает скорость разработки: новые SOTA-алгоритмы эмбеддингов внедряются правкой конфигов, а не реструктуризацией шаблонной машинерии.

Второе следствие это переносимость. Тела ядер одинаковы на Blackwell, Hopper и AMD: специфика платформы (CLC, TMA-редукции, device-scope fences) подключается аддитивными флагами, а не форками кода. Для мира с гетерогенным железом это не удобство, а способ выживания бэкенд-команды.

Третье это заготовка под «мега-спарс-ядро». Раз всё на Triton, открываются фьюжены forward с backward и пролога с эпилогом: оптимизатор становится эпилогом backward-прохода и убирает лишние проходы по памяти, а расширения вроде FP8 momentum scaling превращаются в несколько строк накопления вместо отдельной подсистемы.

И общий тренд здесь важнее частного кейса. Triton из экспериментального DSL превратился в инструмент, которым большая компания заменяет рукописный CUDA в самой нагруженной части рекомендательных систем, обслуживающих миллиарды запросов. CUDA не исчезает: на сверхкоротких сериях она всё ещё держит оборону. Но границы владений сдвигаются, и в этот раз переезд санкционировал сам PyTorch.

Вопросы и ответы

Что такое TBE простыми словами?

TBE это оператор, который за один запуск на GPU находит строки в множестве embedding-таблиц и суммирует их в компактные представления. Он нужен рекомендательным системам, где профиль пользователя и сигналы о кандидатах разложены по тысячам таблиц, а каждый запрос требует тысяч обращений к строкам.

Почему Triton обогнал CUDA, если CUDA годами считалась оптимумом?

Выигрыш дала не магия компилятора, а гибкость настройки. Каждый инженерный компромисс (порог эскалации, ширина gather, размер блока) в Triton это значение конфига, а в шаблонной CUDA переписывание архитектуры ядра. Плюс Triton избегает кооперации там, где она не окупается: в полосе серий между 32 и 256 у него просто нет накладных расходов, которые CUDA платит на синхронизации CTA.

Применимо ли это вне Meta?

Да. TBE это открытый компонент библиотеки TorchRec, а Triton доступен всем. Подход напрямую касается любой системы с большими разреженными эмбеддингами: от рекомендаций до поисковых ранжировщиков и лент новостей. Если у вас есть собственные ядра для lookup и pooling, отчёт стоит прочитать целиком: скорее всего, ваши узкие места выглядят похоже.

Итог

FBTriton это редкий случай, когда переписывание работающего низкоуровневого кода дало одновременно скорость и простоту: медианный прирост 28% на forward, до 4,3× в полосе длин серий между 32 и 256, меньше регистров, выше occupancy и код на Python вместо шаблонного CUDA. Если вы поддерживаете собственные ядра для embedding-операций, посмотрите на этот отчёт как на карту миграции.

А если хотите разобраться, как Triton вообще устроен изнутри, начните с разбора Plugin Extensions: это хорошая точка входа в экосистему, которая уже начала теснить CUDA.