Google adiciona profiling de kernel em nível de ciclo ao XProf para TPUs

O Google publicou, em 23 de setembro de 2026, via post técnico do time de AI↳Inteligência artificial440 conteúdosUX e IA: Transformando Experiências Digitais com Inteligência ArtificialProduto & UX · jan 2025MCP: O que é e por que você vai ouvir falar disso em breve?AI · jul 2025IA generativa e a urgência de reconstruir nossa relação com a verdadeAI · jun 2025Ver tudo em AI → Infra, uma nova suíte de Kernel Profiling dentro do XProf, seu profiler open-source para cargas de trabalho em TPU. A novidade, reportada pela InfoQ, resolve um problema específico de quem escreve kernels customizados com Pallas, Mosaic ou Triton: até agora, esses kernels apareciam como blocos opacos nos traces, sem visibilidade do que acontecia dentro deles ciclo a ciclo.

Para quem constrói infraestrutura pesada de ML, isso importa menos pelo anúncio em si e mais pelo motivo por trás dele: as ferramentas de profiling de TPU dependiam, em boa parte, dos modelos de custo estático do compilador XLA, e esses modelos simplesmente não enxergam o que acontece dentro de um kernel que pula o pipeline padrão de compilação.

Por que os números do XLA mentiam para kernels customizados

Segundo Yogesh SY, do time de AI Infra do Google, kernels feitos com Pallas, Mosaic ou Triton contornam os passes padrão do XLA. Isso distorce os modelos de custo estático usados em métricas como "FLOPs ótimos" e eficiência de goodput, que podem simplesmente ficar erradas ou parar de funcionar. O exemplo dado é direto: uma ferramenta estática pode marcar um bloco de instruções da MXU (a unidade de matmul da TPU) como totalmente utilizado quando, na realidade, a unidade está ociosa, esperando os dados chegarem da HBM. A análise estática não enxerga tempo, só a estrutura do código.

É o tipo de problema que qualquer engenheiro de performance conhece do mundo GPU (profilers como o Nsight Compute existem justamente para fechar essa lacuna entre teoria e execução real), mas que, no ecossistema TPU/JAX, não tinha uma resposta tão granular até agora.

Como a suíte funciona, em três camadas

A implementação expõe informações em três níveis diferentes, cada um com seu próprio mecanismo de ativação.

No nível do compilador, duas flags habilitam a inspeção:

--xla_enable_custom_call_region_trace=true --xla_xprof_register_llo_debug_info=true

Com elas ativas, o Graph Viewer do XProf passa a mostrar um painel "Custom Call Text" com o MLIR já rebaixado (lowered) de cada custom call, o que permite checar se as operações estão de fato fundidas (fused) e se os tiles de memória estão estruturados como o dev pretendia.

No nível estático de execução, o Trace Viewer exibe dados de bundles de Low-Level Operations (LLO): instruções de máquina por ciclo de clock, em tracks alinhadas no tempo para a MXU, as ALUs escalares e vetoriais, vector fills, loads, spills, stores e a cross-lane unit (XLU).

No nível de telemetria de runtime, o XProf amostra contadores de hardware periodicamente. Por padrão, essa amostragem tem resolução mínima de 1 µs, ditada por um timer do host. É o suficiente para a maioria das análises, mas não para eventos muito rápidos dentro de um único kernel.

Sub-microssegundo via trigger externo

A parte mais nova é um modo de captura disparado por evento externo, que elimina esse piso de 1 µs. Nesse modo, o sampler captura instruções de trace da TPU e gatilhos de fronteira, como a entrada e a saída de escopos de custom call, o que permite atribuição em escala sub-microssegundo, ou seja, saber exatamente qual trecho do kernel consumiu qual fração do tempo.

A configuração é feita via jax.profiler.ProfileOptions, usando tpu_enable_periodic_counter_sampling e tpu_tc_perf_counter_sampling_options com is_external_trigger:true. No modo periódico tradicional, o parâmetro correspondente é interval_us. É possível configurar até 28 contadores por core, em até quatro SparseCores (uma grade de 4×28 posições). Como o orçamento é limitado, o engenheiro precisa escolher com cuidado quais contadores acompanhar em cada investigação.

O caso do matmul: de 125,5 µs para 88 µs

A InfoQ reporta um estudo de caso interno do Google usado para demonstrar o fluxo de trabalho: um kernel de matmul tiled (multiplicação de matrizes em blocos) apresentava um stall de memória. Com os contadores de hardware ativos, a equipe viu picos grandes em eventos sync_wait, sinal de que o kernel estava esperando dados da HBM em vez de computar. A correção aplicada foi triple buffering, técnica que sobrepõe os loads da HBM ao cômputo da MXU. O resultado: o tempo do kernel caiu de 125,5 µs para 88 µs, uma redução de cerca de 30%.

Vale o alerta: esse número vem de um único kernel de demonstração escrito pelo Google, não de um benchmark amplo. Ele ilustra o fluxo de trabalho de otimização, não promete ganho equivalente em qualquer kernel Pallas.

Uma hierarquia de confiança para métricas

O post também estabelece o que chama de "hierarquia de confiança" para métricas de performance em TPU. Valores lidos diretamente de registradores de hardware, como utilização de HBM e métricas de TPO, contam como referência real (ground truth) para kernels customizados. Já as estimativas do modelo de custo do XLA precisam ser tratadas com cautela quando o kernel em questão não passou pelo pipeline padrão de compilação.

Uma nova Perf Counters View lista mais de 16.000 contadores brutos em formato tabular. Um detalhe técnico que vale registrar: a altura da track no trace reflete o valor máximo do contador bruto naquele intervalo, não uma porcentagem normalizada. A fonte dá um exemplo de cálculo: um incremento de 100 ciclos numa janela de 500 ns, num core rodando a aproximadamente 2,0 GHz, equivale a 10% de utilização daquela unidade. É uma conta que quem for ler esses gráficos crus vai precisar fazer manualmente.

O que fica em aberto

O XProf faz parte do projeto OpenXLA e se integra ao profiling do JAX, mas a InfoQ é explícita: o post do Google não define status de release nem versão para a suíte de Kernel Profiling. Não há, portanto, garantia de estabilidade de API neste momento.

A amostragem de contadores está documentada para a TPU v7 (Ironwood). Times que rodam gerações anteriores de TPU não devem presumir a mesma cobertura de contadores nem o mesmo comportamento.

Para quem já trabalha com Pallas e quer começar, o caminho indicado pela fonte é o notebook "Pallas Matmul with Perf Counters", no repositório do XProf, junto com as instruções de kernel profiling do OpenXLA.

O recorte para quem desenvolve no Brasil

TPU ainda é infraestrutura de nicho por aqui: a esmagadora maioria dos times brasileiros que otimizam kernels de baixo nível para ML está no mundo CUDA/GPU, com ferramentas como o Nsight Compute já maduras para esse tipo de análise. O público direto desta notícia é menor: equipes que rodam treinamento ou inferência em TPU via Google Cloud↳Google Cloud13 conteúdosPrograma da Google Cloud no Brasil projeta formar de forma gratuita 30 mil universitáriosDev (Back & Front) · mar 2026Google Cloud OnBoard capacita estudantes e desenvolvedores de TIGestão Dev & TI · mai 2019Desenvolvedores poderão participar de treinamento gratuito do Google CloudGestão Dev & TI · mai 2019Ver tudo em DevSecOps →, geralmente startups e times de pesquisa que usam JAX como stack principal e escrevem kernels customizados em Pallas para espremer desempenho além do que o XLA entrega por padrão.

Para esse grupo, a mensagem prática é clara: desconfie dos números de eficiência derivados do XLA para kernels customizados e ancore a otimização em contadores lidos direto dos registradores. É uma mudança de disciplina de trabalho, não só uma feature nova, e, por enquanto, só está documentada para a TPU v7. Para quem consome TPU apenas via APIs de alto nível, sem escrever kernel próprio, o impacto direto é praticamente nulo. Ainda assim, a novidade serve de termômetro: mostra que o Google está levando o ecossistema JAX/Pallas em direção ao controle fino que, até hoje, só o mundo GPU oferecia de forma consolidada.

Fonte: InfoQ