chapter two

2 Write your first kernel in Triton

 

This chapter covers

  • Writing a full Triton kernel
  • Mapping work with program IDs
  • Vectorized loads, stores, and masks
  • Launching grids from Python
  • Tracing Triton’s JIT compilation
  • Benchmarking kernel performance

You now have a picture of how GPUs organize computation: blocks of threads running across Streaming Multiprocessors, a layered memory system, and a compiler that handles the low-level details for you. That foundation is enough to start writing real code. In this chapter, we put it to use and write your first Triton kernel. We begin with a complete, working kernel for a simple task: vector addition. Then we break it down line by line to explain Triton’s syntax, its programming model, and the compiler that drives it. We finish by benchmarking the kernel against pure Python, NumPy, and PyTorch to see the difference that parallel execution makes. By the end of this chapter, you will know how to write, launch, debug, and benchmark a basic Triton kernel, and how that kernel flows through Triton’s compilation pipeline.

Note

This chapter is a foundation for everything that follows, so it deserves careful attention. We will walk through every line of code together. You only need an intermediate understanding of Python to follow along. If you need a refresher, The Quick Python Book by Naomi Ceder is a reliable resource.

2.1 The Triton programming model

2.1.1 Thinking in blocks, not threads

2.1.2 The Single Program, Multiple Data model

2.2 A complete Triton kernel: vector addition end to end

2.2.1 Imports and the JIT decorator

2.2.2 Function signature and pointers

2.2.3 Getting a unique program ID

2.2.4 Calculating the data chunk

2.2.5 Computing the element offsets

2.2.6 Handling edges with a mask

2.2.7 Loading data and performing the computation

2.2.8 The launch grid

2.2.9 Kernel execution flow

2.2.10 Benchmarking the kernel

2.3 A high-level language with a low-level compiler

2.3.1 The compilation funnel (advanced)

2.3.2 Practical inspection of compilation stages

2.4 Summary