Vibecoding GPU Kernels

Escribir kernels de GPU a mano requiere paciencia, esfuerzo y también tu cordura. He estado experimentando con LLMs modernos (Claude Opus 5, GPT 5.6 Sol) para que literalmente me impriman los kernels. Voy a hablar de lo que aprendí haciendo eso.

Sinceramente, escribir kernels de GPU es una tarea vergonzosamente verificable.

Primero, decides qué operación(es) quieres implementar con el kernel. Segundo, escribes la primera versión y te aseguras de que compile sin errores evidentes. Tercero, verificas la corrección con una implementación de referencia lenta. Si no coincide, intentas corregir el problema de corrección. Una vez hecho eso, mides el tiempo de ejecución del kernel con benchmarks. Las versiones posteriores se construyen sobre esta base, y sigues optimizando hasta alcanzar las métricas de roofline o hasta que quedes satisfecho.

Bucle verificable: desarrollo de kernels de GPU

La imagen de arriba lo muestra como un bucle verificable:

• La verificación de compilación es un bucle local cerrado (B <-> C) antes de que siquiera pienses en la corrección.
• La verificación de corrección (D <-> E <-> F) es el núcleo del bucle de recompensa verificable. Esta es la parte que hace que escribir kernels sea un buen problema para la verificación automatizada, ya que tienes una referencia de verdad absoluta contra la cual comparar.
• La optimización (G -> H -> regreso a D) reutiliza el mismo bucle de corrección para cada nueva versión, ya que un kernel rápido pero incorrecto no vale nada; la corrección debe verificarse en cada versión.
• La verificación de roofline/satisfacción (I) es el bucle externo que decide si seguir optimizando o detenerse.

Primera versión

La imagen de arriba sigue siendo una vista de alto nivel, y el diablo está en los detalles. Necesitamos asegurarnos de que nuestro agente LLM tenga todo el contexto requerido para siquiera empezar a escribir una buena primera versión.

Los DSL de CUDA entran en acción aquí. Triton, CuTeDSL y Tilelang son los más fáciles para empezar, en Python. La curva de aprendizaje es menos pronunciada en comparación con CUDA C++; sin embargo, las abstracciones en esos DSL pueden confundir aún más a nuestro agente. Necesitamos una forma de pasarle el contexto de esas abstracciones al agente.

Los LLMs modernos ya saben cómo escribir Triton "bueno". Pueden trabajar bien con las abstracciones de Triton incluso sin ningún contexto. Sin embargo, para otros DSL como CuTeDSL (que proporciona mucho más control que Triton), he encontrado que tener un directorio de contexto donde el agente pueda buscar para entender las abstracciones del DSL ayuda mucho.

Por ejemplo, clonar el

repositorio de NVIDIA cutlass

en el directorio de contexto es una buena manera de permitir que el agente busque abstracciones relacionadas con

Álgebra de Layout, átomos de Copy/GEMM, jerarquía de memoria, kernels de ejemplo

, y así sucesivamente mientras escribe kernels en CuTeDSL.

En mi experiencia, una buena primera versión del kernel compila sin errores evidentes y pasa la prueba de corrección de la que hablaré a continuación.

Pruebas, Benchmark y Profiling

Dado suficiente contexto al agente, el verdadero cuello de botella ahora se traslada a la validación. La implementación de referencia en sí misma, y la validación contra ella, se vuelven cada vez más importantes. A esta fase la llamo pruebas de corrección o simplemente testing. La velocidad de la implementación de referencia no importa tanto como su intención. Lo que pretendes medir y verificar es lo que tu agente optimizará.

Generalmente, cuando la computación no debe realizarse en precisión reducida, mido el Error Máximo Absoluto/Relativo (MAE), el Error Cuadrático Medio (MSE/RMSE) y el PSNR (Peak Signal to Noise Ratio / Relación de pico de señal a ruido). Cuando hay precisiones reducidas involucradas, tiendo a medir PSNR y similitud coseno (cossim).

La forma en que haces que las versiones del kernel realmente se ejecuten en una GPU depende de si la GPU está disponible localmente o en la nube. Independientemente de eso, nuestro agente debería tener la capacidad de acceder a sus resultados de una forma u otra.

Encuentro que la siguiente metodología de rungs es una buena manera de tener N funciones de prueba:

1def rung(name):
2    def deco(fn):
3        try:
4            out = fn()
5            results[name] = {"ok": True, **(out or {})}
6            print(f"[{name}] ok  " + "  ".join(f"{k}={v}" for k, v in (out or {}).items()))
7        except Exception as e:
8            results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
9            print(f"[{name}] FAILED  {type(e).__name__}: {e}")
10            traceback.print_exc()
11        return fn
12    return deco

que puedes llamar así:

1out = {}
2
3@rung("pre-checks")
4def _():
5    run_pure_checks()
6    run_dsl_checks()
7
8@rung("run")
9def _():
10    out["o"] = custom_kernel(*inputs)
11    torch.cuda.synchronize()
12    return {"shape": tuple(out["o"].shape),
13            "finite": bool(torch.isfinite(out["o"]).all())}

Para el rung de benchmarks, puedes hacer varias cosas:

• Medir el tiempo de ejecución del kernel de principio a fin para el tiempo total invertido
• Usar tracing intra-kernel para medir secciones dentro de un kernel y volcarlas en la salida (usando un tracer personalizado o CUPTI)
• Volcar el IR, PTX, SASS y CUBIN generados a un directorio de dumps y dejar que el agente los revise

El último punto se puede expandir un poco más. A veces, puede suceder que el DSL haga lowering y genere PTX subóptimo (y eventualmente SASS) y encuentres mejores instrucciones o formas que se puedan usar en su lugar. Nuestro agente puede leer los archivos de texto de PTX/SASS e insertar código de más bajo nivel en lugar de dejar que el DSL maneje la parte subóptima. Nuevamente, pasar la documentación de PTX "buscable" como contexto es muy útil aquí.

Lo último para unir todo es el Profiling. Si tu agente puede acceder al CLI de NCU (Nsight Compute Systems), puedes pedirle que haga profiling y genere un reporte para tu kernel como parte del bucle de retroalimentación verificable de arriba.

Reflexiones finales

"¿Entonces el desarrollo de kernels de GPU está muerto?"

"Pues sí, pero en realidad no"

Sí, porque la parte difícil de los layouts, el indexado, las abstracciones y la estructura general puede ser resuelta en gran medida por agentes con suficiente contexto. Puedes reducir fácilmente el trabajo de 2-3 semanas a 1-2 días. No, porque el verdadero cuello de botella se ha desplazado de los kernels a las validaciones. Ya no hay una única forma de hacer las cosas: cuanto mejor sea tu contexto y tu harness, más rápido será el proceso. Los casos especializados se beneficiarán aún más, y un buen harness es todo lo que necesitarás.

Finalmente, en lugar de tratar a los agentes como autónomos, todavía se necesita tratarlos como asistentes realmente inteligentes que puedes guiar. Aquí es donde tu comprensión fundamental de las GPUs y los kernels resulta útil. La parte humana (tú) todavía es necesaria aquí.

Es un sentimiento agridulce, lo sé :)