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 data centre 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%.

Each number is the mean of 200 back-to-back launches (100 for the GEMM below) after at least 1.5 s of clock warm-up, built with nvcc from the CUDA 13.3.0 container, as in the published files. Driver 610.47 was logged by imp’s roofline runs on June 10 and 12; power cap and per-launch spread were not recorded. The throughput numbers (113.7 / 157.7 / 186.8 TFLOP/s at 4k / 8k / 16k) are dense FP16 with QK^T accumulated in FP32: 186.8 is 74% of the 253 TFLOP/s the 5090 measures for that mode through mma.sync (FP32 accumulate runs at a quarter of the FP16-accumulate rate per clock).

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 2019 TOPS, the FP4 peak measured with a saturated dense mma.sync block-scale microbenchmark (FP32 accumulation, like the kernel itself); the 3354 TOPS datasheet figure is not reachable through mma.sync on this chip.

Correction (2026-06-17): this post first said the kernel also beats imp’s production CUTLASS path. It does not. That comparison set this square-shape benchmark against production’s ~41% on small-M model shapes. Measured head to head at the same square shapes, production CUTLASS reaches 903 / 1284 TFLOP/s at 4k / 8k against this kernel’s 813 / 972, and it stays the production GEMM.

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. The file header records the pair but not the problem size it was measured at, and 629 matches none of the 4k/8k 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.