B200アテンションカーネルをゼロから約SOTAレベルまで60の図で解説

Contents Open Hide + - Download the videoYou’ll generate these videos with B200 attention kernel built belowGenerated by video model using the Blackwell B200 attention kernel implemented in this post. See Capstone Project below.Percent of same-run Stock FA4 across 14-kernel progression and minor optimizations covered at the end. Each implementation is timed on its preferred contiguous layout; see Benchmark Calibration.What this blog is aboutIn this blog, we build a dense B200 attention kernel from scratch in CUDA and a little PTX, from baseline to 94.4% of FlashAttention-4 performance on 4K, 8K, and 16K shapes used in the FA4 paper.The main contribution is the visual guide: a beginner-friendly progression built around 60 diagrams. We first build an intuitive understanding of how the naive kernel works, then add one optimization at a time, with detailed diagrams, concise explanations, and code.For the capstone project, we plug the final kernel into a video-generation model.The focus here is not squeezing every last percent of performance, that’ll be the focus of my next blog.It’s one of the hardest kernels out there, running on the latest hardware, so it’ll be fun.Research: This is meant to give you a foundation for doing your own GPU kernel research on the latest hardware. We focus on B200 attention, but many of the concepts and mental models apply beyond this kernel. By the end, you’ll understand this kernel and be able to come up with your own optimization ideas.Prerequisites: I included a beginner on-ramp. If you are a complete beginner, or feel your CUDA foundations are shaky, follow this footnote1 before continuing. I still assume basic CUDA familiarity, but no prior knowledge of Blackwell. New concepts are introduced visually, one piece at a time, and only when they become necessary. So if Blackwell is new to you, just read on. You should be able to follow the progression.I assume you know what attention is. I also assume some familiarity with ideas behind online-softmax. This part is not Blackwell specific, and there are plenty resources on it.You do not need to have a B200 at home to follow along, I don’t have one either.Code: All code is available in my B200 Attention repo. After the baseline, each chapter adds one main optimization. The source code is organized so neighboring kernels are mostly easy to diff. So you can see what each optimization changes, one at a time.Capstone: At the end, we will plug the kernel you understand into a video-generation model and generate beautiful videos. See the Capstone Project after the main chapters.DSLs vs CUDA: The original FA4 is written in CuTe, I personally find raw CUDA + a bit PTX simpler to understand (less abstraction layers), so we’re going to implement our kernel in cuda. We will not translate FA4’s CuTe implementation into CUDA syntax, but understand a fast B200 attention kernel in general, through a clean and intuitive progression. Still, FA4 is one of the main references for this work, and most optimization ideas are adapted from it (and FA4 itself is based on cutlass and cute-dsl kernels, see Acknowledgments).Existing Resources: There are excellent resources explaining optimized matmuls on H200 and B200 (see Acknowledgments). But for B200 attention I haven’t found a deep dive explanation I wanted. Some resources cover the final resulting kernel and don’t explain the progression or lower level motivations behind most optimizations. Others stay high-level and superficial, like summarizing the pipeline and warp roles, but skip most of the work and handholding needed to actually understand the kernel. None gave me the deep explanation I wanted.Scope: The kernel we gonna be optimizing is dense, head dim 128, non-causal, BF16.I all you have is AI, we have the same AI as you and are probably better at using it– tomcr00sePart I — BasicsChapter 1 – Baseline Blackwell B200 Attentioni. Roadmap for this chapterWe’ll first study what work each CTA does, and how the work gets assigned to different CTAs of B200 (Work Parallelization section).Then we’ll take an optional detour for beginners, covering logical vs physical representation, pointers, and tiled matmul.Then we’ll zoom into our b200 attention kernel and discuss what happens inside each CTA.All that will be visually explained in much more details later. Just showing the lay of the land for now. Then I will link the code (maps directly to our diagrams).Baseline attention kernelLet’s start understanding the first kernel. Our later kernels mostly use the same math and the same tcgen05 concepts introduced here. So, in this chapter I’m covering the foundations we’ll use throughout. That’s why the first chapter is longer than later chapters.The first kernel exists as a starting point, produces correct numeric results, but not nearly as efficient as our later kernels. This first kernel is already nontrivial and uses many Blackwell-specific features, we will gradually cover them below.The main bottleneckAttention is basically two matmuls with a softmax in between. Softmax does far fewer FLOPs, but runs on the ALU and MUFU units, not on tensor cores. And on Blackwell B200, tensor-core throughput roughly doubled while the exp units stayed mostly unchanged (the FA4 paper calls this asymmetric hardware scaling). So at our tile sizes, softmax takes about as many cycles as the very beefy MMAs, so it’s the main bottleneck of this kernel. Most of our optimizations attack this from two sides: making that softmax work cheaper, and overlapping it with the matmuls (ie hiding it in the MMA’s shadow).ii. Work ParallelizationBefore discussing the B200 specific details, let’s first look at how the work is partitioned.Think of Q, K, V, and the output tensors all having the same shape B, num_heads, seq_len, head_dim.We split our tensors into tiles, so that they can fit into fast but small on-chip memory. Our tile sizes are [128, 128], as shown above.We schedule as many CTAs as there are O (Output) tiles, each CTA produces a single O tile. And collectively all CTAs produce the entire Output tensor (all of its tiles). CTAs execute in parallel (purple arrow).Within each CTA, the K/V-loop work is sequential (orange arrow). Let’s zoom into one CTA, as shown above. To produce its O tile, a CTA loads the corresponding Q tile once, then loops over all K/V tile pairs. At each iteration of its loop, it computes S_tile = Q_tile @ K_tile.T, applies online softmax to produce P_tile, and accumulates P_tile @ V_tile into this CTA’s private buffer O-tile. Inside a CTA, that CTA-private O tile is used as a “running accumulator” (the CTA updates it at each iteration of its K/V loop). After the final iteration of the K/V loop, CTA normalizes its O-accumulator and stores this completed output tile to global memory.Don’t worry if some of this doesn’t make sense yet, I’ll explain each step in detail below.Basically, this split is similar to tiled matmul: different CTAs produce different output tiles, while the loop for one output tile stays inside its CTA. But unlike matmul, attention additionally carries online-softmax state across that K/V loop.In pseudocode, the high-level flow looks like this: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 # purple arrow: independent work items mapped across CTAs parallel_for batch, head, q_tile: # one CTA starts here Q = load_q_tile(batch, head, q_tile) # orange arrow: sequential loop inside this CTA for K, V in kv_tiles(batch, head): # Section A: Q @ K.T -> S S = Q @ K.T # Section B: S -> P, update rowmax/rowsum, and correct old O P, softmax_state, O = online_softmax_update( S, softmax_state, O ) # Section C: add the current P @ V contribution O += P @ V # final normalization of this single O tile output[batch, head, q_tile] = O / softmax_state.rowsum The outer loop over Q tiles becomes the CUDA grid. The inner loop over KV tiles remains inside each CTA.Sections labeled in the pseudocode above map directly to B200 code of our first kernel, an