TPU (Tensor Processing Unit)

Uma Unidade de processamento tensorial (TPU) é um circuito integrado de aplicação específica (ASIC) projetado pelo Google para acelerar o treinamento e a inferência de redes neurais. Uma TPU é construída em torno de grandes unidades de multiplicação de matrizes, porque a multiplicação de matrizes domina o trabalho de aprendizado profundo. O Google usa TPUs em seus data centers desde 2015, treina seus modelos Gemini nelas e as aluga a outras empresas como Cloud TPU. A oitava geração, anunciada em abril de 2026, é composta por dois chips: TPU 8t para treinamento e TPU 8i para inferência.

Como funciona uma TPU#

O núcleo de cada TPU é uma unidade de multiplicação de matrizes (MXU), organizada como um arranjo sistólico: uma grelha de células de multiplicação e acumulação que passam dados diretamente às células vizinhas. Os pesos são carregados nas células uma vez, as entradas atravessam a grelha um passo por ciclo de relógio e as somas parciais descem pelas colunas, pelo que cada valor lido da memória é reutilizado por muitas células. A leitura da memória consome muito mais tempo e energia do que a aritmética, e é aí que o projeto ganha eficiência. Nas Cloud TPU v2 e v3, cada MXU era uma matriz de 128×128 que multiplicava em bfloat16 e acumulava em FP32. Uma breve simulação mostra o fluxo de dados numa grelha de 4×4:

import numpy as np

rng = np.random.default_rng(0)
W = rng.integers(-5, 5, (4, 4))  # weights, preloaded one per cell of a 4x4 grid
X = rng.integers(-5, 5, (6, 4))  # six input rows stream through the grid
Y = np.zeros((6, 4), dtype=int)

# Cell (k, j) holds W[k, j]. Input X[i, k] enters row k from the left and reaches
# cell (k, j) at cycle i + k + j, where it is multiplied by the stored weight and
# added to the partial sum for Y[i, j] passing down column j
cycles = 6 + 4 + 4 - 2
for t in range(cycles):
    for k in range(4):
        for j in range(4):
            i = t - k - j
            if 0 <= i < 6:
                Y[i, j] += X[i, k] * W[k, j]

print(f"{X.size * 4} multiply-adds in {cycles} cycles, equal to X @ W: {np.array_equal(Y, X @ W)}")

Imprime 96 multiply-adds in 12 cycles, equal to X @ W: True: as 16 células produzem o produto matricial exato, enquanto cada peso é lido da memória uma vez e reutilizado nas seis linhas de entrada. A TPU de primeira geração executava este padrão numa grelha de 256×256 células de 8 bits, com 65.536 operações de multiplicação e acumulação por ciclo e um pico de 92 biliões de operações por segundo. A sua execução simples e determinística era adequada ao objetivo da Google de cumprir um tempo de resposta alvo no percentil 99 durante a inferência.

As TPUs executam código compilado pelo XLA, que recebe o grafo computacional completo do JAX, do TensorFlow ou do PyTorch e mapeia-o para as unidades matriciais. Os chips são ligados por interconexões dedicadas em pods com milhares de unidades, o que permite treinar os maiores modelos fundamentais com muitas TPUs em simultâneo.

TPU vs. GPU vs. CPU#

• CPU: Alguns núcleos potentes para ramificações e código de uso geral, que tratam do carregamento e pré-processamento de dados na maioria dos pipelines de treino.
• GPU: Milhares de núcleos programáveis e Tensor Cores para operações matriciais. As GPUs executam quase todos os modelos e frameworks e são vendidas por muitos fornecedores e serviços de cloud.
• TPU: Troca essa flexibilidade por eficiência em grandes operações matriciais densas. As cargas de trabalho com formas de tensores estáticas funcionam melhor, porque uma alteração de forma obriga o XLA a recompilar; além disso, as operações personalizadas, as muitas ramificações ou a aritmética de alta precisão não são adequadas às TPUs.
• NPU: Termo genérico para aceleradores de redes neurais presentes em telemóveis, portáteis e chips de edge, categoria que inclui o Edge TPU.

O Edge TPU#

O Edge TPU é uma TPU pequena, dedicada apenas à inferência, presente no acelerador USB Coral da Google, nos módulos M.2 e PCIe e na Dev Board. Executa 4 biliões de operações por segundo com 2 watts e só suporta operações com números inteiros de 8 bits, pelo que as redes têm de ser totalmente convertidas para INT8 por meio da quantização de modelos e, em seguida, compiladas com o compilador Edge TPU da Google, que coloca no chip as operações compatíveis e deixa as restantes na CPU anfitriã. O Ultralytics YOLO automatiza ambos os passos com format="edgetpu", que ativa a quantização INT8 e executa o compilador numa máquina Linux x86-64, produzindo um modelo que funciona em placas como uma Raspberry Pi com um acelerador USB Coral:

from ultralytics import YOLO

# Load the smallest YOLO26 detection model
model = YOLO("yolo26n.pt")

# Quantize to INT8 and compile for the Edge TPU (x86-64 Linux only)
path = model.export(format="edgetpu")
print(f"Edge TPU model: {path}")

A plataforma Coral da Google centra-se agora no Coral NPU, um projeto de acelerador RISC-V de código aberto anunciado em outubro de 2025 para dispositivos vestíveis, que visa atingir 512 mil milhões de operações por segundo com alguns miliwatts.