Google lança suite de microbenchmarks para medir desempenho real de TPU

Especificações de hardware mentem por omissão. Uma TPU pode prometer determinado throughput teórico em TFLOPs, mas o que importa pra quem coloca modelo em produção é quanto disso o workload real consegue arrancar. É esse buraco entre folheto e runtime que o Google tenta fechar com uma suite de microbenchmarks para TPUs, apresentada em post no Google Developers Blog por engenheiros da empresa. O código está aberto no repositório accelerator-microbenchmarks no GitHub.

O que a suite mede

A proposta é segmentar a avaliação em áreas funcionais, medindo cada componente isoladamente. Segundo os autores, isso transforma otimização de tentativa e erro em "disciplina de engenharia empírica". São cinco categorias:

• Network: caracteriza o interconnect entre chips (ICI) em operações coletivas como all-gather, all-reduce, reduce-scatter e all-to-all. Métricas: throughput (GB/s) e latência.
• Compute: mede a multiplicação de matrizes (GEMM/Matmul) em TFLOPs e, principalmente, o MFU (Model FLOPs Utilization), que indica o quanto o workload realmente aproveita a Matrix Unit (MXU).
• HBM: banda efetiva de memória on-chip em operações vetoriais grandes, em GB/s.
• Host Transfer: eficiência de I/O entre CPU do host e a HBM do acelerador via PCIe, medindo transferências Host-to-Device e Device-to-Host.
• Ragged-Paged Attention (RPAv3): benchmarks voltados a modelos transformer, com foco em atenção e batch matmul, úteis para prever latência de Time-to-First-Token.

Por que isso importa pra quem otimiza ML

O valor central, segundo o post, está em alimentar o modelo Roofline. Medindo o pico de MXU e a banda de HBM, dá pra classificar o gargalo de um workload em três categorias:

• Compute-bound: bateu no teto da MXU. Otimizar acesso à memória não adianta nada; o caminho é eficiência de kernel ou reduzir FLOPs.
• Memory-bound: preso na banda de HBM. Melhora com locality de dados (VMEM) ou menos tráfego de memória.
• Network-bound: stalls no ICI ou DCN impedem os chips de chegar ao potencial de cálculo. Aqui o remédio é melhor sharding ou overlap de comunicação.

Sem essa classificação, é comum gastar dias otimizando a coisa errada. O microbenchmark serve de bússola diagnóstica antes de mexer no código.

O detalhe de hardware que vira decisão de arquitetura

Um ponto interessante do texto é como o benchmark expõe a "personalidade" do silício. Na arquitetura Ironwood (TPU7x), o systolic array de 256x256 impõe restrição física sobre o formato dos operandos. Os testes confirmam que modelos com head_dim de 128 (comum em variantes mais antigas do Llama) atingem MFU subótimo. A conclusão prática: dá pra co-projetar o modelo com dimensões alinhadas a fronteiras de 256 bytes para extrair mais da MXU. É o tipo de insight que só aparece quando você mede, não quando lê a spec.

A suite também ajuda a escolher a alavanca de software certa depois de identificar o gargalo: comparar operações JAX padrão contra kernels Pallas ou Tokamax Splash Attention, ajustar sharding e mesh (topologias tipo 2x2x2 vs 4x2x1) para esconder comunicação atrás de computação, e calcular o trade-off de rematerialização (recalcular dado vs. guardar na HBM).

O caso concreto: MoE de 110B

A parte mais convincente é o estudo de caso. Os engenheiros analisaram um treino de Mixture-of-Experts de 110B em configuração 4x4x4 de TPU 7x. O diagnóstico via Roofline separou os gargalos: forward-pass e operações de dense-core eram compute-bound, chegando a 1,85 PFLOPS, enquanto roteamento e atenção da arquitetura MoE ficavam memory-bound, operando a 30-60% do teto teórico por saturação de HBM durante o dispatch dos experts.

O resultado: os benchmarks guiaram intervenções que reduziram o tempo de step de treino em 21,2%. Os testes de coletivas ICI apontaram stalls de comunicação resolvidos com offload de coletivas na SparseCore, e os testes de banda de HBM revelaram os primitivos de atenção memory-bound, motivando a adoção do Tokamax Splash Attention.

Rodando na prática

A execução é orientada a YAML e roda via Kubernetes. Um exemplo de compute microbenchmark:

kubectl apply -f Ironwood/guides/collectives/tpu7x-2x2x1-ici-all-gather-microbenchmark.yaml

A estrutura do repo separa src/ (implementações) de configs/ (parâmetros de tuning: tamanho de buffer, número de iterações e tipos de dado como bf16, fp8, fp32).

O que fica

Para quem opera ML sobre infraestrutura Google no Brasil, o ponto pragmático é este: os microbenchmarks permitem construir modelos analíticos que preveem throughput, TTFT, TPOT e MFU antes de escalar para slices grandes e caras de TPU. Testar num 2x2x1 e extrapolar sai muito mais barato que descobrir o gargalo num cluster de produção. Vale clonar o repo, rodar as coletivas e comparar com o que a spec promete, porque a distância entre os dois números costuma ser a conta que chega no fim do mês.

Fonte: Google Developers Blog