I Tried Getting Closer to the GPU With Triton

A developer explores Triton, a Python-based language and compiler for writing GPU kernels, to understand low-level GPU programming and performance. The developer explains how Triton abstracts GPU threads into blocks of data, reducing memory movement and improving efficiency compared to eager PyTorch operations. The article details the developer's first experiment with vector addition and the importance of GPU memory hierarchy.

I've been trying to understand what actually happens when a piece of ML code reaches the GPU. Not the usual model.cuda and GPU go brrrr kind of understanding, but the actual stuff underneath it — threads, warps, memory, kernel launches, and why some operations are ridiculously fast while others suddenly become slow. That rabbit hole eventually led me to Triton. At first, I was wondering why anyone would even bother writing custom GPU kernels when PyTorch already gives us highly optimized operations. But the more I looked into it, the more I realized that GPU performance isn't always about how much computation you're doing. A lot of the time, it's about how you're moving data around. CUDA gives you an insane amount of control over the GPU, but that control comes with a lot of complexity. You have to think about threads, warps, synchronization, memory access, and a bunch of other low-level details. That's great if you're comfortable with CUDA, but if you're coming from Python and PyTorch, it can be a pretty steep jump. PyTorch makes this much easier through its higher-level abstractions. You can just write something like torch.softmax and let the framework handle everything underneath. The problem is that when you're executing operations eagerly, individual operations can result in separate kernel launches and repeated reads and writes to global GPU memory. The code is easy to write, but there can be a lot of unnecessary movement of data happening underneath. And this is where Triton becomes interesting. Triton is a Python-based language and compiler for writing GPU kernels. What I found interesting about it is that instead of forcing you to think about individual GPU threads, it lets you think in terms of blocks of data. You basically describe the work that one program instance should perform, and Triton takes care of launching many of those program instances across the GPU. That changed the way I started thinking about GPU programming. Before getting into Triton, I had to understand the GPU itself a little better. A GPU is built around massive parallelism. Instead of having a small number of powerful CPU cores, GPUs have huge numbers of lightweight threads that can execute work simultaneously. Those threads are organized into warps and blocks, and the actual hardware that executes these warps is organized into Streaming Multiprocessors. But the part that really matters for performance is the memory hierarchy. GPU memory isn't just one big pool where everything costs the same to access. You have registers, shared memory, caches, and global memory. Registers are extremely close to the computation and very fast, while global memory is much larger but significantly more expensive to access. So if your algorithm keeps loading data from global memory, performing a tiny amount of work, writing it back, and then loading the same data again, you're potentially wasting a huge amount of time. The basic idea I started taking away was pretty simple: move data as little as possible and reuse it as much as possible. My first Triton experiment was just vector addition. Nothing fancy. Given two arrays, add them together. But even this simple example helped me understand the programming model. Instead of launching a thread for every individual element, I could define a block size and let each Triton program instance handle a chunk of the input. For example, if I have 1024 elements and a block size of 128, I can think of the workload as eight program instances. The first handles elements 0 to 127, the second handles 128 to 255, and so on. Triton gives each program instance an ID, and I can use that ID to calculate which section of memory it needs to work on. The next thing that confused me initially was pointers. Triton kernels don't work with tensors in the same high-level way normal PyTorch code does. They work with pointers to GPU memory. You calculate offsets from those pointers, load the values you need, perform the computation, and then store the results back. This is also where masks become important. If your input size isn't perfectly divisible by your block size, the final program instance might be assigned elements that don't actually exist. You obviously don't want the GPU trying to access memory outside the tensor, so you use a mask to make sure only valid elements are loaded and stored. After vector addition, I moved on to something more interesting: softmax. A straightforward softmax implementation involves several operations. You find the maximum value, subtract it for numerical stability, calculate the exponential, sum the results, and finally divide by the sum. Each of these operations can involve reading and writing data. So conceptually, you can end up with something like: load data, calculate something, write it back, load it again, calculate something else, write it again, and repeat. The data is constantly travelling between global memory and the compute units. Instead of doing that, we can fuse the operations into one kernel. The idea is to load the data once, keep it on-chip while performing the different operations, and only write the final result back to global memory. This is one of the things that made Triton click for me. The code isn't necessarily about doing more computation. It's about avoiding unnecessary memory traffic. Then came matrix multiplication, which made the whole concept of tiling much more obvious. When multiplying two matrices, a naive implementation can repeatedly load the same pieces of the input matrices from global memory. But if you're computing multiple nearby output values, you're often reusing the same data. So instead of calculating one output element at a time, you divide the matrices into smaller tiles. You load a small tile of each input matrix into faster on-chip memory, perform as much computation as possible using those tiles, accumulate the result, and then move on to the next tile. The fundamental idea is basically: load less often, reuse more. That sounds ridiculously simple, but it is one of the core ideas behind high-performance matrix multiplication on GPUs. Once I understood tiling, some of the other Triton concepts started making more sense too. Program IDs can be used to determine which output tile a program instance is responsible for. Strides let you calculate where elements actually live in memory. Masks handle boundaries. The accumulator keeps partial results while working through different chunks of the reduction dimension. Then there is scheduling and cache locality. If multiple output tiles reuse the same pieces of input data, it makes sense to execute those tiles close together so that the data has a better chance of remaining in cache. Triton provides mechanisms for grouping program instances to take advantage of this kind of locality. And then you get to autotuning. Choosing the perfect block size, number of warps, pipeline stages, and other parameters manually isn't always easy. Different workloads and different GPUs can behave differently. Triton's autotuning system lets you provide multiple configurations and benchmark them to find a better configuration for a particular workload. This is another thing I found pretty interesting. Instead of assuming there is one perfect configuration, you can basically let the system search through several possibilities and keep the one that performs best. Eventually I started benchmarking the kernels instead of just looking at the code and assuming the optimized version was faster. And that was probably another important lesson. GPU optimization isn't about writing code that looks complicated. It's about understanding what the hardware is actually spending time doing. Sometimes you're compute-bound. Sometimes you're memory-bound. Sometimes kernel launch overhead matters. Sometimes you're repeatedly moving data that could have stayed on-chip. And sometimes the library implementation you're trying to beat is already so optimized that your custom kernel isn't going to magically win. That's probably the biggest thing I took away from learning Triton. Before this, I used to think about GPU optimization mostly as “How do I make the GPU perform more computation?” Now I find myself asking a slightly different question: “Why am I moving this data in the first place?” That shift in perspective is probably more valuable than any individual Triton kernel I wrote. I'm definitely not claiming to be a GPU programming expert after doing this. If anything, the more I learn, the more I realize how much more there is to understand — occupancy, register pressure, memory coalescing, Tensor Cores, persistent kernels, FlashAttention, compiler behavior, and a lot more. But at least now those topics don't feel completely alien. And that's honestly why I wanted to write this. I didn't want this to be another post that throws a complicated Triton kernel at you and says, "look how fast this is." I wanted to understand why the kernel is written that way and what the GPU is actually doing underneath it. Because once you start thinking about GPUs in terms of data movement, reuse, parallelism, and locality, GPU programming starts making a lot more sense.