Google estendeu o XProf, seu profiler de código aberto para cargas de trabalho em TPU, com criação de perfil de kernel em nível de ciclo que expõe os detalhes de execução de kernels Pallas customizados. Antes dessa adição, kernels customizados apareciam como blocos opacos em capturas de rastreamento; agora desenvolvedores podem ver contadores de desempenho de hardware amostrados com resolução de microssegundos e, em modo de gatilho externo, com granularidade sub-microssegundo.
O conjunto de criação de perfil opera em três níveis de abstração. Para acessar o MLIR reduzido do Graph Viewer para cada chamada customizada, desenvolvedores devem primeiro passar as flags do compilador --xla_enable_custom_call_region_trace=true e --xla_xprof_register_llo_debug_info=true, após o que um painel "Custom Call Text" exibe o MLIR reduzido, permitindo que engenheiros verifiquem se as operações estão fundidas e os tiles de memória estão estruturados conforme pretendido. O Trace Viewer mostra dados do bundle de Low-Level Operations (LLO) com instruções de máquina por ciclo de clock, apresentando faixas alinhadas no tempo para a MXU, ALUs escalares e vetoriais, preenchimentos vetoriais, carregamentos, derramamentos, armazenamentos e a unidade cross-lane. Para telemetria em tempo de execução, o XProf amostra contadores de hardware periodicamente com piso de resolução de 1µs em modo padrão, ou remove esse piso inteiramente em modo de gatilho externo capturando instruções de rastreamento de TPU e gatilhos de limite na entrada e saída de chamada customizada.
A coleta de contadores é habilitada através de jax.profiler.ProfileOptions usando tpu_enable_periodic_counter_sampling e tpu_tc_perf_counter_sampling_options. Para modo de gatilho externo, desenvolvedores definem is_external_trigger:true para captura sub-microssegundo; para modo periódico, interval_us substitui a flag de gatilho. Desenvolvedores podem configurar até 28 contadores por núcleo em até quatro SparseCores. Uma nova Perf Counters View lista mais de 16.000 contadores brutos em forma tabular, com altura de faixa refletindo o valor máximo do contador bruto em um intervalo em vez de uma porcentagem normalizada.
O estudo de caso do Google em um kernel matmul em tiles demonstra o fluxo de trabalho. A variante vinculada à memória mostrou grandes picos em contadores sync_wait. Ao sobrepor carregamentos de HBM com computação de MXU através de triple buffering, a equipe reduziu esses eventos e cortou o tempo de kernel de 125.5µs para 88µs, aproximadamente 30% de melhoria. O post observa que essa figura vem de um único kernel de demonstração escrito pelo Google, não de um benchmark amplo, portanto ilustra o processo de otimização em vez de ganhos típicos em outros lugares.
A adição aborda uma lacuna fundamental na criação de perfil de kernels customizados. De acordo com Yogesh SY da equipe de AI Infra do Google, kernels feitos com Pallas, Mosaic ou Triton pulam passes padrão de XLA, o que pode distorcer modelos de custo estático em tempo de compilação. Métricas como "FLOPs ótimos" e eficiência de goodput podem ser imprecisas ou estar completamente ausentes. Uma ferramenta estática pode sinalizar um bloco de instrução de MXU como totalmente utilizado enquanto a unidade fica ociosa aguardando HBM, porque análise estática ignora tempo. O post estabelece uma "hierarquia de confiança" para métricas: valores lidos diretamente de registradores de hardware, como utilização de HBM e métricas de TPO, contam como verdade fundamental para kernels customizados, enquanto estimativas de modelo de custo de XLA requerem cautela.
O conjunto de criação de perfil faz parte do projeto OpenXLA e se integra com criação de perfil de JAX. O post não declara um status de lançamento ou versão para o conjunto de Kernel Profiling. Amostragem de contadores é documentada para TPU v7 (Ironwood), portanto equipes em gerações anteriores de TPU não devem assumir a mesma cobertura. O orçamento limitado de contadores 4x28 requer que equipes selecionem contadores para cada investigação, e engenheiros devem ancorar otimização em contadores em nível de registrador em vez de tratar números de eficiência derivados de XLA como confiáveis para kernels customizados.
Para equipes ajustando kernels Pallas no TPU v7, o fluxo de trabalho começa com o notebook Pallas Matmul with Perf Counters no repositório XProf e as instruções de criação de perfil de kernel OpenXLA. O resultado: visibilidade em nível de ciclo na execução de kernel customizado agora está disponível no Ironwood, mas apenas se você habilitá-la explicitamente e souber quais contadores observar.