Step 1: Your first CUDA kernel: vector add
Every fast AI library you use — cuBLAS (matmul in PyTorch), FlashAttention (transformers), vLLM's PagedAttention, TensorRT — is, under the hood, handwritten CUDA C++. If you only ever write torch.matmul(a, b) you will never understand why your model is slow. This lab is the layer below.
CUDA is NVIDIA's parallel programming model, introduced in 2007 with the Tesla architecture. The core idea is radical but simple: tens of thousands of threads execute the same tiny program in parallel. You don't write one thread doing a loop — you write the body of the loop once, and the hardware spawns millions of threads to cover the whole input.
The canonical reference stack (while this runs)
- Mark Harris — An Even Easier Introduction to CUDA (NVIDIA blog). The shortest path from zero to a working kernel. Read it alongside this lab.
- NVIDIA CUDA C++ Programming Guide. The canonical reference. You'll want Sections 2 (Programming Model) and 5 (Performance Guidelines).
- Vasily Volkov — Better Performance at Lower Occupancy (2010). Fifteen years old and still the single most important optimization talk. Core lesson: instruction-level parallelism, not thread count, dominates throughput.
- Andrej Karpathy — llm.c. Karpathy's from-scratch GPT-2 training in raw CUDA + C. ~4,000 lines. The clearest real-world educational CUDA-for-ML codebase.
What you'll build
Four CUDA kernels, each more sophisticated than the last, all compiled on the fly from this notebook via PyTorch's inline C++ extension builder:
- Vector addition — the "hello world" of CUDA. 1D thread indexing, CPU-vs-GPU timing.
- Matrix addition — 2D thread grids. Where you learn to index (row, col) correctly.
- Tiled matrix multiplication — shared memory, the single most important CUDA optimization. Goes from memory-bound to compute-bound.
- A custom autograd op — bind your CUDA kernel to
torch.autograd, so it composes with every other PyTorch operation.
Hardware context
Your RTX 3060 Ti has 4,864 CUDA cores, 152 tensor cores, and runs up to 38 warps (of 32 threads each = 1,216 threads) per streaming multiprocessor (SM) simultaneously. There are 38 SMs, giving a theoretical ~46,000 in-flight threads. For context, the H100 has 16,896 CUDA cores and 456 tensor cores across 132 SMs.
c[i] = a[i] + b[i]for 1 million elements
In C this is for (int i = 0; i < N; i++) c[i] = a[i] + b[i]; — a loop. In CUDA, we write the body of the loop once, and spawn one thread per iteration.
__global__ void vec_add(const float* a, const float* b, float* c, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x; // thread's global index
if (i < n) c[i] = a[i] + b[i]; // bounds check
}
Three CUDA keywords to learn:
__global__— this function runs on the GPU and is launched from the CPUblockIdx.x,blockDim.x,threadIdx.x— built-in variables giving every thread a unique global identity<<<blocks, threads_per_block>>>— the launch configuration syntax you'll see in the launcher cell
🐛 Common mistake: off-by-one at the boundary
Launching (n + 255) / 256 blocks of 256 threads covers every element. But the last block has extra threads past n. Without the if (i < n) bounds check, those threads write past the end of the array — memory corruption that sometimes crashes and sometimes just produces quietly-wrong results. It's among the top-3 bugs in production CUDA code.
import torch
from torch.utils.cpp_extension import load_inline
import time
# PyTorch's load_inline compiles CUDA source on the fly, caches by source hash, and binds it as a Python function.
vec_add_source = r'''
__global__ void vec_add_kernel(const float* a, const float* b, float* c, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) c[i] = a[i] + b[i];
}
torch::Tensor vec_add(torch::Tensor a, torch::Tensor b) {
auto c = torch::empty_like(a);
int n = a.numel();
int threads = 256;
int blocks = (n + threads - 1) / threads;
vec_add_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), c.data_ptr<float>(), n);
return c;
}
'''
_mod = load_inline(
name='vec_add_mod',
cpp_sources=['torch::Tensor vec_add(torch::Tensor a, torch::Tensor b);'],
cuda_sources=[vec_add_source],
functions=['vec_add'],
verbose=False,
)
vec_add = _mod.vec_add
print('Kernel compiled and bound as vec_add()')
# Time CPU vs GPU on a million-element vector.
N = 1_000_000
a_cpu = torch.randn(N)
b_cpu = torch.randn(N)
a_gpu = a_cpu.cuda(); b_gpu = b_cpu.cuda()
# CPU reference + timing
t0 = time.perf_counter(); _ = a_cpu + b_cpu; cpu_ms = (time.perf_counter() - t0) * 1000
# GPU kernel timing (warmup first, then time with torch.cuda.synchronize)
_ = vec_add(a_gpu, b_gpu); torch.cuda.synchronize()
t0 = time.perf_counter(); c_gpu = vec_add(a_gpu, b_gpu); torch.cuda.synchronize(); gpu_ms = (time.perf_counter() - t0) * 1000
# Correctness
assert torch.allclose(c_gpu.cpu(), a_cpu + b_cpu)
print(f'CPU: {cpu_ms:.2f} ms')
print(f'GPU: {gpu_ms:.3f} ms')
print(f'Speedup: {cpu_ms / max(gpu_ms, 0.01):.1f}x')
print()
print('The speedup is modest for vector add (memory-bound op).')
print('The REAL win comes with compute-heavy ops like matmul (Step 3).')