{"slug": "b200-attention-kernel-from-scratch-to-near-sota-in-60-diagrams", "title": "B200 Attention Kernel from Scratch to Near-SOTA in 60 Diagrams", "summary": "A new technical blog post by Iaroslav Elistratov presents a visual guide to building a B200 attention kernel from scratch in CUDA and PTX, achieving 94.4% of FlashAttention-4 performance on 4K, 8K, and 16K sequence lengths. The post includes 60 diagrams, a 14-kernel progression, and a capstone project integrating the kernel into a video-generation model.", "body_md": "# B200 Attention Kernel from Scratch to Near-SOTA in 60 Diagrams\n\nBuild and understand one of the most complex GPU kernels on the latest hardware\n\n60 diagrams · 14-kernel progression · 94.4% of FlashAttention-4 · CUDA/PTX · video-generation capstone\n\n## Contents Open Hide\n\n# What this blog is about\n\nIn 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.\n\nThe main contribution is the *visual guide*: a beginner-friendly progression built around 60 diagrams.\nWe first build an intuitive understanding of how the naive kernel works, then add one optimization\nat a time, with detailed diagrams, concise explanations, and code.\n\nFor the capstone project, we plug the final kernel into a video-generation model.\n\nThe focus here is not squeezing every last percent of performance, that’ll be the focus of my next blog.\n\n*It’s one of the hardest kernels out there, running on the latest hardware, so it’ll be fun.*\n\n**Research:**\nThis is meant to give you a foundation for doing your own GPU kernel research on the latest hardware.\nWe focus on B200 attention, but many of the concepts and mental models apply beyond this kernel.\nBy the end, you’ll understand this kernel and be able to come up with your own optimization ideas.\n\n**Prerequisites:**\nI included a [beginner on-ramp](#iii-beginner-onramp--optional).\nIf you are a complete beginner, or feel your CUDA foundations are shaky, follow this footnote 1 before continuing.\nI still assume basic CUDA familiarity, but no prior knowledge of Blackwell.\nNew concepts are introduced visually, one piece at a time, and only when they become necessary.\nSo if Blackwell is new to you, just read on. You should be able to follow the progression.\n\nI 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.\n\nYou do not need to have a B200 at home to follow along, I don’t have one either.\n\n**Code:** All code is available\nin my\n[B200 Attention repo](https://github.com/IaroslavElistratov/b200-attention).\nAfter the baseline, each chapter adds one main optimization.\nThe source code is organized so neighboring kernels are mostly easy to diff.\nSo you can see what each optimization changes, one at a time.\n\n**Capstone:** At the end, we will plug the kernel *you understand* into a video-generation model and generate beautiful videos. See the [Capstone Project](#capstone-project--generate-videos) after the main chapters.\n\n**DSLs vs CUDA:**\nThe 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.\nWe 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.\nStill, 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).\n\n**Existing Resources:** There are excellent resources explaining optimized *matmuls* on H200 and B200 (see Acknowledgments).\nBut for *B200 attention* I haven’t found a deep dive explanation I wanted.\nSome resources cover the final resulting kernel and don’t explain the progression or lower level motivations behind most optimizations.\nOthers 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.\nNone gave me the deep explanation I wanted.\n\n**Scope:**\nThe kernel we gonna be optimizing is dense, head dim 128, non-causal, BF16.\n\nIf all you have is AI, we have the same AI as you and are probably better at using it\n\n– tomcr00se\n\n# Part I — Basics\n\n## Chapter 1 – Baseline Blackwell B200 Attention\n\n### i. Roadmap for this chapter\n\nWe’ll first study what work each CTA does, and how the work gets assigned to different CTAs of B200 (Work Parallelization section).\n\nThen we’ll take an *optional* detour for beginners,\ncovering logical vs physical representation,\npointers, and tiled matmul.\n\nThen we’ll zoom into our b200 attention kernel and discuss what happens inside each CTA.\n\nAll 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).\n\n**Baseline attention kernel**\n\nLet’s start understanding the first kernel. Our later kernels mostly use\nthe same math and the same `tcgen05`\n\nconcepts introduced here.\nSo, in this chapter I’m covering the foundations we’ll use throughout.\nThat’s why the first chapter is longer than later chapters.\n\nThe 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.\n\n**The main bottleneck**\n\nAttention is basically two matmuls with a softmax in between.\nSoftmax does far fewer FLOPs, but runs on the ALU and MUFU units, not on tensor cores.\nAnd on Blackwell B200, tensor-core throughput roughly doubled while the exp units stayed mostly unchanged\n(the FA4 paper calls this *asymmetric hardware scaling*).\nSo at our tile sizes, softmax takes about as many cycles as the very beefy MMAs, so it’s the main bottleneck of this kernel.\nMost of our optimizations attack this from two sides: making that softmax work cheaper,\nand overlapping it with the matmuls (ie hiding it in the MMA’s shadow).\n\n### ii. Work Parallelization\n\nBefore discussing the B200 specific details, let’s first look at how the work is partitioned.\n\nThink of Q, K, V, and the output tensors all having the same shape `B, num_heads, seq_len, head_dim`\n\n.\n\nWe split our tensors into tiles, so that they can fit into fast but small on-chip memory.\nOur tile sizes are `[128, 128]`\n\n, as shown above.\n\nWe 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).\n\nWithin each CTA, the K/V-loop work is sequential (orange arrow).\nLet’s zoom into one CTA, as shown above.\nTo produce its O tile, a CTA loads the corresponding Q tile once, then loops over all K/V tile pairs.\nAt each iteration of its loop, it computes `S_tile = Q_tile @ K_tile.T`\n\n,\napplies online softmax to produce `P_tile`\n\n, and accumulates `P_tile @ V_tile`\n\ninto this CTA’s private buffer O-tile.\nInside 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).\nAfter the final iteration of the K/V loop, CTA normalizes its O-accumulator and stores this completed output tile to global memory.\n\nDon’t worry if some of this doesn’t make sense yet, I’ll explain each step in detail below.\n\nBasically, 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.\n\nIn pseudocode, the high-level flow looks like this:\n\n```\n 1\n 2\n 3\n 4\n 5\n 6\n 7\n 8\n 9\n10\n11\n12\n13\n14\n15\n16\n17\n18\n19\n20\n21\n22\n```\n\n | \n\n```\n# purple arrow: independent work items mapped across CTAs\nparallel_for batch, head, q_tile:\n\n    # one CTA starts here\n    Q = load_q_tile(batch, head, q_tile)\n\n    # orange arrow: sequential loop inside this CTA\n    for K, V in kv_tiles(batch, head):\n\n        # Section A: Q @ K.T -> S\n        S = Q @ K.T\n\n        # Section B: S -> P, update rowmax/rowsum, and correct old O\n        P, softmax_state, O = online_softmax_update(\n            S, softmax_state, O\n        )\n\n        # Section C: add the current P @ V contribution\n        O += P @ V\n\n    # final normalization of this single O tile\n    output[batch, head, q_tile] = O / softmax_state.rowsum\n```\n\n |\n\nThe outer loop over Q tiles becomes the CUDA grid. The inner loop over KV tiles remains inside each CTA.\n\nSections labeled in the pseudocode above map directly to B200 code of our first kernel, and will be addressed in detail in this chapter:\n\n- Section A)\n`Scores = Query @ Key.T`\n\n(tcgen05 in microtiles, K-major layout) - Section B)\n`Scores -> Probabilities`\n\n(reading from TMEM) - Section C)\n`O += Probabilities @ Values`\n\n(tcgen05 in microtiles, MN-major layout)\n\n### iii. Beginner Onramp — Optional\n\n*This section is optional.*\nSkip if you’re already familiar with tiled matmul,\nthe concept of memory layouts,\ndistinction between logical and physical view of the data,\nand how to use pointers to index into memory.\n\n**Similarity to tiled matmul:**\n\nLet’s first build one basic tiled-matmul intuition which we’ll use later.\n\nIn a regular tiled matmul, we’d split both operand matrices into tiles so that the smaller tiles fit into fast but small on-chip memory, allowing us to reuse each loaded tile multiple times directly from on-chip memory. At each iteration, we’d march along the reduction dim (K, in the diagram), matmul the matching A and B tiles, and accumulate the resulting partials into the same C tile. After processing all K tiles, adding these partials is mathematically equivalent to doing the full matmul without splitting K.\n\nAttention is similar in shape but not exactly:\n\n- generic tiled matmul keeps one C tile fixed and walks along its reduction dimension\n- attention keeps one Q tile fixed and walks along the K/V sequence dimension\n\n**Introducing layouts:**\n\nSo we have our tiles of data in global memory.\n\nSuppose we want to matmul a tile of A and a tile of B.\nWe do not want to implement this matmul in software (by manually looping over individual elements and computing dot products with scalar instructions),\nas this will be extremely slow. Instead we want to use tensor cores.\nTo use tensor cores on Blackwell, we must use the `tcgen05.mma`\n\nfamily of instructions.\n\nThese `tile_a`\n\nand `tile_b`\n\nI visualized above look like 2d matrices,\nbut of course these are just logical views of the stored data, actual memory is 1d.\n\nSo there’s always some organization of how our data is “laid out” in physical memory.\n\nLet’s say our A operand in GMEM is row-major. In which case physically rows are stored back to back in the 1d memory (shown above).\n\nTo find one element in that linear memory, for one [row, col] element, we start from the base pointer and use the strides to compute its offset.\n\nSuppose we have a tiny `[2,2]`\n\nmatrix, and we want to get element at `my_matrix[1,1]`\n\n.\n\nSo we have these 2d tiles which are somehow (one way or another) laid out in physical memory.\nA **layout** is the rule that maps logical coordinates, such as `[row, col]`\n\n,\nto physical memory addresses. Row-major is one such layout.\n\nBut, that `tcgen05.mma(tile_A, tile_B)`\n\nwhich we want to call doesn’t work on\nthat simple row-major physical layout. It only supports a small number of very specific\nlayouts (not just plain row-major or col-major; and what specifically\nare these layouts expected by the mma is not too important for now, discussed later).\n\nBasically tcgen05 expects data organized in one of specific layouts (ie how the data is laid out in memory). And additionally, tcgen05 does not support using GMEM operands (because GMEM bandwidth is too low, operands need to be written to fast on chip memory first).\n\nThe takeaway from this section is, we need to lay out these A and B operands in SMEM, in some certain layout supported by tcgen05.\n\n**Back to the main flow (and out of the onramp):**\n\nAs we saw in the first two diagrams, and the pseudocode, in each CTA, Q stays fixed, while each iteration of the K/V loop computes:\n\n- A)\n`Q @ K.T -> S`\n\n- B)\n`S -> P`\n\n- C)\n`P @ V -> O`\n\nThese are the 3 main computations that each iteration of the KV loop (along the SEQ_LEN) does. Let’s zoom into a single iteration of the K/V loop, and cover these 3 steps one by one.\n\n### A) Q@K.T (explaining K-major)\n\n*Inside one K/V-loop iteration — Step 1 of 3.*\n\nBefore running QK, we need Q and the current K tile in SMEM. Q is loaded once for the CTA, while a new K tile is loaded on every K/V iteration. Both Q and K use the same layout, so below I draw only the Q copy and omit the K copy.\n\nNote: In the QK matmul diagrams below, I name dims using conventional matmul dim names: Q[M,K] @ K[N,K].T = S[M,N]. These are local matmul-axis names, not axis names of our kernel’s input global tensors.\n\nAs mentioned earlier, Tensor Cores on Blackwell represented by the `tcgen05`\n\nfamily of instructions.\nAnd `tcgen05.mma`\n\nexpects its operands be organized in a particular layout\n(we cannot just copy our tiles from global to shared in any arbitrary layout and run MMA on that).\nGiven that constraint, we must “lay out” that data in one of these layouts MMA supports.\n\nSo how specifically we should lay out the data in SMEM, so that it’s one of the layouts supported by MMA?\nFor simplicity, let’s only describe the process of copying Q tile from GMEM to SMEM (the same logic applies separately to the K tile).\nLogically we’re splitting that Q tile vertically (as shown in the picture – here pictorially 4 slices total).\nAnd then copy each slice individually from GMEM to SMEM with TMA 2 (see the diagram above).\n\nWhy do we need to slice the tile this way instead of copying the whole tile at once?Good question, hold this thought for a moment. I don’t want to introduce too many things at once, so let’s return to this later when it matters. For now just remember we’re copying in slices.\n\nAfter copying all slices, it seems nothing changed going from GMEM to SMEM\n(as visualised, we see the same logical `[BLOCK_M, HEAD_DIM]`\n\nmatrix but now in SMEM instead of GMEM).\nBut in fact physically, the representation *did change*. See “Physical view” green annotation in the picture above.\nSo the resulting data in SMEM is no longer the row-major layout we had in GMEM.\nIn the physical SMEM layout we now have **all rows of one slice**,\nonly then all rows of **another slice** and so on.\n\nSo this is row-major inside each slice, but overall\n(the layout of the entire `[BLOCK_M, HEAD_DIM]`\n\nin SMEM)\nis not merely a simple row-major or column-major.\nRather, this is one of the so called *“canonical layouts”* that `tcgen05.mma`\n\nsupports.\nRemember this term, we will see it often.\n\nAnother visualisation (above) trying to explain differences in the physical representation between\nthe same tile of data (Q tile) in GMEM and in SMEM.\nPink arrows denote physical order of elements.\nGMEM had regular row-major, SMEM has a canonical layout\nwhich tcgen05 knows how to work with (this specific\nlayout drawn is so called **K-major no-swizzle** – we’ll return to this later).\n\n“K” in K-major does not refer to the Key tensor. “K” in K-major means the reduction dimension (in the canonical matmul naming convention), in our case visualized running horizontally across the logical tile. Each vertical slice above covers a small part of that K dimension.\n\nFor now can ignore this part, and for simplicity just think of this layout as the only layout tcgen05 supports. Again, I don’t want to introduce too many things at once. We’ll discuss other layouts later when it matters. That’s just to emphasise how the physical organization of the data changed as a result of us copying vertical slices with TMA.\n\nWhy physical layout should be this way but not some other arbitrary way?This does seem unintuitive, but I don’t belive there’s a deeper intuition than to say this is one of the canonical layouts tcgen05 supports (according to the nvidia docs). Some things are ultimately a hardware contract we have to fallback on as ground truth.\n\nAbove we’ve covered the TMA part: copying data from GMEM to SMEM. Now let’s also discuss the MMA part: how can we actually now run MMA on that SMEM data we’ve prepared.\n\nWe used TMA to copy data from Global memory into SMEM, but it’s just a copying mechanism and is not inherently aware about the downstream tcgen05.MMA we’re about to call. And separately, tcgen05.MMA knows only how to interpret a small set of specific layouts, and it doesn’t know how the data it’s about to read got into SMEM in the first place. So, we can think of it as: load part (discussed earlier), and MMA part (discussed below). These are two independent steps and are not aware about each other, so it’s our job to make sure they agree on the same layout.\n\nDiagram simplification: diagram uses BLOCK_K = HEAD_DIM = 32 BF16 elements (64B) to keep it readable. The code and the text below use BLOCK_K = HEAD_DIM = 128 BF16 elements (256B). Everything explained below works the same, the actual code just contains 4 times as many slices.\n\nTo recap, we earlier copied both of our Q_tile and, separately K_tile into SMEM (in 16B slices along K).\nNow we need to take this `(block_m, block_k)`\n\nQ tile, and matmul it with the `(block_n, block_k).T`\n\nK tile.\nFor now let’s only focus on the Q tile (1st input to the tcgen05.mma), while for now omitting K tile (2nd input to the tcgen05.mma).\n\nOur tiles have `BLOCK_K = 128`\n\nelements (which is 256 bytes, given our bf16 dtype).\nBut one `tcgen05.mma`\n\ninstruction consumes `MMA_K=16`\n\nBF16 elements (32B) along K.\nWe’ll return to where this `32B`\n\ncame from in a moment.\nSo, the mma_K width is smaller than our tiles K width, thus we need a for loop which walks over the BLOCK_K dim of our Q and K tiles\nin steps of 32 bytes (software loop over the data in smem, shown in yellow above).\n\nEach of our Q slices is 16B wide along K, while one tcgen05.mma consumes 32B along K. So two neighboring Q slices form one Q micro-tile. In the diagram above, each same-colored pair of slices forms one micro-tile.\n\nAnd when I say “micro-tile” I mean a `BLOCK_M=128, MMA_K=16`\n\nchunk which tcgen05 consumes directly.\nA “micro-step” is one tcgen05.mma call consuming one Q micro-tile and one corresponding K micro-tile.\nSince BLOCK_K=128 and MMA_K=16, 128 / 16 = 8 micro-steps cover the full reduction-K dimension.\nThe simplified diagram shows two of those eight steps.\n\nThese micro-tile shapes I used above come from PTX docs [Table 39](https://docs.nvidia.com/cuda/archive/13.0.0/parallel-thread-execution/#tcgen05-kind-shapes),\nwhich defines MMA shapes as `MxNxK`\n\n, with A shape `M,K`\n\n, B shape `K,N`\n\n, D shape `M,N`\n\n:\n\n```\n// The below numbers mean num elemnts in bf16, not bytes.\nFor .kind::f16, no .ws, cta_group = 1, dense, it lists supported shapes:\n      - 64xNxK\n      - 128xNxK\n      - N = {8, 16, 24, ... 256}\n      - K = 16\n```\n\nTo try to visualise this further, I’ve drawn K_tile as well. Which was independently split on 16B wide slices and copied with TMA, then each pair of its slices is treated as a micro-tile (same as for the Q tile we covered earlier, so I omitted showing copying the K_tile from the diagrams).\n\nThen, to matmul the Q_tile and K_tile tougether, we basically need to matmul their corresponding micro-tiles (shown in the red box in the diagram above):\n\n- 1st Q micro-tile @ 1st K micro-tile.T\n- 2nd Q micro-tile @ 2nd K micro-tile.T\n\nThe `output_accum`\n\nshown in the picture above is our Scores tile (S).\nOn Blackwell, `tcgen05.mma`\n\naccumulates its output in Tensor Memory (TMEM),\na dedicated on-chip memory separate from SMEM and registers.\n(on earlier GPUs, tensor-core accumulators lived in registers instead).\nAll QK micro-steps accumulate into the same S tile in TMEM.\nFor now, just keep in mind that MMA outputs live in TMEM. We’ll explain how\nTMEM is allocated and accessed in Section B, when softmax needs to read S.\n\nWe turn these straight-line MMA calls into a loop which iterates over BLOCK_K dim of Q_tile and K_tile using MMA_K=16 micro-steps. In pseudocode:\n\nThis is our “yellow” loop from the actual kernel code, slightly simplified (executed by a single thread):\n\n```\n// tcgen05 mma_k width along K-dim\nint MMA_K_BYTES = 32; // in bytes\nint MMA_K = 16; // in BF16 values\n\n// for now, think of each K-major SMEM \"descriptor\" as a\n// pointer-like value selecting current micro-tile\n// (it also contains some layout metadata, covered later)\nuint64_t q_desc = noswizzle::desc_kmajor(Q_smem, BLOCK_M);\nuint64_t k_desc = noswizzle::desc_kmajor(K_smem, BLOCK_N);\n\n// offsets to advance the descriptors to next Q and K micro-tiles\n// (minor detail, divide by 16 because descriptor addresses are encoded in 16B units)\nuint64_t q_desc_step = (BLOCK_M * MMA_K_BYTES) / 16;\nuint64_t k_desc_step = (BLOCK_N * MMA_K_BYTES) / 16;\n\n// BLOCK_K= HEAD_DIM = 128, MMA_K=16, so 8 micro-steps\nfor (int k = 0; k < HEAD_DIM / MMA_K; ++k) {\n\n    // one micro-step\n    tcgen05::mma_f16(\n        taddr_s, q_desc, k_desc,\n        /*MMA shape and dtypes=*/i_desc_qk,\n        /*enable_input_d=*/(k == 0 ? 0 : 1));\n\n    // advance descriptors to the next pair of micro-tiles\n    q_desc += q_desc_step;\n    k_desc += k_desc_step;\n}\n\n// make mbarrier track completion of all 8 micro-steps\ntcgen05::commit_arrive(mbar_addr);\n```\n\nEach iteration calls tcgen05.mma on the next MMA_K=16 chunk and accumulates into the same S tile.\nFirst MMA micro-step initializes S, remaining seven MMA micro-steps accumulate into that same S tile.\n**After the loop over all micro-tiles, that S_tile contains the result of Q_tile @ K_tile.T.**\n\n**No explicit K_tile transpose is needed.** Both Q and K use K-major layouts.\nThe MMA A operand reads Q[m,k] directly. For the B operand, tcgen05.mma interprets\nthe stored K tile as the [k,n] operand required for Q @ K.T.\nK-major itself does not mean “transpose this operand”: A and B are different MMA operand slots with different logical shapes, [M,K] and [K,N], and are therefore interpreted differently by MMA.\n\nAt this point, you already know enough to follow the chapter-level dataflow:\neight MMA micro-steps accumulate partial products into S, collectively producing\n`S = Q @ K.T`\n\n.\n\nTurns out, there’s one more level of granularity here.\n**Let’s zoom on a Q micro-tile** (BLOCK_M=128, MMA_K=16), though the same logic applies to K micro-tile as well.\n\nThese micro-tiles themselves consist of lower level units. Called atoms\n(or “swizzle atoms” as per the cuda docs, note we are NOT doing any swizzling yet,\nNVIDIA uses “swizzle layout atom” as the umbrella term even when the swizzling mode is None\n([table 53](https://docs.nvidia.com/cuda/archive/13.0.0/parallel-thread-execution/#tcgen05-smem-swizzle-mode))).\n\nThe size of these atoms depends on the specific canonical layout we’re using, and in our case (the layout we’re using is called K-major no swizzle, but for now we can ignore these details and assume there’s only a single layout that tcgn05.mma supports), these atoms are (8, 8). Below, I’ll show where this shape comes from.\n\nThe top-down story so far (see legend in the diagram above): a micro-tile consists of 2 slices, each of which consists of atoms. We need to understand how atoms are arranged in SMEM. But they are not independently programmable MMA units. When issuing MMA we’re still operating on the level of micro-tiles, not individual atoms.\n\nNow we can answer the question from earlier:\n\nwhy did we copy Q and K with TMA in 8-element-wide vertical slices instead of copying the whole tiles at once?The slice width we used earlier wasn’t arbitrary: it was actually one atom wide. For Q, each (128,8) slice we used, is a vertical stack of (8,8) atoms. So we split the tile into these slices and use TMA to copy one complete stack at a time, placing the stacks one after another in SMEM. This produces the K-major physical atom order that tcgen05.mma supports.\n\nA small, optional, derivation of where the (8,8) atoms shape comes from:\nPTX defines atom shapes in 128-bit units, not in BF16 elements.\nFor K-major with no swizzling, [PTX Table 53](https://docs.nvidia.com/cuda/archive/13.0.0/parallel-thread-execution/#tcgen05-smem-swizzle-mode)\ndefines one atom as `8×1`\n\n(eight rows with one 128-bit chunk per row).\nSince one 128-bit chunk contains 8 BF16 values, the atom shape is (8,8) BF16 elements.\n\nRemeber as discussed above, at each iteration of the “yellow” loop we feed to tcgen05.mma two micro-tiles of data (one Q micro-tile and one K-micro tile).\n\nFor a single micro-tile we don’t need to provide to mma the addreses of all its atoms. Instead, we need to provide the base pointer to the first atom, highlighted in the picture above.\n\nAs the yellow loop progresses, we advance the base pointer from one micro-tile to the next. In the first iteration, it points to cyan atom, the start of micro-tile 0. In the second iteration, it points to purple atom, the start of micro-tile 1, and so on across the K dimension. The same happens for the K operand descriptor.\n\nBut becuase the base points only to the first atom,\nintuitively, in addition to that base address, *MMA somehow needs to\nknow how to step from the first atom in a given\nmicro-tile to other atoms in the same micro-tile*.\n\nTo that end, there’s a concept of SBO and LBO (shown in the figure above).\nThese do not create extra software loop iterations, they describe where the other atoms in the same micro-tile are relative to its base pointer.\nSBO tells MMA how to step to the next atom **in the same slice**,\nand LBO tells MMA how to get to the atom (at the same index) but **in the next slice**.\n\nIn our case, for the SMEM memory layout we’re discussing (k-major no swizzle), each atom, and therefore each vertical slice, is 8 BF16 elements (16B) wide, so:\n\n```\nSBO = 8 rows * 16B per row = 128B   // step over one atom\nLBO = BLOCK_M rows * 16B per row = 2048B    // step over one complete slice of atoms\n```\n\nPutting it together, our software loop advances the base pointer from one micro-tile to the next, while SBO and LBO describe the offsets to the other atoms inside the current micro-tile.\n\n### B) Online Softmax (explaining S->P)\n\n*Inside one K/V-loop iteration — Step 2 of 3.*\n\nAt this point we finished going through the first matmul (Q@K.T), and now we’re appraoching the online softmax part of our kernel (S->P).\n\nThe online-softmax math is standard and not specific to B200, and there are plenty of explanations elsewhere, so I won’t rederive it here. I’ll only summarize the state update we need, then focus on how our B200 kernel implements it.\n\nWe already briefly saw TMEM in Section A, our QK micro-steps accumulated S there. Now let’s zoom in on TMEM itself and see how online softmax reads S and converts it into P.\n\nTMEM is (128 rows, 512 columns), and it must be explicitly allocated (allocation is in columns, each column contains 128 elements). Kernel 1 allocates 128 columns for S and another 128 for O.\n\nS is rewritten at each iteration of the K/V loop. The running output accumulator O stays in TMEM across the entire KV loop. For a single CTA: every K/V iteration updates the same O tile, after the entire loop, that O eventually becomes the final output tile of this CTA.\n\nNow we need to compute online softmax update, for one K/V iteration, the state changes like so:\n\n```\n// Section A, covered earlier\n// S = Q @ K.T;\n\n// Section B:\n\ntile_rowmax = row_max(S);\n// new basis\nnew_rowmax  = max(running_rowmax, tile_rowmax);\n\n// rescale old state to the new basis\nrescale     = exp(running_rowmax - new_rowmax);\nrunning_rowsum *= rescale;\nO              *= rescale;\n\n// produce P and update denom\nP               = exp(S - new_rowmax);\nrunning_rowsum  += row_sum(P);\nrunning_rowmax   = new_rowmax;\n\n// Section C, covered later:\n// O += P @ V;\n// after complete K/V loop:\n// output = O / running_rowsum;\n```\n\nThis is generic online softmax update, not B200 specific so I will not dwell on it. B200 question is how the kernel reads S from TMEM and produces P.\n\nOn high level, our goal in this section is to turn S into unnormalized 3 P,\nto be later used in the PV matmul.\n\nTo do that, Kernel 1 traverses each row of S twice:\n\n- First, to compute rowmaxes (to be used for numerical stability in the 2nd step).\n- Second, to produce unnormalized P, by subtracting rowmaxes and exponentiating the result\n\nEach pass over TMEM is explained in the diagram below.\n\nFor the first pass, we cannot form P until we know the maximum over all 128 scores in this S-tile row.\nNote in the 2nd pass we re-read from TMEM the same `S`\n\nwhich we already read before in the 1st pass.\nThe reason for re-reading is our 4 warps read S effectively 128x8 chunks at a time (these chunks were stored in regs).\nEach time we loaded new chunk we overwrote the previous chunk stored at the same registers.\nThe complete 128x128 S was never stored in registers, so if later we need one of earlier 128x8 chunks, we have to reread them from TMEM again.\n\nLet’s discuss more concretely how do we traverse TMEM, and read of the S values.\n\nOnly `tcgen05.`\n\nfamily of instructions can access TMEM, regular load/store instructions can’t.\nSo to compute the per-row max values, we first need to load values from TMEM (before computing their max),\nso I’m using `tcgen05.ld`\n\n(stands for load) insturction to load values from TMEM into registers.\n\nLet’s read this picture from left to right.\n\nFrom the CTA perspective, the four warps split the 128 rows of S between them.\nEach warp owns one `32, 128`\n\nrow-chunk. That is why our Kernel 1 has 4 warps\n(to cover all 128 rows of S in 32-row chunks, so all 128 rows can be processed in parallel).\n\nZooming into one warp: on each loop iteration, the warp collectively loads one\n`(32, 8)`\n\nchunk from TMEM (`tcgen05.ld`\n\nis warp collective).\nEach of its 32 lanes gets eight values from one row.\nAcross all four warps, one iteration therefore processes a `(128, 8)`\n\nchunk of S.\n\nThen the loop moves eight columns to the right and repeats. After 16 iterations, the four warps have traversed all 128 columns of S.\n\nFrom one thread’s perspective, it keeps following the same row and receives eight values at a time.\n\nAs mentioned, overall we traverse S twice. From the perspective of TMEM traversal, both look nearly identical. The differences in the traversals happen only after we read a chunk of S into regs (lower right corner of the figure).\n\nIn the first S pass, it reduces those values into its `tile_rowmax`\n\n(as shown in the diagram).\nThe second S pass repeats the same TMEM traversal, but this time subtracts the\nupdated running rowmax, exponentiates the scores to form P, and writes P into SMEM.\nThe SMEM write is shown in the next figure.\n\nBetween the two S passes, there is also a separate traversal over O in TMEM.\nOnce the new running rowmax is known, the same four warps load O in\n8-column chunks, multiply it by `rescale`\n\n, and write the corrected O back\nto TMEM before PV.\n\nSo Kernel 1 reads S twice, and separately performs one load-rescale-store traversal over O.\n\n```\nSOURCE WALKTHROUGHExplanation above is enough to follow the rest of the blog, so you can continue with the next figure.If you want to go slower and deeper, see\nHow Kernel 1 Implements Online Softmax with TMEM\nin the Appendix. There I map this picture directly on\n1_baseline.cu source code.\nWe start from the first S pass, follow the O correction, and then reach the\nsecond S pass which produces P.\n```\n\nWe then write P (in chunks, because as explained earlier S is read and processed in chunks, thus P is produced in chunks accordingly) into one of the canonical layouts (K-major). That layout is needed so later PV MMA can understand how to read and interpret our data (ie, so that the P micro-tiles are consumable by the later PV MMA).\n\nS itself stays unchanged in TMEM; each S chunk we read produces a corresponding P chunk in SMEM.\n\nWe store P in the same SMEM buffer that previously held K. K is already dead after QK, so it is safe to reuse that storage for P.\n\n### C) P@V (explaining MN-major)\n\n*Inside one K/V-loop iteration — Step 3 of 3.*\n\nAbove we covered how P tile is produced, now we need to matmul it with V tile. Turns out, unlike for the Q, K, and P (which all use the same K-major SMEM layout, as dicussed above), V uses a different SMEM layout. So that’s our next topic of discussion.\n\nFor our Kernel 1, `P_tile`\n\nis produced into SMEM by softmax, so it is not shown in the global-memory view below.\n\nNote: In the PV matmul diagrams below, I name dims using conventional matmul dim names:\n\nP[M,K] @ V[K,N] = O[M,N]\n\nQ, K, and V start in the same global-memory layout, but in the kernel, we lay them out differently in SMEM depending on how each MMA will consume them.\n\n```\nQK: Q[M,K] @ K[N,K].T\nPV: P[M,K] @ V[K,N]\n```\n\nJust let me step back for a moment and explain why I name these axes M, N, and K.\nFor PV, `M = BLOCK_M`\n\n, `K = BLOCK_N`\n\n, and `N = HEAD_DIM`\n\n.\nI use conventional M/N/K names separately for each matmul, because for each matmul standalone,\nit makes its input shapes easier to read. So these M/N/K are a matmul-centric naming choice,\nnot inherent properties of global tensor (Q, K, P, or V).\nThat is why the same source dimension have a different local name in QK and PV.\nSo, K-dim means the reduction dim of the current matmul: `HEAD_DIM`\n\nfor QK,\nbut the K/V tile rows (`BLOCK_N`\n\n) for PV. Now back to our larger, PV matmul discussion.\n\nWhy do I use k-major layout for Q tile, K tile, and P tile, but use a different layout (discussed below) for the V tile? Becuase QK needs K viewed transposed, while for PV the inner dimensions already match, so V does not need a transpose. More concretely:\n\nFor QK, K is stored as `[N,K]`\n\n, MMA consumes its B operand as logical `[K,N]`\n\n.\nSo I use K-major for K. This lets MMA interpret the stored K tile as the logical\ntransposed operand, without first constructing a separate transposed K tile.\n\nFor PV, V is already in the [K,N] orientation needed by the MMA B operand, so I use MN-major for it.\nThis lets MMA consume V directly, without first transposing V tile.\nP is the `[M,K]`\n\nA operand, so it remains K-major.\n\nIn other words, for the stored tile views used in these diagrams, my shortcut is:\n\n- K-major: reduction K is on the right, as in Q[M,K], K[N,K], or P[M,K].\n- MN-major B: reduction K is on the left, as in V[K,N].\n\nIt’s only a shortcut I use for reading these operand views, not formal definition of K-major and MN-major layouts.\n\nFor now, this is enough to understand the MN-major V copy. In Chapter 3, we’ll compare it with K-major and look more closely at why we place V atoms one at a time.\n\nAt this point we covered TMA part for the PV matmul. We have P produced in SMEM (by our previous section B), and we’ve just discussed how to separately copy V into SMEM. At this point we have both operands for the P@V ready, so let’s start discussing that 2nd matmul itself.\n\nThese are the two operands of the same PV MMA. tcgen05.mma lets us choose the major mode of A and B independently, so P can be K-major and V can be MN-major.\n\nSame basic idea as for QK (explained in earlier section A):\nour reduction dim labeled BLOCK_K in this diagram (BLOCK_N in the source)\nhas 128 BF16 elements,\nbut one `tcgen05.mma`\n\nissue consumes only MMA_K=16 BF16 elements.\nSo PV is split into `128 / 16 = 8`\n\nMMA microsteps.\nFor simplicity, I only show 2 of these microsteps.\n\nIn this V view, reduction K is vertical. So one microstep consumes two neighboring 8-row slices.\n\nAlong BLOCK_K, each microstep consumes 16 BF16 elements. Since each BF16 element is 2 bytes, that is the 32 bytes shown in the diagram.\n\nFor each PV microstep, the V operand descriptor base points at the first `(8,8)`\n\nV atom\nin that 16-row chunk.\n\nFor microstep 0, the V descriptor base is the green highlighted atom. Our loop over micro-tiles then advances descriptor to the next 16-row chunk. For microstep 1, the V descriptor base becomes the purple highlighted atom.\n\nThe next figure shows how LBO and SBO cover the remaining atoms in the same micro-tile from that descriptor base.\n\nFor each PV microstep, descriptor base points at the first atom in the first 8 row V slice.\n\nLBO reaches the corresponding atom in the second 8 row slice, while SBO steps across the 8 column atoms along N.\n\nUsing that starting address, the two offsets, and the MMA shape, one\n`tcgen05.mma`\n\nissue covers complete 16row V micro-tile.\n\nThen we advance both P and V descriptors to the next 16-wide reduction chunk.\n\n**Zooming back out: one K/V iteration**\n\nIn Chapter 1’s Sections A, B, and C, we zoomed in the individual steps inside one K/V-loop iteration. Now let’s zoom back on the level of a complete KV-loop iteration.\n\nOne thing the diagrams above omit is synchronization.\n\nTMA loads and `tcgen05.mma`\n\nare asynchronous.\nBut in Kernel 1, we immediately wait on every major handoff:\nload the current K/V tile and wait,\nrun QK and wait,\ndo the row work,\nthen run PV and wait.\nOnly after PV finishes we start loading the next K/V tile.\n\nSo this baseline still runs one semantic stage after another. Only one K/V tile is live, and there is no useful overlap between the K/V load, QK, row work, and PV yet. Later kernels start overlapping these stages.\n\nIn the source, mbarriers tell us when the asynchronous TMA and MMA work has completed.\n`__syncthreads()`\n\nonly waits until all CTA threads reach that point;\nit does not wait for the asynchronous TMA or MMA work to finish.\n\n```\nKERNEL CHECKPOINTThis blog is not diagrams only, the code is equally important.See 1_baseline.cuIt follows the same exact sequence as we studied:each CTA holds a single Q tile; loops over all K, V tilesS=Q@K.T (tcgen05 in microtiles, K-major layout)S->P (accessing TMEM)O+=P@V (tcgen05 in microtiles, MN-major layout)Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\nThis table will grow by one row after every chapter. Stock FA4 (cute DSL) is 100%. Each value is the median of six same-run ratios; every benchmark invocation measured this kernel and stock FA4 on the same B200.\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n1 | Baseline | 14.2% | 14.1% | 13.7% |\n\n⭐ [Star the repo](https://github.com/IaroslavElistratov/b200-attention) so you have all the kernels at hand as we go through the later chapters.\n\nMore ML systems articles, videos, and code are coming:\n[ LinkedIn](https://www.linkedin.com/in/iaroslav-elistratov/) ·\n\n[·](https://x.com/iaro_e)\n\n**X**\n\n**YouTube**## Chapter 2 – move P from SMEM to TMEM\n\nIn kernel 1, the second softmax pass reads S from TMEM, converts it into the unnormalized probabilities P, and writes P into the SMEM buffer that previously held K (K buffer is dead after QK, so we safely re-used it for P).\n\nKernel 2 changes where P lives. Instead of writing P fragments to SMEM, we write them back into the TMEM. That’s the main change introduced by kernel 2.\n\nThis change doesn’t reduce SMEM usage, but it prepares for later optimizations. By storing P into K’s dead SMEM slot Kernel 1 coupled their lifetimes: K is no longer needed after QK, but its slot remains used by P until PV consumes it, which is suboptimal. We could give P a separate SMEM buffer, but we’ll need almost entire SM’s SMEM budget for our future optimizations. Instead, kernel moves P to TMEM, separating their lifetimes without allocating more SMEM. In later chapters, this lets us use that decoupled P and K lifetime for additional optimizations.\n\nThe allocated SMEM does not reduce, kernel still needs Q,K,V storage.\n\nAs an optimization we write P directly into the allocation that held S so P needs no additional TMEM allocation. (because we can safely overwrite the elements of S which the second pass over TMEM already consumed).\n\nSo, in kernel 2, the same TMEM slot changes meaning over time:\n\n- after QK: the slot holds FP32 S\n- softmax pass 2: consumed S chunks are overwritten by packed BF16 P\n- PV: the slot is consumed as P\n\nThis does create the coupling of S and P lifetimes, but as you will see later\nis not as problematic as the coupling of P and K lifetimes.[4](#fn:4)\n\nTMEM is organized as 32-bit memory cells yet our P values are bf16, so 2 of our P values can be packed into one TMEM cell. Likewise, S contains FP32 values, so one score (one value in the S matrix) occupies one b32 TMEM cell. So, P uses half of TMEM columns as S.\n\nThe figure shows one logical 16-score step, shown as differently colored arrows.\nIn code, that step is implemented using two `tcgen05.ld.x8`\n\nloads (shown as orange arrows),\nfollowed by one `tcgen05.st.x8`\n\nstore (shown as green arrows).\nKernel 1 processes P in width-8 chunks. Kernel 2 groups two of these chunks into one iteration,\nper row, loading 16 FP32 S values and producing 16 BF16 P values.\nThe width-16 grouping is not fundamental.[5](#fn:5)\n\nRemember from kernel 1, when producing P we iterate over TMEM in chunks (see the relevant diagram from the previous chapter),\nthis is very similar to chapter 1.\nWhat’s new in ch-2, is we store the resulting P chunk in TMEM (by packing 16 P values into 8 b32 TMEM cells).\nSo S-read pointer advances by 16 columns, while the packed-P-write pointer advances by only eight.\nThus, each P write remains behind the read pointer and overwrites only scores already consumed by this second TMEM pass.\nSo when producing P and writing it into TMEM (in chunks) we are not at the risk of clobbering unread S values,\nas also shown on the diagram.[6](#fn:6)\n\nOur S takes up 128 TMEM columns, and as mentioned we read in 16-column chunks, so we need 128/16=8 iterations to cover all TMEM S columns. After all eight steps, the lower 64 b32 columns contain complete 128-column BF16 P tile. So P doesn’t overwrite the entire S buffer, but only half of it. The remaining half contains stale S values that are never read again. The next QK (in the next iteration of the KV for-loop) fully overwrites the entire S/P slot with S(i+1).\n\nPV matmul doesn’t change:\n`P[BLOCK_M, BLOCK_N] @ V[BLOCK_N, HEAD_DIM] -> O[BLOCK_M, HEAD_DIM]`\n\nOnly the source of P (the A operand) changes. In kernel 1 both P and V were consumed from SMEM. In kernel 2, P is consumed from TMEM, V from SMEM.\n\nIn kernel 1, P needed a K-major SMEM layout and SMEM descriptor. In kernel 2, P is identified directly by a TMEM address, and we no longer use the SMEM descriptor for P.\n\nK-major and MN-major are SMEM-only-layouts and are not used for TMEM operands. V remains the MN-major SMEM B operand.\n\nAt this point we’ve discussed how P is produced and written to TMEM (as oppose to written in SMEM, as in Chapter 1). Let’s now discuss the P@V matmul again, but now one of the operands (P) will come from TMEM.\n\nKeep in mind all my MMA-related dims in the diagrams use conventional local matmul notation: A[M,K] @ B[K,N] = C[M,N]\n\nSo, PV reduces over `BLOCK_K = 128`\n\n.\nSame as in chapter 1, each MMA issue consumes matching `MMA_K=16`\n\nmicro-tiles of P and V:\n`P[:, 16k : 16(k+1)] @ V[16k : 16(k+1), :]`\n\n.\nEight `MMA_K=16`\n\nissues cover the complete `BLOCK_K=128`\n\nreduction.\nTo avoid clutter, the figure shows only the first two.\n\nThe loop follows the same logical MMA_K=16 micro-tile order as kernel 1, but selects P through TMEM addresses instead of advancing an\nSMEM descriptor. `taddr_p + k * 8`\n\npoints to the base of the next 128x16 P micro-tile.\nWithin each row, its 16 BF16 values take eight b32 TMEM columns.\nThe V descriptor advances to the matching 16 rows, and every issue accumulates into the same O tile.\n\nBecause P is now produced with `tcgen05.st`\n\n, the row threads must wait\nfor those TMEM stores and fence before PV begins.\nThis is just a minor change to the publication mechanism.\n\nSo kernel 2 leaves the softmax math, PV reduction order, and high level schedule unchanged. Its only semantic change is: QK produces S in TMEM, softmax turns that storage into P, PV consumes P directly from TMEM.\n\n```\nKERNEL CHECKPOINTSee 2_p_to_tmem.cu and diff it against Kernel 1.It follows the same exact sequence as we studied:softmax overwrites consumed S with packed BF16 PPV reads P from TMEM and V from MN-major SMEMP doesn’t reuse K’s SMEM slotthe math and schedule unchangedNow, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n\nThis didn’t lead to a performance improvement, but these changes did uncouple K and P lifetimes, which will enable our later major optimizations.\n\nI kept this as a separate step instead of folding it into Kernel 1 because writing P to TMEM introduces additional addressing, packing, and synchronization details. Kernel 1 already introduces enough new concepts, so I didn’t want to complicate it unnecessarily, and tried to keep the initial baseline simpler.\n\n## Chapter 3 – Swizzling\n\nWe already chose Q/K as K-major and V as MN-major based on how each source tile is consumed by its MMA (we covered this in chapter 1, revisit if you need a refresher). Kernel 3 keeps the same schedule and major modes, and only changes the swizzle mode (of each of the tiles independently): Q/K/V moves from no-swizzle to SW128 in SMEM, to reduce SMEM bank conflicts. P stays in TMEM, so swizzling is not applied to it.\n\nThere’s my resources explaining the motivation behind swizzling, I will not repeat it here, because our goals is understanding B200 related concepts, suffice to say swizzling permutes SMEM addresses within each atom so tensor core accesses map better across SMEM banks and avoid bank conflicts. I cover B200 specific details below (but not the generic swizzling expaliner).\n\nMajor mode and swizzling are two separate things (dicussed in chapter 1, revisit if you need a refressher). Major mode says which matmul dim is packed inside the 16B elements: K for K-major, or M/N for MN-major. Swizzling adds address permutation inside fixed-size atoms of that layout. So we first choose the major mode based on the operand’s role in the matmul. Once we also decided to use swizzling (SW128), the required atom shapes follow from the tcgen05 layout table below.\n\nLet’s translate that table into BF16 elements, which is the dtype of our input tensors.\n\nWhen we added swizzling, two points worth covering: (i) atoms got wider, and (ii) became swizzled inside. These two points are adresed below in order.\n\n**(i) atoms get wider**\n\nNVIDIA describes the [tcgen05 shared-memory layout table](https://docs.nvidia.com/cuda/archive/13.0.0/parallel-thread-execution/#tcgen05-smem-swizzle-mode) in terms of 128-bit elements (16 bytes).\n, not BF16 values.\nOne BF16 value is 16 bits, so one 128-bit element contains `128/16 = 8 BF16 values`\n\n.\n\nThe docs say to expand only the **leading dim** in the atoms.\nIn the table, atom dims are ordered as `(M/N, K)`\n\n, though this wasn’t immediately obvious 7.\n\nSo for K-major, K is the leading dim, which is 2nd axis:\n\n``` php\n(8,8) -> (8, 8*8) = (8,64) BF16\n```\n\nFor Q/K, we already draw source tiles in `[M/N,K]`\n\ncoordinates,\nso their K-major atoms appear as `(8,64)`\n\n.\n\nFor MN-major, M/N is the leading dim, which is the 1st axis:\n\n``` php\n(8,8) -> (8*8, 8) = (64,8) BF16\n```\n\nFor V, the MN-major atom is `(64,8)`\n\nin NVIDIA’s `[N,K]`\n\ncoordinates.\nIn the source `V[K,N]`\n\nlogical views I draw, those axes are swapped, making it `(8,64)`\n\n.\nThat’s why both atom types look `(8,64)`\n\nin my source-coordinate diagrams,\neven though they are not the same layout.\n\nFor now, let’s first zoom into the K-major layout (used Q@K). We’ll later return to the MN-major layout (used for V in P@V).\n\n**(ii) swizzled inside**\n\nAs far as I understand, with swizzling added, elements inside a 16B element are not permuted, 16B elements themselves are permuted.\n\nWe don’t care about exact row-by-row permutation, because we do not compute it manually. Instead, on blackwell we rely on a) TMA writes in swizzle-128B layout, and b) tcgen05 reads expecting swizzle-128B layout.\n\nSMEM layout and the downstream MMA have to agree on layout, as we know from Chapter 1.\nConcretely, we set `CU_TENSOR_MAP_SWIZZLE_128B`\n\nin TMA tensor maps, so TMA writes Q/K/V swizzled.\nAnd then set the matching SW128 mode in the tcgen05 operand descriptors, so tcgen05 knows to consume the swizzled layout.\n\nLet’s zoom into K-major layout first (used for both operands of Q@K.T).\n\nWhen copying into SMEM we want to lay out our data according to (8,64) atoms as explained above.\nOne K-major SW128 atom is 64 BF16 elements wide (ie 128B).\nBut our full Q/K tile is `HEAD_DIM=128`\n\nelements wide (ie 256B).\n\nWith a simple 2D SW128 TMA map that we use, we can’t copy the full 128-wide tile at once and ask TMA to treat it as two separately swizzled 128B regions. CUDA requires the inner dim of one SW128 TMA box to be at most 128B. In our Q/K tensor maps, that inner dim corresponds to HEAD_DIM (the horizontal dim in the figure). So we split HEAD_DIM into two 64-wide slices.\n\nBut it doesn’t mean we need one TMA issue per `(8,64)`\n\natom.\nThe 128B restriction is only on the inner width.\nOnce the width is 64, the same TMA box can extend down all rows of the tile:\n\n- Q copy:\n`(BLOCK_M, 64)`\n\n- K copy:\n`(BLOCK_N, 64)`\n\nIn this kernel, both tile heights are 128.\nSo one TMA issue copies one `(128,64)`\n\nslice,\ncontaining 16 vertically stacked `(8,64)`\n\natoms.\n\nAbove we discussed how we tell TMA to use swizzling when copying our data to K-major swizzled SMEM layout. But there’s the second part, which is: telling MMA how to consume the K-major siwzzled layout.\n\nNow let’s zoom into what one of these 64-wide slices consists of.\n\nOne K-major SW128 atom is (8,64) BF16 (1024B).\nOne 64-wide Q slice consists of these atoms stacked along `BLOCK_M`\n\n.\n\nMinor note, here I draw BLOCK_M=32 only to keep it readable; the real kernel has 16 atoms in each slice.\n\nOne MMA issue still operates on width MMA_K=32B (16 BF16 elements). But our BLOCK_K is larger than MMA_K, so (similarly as in no-swizzle case explained in ch1) we gonna need a for-loop which steps over BLOCK_K in MMA_K=16 element steps.\n\nThe difference here (compared to the non-swizzle case discussed earlier) is that logically,\nnow we gonna be splitting individual atoms, as each atom is of width=64, but MMA_K is only 16.\n`128B atom width / 32B mma width = 4 micro-tiles`\n\n.\n\nSo whereas in the non-swizzle case we were taking 2 columns of atoms and calling it a micro tile, here our micro-tiles are result of splitting (and not combining) the atoms.\n\nInside one K-major SW128 slice, I think of SBO field as telling hardware how to reach the corresponding fragment in the next 8-row atom.\n\nLBO is unused here, beucase as far as I understand:\n\n- SBO already tells how to reach the other atom chunks in the same micro-tile, and separately,\n- advancing the descriptor base(our software loop around MMAs) tells how to move to the next micro-tile. So these two together already cover both axis of the data.\n\nSo in this case, no additional offset (like LBO) is needed to specify another independent direction for stepping through the data.\n\nTo recap:\n\nOnce the major mode is chosen, SW128 tells us the shape of one atom and how its 16B elements are permuted inside it. But SW128 alone does not tell us the order of all atoms in the full tile.\n\nOur TMA copy destinations decide where each atom lands in SMEM. On the MMA side, the descriptor base (which our software loop advances on each step over micro-tiles), and SBO, and if applicable LBO, must describe that same placement. Otherwise, tcgen05 will read the wrong data.\n\nSo far, we covered K-major layout with swizzling applied. Now let’s discuss the MN-major layout with swizzling applied, used for V in the PV matmul.\n\nEach PV MMA microtile reads a `16×128`\n\nslice of V:\n(16 rows along reduction-K and 128 columns along output-N).\nSince each atom is `8×64`\n\n, this slice contains four complete atoms.\n\nIn this `V[K,N]`\n\nview, LBO finds the next atom to the right,\nwhile SBO finds the next 8-row stripe below.\n\nThe diagram shows two of the eight PV microtiles to keep it readable.\n\n```\nKERNEL CHECKPOINTSee 3_swizzle.cu and diff it against Kernel 2.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n3 | Swizzling | 25.9% | 25.9% | 26.2% |\n\n## Chapter 4 – Warp Specialization\n\nSo far, our kernel uses 4 warps, each of which is independently schedulable.\nBut all of the 4 warps (128 threads) are synchronized multiple times throughout the kernel, so our synchronization is effectively CTA wide.\nEarlier kernels already use `mbarriers`\n\n, which can be used more granularly,\nbut we still make all the warps wait on them, and several `__syncthreads()`\n\ncalls still synchronize the entire CTA (as shown on the diagram above).\n\nWe would like to make workers (warps) independent so that later we can parallelize and overlap the work among them; but first we need to make the warps independent (not yet giving them independent work but sort of giving them capacity to run independently).\n\nFor example, if we have 128 threads total (4 warps), and at present we synchronize all of them at each arrow\n`load K/V(i) -> QK(i) -> softmax(i) -> PV(i) -> load K/V(i+1)`\n\n.\nBut we can rearrange our kernel in such a way that not all of the threads are synchronized at each arrow,\nbut only the ones which are *need to synchronize for a given transaction* (e.g. only 2 of the 4 warps can synchronize, without stalling the remaining 2 warps).\n\nFor example load producer and mma consumer handshake through a barrier, ie one producer warp arrives on the barrier signaling data is ready[^20], another consumer warp waits on the barrier, while other 2 warps don’t need to participate in the handshake, so in the future they can progress chugging along with their own work, without being stalled on a global handshake. This allows other warps (which are not needed for the transaction) to not be blocked/synchronized by it, and therefore to not pause their work unnecessarily.\n\nKernel 4 does that kind of transformation. Instead of synchronizing all 128 threads at every stage of the work (as in earlier kenels), we give each warp it’s work and we synchronize only the subset of them, which is actually relevant for a given transaction.\n\nNotably, Kernel 4 gives warps independent roles, but it does not give them enough independent work to overlap **yet**.\nIe the warps that don’t participate in a given handshake don’t overlap their own independent work yet – that’s what some of my later chapters do.\nSo even though the warp roles now have separate instruction streams and the hardware operations are asynchronous,\nthe major stages are still effectively serialized (load -> MMA -> softmax -> MMA).\nConcretely, for my earlier example of load warp and mma warp handshake,\nfor now the rowwarps don’t have outstanding work to do so they are still\nsitting idle even though they don’t participate in the load-mma handshake directly.\n\nThis is clearly suboptimal (we’ll want to give these idle warps some work) and the topic of some of the future chapters. Which I did for teachability, rather than throwing all optimizations at once, so it’s easier to follow.\n\nAnd this change alone is not expected to increase perf (it does not create more parallel work yet), but it enables our future optimizations. So basically, kernel 4 conceptually splits the kernel on individual workers, (where some do syncrize, but others now have the capacity to progress independently), and subsequent kernels given these workers independent work, turning that capacity into actual work overlap.\n\nTo create these independent workers, kernel 4 needs more warps than the previous kernel:\n\nFirst, we’ll need 4 warps only for parallel processing of entire S->P the TMEM\n(because we have 128x128 S, and a single warp can only load 32-rows at a time, we want\nto load and process all 128 rows in parallel, thus we’d need 4 warps)\nand this will be our “row warps” – tasked with doing softmax, and O-rescaling.[8](#fn:8)\n\nSecond, if want want load warps be independent worker, we need one warp for that (cos the SMs scheduling granularity is in warps and not individual threads – so even though load worker only needs a single warp to launch TMA, we will allocate whole warp for it) so we have 5 warps at this point.\n\nThen if want MMA to be a separate worker we need a separate warp for that also (by same the logic as in the TMA case above) resulting in 6 warps total.\n\nSo kernel 4 uses 6 warps in total: 4 row warps, 1 load warp, 1 MMA warp.\n\n```\nKERNEL CHECKPOINTSee 4_warp_specialization.cu and diff it against Kernel 3.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n\n# Part II — Coarse Scheduling\n\n## Chapter 5 – Two Q Tiles\n\nKernel 4 assigns one Q tile to each CTA. So, two neighboring CTAs (ie CTAs which are assigned consecutive Q tiles for the same batch/head) independently loop over the same K/V tiles sequence and each of the CTAs loads the same K/V tiles. Though these are the same K and V operands which each CTA loads, so we can avoid loading them twice (once per Q tile / CTA) and instead re-use them (in each kv-loop iteration in a single CTA). Our next optimization is to let one CTA process two Q tiles and reuse each loaded K/V tile for both of them. This amortizes each K/V load across two Q tiles worth of work.\n\nSimply moving both Q tiles into one CTA (left on the diagram above) is not enough. If we naively make one CTA process two Q/Output tiles, it wouldn’t create direct SMEM K,V re-use (if a CTA completes the entire KV loop for the first Q tile and then repeats the loop for the second, it still loads the same K/V tiles twice).\n\nKernel 5 instead loops over K/V once, using each loaded tile for both Q tiles.\n\n```\nload Q0 and Q1 once\n\nfor each KV tile i:\n    load once K(i) and V(i)\n\n    use K(i) for Q0K(i) and Q1K(i)\n    use V(i) for P0V(i) and P1V(i)\n```\n\nFor the same 2 Q tiles, this reduces the number of issued K/V TMA loads by half compared to our earlier kernels. We reuse the same KV tiles directly from SMEM (not L2 re-use: we’re not relying on the data staying in L2 across separate CTAs).\n\nIn the notation below, 0/1 identifies the Q stage, while i identifies the K/V tile, for example:\n`Q0K(i) = Q0 @ K(i)^T`\n\n, and `P0V(i) = P0(i) @ V(i)`\n\n. **We will also use this naming convention for future kernels.**\n\nOne CTA now owns two consecutive Q/O tiles of work (holds 2Q tiles for the duration of the kv-loop, and eventually produces 2-Output tiles).\n\nBoth Q stages reuse the same K/V tiles while they are still in SMEM, but QK, online softmax, and PV still run separately for each Q tile. Each Q stage keeps its own S/P tile, O accumulator, row maxes, and row sums.\n\nThe second Q stream also needs its own TMEM state, so the layout becomes as illustrated above.\n\nAs before, each BF16 P tile consists of two packed values per each 32-bit TMEM cell and overwrites the lower half of it’s FP32 S allocation (as explained in Chapter 2).\nSo, P0 doesn’t need additional storage beyond S0, and P1 doesn’t need additional storage beyond S1.\nThis relies on the Chapter 2 change: if we didn’t write P into the dead-S’ buffer (and instead stored it in a separate TMEM region),\nwe wouldn’t have enough TMEM to naively 9 store S/P/O for two Q tiles.\n\nAnother angle to look at this optimization is: kernel 5 uses TMEM spatially (within the same KV-loop iteration), but not temporally (across KV-loop iterations).\nMeaning, two Q/Output streams are live at the same time and reuse K/V across the two live\nQ tiles (at the same KV-loop iteration), but not temporal compute pipelining (across different KV-loop iterations).\nAll 512 TMEM columns are used, so there is no available TMEM to hold additional full stage S(i+1)/P(i+1) slots for future kv-loop iteration[s] on top of the two currently live Q-stage slots 10.\n\nThis does not implement load pipelining yet, Kernel 5 adds two-Q reuse, but it still has only one current K slot and one current V slot.\n\nPrevisoly we had 6 warps (see chapter 4), now the CTA has 10 warps:\n\n```\nwarps 0..3:     row/softmax work for Q0\nwarps 4..7:     row/softmax work for Q1\nwarp 8:         TMA loads\nwarp 9:         QK and PV MMA issue\n```\n\nThe two row-warp groups are independent: Q0 row work can start while the MMA warp processes Q1K.\n\nMinor note on naming, I call them row warps and not softmax warps because these warps don’t just do softmax, but also O-rescaling (the same naming was used for chapter 4, so this is not new).\n\nThere is still only one MMA warp though, so Q0K, Q1K, P0V, and P1V are being issued by one ordered MMA stream rather than two parallel tensor-core streams. Adding another MMA warp would not create additional hardware tnesor-core pipeline because both target the same tensor-core hardware.\n\nAnd tcgen05.mma is asynchronous, one issuer (selected lane) is enough to submit the MMA work. Instead, the bottleneck generally occurs when the MMA stream has no MMA to issue, for example operand is not ready. So the question is not whether we have another MMA warp, it’s whether the schedule exposes enough work to keep existing tensor core pipeline busy (creates operands fast enough so that MMA is not stalling for work; and MMA issue order doesn’t unnecessarily block MMAs from being issued on existing operands). Later optimizations address this by improving operand readiness and the schedule, not by adding another MMA issuer. That’s why we don’t add another MMA warp here, and keep only 1 MMA warp for the rest of the lineage.\n\nAs we added 2Q stages, this required adding new barriers. Some of the added barriers are shared by both Q stages, while others are separate for each stage:\n\n- q_ready, kv_ready: shared by both Q stages. q_ready is used once before the KV loop starts; kv_ready is reused once per KV iteration.\n- qk_done[q], softmax_done[q], pv_done[q]: separate stream for each Q stage\n\nEach softmax_done[q] expects 4 arrivals, one for every row warp assigned to that Q stage. These row warps do both softmax/P production and O rescaling. Once all 4 arrived, it signals the P tile is ready and the old O accumulator is safe for PV. So this barrier represents both P ready and O-safe. MMA waits on this barrier before issuing PV.\n\nAs explained in the figure, becuase the K/V operands are now shared between two Q/O work items,\nwe must ensure both of them finished using these buffers before re-using them for the next K/V tiles.\nSo we re-use K/V storage only after P1V completes.[11](#fn:11)\n\n```\nKERNEL CHECKPOINTSee 5_two_q_tiles.cu and diff it against Kernel 4.It follows the same optimizations we discussed.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n\n## Chapter 6 – Load Pipeline\n\nDon’t try to parse the diagram above yet, follow the text for now, I will draw your attention to the diagram later when needed.\n\nSo far, when we load K(i),V(i) tiles at each iteration of the loop, we stall and MMA warp waits until both are loaded. Only then QK can proceed, so the rest of the work downstream of QK (S->P, and eventually PV) is also blocked until both K and V operands are loaded.\n\nThere are two obvious inefficiencies there:\n\nQKs only need K tile (Q@K needs no V tile) – so we can cut our load waiting time by letting QK proceed after only K is loaded (without waiting for V). This also unlocks work downstream of QK (S->P), and lets softmax begin earlier (less time idle).\n\nWe can significantly reduce idle time waiting for individual K/V loads, by pipelining the loads.\n\nThe second change is the main optimization of this chapter:\n\nI like to think about this as producer (load warp) and consumer (mma warp), the idea behind pipelining is producer loading operands ahead of what the consumer needs. This way, the load generally has enough time to finish before the consumer needs it. When the consumer finally rolls around to that operand, it’s likely already loaded so MMA can just use it avoiding the stall. And when a future K or V operand is being loaded asynchronously, the consumer is busy working on operands already loaded in SMEM. This way we overalp load and compute.\n\nThe first change makes the pipeline finer grained:\n\nWe’re treating individual K or V as seprate operands being pipelined,\ninstead of treating K(i) and V(i) as one pair, which would require larger physical SMEM buffers. 12\nEach of the 3 physical pipeline slots holds one K or one V tile.\n\n`KV[3]`\n\nis one shared 3 slot pipeline buffer, and any slot may hold either K or V.\nSo, K and V are separated logically, but not placed in separate pipelines.\nThis lets the load warp run ahead instead of waiting for an entire K/V pair to become reusable.This also relies on Chapter 2’s optimization, P no longer occupies K’s SMEM storage. K can therefore be recycled after QK while P remains alive in TMEM for PV.\n\nIn more details:\n\nQ0 and Q1 are loaded once for the whole KV loop, so only the K/V operands (the ones which are changing each KV-loop iteration) participate in the pipeline.\n\nOur load pipeline depth is 3 (3 physical buffers).\n\nBefore the KV loop begins, we pre-load K0, V0, and K1, filling all 3 pipeline buffers. K0 and V0 belong to the first KV-loop iteration, while K1 is loaded ahead for the next iteration. This preloads the data ahead of what the first consumer MMA will need (which is Q0@K0, Q0 is already loaded once before the for-loop, so K0 is the only operand), giving the producer a head start over the consumer.\n\nIn the for loop over KV tiles, each QK or PV uses a preloaded K or V operand,\nonce the 2nd Q stage finishes using that operand,\nthe corresponding consumer Q1K(i) or P1V(i) signals that the old token is no longer needed,\nthe load warp, waiting in its own token loop, then re-uses that slot, issuing load for a future operand\nwhich is 3 (pipeline depth) logical positions ahead (**now, please see the diagram ABOVE**, it illustreate how it’s implemented in code).\n\nWith 3 token pipeline, there are 2 K/V tokens before the newly loaded token is needed. Each token is consumed by both Q stages, so MMA stream processes 4 full QK/PV matmuls before using the new token. This way a given load can be compleated in the background, and generally has enough time to finish loading that operand by the time the consumer needs it later. This effectively will overlap the load with compute.\n\nOf course in hardware buffer isn’t circular, but because the indices wrap around, and tokens separated by pipeline_depth (3) positions map to the same physical slot – it can be treated logically as a circular buffer. The data does not physically move around a circle.\n\nNow covering both diagrams above, for this chapter, in more detail.\n\n`token`\n\nidentifies the logical K or V operand.\n`stage`\n\nmaps that operand onto one of 3 physical SMEM slots. It does not tell us whether that slot is ready or safe to overwrite.\n\nTokens `t`\n\nand `t+3`\n\nuse the same physical slot.\nSo, before loading token `t`\n\n, the loader must wait until the old token `t-3`\n\nis used by its final consumer.\n\nBecause we have 2Q tiles and both Q0 and Q1 reuse the same K/V operand, we can recycle a given K or V token only after the last of the two Q tiles finished using it.\n\nSo:\n\n- for a K token, the final consumer is the last-Q QK\n- for a V token, the final consumer is the last-Q PV\n\nLet’s see the logic behind the calculations we do on a concrete example (see diagram **below**):\n\nFor a visualization of how the pipeline progresses (see the diagram **below**):\n\n```\nKERNEL CHECKPOINTSee 6_load_pipeline.cu and diff it against Kernel 5.It follows the same optimizations we discussed.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n\n## Chapter 7 – Compute Pipeline\n\nLet’s recap the sheudle we have so far (in the earlier kernels), which is relevant for understanding our next optimization.\n\nIn our earlier kernels, MMA warp works on the granularity of kv-loop iterations, ie issues Q0K(i), Q1K(i), then P0V(i), P1V(i). And only then advances to the next kv-loop iteration.\n\nSo once PV for the 1st-Q-tile was issued, its row warps (softmax + O-rescaling) sit idle (unable to advance to the next kv-loop iteration) because the MMA warp does not reach Q0K(i+1) (ie 1st-Q-tile’ QK for the next iteration of the K/V loop) until after it issues 2nd-Q-tile’ PV.\n\nThis is clearly suboptimal, because 1st-Q-tile QK(i+1), and even more importantly its softmax(i+1) are blocked behind the 2nd-Q-tile PV(i).\n\nYet, as dicussed in the chapter 5 where we introduced the “2Q tiles” optimization,\neach of the 2 Q streams of work (which a given CTA processes) are mostly independant\nstreams, with their own buffers (like S/P/O).\nThey do share K/V, but that is not what forces Q0K(i+1) to wait behind P1V(i) 13.\nSo there’s no cross-Q-stage data dependency holding us from letting 1st-Q-tile to advance to the next\nkv-loop iteration (ie no dependency requiring Q0K(i+1) to remain behind P1V(i)).\n\nInstead we want to issue 1st-Q-tile QK(i+1) and later its softmax work, as soon as 1st-Q-tile PV(i) is issued. This will allow the 1st-Q-tile softmax(i+1) to overlap with 2nd-Q-tile PV(i), which is highly desirable because as I said in the first chapter, the main bottleneck of this entire kernel is softmax work so the more of it we can overlap with MMAs the better.\n\nFor the 1st Q stage, if we simply move its QK(i+1) to right after its PV(i) in the MMA schedule (once K(i+1) is ready), conceptually, we’ll achieve our desired result of unblocking this Q stage’s stream of work, letting it proceed to its next QK without being blocked by the 2nd-Q-tile PV.\n\nBut notice if we naively move 1st Q stage’s QK(i+1) to the i-th iteration of the kv-loop, we would accidentally issue the same Q0K(i+1) twice: once near the end of iteration i, and again at the beginning of iteration i+1. And similar problem for the 2nd Q stage.\n\nSo there’s a few modifications we need to do to the kernel, discussed below.\n\nSo we peel-off first QK (for each of Q stage) out of the for loop. Analogous to prefilling load pipeline as we done in kernel 6, but now we’re sort of prefilling the compute pipeline.\n\nNow in the steady state of the loop we no longer issue the same next-tile QKs twice.\n\nBut as a consequence of the prefill before the loop, we also need a corresponding drain logically after the steady state loop (to exhaust the two Q stages work streams, because at the end of the loop over kv tiles there’s no subsequent iteration so unelss we add a drain, the final two PV tiles will not be executed).\n\nAnd that’s how we arrived to our final design for this chapter.\n\nThe schedule visually looks like we reversed QK and PV, but this is not the case.\nHere q is the Q-stage index (0 or 1).\nIn the steady state, each loop iteration now finishes the current KV tile with PqV(i)\nand then starts the next KV tile with QqK(i+1).\nFor any individual KV tile, the order remains QqK(i) -> softmax-q(i) -> PqV(i).\nSo it only *looks* like we reversed QK and PV (for each Q tile), but we only moved the loop boundary; within each KV tile, nothing is reversed.\n\n```\nKERNEL CHECKPOINTSee 7_compute_pipeline.cu and diff it against Kernel 6.It follows the same optimizations we discussed.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n\n# Part III — Hot-Loop Optimizations\n\n## Chapter 8 – Hardware Approximations\n\nThis chapter doesn’t have anything meaningful for me to draw, so we’ll go over a little pseudocode instead. In chapter 10 we will continue with diagrams again.\n\nSo far, most of our optimizations changed the schedule. But the row-side softmax work is still the main bottleneck: this exp (part of S->P logic) runs per Q stage, per KV tile, per row, per 8-score chunk, which is what PV waits on. Extra work here gates the tensor cores. Our next few kernels (including Kernel 8) start making that hot path cheaper. So basically, we want to make the exp in our softmax hot path faster.\n\nThe GPU provides an approximate base-2 exponential instruction `ex2.approx.ftz.f32`\n\n.\nThe “approximation” in the name of this chapter, comes from the hardware computing exp2 approximately 14.\nSo, to try to speedup our hot path, we can use this instruction (hardware approximation)\ninstead of our\n\n`__expf`\n\nintrinsic we’ve been using.But our online softmax still needs exp (and not exp2), while the hardware instruction computes exp2.\nTo produce the same value as regular exp would, we express the exponent in base-2,\nusing basically a math formula `exp(x) = exp2(x * log2(e))`\n\n.\nSo if we convert the input into the units expected by exp2 (by multiplying the input we’ve been feeding to exp by log2(e)),\nits result is already the exp(x) value we need.\nThis identity itself is exact.\n\nActually … it turns out, __expf intrinsic we’ve been using in all our previous\nkernels already lowers to `ex2.approx.ftz.f32`\n\n(plus some wrapper instructions discussed later) !!!\nSo turns out we unknowingly been using that all along!\nSo at first, this optimization (of using `asm ex2.approx.ftz.f32`\n\ndirectly), introduced in this chapter, looks redundant.\nBut in fact it’s not redundant and actually helps improve our kernel perf about 10%.\nThis sounds contradictory: if we were already using MUFU.EX2, where does Kernel 8’s speedup come from?\n\nFirst, I’ll give the punchline, and then we’ll undersntad this deeper.\nTurns out yes, __expf lowers to the same MUFU.EX2 core as `ex2.approx.ftz.f32`\n\n,\n**but with additional conversion and underflow-handling code around it**.\nAnd removing this additional wrapper code is what helps us improve the perf.\n\nNow when we got the high level idea, let’s understand deeper.\nThe compiler expands that `__expf`\n\nintrinsic into something like:\n\n```\n// natural-exp argument originally passed to __expf\n// float x = s * softmax_scale - rowmax;\n\n// convert into base-2 exponent expected by exp2\n// (see the math identity dicussed above)\nfloat u = x * LOG2_E;\n\nbool small = u < -126.0f;\n\n// preserve results that'd underflow\n// inside FTZ exp2 instruction\nif (small) {\n    u *= 0.5f;\n}\n\nfloat y;\nasm(\"ex2.approx.ftz.f32 %0, %1;\" : \"=f\"(y) : \"f\"(u));\n\nif (small) {\n    y *= y;\n}\n```\n\nSo, `__expf`\n\ndesugars into an `ex2.approx.ftz.f32`\n\ncore *plus some wrapper code around it*.\nFor our use case, that wrapper has two pieces of avoidable work **in the hot loop**:\n\n- repeated conversion to the log2 basis\n- underflow behavior preservation (comparison in the predicate)\n\nThe reason for that wrapper code in the first place, is that the compiler cannot\ninfer our larger semantics, so conservatively chooses to do this\nconversion (to the log2 basis, which the exp2 expects) **locally for every expf in our hot loop**,\nand the underflow-handling around the EX2 (`ex2.approx.ftz.f32`\n\n) core is also emitted **for\nevery __expf in our hot loop**, and this additional work takes a noticeable bite out of performance.\n\nKernel 8 keeps the same num of MUFU exponentials, but basically removes the whole `__expf`\n\nwrapper:\n\n- hoist that repeated conversion (to the log2 basis) out the hot loop and fold it into the existing S-scale, and\n- remove the underflow checks and instead rely on the FTZ.\n\nLet’s cover each of the two points above in turn.\n\n**1. Hoist repeated conversion to log2 basis out of the hot loop**\n\nIn earlier kernels, we express the softmax directly in terms of natural-exp:\n`p = __expf(s * softmax_scale - rowmax);`\n\n.\n\nThe compiler handles each `__expf`\n\ncall locally.\nIn the generated code, it does not globally rewrite our online-softmax so that its\nexponent inputs remain in exp2-friendly units.\nSo, it repeats the conversion to the log2 basis for every value in the hot loop.\n\nKernel 8 does that rewrite manually:\n\n```\n// fold the usual 1/sqrt(HEAD_DIM) attention\n// scale together with log2(e)\nsoftmax_scale_log2 = softmax_scale * log2(e);\n// call ex2.approx.ftz directly\n// (not the __expf which would add the wrapper)\np = fast_exp2(\n        // scale S by that scaling factor, which includes log2,\n        // so now the inputs are in exp2 input-friendly units;\n        // subtract rowmax as usual\n        s * softmax_scale_log2 - rowmax\n        );\nrescale = fast_exp2(old_rowmax - new_rowmax);\n```\n\nRowmax actually also need to be in log2 units.\nRecall from chapter 1, our kernel does two passes over TMEM to convert S into P.\nThe first pass over S still finds the maximum of the raw S scores, this stays unchanged.\nAfter the reduction, in kernel 8, we multiply that one scalar (max per row) by `softmax_scale * log2(e)`\n\n, which puts rowmax directly into exp2-input units.\nSo both sides of the subtraction above are in the units exp2 expects:\n\n```\nscore:  s * softmax_scale * log2(e)\nrowmax: already stored in those same units\n```\n\nP, rowsum, and O are not stored in log2 space. Only rowmax and the scaled S (which is passed to exp2) use these units.\n\n**2) Remove underflow handling and rely on FTZ**\n\nSecond, we remove `__expf`\n\n’s underflow handling. So we remove the comparison and branching logic.\n\nThis part, unlike the first part above, is a real numerical apprixmation.\nWhen the result is smaller than the smallest FP32 number (subnormal values), `__expf`\n\nwrapper tries to\npreserve it by computing `exp2(u / 2)`\n\nand then squaring the result.\n\nThe __expf wrapper always computes the predicate of the if/else branches, even when, for almost all cases, the S value is large enough so that the branching logic doesn’t apply – but we still always pay for that compare in the predicate.\n\nTo decide “is this input in the underflow range?” we had to actually do the comparison, there’s no way to\nskip the check that determines whether the check’s consequences apply.\nSo this compare instruction runs on every `__expf`\n\ncall regardless of the score values.\nBut these two predicated fixup multiplies only matter when the result of exp2 is below 2^-126 (smallest FP32 value).\n\nSo, Kernel 8 deliberately flushes subnormal P and rescale values to zero. We make a call that this is acceptable here. This is what compiler could not have done for us, because this is not semantics preserving modification, and compilers are conservative to make such kind of modifications.\n\nOnce we stop trying to preserve subnormal results, there is no fixup to select, so we can remove the predicate and comparison too.\nDirect `ex2.approx.ftz.f32`\n\nskips that handling and simply\nturns the value into zero (denoted by the FTZ in the instruction name).\n\nTogether, these two choices let us remove the per-call `__expf`\n\nwrapper, its\nbase2-conversion, range check, and predicated underflow handling.\nOfficial FA4 makes the same direct-FTZ choice.\n\nSASS comparison:\n\n``` php\nMUFU.EX2:             129 -> 129\nstatic instructions: 3296 -> 2776\npredicated inst:       368 -> 110\nFMUL:                   644 -> 257\nFSETP.GEU.AND:          129 -> 0\n```\n\nSo Kernel 8 changes didn’t reduce MUFU count, but made the work around each MUFU cheaper.\n\n```\nKERNEL CHECKPOINTSee 8_hardware_approx_exp2.cu and diff it against Kernel 7.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n\n## Chapter 9 – Software Approximations\n\nThis chapter doesn’t have anything meaningful for me to draw, so we’ll go over a little pseudocode instead. We will continue with diagrams in the next chapter.\n\nKernel 8 removed extra wrapper work around every hardware exponential,\nbut all P exponentials still went through the same `MUFU.EX2`\n\npipeline.\n\nKernel 9 moves some of these exponentials off MUFU and approximates them in software using regular ALU/FMA instructions.\n\nSo we do not reduce num of exp computations (we still need to exp each value in S, as in all our previous kernels). But now we speread the workload across two different pieces of hardware, reducing pressure on the MUFU pipeline.\n\nThe split is:\n\n``` php\nfor each 16 P values:\n    first 12 -> hardware exp2\n    final 4  -> software exp2\n```\n\nThus, 25% of the P exponentials use the software path. This selection depends only on the column position, not on the individual P values themselves.\n\nThe exact split (25% for ALU and 75% for MUFU) is determined experimentally, and is not a hard requirement. Software exp2 needs more instructions, temporary values, and registers. Moving everything to software would move the bottleneck onto other resources.\n\nWhy can several software instructions be faster than one hardware instruction? Because software exp2 uses different execution resources. The hardware path is shorter, but all those exponentials compete for the relatively limited MUFU pipeline. Moving some of them to the regular ALU pipelines gives GPU more independent resources with which to do the work.\n\n```\nSOFTWARE EXP2 WALKTHROUGHExplanation above is enough to follow the rest of the blog.For an optional walkthrough of the approximation math, see\nHow Software exp2 Works in the Appendix.\n```\n\nThe helper looks like:\n\n```\nfloat software_exp2_scalar(float x) {\n    // polynomial coefficients\n    float C3 = 0.077f;\n    float C2 = 0.227f;\n    float C1 = 0.695f;\n\n    // input is in base-2 units, and we want to compute 2**x,\n    // if x is below -127, that means 2**x will be extremely small,\n    // so small that we don't care (FTZ like)\n    x = fmaxf(x, -127.0f);\n\n    int int_part = floor(x);\n    float frac_part = x - int_part;\n\n    // approximate 2^frac_part using three chained FMAs;\n    // as an optimization, feed each result into the next FMA,\n    // avoiding separate frac_part^2 and frac_part^3 computations\n    float frac_exp2 = C3;\n    frac_exp2 = fma(frac_exp2, frac_part, C2);\n    frac_exp2 = fma(frac_exp2, frac_part, C1);\n    frac_exp2 = fma(frac_exp2, frac_part, 1.0f);\n\n    // expose the bits of our approximate 2^frac_part\n    int frac_exp2_bits = __float_as_int(frac_exp2);\n\n    // FP32 has 23 fraction bits, so 1 << 23 is one exponent step\n    int exponent_adjustment = int_part * (1 << 23);\n\n    // combine approximate 2^frac_part with 2^int_part\n    int result_bits = frac_exp2_bits + exponent_adjustment;\n    return __int_as_float(result_bits);\n}\n```\n\nThis algorithm is adapted from FA4’s `ex2_emulation_2`\n\n.\nThis is a simplified version, processes one value for clarity.\nThe [actual helper](https://github.com/IaroslavElistratov/b200-attention/blob/master/kernels/approximations.cuh)\nperforms the same algorithm on two values together using packed f32x2 instructions.\n\nBoth paths (hardware exp2, added in chapter 8; and this chapter’s software exp) are approximate. This optimization changes only how selected P values are computed and does not reduce number of exp computations.\n\n```\nKERNEL CHECKPOINTSee 9_selective_software_approx_exp2.cu and diff it against Kernel 8.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n| 8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n9 | Software `exp2` | 69.3% | 70.1% | 70.6% |\n\n## Chapter 10 – Register Caching\n\nIn the earlier kernels, for each Q stage, we traversed S in TMEM and read it twice (explained in chapter 1). First, to compute rowmaxes (to be used for numerical stability in the 2nd step). Second, to subtract rowmaxes and exponentiate the result (so we first need to have computed rowmaxes before we can subtract them in this step, which is why we traversed same S in the TMEM twice).\n\nOur next optimization reduces 2nd-pass TMEM re-reads by retaining part of S in registers.\n\nSo we want to reduce register lifetimes, so we:\n\n- (a) in the 1st pass over TMEM, cache the most recent values of S that we read (because we read S low to high, it’s the high 96 cols which we encountered latest)\n- (b) change the P-production pass (previously the 2nd TMEM pass) to traverse high to low (so these recently cached values are consumed first)\n\nIn the previous kernels (beginning from Chapter 2), we stored P in the lower part of S region in TMEM. In this chapter, since only the high 96 are cached (as explained earlier), the low 32 must be re-read later (128 cols total). So we cannot store the high 96 P at the lower S anymore (storing P at lower S, would clobber the low 32 S before that reread), so begining from Chapter 10, we store P at the upper half of S.\n\nOne more detail: the re-read S low32 are needed to produce only 1/4 of P. The rest 3/4 of P is produced using the S values cached in regs.\n\nTherefore, the second pass over S remains, but it rereads only 32 of 128 scores from TMEM, removing 3/4 of the re-reads.\n\n```\nKERNEL CHECKPOINTSee 10_cache96_reread32.cu and diff it against Kernel 9.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n| 8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n| 9 | Software `exp2` | 69.3% | 70.1% | 70.6% |\n10 | Register caching | 72.9% | 73.7% | 73.6% |\n\n## Chapter 11 – Skip O rescale\n\nThis chapter doesn’t have anything meaningful for me to draw, so we’ll go over a little pseudocode instead. In the next chapter we will continue with the diagrams.\n\nIn regular softmax, we see the entire row, find the final rowmax once, subtract it, and then compute the exponentials.\n\nOnline softmax processes one K/V tile at a time, so we don’t know the final rowmax in advance, it discovers the row tile by tile, so a later tile may contain a larger maximum. So the unique decision online softmax has to make (unlike the reguaral full softmax) is:\n\n- Do we immediately switch to every newer maximum?\n- Or can we keep using the old rowmax?\n\nIn earlier kernels, whenever we found a larger maximum, we immediately started using it:\n\n```\nnew_rowmax = max(old_rowmax, tile_rowmax)\n```\n\nThen they expressed all the old state in that new basis:\n\n```\nold O      *= rescale\nold rowsum *= rescale\n```\n\nThe rowsum update is cheap, but rescaling O requires another full traversal\nof its 128 values per row in TMEM.\nKernel 10 ran this O pass after every iteration of the kv-loop, even when the scale was 1.\nSo, Kernel 10 traversed the entire O tile when the maximum did not change and `rescale == 1`\n\n.\nSo it loaded 128 O values, multiply them by one, and store them back.\nThat is numerically safe but unnecessarily conservative.\n\nIn Kernel 11, if the new tile maximum is only slightly above our currently kept rowmax, we keep using the old rowmax and leave O and rowsum unchanged. If the difference is large enough, only then we switch to the new rowmax and rescale the old O and rowsum into the new basis.\n\n**Skipping rescaling doesn’t erase information**\n\nThis was not immediately obvisous to me 15,\nbut let’s see what happens when we skip rescaling O on a little example,\nto convince ourself that no information gets lost\n(as a reuslt of us skipping the rescale).\n\nSuppose we decided to skip O rescale for a given tile\n(we gonna discuss the specific threshold value later).\nFor now, let’s just say our current tile’s maximum is `7.78`\n\nabove our old rowmax.\n\nWe still compute P relative to the old rowmax:\n\n```\nP = exp2(scaled_S - kept_rowmax)\n```\n\nSo, the newer maximum produces:\n\n```\nexp2(7.78) ~= 219\n```\n\nSo even though we skipped rescaling O, the newer maximum is not erased and no info is lost. Because we still compute P using the old rowmax, that larger score turns into an unnormalized P value about 219x as large. That larger P value is where the difference remains preserved.\n\nEven better, compared to the world where choose to rescale instead, the final result of our entire kernel does not change at all. Because the larger P enters both parts of softmax, the larger P values enter both the numerator O and the denominator rowsum.\n\nThe same P values enter both parts of attention:\n\n```\nO      += P @ V\nrowsum += sum(P)\n```\n\nSo, for my earlier `219`\n\nexample:\n\n```\nO      += 219 * V_new\nrowsum += 219\n```\n\nThen, after the entire K/V loop:\n\n```\noutput = O / rowsum\n```\n\nSo the newer score is still treated as about 219x larger in both O and rowsum. The final division does not erase that 219x relative difference; it only cancels the common scale shared by O and rowsum.\n\n**Then why rebase at all?**\n\nIf the final `O / rowsum`\n\nhandles the common scale,\nit seems like we could keep the first rowmax forever.\nMathematically seems, we can, as long as all the intermediate values don’t overflow.\nFor example, if the old rowmax became very stale:\n\n```\ngap = 20\nlargest P = exp2(20) = 1,048,576\n```\n\nThe final answer will still mathematically be equivlant, but P, rowsum, and O would now use huge intermediate values. If these overflow, obviously the final normalization (dividing by the rowsums) cannot undo that overflow.\n\n**How the per-row decision maps to the warp-wide TMEM O pass**\n\nThe decision to rebase is made separately for each row, but the TMEM O loads and stores are warp-collective. So the actual implementation looks like this:\n\n```\n// convert the current tile's raw rowmax into base-2 exponent units\ntile_rowmax *= softmax_scale_log2;\n\n// kept_rowmax is already stored in the same units\nfloat candidate_rowmax = max(kept_rowmax, tile_rowmax);\nfloat rowmax_gap = candidate_rowmax - kept_rowmax;\nbool row_needs_rebase = rowmax_gap >= REBASE_THRESH_LOG2;\n\nfloat rescale = 1.0f;\n\n// 1) if needed, rescale this row's state and start using the new rowmax\nif (row_needs_rebase) {\n    // this factor will be applied to both rowsum and O\n    rescale = fast_exp2(kept_rowmax - candidate_rowmax);\n\n    // move this row's scalar accumulator into new rowmax basis\n    rowsum *= rescale;\n\n    // use new rowmax when producing this P tile and future P tiles\n    kept_rowmax = candidate_rowmax;\n}\n\n// 2) apply same rescale factor to O\n//\n// TMEM loads/stores are warp-collective. If any row needs rescaling,\n// entire warp must process its 32x128 O slab.\n// Rows that don't rebase still participate but with rescale 1\nif (__any_sync(FULL_WARP_MASK, row_needs_rebase)) {\n    float o8[8];\n\n    // same O traversal used in earlier kernels\n    for (int col = 0; col < HEAD_DIM; col += 8) {\n        int taddr = taddr_o_stage + trow + col;\n\n        tcgen05::ld_32x32b_x8(taddr, o8);\n\n        for (int i = 0; i < 8; ++i) {\n            o8[i] *= rescale;\n        }\n\n        tcgen05::st_32x32b_x8(taddr, o8);\n    }\n\n    asm(\"tcgen05.wait::st...\");\n    asm(\"tcgen05.fence::before_thread_sync...\");\n}\n\n// P production and the later PV proceed as in earlier kernels,\n// using kept_rowmax as the rowmax for this tile\n```\n\n**Actual thresh value**\n\nSince Kernel 8 we express our exponent inputs in base-2 units. Kernel 11 uses threshold of 8 base-2 exponent units, matching official FA4. As far as I understand, there is no deep mathematical cliff at 8. Larger thresholds such as 12 or 15 numerically plausible for our BF16-P/FP32-accumulator path. 8 is a conservative point that already appears to skip most useful O-rescale opportunities not an overflow boundary.\n\n```\nexp2(8) = 256\n```\n\nSo we keep using the old rowmax while the largest unnormalized P value, computed relative to that rowmax, stays below roughly 256. And I say “relative to that rowmax” because we subtracted the earlier rowmax before exponentiating the new value. Once it grows past that, we switch to the newer rowmax and rebase the old state.\n\nFor intuition, let’s see what this base-2 threshold of 8 corresponds to, in the other units used by the kernel, with HEAD_DIM = 128:\n\n```\nbase-2 exponent gap:  8\nnatural-exp gap:      8 / log2(e) ~= 5.545\nraw QK-score gap:     5.545 * sqrt(128) ~= 62.7\n```\n\nThe 62.7 is the gap between the current tile’s per-row maximum and the raw QK score corresponding to kept_rowmax. We don’t use 62.7 anywhere in the kernel (the threshold is 8, expressed in the log2 basis). This conversion above is just to give us a sense of how much slack the threshold allows in raw QK-score units before we rebase.\n\n```\nKERNEL CHECKPOINTSee 11_skip_o_rescale.cu and diff it against Kernel 10.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n| 8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n| 9 | Software `exp2` | 69.3% | 70.1% | 70.6% |\n| 10 | Register caching | 72.9% | 73.7% | 73.6% |\n11 | Skip O rescale | 81.7% | 83.0% | 83.3% |\n\n## Chapter 12 – Split Correction from Softmax\n\nIn earlier kernels, the same row warps:\n\n- compute the current tile rowmax and the O-rescale factors;\n- rescale O, if the new tile rowmax exceeds the running rowmax by the threshold (see kernel 11);\n- subtract the rowmaxes from S and exponentiate to produce P.\n\nSo (2) delays completion of the row warp’s work (1-3). So O correction and P production are serialized in the same row warps.\n\nThe O-rescaling (2) does not depend on computing the unnormalized-P (3), so we can move O-rescaling to separate warps which can do this work in parallel to our row warps computing the unnormalized-P (3). Therefore reducing the main bottleneck: P production (removes O correction from in front of P production, allowing the two branches to overlap). The row warps are mostly doing the softmax work now, so I’ll call them “softmax warps” from now on.\n\nPV still requires both things to be true: P ready and O-safe, for each Q stage, these are combined into p_ready_o_safe barrier, so PV mma waits on the barrier that has 8 arrivals: 4 from softmax warps (which indicates P ready) and 4 from correction warps (which indicate O-safe). So the slowest of the two branches delays PV.\n\nNote I deliberaly say o-safe and not o-rescaled, becuase as per our optimizations in Chapter 11, sometime we skip O rescale (so we do not always rescale O).\n\nThe softmax warps still compute the rescale factors, write the rescale factors to SMEM and signal stats_ready, and then continue producing P. The correction warps then wait for that signal, read those factors, decide whether to rebase O (if any published scale is not 1), and if so, rescale O.\n\nThe softmax warps still compute the rescale factors (as oppose to letting rescale warps do that) because the softmax warps already have the old running rowmax and the new tile rowmax in registers. Otherwise, we’d need to send that state to the correction warps and repeat the rebase logic there. So, softmax sends only one final scale value per row, and correction simply applies it to O.\n\nThe same four correction warps are shared across Q0 and Q1. For each Q stage, they read the per-row scale published by the softmax warps and use it to rescale that stage’s O tile.\n\n```\nKERNEL CHECKPOINTSee 12_split_correction_from_softmax.cu and diff it against Kernel 11.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n| 8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n| 9 | Software `exp2` | 69.3% | 70.1% | 70.6% |\n| 10 | Register caching | 72.9% | 73.7% | 73.6% |\n| 11 | Skip O rescale | 81.7% | 83.0% | 83.3% |\n12 | Split correction | 84.9% | 85.8% | 86.3% |\n\n## Chapter 13 – Early PV\n\nRemember P is produced right to left (written from the high end of its TMEM region towards the low end), as a consequence of our earlier optimizations (see kernel 10). So the right side of P becomes ready earlier than the entire P (by the time green box is produced, red box is not produced yet).\n\nKernel 13 exploits that P production order, once the\nhigh96 P region is produced and old O is safe, a part of PV MMA can begin.\nThe idea is to avoid blocking PV until the entire P is ready, and instead:\ndo `p_high96 @ V_slice1`\n\nfirst;\nfollowed by `P_low32 @ V_slice2`\n\nlater (when P_low32 is produced).\n\nThis relies on the fact that by the time we start producing P chunks, we have already computed rowmaxes for the entire S tile, so the the running rowmax/reference is fixed, so later producing P_low32 cannot change the rowmax, and therefore cannot change already-produced P_high96.\n\nThe fact that P split visually mirrors one of our earlier optimizations from kernel 10,\nnamely the S register caching split (both are split 32/96) is not a fundamental requirement.\nThese two are largely orthogonal optimizations, that happen to align in our implementation.[16](#fn:16)\n\nNote I omitted P1 (ie P of the 2nd Q stage) from the diagram cos the logic there is identical.\n\nAfter the first barrier releases, all p_high96 is ready (ie all 6 high96 microtiles are ready, and O is safe).\nOnce the complete high96 P region has been published, MMA does not need to consume its microtiles in the same order that the softmax warps produced them.\nWe could have equivalently issued k=7..2; this would not meaningfully change numerics 17.\nI kept the k=2..7 PV microtile traversal order because incrementing the corresponding\nV SMEM descriptor looked cleaner (so this is not a fundamental decision).\n\nWhen doing the matmul with p_high96, we also need to select corresponding V slice (ie offset P’s micro-tile address and V’s descriptor by 2 microtiles). Green part of the diagram.\n\nLater, once the softmax warps produced P_low32 and signaled that full P is ready, we matmul the remaining 2 micro-tiles (the ones we skipped earlier), P_low32 @ V_slice2. Red part of the diagram.\n\nBoth partial matmuls accumulate into the same O tile. We are only splitting P and V along the reduction axis, each into 2 chunks, together they still compute the original full P@V.\n\n```\nKERNEL CHECKPOINTSee 13_early_pv_96_32.cu and diff it against Kernel 12.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n| 8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n| 9 | Software `exp2` | 69.3% | 70.1% | 70.6% |\n| 10 | Register caching | 72.9% | 73.7% | 73.6% |\n| 11 | Skip O rescale | 81.7% | 83.0% | 83.3% |\n| 12 | Split correction | 84.9% | 85.8% | 86.3% |\n13 | Early PV | 87.2% | 88.1% | 88.4% |\n\n# Part IV — Persistent Kernels\n\n## Chapter 14 – Persistent\n\nOur kernels have at most one resident CTA per SM (each CTA consumes enough resources that only one can fit on an SM at a time), B200 GPU has 148 SMs, so at any given time we have up to 148 CTAs working. Since our Chapter 5 optimizations, each CTA independently consumes 2Q tiles, loop over all KV tiles, and produces 2 Output tiles.\n\nLet’s call this a “work item” and abstract this visually, like so, and our later diagram will illustrate on the granularity of work items. Our next optimization changes the lifetime around the attention body, not the body itself. So we can safely abstract it away.\n\nFor our kernels so far, if num work items > num SMs (so there’s not enough SMs to process all our work_items at once), each SM cycles through multiple CTAs sequentially: one CTA finishes it’s work item and gets de-scheduled, another new CTA (processing a different work item) gets scheduled in its place.\n\n**These are “waves” of work.**\n\nWe can compute the number of waves, by dividing total work count by the amount of work each CTA does:\n\nFor input shape `SEQ_LEN=4096`\n\n, for a single batch and head, and given our square tile sizes of 128,\nQ and O have 4096/128=32 tiles, and because each CTA processes 2 tiles, 32/2=16 there’s only 16 work items.\nSo,\n\n```\n16 work items per batch-head\n8 batches * 16 heads * 16 = 2048 total work items\n```\n\nFor a different input shape of `SEQ_LEN=16384`\n\n, for a single batch, head, and given our squre tile sizes of 128,\nQ and O have 16384/128=128 tiles, each CTA processes 2 tiles, so 128/2=64 there’s only 64 work items.\n\n```\n64 work items per batch-head\n2 batches * 16 heads * 64 = 2048 total work items\n```\n\nSo, for the 3 shapes we benchmark our kernels on,\n2048 separate CTAs are launched, about 148 can be resident at once,\ntherefore they execute in roughly 14 CTA waves – `ceil(2048 / 148 SMs) = 14 waves`\n\n.\n\nThese are typically called waves. CUDA does not wait for all 148 CTAs in one wave to finish simultaneously. As soon as one CTA finishes, another waiting CTA may become resident on the freed SM, so the CTAs can drift relative to one another.\n\nSo, each SM processes roughly 14 separate CTAs sequentially: one CTA finishes, its resources are released, another CTA becomes resident. But each one of these 14 repetitions incurs the same cost: TMEM alloc/dealloc, CTA-lifetime state (setup, pointer offsets), barrier init, CTA scheduling, etc.\n\n**Fixed CTA overhead matters more at shorter SEQ_LEN**\n\nThis “repeat 14 times” is the same for both 4k and 16k shapes. All 3 shapes we benchmark on (from the FA4 paper) each of them happens to have 2,048 work items (so, same fixed costs of 14 waves).\n\nBut, the KV loop grows from 32 iterations at 4K to 128 iterations at 16K, so each wave contains 4X more useful work at 16K:\n\n```\nKV-loop iterations = len_kv / BLOCK_N = 4096 / 128 = 32\nKV-loop iterations = len_kv / BLOCK_N = 16384 / 128 = 128\n```\n\nSo, same number of output work_items; different amount of hot-loop work inside each work_item. Thus the fixed overhead matters more for smaller-SEQ_LEN shapes: they still pay the same setup cost 14 times, but each KV loop is shorter. So, compared with longer loop lengths, the setup cost is proportionally larger for smaller SEQ_LEN shapes. Therefore, all the fixed per-CTA costs mentioned earlier take up larger part of the runtime for the shorter 4K work items.\n\nIf instead, we give each CTA more work than a single work item, we can pay that cta-setup cost once but amortize it among multiple work items. And ideally the more we can amortize the better. To amortize the setup cost as much as possible, we keep each CTA alive and let it process its own sequence of 14 work items.\n\nNote we cannot amortize: 2Q loads, KV traversal, output store, work-ID and pointer calculations. Because these legitimately differ per output work-item, even if we let one CTA process multiple work-items it still needs to re-do all these for each of the work-items it computes.\n\nCTAs do not execute in lockstep, which is why their arrows have different lengths in the figure. So the completion of work items in one CTA can drift relative to other CTA’s progression.\n\n**Reusing barriers across work items**\n\nIn earlier kernels, barriers were initialized once when the CTA started and were used for that CTA’ lifetime. Kernel 14 keeps barriers work for multiple work_items.\n\nMy first naive implementation stopped all warps between work items, waited until everything had finished, then invalidated and reinitialized the barriers before starting the next work item. That invalidation step is required by PTX, calling mbarrier.init on a barrier that is still valid is undefined. So before reinitializing, we must first make sure no warp can still be using it. It’s correct, but it creates a hard stop between work items. Every warp waits for the slowest one, all work fully drains, and only then CTA can begin its next work item.\n\nTo avoid repeated barrier init and invalidation at each of 14 work_item boundaries,\nkernel 14 instead initializes all barriers once when the CTA starts,\nand keeps cycling through the barrier phases as the CTA moves from one work item to the next.\nThe old work_done sync (where every warp waited for the slowest one and the current work fully drained)\nmade that barrier re-init safe, and conservatively prevented any buffer from being reused until the entire work item drained.\nKernel 14 uses the existing K/V recycle handoffs across work-item boundaries and adds `q_free`\n\nfor Q SMEM.\nTogether with keeping the barriers initialized, this makes the all-warp `work_done`\n\nsync unnecessary.\n\nSo each buffer is released after its actual final consumer for a given work_item:\n\n``` php\nQ SMEM:\n    final QK reads Q(i)\n    -> q_free\n    -> load Q(i+1)\n\nK/V ring:\n    final QK reads K stage\n    -> refill K stage\n\n    final PV reads V stage\n    -> refill V stage\n```\n\nSo, there is no single point where work_item(i) ends and work_item(i+1) begins. While softmax warps are still normalizing and writing O(i), load warp can begin loading next Q/K/V tiles. Once Q and K are ready, MMA warp can begin QK(i+1).\n\nSo Kernel 14 adds both persistent CTA lifetime and cross-work pipelining.\n\nIn kernel 14, each persistent CTA does not process consecutive work IDs (as illustrated in the drawing).\n\nIn code, this is just `work_id += gridDim.x; // 148 CTAs`\n\nThis static grid-stride assignment gives every work item to one of the 148 persistent CTAs.\nEach work item covers two consecutive Q/O tiles.\nSo when a CTA advances by 148 work items, its tile indices advance by 296 (CTA 0 handles tiles 0–1, then 296–297, then 592–593).\n\nI also tried assigning each CTA a consecutive range of work items. The strided assignment explained above performed better overall, though the results were shape-dependent and I haven’t tried to isolate exactly why. My guess is that shared-L2 K/V reuse helps when CTAs process different Q pairs from the same batch-head close together in time.\n\nThe figure shows which work items each CTA owns, not hard execution boundaries. As explained above, different warp roles can still overlap the end of one work item with the beginning of the next.\n\n```\nKERNEL CHECKPOINTSee 14_persistent.cu and diff it against Kernel 13.Now, familiarize yourself with the code.\n```\n\n**Performance so far**\n\n| Kernel | Change | 4k | 8k | 16k |\n|---|---|---|---|---|\n| 1 | Baseline | 14.2% | 14.1% | 13.7% |\n| 2 | P in TMEM | 15.1% | 15.0% | 15.0% |\n| 3 | Swizzling | 25.9% | 25.9% | 26.2% |\n| 4 | Warp specialization | 26.6% | 26.4% | 26.5% |\n| 5 | Two Q tiles | 40.5% | 40.3% | 41.3% |\n| 6 | Load pipeline | 48.6% | 48.6% | 48.4% |\n| 7 | Compute pipeline | 58.2% | 58.5% | 58.0% |\n| 8 | Hardware `exp2` | 64.1% | 64.1% | 63.8% |\n| 9 | Software `exp2` | 69.3% | 70.1% | 70.6% |\n| 10 | Register caching | 72.9% | 73.7% | 73.6% |\n| 11 | Skip O rescale | 81.7% | 83.0% | 83.3% |\n| 12 | Split correction | 84.9% | 85.8% | 86.3% |\n| 13 | Early PV | 87.2% | 88.1% | 88.4% |\n14 | Persistent CTAs | 90.7% | 90.3% | 89.8% |\n\n# Minor Optimizations\n\nKernel 14 is the end of the main 14-kernel progression. The remaining four changes are smaller local optimizations, so I group them here. In the repo, I still keep these as separately numbered source checkpoints 15–18, so each change remains easy to diff. Together, they get us the last few percent of performance.\n\n[ Deferred score acquisition](https://github.com/IaroslavElistratov/b200-attention/blob/master/kernels/5_minor/15_wide_deferred_score_loads.cu). Reading S from TMEM previously used 16\nloads with a wait after each load. Here, we use\neight wider loads with no immediate waits, four independent rowmax chains, and\na single deferred wait just before P overwrites that TMEM region, allowing\nload latency to overlap useful work.\nMost of the performance improvement comes from using four independent rowmax chains.\nDeferring the wait helps too, making the loads wider by itself had almost no effect.\n\n[ Split rowsum accumulation](https://github.com/IaroslavElistratov/b200-attention/blob/master/kernels/5_minor/16_split_rowsum_accumulation.cu). We already produce P in two parts: 96\ncolumns from scores cached in registers and 32 columns from scores reread from\nTMEM. Instead of feeding both parts through one long rowsum dependency chain,\nhere we accumulate cached 96 into two partial sums, use scalar sum for the\nreread 32, and combine them once.\n\n[ Phase bitmasks](https://github.com/IaroslavElistratov/b200-attention/blob/master/kernels/5_minor/17_phasebits.cu). Every barrier stage only needs 0/1 phase, but\nstoring these phases in small runtime-indexed arrays made the compiler place\nthem in thread-local memory. Here we pack them into integer bitmasks instead,\nkeeping them in registers and removing the spills.\n\n[ TMA L2 promotion](https://github.com/IaroslavElistratov/b200-attention/blob/master/kernels/5_minor/18_tma_l2_promotion.cu). Each TMA row fetch uses one 128-byte half of an\naligned 256-byte region, and another nearby fetch soon needs the other half.\nWe ask TMA to fetch the full 256 bytes into L2, making the neighboring half\nmore likely to already be cached later when we need it;\nthe kernel body itself does not change.\n\n| After adding | 4k | 8k | 16k |\n|---|---|---|---|\nDeferred score acquisition | 92.3% | 92.5% | 92.0% |\nSplit rowsum accumulation | 93.4% | 93.7% | 93.7% |\nPhase bitmasks | 93.7% | 94.0% | 94.1% |\nTMA L2 promotion | 94.4% | 94.4% | 94.1% |\n\n# Capstone Project – Generate Videos!\n\nNow let’s use the attention kernel we just built for something cool: generating videos.\nI highly recommend doing it; you will get a kick out of seeing the kernel **you understand** generate beautiful videos for you.\n\n[ The capstone project](https://github.com/IaroslavElistratov/b200-attention/tree/master/capstone-project)\ncontains the integration of\n\n[our final kernel](https://github.com/IaroslavElistratov/b200-attention/blob/master/kernels/5_minor/18_tma_l2_promotion.cu)from the article (a dense BF16 attention kernel for NVIDIA B200) into the LTX-2.3 video-generation model.\n\nTo generate each video, the model will call our attention kernel 1,056 times.\n\nYou’ll need about $7 for hour of B200 compute, depending on your cloud provider. It’s enough time to set up the repo, download the weights, and generate multiple videos.\n\nA 10.4-second video (249 frames) takes about 41 seconds to generate on B200. Runtime depends on the complete model architecture, not only on our kernel.\n\n**Don’t bother reading the video-model integration code.**\n**Your time is better spent understanding the kernel itself and its optimizations, which I already explained in this article.**\n\n**The capstone code simply integrates our kernel into a video-generation model and runs it.**\n\nBasically:\n\n- Rent the compute, SSH into the machine, clone\n[the repository](https://github.com/IaroslavElistratov/b200-attention), and enter`b200-attention/capstone-project`\n\n. - Follow the\n[setup and generation instructions](https://github.com/IaroslavElistratov/b200-attention/blob/master/capstone-project/GETTING_STARTED.md). - Get a kick out of seeing the near-SOTA kernel\n**you understand** generate beautiful videos for you.\n\nThe last one is the main requirement o_0.\n\n# The Graveyard 🪦 … and Some Hope ✨\n\nSooo… I’ve been on this problem of understanding and trying to improve over FA4 for some time.\n\nI haven’t just tried to explain the kernel, but also developed probably tens and tens of my own new and, what seemed to me, exciting optimization ideas (not counting thousands of other ideas agents separately tried).\n\nImproving over such heavily optimized kernel is hard. Surprise surprise :)\n\nA couple did appear to potentially work though, but I decided to pause and finish this article before continuing my research.\n\nI plan to share different families of my own ideas I tried, and post mortem and the learnings from there, in its own follow-up post. I’m deliberately not covering any of them here yet, until I do some more research.\n\n# Stay in touch for more ML sys\n\nAll kernels and the capstone project are available in the\n[ B200 Attention repository](https://github.com/IaroslavElistratov/b200-attention).\n\nVideo of my walkthrough the code is coming.\n\nI’m also releasing a free e-book soon.\n\nMore ML systems articles, videos, and code are coming.**Follow:**\n[ X](https://x.com/iaro_e) ·\n\n[·](https://www.linkedin.com/in/iaroslav-elistratov/)**LinkedIn**\n\n[·](https://www.youtube.com/@IaroslavElistratov)\n\n**YouTube**\n\n**GitHub**# Appendix\n\n### Benchmark Calibration\n\nFA4 is timed on contiguous BSHD, while my kernels timed on contiguous BHLD.\nBoth receive inputs already stored in the layout they expect, so neither timing includes input-layout conversion.\nA [direct-BSHD variant of the final kernel](https://github.com/IaroslavElistratov/b200-attention/tree/master/benchmarks#input-layout-and-direct-bshd)\nis provided separately. It consumes and produces contiguous BSHD without conversion copies and reaches 92.1–92.7% of same-run stock FA4.\n\nOfficial stock FA4 beta4 (CuTe DSL) against the paper numbers:\n\n| Calibration | 4k | 8k | 16k |\n|---|---|---|---|\n| FA4 paper TFLOPS | 1532 | 1579 | 1601 |\n| Measured stock FA4 TFLOPS | 1484 | 1528 | 1554 |\n| Measured stock FA4 / paper | 96.9% | 96.8% | 97.1% |\n\nThis calibration uses 342 B200 invocations. Taking the median stock-FA4 result for each shape. On the B200s I used, even stock official FA4 itself reaches only around 97% (of their TFLOPs reported in the paper). So I mainly compare my kernels against FA4 measured in the same run (as opposed to comparing to the paper numbers).\n\n[Benchmark harness and instructions to reproduce](https://github.com/IaroslavElistratov/b200-attention/tree/master/benchmarks).\n\n### How Kernel 1 Implements Online Softmax with TMEM\n\n*Optional background for Chapter 1 Section B – Online Softmax.*\n\nSection B showed the high-level TMEM picture: four row warps cover all 128 rows, processing eight columns at a time. Let’s map that into the actual code in Kernel 1.\n\n**iterating in 8 col chunks**\n\nThe x8 here is the TMEM column width of tcgen05.ld; unrelated to the 8-element SMEM slices used by the MMA descriptors earlier.\n\nNow onto understanding `32x32b.x8`\n\n(part of the instruction name):\n\n- The\n`.32x32b`\n\npart means:**32 lanes × 32 bits** per repetition: the warp covers 32 TMEM lanes (rows) - The\n`.x8`\n\npart means: repeat that load**8 times along columns**, so each thread gets** 8 registers**(8 FP32)\n\n`tcgen05.ld`\n\nsupports loading wider chunks e.g. 32 columns, not only 8 as we do here.\nThere’s a tradeoff though as it would require less iterations along the TMEM columns\nbut more regs to hold the wider chunks. For now I somehwat arbitrarily picked width 8.\n\nThat’s why intuitively with 4 warps, we’re marching to the right along TMEM columns (the for-loop in the diagram), loading 8 columns at a time.\n\nAt each iteration of that for loop, each warp loads 32x8 chunk computing per row max. Where each thread holds 8 elements, so logically 32 threads in a warp hold a 32x8 tile. So, each thread loads 8 FP32 values from TMEM.\n\n**How four row warps cover all 128 TMEM rows**\n\n`tcgen05.ld/st`\n\nis a whole-warp instruction.\nOne warp covers 32 TMEM rows at a time.\nTo process all 128 rows at the same time, we use 4 warps.\nThat’s why our baseline kernel has 4 warps.\n\n```\nwarp 0 handles rows   0..31\nwarp 1 handles rows  32..63\nwarp 2 handles rows  64..95\nwarp 3 handles rows  96..127\n```\n\nS is (block_m, block_n), where block_m maps to the TMEM rows, and block_n to TMEM columns. Here we’re iterating over columns.\n\nA TMEM address is 32 bits: the upper 16 bits contain the starting TMEM lane (which I draw as a row), and the lower 16 bits contain the starting column.\n\nKernel encodes the first row owned by this warp as:\n\n``` js\nconst int row_base = warp_id * 32;\nconst int trow = row_base << 16;\n```\n\nIn the S and O loops below, `taddr_s`\n\nor `taddr_o`\n\nsupplies the TMEM base column.\n`trow`\n\nstays fixed for the warp, while the loop’s column offset advances by 8\non each iteration. So the warp keeps processing the same 32 rows while marching\nacross the TMEM columns.\n\nBecause tcgen05.ld is warp-collective, all 32 threads pass the same taddr. That address selects the complete 32-row region, while lane ID implicitly selects row_base + lane_id inside it. So we do not compute a separate TMEM row address for every thread.\n\n**computing max**\n\nThis first pass over TMEM’s `S`\n\ncomputes the row max for the current iteration of the KV-loop.\nWe cannot form P yet because we first need the updated running rowmax\n(to be subtracted from S before exponentiation, for numerical stability).\n\nEach of the 4 warps loads its own 32x8 chunk of S, so each thread in a warp holds 8 elements in its registers. Each thread independently computes max over the 8 elements it’s holding in its regs.\n\n```\nfloat tile_rowmax = -FLT_MAX;\n\n#pragma unroll\nfor (int score_n8 = 0; score_n8 < BLOCK_N / 8; ++score_n8) {\n    // 8 FP32 temporaries per lane\n    float s8[8];\n\n    const int taddr = taddr_s + trow + score_n8 * 8;\n    tcgen05::ld_32x32b_x8(taddr, s8);\n\n#pragma unroll\n    for (int i = 0; i < 8; ++i) {\n        tile_rowmax = fmaxf(tile_rowmax, s8[i]);\n    }\n}\n\n// apply standard attention scale, 1 / sqrt(HEAD_DIM), to the scalar\ntile_rowmax *= softmax_scale;\n```\n\nAnd by the time we iterate to end of BLOCK_N number columns in TMEM (in the 8 element chunks, computing max at each chunk), we effectively computed S row maxes (each thread in a warp now holds the max value for a given TMEM row).\n\n**Updating running basis**\n\nNow, once we computed `Score`\n\n’ per-row maxes, we need to update the running counters as per online softmax definition.\n\n``` js\nconst float new_rowmax = fmaxf(rowmax, tile_rowmax);\nconst float rescale = __expf(rowmax - new_rowmax);\n\nrowmax = new_rowmax;\nrowsum *= rescale;\n```\n\n`rowmax`\n\nis the per-row maximum across all previous S tiles.\n`tile_rowmax`\n\nis the per-row maximum of the current S tile.\n\nAs mentioned in Chapter 1, from the thread perspective, these are scalars. Because we’re accessing tmem with 4 warps, each of which accesses 32 rows, and has 32 threads – so, for a single thread there’s a max and a sum scalar.\n\nThis is just standard online-softmax stuff, and not really focus of this blog.\n\nSoftmax is: take values, exponentiate them, then divide by the sum of exponentiated values. Online softmax needs a running unnormalized denominator (“unnormalized” because it hasn’t yet been divided, as this is done once at the end of the kernel).\n\nAnd it is computed by accumulating, for each scalar score in the tile:\n\n```\nrunning_rowsum += exp(this_tile_S - running_rowmax)\n```\n\nThe above is just a sum of exponentiated values.\nAnd that `- running_rowmax`\n\nshift is just the max trick so exp doesn’t blow up.\n\nIf the current tile raises the max, the previous tiles have been accumulated using an older max, and are still expressed relative to the old max. So we need to rescale – both the old rowsum, and the old O accumulator – by the delta between old max and new max:\n\n```\nrunning_rowsum *= exp(old_rowmax - new_rowmax)\n```\n\nIn addition to rescaling rowsum (shown above) we also need to resclae the partial O (ie update it to the new basis).\n\n**Correcting O in TMEM**\n\nAt this point, we know the new max basis. But the old O tile is still in the old basis. So before we do `O += P @ V`\n\nwe need to rescale O.\nIn the online-softmax algorithm (again not unique to Blackwell B200), PV is generally unsafe to accumulate into O until old O is in the current max basis.\n\nLet’s orient ourselves in the bigger online-softmax picture.\nBasically we computed per `Scores`\n\ntile maxes, and we need to rescale the partial output accumulator `O`\n\nto the new basis.\n\nAgain these math semantics is standard online-softmax formulation – the only changes here is that I’m mapping it to the Blackwell B200 hardware.\n\nAs mentioned earlier, the rowsum and rowmax are per-thread scalars. O is different, it’s a full [BLOCK_M, HEAD_DIM] partial accumulator tile stored in TMEM.\n\nSo, to rescale these values: we need to read columns of TMEM holding `O`\n\n,\nthen rescale them (by the rescale factor we just computed from the old and new rowmaxes),\nthen write the rescaled values back to TMEM.\nWe need to write O back because this is still a running partial O, not the final output.\nSo that later PV matmul can accumulate its partial into the udpated basis.\n\n```\nfloat o8[8];\n\n#pragma unroll\nfor (int out_n8 = 0; out_n8 < HEAD_DIM / 8; ++out_n8) {\n    const int taddr = taddr_o + trow + out_n8 * 8;\n    tcgen05::ld_32x32b_x8(taddr, o8);\n\n#pragma unroll\n    for (int i = 0; i < 8; ++i) {\n        o8[i] *= rescale;\n    }\n\n    tcgen05::st_32x32b_x8(taddr, o8);\n}\n\n// wait until TMEM stores complete\nasm volatile(\"tcgen05.wait::st.sync.aligned;\\n\" ::: \"memory\");\n```\n\nThis uses the same 4 warps, row ownership, and 8-column chunks as the S traversal covered above. The new part is that after rescaling the values in registers, we use tcgen05.st to write the corrected O back into TMEM.\n\n**Back to the main flow: With the running rowmax updated and the old O corrected, the second S pass can now produce P in chunks. Continue\nin Chapter 1 Section B – Online Softmax.**\n\n### How Software exp2 Works\n\n*Optional background for Chapter 9 – Software Approximations.*\n\nAfter Kernel 8, our exponent input is already in base-2 units, and we compute exp2(x).\n\nMathematically, exp2(x) = 2^x.\n\nSo we want to compute some portion of 2^x with ALUs.\n\nWe split `x`\n\ninto its integer part and fractional part,\nraise `2`\n\nto each part separately, and then multiply the two results.\nMathematically `2^(a+b) = 2^a * 2^b`\n\n, so it’s equivalent to computing `2^x`\n\ndirectly.\n\n```\ninteger_part = floor(x)\n// always between zero and one\n// by our definition above\nfractional_part = x - integer_part\n\n2^x = 2^integer_part * 2^fractional_part\n```\n\nFor example:\n\n```\nx                   = -3.25\ninteger_part        = -4\nfractional_part     = 0.75\n```\n\nThat split on the integer_part is convenient because FP32 already has a base-2 exponent\nfield. So, we can apply `2^integer_part`\n\nby adjusting the exponent bits of the polynomial result.\n\nBut we cannot compute `2^fractional_part`\n\nas conveniently as\nthe `2^integer_part`\n\nbecause FP32’s exponent field can move only in whole steps\n(we can move it by -4, but not by 0.75 of a step).\nSo we use a different method to raise 2 to the fractional_part: approximate it with a polynomial.\nBecause `fractional_part`\n\nis always between zero and one (as by our definition),\nwe only need to approximate one small section of the 2^x curve where x is in `[0, 1)`\n\n.\n\nSo the exponent bits handle `2^integer_part`\n\n, while the\npolynomial handles `2^fractional_part`\n\n.\n\nWe approximate that curve with a cubic polynomial:\n\n```\n2^fractional_part\n    ~= 1\n    + c1 * fractional_part\n    + c2 * fractional_part^2\n    + c3 * fractional_part^3\n```\n\nWe can think of this polynomial as weighted blend of progressively more curved shapes. A line is too simple, while a quadratic is still not accurate enough. A cubic is flexible enough while remaining cheap to evaluate:\n\n```\n1                    constant, flat component\nfractional_part      straight-line component\nfractional_part^2    quadratic curve\nfractional_part^3    cubic curve\n```\n\nThe coefficients c1, c2, and c3 control how strongly each component contributes. They are not simply sampled values from the exponential curve. They were selected ofline to fit that curve (so that the weighted curve stays close to 2^fractional_part, on the interval between zero and one).\n\nA small optimization we use here.\nNaive implementation would first compute\n`fractional_part**2`\n\nand `fractional_part**3`\n\n,\nthen multiply each term by its coefficient and add everything together.\nBut explicitly constructing those powers costs two extra multiplications.\nInstead, we keep feeding the result of one multiply-add into the\nnext one. So the whole cubic takes only 3 FMA instructions.\n\n```\n2^f ~= ((c3*f + c2)*f + c1)*f + 1\n```\n\n[Back to Chapter 9 – Software Approximations](#chapter-9--software-approximations)\n\n# Acknowledgments\n\nThanks to the\n[FlashAttention-4 authors](https://arxiv.org/abs/2603.05451), their paper and\n[open-source CuTe DSL kernel](https://github.com/Dao-AILab/flash-attention/blob/fa4-v4.0.0.beta4/flash_attn/cute/flash_fwd_sm100.py).\nThanks to CUTLASS and CuTe DSL teams, and everyone who\ncontributed to\n[CUTLASS Blackwell FMHA example](https://github.com/NVIDIA/cutlass/tree/v4.4.1/examples/77_blackwell_fmha)\nand\n[CuTe DSL Blackwell FMHA example](https://github.com/NVIDIA/cutlass/blob/v4.4.1/examples/python/CuTeDSL/blackwell/fmha.py).\n\nThese blogs directly inspired me to try similar format but for a different kernel:\n\nThien Tran: your\n\n[Matmul blog](https://gau-nernst.github.io/tcgen05/)is truly amazing! Highly recommend to everyone.Modular team:\n\n[Matrix Multiplication on Blackwell B200](https://www.modular.com/matrix-multiplication-on-blackwell)Aleksa Gordić\n\n[Anatomy of high performance Matmul kernels](https://www.aleksagordic.com/blog/matmul)Mark Saroufim: thank you for building GPU MODE and its community.\n\nThese are some of the people who inspired me to do DL and ML sys for the last 7 years.\n\nPhilippe Tillet: our chat about\n\n[my earlier project](https://github.com/IaroslavElistratov/triton-autodiff)inspired me a lot. I’m deeply grateful.Edward Yang: you made me love ML sys. Few years ago, I found your\n\n[PyTorch internals](https://blog.ezyang.com/2019/05/pytorch-internals/)blog and liked it oh so much, and the[torch dev podcast](https://pytorch-dev-podcast.simplecast.com/episodes)I’ve been listening for like 4 years.Horace He: I appreciate our chat and your honest feedback about my earlier project; I enjoy your work\n\n[Thonk From First Principles](https://www.thonking.ai/),[Building ML Sys for a Trillion Trillion FLOPs](https://youtu.be/139UPjoq7Kw). I hope we get to see more blogs from you someday, I’m sure it will be a banger.Aleksa Gordić: I still remember your GNN project and the code you open sourced… it’s been like 6 years … oh time flies. Your matmul blog I cite above is also great.\n\nAlexander Amini: Your\n\n[MIT 6.S191 lectures](https://introtodeeplearning.com/)kick-started my programming career.Justin Johnson: your Michigan lectures are phenomenal.\n\nYannic Kilcher: 7 years ago, your videos made me love DL.\n\nWen-mei Hwu and coauthors, PMPP book: GPU programming starts here.\n\nPyTorch Team: thank you for releasing Composability Syncs, I’ve enjoyed them over the years.\n\nDo not watch random “learn CUDA” videos on YouTube. Either (1) watch the first 5 videos in\n\n[this playlist](https://www.youtube.com/playlist?list=PLRRuQYjFhpmvu5ODQoY2l7D0ADgWEcYAX). Don’t worry that these videos are old, they teach GPU-programming fundamentals that still apply today. Or (2) read the first few chapters of the Programming Massively Parallel Processors (PMPP) book.[↩︎](#fnref:1)Hopper added TMA (separate hardware unit) which is more restricted than the ldgsts – which in turn, can be thought of as generalized gather instructions, ie each thread in a warp can compute a different offset in to the memory, which for each offset takes registers to hold the offset calculations and additional instructions (to load each thread’s element independently) – so they made a more restricted alternative to ldgsts. This saves the per-thread address-calculation instructions and the registers needed to hold these addresses. MMA already supports only a limited set of structured SMEM layouts, so giving up arbitrary per-thread addressing is not much of a restriction here.\n\n[↩︎](#fnref:2)These are “unnormalized” probabilities, because in the online softmax formualtion the final division by the sum of the exponentiated values happens later, after the kv-for-loop, ie after traverals of the KV tiles. So untill we divided by the sum of the exponentiated values, this is not technically probabilities (P), but rather unnormalized probabilities. Sometimes for brevity I’ll refer to it as P, just keep in mind this is the unnormalized P. This is relatively a minor detail, not worth focusing on at the moment.\n\n[↩︎](#fnref:3)Because P re-uses S TMEM buffer, the next QK cannot reuse that S/P slot until PV consumes P. P is produced from S, so each S region becomes dead soon after. The rest of our lineage still reaches near-SOTA perf while keeping this S/P coupling. By contrast, K is already dead after QK, while P must remain alive until PV, which is a much longer time window.\n\n[↩︎](#fnref:4)We could load the 16 FP32 scores with one tcgen05.ld.x16 and still write the packed P with one tcgen05.st.x8. I use two ld.x8 calls only to reuse the existing width-8 load helper.\n\n[↩︎](#fnref:5)The reason why we store P in the lower half of S and not in the upper half is not accidental: if we were to instead write P into upper half of S then we’d be overwriting unread S values, unless we change the loop over TMEM to traverse from right to left.\n\n[↩︎](#fnref:6)NVIDIA docs tells us to expand 128b along the leading dim (based on the selected majorness). But for the\n\n`(8,8)`\n\natom shapes which the table specifies, they don’t say which of the dims, in their convention, is M/N and which one is K. So this next part is my inference from the figures. They draw M/N vertically and K horizontally. So, I read their atom shapes as (M/N, K).[↩︎](#fnref:7)Note I call them “row warps” and not softmax warps, because they don’t only compute online-softmax but also do O rescaling.\n\n[↩︎](#fnref:8)Incidentally, that’s one of the main reasons why we use 2Q stages and not, for example, 3Q stages. Hypothetically, it seems if we let one CTA process even more Q tiles, the K/V tile SMEM loads would be amortized even more, but we’re already at TMEM limit so we have no TMEM to directly store S/P/O for the 3rd Q-tile. I did try workarounds like trying to reduce S lifetime as much as possible – by for example writing S to SMEM right after matmuls are done, and producing its P in SMEM also, thus letting the subsequent PV consume it from SMEM directly – this way the TMEM would be free for longer, potentially allowing more Q stages. But juggling data in and out of TMEM caused more overhead than the benefit of the K/V re-use introduced by 3Q tiles.\n\n[↩︎](#fnref:9)that is, there isn’t enough TMEM for processing 2Q tiles of work for kv-loop iteration, and simultaneously processing some matmuls for iteration i+1, because these matmuls will also need to reserve TMEM and we’re already at capacity (we’re using all 512 TMEM columns)\n\n[↩︎](#fnref:10)This is somewhat stricter than needed, K is dead after Q1K so we could re-use without waiting for P1V. But in kernel 5, the producer (load warps) still waits until P1V(i) because this kernel still recycles the K and V storage as one pair, and V remains live until that final PV. We will improve this in the later kernels.\n\n[↩︎](#fnref:11)Three stages of complete K/V pairs would require six tile-sized KV slots. Together with Q0 and Q1, that would require 256 KiB of SMEM, which exceeds SM’ SMEM 227 KiB budget.\n\n[↩︎](#fnref:12)Kernel 6 already has a multi-stage K/V load pipeline, so K(i+1) must be resident in another SMEM buffer, while V(i) remains available for P1V(i).\n\n[↩︎](#fnref:13)the instruction is\n\n`ex2.approx.ftz.f32`\n\n, FTZ means turn values smaller than FP32’s regular full-precision range directly into zero. Which is the additional source of approximations (in addition to hardware computing ex2.approx itself approximately). Every flushed weight is below 2^-126.[↩︎](#fnref:14)I guess I got so used to seeing rebasing after each tile in online softmax kernels, that it didn’t occur to me to question this. Untill of course, I saw FA4 source.\n\n[↩︎](#fnref:15)There are two distinct 96/32 choices here. Kernel 10 uses 96/32 for S caching: cache high96 in regs, reread low32 from TMEM. Kernel 13 uses 96/32 for P publication: high96, start PV, then publish low32. These optimization are standalone implementaions, that don’t fundamentally require each other. Early PV itself doesn’t require the 96/32 S-caching. Mathematically, we could have published P after 16, 32, 48, 64, or 80 columns. This implementation chooses 96/32, but could have been other ratio. I tried other early_PV splits like 64/64 and others, they performed worse.\n\n[↩︎](#fnref:16)Apart the small rounding differences caused of floating point non-associativity.\n\n[↩︎](#fnref:17)", "url": "https://wpnews.pro/news/b200-attention-kernel-from-scratch-to-near-sota-in-60-diagrams", "canonical_source": "https://iaroslavelistratov.github.io/b200-attention/", "published_at": "2026-09-02 12:30:27+00:00", "updated_at": "2026-09-02 12:53:04.003053+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-research", "ai-infrastructure", "developer-tools"], "entities": ["Iaroslav Elistratov", "B200", "CUDA", "PTX", "FlashAttention-4", "Blackwell"], "alternates": {"html": "https://wpnews.pro/news/b200-attention-kernel-from-scratch-to-near-sota-in-60-diagrams", "markdown": "https://wpnews.pro/news/b200-attention-kernel-from-scratch-to-near-sota-in-60-diagrams.md", "text": "https://wpnews.pro/news/b200-attention-kernel-from-scratch-to-near-sota-in-60-diagrams.txt", "jsonld": "https://wpnews.pro/news/b200-attention-kernel-from-scratch-to-near-sota-in-60-diagrams.jsonld"}}