“Los benchmarks rigurosos han impulsado el progreso en la optimización autónoma del rendimiento de kernels GPU… pero no existe un equivalente para los TPUs.”
Es una admisión demoledora, la verdad. Mientras el ecosistema de NVIDIA ha pasado años construyendo un campo de pruebas donde los desarrolladores pueden iterar en kernels con retroalimentación real, el mundo de los TPUs se ha limitado esencialmente a decir “confíen en nosotros, es rápido.” Durante buena parte de la última década, optimizar para el hardware de Google ha sentido menos como ingeniería y más como intentar resolver un cubo de Rubik en una habitación oscura. Cambias una línea de código, esperas a que el compilador termine su magia y luego rezas para que el throughput haya subido de verdad.
La frustración no es solo por la falta de herramientas; es por la cultura de opacidad. En el mundo de las GPU, si un kernel es lento, puedes perfilarlo, encontrar el cuello de botella y corregirlo. En el mundo de los TPUs, a menudo solo te quedas mirando una pared de logs de XLA que parecen escritos por un alienígena confundido. Nos han condicionado a aceptar esto porque los TFLOPs brutos son impresionantes, pero la potencia bruta es inútil si no puedes dirigirlo hacia una carga de trabajo específica sin un doctorado en la arquitectura del compilador interno de Google.
El problema es que el TPU es una caja negra envuelta en un enigma, envuelta en una API propietaria. Si trabajas en JAX, estás mayormente a merced de XLA. Pero XLA es un compilador de propósito general; no siempre es la forma más eficiente de implementar una operación específica y extraña que solo existe en tu artículo de investigación de nicho. Cuando los desarrolladores de GPU se topan con un muro, escriben un kernel CUDA personalizado. Cuando los desarrolladores de TPU se topan con un muro, usualmente solo compran más pods de TPU.
Afinar un kernel de TPU sin un benchmark es como intentar hornear un pastel donde la temperatura del horno es un secreto y el temporizador está roto. Solo sigues probando diferentes configuraciones y esperando que el resultado no sea una masa quemada. (Y Dios sabe que todos hemos pasado demasiadas horas mirando el HLO de XLA).
Entra JAXBench. La idea aquí no es solo proporcionar un conjunto de pruebas, sino crear un marco para la optimización autónoma de kernels. Estamos hablando de que la IA escriba el código para hacer que el hardware de IA corra más rápido. ¿Quién realmente disfruta escribiendo kernels de TPU a mano? El proceso es tedioso, propenso a errores y, francamente, un poco miserable.
La falta de un benchmark estandarizado para kernels de TPU ha sido un lastre masivo para el ecosistema. Es un fallo estructural. Al no proporcionar una forma de medir ganancias incrementales, Google esencialmente aseguró que solo un puñado de personas—los “sumos sacerdotes” de la optimización de TPUs—supieran realmente cómo exprimir el hardware por cada gota de rendimiento.
Al proporcionar un “objetivo compartido para escalar colinas,” JAXBench finalmente le da a los LLMs—los encargados de escribir los kernels—una forma de verificar su trabajo. El ciclo es simple: la IA propone un kernel, el benchmark lo prueba en hardware TPU real, y el resultado se retroalimenta a la IA para refinar el código. Este es el mismo ciclo que ya ha hecho que la generación de kernels GPU sea significativamente más viable.
Pero seamos honestos con la fricción aquí. El acceso a los TPUs sigue siendo un círculo cerrado. A menos que estés en un laboratorio de nivel 1 o tengas una línea de créditos masiva en GCP, no puedes simplemente levantar un cluster para ver si tus resultados de JAXBench se mantienen. Incluso así, la latencia del ciclo compilar-probar-repetir es suficiente para hacer que cualquiera quiera renunciar. No solo estás peleando contra el código; estás peleando contra la infraestructura en la nube.
Aun así, el cambio hacia la optimización autónoma es el único camino lógico. La complejidad de los aceleradores modernos ha superado oficialmente la capacidad humana para afinarlos a mano. Hemos llegado a un punto donde el hardware es demasiado intrincado para que un humano lo optimice eficientemente sin pasar seis meses en un agujero de conejo de alineación de memoria y estrategias de tiling.
La era del experto humano en TPUs ha terminado.
Si podemos automatizar realmente el proceso de escritura de kernels, el cuello de botella se desplaza de “cómo hago que esto sea rápido” a “qué es lo que realmente quiero computar.” La verdadera victoria aquí no es el benchmark en sí, sino la admisión de que necesitamos máquinas para optimizar las máquinas. Es una rendición, pero es una práctica.
Dentro de seis meses, veremos un ranking público para kernels de TPUs autónomos que obligue a Google a liberar especificaciones de hardware más detalladas para mantener la competencia honesta. Hasta entonces, solo estamos adivinando.
