Level 3 · Performance engineering with a public bar
Can you write a correct Triton kernel for one real operation that beats PyTorch’s implementation – and explain where the speed comes from?
Learn GPU programming through Triton’s official tutorials, then pick one operation (a fused normalisation, softmax or a small fused block), beat PyTorch on it, and explain the result with profiling.
Performance engineers are some of the most sought-after people at labs. A correct, profiled kernel with an honest explanation is the warm-up for the Anthropic take-home and the speedrun.
01Do the official Triton tutorials in order
Vector add, fused softmax, matrix multiplication, layer norm. Type them out and benchmark each against PyTorch.
$ pip install triton02Learn to read a profile
Use the PyTorch profiler or NVIDIA Nsight to see time per kernel. Work out whether each op is limited by memory bandwidth or by compute.
03Pick one target operation
Something real from your GPT-2 training loop where fusing steps should save memory traffic. Write the PyTorch reference first.
04Write correctness tests before speed
Compare against the reference across shapes, dtypes and edge cases (non-power-of-two sizes) with sensible tolerances.
05Write and tune the kernel
Get it correct, then tune block sizes. Benchmark properly: warm-up, many repeats, report the median and spread.
06Compare fairly
Also compare against torch.compile, not just eager PyTorch, and note settings like TF32 – benchmarks have been criticised for unfair baselines.
07Optional: try KernelBench problems
KernelBench’s fast_p metric counts kernels that are correct and faster than the baseline by a factor p. Try a few problems from its easier levels.
08Write it up
Speedup table across shapes, the profile that explains it, and where your kernel loses.