Optimización de núcleos GPU generalmente requiere un profundo conocimiento de la arquitectura del hardware y de las instrucciones de bajo nivel. Los modelos de lenguaje todavía manejan esto de manera inestable: pueden generar código plausible, pero no siempre resulta correcto y rápido. PTXBench es un nuevo benchmark que intenta aportar medibilidad y claridad a este proceso.
Qué es PTXBench y por qué es necesario
PTXBench es un banco de pruebas para grandes modelos de lenguaje que intentan utilizar PTX (Parallel Thread Execution) específico de la arquitectura para optimizar núcleos GPU. PTX es una representación intermedia de instrucciones para GPU de NVIDIA, y es a este nivel donde se puede controlar finamente los registros, la memoria y el planificador. Los benchmarks habituales solo comprueban el código final, pero PTXBench mira más profundo: si el modelo aplica realmente las instrucciones de bajo nivel objetivo, o simplemente genera código «parecido a C».
Los autores se centraron en dos cargas de trabajo clásicas: la multiplicación de matrices (GEMM) y el mecanismo de atención (attention), que es crítico para los transformers. Las pruebas se realizaron en aceleradores actuales: H100 y B200. Esta elección permite comprobar no solo las habilidades básicas del modelo, sino también su capacidad de adaptarse a arquitecturas modernas con sus características específicas.
Qué mide el benchmark
PTXBench evalúa los modelos según tres parámetros clave:
- Corrección funcional — el resultado del núcleo generado debe coincidir con el de referencia.
- Ejecución de instrucciones objetivo — se comprueba si en tiempo de ejecución aparecen esas instrucciones PTX que el modelo debía utilizar.
- Aceleración — cuánto más rápido es el núcleo obtenido en comparación con las bibliotecas frontier, es decir, las mejores implementaciones ya existentes.
Este enfoque separa tres problemas distintos. El modelo puede escribir código correcto, pero no usar las instrucciones requeridas. O usarlas, pero obtener un rendimiento peor que el de la biblioteca estándar. PTXBench permite ver en qué etapa exacta se produce el fallo.
Resultados: la optimización todavía se le da a los modelos de manera desigual
La evaluación mostró que la capacidad de los LLM para trabajar con PTX depende en gran medida de la tarea. En casos simples, los modelos demuestran resultados decentes, pero la proporción de soluciones exitosas cae drásticamente en tareas backward complejas para attention. Esto es esperable: el paso inverso requiere un manejo más cuidadoso de los gradientes y la memoria.
Otra conclusión importante: la ejecución de instrucciones objetivo por sí sola no garantiza velocidad. El modelo puede «acertar» con las instrucciones PTX necesarias, pero elegir una mala estrategia de paralelización o una mala colocación de las sincronizaciones. Al final, ninguno de los modelos probados pudo competir de manera estable con las bibliotecas frontier en todo el conjunto de tareas.
Cómo se ajustó el modelo y qué funcionó
Para entender si se pueden mejorar estas habilidades, los autores realizaron un ajuste fino supervisado del modelo Qwen3.6-27B. Además, utilizaron una técnica que se podría llamar aprendizaje condicionado a la corrección: el modelo no solo aprende a generar código, sino a corregirlo, recibiendo retroalimentación sobre los pasos incorrectos.
Este enfoque mejoró varias tareas, pero no proporcionó un progreso uniforme. La generalización siguió siendo débil: el modelo podía manejar bien algunos tipos de núcleos y fallar en otros. El análisis mostró que no solo importa el tamaño del conjunto de datos de entrenamiento. Más importante es su cobertura, el equilibrio entre diferentes escenarios y la calidad del «profesor» que genera ejemplos para las cadenas de razonamiento. Si los datos están sesgados hacia casos simples, el modelo no aprenderá la complejidad.
Conclusión principal
PTXBench ofrece a la comunidad un banco de pruebas verificable donde se puede rastrear el progreso de los LLM en la optimización de núcleos GPU. No solo da una calificación de «rápido/lento», sino que muestra en qué nivel se rompen los planes del modelo: en la sintaxis, en la aplicación de instrucciones específicas o en la elección de la estrategia. Es un paso importante para hacer que la optimización automática de núcleos sea predecible y práctica.
