How to get a matmul kernel into the top percentile (worklog)

Around Christmas last year I was looking for a fun side project and ended up looking more into CUDA. I’ve found LeetGPU to be extremely helpful in having a sandbox to test kernels for different problems on different hardware without having to spin up your own VM. I ended up with a challenge to get to the top percentile of a mat mul and this is what I learned.

More than anything, this is a worklog. The goal is not to present a definitive GEMM optimization guide, but to document how I learned to reason about GPU performance by iteratively improving a kernel and understanding the bottlenecks encountered along the way. The focus is less on absolute performance and more on building intuition for how memory hierarchy, data reuse, and execution resources interact to determine performance.

There are already excellent deep dives into CUDA GEMM implementations such as How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance by Simon Boehm and Matrix Multiplication on GPU by Aleksa Gordic. Rather than focusing on the CUDA programming model itself, this post focuses on the performance bottlenecks that shaped each optimization step.

Note: As you can see the kernels were written around Christmas last year, I’m sure Claude/Codex will one shot a better one now.

GPU bottlenecks

Before we start implementing the kernels it’s a good idea to reason about the hardware.

The two fundamental properties of a compute system are:

Operation throughput (how many operations per second it can execute)

Data supply rate (how quickly operands can be delivered to those operations)

The time of your algorithm is dominated by:

\[\text{Time} \geq \max(\text{compute_time},\ \text{data_movement_time})\]

In other words, performance is bounded by either computation or communication.

If we look at modern GPUs, they have a ridiculous amount of FLOPs and it’s growing much faster than the amount of memory and memory bandwidth. This phenomenon has been dubbed the “memory wall” as characterized by this paper. Their graph shows why this is the case:

Stephen Jones from Nvidia gives a good explanation in this lecture. Fundamentally, SRAM needs 6 transistors per bit so fast memory takes up a lot of space on the die and memory bandwidth is increasingly constrained by physical limits, the speed of light becomes a factor when moving data across chips. There’s no reason at this point to believe this is going to change soon so the focus on memory optimization will stay relevant.

This has been a known problem in computing for a while and memory hierarchy systems have been devised to partly overcome this issue. CPUs rely on hardware-managed caches (L1/L2/L3), while GPUs combine caches with programmer-managed shared memory to keep data close to the compute units. TPUs take this to the extreme and drop caches almost entirely in favor of a large compiler-managed scratchpad, with the compiler (XLA) scheduling all data movement so operands stream deterministically into the compute units. Given the predictable nature of data movement for the dominant ML operations, this makes a lot of sense (and saves a lot of die area).

As we’ll mostly be focusing on optimising the use of the memory hierarchy, let’s have a look at the hierarchy of an Nvidia GPU. The numbers shown are representative of a modern data-center GPU, exact capacities vary per architecture (the T4 we’ll be using has 16 GB of GDDR6 and 64 KB of shared memory per SM for example).

As you can see there is quite an elaborate memory hierarchy divided into programmer and hardware managed levels. On one side you have the caches, where the hardware stores recently used data. On the other side you’re provided with shared memory which is under our control, and registers which the compiler allocates for us, to minimize round trips to very expensive global memory. There’s an inverse relationship between size and speed and the difference between global memory to shared memory to registers is 20-30x per step!

Given this hierarchy, the challenge becomes restructuring algorithms so data remains in fast memory as long as possible. This typically means breaking problems into tiles that fit into shared memory and registers. Matrix multiplication is a classic example, and modern algorithms like FlashAttention apply the same principle by restructuring attention into SRAM-sized tiles.

However, locality alone is not enough. Some memory accesses are unavoidable, and GPUs address this through massive parallelism. By keeping many more warps in flight than execution units, the scheduler can switch to ready work whenever others stall on memory.

Together, locality and parallelism form the two fundamental strategies for achieving high performance on GPUs:

Reduce memory traffic through locality and reuse

Hide remaining latency through massive parallelism

Setting the target

Before diving in, let’s establish some targets. Performance is measured with...