ALT What is this document? I was going to start an AI chip company and this document contains my thinking about how to make an AI chip, both software and hardware. It’s written for someone else who would be responsible for or work on an AI chip project - it is some of what you’d need to know. I thought others might be interested in understanding the area of AI chips, as well, so most of this should be easily understandable at a hobbyist level. Some parts are more technical and require a deeper background knowledge - those are marked in bold as “for nerds”. This document also contains anecdotes from my time in Silicon Valley that I hope will be entertaining and perhaps illuminating. They are clearly marked and so can be skipped if that doesn’t interest you. Who am I to write this I have worked at Google on TPUs and at Nvidia on GPUs in positions of technical leadership and also at other companies that produce AI software and hardware like Amazon and Facebook. I was the technical software

ALT TOC

ALT Proposal for a hardware design (summary) Here’s a summary of my proposed design. I recommend reading this section last, but you can read it up front if you wish. The defining features are that this design is easy to do in hardware, as easy to program as a CPU and is very cheap on a computation / watts basis. I doubt anyone making an AI chip will simply adopt my plan as-is, but I would suggest going through this design and understanding, for each deviation, exactly why you think deviating from this is a good idea. I would surely also have modified the design over time if I had embarked on actually building this. “Hold on, this is just a CPU with a systolic array attached, you mean to say you can outcompete sophisticated GPUs and AI ASICs with just this?” Yes, that’s what I believe. Most of this chip is cache and systolic array, so you do not have to deviate from a CPU approach for scalar and vector compute. CPUs are the most programmable thing and both hardware and software engineers a

ALT My experience with co-design This section is a bit specific to my own work, but it has some ideas for co-design overall. For a period of some years, I was the software person that would be consulted on most hardware matters for TPUs. I eventually transitioned myself out of this by connecting people more directly as the XLA software team became larger with more experts in various areas. A few examples of what I worked on, so you can get a sense for what I was doing: I designed the interface for DMAs for TPUv4 to be much more software friendly than previously (the previous functionality required some kind of bit-pattern PhD to understand and also wasn’t general enough), I completely changed how memory protection worked on TPUs (the previous solution was so complex that we never used it, so I simplified it), I got quite a bit into the details of the performance properties of HBM on TPUs and I rejected some changes that otherwise would have made it into the chip on the grounds that it wou