Home › Blog › Frontier Lab Engineering — Custom Kernels
Frontier Lab Engineering Practicum · Article 5 of 9 🧵

Writing Custom Kernels and Why You Occasionally Have To

Article 4 introduced Model FLOPs Utilization as the honest measure of whether a GPU is doing useful work. This article covers what happens after you find a job with stubbornly low MFU that framework-level fixes can't touch: why PyTorch's built-in operators sometimes aren't fast enough, what a GPU kernel actually is, why Triton has become the practical entry point for most engineers instead of raw CUDA, and the profiling discipline that separates a justified custom kernel from wasted engineering effort.

FL
FrontierAGI Team

The Layer Almost Nobody Touches, Until Someone Has To

Most research engineering work in this series so far happens entirely in Python: writing training loops, wiring up configs, debugging distributed jobs, and reasoning about scheduler queues. All of that sits on top of an assumption — that when your code calls a matrix multiply or an attention operation, the underlying GPU implementation is already about as fast as it can be. Most of the time that assumption holds, because PyTorch's built-in operators are themselves backed by heavily-optimized vendor kernels. Occasionally it doesn't, and understanding why is what separates an engineer who can only tune hyperparameters from one who can actually make a slow model fast. This article is deliberately narrow: it is not a CUDA programming course, and it will not make you a kernel engineer. It is the honest map of when this level of the stack matters, what the entry point actually looks like today, and — just as importantly — when reaching for it is the wrong move.

2 Fundamental bottleneck types every kernel decision reduces to: compute-bound, memory-bound
1st Step before writing any kernel, always: profile the actual bottleneck
2022 Year FlashAttention showed a rewritten kernel — not a new algorithm — could be the real speedup
Part 1 — Compute-Bound vs. Memory-Bound

The Question That Comes Before Any Optimization

Horace He's widely-cited engineering post, "Making Deep Learning Go Brrrr From First Principles", frames the entire problem cleanly: every GPU operation is either compute-bound (the GPU's arithmetic units are the bottleneck — it's actually crunching numbers as fast as it can) or memory-bound (the GPU is mostly waiting on data moving between its slow high-bandwidth memory and its fast on-chip memory, with the arithmetic units sitting comparatively idle). This is the same "is the GPU on vs. is the GPU doing useful work" distinction from Article 4's MFU discussion, pushed one level deeper: a memory-bound operation can show high GPU utilization in a naive sense while still wasting most of its theoretical compute capacity, because the numbers it's waiting to receive are the actual constraint, not the arithmetic itself.

"Flops don't matter if you're not compute-bound, and reducing flops can make a memory-bound op slower, not faster, if it doesn't also reduce memory movement." — the core insight behind most real-world custom-kernel work

This matters because it inverts a common intuition. Many engineers assume the path to a faster model is always "fewer floating-point operations" — a smaller, more elegant formula. But if an operation is memory-bound, an algebraically simpler formula that still moves the same amount of data through memory won't help at all. This is precisely why operations like attention — mathematically simple, but involving large intermediate tensors that don't fit in fast on-chip memory — became one of the first and most famous custom-kernel success stories, covered in Part 4 below.

Part 2 — Profile Before You Optimize

The Discipline That Prevents Wasted Weeks

The single most common mistake at this level of the stack is writing a custom kernel for an operation that was never actually the bottleneck. PyTorch's own torch.profiler documentation exists precisely because "which operation is slow" is an empirical question, not something you can reliably guess from reading code — a training step involves dozens of operations, and intuition about which one dominates wall-clock time is wrong often enough that skipping this step routinely wastes engineering time on the wrong target.

# the discipline, roughly, before touching a single line of kernel code with torch.profiler.profile(activities=[ProfilerActivity.CUDA]) as prof: train_step(batch) print(prof.key_averages().table(sort_by="cuda_time_total")) # only now do you know which op actually dominates — and whether it's compute- or memory-bound

The pattern that recurs across real engineering accounts is: profile first, identify the specific operation and its bottleneck type, estimate the realistic ceiling on improvement (a memory-bound op is bounded by memory bandwidth, not by however clever the replacement code is), and only then decide whether a custom kernel is worth the engineering and maintenance cost — a cost this article returns to directly in Part 5.

Part 3 — Triton vs. Raw CUDA

Why Most Engineers Never Touch Raw CUDA

OpenAI's "Introducing Triton: Open-Source GPU Programming for Neural Networks" describes the actual reason Triton has become the practical entry point for this work: writing efficient raw CUDA requires reasoning explicitly about memory hierarchies, thread scheduling, and hardware-specific tuning that historically took specialist GPU engineers years to master. Triton lets an engineer write Python-like code that the compiler lowers to efficient GPU code automatically handling much of that low-level complexity, trading some of the theoretical peak performance a hand-tuned CUDA expert could reach for a dramatically shorter path to a working, meaningfully-faster kernel.

ApproachLearning CurveTypical Use Case
PyTorch built-in opsNone — this is the defaultThe vast majority of code; vendor-optimized already
TritonModerate — Python-like syntax, GPU concepts still requiredA specific profiled bottleneck where a fused custom op would help
Raw CUDA / C++Steep — genuine specialist skillExtracting the last percentage points of peak hardware performance, novel hardware features
torch.profiler Triton NVIDIA Nsight Systems torch.compile

It's also worth naming that a large share of what used to require hand-written kernels is now handled automatically by torch.compile, which can fuse operations and generate Triton kernels behind the scenes without an engineer writing GPU code directly at all — meaning the honest first question, before writing anything by hand, is whether the compiler can already close the gap.

Part 4 — The FlashAttention Case Study

A Real Kernel That Changed the Field

Dao et al.'s "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (2022) is the canonical real-world example of everything in Parts 1 through 3 combined: attention was identified as memory-bound (Part 1) rather than compute-bound, because standard implementations materialize large intermediate attention matrices that don't fit in fast GPU memory; the fix was not a new attention algorithm producing different numbers, but a kernel-level rewrite that restructures the computation to keep more of it in fast on-chip memory and avoid ever materializing the full intermediate matrix — the same mathematical result, computed with far less memory traffic. The paper reports substantial wall-clock speedups and memory reductions on real training and inference workloads, and its core technique was significant enough to be adopted directly into mainstream frameworks rather than remaining a standalone library.

Standard Attention Compute full N×N score matrix Write full matrix to slow HBM memory Read it back for softmax + output FlashAttention Process in small on-chip tiles Fuse softmax + output, stay on-chip Never materializes full matrix
Same mathematical result, radically less memory traffic — the exact shape of a justified custom kernel.
Part 5 — When Not To

The Cost Side of the Ledger

A custom kernel is not a free win. It is code that must be maintained across new GPU hardware generations, verified for numerical correctness against the reference implementation (an easy place for silent bugs, echoing Article 3's warning about training that runs without crashing while producing wrong numbers), and re-tuned as workload shapes change. Sculley et al.'s "Hidden Technical Debt in Machine Learning Systems" — already cited in Article 2 of this series for configuration sprawl — applies just as directly here: a hand-written kernel is exactly the kind of specialized, hard-to-read, high-maintenance code the paper warns accumulates debt fastest, and it should only be taken on when the profiled, quantified speedup clearly justifies that ongoing cost.

1
You've profiled and confirmed the bottleneck, not guessed at it from reading code.
2
You've checked whether torch.compile or an existing library already solves it before writing anything by hand.
3
You've estimated the realistic ceiling — a memory-bound op won't get faster no matter how clever the arithmetic is.
Part 6 — Real Scenarios

Real Scenario Walkthroughs

🐢Scenario A — Optimizing the Wrong Operation for a Week
An engineer notices a custom loss function is "clearly" the slow part of a training step by reading the code, and spends a week hand-optimizing its arithmetic. A profiler run (Part 2) — done only afterward — reveals the loss function was under 2% of step time all along; the real cost was a data-loading stall (echoing Article 4's Scenario A on low-MFU jobs being starved, not compute-bound). The week of kernel work produced a measured 0% end-to-end speedup.
The lesson: intuition about which operation is slow is wrong often enough that skipping the profiling step is never actually faster in expectation, even though it feels faster to start coding immediately.
🧪Scenario B — A Custom Kernel That Was Numerically Wrong
A hand-written Triton kernel replacing a standard normalization operation passes a quick visual smoke test and ships. Weeks later, a training run shows subtly worse final accuracy than an equivalent baseline; the root cause turns out to be a small numerical-precision difference in the custom kernel versus the reference implementation, invisible in casual testing but real over millions of training steps.
The lesson: a custom kernel needs the same numerical-correctness rigor as any other change to core training math — a fast wrong answer is worse than a slow correct one.

Readiness Checklist

1
Can you explain the difference between a compute-bound and a memory-bound operation, and why it changes what "optimization" even means?
2
Have you ever profiled a training step before deciding what to optimize, rather than guessing from reading the code?
3
Can you explain, in one sentence, why FlashAttention is faster despite computing the exact same mathematical result?
4
Would you know how to numerically validate a custom kernel against a reference implementation before trusting it in a real training run?

⚠️ What's Missing or Uncertain

This article intentionally does not teach CUDA or Triton syntax. Writing a correct, fast kernel is a genuine specialist skill that takes real practice beyond what any single article can convey; the goal here is the decision framework — when this level of the stack is worth engaging with at all, and how to reason about it responsibly — not a tutorial. The specific tools (Triton, torch.compile, profiler APIs) also change rapidly as frameworks evolve.

Where This Series Goes Next

Article 6 moves from the deepest, most specialized layer of the stack to the opposite end of the daily experience: the on-call reality of monitoring, alerts, and 2am incidents — what happens when a job like the ones profiled and optimized in this article fails in production, at an hour when the person paged is not the person who wrote the kernel.

🎥 Recommended Videos

🧭 Closing — A Specialist Skill, Not a Default Habit

🎯 The Bottom Line
Writing a custom GPU kernel is one of the few pieces of frontier-lab engineering work that most research engineers will do rarely, and a smaller number will do often — but understanding when framework-level code stops being fast enough, why compute-bound and memory-bound bottlenecks demand entirely different fixes, and why profiling always comes before optimizing is universal knowledge, not specialist knowledge. FlashAttention remains the clearest proof that this level of the stack can matter enormously; the discipline in this article is what keeps that possibility from turning into wasted weeks on the wrong operation.