Google incorpora perfilado a nivel de ciclo al profiler XProf para TPUs
XProf ahora muestra métricas de hardware por ciclo en kernels personalizados de Pallas, permitiendo reducir tiempos de ejecución hasta un 30% en TPU v7.
Google ha ampliado XProf, su profiler open‑source para cargas de trabajo en TPU, con una suite de perfilado a nivel de ciclo. Con la nueva funcionalidad, los kernels escritos en Pallas, Mosaic o Triton dejan de aparecer como bloques opacos y se pueden inspeccionar instrucción por instrucción.
En la arquitectura TPU v7 (Ironwood) XProf captura contadores de rendimiento hardware en tiempo real. Un caso de estudio con una multiplicación de matrices mostró cómo un cuello de botella de memoria, identificado mediante los contadores de sync_wait, se resolvió con triple buffering, bajando el tiempo de kernel de 125.5 µs a 88 µs, es decir, una mejora cercana al 30 %.
Para activar el perfilado, los desarrolladores deben pasar dos flags al compilador: --xla_enable_custom_call_region_trace=true y --xla_xprof_register_llo_debug_info=true. Con ello, el Graph Viewer incluye un panel “Custom Call Text” que muestra el MLIR reducido de cada llamada personalizada, y el Trace Viewer expone datos de bajo nivel (LLO) con instrucciones por ciclo para unidades MXU, ALU, XLU, entre otras.
El muestreo de contadores tiene una resolución mínima de 1 µs, pero un modo externo de disparo permite capturas sub‑micrósecond, registrando entradas y salidas de los ámbitos de llamadas personalizadas. Se pueden habilitar hasta 28 contadores por núcleo, distribuidos en hasta cuatro SparseCores (matriz 4 × 28). La configuración se realiza mediante jax.profiler.ProfileOptions, usando los parámetros tpu_enable_periodic_counter_sampling y tpu_tc_perf_counter_sampling_options con is_external_trigger:true.
Google advierte que los números de coste de XLA pueden distorsionarse para kernels personalizados, pues se omiten pasos de compilación estándar. Por eso, la suite propone una “jerarquía de confianza”: los valores leídos directamente de registros hardware son la verdad, mientras que los estimados del modelo de coste deben tomarse con cautela.
XProf forma parte del proyecto OpenXLA y se integra con el ecosistema JAX. Aunque la nota no indica una versión específica, la documentación está disponible en el repositorio y en los notebooks de ejemplo. Los equipos que trabajen con TPUs anteriores a la generación v7 no deben asumir cobertura completa.
Implicaciones: los ingenieros pueden ahora validar fusiones, estructuración de tiles y utilización de memoria con datos de hardware, lo que reduce la dependencia de métricas estáticas y ayuda a evitar optimizaciones basadas en información incompleta. La capacidad de inspeccionar contadores a nivel de ciclo abre la puerta a ajustes finos de rendimiento que antes requerían pruebas manuales extensas.
