hands-on labon a real NVIDIA GPU

CUDA Programming Fundamentals

Write four real CUDA C++ kernels and run them from PyTorch: vector add, 2D matrix add, tiled matmul with shared memory, and a custom autograd op.

time
45 min
checked steps
4
level
Advanced
setup
None
NCP-GENLNCA-AIIONCP-ADSNCA-GENMNCP-AII

$29.99a month with Pro

Pro opens this lab and every other one. Cancel any time.

4 steps, each checked when you finish it

A check runs your work at the end of every step. Hints and the full solution are there if you get stuck.

  1. 1

    Your first CUDA kernel: vector add

    read this step below ↓

  2. 2

    2D thread grids: matrix add

  3. 3

    Tiled matmul with shared memory

  4. 4

    Custom autograd op: wire a kernel into PyTorch

Step 1 as it appears in the lab

The lab’s own text. The hint and the solution stay inside the lab.

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)

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:

  1. Vector addition — the "hello world" of CUDA. 1D thread indexing, CPU-vs-GPU timing.
  2. Matrix addition — 2D thread grids. Where you learn to index (row, col) correctly.
  3. Tiled matrix multiplication — shared memory, the single most important CUDA optimization. Goes from memory-bound to compute-bound.
  4. 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 CPU
  • blockIdx.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).')

what you get

  • A JupyterLab notebook on a real NVIDIA GPU
  • A check at the end of every step
  • Hints and the full solution
  • Nothing to install or configure

before you start

  • Comfortable with PyTorch tensors and CUDA basics
  • Basic C/C++, pointers, arrays, function signatures
  • Understanding of matrix multiplication

Exam domains covered

GPU Acceleration & Distributed TrainingModel Deployment & Inference Optimization

Part of the GPU and Kubernetes track

Lab 14 of 20. GPU scheduling on your own Kubernetes cluster, the GPU Operator, CUDA kernels and profiling.

Questions about this lab

Do I need prior CUDA experience to take this lab?

No. Step 1 introduces __global__, thread indexing (blockIdx.x * blockDim.x + threadIdx.x), and the <<<blocks, threads>>> launch syntax from scratch. You do need basic C/C++ fluency, reading pointer arithmetic, function signatures, and loops, and comfort with PyTorch tensors. If you can read a short .cu file and roughly follow what it does, you’re ready.

Why write a matmul kernel when cuBLAS already exists?

The point isn’t to replace cuBLAS, it’s to understand why cuBLAS wins. Step 3 benchmarks your tiled kernel against torch.matmul and you’ll typically see a 5, 10× gap. That concrete comparison teaches when writing a kernel pays off (custom fused ops, unusual tensor shapes, operations vendor libraries don’t cover) versus when you should lean on cuBLAS / cuDNN / FlashAttention instead.

Does the custom autograd op in Step 4 actually work with loss.backward()?

Yes. You subclass torch.autograd.Function, implement forward that calls your CUDA ReLU kernel, and backward that calls a second kernel multiplying grad_output by the ReLU mask. After the step, gradients flow through your kernel cleanly and match the reference torch.nn.functional.relu. This is the exact pattern real libraries use, FlashAttention, xformers, and most of torchao ship custom CUDA ops bound this way.

What is torch.utils.cpp_extension.load_inline and why use it instead of a setup.py build?

load_inline takes a string of CUDA source, JIT-compiles it with nvcc, caches the compiled module by source hash, and returns a Python-callable binding, no setup.py, no build system, no .cu files to manage. It’s the fastest way to iterate on kernels during research and prototyping, which is why the PyTorch team added it. For shipping production ops you’d graduate to a proper extension build, but the kernel code itself is identical.

What do I need installed locally?

Nothing. The lab runs on a real NVIDIA GPU pod we provision on demand, nvcc, the CUDA toolkit, PyTorch, and the compile cache are preinstalled. You only need a browser. Your session persists between visits, so compiled modules and workspace files stay put if you pause mid-lab.

How is each step auto-graded?

By running your compiled kernel against a correctness reference. Step 1 checks vec_add output matches a + b on 1024 elements and that you measured CPU vs GPU timing. Step 2 verifies mat_add on a 256×384 matrix. Step 3 confirms your tiled matmul matches torch.matmul within floating-point tolerance and that the benchmark cell ran. Step 4 checks that your PreporatoReLU.apply produces correct forward values and correct gradients through .backward().

What you'll build in this CUDA lab

Writing CUDA kernels is the skill that separates AI engineers who consume PyTorch from the ones who extend it, and in 2026, with every inference optimisation (FlashAttention-3, PagedAttention, fused MoE gating, custom quantisation kernels) shipping as bespoke CUDA, it's the highest-leverage low-level skill an ML engineer can own. In about 45 minutes on a real NVIDIA GPU we provision, you'll walk away with four working kernels you wrote from scratch, vector add, 2D matrix add, a shared-memory tiled matmul, and a custom ReLU with a backward pass that drops into torch.autograd like any native op, plus a concrete mental model of the grid-block-warp-thread hierarchy and where your handwritten code still loses to cuBLAS.

The substance is the CUDA you're usually shielded from: __global__ kernels, thread indexing with blockIdx.x * blockDim.x + threadIdx.x, the <<<blocks, threads>>> launch syntax, 2D dim3 grids, and, for the tiled matmul, declaring __shared__ float sA[TILE][TILE], staging tiles, and synchronising with __syncthreads(). The matmul step benchmarks your kernel against torch.matmul so you see the 5-10x gap with cuBLAS yourself; that's the concrete answer to "when is it worth writing a kernel?" (custom fused ops, unusual shapes, ops vendor libs don't cover) versus "when should I just call cuBLAS / cuDNN / FlashAttention?" The autograd step uses torch.utils.cpp_extension.load_inline, the same JIT path FlashAttention, xformers, and most of torchao ship on, to compile CUDA source and bind a forward + backward kernel that loss.backward() flows through cleanly.

Prerequisites: comfort with PyTorch tensors plus enough C/C++ to read pointer arithmetic and function signatures. Prior CUDA experience isn't assumed, launch syntax and thread indexing are introduced from scratch at the start. The sandbox is a real NVIDIA GPU pod with nvcc, the CUDA toolkit, PyTorch, and the load_inline compile cache preinstalled, so iteration is measured in seconds per kernel edit rather than minutes of build setup. The lab also feeds the exam prep track for NVIDIA's NCP-AII, NCP-ADS, and NCP-GenL certifications, where low-level GPU programming shows up.