From-scratch CUDA attention kernel reaches 94.4% of FlashAttention-4 on NVIDIA B200

A 14-step climb toward FlashAttention-4 speed

A blog post that reached the Hacker News front page walks through building a dense attention kernel for NVIDIA's Blackwell B200 GPU from scratch, using plain CUDA with a small amount of PTX. According to the author, the finished kernel delivers 94.4% of FlashAttention-4's performance on the 4K, 8K and 16K sequence-length shapes used in the FA4 paper. The stated goal is teaching rather than record-setting; squeezing out the last few percent is reserved for a follow-up post.

The material is organized as a progression through 14 kernel versions, with each step adding a single optimization and roughly 60 diagrams carrying the explanation. The scope is intentionally narrow: dense, non-causal attention with a head dimension of 128, computed in BF16. As a capstone, the final kernel is plugged into a video-generation model to produce sample videos.

How the kernel partitions its work

Per the post, Q, K, V and the output tensor all share the shape of batch, heads, sequence length and head dimension, and are cut into 128 by 128 tiles. One CTA is assigned to each output tile, and those CTAs run in parallel across the grid. Inside a CTA the work is sequential: it loads its Q tile once, then loops over the K/V tile pairs. Each iteration computes S = Q @ K^T, applies an online-softmax update to produce P while maintaining running row maxima and sums, and accumulates P @ V into a private output buffer. Once the loop finishes, that buffer is normalized and written to global memory. The author compares the structure to a tiled matmul, with the added complication that softmax state must survive across the K/V loop.

Softmax, not matmul, is the bottleneck

The post's central technical point is that attention amounts to two matrix multiplications with a softmax in between, and on Blackwell the softmax is what limits throughput. Although it performs far fewer FLOPs, softmax runs on the ALU and MUFU units rather than the tensor cores. On the B200, tensor-core throughput roughly doubled while the exponent units stayed largely the same, a mismatch the FA4 paper calls "asymmetric hardware scaling." At the author's tile sizes, the softmax consumes about as many cycles as the matmuls. The optimizations in the progression therefore attack from two directions: making the softmax work cheaper and overlapping it with the tensor-core math.

Raw CUDA instead of CuTe

FlashAttention-4 itself is written in CuTe, but the author finds raw CUDA plus PTX easier to follow because it has fewer abstraction layers. The result is not a line-by-line translation of FA4 but an independent construction, with FA4 credited as a main reference and the source of most optimization ideas. A companion repository arranges the code so neighboring kernel versions can be diffed, isolating what each step changes.

The on-ramp is deliberately gentle: readers need basic CUDA familiarity plus an understanding of attention and online softmax, but no prior Blackwell knowledge, since hardware-specific machinery such as tcgen05 tensor-core concepts is introduced only when needed. The author also notes that no B200 is required to follow along, not owning one either.

Why it matters

Attention kernels sit on the critical path of training and inference, and Blackwell changes their economics: with tensor cores scaling much faster than the special-function units that handle exponentials, kernels written for earlier GPUs leave throughput stranded. According to the author, existing explanations of B200 attention either document a finished kernel without the progression that produced it or stay at a high level. A step-by-step account that lands within six percent of FlashAttention-4, complete with diffable code, gives kernel engineers a working template for reasoning about Blackwell's trade-offs, and the mental models carry over well beyond this one kernel.