TPUs, Trainium and Custom Accelerators: Where Non-GPU Silicon Wins

General-purpose Graphics Processing Units (GPUs) excel at matrix multiplication due to their programmable nature, but a purpose-built accelerator is essentially a matrix multiplier that is tailored for efficient computation. This distinction leads to the existence of non-GPU silicon, which is driven by generality costing area and power.

A GPU contains instruction fetch and decode for every warp, a register file optimized for arbitrary programs, caches with coherence machinery, schedulers for divergent control flow, and graphic heritage. However, these features are unnecessary when the workload consists of a fixed sequence of large dense multiplications with a few element-wise operations in between.

By reclaiming area and on-chip memory, a custom accelerator can provide better performance per watt and per unit of silicon, where power delivery and cooling are hard limits rather than preferences.

Accelerators are primarily designed by organizations with high internal demand due to the large upfront cost of designing and taping out a chip on a leading process. This cost must be recovered through the number of units built, making it economical only for organizations deploying a large fleet of one workload. For a buyer, custom silicon is often part of a platform rather than a standalone product, making evaluation usually an assessment of a service rather than a product.

Systolic arrays and dataflow are recurring architectural ideas that aim to maximize FLOPs performed per byte moved, thus improving the ratio that determines whether memory or arithmetic units bind. These designs can sit further up the roofline on the same memory bandwidth, as reusing operands in the fabric allows for more FLOPs to be performed.

However, there are trade-offs to consider when utilizing custom accelerators. Portability is lost when the work done to optimize a model for one accelerator cannot be easily transferred to another. New techniques in research code may take time to be ported to an accelerator, leading to a long tail of support. Debugging surface tends to be thinner in areas where the architecture is most unusual, and suppliers may be concentrated, making hardware choices dependent on the chosen cloud provider.

Deciding on a custom accelerator without a benchmark that can be run is challenging, as the useful comparisons require access that most readers do not have. To ensure compatibility, it is essential to check that the operation set compiles before anything else. Additionally, comparing on different hardware is crucial, as a single generation's results may not be representative of all generations.