zihaomu/gpu_notebook

gpu_notebook is a hands-on learning project that walks through how to make AI computations run faster on GPUs. It's structured as a series of exercises or "tasks," each tackling a different piece of the puzzle that goes into running large AI models efficiently. The project covers four main areas. Task 0 explores how to speed up matrix multiplication, which is the backbone of most AI math. Task 1 looks at quantization, a technique for representing numbers in more compact formats (like 8-bit integers or floating-point) to save memory and speed things up. Task 2 dives into "attention," a key mechanism behind modern AI models, and compares different optimized versions of it. Task 3 walks through softmax, a common operation, showing how a basic approach can be improved into a more memory-efficient version. Finally, Task 4 brings it all together by building a small inference engine, essentially a lightweight program that runs AI models, using a mix of Python, CUDA, and C++. This project is aimed at people who already understand AI concepts at a high level and want to peek under the hood at how the actual computations are optimized for GPU hardware. A student learning about GPU programming, an engineer curious about how tools like FlashAttention work, or a founder who wants to understand why AI inference can be slow or expensive could all get value from working through these exercises. The progression from basic to optimized approaches makes it a structured learning path rather than a reference tool. The README is quite sparse, it lists the tasks but doesn't provide step-by-step instructions or deep explanations, so you'd likely need some prior familiarity with GPU programming concepts to get the most out of it. The inclusion of CUDA and C++ alongside Python signals that this is for people willing to work close to the metal, not just call high-level libraries.