# Three kernels for a chip the ecosystem skipped
> FlashAttention-2 and NVFP4 GEMM, tuned from scratch for the RTX 5090, and what the profiler taught me when every textbook optimisation turned out to be a red herring.
Source: https://rfriedmann.de/blog/optimizing-kernels-consumer-blackwell/
Published: 2026-06-11 · Track: log · Level: Expert
There are three files in [imp](https://github.com/kekzl/imp) you can compile and
run on their own, no engine, no CUTLASS, no cuDNN. Each is a single `.cu` file:
the kernel, a CPU reference, a correctness check, and a timing loop. They live in
[`tools/standalone`](https://github.com/kekzl/imp/tree/main/tools/standalone),
and they exist because the interesting part of building a fast inference engine
isn't the engine, it's the handful of kernels underneath, and the profiler runs
that shaped them.
This post is about those runs. And about the thing nobody warns you about when
you read a GPU optimisation guide: most of the "production-standard" tricks are
answers to a problem you might not actually have.
## The chip the playbook wasn't written for
Consumer Blackwell, the `sm_120a` chip in a 5090, looks like a small B200 and
behaves like a different animal. It has FP4 tensor cores and fast memory, but
it's missing the entire async-matrix-multiply instruction family that modern
data centre kernels are built on: no `wgmma` (the Hopper `sm_90` instruction)
and no `tcgen05`/UMMA (the datacenter `sm_100` path). What it does have is
`mma.sync` (the Ampere-lineage synchronous family), block-scale variant, which
runs the FP4 matrix multiply but blocks the issuing warp. TMA (`cp.async.bulk.tensor`)
is present, so the missing piece is the async matrix multiply itself, not the
async copy. The flagship FlashAttention-4 design, the warp-specialised CUTLASS
GEMM, both assume hardware that isn't on the chip.
So you can't port. You write for the instructions you actually have, and you let
the profiler tell you where it hurts. Three kernels, three different lessons.
## Kernel one: FlashAttention-2, and the register that can't move
The thing that makes FlashAttention-2 fast isn't the tiling, it's that the
running output lives in registers and never gets written to shared memory during
the whole attention loop. That one constraint dictates everything else, right
down to which tensor-core instruction you're allowed to use. (The convenient
high-level one hides exactly the detail you need; you have to drop to the raw
instruction.)
From there it's eight profiler-driven steps, each aimed at whatever the previous
run flagged as the bottleneck, measured, never guessed. A couple of the
highlights:
- **Pad a shared-memory buffer by 8 elements.** One line. It killed a
7.8-million-event memory-bank conflict and bought **+25%**.
- **Load values through an async copy instead of a transposing read**, hiding a
stall the profiler had been complaining about. Another **+25%**.
- **Double-buffer the key/value tiles**, prefetching the next one while computing
the current. **+13%.**
FlashAttention-2 forward - TFLOP/s on a single RTX 5090
[diagram omitted — see the page for the chart]
Same kernel, eight optimization steps apart. Roughly 2.2× the untuned reference (the from-scratch kernel before any profiler step), and past imp's production FA2 at long sequences. TFLOP/s at S = 4k / 8k / 16k.
These are single measurements unless noted. The throughput numbers (113.7 / 157.7 / 186.8 TFLOP/s
at 4k / 8k / 16k) need their input dtype and accumulator stated to be read
correctly: 186.8 sits above the 5090's ~105 TFLOP/s dense FP16 peak with FP32
accumulation, so it is only physical against an FP16-accumulate or a sparse peak.
That's roughly 2.2× the untuned reference kernel (the from-scratch starting
point, before any of the eight profiler steps), and past imp's own production
attention at
long sequences. Not every idea survived contact: f16 accumulation in the score
step *looked* obviously faster, got built, and was reverted when it turned out
the unpacking cost more than it saved. The trail has nearly as many reverts as
wins, which is the point.
And then it stops. The remaining stall is the async-copy wait, and the textbook
fix, a deeper pipeline, doesn't fit: three buffers need 102 KB of shared
memory against a 99 KB ceiling. The deeper reason it wouldn't help anyway is that
without an async matrix-multiply instruction, the dependency chain caps how much
work can overlap no matter how much you raise occupancy. The gap to a B200-class
kernel here is silicon, not code.
## Kernel two: the GEMM, where the playbook lied
This is the matrix-multiply that dominates prefill, running on the chip's peak
FP4 path: 4-bit weights with a hardware block scale. The optimisation trail goes
from a 1.8 TFLOP/s untuned reference kernel to **972**, and the *shape* of that
climb is the whole story. The 540× figure is the distance from that untuned
starting point, not a claim about the engineering, the interesting work is in the
later steps where the easy gains are already spent.
[diagram omitted — see the page for the chart]
1.8 → 972 TFLOP/s, where 1.8 is the untuned reference kernel. The dip at step 4 is the real block scales; the jump at step 5 is the re-diagnosis (request-bound, not bandwidth-bound) plus a layout fix. The final number beats imp's production CUTLASS path. Single measurements unless noted.
The first four steps are the ones you'd expect: pack the weights tightly,
double-buffer the loads, use a bigger tile with register blocking, wire in the
real block scales. That gets you to ~680 TFLOP/s and a clean diagnosis: the
kernel is bound on the L2 cache.
Now, the standard responses to an L2 bottleneck are well known, and I tried them.
Every single one was a dud:
- **Threadblock swizzle**, reverted, −7%. It improves cache *hit rate*; the same
bytes were flowing through the cache *pipe* regardless.
- **Three-stage pipeline**, no gain. That's a fix for *latency*, and this wasn't
a latency problem.
- **TMA + warp-specialisation**, the full data centre answer, built from scratch.
It worked, it was bit-exact, and it was **slower**. (More on that shortly.)
Three production-grade levers, three red herrings, because each one answered a
question the kernel wasn't asking. The actual fix came from distrusting the
diagnosis itself. Reading the profiler more carefully: the cache was at 82% on
*requests* but only 44% on *sectors*. That's not "out of bandwidth", that's
"wasting most of every cache line". The data layout had scattered each tile's
rows across separate cache lines at 25% fill. Two layout changes, store each
tile as one contiguous block, then interleave the packed values so a fragment
loads in one shot, and throughput jumped from 683 to **972 TFLOP/s**. That's
about 48% of the card's raw FP4 peak as quoted on the 2:4-sparse class number
(972/0.48 lands near the ~2025 TFLOP/s sparse figure); against the 5090's dense
NVFP4 throughput, which is roughly half the raw/sparse rate and lower still with
FP32 accumulation, the same 972 is a much larger fraction of what dense actually
delivers. It also beats imp's own production
CUTLASS path.
The simple layout change beat the brute-force kernel rewrites by ~40%.
The lesson I keep having to relearn: the optimisation you reach for encodes an
assumption about your bottleneck. Check the assumption before you spend a week on
the fix.
## Kernel three: the one I kept because it lost
The third file is a documented failure, and its header says so in capital
letters. It's the full TMA + warp-specialisation implementation, the
data-centre-grade approach everyone would reach for first. It's correct, it
doesn't deadlock, and it loses: 509 TFLOP/s against the simpler kernel's 629,
because the elaborate version starves the chip of parallelism to feed its one clever
loader. These two are a head-to-head on a different shape/config than the 683/972
numbers above, so compare the 509 and 629 to each other, not across sections.
I kept it in the repo on purpose. A benchmark table that only shows wins is
marketing; a negative result you can compile and run is evidence. On this chip,
for this kernel, the technique everyone would reach for is the wrong one, and
the file is there so the next person doesn't have to spend a week proving it.
## Why they're standalone
These three could have stayed buried in the engine. They're carved out as single,
runnable files because that's the honest shape of the knowledge: not "imp is
fast", but "here's the exact kernel, here's the profiler trail that shaped it,
here's the reference that proves it's correct, compile it yourself". Each one
builds in one command and checks its own answer before it prints a number.
And, like everything else in imp, every line was written by [Claude
Code](https://claude.ai/claude-code). My part was reading profiler output and
deciding which bottleneck to chase next; the agent wrote the kernels, the
references, and the reverts. The reverts matter as much as the wins. A surprising
amount of fast GPU code is just the discipline to delete the clever thing that
didn't pay off.
If you have a 5090, the files are in
[`tools/standalone`](https://github.com/kekzl/imp/tree/main/tools/standalone):
MIT, self-contained, and each one prints its own TFLOP/s when you run it.