Designing my own TPU for transformer inference
This was a summer project I made to learn the basics of ML performance optimization and hardware-software co-design. The overall goal of this project is to run a small transformer on my Cmod A7 FPGA board. Then, I want to derive some of the common results you’ll find in LLM scaling (limited to a single device for now), and optimize my kernels to see how much performance I can squeeze out of my hardware design. I’m assuming you have basic knowledge on how transformer inference works (prefill, decode), but not the hardware side of things. All the code can be found in the project’s github repo. Much of the analysis I do (including the next section), along with the TPU architecture itself, is heavily inspired by the Scaling Book. If you find this article interesting, you might enjoy reading the scaling book! Before we get into hardware, we need to ask a seemingly obvious question: what makes a model/algorithm faster? In ML, this problem can be analyzed in 2 domains: compute and communicati