SCALE-Sim TPU: Validating and Extending SCALE-Sim for TPUs

TL;DR

This paper introduces SCALE-Sim TPU, an enhanced cycle-accurate simulator validated against Google TPU v4 hardware, with improved modeling of latency and integration with modern ML frameworks for accurate TPU performance analysis.

Contribution

The paper extends SCALE-Sim v3 with validated hardware models, learned latency estimations, and a modern frontend for seamless integration with ML frameworks.

Findings

Strong linear correlation between simulated and real hardware latency.

Median relative error below 3% for latency predictions.

Unified simulation of workloads from JAX and PyTorch.

Abstract

Cycle-accurate simulators are widely used to study systolic accelerators, yet their accuracy and usability are often limited by weak validation against real hardware and poor integration with modern ML compiler stacks. This paper presents SCALE-Sim TPU, a validated and extended version of SCALE-Sim v3 for TPU-style accelerators. Specifically, we make three contributions: (1) We validate SCALE-Sim's systolic GEMM model against measurements on Google TPU v4 and show that simulated cycle counts exhibit a strong linear correlation with hardware latency, enabling a simple cycle-to-latency mapping. (2) We introduce lightweight learned latency models for non-systolic elementwise operations, achieving median relative errors below 3 percent using only tensor size and shape, substantially improving end-to-end latency estimation. (3) We integrate a StableHLO-based frontend that allows workloads from modern ML frameworks such as JAX and PyTorch to be simulated directly via a unified compiler IR. Together, these contributions improve the fidelity, coverage, and practicality of cycle-accurate simulation for whole-model performance analysis on TPUs.