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

Key Concepts

• GPU Architecture: Streaming Multiprocessors (SMs), Registers, L1 Cache, Shared Memory, L2 Cache, and High-Bandwidth Memory (HBM).
• Programming Model: Threads, Thread Blocks (Concurrent Thread Arrays/CTAs), and Grids.
• Performance Metrics: Occupancy, Bank Conflicts, Memory Coalescing, and Arithmetic Intensity.
• Triton: A domain-specific language for writing GPU kernels that abstracts away low-level thread synchronization and memory management.
• Kernel Fusion: Combining multiple operations into a single kernel to minimize HBM read/write overhead.
• Tiling: A strategy to break large computations into smaller blocks that fit into fast, local memory (Shared Memory).
• PTX (Parallel Thread Execution): The low-level intermediate assembly language for NVIDIA GPUs.

---

1. GPU Architecture and Programming Model

The lecture emphasizes that while the programming model (threads, blocks, grids) is elegant, performance is dictated by hardware constraints.

• Memory Hierarchy: Registers and L1/Shared memory are fast but small and local to the SM. HBM is large but slow.
• Warps: Groups of 32 threads that execute instructions in lockstep. Control divergence (branching) causes serialization, reducing efficiency.
• Latency Hiding: SMs switch between resident warps with zero cost to hide the latency of slow HBM memory accesses.
• Occupancy: The ratio of active warps to the maximum supported. High occupancy isn't always optimal; "fattening" threads (thread coarsening) to do more work per thread can sometimes improve performance despite lower occupancy.

2. Benchmarking and Profiling

The instructor advocates for a "measure-modify-measure" cycle:

• Benchmarking: Use CUDA events (start_event, end_event) with torch.cuda.synchronize() to ensure accurate timing. Always include a "warm-up" phase to account for lazy compilation.
• Profiling: Essential for identifying bottlenecks. PyTorch’s profiler reveals which specific CUDA kernels are invoked. Different tensor dimensions trigger different optimized kernels (e.g., CUTLASS libraries for matrix multiplication).

3. Kernel Optimization: GELU Example

• Naive Implementation: Multiple PyTorch operations result in multiple kernel launches, forcing data to be read from and written to HBM repeatedly.
• Kernel Fusion: Built-in or torch.compile versions fuse operations into a single kernel, reading from HBM once and writing once, significantly reducing memory traffic.

4. Triton Programming Methodology

Triton allows developers to think in terms of Thread Blocks rather than individual threads.

• Workflow:
  1. Allocate output tensor in HBM.
  2. Define a grid of thread blocks.
  3. Inside the kernel: Identify the block (program_id), calculate memory offsets, load data from HBM into local/shared memory, perform computation, and store results back to HBM.
• Handling Large Data: When data exceeds block size, use a loop within the thread to iterate over "tiles" of data, accumulating results in registers or shared memory before writing back.

5. Matrix Multiplication (MatMul) and Tiling

• The Problem: Naive MatMul has low arithmetic intensity (constant ratio of operations to memory reads), leading to HBM bottlenecks.
• The Solution (Tiling):
  • Load tiles of matrices A and B into shared memory.
  • Perform the dot product on these tiles.
  • Accumulate partial sums in shared memory.
  • Write the final tile to HBM.
• Arithmetic Intensity: Tiling increases arithmetic intensity to the order of the tile size, drastically improving performance by reducing redundant HBM reads.

Notable Quotes

• "The programming model is actually very nice... but in practice, the performance is very sensitive to the hardware."
• "You should always just measure what's going on in your code and figure out what the bottlenecks are before you start writing kernels."
• "Triton makes this very easy... if you can fit anything through a block, you can just write normal PyTorch almost."

Synthesis

The core takeaway is that high-performance GPU programming requires balancing the abstraction of the programming model with the physical realities of the hardware. By using Triton to manage tiling and kernel fusion, developers can minimize expensive HBM transactions and maximize the utilization of fast, local memory (Shared Memory/Registers). Benchmarking and profiling are not optional; they are the primary tools for understanding how high-level code maps to the underlying hardware execution.