Lecture 6: Kernels, Triton

What this lecture covers

The lecture teaches the practical workflow for writing high-performance GPU code: benchmark end-to-end runtime, profile to find bottlenecks, then write fused kernels to remove them. It uses a single running example, a GELU nonlinearity, implemented four ways (naive PyTorch, hand-written C++/CUDA, Triton, and torch.compile), comparing their speed and profiler output at each step.

This lecture follows directly from the GPU architecture lecture and prepares students for assignment two, where they write a Triton kernel for Flash Attention 2. After watching, you should be able to write a correct benchmarking harness, read PyTorch's profiler and Nsight Systems output, and understand why writing a fused kernel (in CUDA or Triton) can turn several slow, separate operations into one fast one.

Key ideas

• Always profile before optimizing: guessing where the bottleneck is wastes effort; a profiler shows exactly which operations and kernels consume time.
• Correct benchmarking requires warm-up and synchronization: the first run of any GPU code pays a one-time compilation and initialization cost, and because the CPU dispatches work to the GPU asynchronously, you must call torch.cuda.synchronize() to measure actual GPU execution time rather than CPU dispatch time.
• CPU and GPU run asynchronously: the CPU can queue many CUDA kernels ahead of the GPU's actual execution, which is why Python's overhead usually doesn't bottleneck GPU-bound code, unless something (like printing a loss value) forces a synchronization point.
• Naive multi-op code launches many separate kernels: an expression like a GELU approximation written out with several tanh, multiply, and add calls dispatches one CUDA kernel per operation, each paying memory round-trip and launch overhead.
• Kernel fusion collapses many operations into one: writing a custom CUDA or Triton kernel that computes the whole formula in one pass eliminates the redundant memory traffic, closing most of the gap to PyTorch's built-in fused implementation.
• Triton trades some low-level control for productivity: you program at the level of thread blocks rather than individual threads, and Triton automatically handles memory coalescing and shared-memory management, while still reaching performance close to hand-written CUDA.
• torch.compile often matches hand-written kernels: PyTorch's JIT compiler can automatically fuse simple elementwise chains and pick optimized matrix-multiply kernels, frequently making manual kernel-writing unnecessary except for genuinely novel operations.
• Nsight Systems reveals CPU/GPU interleaving: it shows that the CPU can run many steps ahead of GPU execution, and that operations like a print statement mid-training-loop force a synchronization that can stall the pipeline.

Walkthrough

Review and arithmetic intensity (2:05)

The lecture briefly reviews SMs, thread blocks, warps, and the memory hierarchy from the prior lecture, then introduces arithmetic intensity: since compute has scaled much faster than memory bandwidth, code should aim for more flops per byte moved. Matrix multiplication, done well, is compute-bound; most other operations tend to be memory-bound.

Benchmarking correctly (7:07)

Using a simple MLP as a running example, the lecture builds a benchmarking helper that performs warm-up iterations and calls torch.cuda.synchronize() before and after timing, explaining that skipping either step gives misleadingly fast or meaningless numbers. It demonstrates the expected super-linear-then-linear scaling of matrix multiply time with size, and linear scaling of MLP runtime with the number of layers or steps.

Profiling with PyTorch's built-in profiler (17:13)

The lecture switches from coarse benchmarking to fine-grained profiling, showing how a simple add or matrix multiply dispatches through PyTorch's C++ interface (ATen) down to specific CUDA kernels (such as a CUTLASS matmul kernel), with visible time splits between CPU dispatch, kernel launch, and GPU execution. More complex operations like cdist and GELU are shown decomposing into multiple underlying kernels, letting you see exactly which sub-operation dominates GPU time.

Nsight Systems and the CPU/GPU relationship (33:22)

Using Nvidia's Nsight Systems profiler with NVTX code annotations, the lecture visualizes an MLP training loop and shows the CPU queuing CUDA kernels for later steps well before the GPU has finished executing earlier ones. It demonstrates concretely how adding a print statement that reads a loss value forces a cudaStreamSynchronize call, which stalls the CPU's ability to run ahead and can bottleneck training if overused.

Writing a fused GELU kernel in C++/CUDA (44:25)

The lecture shows that a "manual" GELU written with separate PyTorch operations (multiply, tanh, cube, add) is about eight times slower than PyTorch's built-in fused GELU because each operation is a separate kernel launch and memory round-trip. It then walks through writing a CUDA kernel from scratch: a host-side wrapper function that checks the tensor is contiguous and on the GPU, computes grid and block dimensions, and a __global__ kernel function where each thread computes its own index and applies the full GELU formula in one pass. This closes most of the performance gap, from about 8.1ms down to roughly 1.8ms versus PyTorch's 1.1ms.

The same kernel in Triton (1:01:40)

The lecture rewrites the identical GELU kernel in Triton, a Python-embedded DSL where you program at the granularity of a thread block rather than individual threads, using vectorized loads and a boundary mask instead of manual index arithmetic. Triton achieves essentially the same performance as the hand-written CUDA kernel with substantially less code, and the lecture briefly shows the compiled PTX output to illustrate how Triton automatically batches memory loads for coalescing.

torch.compile and a softmax kernel (1:11:47)

torch.compile is shown automatically fusing the naive multi-op GELU into a single efficient kernel, reaching performance close to the hand-written versions with no manual kernel-writing at all; the lecture suggests it's usually the first thing to try before writing custom kernels. The lecture closes with a Triton softmax kernel, which requires a reduction (summing across a row) rather than a pure elementwise map: the design assigns one row per thread block so the whole row fits in fast on-chip memory, letting the max-subtract-exponentiate-sum-divide sequence run without any extra global memory round trips.

Before you watch

• Watch the prior GPU architecture lecture in this course, since this lecture assumes familiarity with SMs, thread blocks, warps, and the memory hierarchy.
• Be comfortable reading basic CUDA-style pseudocode (pointers, thread and block indices) and Python/PyTorch tensor operations.
• Familiarity with the GELU activation function and softmax will help follow the worked kernel examples.

Check your understanding

1. Why must you call torch.cuda.synchronize() before measuring elapsed time in a GPU benchmark?
2. Why is a GELU implementation written as several separate PyTorch operations much slower than a single fused kernel, even though they compute the same result?
3. How does the CPU "running ahead" of the GPU affect the perceived cost of Python overhead, and what can force the CPU to stop running ahead?
4. What is the key structural difference between writing a kernel in raw CUDA C++ versus in Triton?
5. Why does the softmax kernel assign one entire row to a single thread block instead of one element per thread?