Tensor processing unit (TPU)

Google begyndte at tage den første TPU i brug i sine datacentre i 2015 og omtalte den offentligt i 2016; ISCA 2017-artiklen af Jouppi m.fl. gav detaljerne. TPU v1 var en coprocessor kun til inferens, tilsluttet via PCIe og bygget op om en matrixenhed med 65.536 8-bit multiply-accumulate-celler (et 256 × 256-array), der leverede 92 TOPS i maksimal heltalsydelse. Det centrale designvalg var det systoliske array: operanderne strømmer rytmisk gennem et gitter af simple processorelementer, som hver sender delsummer videre til naboen, så en stor matrixmultiplikation gennemføres med meget få læsninger og skrivninger i hukommelsen. Det sparer den energi, en generel processor bruger på caches, instruktionshentning og registerfiler, og det er derfor, TPU'er opnår høj ydelse pr. watt på tæt lineær algebra, men er dårlige til uregelmæssig kode med meget kontrolflow.

Senere generationer gjorde chippen til en træningsplatform. TPU v2 (2017) tilføjede high-bandwidth memory, understøttelse af flydende kommatal og bfloat16-formatet - 8 eksponentbit som FP32, men kun 7 mantissebit - der bevarer FP32's dynamiske område og siden er taget i brug bredt af anden hardware. Fra v2 og frem forbindes chippene med en dedikeret inter-chip interconnect (ICI) til "pods"; TPU v4 indførte optisk omkoblede 3D-torus-topologier til store slices, og senere generationer fortsatte med v5e og v5p, Trillium (v6e) og Ironwood (TPU7x), annonceret i april 2025, som Google dokumenterer med 192 GB HBM pr. chip, indbygget FP8-understøttelse og pods på op til 9.216 chips. Flere generationer kombinerer TensorCores, der rummer matrixenhederne samt vektor- og skalarenheder, med SparseCores til arbejdsbelastninger med mange embeddings, fx anbefalingsmodeller.

TPU'er programmeres gennem XLA-compileren, der sporer en beregningsgraf og kompilerer den til den specialiserede hardware. JAX er den mest naturlige frontend, og TensorFlow og PyTorch/XLA understøttes også. Denne kompilér-først-model har praktiske konsekvenser: statiske tensorformer foretrækkes klart, fordi ændrede former udløser genkompilering; brugerdefinerede CUDA-kernels kan ikke flyttes og skal skrives om (fx med Pallas) eller undgås; og fejlfinding og profilering foregår med andre værktøjer end i GPU-verdenen. Det er som regel dette og ikke den rå ydelse, der er den største omkostning ved at flytte.

Sammenlignet med GPU'er bytter TPU'er alsidighed og valg af leverandør for effektivitet på neurale netværk og tæt integration med Googles infrastruktur. De bruges primært som kapacitet i Google Cloud og i Googles egne produkter, så at tage dem i brug betyder at acceptere denne cloud som afviklingsmiljø, hvilket er et spørgsmål om lock-in og dataplacering. TPU'en er også ét eksempel på en bredere klasse af AI-acceleratorer (ASIC'er som AWS Trainium og Inferentia), og den skal ikke forveksles med Googles Edge TPU, en lille, strømbesparende inferenschip til enheder, der deler navnet, men ikke datacenterarkitekturen.