Writing CUDA Kernels from Scratch: A Beginner's Guide

Ever wondered how PyTorch or TensorFlow make your neural networks run so fast on GPUs? In this post, we'll build matrix multiplication kernels from scratch, starting naive and slow, then making them fast step by step. By the end, you'll understand exactly what makes GPU code fast (or slow).

All the code here is from Tensorax, a tensor library I built to learn these concepts. Feel free to explore the full implementation there.

Why GPUs?

Your CPU has 8-16 powerful cores. A GPU has thousands of tiny cores. Each GPU core is weaker than a CPU core, but when you have 10,000 of them doing the same thing to different data? That's where the magic happens.

Deep learning is mostly matrix math. Matrix multiplication, adding vectors, applying functions to every element. These operations are embarrassingly parallel, each output element can be computed independently.

CUDA Basics: Threads, Blocks, and Grids

In CUDA, you write a function called a kernel that runs on the GPU. But here's the twist: it runs thousands of times in parallel, once per thread.

Threads are organized into:

• Blocks: Groups of threads (up to 1024) that can share fast memory and synchronize
• Grid: All the blocks together

Each thread knows its position via built-in variables: threadIdx, blockIdx, and blockDim.

Let's Start Simple: Vector Addition

Adding two arrays is the "Hello World" of CUDA. We want: C[i] = A[i] + B[i] for every element.

1__global__ void add_kernel(const float* a, const float* b, float* out, int size) {
2    int idx = blockIdx.x * blockDim.x + threadIdx.x;
3    if (idx < size) {
4        out[idx] = a[idx] + b[idx];
5    }
6}

That's it. Each thread computes one addition. With 1 million elements and 256 threads per block, we launch ~4000 blocks. All running simultaneously.

The Real Challenge: Matrix Multiplication

Now let's tackle something harder, matrix multiplication. This is the backbone of neural networks. Every linear layer, every attention head uses it.

For matrices A (m×k) and B (k×n), each output element C[i,j] is:

C[i,j] = A[i,0]*B[0,j] + A[i,1]*B[1,j] + ... + A[i,k-1]*B[k-1,j]

In simple words: to get one output element, you take a row from A and a column from B, multiply them element-by-element, and add everything up. That's k multiplications and k-1 additions for each output element.

Attempt 1: The Naive Approach

The simplest idea: assign one thread to compute one output element. If our output matrix is 1024×1024, we launch about 1 million threads. Each thread loops through k values, multiplying and summing.

Let's say thread #42 is responsible for computing C[5,7]. It will read the entire 5th row of A (that's k values), read the entire 7th column of B (another k values), multiply them pair-wise, and sum everything up. Simple, right?

1__global__ void matmul_naive(const float* A, const float* B, float* C, 
2                               int m, int n, int k) {
3    int row = blockIdx.y * blockDim.y + threadIdx.y;
4    int col = blockIdx.x * blockDim.x + threadIdx.x;
5    
6    if (row < m && col < n) {
7        float sum = 0.0f;
8        for (int i = 0; i < k; ++i) {
9            sum += A[row * k + i] * B[i * n + col];
10        }
11        C[row * n + col] = sum;
12    }
13}

This works. But it's slow. Really slow.

The Problem: Memory is the Bottleneck

Here's what's happening: for a 1024×1024 matrix multiply, each output element needs to read 1024 values from A and 1024 from B. That's over 2 trillion memory reads total.

GPU global memory has high bandwidth (~900 GB/s on modern GPUs), but also high latency (~400 cycles). When every thread is hammering global memory independently, things get congested.

Can you think of a way to avoid this redundant memory access?

Attempt 2: Shared Memory to the Rescue

Here's an analogy: imagine you need to look up facts from books to write an essay. You could walk to the library (10 km away) every time you need a fact. Or, you could bring a few books to your desk and look things up from there.

The library is global memory, huge but far away. Your desk is shared memory, tiny but right next to you. It's about 100x faster to read from shared memory than global memory.

The catch? Your desk (shared memory) is small, only about 48KB per block. You can't fit the whole matrix there. So we work in tiles: bring a small chunk of data to the desk, do all the work we can with it, then bring the next chunk.

The strategy:

1. All threads in a block cooperate to load a "tile" of A and B into shared memory
2. Everyone computes their partial results using that fast, local data
3. Repeat for the next tile until we've processed all of k

1__global__ void matmul_tiled(const float* A, const float* B, float* C, 
2                               int m, int n, int k) {
3    const int TILE = 32;
4    __shared__ float tileA[TILE][TILE];
5    __shared__ float tileB[TILE][TILE];
6    
7    int row = blockIdx.y * TILE + threadIdx.y;
8    int col = blockIdx.x * TILE + threadIdx.x;
9    float sum = 0.0f;
10    
11    // Process tiles one at a time
12    for (int t = 0; t < (k + TILE - 1) / TILE; ++t) {
13        // Each thread loads ONE element into shared memory
14        if (row < m && t * TILE + threadIdx.x < k)
15            tileA[threadIdx.y][threadIdx.x] = A[row * k + t * TILE + threadIdx.x];
16        else
17            tileA[threadIdx.y][threadIdx.x] = 0.0f;
18            
19        if (col < n && t * TILE + threadIdx.y < k)
20            tileB[threadIdx.y][threadIdx.x] = B[(t * TILE + threadIdx.y) * n + col];
21        else
22            tileB[threadIdx.y][threadIdx.x] = 0.0f;
23        
24        __syncthreads();  // Wait for everyone to finish loading
25        
26        // Now compute using fast shared memory
27        for (int i = 0; i < TILE; ++i)
28            sum += tileA[threadIdx.y][i] * tileB[i][threadIdx.x];
29        
30        __syncthreads();  // Wait before loading next tile
31    }
32    
33    if (row < m && col < n)
34        C[row * n + col] = sum;
35}

Can We Do Better?

The tiled version is good, but we're still limited. Each thread computes just one output element. What if one thread could compute multiple elements?

Attempt 3: Register Blocking

We moved data from the library (global memory) to our desk (shared memory). But there's an even faster place: your hands. If you're copying numbers from a book, you don't look at each digit separately, you remember a few digits at a time in your head.

Registers are like your short-term memory. They're the fastest storage on the GPU, even faster than shared memory. But each thread only gets a small number of them.

The key insight: if one thread computes multiple output elements (say, an 8×8 tile), it can load some values into registers and reuse them many times. Instead of one thread doing one output, one thread does 64 outputs and reuses data aggressively.

1// Each thread computes a TM×TN tile of output
2template<int BM, int BN, int BK, int TM, int TN>
3__global__ void matmul_register_blocked(const float* A, const float* B, float* C,
4                                        int m, int n, int k) {
5    __shared__ float sA[BM * BK];
6    __shared__ float sB[BK * BN];
7    
8    // Each thread's results stored in registers
9    float results[TM][TN] = {0.0f};
10    float regA[TM];
11    float regB[TN];
12    
13    // ... load tiles into shared memory ...
14    
15    for (int dot = 0; dot < BK; ++dot) {
16        // Load into registers
17        for (int i = 0; i < TM; ++i)
18            regA[i] = sA[/* thread's row */ * BK + dot];
19        for (int i = 0; i < TN; ++i)
20            regB[i] = sB[dot * BN + /* thread's col */];
21        
22        // Compute outer product - maximum register reuse!
23        for (int i = 0; i < TM; ++i)
24            for (int j = 0; j < TN; ++j)
25                results[i][j] += regA[i] * regB[j];
26    }
27}

The key trick: if we load 8 values from A (call them regA) and 8 values from B (call them regB), we can compute 8×8 = 64 results! Each value in regA gets multiplied with all 8 values in regB, and vice versa. That's the "outer product" approach, maximum reuse from minimum loads.

To summarize our memory hierarchy journey:

• Global memory (Library): Huge, slow. Avoid when possible.
• Shared memory (Desk): Medium, fast. Share between threads.
• Registers (Your brain): Tiny, fastest. Reuse within a thread.

How Fast Did We Get?

We went from 3.37s to 0.95s, about 3.5× faster. PyTorch (using cuBLAS) is still faster because they have years of optimization.

The Takeaways

GPU optimization boils down to a few key principles:

1. Memory is the bottleneck: Compute is cheap; memory access is expensive. Minimize global memory reads.
2. Share data between threads: Use shared memory when multiple threads need the same data.
3. More work per thread: Having each thread compute multiple outputs means better data reuse.
4. Registers are your friend: Keep frequently used values in registers, not shared memory.

Bonus: Automatic Differentiation

Fast matrix multiply is great, but neural networks need gradients. When you call loss.backward(), how does PyTorch know what gradients to compute?

The trick is building a computational graph during the forward pass. Each operation records its inputs and what operation was performed.

1# Forward pass builds the graph
2x = Tensor([[2.0]], requires_grad=True)
3w = Tensor([[3.0]], requires_grad=True)
4y = w * x  # Records: y = mul(w, x)
5
6# Backward traverses it in reverse
7y.backward()
8print(w.grad)  # dy/dw = x = 2
9print(x.grad)  # dy/dx = w = 3

For matrix multiply specifically, if C = A @ B, the gradients are:

• dL/dA = dL/dC @ B.T
• dL/dB = A.T @ dL/dC

So the backward pass is... just more matrix multiplies! The same optimized kernels we wrote work for both forward and backward passes.

What's Next?

If you want to go deeper:

• Simon Boehm's goated CUDA matmul tutorial
• tensorax source code
• Lei Mao's detailed exploration of CUDA matmul optimizations
• NVIDIA's official CUDA programming guide

Try implementing your own kernel! Start with the naive version, verify it's correct, then optimize. There's no better way to learn than by doing.