Modern GPU Programming for MLSys

The performance of modern AI workloads is determined by the quality of core GPU kernels such as Attention kernel, LLM prefill/decode, and MoE layers. This guide goes beyond simple optimization techniques to cover how to understand and leverage the complex memory structures and execution units of the latest GPU architectures.

Key learning content includes the following:

• Understanding GPU hardware: Hardware-level concepts such as data layout, asynchronous data movement (TMA), and asynchronous coordination (mbarriers)
• Using TIRx: Using a Python DSL that allows writing executable code while retaining hardware control
• Implementing core kernels:
  • GEMM: Optimization applying TMA pipelining, persistent scheduling, warp specialization, and more
  • Flash Attention 4: Implementation of a state-of-the-art attention kernel including online softmax rescaling, GQA (Grouped Query Attention), and more

This material primarily targets Blackwell-generation GPUs and provides a systematic, step-by-step learning path for building high-performance kernels.