Stanford CS336 Language Modeling from Scratch | Spring 2025 | Lecture 6: Kernels, Triton

In this lecture, the speaker details how to write high-performance code for GPUs, particularly for Assignment 2, which involves profiling and writing a Triton kernel for FlashAttention2. The lecture includes a review of GPU components, benchmarking, and profiling techniques using CUDA kernels in C++, Triton, and PyTorch's JIT compiler. The speaker emphasizes the importance of benchmarking and profiling code to identify bottlenecks, demonstrating with a simple MLP example. The lecture further explores profiling with PyTorch's built-in profiler and NVIDIA's Nsight Systems, illustrating how to analyze GPU behavior and performance, including the interaction between the CPU and GPU. The speaker also covers kernel fusion, writing CUDA kernels in C++, and using Triton for GPU programming, comparing performance across different implementations and concluding with a demonstration of writing a fast Triton implementation of Softmax.