I'm Richael (dh8116), a Year 11 student in Auckland. Since late July I've written one Triton kernel a week and benchmarked each against PyTorch, publishing the numbers even when PyTorch won, which was often. Here's what nine weeks taught me, kernel by kernel.
The ones that won
Fused softmax (week 2). Naive softmax makes three passes over memory. Fusing them into one took it from ~55 GB/s to ~230 GB/s, a 4x jump, and it held a tighter line than PyTorch's own kernel. Lesson one: on a memory-bound op, the passes you don't make are the speedup.
Fused cross-entropy (week 7). Forward and backward in one pass, with the gradient written straight over the logits buffer, so PyTorch's second [N, V] tensor never exists. On a T4 at vocab 131,072 in fp16: 15.90 ms vs 24.03 ms (1.51x), and 1.67x less peak memory.
The first version of this benchmark said 2.5x. That was wrong: I was comparing my fused kernel against an unfused PyTorch baseline. Comparing fused against fused is how I caught myself handicapping the baseline.
RoPE (week 6). Not the speedup, the elegance: backward reuses the exact same kernel as forward. You just flip the sign on sin.
The ones that lost, and why that was the useful part
Flash attention (week 3). Correct, but about 25x slower than PyTorch on my T4.
Matmul (week 4). The output matched, but throughput sat flat around 1 TFLOPS while cuBLAS climbed to 38. Reading the compiled PTX showed zero tensor-core instructions: the compiler never emitted them for my layout.
Fused linear + cross-entropy (week 8). Chunking the lm_head projection into the loss means the [batch, vocab] logits never exist at all. That's about 10x less activation memory, and it was slower at every single size. Sometimes memory is the thing you're buying, and you pay for it in time.
Fused SwiGLU MLP (week 9). Level with torch.compile on time. The win is memory: it holds two big activation tensors between forward and backward, where eager holds four. A compiler will fuse for you, but it won't throw an activation away.
What I'd tell anyone starting
- Benchmark fused against fused. An unfused baseline makes every kernel look like a win.
- Read the PTX when the numbers are flat. The matmul answer was sitting in the compiled output the whole time.
- Memory is a result too. Two of my best kernels "lost" on time and won on memory, and that trade decides batch size.
- Publish the losses. They taught me more, and they're the posts people actually argue with.
Code: https://github.com/dh8116/triton-kernels
Every write-up, with methods: https://dh8116.github.io/blog
Next on my list is revisiting flash attention on Ampere, because the T4 result deserves a second look.
Top comments (0)