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 kernels. Voy a hablar de lo que aprendí haciendo eso.
Honestamente, escribir kernels de GPU es una tarea vergonzosamente verificable.
Primero, decides para qué operación(es) quieres escribir 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 esto, mides el tiempo de ejecución del kernel. Las versiones posteriores se construyen sobre esto, y sigues optimizando hasta que alcanzas las métricas de roofline o quedas satisfecho.
Bucle verificable: desarrollo de kernels de GPU
La imagen anterior 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 -> de vuelta 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 anterior 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 necesario 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 que resultan muy fáciles de empezar a usar 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 pasar 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
NVIDIA cutlass
repositorio en el directorio de contexto es una buena forma de permitir que el agente busque abstracciones relacionadas con
Álgebra de layouts, átomos de Copy/GEMM, jerarquía de memoria, kernels de ejemplo
, y demás 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 perfilado
Con suficiente contexto para el agente, el verdadero cuello de botella ahora se traslada a la validación. La implementación de referencia en sí, 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á.
Normalmente, cuando la computación no debe realizarse en precisión reducida, mido el Error Absoluto/Relativo Máximo (MAE), el Error Cuadrático Medio (MSE/RMSE) y el PSNR (Relación Señal-Ruido de Pico). Cuando hay precisiones reducidas de por medio, tiendo a medir el PSNR y la similitud coseno (cossim).
La forma en que haces que las versiones del kernel se ejecuten realmente 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.
Me parece que la siguiente metodología de rungs es una buena forma 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 fn12 return deco
que puedes llamar así:
1out = {}23@rung("pre-checks")4def _():5 run_pure_checks()6 run_dsl_checks()78@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 benchmark, puedes hacer varias cosas:
- Medir el tiempo de ejecución del kernel de extremo a extremo para el tiempo total empleado.
- Usar el rastreo 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 dumps y dejar que el agente los revise.
El último punto se puede ampliar un poco más. A veces puede suceder que el DSL genere un 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 PTX/SASS e insertar código de más bajo nivel directamente en lugar de dejar que el DSL maneje la parte subóptima. De nuevo, pasar la documentación de PTX "buscable" como contexto es muy útil aquí.
Lo último para unirlo todo es el perfilado. Si tu agente puede acceder a la CLI de NCU (Nsight Compute Systems), puedes pedirle que perfile y genere un informe de tu kernel como parte del bucle de retroalimentación verificable anterior.
Reflexiones finales
"¿Entonces el desarrollo de kernels de GPU está muerto?"
"Bueno, 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 trasladado 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.
Por último, en lugar de tratar a los agentes como autónomos, todavía es necesario tratarlos como asistentes realmente inteligentes que puedes guiar. Aquí es donde tu comprensión fundamental de las GPU y los kernels resulta útil. La parte humana (tú) todavía es necesaria aquí.
Es un sentimiento agridulce, lo sé :)
