Triton Basics: Implementing Element-wise Addition & Simplified GEMM

Implementing an Element-wise Kernel

Let's start by implementing a kernel for an operation slightly more complex than the vector addition we've seen before: the GELU (Gaussian Error Linear Unit) activation function. The exact formula is less important than the pattern of implementation, which applies to any element-wise operation (e.g., ReLU, SiLU, bias-add).

The logic is identical to the vector addition pattern from our previous lesson:

1. The host wrapper calculates the grid size and launches the kernel.
2. Each program instance gets a unique pid using tl.program_id.
3. Each instance computes the starting offset and creates a block of offsets for its tile using tl.arange.
4. It loads the data for that tile from global memory into on-chip registers/SRAM, using a mask for boundary checks.
5. It performs the computation on the tile. Since the data is now on-chip, this step is very fast.
6. It stores the resulting tile back to global memory, again using a mask.

The following video provides an excellent walkthrough of implementing a GELU kernel in Triton, following this exact pattern.

As you can see, once you understand the basic template, writing custom element-wise kernels becomes a highly productive task. You simply change the mathematical expression in the middle of the kernel, and Triton's compiler handles the low-level hardware mapping.

From Element-wise to Tiled GEMM

Element-wise operations are memory-bound; their performance is limited by the speed at which data can be fetched from global memory. The arithmetic intensity is low because each loaded element is used for only one or two calculations.

Matrix multiplication is different. It is a compute-bound operation with high arithmetic intensity, meaning many computations are performed for each element loaded. A naive implementation where each thread computes one element of the output matrix, $C_{ij}$, would be disastrously inefficient. Each thread would repeatedly load the same rows of matrix A and columns of matrix B from slow global memory, failing to exploit the GPU's fast on-chip memory.

The solution is tiling (or blocking). Instead of computing one element at a time, each Triton program instance will compute an entire tile of the output matrix C.

Here's the core idea:

1. Partition: The output matrix C is partitioned into 2D tiles of size BLOCK_SIZE_M x BLOCK_SIZE_N.
2. Assign: Each Triton program instance is responsible for computing a single tile of C.
3. Iterate and Accumulate: To compute its tile, the program instance iterates along the shared K dimension in chunks of size BLOCK_SIZE_K. In each step of the loop, it:
   • Loads a BLOCK_SIZE_M x BLOCK_SIZE_K tile from A into fast on-chip memory.
   • Loads a BLOCK_SIZE_K x BLOCK_SIZE_N tile from B into fast on-chip memory.
   • Performs a matrix multiplication on these two small, on-chip tiles.
   • Adds the result to an accumulator tile that lives in registers.
4. Store: After the loop over K is finished, the final accumulator tile is written to its corresponding position in the global memory of matrix C.

This strategy maximizes data reuse. A tile loaded from matrix A is used for all BLOCK_SIZE_N columns of the output tile, and a tile from B is used for all BLOCK_SIZE_M rows. This keeps the compute units fed with data from fast memory, which is the key to high performance.

This next resource provides a fantastic animated visualization and explanation of why tiling is so effective.

Implementing Tiled GEMM in Triton

Now, let's see how this tiled algorithm translates into Triton code. The implementation requires two key steps beyond what we saw in the element-wise kernel: multi-dimensional pointer arithmetic and an accumulation loop.

The official Triton documentation provides a full, high-performance GEMM kernel. We will walk through the key parts of this implementation.

Let's break down the most crucial components from that code:

• 2D Pointers: Unlike the 1D offsets in our GELU kernel, for a tile of matrix A, we need a 2D block of pointers. This is constructed using NumPy-style broadcasting:

  
  
  
  # From the Triton tutorial
  offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
  offs_k = tl.arange(0, BLOCK_SIZE_K)
  a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
  

  Here, offs_am[:, None] (a column vector of row offsets) is broadcast-added to offs_k[None, :] (a row vector of column offsets) to create a 2D grid of offsets, which are then added to the base pointer a_ptr.

• The Accumulator and tl.dot:

  
  
  
  # From the Triton tutorial
  accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
  for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
      a = tl.load(a_ptrs, ...)
      b = tl.load(b_ptrs, ...)
      accumulator = tl.dot(a, b, accumulator) # Fused-Multiply-Add
      a_ptrs += BLOCK_SIZE_K * stride_ak
      b_ptrs += BLOCK_SIZE_K * stride_bk
  

  The accumulator is initialized to zeros in float32 for better precision. Inside the loop, tl.dot performs the matrix multiplication on the local tiles a and b and adds the result to the accumulator. This single function call is where the bulk of the computation happens. It is a highly optimized intrinsic that can map directly to the GPU's Tensor Core hardware. After the dot product, the pointers for A and B are advanced to the next slice in the K dimension.

The Triton Compilation Pipeline

Your interest in low-level systems warrants a brief look at what happens when you run a @triton.jit function. Triton is not an interpreter; it's a real compiler. The Python code you write is parsed and progressively "lowered" through several stages of intermediate representations (IR) until it becomes machine code.

This process ensures that the high-level, Pythonic code you write is translated into the low-level, highly-optimized instructions that the GPU actually executes. The compiler handles the complex register allocation, instruction scheduling, and memory access patterns that you would otherwise need to manage by hand in CUDA.

Conclusion

In this lesson, we have successfully translated Triton's programming model into practical, working AI kernels. We started with a familiar element-wise pattern for a GELU activation and then progressed to the far more complex and performance-critical tiled matrix multiplication.

Key Takeaways:

• Element-wise Kernels: Follow a simple template of defining a tile, loading it, computing on it, and storing the result.
• Tiled GEMM: The key to high performance is maximizing data reuse from fast on-chip memory. This is achieved by having each program instance compute a tile of the output matrix C while iterating through the K dimension in blocks.
• Core GEMM Components: The implementation relies on 2D pointer arithmetic to address tiles, a register-based accumulator to hold partial sums, and the tl.dot intrinsic to perform the on-chip matrix multiplication.
• Triton as a Compiler: Triton is not just a library; it's a sophisticated compiler that transforms high-level Python code into optimized low-level machine code, abstracting away immense complexity.

We have now written a high-performance GEMM kernel in about 25 lines of Python. In the next lesson, we will do a direct, critical comparison of this implementation against a raw CUDA C++ version. We will analyze the code complexity, developer productivity, and benchmark the performance to fully appreciate the trade-offs Triton offers.