В мире высокопроизводительных вычислений на графических процессорах (GPU) существует одна фундаментальная проблема, с которой сталкиваются разработчики: вычислительная мощность современных GPU растет гораздо быстрее, чем пропускная способность памяти. Даже самая быстрая видеопамять HBM может стать узким местом, если ядро (kernel) выполняется слишком быстро, ожидая данных. Это явление известно как «memory wall» (стена памяти). Когда вы пишете код для GPU, часто возникает ситуация, когда промежуточные результаты вычислений выгружаются в глобальную память, а затем считываются обратно для следующей операции. Это не только тратит драгоценную пропускную способность, но и требует дополнительных запусков ядер, что добавляет накладные расходы со стороны хоста.
Решением этой проблемы является техника, известная как Kernel Fusion (слияние ядер). Суть метода проста: вместо того чтобы запускать несколько отдельных ядер, которые пишут и читают промежуточные данные из глобальной памяти, мы объединяем их логику в одно большое ядро. Благодаря этому промежуточные результаты остаются в регистрах или разделяемой памяти (shared memory) внутри блока, а доступ к глобальной памяти происходит только для чтения исходных данных и записи финального результата. В этой статье мы подробно разберем, как работает Kernel Fusion, какие существуют подходы к его реализации — от ручной оптимизации на C++ до автоматического слияния в PyTorch и явного управления через CUDA Compute, и как это влияет на производительность ваших приложений.
01Почему Kernel Fusion важен: проблема накладных расходов и трафика памяти
Чтобы понять ценность слияния ядер, нужно взглянуть на типичный сценарий оптимизации. Часто разработчики пытаются ускорить запуск ядер с помощью CUDA Graphs. CUDA Graph — это механизм, который захватывает последовательность запусков ядер, копирования памяти и синхронизаций, превращая их в один переиспользуемый объект. Это позволяет отправить всю последовательность на GPU одним вызовом, что значительно снижает задержку на стороне хоста (CPU).
Однако важно понимать: CUDA Graph не сливает тела ядер. Ядра внутри графа все равно выполняются отдельно, и промежуточные данные по-прежнему проходят через глобальную память. Если у вас есть операция, требующая записи промежуточного результата в память и его последующего чтения, обертывание этого процесса в граф сэкономит микросекунды на стороне CPU, но не решит проблему трафика памяти. Два подхода — Kernel Fusion и CUDA Graphs — являются комплементарными. Fusion оптимизирует работу с памятью и вычислениями внутри GPU, а Graphs оптимизируют управление этими вычислениями со стороны CPU. Для максимальной производительности их часто используют вместе.
Простой пример: сумма модулей массива
Давайте рассмотрим классический пример, который наглядно демонстрирует проблему. Представьте себе операцию sum(abs(x)). Нам нужно прочитать массив чисел, вычислить модуль каждого элемента и затем просуммировать все полученные значения. Наивная реализация этой задачи в CUDA обычно состоит из двух шагов:
- Ядро вычисления модуля (abs_kernel): Читает массив
x, вычисляет абсолютное значение каждого элемента и записывает результат во временный буфер (intermediate buffer) того же размера, что и входные данные. - Ядро суммирования (sum_kernel): Читает этот временный буфер, выполняет редукцию (суммирование) всех элементов и записывает итоговое число в выходной буфер.
Этот подход работает корректно, но он неэффективен. Мы тратим пропускную способность памяти на запись промежуточного результата, который сразу же читается обратно. Кроме того, мы дважды запускаем ядро, что добавляет накладные расходы. Давайте посмотрим, как это выглядит в коде на современном C++ с использованием интерфейсов NVIDIA CCCL (CUDA C++ Core Libraries), которые стали стандартом в CUDA 13.2+.
template <typename Config>
__global__ void abs_kernel(Config config,
cuda::std::span<const float> in,
cuda::std::span<float> tmp)
/* Базовая реализация SIMT с использованием fsabs */
template <typename Config>
__global__ void sum_kernel(Config config,
cuda::std::span<const float> tmp,
cuda::std::span<float> out)
/* CUB block-reduce + atomic add в out[0] */
int main()
/* Настройка устройства, потока и буферов через cuda::make_buffer
* См. пост CCCL Runtime для полного примера настройки. */
auto config = cuda::distribute<BLOCK_THREADS>(in.size());
cuda::launch(stream, config, abs_kernel<decltype(config)>, in, tmp);
cuda::launch(stream, config, sum_kernel<decltype(config)>, tmp, out);В этом коде мы явно разделяем логику. abs_kernel пишет в tmp, а sum_kernel читает из tmp. Если мы профилируем этот код в Nsight Systems, мы увидим два отдельных блока выполнения, разделенных операцией записи/чтения в глобальную память.
02Ручное слияние ядер (Manual Kernel Fusion)
Первый и самый прямой способ оптимизации — это ручное объединение двух ядер в одно. Мы пишем новое ядро, которое выполняет обе операции: вычисление модуля и суммирование. Ключевое изменение здесь заключается в том, что временный буфер tmp больше не нужен. Он исчезает, так как данные обрабатываются «на лету».
Как это работает технически
В новом ядре sum_abs_kernel мы используем так называемый «grid-stride loop» (цикл с шагом сетки). Это позволяет каждому потоку обрабатывать несколько элементов массива. Внутри цикла мы вычисляем fabsf(x[i]) и сразу же добавляем результат к локальной переменной thread_sum, которая хранится в регистре потока. Таким образом, промежуточные суммы никогда не записываются в глобальную память.
После того как поток обработал все свои элементы, мы выполняем редукцию на уровне блока (block-level reduction) с использованием библиотеки CUB. Результат сохраняется в разделяемой памяти (shared memory). Наконец, один поток из каждого блока атомарно добавляет свой частичный результат к глобальному результату out[0].
template <typename Config>
__global__ void sum_abs_kernel(Config config,
cuda::std::span<const float> x,
cuda::std::span<float> out)
using BlockReduce = cub::BlockReduce<float, BLOCK_THREADS>;
__shared__ typename BlockReduce::TempStorage temp_storage;
// Grid-stride loop. abs() вычисляется inline, без временного буфера.
float thread_sum = 0.0f;
const auto tid = cuda::gpu_thread.rank(cuda::grid, config);
const auto stride = cuda::gpu_thread.count(cuda::grid, config);
for (size_t i = tid; i < x.size(); i += stride) {
thread_sum += fabsf(x[i]);
}
// Редукция на уровне блока в shared memory через CUB.
float block_sum = BlockReduce(temp_storage).Sum(thread_sum);
// Один поток на блок атомарно накапливает результат в глобальной памяти.
if (threadIdx.x == 0) {
cuda::atomic_ref<float, cuda::thread_scope_device> r(out[0]);
r.fetch_add(block_sum, cuda::memory_order_relaxed);
}
int main()
/* Настройка устройства, потока и буферов */
// Фиксированная сетка из 1024 блоков для паттерна grid-stride.
auto config = cuda::make_config(cuda::make_hierarchy(
cuda::grid_dims(1024),
cuda::block_dims<BLOCK_THREADS>()));
cuda::launch(stream, config, sum_abs_kernel<decltype(config)>, x, out);Результаты производительности
Давайте сравним метрики производительности на примере NVIDIA GeForce RTX 4090. В таблице ниже приведены результаты сравнения наивной реализации и ручной оптимизации.
Как видно из данных, ручное слияние сокращает объем данных, передаваемых через глобальную память, с 3 ГБ до 1 ГБ. Это происходит потому, что мы убрали запись промежуточного буфера. Время выполнения упало с 3.51 мс до 1.18 мс, что дает ускорение в 3 раза. При этом эффективная пропускная способность памяти остается на уровне ~850 ГБ/с, что составляет около 90% от теоретического пика видеокарты. Это означает, что мы достигли предела, ограниченного пропускной способностью памяти (bandwidth-bound), и дальнейшее ускорение возможно только за счет уменьшения объема передаваемых данных, а не увеличения вычислительной мощности.
03Имплицитное слияние через компиляторы (torch.compile)
Написание собственных CUDA-ядер — это мощный инструмент, но он не всегда практичен. Поддержка такого кода сложна, а время разработки велико. Здесь на помощь приходят компиляторы. Фреймворк PyTorch предлагает инструмент torch.compile, который использует компилятор Torch Inductor для автоматической генерации оптимизированного кода.
Когда вы оборачиваете функцию в torch.compile, компилятор анализирует граф вычислений. Он видит, что операция abs сразу же передается в sum, и решает объединить их в одно ядро. Это называется имплицитным (скрытым) слиянием.
import torch
x = torch.randn(N, device="cuda")
def sum_abs(t):
return t.abs().sum()
compiled = torch.compile(sum_abs)
result = compiled(x)Что происходит под капотом?
Если вы посмотрите на профилирование скомпилированного кода, вы можете увидеть два ядра, например, triton_red_fused_abs_sum_0 и triton_red_fused_abs_sum_1. Это может сбить с толку, так как кажется, что слияния не произошло. Однако на самом деле компилятор сгенерировал ядро на языке Triton, которое выполняет слияние логики. Операция abs происходит в регистрах сразу после загрузки элемента, без создания промежуточного буфера.
Компилятор может разделить редукцию на два этапа: первое ядро снижает входные данные до набора частичных сумм по блокам, а второе ядро суммирует эти частичные суммы. Обе версии (ручная и компиляторная) перемещают 1 ГБ данных через глобальную память, обеспечивая аналогичное ускорение. Однако компилятор не гарантирует, что стратегия слияния останется неизменной. Изменение типа данных, формы тензора или добавление небольшой операции может привести к тому, что компилятор решит не сливать ядра или создаст совершенно другую структуру. Это дает свободу, но лишает предсказуемости.
04Эксплицитное слияние: cuda.compute
Что если мы хотим получить предсказуемость ручного написания кода, но с удобством Python? Для этого NVIDIA представила cuda.compute. Этот модуль предоставляет высокоуровневые алгоритмы, такие как reduce, scan, sort и transform, а также итераторы, которые позволяют явно контролировать, как данные читаются и преобразуются.
В cuda.compute вы не пишете ядро напрямую. Вместо этого вы описываете композицию вычислений. Например, вы создаете итератор, который говорит: «Когда кто-то запросит элемент, прочитай его из массива и примени к нему функцию abs». Затем вы передаете этот итератор в функцию редукции. Слияние происходит потому, что вы явно запросили его через композицию.
import torch
import numpy as np
from cuda.compute import (
Determinism, OpKind, TransformIterator, reduce_into,
)
x = torch.randn(N, device="cuda")
out = torch.empty(1, dtype=torch.float32, device="cuda")
h_init = np.array([0.0], dtype=np.float32)
# TransformIterator не выделяет память и не запускает ядро.
# Это ленивое представление, которое применяет abs «на лету».
abs_it = TransformIterator(x, lambda a: abs(a))
# reduce_into выполняет однопоточную редукцию, используя итератор.
reduce_into(
d_in=abs_it, d_out=out, num_items=N,
op=OpKind.PLUS, h_init=h_init,
determinism=Determinism.NOT_GUARANTEED,
)Преимущества cuda.compute
Ключевое преимущество cuda.compute заключается в детерминизме и предсказуемости. В отличие от компилятора, который может изменить стратегию оптимизации, здесь вы явно решаете, что должно быть слито. TransformIterator не выделяет память и не запускает ядро — это просто ленивое представление. Фактическое слияние происходит внутри алгоритма reduce_into, который использует оптимизированную библиотеку CUB под капотом.
Это дает вам лучшее из двух миров: эргономику Python и производительность, сравнимую с ручным написанием CUDA C++. Код становится портативным, композируемым и идеальным для авторов библиотек. Кроме того, cuda.compute поддерживает PyTorch Tensors, CuPy и другие библиотеки, реализующие CUDA Array API, что делает его универсальным инструментом для оптимизации пайплайнов данных.
05Сравнение подходов: Кто решает, что сливать?
Давайте обобщим три рассмотренных метода слияния ядер. Все они приводят к одному и тому же результату: снижению трафика памяти до 1 ГБ и ускорению примерно в 3 раза по сравнению с наивным подходом. Однако пути к этому результату различаются.
- Ручное слияние (Manual Fusion): Вы пишете ядро. Полный контроль, максимальная производительность, но высокая стоимость разработки и поддержки. Предсказуемость 100%.
- Компилятор (torch.compile): Компилятор решает за вас. Быстро, удобно, но непредсказуемо. Стратегия слияния может меняться в зависимости от входных данных и версии компилятора.
- cuda.compute: Вы определяете композицию. Предсказуемо, как при ручном написании, но с удобством Python. Использует проверенные библиотеки (CUB) под капотом, обеспечивая высокую производительность без написания низкоуровневого кода.
06Что это значит на практике
Для разработчиков, работающих с AI-инференсом и обучением моделей, понимание Kernel Fusion критически важно. Если вы используете PyTorch или другие фреймворки, torch.compile должен быть вашим первым шагом к оптимизации. Он часто дает значительный прирост производительности «из коробки» без изменения логики модели. Однако, если вы сталкиваетесь с узкими местами, где компилятор не может оптимизировать специфическую операцию, или если вам нужна строгая детерминированность, стоит рассмотреть cuda.compute или ручную оптимизацию.
Для создателей библиотек и высокопроизводительных приложений cuda.compute открывает новые возможности. Он позволяет писать чистый, читаемый код на Python, который компилируется в эффективные CUDA-ядра. Это снижает порог входа для использования передовых техник оптимизации памяти. Помните, что Kernel Fusion — это не серебряная пуля для всех проблем, но это один из самых эффективных инструментов для борьбы с ограничениями пропускной способности памяти, которые остаются главным узким местом в современных GPU.
Экспериментируйте с профилированием в Nsight Systems, анализируйте трафик памяти и выбирайте тот уровень абстракции, который подходит именно для вашей задачи. От простых скриптов на Python до сложных кастомных ядер — оптимизация памяти через слияние остается ключом к раскрытию полного потенциала GPU.
Источник: NVIDIA Developer ↗