Three kernels for a chip the ecosystem skipped

There are three files in 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, 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%.

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.

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. 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: MIT, self-contained, and each one prints its own TFLOP/s when you run it.