PyTorch Foundation описала ускорение embedding-ядер на Triton

PyTorch Foundation опубликовала разбор FBTriton — набора ядер Triton для forward- и backward-операций Table Batched Embedding. TBE выполняет поиск embedding-векторов и pooling для нескольких таблиц в одном запуске GPU, что предназначено для нагрузок рекомендательных систем. По утверждению авторов материала, новая реализация превосходит прежние CUDA-ядра на рассматриваемых сценариях.

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

В качестве примера PyTorch Foundation приводит конфигурацию на B200. В ней время forward выросло с 22.844 до 33.252 мс, а backward сократилось с 56.693 до 32.931 мс. Совокупная задержка снизилась с 79.537 до 66.183 мс, что соответствует изменению на −16.8%. Отдельно проверка границ дала ускорение до 1.24 раза на рассматриваемых нагрузках.

Авторы отмечают, что backward не использует tensor cores и испытывает значительные затраты на перемещение и редукцию данных. Архитектура FBTriton пытается компенсировать это распределением работы по характеристикам сегментов и переносом части предварительной обработки в forward-проход.