Проблема: узкое место пропускной способности памяти
Основная проблема оптимизации GPU-кода заключается в том, что вычислительная мощность современных графических процессоров настолько высока, что даже память с высокой пропускной способностью (HBM) не всегда успевает обеспечивать данными ядра. Классический подход, когда промежуточные результаты записываются в глобальную память и читаются следующим ядром, создает лишние циклы передачи данных.
Kernel Fusion решает эту задачу, объединяя несколько операций GPU в одно устройство (single device kernel). Это позволяет промежуточным результатам оставаться в регистрах, исключая необходимость их записи в глобальную память и отдельного запуска ядра.
Практический пример: суммирование модулей массива
Для демонстрации эффективности NVIDIA рассмотрела задачу вычисления суммы абсолютных значений массива: sum(abs(x)). Наивная реализация использует два отдельных ядра: одно для вычисления модуля (abs_kernel) с записью результата во временный буфер, и второе для редукции (sum_kernel).
При использовании новых интерфейсов CUDA 13.2, таких как cuda::std::span, cuda::launch и cuda::make_buffer, разработчики могут вручную слить эти ядра. В слитом варианте (sum_abs_kernel) операция fabsf вычисляется инлайн в регистре, а частичные суммы накапливаются без обращения к глобальной памяти до финального шага.
Результаты бенчмарков на RTX 4090
Сравнение производительности на базе NVIDIA GeForce RTX 4090 показывает значительное улучшение метрик при ручном слиянии ядер. Пропускная способность памяти остается на уровне ~850 GiB/s (90% от теоретического пика в 1,008 GB/s), что подтверждает: GPU работает на пределе физических возможностей, но с меньшим объемом лишней работы.
| Метрика | Наивная реализация (2 ядра) | Ручное слияние (1 ядро) |
|---|---|---|
| Количество запусков ядер | 2 | 1 |
| Промежуточный буфер | Да | Нет |
| Объем данных через глобальную память | 3 GB | 1 GB |
| Время выполнения | 3.51 мс (2.28 + 1.23) | 1.18 мс |
| Ускорение | — | 3x |
Автоматическое слияние через PyTorch и Triton
Помимо ручной оптимизации, NVIDIA подчеркивает важность автоматического слияния на уровне компилятора. Использование torch.compile в PyTorch позволяет компилятору Inductor генерировать оптимизированные ядра Triton. Даже если профилирование показывает два ядра (например, triton_red_fused_abs_sum_0 и 1), внутренняя логика Triton выполняет операцию abs в регистрах сразу после загрузки данных, избегая создания перэлементных промежуточных буферов.
Этот подход снижает нагрузку на разработчика, так как не требует написания низкоуровневого CUDA C++ кода, сохраняя при этом высокую эффективность за счет использования библиотек CUB, libcu++ и NVSHMEM.
Источник: NVIDIA dev blog ↗