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

High-Performance Code for GPUs: A Deep Dive

Key Concepts:

• Streaming Multiprocessors (SMs): The fundamental compute units within a GPU.
• Thread Blocks: Collections of threads scheduled on a single SM.
• Warps: Groups of 32 threads executed simultaneously within an SM.
• Arithmetic Intensity: The ratio of floating-point operations (FLOPs) to bytes of memory movement. High arithmetic intensity is desirable for performance.
• Benchmarking: Measuring the wall-clock time of code execution.
• Profiling: Analyzing code execution to identify performance bottlenecks.
• Kernel Fusion: Combining multiple operations into a single GPU kernel to reduce memory transfers.
• CUDA: NVIDIA's parallel computing platform and programming model.
• Triton: A domain-specific language for writing high-performance GPU kernels in Python.
• PTX: Parallel Thread Execution, a low-level assembly language for NVIDIA GPUs.
• Torch Compile: PyTorch's just-in-time (JIT) compiler for optimizing PyTorch code.

1. GPU Architecture Review

• GPUs like A100 and H100 consist of multiple SMs.
• Each SM contains numerous compute units (INT32, FP32).
• SMs launch many threads.
• Memory hierarchy: DRAM (global memory - large, slow), caches (faster), register file (very fast, per-thread).
• Execution model: Thread blocks are scheduled on SMs. Threads within blocks perform computations.
• Thread blocks can communicate via shared memory within an SM (fast). Communication across blocks is expensive.
• Threads are grouped into warps (32 threads) executed simultaneously. Aim for equal computation across warps.
• Ideally, the number of thread blocks should be much greater than the number of SMs.
• Arithmetic Intensity: Maximize FLOPs per byte of memory movement. Matrix multiplication is compute-bound (if done cleverly); most other operations are memory-bound.

2. Benchmarking and Profiling: The Foundation of High-Performance Code

• Benchmark and profile your code! Don't guess bottlenecks.
• Profiling reveals bottlenecks and hardware behavior.
• End-to-end benchmarking is crucial due to library versions, hardware variations, and microcode.

2.1 Benchmarking Methodology

• Purpose: Compare implementations (Triton vs. C++ vs. PyTorch), understand scaling behavior.
• Benchmark Function:
  • Takes a function to benchmark (run).
  • Performs warm-up iterations to avoid measuring startup costs (compilation, initialization).
  • Executes trials to get stable measurements.
  • Uses torch.cuda.synchronize() to ensure the CPU and GPU are in the same state before and after timing. The CPU runs ahead and dispatches CUDA kernels to the GPU. torch.cuda.synchronize() forces the CPU to wait for GPU execution to complete.
  • Averages the results of multiple trials.
• Example: Benchmarking matrix multiplies shows super-linear scaling with matrix size. Small matrices show constant overhead due to CPU-GPU transfer and kernel launch.
• Example: Benchmarking an MLP shows linear scaling with the number of steps and layers.

2.2 Profiling Methodology

• Purpose: Identify where time is spent within a function.
• PyTorch's built-in profiler provides basic profiling within the Python/PyTorch environment.
• Profiling reveals low-level CUDA calls.
• Example: Profiling torch.sleep shows 100% of the time spent on cudaDeviceSynchronize.
• Example: Profiling adding two matrices reveals calls to a10 (PyTorch's C++ interface), a vectorized element-wise kernel, cudaLaunchKernel, and cudaDeviceSynchronize.
• Example: Profiling matrix multiplication reveals calls to a10 map mole and a specific cutlass kernel with tile sizes and block counts.
• Different matrix sizes lead to different CUDA kernels being dispatched.
• Example: Profiling torch.cdist (pairwise Euclidean distance) shows decomposition into a10 mm mole, a10 pow, and sum. Matrix multiplication (gmm) dominates GPU time.
• Example: Profiling GELU and Softmax shows fused operators with dedicated CUDA kernels.

2.3 Advanced Profiling with NVIDIA Nsight Systems (NSIS)

• NSIS provides a detailed view of GPU behavior, including CPU and GPU activity.
• nvtx markers annotate code for easier profiling.
• NSIS reveals the asynchronous execution model: the CPU runs ahead, queuing CUDA kernels to the GPU.
• Printing losses between iterations forces CPU-GPU synchronization, potentially introducing a CPU bottleneck.
• NSIS allows analyzing kernel execution times and identifying performance bottlenecks.
• NSIS shows that Python's performance is not a bottleneck because the CPU can queue commands to the GPU.

3. Kernel Fusion: Optimizing GELU

• Kernel fusion combines multiple operations into a single GPU kernel to reduce memory transfers.
• Example: Implementing GELU naively in PyTorch leads to multiple CUDA kernel launches and poor performance.
• The fused PyTorch implementation of GELU is significantly faster.

3.1 CUDA Implementation of GELU

• CUDA (C++ API) allows writing custom GPU kernels.
• CUDA kernels are functions that are automatically called on all elements of a vector or matrix.
• CUDA uses a grid of thread blocks, each containing multiple threads.
• Each thread needs to compute its coordinate within the matrix or vector.
• Out-of-bounds checking is crucial in CUDA kernels.
• CUDA Code Structure:
  • Kernel Function (__global__ void geluKernel): Performs the GELU computation for each element.
  • Wrapper Function (gluFunction): Orchestrates the kernel launch, allocates memory, and handles input validation.
• The CUDA implementation of GELU significantly improves performance compared to the naive PyTorch implementation.

3.2 Triton Implementation of GELU

• Triton is a domain-specific language for writing high-performance GPU kernels in Python.
• Triton manages memory coalescing, shared memory, and thread scheduling.
• Triton code is written at the SM-centric level.
• Triton allows stepping through and debugging code in Python.
• Triton kernels can outperform PyTorch implementations.
• Triton Code Structure:
  • Kernel Function (@triton.jit def gelu_kernel): Performs the GELU computation for a block of elements.
  • Wrapper Function (gluFunction): Orchestrates the kernel launch, allocates memory, and handles input validation.
• Triton uses vectorized operations to operate on entire blocks at once.
• Triton compiles into low-level PTX code.
• The Triton implementation of GELU achieves similar performance to the CUDA implementation.

3.3 PTX Code Analysis

• PTX code reveals the low-level operations performed by the GPU.
• PTX code shows register allocation, memory loads and stores, and floating-point operations.
• PTX code shows that each thread operates on multiple values at a time.

3.4 Torch Compile Implementation of GELU

• Torch Compile is PyTorch's JIT compiler.
• Torch Compile can automatically perform kernel fusion and other optimizations.
• Torch Compile can generate Triton code under the hood.
• Torch Compile achieves similar or better performance than the manual CUDA and Triton implementations.

4. Triton Implementation of Softmax

• Softmax requires a reduction operation (sum across elements).
• A simple design is to have each SM handle a single row.
• The block size should be the number of columns plus padding.
• The number of blocks should be the number of rows.
• Triton code for Softmax is similar to normal Python code, with load and store operations.
• Torch Compile can often outperform native PyTorch implementations, especially when it knows the shapes and sizes of operations.

5. Conclusion

• Benchmarking and profiling are essential for writing high-performance GPU code.
• Kernel fusion is a crucial optimization technique.
• CUDA, Triton, and Torch Compile provide different levels of abstraction for writing GPU kernels.
• Modern JIT compilers like Torch Compile are often sufficient for simple operator fusion.
• Triton is useful for more complex optimizations and for gaining a deeper understanding of GPU behavior.
• CUDA provides the most control but requires more manual effort.