Snap & Replay: A new way to analyze uarch-scale performance bottlenecks for ML accelerators

As models become larger, ML accelerators are a scarce resource whose performance must be continually optimized to improve efficiency. Existing performance analysis tools are coarse grained, and fail to capture model performance at the machine-code level. In addition, these tools often do not provide specific recommendations for optimizations. We present SnR, a fine-grained methodology for analyzing ML models at the machine-code level that provides actionable optimization suggestions. Our core insight is to use a hardware-level simulator, an artifact of the hardware design process that we can re-purpose for performance analysis. SnR captures traces from production deployments running on accelerators and replays them in a modified microarchitecture simulator to gain low-level insights into the model’s performance. We implement SnR for our in-house accelerator and used it to analyze the performance of several of our production LLMs, revealing several previously-unknown microarchitecture inefficiencies. Based on these findings, we implement optimizations that have decreased token generation latency for our already highly-optimized production LLMs by up to 4.1%.

(*) Ioannis Zarkadas and Amanda Tomlinson are equal-contribution co-authors to this work.

• Software Systems

• Machine Intelligence