Modern GPU Programming For MLSys

Machine learning systems power many of today’s AI workloads. As models grow and deployment settings become more complex, end-to-end performance increasingly depends on a small number of critical GPU kernels. Attention, LLM prefill and decode, low-precision block-scaled GEMM, fused MoE layers, and other large fused kernels directly affect both training and serving speed.

Making these kernels fast requires more than a list of optimization tricks. Recent GPU architectures introduce richer memory spaces, new data-movement mechanisms, and increasingly specialized execution units. Using them effectively requires both a clear understanding of how the hardware executes a program and practical knowledge of how a basic kernel evolves into a high-performance implementation. This book develops both.

The book proceeds from hardware to programming model to complete kernels. It first introduces GPU organization and execution, then presents the programming model used throughout the book, and finally builds high-performance kernels step by step. The main target is NVIDIA Blackwell, and the running examples are General Matrix-Matrix Multiplication (GEMM) and FlashAttention. Along the way, the book develops the key ideas behind GPU optimization: data layout, asynchronous data movement, and asynchronous coordination.

The material grows out of the Machine Learning Systems course series at Carnegie Mellon University. The examples use the TIRx Python DSL so that the ideas can be studied, run, and verified in real kernels. TIRx keeps hardware-level choices explicit, making it possible to reason about control flow, memory access, and synchronization while working with runnable code.

This book is open source. Contributions, corrections, and examples are welcome through the GitHub repository.

How This Book Is Organized

• Part I, Understanding the GPU. This part introduces the overall organization of the GPU, general techniques for writing fast kernels, and key concepts such as data layout, asynchronous memory operations, and coordination. It builds the hardware intuition that the rest of the book relies on.

• Part II, TIRx Overview. This part introduces the key elements of TIRx, which serve as the foundation for the code examples throughout the book.

• Part III, GEMM: Tiled to SOTA. A complete guide to optimizing a tiled GEMM, built up through TMA pipelining, persistent scheduling, warp specialization, and 2-CTA clusters.

• Part IV, Flash Attention 4. A complete attention kernel built from the Part III techniques: two MMAs with softmax between them, online-softmax rescaling, causal masking, and GQA.

• Appendices. TIRx language reference, a reproducible GPU benchmarking and profiling workflow, compiler internals, and a guide to debugging asynchronous kernels.

Part I, Understanding the GPU

• GPU Execution Model
• What Makes a Kernel Fast
• Data Layout and Its Notation
• The Evolution of Tensor Core Data Layouts
• Async Data Movement: TMA
• Blackwell Tensor Core: tcgen05.mma
• Tensor Memory (TMEM)
• Async Coordination: mbarrier
• Advanced Scheduling: Cluster Launch Control

Part II, TIRx Overview

• Introduction to TIRx
• TIRx Layout API

Part III, GEMM: Tiled to SOTA

• Building a Tiled GEMM
  • GEMM
  • Optimization Path
  • Step 1: Sequential Single-Tile GEMM
  • Step 2: K-Loop Accumulation
  • Step 3: Spatial Tiling (Multi-CTA)
  • Exercises
• Pipelining GEMM with TMA
  • Step 4: TMA Async Load
  • Step 5: Software Pipeline (PIPE_DEPTH=2)
  • Step 6: Persistent Kernel + Tile Scheduler
  • Exercises
• Scaling GEMM with Warp Specialization and Clusters
  • Step 7: Warp Specialization
  • Step 8: Two-CTA Cluster
  • Step 9: Multi-Consumer Warp Specialization
  • End-to-End Results
  • Exercises

Part IV, Flash Attention 4

• Flash Attention 4
  • Algorithm Structure
  • Tile Primitive Data Flow
  • Warp Roles and Scope
  • Conventions for Reading the Code
  • QKᵀ MMA and PV MMA
  • TMEM Layout and Reuse
  • Key Barrier Protocols
  • Pipeline Timeline
  • Rescaling and Writeback
  • Causal Masking
  • GQA Support
  • Tile Scheduling
  • Compile and Verify
  • Exercises

Appendices

• Overview
• TIRx Language Reference
• Measuring and Analyzing GPU Kernel Performance
• Compiler Internals
• Debugging Warp-Specialized Kernels