Existen muchas formas de optimizar código para GPU. Esta guía técnica de NVIDIA se centra en una de las más efectivas: la fusión de kernels (kernel fusion), una técnica que combina varias operaciones de GPU en un único kernel de dispositivo para que los resultados intermedios no tengan que viajar de ida y vuelta por la memoria global ni requerir lanzamientos de kernel separados.

Un cuello de botella habitual al escribir código de GPU es que el cómputo es tan rápido que incluso la memoria de dispositivo de alto ancho de banda no logra mantener ocupado al kernel. La fusión de kernels ataca ese problema fusionando los cuerpos de los kernels para que los resultados intermedios queden en registros y se reduzcan las transferencias de memoria.

Conviene distinguirla de los CUDA Graphs, que ofrecen otro tipo de fusión pero en una capa distinta. Un grafo captura una secuencia de lanzamientos de kernel, copias de memoria y sincronizaciones en un único objeto reutilizable que el host puede despachar con una sola llamada. No fusiona los cuerpos de los kernels: estos siguen ejecutándose por separado y pasando resultados intermedios por la memoria global. Ambos enfoques son complementarios y pueden usarse juntos.

¿Qué resuelve la fusión de kernels?

Tomemos un ejemplo simple: sum(abs(x)). Leemos un arreglo, calculamos el valor absoluto de cada elemento y devolvemos la suma.

Una implementación ingenua usa dos kernels: uno calcula abs sobre un búfer temporal del mismo tamaño que la entrada, y otro reduce ese búfer a un único número. Funciona, pero es lento de una manera que revela el problema de fondo.

Código
template <typename Config>
__global__ void abs_kernel(Config config,
                           cuda::std::span<const float> in,
                           cuda::std::span<float> tmp)
{
    /* Base SIMT implementation using 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 into out[0] */
}

int main()
{
    /* device, stream, and buffers via cuda::make_buffer */

    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);
}

Los ejemplos en C++ de este artículo usan las interfaces más recientes del runtime NVIDIA CCCL, incluidas cuda::std::span, cuda::launch y cuda::make_buffer, que debutaron en CUDA 13.2. Para profundizar en estas primitivas, la referencia es CCCL Runtime: A Modern C++ Runtime for CUDA.

Un perfil de la implementación no fusionada en Nsight Systems se ve así:

Figura 1. Perfil en Nsight Systems de los kernels no fusionados para sum(abs(x)). Cada kernel se ejecuta de forma independiente y se necesita un resultado intermedio.
Figura 1. Perfil en Nsight Systems de los kernels no fusionados para sum(abs(x)). Cada kernel se ejecuta de forma independiente y se necesita un resultado intermedio.

Fusión manual del kernel

Lo primero que podemos hacer es reescribir esos dos kernels como un único sum_abs_kernel. Esto elimina el búfer intermedio, que existía solo para pasar datos entre kernels. Si ambas operaciones viven en el mismo kernel, el búfer desaparece y el rendimiento mejora.

Código
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() se calcula inline, sin buffer tmp.
    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]);
    }

    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);
    }
}

Las claves de esta versión:

  • fabsf(x[i]) se calcula inline, en un registro, en vez de leerse desde un búfer intermedio que otro kernel escribió.
  • Un bucle grid-stride permite que cada hilo acumule varios elementos en un único registro (thread_sum), de modo que las sumas parciales nunca tocan la memoria global.
  • cub::BlockReduce maneja la reducción por bloque en memoria compartida, y cuda::atomic_ref deja que un hilo por bloque sume su parcial al resultado global.

El perfil de la versión fusionada colapsa en consecuencia:

Figura 2. Perfil en Nsight Systems de un kernel fusionado manualmente. sum_abs_kernel es un solo kernel que hace toda la operación.
Figura 2. Perfil en Nsight Systems de un kernel fusionado manualmente. sum_abs_kernel es un solo kernel que hace toda la operación.

Estas cifras corresponden a una NVIDIA GeForce RTX 4090. El ancho de banda de memoria se sitúa en unos 850 GiB/s, cerca del 90% de su máximo teórico de 1.008 GB/s. Estamos limitados por ancho de banda: la GPU hace tanto trabajo de memoria como la física permite. La versión fusionada es más rápida porque hay menos trabajo de memoria que hacer, ya que mueve 1 GiB por la memoria global en vez de 3 GiB.

Un kernel escrito a mano en CUDA C++ como este es a veces la solución correcta, porque entrega control total y se acerca lo más posible al rendimiento máximo. La desventaja es el esfuerzo de escribirlo, el alto costo de mantenimiento y la portabilidad limitada entre arquitecturas de GPU. Bibliotecas del lado del dispositivo como CUB, libcu++, nvmath.device y NVSHMEM pueden simplificar la fusión manual aportando primitivas reutilizables para FFT, multiplicación de matrices, reducciones, escaneos, ordenamientos y comunicación.

Fusión implícita con torch.compile

Escribir tu propio kernel funciona. En este caso, la operación es esencialmente "aplicar abs mientras se reduce". En vez de expresar eso a mano en CUDA C++, podemos describir el cómputo en Python plano y dejar que el compilador genere el kernel. Eso es lo que hace torch.compile: le das una función, traza lo que hacen los tensores y Torch Inductor, el compilador de PyTorch, genera un kernel que corre en la GPU.

Código
import torch

x = torch.randn(N, device="cuda")

def sum_abs(t):
    return t.abs().sum()

compiled = torch.compile(sum_abs)
result = compiled(x)

En la versión eager (sin torch.compile) obtenemos lo esperado: dos kernels, uno elemento a elemento (abs) y otro de reducción (sum). En la versión compilada también obtenemos dos kernels: triton_red_fused_abs_sum_0 y triton_red_fused_abs_sum_1, porque eso es lo que genera el compilador Inductor.

Figura 3. Perfil de PyTorch en Perfetto.dev para una operación eager. Se ven dos kernels, uno por operación.
Figura 3. Perfil de PyTorch en Perfetto.dev para una operación eager. Se ven dos kernels, uno por operación.

Para entender qué hizo Inductor, se pueden inspeccionar los kernels Triton generados con TORCH_LOGS="output_code". La parte interesante del primer kernel:

Código
tmp0 = tl.load(in_ptr0 + ...)   # carga un trozo de la entrada
tmp1 = tl_math.abs(tmp0)        # aplica abs en registros, inline
tmp2 = _tmp3 + tmp1             # acumula en la suma parcial

El abs ocurre en un registro, justo después de la carga, dentro del kernel de reducción. No hay un kernel abs separado ni un búfer intermedio por elemento. Un segundo kernel, corriendo en un solo bloque, lee los 2.048 resultados parciales y los reduce a un escalar final. Hay más detalle en este blog de PyTorch.

Figura 4. Perfil de la versión compilada en Perfetto.dev. Los dos kernels los genera PyTorch aunque la operación esté fusionada por el compilador.
Figura 4. Perfil de la versión compilada en Perfetto.dev. Los dos kernels los genera PyTorch aunque la operación esté fusionada por el compilador.

Ambas versiones mueven 1 GiB por la memoria global en vez de 3 GiB. El compilador eligió una estrategia razonable, pero no la misma que habríamos elegido a mano. Esto es bueno y malo a la vez: el compilador nos da un kernel fusionado gratis, pero cedemos algo de previsibilidad. Un dtype distinto, una forma distinta o una pequeña operación añadida en el medio pueden producir un conjunto de kernels completamente diferente. A veces obtienes un kernel, a veces cuatro, y a veces una fusión con la que contabas deja de ocurrir sin aviso.

Fusión explícita con cuda.compute

¿Y si queremos la productividad de Python con la previsibilidad de un kernel escrito a mano? Para eso está cuda.compute. Expone algoritmos de dispositivo como reduce, scan, sort y transform, junto con iteradores que se pueden componer para controlar cómo esos algoritmos leen su entrada. No escribes un kernel directamente: compones un cómputo y lo pasas a un algoritmo, y la fusión ocurre porque la pediste explícitamente.

Código
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)

abs_it = TransformIterator(x, lambda a: abs(a))
reduce_into(
    d_in=abs_it, d_out=out, num_items=N,
    op=OpKind.PLUS, h_init=h_init,
    determinism=Determinism.NOT_GUARANTEED,
)

TransformIterator no reserva memoria ni lanza un kernel. Es una vista perezosa que dice: "cuando alguien me pida el elemento i, leo x[i] y le aplico abs". Luego reduce_into realiza una única reducción de dispositivo, tirando de los elementos a través del iterador. El abs corre inline, en registros, a medida que cada elemento se carga desde la memoria global, sin ningún búfer intermedio por elemento. Además, cuda.compute funciona con tensores de PyTorch y admite entradas que implementen DLPack o la CUDA Array API, como los arreglos de CuPy.

Una vez que separas "qué son los datos" de "cómo recorrerlos", producir el elemento i se vuelve un cómputo inline arbitrario: una carga, una transformación, un gather o el desempaquetado de una estructura. El algoritmo enhebra esas operaciones en una sola pasada sin almacenamiento intermedio. Esta solución es determinista, portable y componible, ideal para autores de bibliotecas y aplicaciones de alto rendimiento.

¿Cuál conviene usar?

Los tres caminos llegan al mismo destino (mover 1 GiB en vez de 3 GiB), pero con distintas concesiones:

  • CUDA C++ manual: máximo control y rendimiento, a costa de esfuerzo, mantenimiento y portabilidad limitada.
  • torch.compile: fusión "gratis" desde Python, a costa de previsibilidad sobre qué kernels se generan.
  • cuda.compute: productividad de Python con fusión explícita y determinista, pensada para bibliotecas.

La lección de fondo se sostiene en cualquiera de los tres: en cargas limitadas por memoria, cada viaje evitado a la memoria global se traduce directamente en velocidad.