Lesson 17.2 · 28 min

CUDA Kernels: Writing Parallel Code for NVIDIA GPUs

When PyTorch multiplies two matrices on a GPU, a million tiny threads each run the same short function — so how does each one know which numbers are its job?

In short: A CUDA kernel is a function written once and executed in parallel by thousands or millions of GPU threads. Threads are grouped into blocks, and blocks into a grid; each thread computes its own global index from built-in variables and works on the matching piece of data. The CPU (host) copies data to GPU memory and launches kernels; the GPU (device) schedules blocks onto its streaming multiprocessors in groups of 32 threads called warps. Fast kernels depend as much on smart memory use as on raw arithmetic.

Why do we need a GPU?

In the previous lesson we saw that deep learning is dominated by operations like matrix multiplication and element-wise maths over huge arrays. A CPU has a handful of powerful cores and works through such arrays a few elements at a time. A GPU has thousands of simpler cores and can process thousands of elements at the same moment — if we give it the work in the right shape.

Running example: our appliance shop's chatbot. Every generated token involves adding bias vectors, applying activation functions, normalising, and multiplying by weight matrices with millions of numbers. Each of those is one or more CUDA kernels. Understanding kernels explains why some operations are fast, why others waste the GPU, and why tools like FlashAttention matter.

What is CUDA?

CUDA (Compute Unified Device Architecture) is NVIDIA's platform for general-purpose programming of its GPUs, first released in 2007. It includes:

• A language extension of C/C++ (compiled with nvcc) that lets us mark functions to run on the GPU and launch them on many threads.
• A runtime and driver that allocate GPU memory, copy data, and schedule work.
• Libraries built on top: cuBLAS (linear algebra), cuDNN (neural-network operations), NCCL (multi-GPU communication), and more.
• Tools: profilers and debuggers (such as Nsight) to see where time goes.

Most AI engineers never write raw CUDA: PyTorch and JAX call pre-written kernels for us. But the programming model — threads, blocks, grids and memory levels — shapes the performance of everything we run, and writing or choosing custom kernels is a key optimisation skill.

What is a CUDA kernel?

A kernel is a function that runs on the GPU and is executed by many threads in parallel. We write the code for one thread, describing the work for one small piece of the data, and CUDA runs it on as many threads as we ask for. In CUDA C++ a kernel is marked with __global__, returns void, and is launched with the special <<<blocks, threads>>> syntax.

Think of it like one instruction sheet for a whole stadium A teacher writes one instruction card: 'Look at your seat number. Take the worksheet with that number, add the two numbers on it, write the answer.' Every student in the stadium gets an identical card. Because each student has a different seat number, they all do different worksheets at the same time. The card is the kernel; the seat number is the thread index.

This style is called SPMD (single program, multiple data); NVIDIA calls its hardware execution model SIMT (single instruction, multiple threads). The program is identical for every thread; only the data each thread touches differs.

Threads, blocks, and grids

CUDA organises threads in a three-level hierarchy:

• Thread: the smallest unit; runs the kernel once with its own registers and its own index.
• Block (thread block): a group of threads (up to 1,024 on current NVIDIA GPUs) that run on the same processor, can share a fast on-chip memory called shared memory, and can wait for each other with __syncthreads().
• Grid: all the blocks of one kernel launch. Blocks are independent: they may run in any order, at the same time or one after another.

Blocks and grids can be 1-D, 2-D or 3-D, which is convenient for images and matrices: a 2-D grid of 2-D blocks maps naturally onto a 2-D output. Built-in variables tell each thread where it sits: threadIdx (position inside its block), blockIdx (position of its block in the grid), blockDim (threads per block) and gridDim (blocks in the grid).

Host and device

CUDA programs have two sides. The host is the CPU and its main memory (RAM). The device is the GPU and its own memory (VRAM). They are separate: a normal CPU pointer cannot be read by a kernel, and vice versa (unless we use unified memory, which moves data automatically, sometimes at a performance cost).

Keep data on the GPU Frameworks keep tensors on the device between operations precisely to avoid the copy steps. A common performance bug in PyTorch is calling .item(), .cpu() or printing a GPU tensor inside a training loop: each forces a device-to-host copy and makes the CPU wait for the GPU.

Writing our first CUDA kernel

The 'hello world' of CUDA is vector addition: c[i] = a[i] + b[i] for about a million elements. Here is a complete program. It needs an NVIDIA GPU and the CUDA toolkit to compile (nvcc vec_add.cu -o vec_add), so we show it for reading and then simulate the same launch in Python with real output.

vec_add.cu (for reading)

#include <cstdio>
#include <cuda_runtime.h>
// __global__ = a kernel: called from the CPU (host), runs on the GPU (device)
__global__ void vec_add(const float* a, const float* b, float* c, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;   // my global index
if (i < n) c[i] = a[i] + b[i];                   // guard + one element of work
}
int main() {
const int n = 1 << 20;                            // about 1 million elements
size_t bytes = n * sizeof(float);
float *h_a = (float*)malloc(bytes), *h_b = (float*)malloc(bytes), *h_c = (float*)malloc(bytes);
for (int i = 0; i < n; i++) { h_a[i] = i; h_b[i] = 100; }
float *d_a, *d_b, *d_c;                           // pointers to GPU memory
cudaMalloc(&d_a, bytes); cudaMalloc(&d_b, bytes); cudaMalloc(&d_c, bytes);
cudaMemcpy(d_a, h_a, bytes, cudaMemcpyHostToDevice);
cudaMemcpy(d_b, h_b, bytes, cudaMemcpyHostToDevice);
int threads = 256;
int blocks = (n + threads - 1) / threads;         // round up: 4096 blocks
vec_add<<<blocks, threads>>>(d_a, d_b, d_c, n);   // launch (asynchronous)
cudaMemcpy(h_c, d_c, bytes, cudaMemcpyDeviceToHost);   // waits for the kernel
printf("c[5] = %.1f\n", h_c[5]);                   // 105.0
cudaFree(d_a); cudaFree(d_b); cudaFree(d_c);
free(h_a); free(h_b); free(h_c);
}

simulate_launch.py

import numpy as np
# Simulate how CUDA launches  vec_add<<<blocks, threads_per_block>>>(a, b, c, n)
n = 10
a = np.arange(n, dtype=np.float32)          # 0, 1, ..., 9
b = np.full(n, 100, dtype=np.float32)
c = np.zeros(n, dtype=np.float32)
threads_per_block = 4                        # blockDim.x
blocks = (n + threads_per_block - 1) // threads_per_block   # round up -> 3
print(f"launch: {blocks} blocks x {threads_per_block} threads = "
f"{blocks * threads_per_block} threads for {n} elements")
def vec_add_kernel(blockIdx, threadIdx, blockDim):
"""The body every thread runs. Only the indices differ."""
i = blockIdx * blockDim + threadIdx      # this thread's global index
if i < n:                                # guard: extra threads do nothing
c[i] = a[i] + b[i]
return f"c[{i}]"
return "idle"
# On a GPU these run at the same time; here we just loop over them.
for blk in range(blocks):
jobs = [vec_add_kernel(blk, t, threads_per_block) for t in range(threads_per_block)]
print(f"block {blk}: threads 0-3 -> {jobs}")
print("c =", c)
# Memory coalescing: addresses touched by 8 neighbouring threads (4-byte floats)
tid = np.arange(8)
print("good  (a[i])     byte offsets:", (tid * 4).tolist())
print("bad   (a[i*32])  byte offsets:", (tid * 32 * 4).tolist())

Output:

launch: 3 blocks x 4 threads = 12 threads for 10 elements
block 0: threads 0-3 -> ['c[0]', 'c[1]', 'c[2]', 'c[3]']
block 1: threads 0-3 -> ['c[4]', 'c[5]', 'c[6]', 'c[7]']
block 2: threads 0-3 -> ['c[8]', 'c[9]', 'idle', 'idle']
c = [100. 101. 102. 103. 104. 105. 106. 107. 108. 109.]
good  (a[i])     byte offsets: [0, 4, 8, 12, 16, 20, 24, 28]
bad   (a[i*32])  byte offsets: [0, 128, 256, 384, 512, 640, 768, 896]

How a thread finds its own work

Thread 2 of block 1, with 4 threads per block

1. Skip earlier blocks: Block 0 already covers indices 0–3, so the blocks before us account for 1 × 4 = 4 elements.
2. Add own position: We are thread 2 in our block: 4 + 2 = 6.
3. Check the guard: 6 < n = 10, so we do work.
4. Do the work: Read a[6] and b[6], write c[6] = 106.

For 2-D data such as a matrix with R rows and C columns, each thread computes a row and a column: row = blockIdx.y blockDim.y + threadIdx.y and col = blockIdx.x blockDim.x + threadIdx.x, then works on element row * C + col (matrices are stored row by row in one flat array).

Grid-stride loops If the data is larger than the number of threads we launch, each thread can loop: start at its index i and keep adding the total thread count (blockDim.x × gridDim.x) until it passes n. One kernel then works for any data size.

Pause and think: A launch uses 128 threads per block. Which global index does thread 5 of block 3 handle?

i = 3 × 128 + 5 = 389.

What happens inside the GPU when a kernel runs

An NVIDIA GPU is built from many Streaming Multiprocessors (SMs) — for example, the A100 has 108 and the H100 SXM has 132. Each SM has its own CUDA cores, Tensor Cores, registers, shared memory/L1 cache and warp schedulers.

From launch to results

1. Blocks are assigned to SMs: The GPU's scheduler hands blocks to SMs that have free resources. A block stays on one SM until it finishes. Several blocks can live on one SM at once if registers and shared memory allow.
2. Blocks are split into warps: The SM divides each block into warps of 32 consecutive threads. A warp is the real unit of execution: its 32 threads execute the same instruction at the same time on different data.
3. Warp schedulers pick ready warps: Each cycle, the schedulers choose warps that are ready to run. When a warp waits for memory (hundreds of cycles), the SM simply switches to another ready warp at almost no cost.
4. Latency is hidden, not avoided: With enough warps in flight (high occupancy), there is always someone ready to compute while others wait for data. This is why GPUs want far more threads than they have cores.
5. Blocks finish, new ones start: As blocks complete, their SM takes new ones until the whole grid is done. Because blocks are independent, the same kernel runs on a small GPU (few SMs, more rounds) or a big one (many SMs, fewer rounds).

Warp divergence If threads in the same warp take different branches of an if, the warp runs both branches one after another, with the inactive threads masked off. A kernel where half of each warp takes each branch can run about half as fast. Arrange data so neighbouring threads usually follow the same path.

Memory in CUDA

Two rules make kernels fast. Coalescing: when the 32 threads of a warp read consecutive addresses (thread k reads element k), the hardware combines them into a few wide memory transactions. If they read scattered addresses, as with a[i*32] in our simulation, each read needs its own transaction and effective bandwidth collapses. Reuse via shared memory: in matrix multiplication, each input value is needed by many outputs. A tiled kernel loads a tile of each matrix from global memory into shared memory once, synchronises the block, and lets every thread reuse those values many times, cutting global-memory traffic by roughly the tile size.

Why CUDA kernels matter for AI

Every PyTorch operation on a GPU tensor launches at least one kernel. A naive model runs a long chain of small kernels — matmul, add bias, GELU, dropout, LayerNorm — and each one reads its input from HBM and writes its output back. For element-wise steps the arithmetic is tiny, so the time goes almost entirely to memory traffic and launch overhead.

• Kernel fusion combines several operations into one kernel, so intermediate results stay in registers or shared memory instead of travelling to HBM and back. Compilers like torch.compile and XLA do this automatically for many patterns.
• FlashAttention (2022) is a custom attention kernel that processes queries, keys and values in tiles held in on-chip memory and never writes the full N×N attention matrix to HBM. It computes exactly the same result, uses far less memory, and runs faster, which helped make long context windows practical.
• Tuned matmul libraries (cuBLAS, CUTLASS) use tiling, Tensor Cores and careful scheduling to get close to the hardware's peak FLOPs.
• Inference engines such as vLLM and TensorRT-LLM rely on specialised kernels for paged attention, quantised matrix multiplication and sampling.
• Triton (from OpenAI) lets engineers write GPU kernels in a Python-like language at the level of blocks rather than individual threads, and is used by torch.compile to generate fused kernels.

Back to our chatbot Swapping a naive attention implementation for a FlashAttention-style kernel, and enabling a fused, quantised matmul kernel in the serving engine, can raise our chatbot's throughput noticeably without changing the model at all. That is the practical payoff of understanding kernels.

Where CUDA kernels work well and where they fail

Common mistakes Forgetting the bounds check if (i < n) (out-of-bounds writes corrupt memory); launching too few threads to fill the GPU; ignoring that kernel launches are asynchronous (timing them without synchronising measures nothing); many threads updating the same location without atomics (race conditions); and assuming a custom kernel beats cuBLAS — it rarely does for plain matmul.

Worked example, step by step

Vector addition is the easy case: every thread writes to its own output slot, c[i]. Many useful operations are not like that. Take the simplest one: adding up all the elements of an array. Now every thread wants to write to the same place. Let us walk through why the obvious kernel breaks and how a correct one is built.

Summing an array on a GPU

1. The obvious kernel: Each thread runs total = total + a[i]. That one line is really three actions: read total, add, write total.
2. The race, with numbers: Say total = 0, a[0] = 1 and a[1] = 2. Thread 0 reads 0. Thread 1 also reads 0, before thread 0 has written. Thread 0 writes 1. Thread 1 writes 2. The final total is 2, not 3. One update is lost. This is a race condition.
3. Fix A: atomic add: An atomic operation does read, add and write as one unbreakable action. The answer is now correct, but the threads must take turns at one memory location. We have built a queue: the work is serial again.
4. Fix B: a tree: Pair the elements up. In round 1, each pair is added by a different thread, and no two threads touch the same slot. That halves the list. Round 2 pairs up the partial sums, and so on.
5. Count the rounds: Each round halves the list, so n elements need about log₂ n rounds. Eight elements take 3 rounds. About a million (2²⁰) take only 20.
6. How real kernels arrange it: A common layout: each block reduces its own slice in shared memory, calling __syncthreads() between rounds, and writes one partial sum. A second, much smaller pass adds the partial sums.

The habit to build: before writing a kernel, ask 'can two threads ever write to the same address?'. If the answer is no, the kernel is safe by construction. If it is yes, we need a tree, atomics, or a different way to split the work. In practice we call a library reduction rather than write our own, but knowing the pattern explains why a sum is not one simple kernel.

Practice: try it yourself

We will simulate three things in plain Python: a grid-stride loop that covers 10 elements with only 4 threads, the broken 'everyone adds to total' sum, and a tree reduction that gets the right answer in a few rounds.

practice_reduction.py

import numpy as np
a = np.arange(1, 11)                             # 1..10, true sum = 55
n = len(a)
# 1) Grid-stride loop: only 4 threads, but 10 elements to cover.
total_threads = 4                                # blockDim.x * gridDim.x
for tid in range(total_threads):
mine = list(range(tid, n, total_threads))    # i = tid, tid+4, tid+8, ...
print(f"thread {tid} handles indices {mine}")
# 2) A broken sum: every thread does  total = total + a[i]  at the same time.
total = 0
seen = [total for _ in range(n)]                 # all threads read total first
for i in range(n):
total = seen[i] + a[i]                       # then each writes its own answer
print(f"racy sum: {total} (should be {a.sum()})")
# 3) A tree reduction: in each round, thread i adds its partner's value.
vals = a.copy()
stride, rounds = 1, 0
while stride < n:
for i in range(0, n - stride, 2 * stride):   # these pairs are independent
vals[i] += vals[i + stride]
stride *= 2; rounds += 1
print(f"round {rounds}: {vals[::stride].tolist()}")
print("tree sum:", vals[0], "in", rounds, "rounds")
# 4) How the number of rounds grows with the data size
for size in (10, 1024, 1 << 20):
print(f"{size:>9,} elements -> {int(np.ceil(np.log2(size)))} rounds")

Output:

thread 0 handles indices [0, 4, 8]
thread 1 handles indices [1, 5, 9]
thread 2 handles indices [2, 6]
thread 3 handles indices [3, 7]
racy sum: 10 (should be 55)
round 1: [3, 7, 11, 15, 19]
round 2: [10, 26, 19]
round 3: [36, 19]
round 4: [55]
tree sum: 55 in 4 rounds
10 elements -> 4 rounds
1,024 elements -> 10 rounds
1,048,576 elements -> 20 rounds

Now change it:

• Set total_threads = 3. Predict the indices thread 0 handles before running. Is every index still covered exactly once?
• Change a to np.arange(1, 17) (16 elements). Predict the number of rounds and the list printed in round 1.
• In the tree, replace vals[i] += vals[i + stride] with vals[i] = max(vals[i], vals[i + stride]). Predict the final value. Which property must an operation have for the tree to work?

Pause and think: Our simulated racy sum printed exactly 10 every time. Would a real GPU also give 10 every time?

No. In the simulation we fixed the order: everyone reads 0, and the last thread (holding a[9] = 10) writes last. On a real GPU, which threads read stale values and whose write lands last depends on timing, so the result can change from run to run and from GPU to GPU. That is what makes race conditions hard to debug: the code can even look correct on a small test.

Pause and think: We launch a grid-stride kernel with 1,048,576 threads, but today the array has only 1,000 elements. Do we need an extra if (i < n) guard to protect the surplus threads?

The loop condition already does that job. A thread starts at i = its own index and only runs while i < n. Threads with an index of 1,000 or more fail the test at once and do nothing, and the first 1,000 threads each handle one element. The launch is wasteful, since most threads are idle, but it is safe.

Key takeaways

• A CUDA kernel is code for one thread, launched on a grid of blocks of threads.
• Each thread finds its data with i = blockIdx.x × blockDim.x + threadIdx.x and a bounds check.
• The host (CPU) manages memory and launches kernels; the device (GPU) runs blocks on SMs in warps of 32.
• Speed depends on memory: coalesced access, shared-memory tiling and keeping data on the GPU.
• Fused and specialised kernels (FlashAttention, quantised matmuls) are a major source of AI speed-ups.

Key terms

• CUDA: NVIDIA's platform and C++ extension for general-purpose GPU programming.
• Kernel: A function that runs on the GPU, executed by many threads in parallel.
• Thread block: A group of up to 1,024 threads that run on one SM and can share memory and synchronise.
• Grid: All the thread blocks of one kernel launch.
• Warp: A group of 32 threads that execute the same instruction together.
• Streaming Multiprocessor (SM): One of the GPU's processing units that runs blocks of threads.
• Coalesced access: Neighbouring threads reading neighbouring addresses so the hardware can combine their memory requests.
• Kernel fusion: Combining several operations into one kernel to avoid round trips to global memory.

← 17.1 GPUs for Deep Learning: Parallelism at the Core · 17.3 Google TPUs: Purpose-Built Hardware for Neural Networks →