Google ha extendido XProf, su perfilador de código abierto para cargas de trabajo de TPU, con perfilado de kernels a nivel de ciclo que expone los detalles de ejecución de kernels Pallas personalizados. Antes de esta adición, los kernels personalizados aparecían como bloques opacos en capturas de traza; ahora los desarrolladores pueden ver contadores de rendimiento de hardware muestreados con resolución de microsegundos y, en modo de disparo externo, con granularidad sub-microsegundo.
La suite de perfilado opera en tres niveles de abstracción. Para acceder al MLIR reducido del Graph Viewer para cada llamada personalizada, los desarrolladores primero deben pasar los flags del compilador --xla_enable_custom_call_region_trace=true y --xla_xprof_register_llo_debug_info=true, después de lo cual un panel "Custom Call Text" muestra el MLIR reducido, permitiendo a los ingenieros verificar si las operaciones están fusionadas y los tiles de memoria están estructurados como se pretendía. El Trace Viewer muestra datos del bundle de Low-Level Operations (LLO) con instrucciones de máquina por ciclo de reloj, presentando pistas alineadas en tiempo para la MXU, ALUs escalares y vectoriales, rellenos vectoriales, cargas, derrames, almacenamientos y la unidad entre carriles. Para telemetría en tiempo de ejecución, XProf muestrea contadores de hardware periódicamente con un piso de resolución de 1µs en modo estándar, o elimina ese piso completamente en modo de disparo externo capturando instrucciones de traza de TPU y disparadores de límite en entrada y salida de llamada personalizada.
La recopilación de contadores se habilita a través de jax.profiler.ProfileOptions usando tpu_enable_periodic_counter_sampling y tpu_tc_perf_counter_sampling_options. Para modo de disparo externo, los desarrolladores establecen is_external_trigger:true para captura sub-microsegundo; para modo periódico, interval_us reemplaza el flag de disparo. Los desarrolladores pueden configurar hasta 28 contadores por núcleo en hasta cuatro SparseCores. Una nueva Perf Counters View lista más de 16,000 contadores sin procesar en forma tabular, con la altura de la pista reflejando el valor máximo del contador sin procesar en un intervalo en lugar de un porcentaje normalizado.
El caso de estudio de Google sobre un kernel matmul en tiles demuestra el flujo de trabajo. La variante limitada por memoria mostró grandes picos en contadores sync_wait. Al superponer cargas de HBM con cómputo de MXU mediante triple buffering, el equipo redujo esos eventos y redujo el tiempo del kernel de 125.5µs a 88µs, aproximadamente 30% de mejora. El post señala que esta cifra proviene de un único kernel de demostración escrito por Google, no de un benchmark amplio, por lo que ilustra el proceso de optimización en lugar de ganancias típicas en otros lugares.
La adición aborda una brecha fundamental en el perfilado de kernels personalizados. Según Yogesh SY del equipo de AI Infra de Google, los kernels hechos con Pallas, Mosaic o Triton omiten pases XLA estándar, lo que puede distorsionar modelos de costo estático en tiempo de compilación. Métricas como "FLOPs óptimos" y eficiencia de goodput pueden ser inexactas o estar completamente ausentes. Una herramienta estática puede marcar un bloque de instrucciones MXU como completamente utilizado mientras la unidad está inactiva esperando HBM, porque el análisis estático ignora el tiempo. El post establece una "jerarquía de confianza" para métricas: los valores leídos directamente de registros de hardware, como utilización de HBM y métricas de TPO, cuentan como verdad fundamental para kernels personalizados, mientras que las estimaciones del modelo de costo de XLA requieren precaución.
La suite de perfilado es parte del proyecto OpenXLA e se integra con perfilado de JAX. El post no indica un estado de lanzamiento o versión para la suite de Kernel Profiling. El muestreo de contadores está documentado para TPU v7 (Ironwood), por lo que los equipos en generaciones anteriores de TPU no deben asumir la misma cobertura. El presupuesto limitado de contadores 4x28 requiere que los equipos seleccionen contadores para cada investigación, e los ingenieros deben anclar la optimización en contadores a nivel de registro en lugar de tratar números de eficiencia derivados de XLA como confiables para kernels personalizados.
Para equipos ajustando kernels Pallas en TPU v7, el flujo de trabajo comienza con el notebook Pallas Matmul with Perf Counters en el repositorio de XProf y las instrucciones de perfilado de kernels de OpenXLA. El resultado: la visibilidad a nivel de ciclo en la ejecución de kernels personalizados ahora está disponible en Ironwood, pero solo si la habilitas explícitamente y sabes qué contadores observar.