11. CUDA programming

One function, many threads

A CUDA kernel is a function that runs once per thread, on thousands of threads at the same time. Every thread executes the same code; what differs is its index, which it uses to decide which data to work on. This is the SPMD style (single program, multiple data), and the hardest part of learning it is to stop thinking “loop over elements” and start thinking “I am one element”.

Threads are organized in two levels:

• A block is a group of up to 1,024 threads that run on the same SM. They can cooperate through shared memory and synchronize with barriers (__syncthreads()).
• A grid is all the blocks of one launch. Blocks are independent: the hardware may run them in any order, in parallel or one after another, depending on how many fit on the GPU. A kernel must never assume an order between blocks.

That independence is what makes CUDA programs scale: the same kernel runs on a 48-SM laptop GPU and a 132-SM H100, just with more blocks in flight on the bigger one. PMPP calls this transparent scalability.

Vector addition

The “hello world” of CUDA adds two vectors. Each thread computes its global index from three built-in variables: blockIdx.x (which block), blockDim.x (threads per block) and threadIdx.x (which thread in the block):

$$ i = \texttt{blockIdx.x} \times \texttt{blockDim.x} + \texttt{threadIdx.x} $$

For $N = 1{,}000$ elements and 256 threads per block, we need $\lceil 1000 / 256 \rceil = 4$ blocks, which is 1,024 threads. The last 24 threads have no element to process, so every kernel includes a bounds check. Here’s the same kernel in three languages:

__global__ void add_kernel(const float* a, const float* b, float* out, int64_t n) {
    int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;   // this thread's element
    if (i < n) out[i] = a[i] + b[i];                              // the grid may overshoot n
}

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    """(Your engine: Chapter 14)"""
    pid = tl.program_id(0)                       # which block of BLOCK elements is mine
    offsets = pid * BLOCK + tl.arange(0, BLOCK)  # a vector of BLOCK indices
    mask = offsets < n                           # the last block may run past the end
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


def vector_add(x, y, block=1024):
    check_device(x, y)
    if x.shape != y.shape or not x.is_contiguous() or not y.is_contiguous():
        raise ValueError("vector_add needs equal-shape contiguous tensors")
    out = torch.empty_like(x)
    n = x.numel()
    add_kernel[(triton.cdiv(n, block),)](x, y, out, n, BLOCK=block)
    return out

#![allow(unused)]
fn main() {
#[cuda_module]
mod kernels {
    use super::*;
    #[kernel]
    pub fn twice(input: &[f32], mut output: DisjointSlice<f32>) {
        let index = thread::index_1d();
        if let Some(value) = output.get_mut(index) {
            *value = input[index.get()] * 2.0;
        }
    }
}
}

The CUDA C++ version is one thread’s view of the work. Triton (Chapter 14) is one block’s view: a program handles BLOCK elements as a vector, and the mask plays the role of the bounds check. NVIDIA’s experimental cuda-oxide compiles Rust to GPU code; its DisjointSlice turns the bounds check into a type-system guarantee that no two threads write the same element.

Launching a kernel

From host code (the CPU side), a launch names the grid and block sizes in triple angle brackets:

int threads = 256;
int blocks = (n + threads - 1) / threads;        // ceiling division
add_kernel<<<blocks, threads, 0, stream>>>(a, b, out, n);

Before that, the data must be in GPU memory. A standalone program allocates device buffers and copies to and from them explicitly (cudaMalloc, cudaMemcpy). The book’s standalone example does this with error checking and CUDA-event timing. Inside PyTorch, tensors already live on the device. You write a small C++ wrapper that checks its inputs, launches the kernel on PyTorch’s current stream, and returns a tensor:

torch::Tensor vector_add(torch::Tensor a, torch::Tensor b) {
    check(a); check(b);                               // CUDA, float32, contiguous
    TORCH_CHECK(a.sizes() == b.sizes(), "shapes differ");
    c10::cuda::CUDAGuard guard(a.device());           // launch on a's GPU
    auto out = torch::empty_like(a);
    int64_t n = a.numel();
    if (n) {                                          // zero-size launches are invalid
        add_kernel<<<(n + 255) / 256, 256, 0, stream()>>>(
            a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), n);
        C10_CUDA_KERNEL_LAUNCH_CHECK();               // surface launch errors immediately
    }
    return out;
}

torch.utils.cpp_extension.load compiles the .cu file the first time you call it and imports the result as a Python module (izh/kernels/cuda.py). The wrapper’s checks are part of correctness. A kernel receives raw pointers and knows nothing about shapes, dtypes or strides; pass it a transposed tensor and it silently reads the wrong elements (Chapter 2).

Warning

Kernel launches are asynchronous, and so are their errors. An out-of-bounds write may only be reported at the next synchronizing call, far from the bug. When something crashes mysteriously, rerun with CUDA_LAUNCH_BLOCKING=1 (launches become synchronous) or under compute-sanitizer (NVIDIA’s memory checker), which pinpoints the offending kernel and thread.

Two-dimensional grids

Grids and blocks can be 2-D or 3-D, which maps naturally onto matrices. For an M × N output with 16×16 blocks:

dim3 block(16, 16);
dim3 grid((N + 15) / 16, (M + 15) / 16);
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

A row-major matrix element (row, col) lives at row * N + col. For a width-5 matrix, (2, 3) is at offset 13.

Warps, divergence and occupancy

The hardware runs each block as warps of 32 consecutive threads that execute one instruction at a time, in lockstep. Two performance consequences follow:

• Divergence. If threads in a warp take different branches, the warp runs both paths one after the other, with some threads masked off each time. The bounds check in vector add diverges only in the last warp, which is harmless. Data-dependent branches in a hot loop can halve throughput or worse.
• Occupancy. An SM keeps many warps resident and switches between them to hide memory latency (Chapter 10). How many fit depends on each thread’s registers and each block’s shared memory: use more, and fewer warps fit. Occupancy is resident warps divided by the hardware maximum. It’s a means, not a goal. A kernel with low occupancy but lots of data reuse can beat one with high occupancy and none.

Coalescing: the rule that matters most

When the 32 threads of a warp load from global memory, the hardware combines their requests into as few memory transactions as possible. If thread $t$ reads element $\text{base} + t$, the 32 four-byte reads cover one contiguous 128-byte segment, served by a single transaction. That’s a coalesced access, and it uses the full bandwidth. If thread $t$ reads element $\text{base} + 32t$ (a stride), the warp touches 32 different segments and wastes most of each. Strided access can be 10-30x slower for the same number of useful bytes.

The practical rule: map threadIdx.x, which varies fastest within a warp, to the index that varies fastest in memory (the last axis of a row-major tensor). GPU Mode L8 measures this effect directly (coalesce.cu in the lecture repository).

A naive matrix multiplication

The direct translation of $C_{ij} = \sum_k A_{ik} B_{kj}$ gives each thread one output element:

// One thread per output element. Every thread re-reads a whole row of A and column of B
// from global memory: 2K loads for 2K flops, an arithmetic intensity of ~0.25 flop/byte.
__global__ void naive_matmul_kernel(const float* a, const float* b, float* c,
                                    int M, int K, int N) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;   // x varies fastest -> coalesced B and C
    if (row < M && col < N) {
        float total = 0.f;
        for (int k = 0; k < K; ++k) total += a[row * K + k] * b[k * N + col];
        c[row * N + col] = total;
    }
}

Check the address formulas against the shapes: A is [M, K] so its row stride is K, B is [K, N] so its row stride is N. A common bug uses one width for both, which works only for square matrices. Test with non-square shapes like [3, 5] @ [5, 7].

Is it coalesced? Within a warp, col varies (x is the fastest thread index) and row is fixed. Reads of B[k * N + col] are consecutive, which is coalesced. Reads of A[row * K + k] are the same address for every thread, a broadcast, which is fine. Writes to C[row * N + col] are consecutive, also good.

So why is it slow? Apply Chapter 10’s roofline. Each thread loads $2K$ values to do $2K$ FLOPs: an arithmetic intensity of about 0.25 FLOP per byte in FP32. Every element of A is re-read by each of the N threads that use it, and every element of B by M threads. The kernel is starved for memory bandwidth, and its speed is a small fraction of what the arithmetic units can do. Fixing that, by loading tiles into shared memory and reusing them many times, is the subject of the next chapter.

Build it

uv pip install ninja setuptools               # build tools for PyTorch extensions
export CUDA_HOME=/usr/local/cuda              # where your CUDA Toolkit lives
pytest tests/test_ch11_cuda.py -k "vector_add or matmuls"
python run.py kernels --backend cuda --impl engine

uv pip install "nvidia-cuda-nvcc==13.3.*" "nvidia-cuda-runtime==13.3.*" "nvidia-cuda-cccl==13.3.*"
NV=$(python -c "import nvidia, os; print(os.path.dirname(nvidia.__path__[0]))")/nvidia/cu13
$NV/bin/nvcc -std=c++17 -arch=sm_90 -I$NV/include -I$NV/include/cccl cpp/kernels.cu -o build/kernels

Important

DGX Spark compiles for -arch=sm_121 with CUDA 13. PyTorch’s extension builder picks the architecture from the GPU it finds, or from TORCH_CUDA_ARCH_LIST="12.1".

Stretch exercises

1. ★ Make vector add process 4 elements per thread with float4 loads (reinterpret_cast<const float4*>). Measure the bandwidth before and after on 2²⁸ elements. Where: add_kernel in engine/kernels/cuda_ops.cu; handle the tail when the length is not divisible by four.
2. ★★ Write a deliberately uncoalesced copy kernel (thread $t$ reads element $32t \bmod N$) and compare its bandwidth with a coalesced copy. Where: add copy kernels and host launchers in engine/kernels/cuda_ops.cu, then expose them in its PYBIND11_MODULE.
3. ★★ Swap the roles of x and y in the naive matmul (map threadIdx.x to rows). Measure the slowdown and explain it with coalescing. Where: naive_matmul_kernel and its launch geometry in engine/kernels/cuda_ops.cu.
4. ★★★ Write transpose_kernel for a [M, N] FP32 matrix. A naive version must choose between coalesced reads and coalesced writes. Then fix it with a 32×32 shared-memory tile (you’ll need the bank-conflict padding from Chapter 12). Where: add transpose_kernel, a host launcher and a binding in engine/kernels/cuda_ops.cu.

Check your understanding

1. Why launch more threads than elements and then bounds-check, instead of launching exactly N threads?
2. Why can’t one block wait for another block to finish?
3. A warp reads x[threadIdx.x * 16] (FP32). How many 128-byte segments does it touch, and how much of the transferred data is useful?
4. What is the naive matmul’s arithmetic intensity, and what does the roofline predict about its speed?

Going deeper

• PMPP Chapters 2-3 (pp. 23-73): data parallelism, CUDA program structure, multidimensional grids and the naive matmul. Chapter 4 (warps, divergence, occupancy) and §6.1 (coalescing).
• GPU Mode L2 (a recap of PMPP Chapters 1-3), L3 (Jeremy Howard, CUDA for Python programmers, writing kernels inline from a notebook), L4 (compute and memory architecture), and L8 (the performance checklist, with runnable coalesce.cu, divergence.cu and occupancy.cu).
• NVIDIA’s CUDA C++ Programming Guide, chapters on the programming model and memory hierarchy; compute-sanitizer documentation.