Flash attention in PyTorch, then in Triton

Intro

0:00 Hello and welcome back to episode thirteen, where we're going to write flash attention, which computes attention without keeping its table of scores, first in PyTorch on the Mac, and then as a kernel on an NVIDIA GPU. Let's begin with a quick look at where we're going to be by the end of the episode.

0:16 By the end of today, we're going to have a flash attention of our own, a kernel written in Triton, which runs about as fast as the one inside PyTorch on an RTX 4090. At eight thousand tokens, ours does about one hundred and forty-six trillion operations a second, against about one hundred and forty-nine for PyTorch's, while attention written out by hand manages about five point six, and runs out of memory past that length.

0:40 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on flash attention today, with all of today's measurements. And each episode has a page here as well, with its chapters and its code.

Attention's table of scores

1:01 Now, before we get to today's work, let's quickly go over what we did last time. We rented four GPUs by the hour and trained a model about six times bigger than the one on my laptop, in the shape of GPT-2 small, with a context of five hundred and twelve tokens. Each GPU trained a copy of its own, and the copies averaged their gradients after each step. The models you use day to day read far longer contexts than that, tens of thousands of tokens and more, and at those lengths, attention becomes the expensive part. It scores each token against the tokens before it, so its cost grows with the square of the length. So today, we're going to measure that cost, and write flash attention, which computes attention without keeping the table of scores in memory. All right, let's get started with attention's table of scores, as the paper writes it.

1:49 In each head, the scores are the queries times the keys transposed, divided by the square root of the head's sixty-four channels, which gives each query's score against each of the keys. For a story of T tokens, that's a table of T by T scores per head, about four million at two thousand tokens, and about sixty-seven million at eight thousand. Then the mask adds minus infinity above the diagonal, so each token sees the tokens up to its own, the softmax turns each row into probabilities, and those weigh the values. In training, each of the model's blocks keeps its tables for the backward pass, so the memory they take grows with the square of the length.

Attention by hand

2:27 We start with a new file, attention dot py, which imports math, for the square root, and torch. Attention by hand takes q, k and v, shaped batch, heads, tokens and sixty-four channels, which is how our model's attention hands them to PyTorch. T is the number of tokens, and over the function, we write down what it computes. It's causal attention as the paper writes it, each query's score against each key, a table of tokens by tokens per head, masked so that each token sees the tokens up to its own, then the softmax of each row, and the values it weighs. The scores are q times k transposed, which swaps its last two dimensions, divided by the square root of sixty-four. The mask is a T by T table of booleans, all ones, where tril keeps the lower triangle, the diagonal and below it. Masked fill puts minus infinity where the mask is false, as e to the minus infinity is zero, so those tokens get no weight in the softmax. And the softmax of each row, times v, weighs the values by the probabilities. Then a test, in tests slash test attention dot py, against PyTorch's scaled dot product attention, which our model calls. Devices is the CPU, along with the Mac's GPU and CUDA where they're available, and q k v makes three tensors of random numbers, two sequences of three heads, from a fixed seed. The test runs on each device for one token, a hundred, two hundred and fifty-six and a thousand, and the largest difference from PyTorch's causal attention has to be under one in a hundred thousand, which leaves room for float32's rounding and no more. And to measure it, attention bench dot py, which imports argparse, json, os, resource, statistics, for the median, subprocess, sys and time. Resource is Python's view of what the operating system counts for a process, and we're going to read the peak memory from it. Then torch, DeviceType from autograd, functional, and from PyTorch's profiler, which we used in episode six, its activities and profile. From our own files, attention by hand, and sync from bench dot py, from episode four, which waits for the GPU to finish before we read the clock. Pytorch is PyTorch's causal attention, called the way the model calls it. Methods maps each name to its function, by hand and pytorch, and dtypes maps the precisions we can ask for, float32, bfloat16 and float16. Inputs makes q, k and v of the shape we ask for, from a fixed seed, so each method gets one set of numbers. Milliseconds starts with a first run, to warm up, because the first run on a GPU pays for setting things up, which isn't what we want to time. Then it times each run between two syncs, and returns the median, which a single slow run can't move much. Peak memory measures the most memory attention takes on top of q, k and v. On CUDA, PyTorch's allocator counts each tensor it hands out, and keeps the largest total it reached. So we reset that peak, note what's allocated before, run attention, and the new peak minus that is what attention took, and a comment at the top gives the unit, in bytes. On the CPU, PyTorch keeps no such count, so we ask the operating system instead. The comment over it says how: a separate process runs attention, and we take the most memory that process held, less what it held before. That number belongs to the whole process, and it can go up but not down, so each measurement needs a fresh process, or the first method's peak would hide the next one's. Subprocess runs this script again, with the device, a hidden peak flag, the method and the length, and reads the number it prints. And the Mac's GPU doesn't say how much memory it held at most, so there it returns none. Max r s s is the process's maximum resident set size, from get r usage, which counts the memory the process holds in RAM, and it's the most it has held so far. Linux gives it in kilobytes, and macOS in bytes. What pytorch runs finds out which attention PyTorch picks on the device. PyTorch has several, and it chooses one from the device, the precision and the inputs. It runs PyTorch's attention once under the profiler, on the CPU and on CUDA too, and keeps the operations named after scaled dot product attention, along with the names of the kernels the GPU ran, cut before their template arguments. Flops counts the arithmetic, for the speed in trillions of operations a second. Attention is two products of a T by T table and the sixty-four channels, two operations each, a multiply and an add, with half of the table masked away. In main, argparse reads the flags, where the device is CUDA if there is one, and the Mac's GPU otherwise. The lengths double from a thousand and twenty-four tokens up to eight thousand one hundred and ninety-two, so that each step shows the table's square, four times the memory for twice the length. The methods default to all of them, and we measure one sequence of twelve heads, like our bigger model's, in float32, with five runs each. Out is a JSON file to write the measurements to, and peak is hidden from the help with argparse dot suppress, as it's for the script's own use. With peak, the script is running as that separate process, started to measure peak memory, so it makes q, k and v, notes the peak memory so far, runs the method once, and prints how much the peak went up. Otherwise it prints the shapes and the device, and the attention PyTorch runs there. For each length and each method, it makes q, k and v, times the method, and deletes them before measuring the memory, so the timing's tensors don't count in it. If the GPU runs out of memory, the error is caught and printed, and empty cache gives what PyTorch's allocator kept back to the GPU, so the next method starts with all of it. Otherwise the speed is the flops over the time, in trillions of operations a second, printed on CUDA, and each row goes in a list. And out writes the rows as JSON, for the charts.

8:21 Let's run the test with pytest, where dash v lists each case, and attention by hand matches PyTorch's on the CPU and on the Mac's GPU, at all four lengths.

8:31 Now the measurements, on the CPU, for one sequence of twelve heads in float32, with dash dash out to keep them, and two greater-than slash dev slash null to hide the profiler's log lines on macOS. The line under the shapes is PyTorch's choice here, scaled dot product flash attention for CPU, so on the CPU, PyTorch's attention is already a flash attention, the kind we're going to write today. By hand, the memory goes from about one hundred and three megabytes at a thousand tokens to about sixty-three hundred at eight thousand, about four times as much each time the length doubles, which is the table, growing with the square of the length. PyTorch's grows far more slowly, from about nine megabytes to about thirty-two, with the length, like q, k and v. It's faster too, at about sixty-nine milliseconds against about three hundred and eighty by hand, at the longest length.

9:24 On the Mac's GPU, PyTorch's choice is scaled dot product attention math for MPS, which is attention as the paper writes it, table and all, as a single operation. It still runs about six times faster than ours by hand, at about twenty-seven milliseconds against about one hundred and fifty-eight at eight thousand tokens. I'd expect that's because a single operation makes fewer trips to memory, where ours runs four, each reading and writing the whole table. The Mac's GPU doesn't report its peak memory, which is why we measured the memory on the CPU.

Flash attention

9:54 The paper is FlashAttention, from 2022, by Tri Dao and others at Stanford and the University at Buffalo. Its abstract starts from what we've just measured: transformers are slow and need a lot of memory on long sequences, as attention's time and memory grow with the square of the length. It argues that attention should be IO-aware, which means counting the trips between the GPU's large main memory, HBM, and its small, fast memory on the chip, SRAM. Figure one shows the idea: the large N by N table, in the dotted box, stays out of main memory, as blocks of K and V, and of Q, are loaded into the fast memory, and the output goes back. On the right, that made the attention in GPT-2 about seven point six times faster than PyTorch's at the time.

10:38 A year later, Tri Dao's FlashAttention-2 split the work better across the GPU, with a program for each block of queries, which is how we're going to write ours, and it's the version built into PyTorch. Its figure one shows the forward pass with the keys and values in two blocks. Attention over each block, with the output rescaled when the second block comes in, ends with the correct output, without the table of scores, S, or of probabilities, P, going to main memory.

11:06 The softmax of a row subtracts the row's largest score before taking the exponentials, so they don't overflow, and divides by their sum, which seems to need the whole row at once. The online softmax, from a paper by Maxim Milakov and Natalia Gimelshein at NVIDIA in 2018, keeps a running maximum and a running sum for each row instead, updated as each block of keys comes in, and flash attention adds a third running value, the output so far. The new maximum, m prime, is the larger of the old one and this block's largest score. L is the sum of e to each score minus m, over the keys seen so far. That sum was computed against the old maximum, so it's multiplied by e to the old minus the new. That turns each of its terms into e to the score minus the new maximum, as if that maximum had been there from the start, and then this block's own terms are added. The values, weighed by those exponentials, are summed the way the probabilities are, with that correction too. And after the last block, dividing by l gives each row's output, the softmax's weighted sum of the values. There's no approximation in any of it, so the result matches the softmax of the whole row, to rounding.

The online softmax, in PyTorch

12:18 Back in attention dot py, flash attention pytorch takes q, k and v, and a block size of sixty-four. T is the number of tokens, and above it, we write what the function does: attention without the table, FlashAttention's algorithm, a block of queries against a block of keys at a time. The keys go past each block of queries in blocks, and a running maximum and a running sum make the softmax come out as if all the scores had been there at once. The scale is one over the square root of sixty-four, and out is an empty tensor like q, for the output. For each block of queries, i goes from zero up to T in steps of sixty-four. Q i is the block's queries, times the scale, so we scale sixty-four queries once, instead of each block of scores, and rows are their positions in the story. M starts at minus infinity, and its comment calls it the largest score so far, per row. L starts at zero, the sum of e to the score minus m so far, and acc at zeros, the values, weighed by e to the score minus m, so far. Over the loop, we note what causal attention means here: the queries in this block see no key past its last row. So j, the first key of each block of keys, goes up to the end of this block of queries, and the blocks after it are skipped, which is about half of them. K j and v j are the block's keys and values. S is the block of scores, q i times k j transposed, in float32, and cols are the keys' positions, so masked fill puts minus infinity where a key comes after its query. M new is the larger of m and each row's largest score in this block. The correction is e to m minus m new, and above it, we write why: a new maximum shrinks what was summed so far by e to the old minus the new, as if it had been there all along. P is e to each score minus the new maximum. L is the old sum, corrected, plus this block's, and acc is the old weighted values, corrected, plus p times this block's values, and m becomes the new maximum. After the last block of keys, this block's output is acc divided by l, in q's precision, and once all the blocks of queries are done, it returns out. The test gets a second case, a copy of the first with flash attention pytorch in place of attention by hand, against PyTorch's attention on each device and at the four lengths, held to that one in a hundred thousand. And attention bench dot py imports it, and adds it to the methods as flash pytorch, so the benchmark measures it next to the other two.

14:50 Both tests pass on the CPU and on the Mac's GPU, so computing the softmax a block at a time gives PyTorch's answer, to float32's rounding.

15:00 On the CPU, our flash attention in PyTorch holds about thirty megabytes at eight thousand tokens, close to PyTorch's thirty-two, against about sixty-three hundred by hand. But it takes about two point eight seconds there, about eight times longer than by hand.

15:18 On a chart, by hand curves up with the square of the length, while ours stays along the bottom, as PyTorch's flash attention for CPUs does, growing with the length alone, and at eight thousand tokens, ours holds about two hundred and ten times less than attention by hand.

15:33 On the Mac's GPU, our flash attention in PyTorch is further behind, at about eleven hundred milliseconds for eight thousand tokens, against about one hundred and seventy-one by hand and twenty-seven for PyTorch's. Each pair of blocks is a handful of small operations, launched one after the other from Python, and at that length there are thousands of pairs, so the GPU spends most of its time waiting for the next one. The algorithm saves the memory, and to make it fast as well, the whole loop has to run on the GPU, as a kernel.

On an NVIDIA GPU

16:04 To write that kernel, I've rented a machine with an NVIDIA GPU, an RTX 4090, the way we did last time, for under fifty cents an hour, and given it the name gpu for ssh. nvidia smi shows the GPU, idle, with twenty-four gigabytes of memory, which is the limit attention by hand is going to run into.

16:24 We copy the code there with git archive, which packs the files of the commit, piped through ssh and unpacked into a folder. Then scp copies the BPE merges and the stories, as the token ids our models on the Mac trained on. And on the machine, pip installs NumPy, TensorBoard, pybind11 and pytest, and setup dot py builds our C++ tokenizer from episode eleven. PyTorch comes with the machine's image, along with Triton, which PyTorch uses for torch compile on NVIDIA GPUs.

16:55 A GPU like this one has a large main memory, the twenty-four gigabytes, which it reads at about a terabyte a second. Its arithmetic happens in a hundred and twenty-eight streaming multiprocessors, and each has a small memory of its own on the chip, a hundred and twenty-eight kilobytes, which is a lot faster to read and write. Attention by hand writes the table of scores to main memory, reads it back to mask it, writes it again, reads it for the softmax, and so on, and at long lengths, that traffic takes far longer than the arithmetic. Flash attention loads a block of queries and a block of keys into the fast memory, does all of its work on them there, and writes back the output and no more.

17:34 First, by hand against PyTorch's, on the 4090, in bfloat16, the precision our bigger model trained in, with four sequences at a time and lengths up to sixteen thousand tokens. Here, PyTorch picks scaled dot product flash attention, which launches pytorch flash, flash forward kernel, the kernel from FlashAttention-2, built into PyTorch. So in the last episode, as our bigger model trained in bfloat16 on four 4090s, PyTorch would have picked this kernel for its attention too. By hand gets about five point six trillion operations a second, and holds about twelve point one gigabytes at eight thousand tokens, and at sixteen thousand, it runs out of memory. PyTorch's gets about one hundred and fifty-seven trillion, in about ninety-nine megabytes.

A kernel in Triton

18:18 Our kernel goes in a new file, flash dot py. Triton, from OpenAI, is a language for GPU kernels, written in Python. A kernel says what a single program does with blocks of numbers, and Triton's compiler turns it into code for the GPU, splitting each block among the program's threads. It runs on Linux, with NVIDIA or AMD GPUs, so not on the Mac, which is why we're on the 4090. The file imports math, torch, triton, and triton dot language as tl, the operations a kernel can use. Log two e is a constant the compiler knows, a tl dot const expr, and above it, we write why we need it. To get e to the x, we can multiply x by log two of e, and raise two to that. A GPU raises two to a power faster than e, so the kernel keeps its scores times log two e. Triton dot jit compiles the function under it for the GPU, the first time it's launched. Our forward kernel takes Q, K, V and O, which arrive as pointers to the first number of each tensor, then the scale and T, and three const exprs, the sixty-four channels and the two block sizes, which the compiler needs to lay out the blocks. Start is the first query of this program's block, as program id zero is which block of queries the program has. Head is program id one, which head of which sequence, as a sixty-four-bit integer, so that offsets into large tensors don't overflow. Above them, we describe how the work is split: a program for each block of block M queries of one head of one sequence, with thousands running at once, and q, k, v and o are sequences times heads, tokens, D, one after the other in memory. Rows are the block's query positions, start plus tl dot arange, the numbers from zero up to block M, and dims are the sixty-four channels. The queries come in with tl dot load, from a block of addresses: Q, plus the head's offset and each row's, times sixty-four, plus each channel. The mask leaves out rows past T, which the last block can have, and loads zeros for them. Its comment notes that the block of queries is read once from the GPU's memory into the program's own, much faster one. M, l and acc start like in PyTorch, in float32, and their comments describe them with two to the score in place of e to the score. The loop goes through the keys up to this block's last query, a block at a time, as causal attention needs, and tl dot minimum stops it at T. Cols are the keys' positions, and k and v come in like the queries, block N rows of sixty-four channels, with the mask leaving out keys past T, which the last block of keys can have. Both blocks now sit in the program's fast memory, next to the queries from before, and k is transposed where it's used. S is tl dot dot of q and k transposed, a product that runs on the GPU's tensor cores, times the scale and log two e. Tl dot where keeps a score where the query comes at or after the key, and puts minus infinity elsewhere. Then the online softmax, in powers of two: m new, the correction, p, l and acc, where p goes into its product in v's precision, so the tensor cores can take it. After the loop, tl dot store writes acc divided by l, in O's precision, to the output's rows, masked like the load. The scores and the probabilities stayed in the program's fast memory, and the output is what goes back to main memory. Block M and block N are a hundred and twenty-eight and sixty-four, a hundred and twenty-eight queries against sixty-four keys at a time. Flash attention is what the rest of our code calls, and its comment sums it up: causal attention with our kernel, a program for each block of a hundred and twenty-eight queries of each head of each sequence, going through the keys sixty-four at a time. It makes q, k and v contiguous, as our model's are views of a transposed tensor, while the kernel's offsets expect each head's tokens one after the other. O is an empty output like q, and the grid is how many programs to run, along two axes: T over a hundred and twenty-eight, rounded up, with triton dot cdiv, and sequences times heads, which are the two program ids the kernel reads. Indexing our forward kernel with the grid launches it, with num warps four, four groups of thirty-two threads to a program, and num stages two, so that Triton loads the next block of keys while the program works on this one. The test, in tests slash test flash dot py, gets triton with pytest's import or skip, so on the Mac, where Triton isn't installed, the file is skipped with a reason. It runs on CUDA, or slowly, on the CPU, with Triton's interpreter, which runs a kernel's programs one after the other in Python, and which is how I checked the kernels before renting the GPU. The precisions are bfloat16 and float16 on CUDA, the two that the tensor cores take, and the lengths are one, seventy-seven, a hundred and twenty-eight, three hundred and a thousand and twenty-four, most of which leave a partial block at the end, so the masks get tested too. Pytorch is PyTorch's causal attention, and the test compares both with PyTorch's attention in float64, as near to the true answer as we can get. Over it, we write the rule, which is the test FlashAttention's authors use: ours may be off by rounding, but by no more than twice as much as PyTorch's own attention is, in that precision. So it computes PyTorch's error and ours, and asserts that ours is at most twice PyTorch's, with a tiny margin. And attention bench dot py adds our kernel as a method, triton, when CUDA is available, as flash dot py needs Triton.

23:59 On the 4090, all ten cases pass, so our kernel is as accurate as PyTorch's attention, within that factor of two, in both precisions and with partial blocks.

24:10 Now our kernel against PyTorch's attention, forwards, with by hand alongside. At sixteen thousand tokens, ours takes about ten point six milliseconds, and PyTorch's about ten point five, which is about one hundred and fifty-six trillion operations a second against one hundred and fifty-seven, and both hold about the size of the output, about ninety-six megabytes. At a thousand tokens, ours is further behind, about forty-seven against fifty-eight, where the work is small enough that launching it counts.

24:40 As the stories get longer, ours gets closer to PyTorch's, from about eighty-one percent of its speed at a thousand tokens to about ninety-nine percent at sixteen thousand, as both climb towards what the GPU's tensor cores can do. By hand, meanwhile, stays at about five point six trillion operations a second at each length that fits in memory, as its time goes into moving the table.

Backwards, and training with it

25:02 To train with it, we need the gradients of q, k and v, how the loss changes with each, given how it changes with the output, d O. Autograd gets them for attention by hand from the tables it kept, which is the memory we're trying to save. Instead, the forward pass keeps each row's log-sum-exp, the log of the sum the softmax divides by, and the backward pass makes the probabilities again from the scores, a block at a time, as P is e to each score minus it. The gradient of the values is then P transposed times d O, and of the probabilities, d O times V transposed. The softmax's backward turns that into the gradient of the scores, P times d P minus delta, where delta is each row's sum of d O times O, computed once. And as the scores are q times k transposed, scaled, d Q is d S times K, and d K is d S transposed times Q, each scaled. Each of these is a product of blocks, so the backward pass can work a block at a time, like the forward pass.

26:02 In flash dot py, the forward kernel takes a pointer more, LSE, and after the loop, it stores each row's log-sum-exp, which in our powers of two is m plus log two of l. The comment over it says what that's for: the backward pass makes the probabilities again from the scores. The first backward kernel, backward kv, has a program for each block of keys and values, with the queries that see them going past in blocks. It takes q, k, v, d O, the log-sum-exps and delta, and pointers for d K and d V, with one block size, sixty-four, for the queries and the keys, as a program now holds the gradients along with the inputs. Start is this program's first key, and head is which head of which sequence, like in the forward kernel. Cols are this program's keys, which it loads once, along with their values, with the masks for the last block, and they stay in the program's fast memory while the queries go past. D k and d v start at zero, in float32, as they add up a little from each block of queries, and in bfloat16, those small additions would get lost to rounding. With causal attention, these keys are seen by the queries from the block's first key on, so i starts at the block and goes up to T. For each block of queries, it loads q and d O, the gradient of the loss with respect to those queries' outputs, a row for each query. It loads the log-sum-exps and delta too, a number for each, all masked for the last block, so the program has what it needs for this block of the table. The probabilities are made again from the scores and the saved log-sum-exp, two to the scores minus it, and the comment notes that this is cheaper than storing them. Where the query comes before the key, or past T, p is zero, and d v adds p transposed times d O. Then comes the softmax's backward, d S equals P times d P minus the row's sum of d O times O, with d P as d O times v transposed, and d k adds d S transposed times q. At the end, it stores d k, times the scale, and d v. The second kernel, backward q, has a program for each block of queries, with the keys they see going past in blocks. The work is split into two kernels so that each program writes its own rows of the gradients, d K and d V in the first, d Q in the second. So no two programs write to one place, and none has to wait for another or add into another's numbers. It loads its queries, d O, the log-sum-exps and delta once, each with the mask for the last block, and d q starts at zero, in float32, like d k and d v in the first kernel. Then it goes through the keys up to its last query, like the forward pass, loading each block of keys and values. For each, it makes p again from the scores and the log-sum-exps, masks it where the key comes after the query, and computes d S from d O times v transposed, as before, and d q adds d S times k. Like the scores in the forward pass, p and d S for each block stay in the program's fast memory, and they're gone once the next block of keys comes in. At the end, it stores d q, times the scale, as the scores were scaled, masked for the rows past T. Both kernels compute the scores and d P, so that work is done twice. FlashAttention-2 does it once, in a single kernel, where the programs add into d Q together with atomic additions, which the GPU keeps correct when several programs add into one place. Those are harder to get right, so we keep our two kernels, and we're going to measure what that costs. FlashAttention is a torch autograd Function, which pairs a forward and a backward of our own, so that autograd calls our kernels for the gradients. Its comment gives the block sizes: a hundred and twenty-eight queries against sixty-four keys at a time forwards, and sixty-four against sixty-four backwards, with more to hold. Forward takes ctx, the context autograd gives the Function to pass things from the forward pass to the backward one, and q, k and v. Its launch is the one from before, along with the log-sum-exps, a float32 for each row of each head, and save for backward keeps q, k, v, the output and the log-sum-exps for later, which is a lot less than the table would be. Backward gets ctx and d O, the gradient of the output. It takes back what forward saved, reads the shapes from q, makes d O contiguous, and computes delta, each row's sum of d O times O, in float32, with PyTorch, as it's a single pass over two tensors. The gradients start empty, shaped like q, k and v, and the grid is a program for each block of sixty-four, for each head of each sequence, with the arguments both launches share in a dictionary. Both kernels launch with four warps and three stages, the fastest of a few settings I tried on the 4090, and backward returns the three gradients. Flash attention now calls the Function's apply, which runs forward and records it for the backward pass. If autocast is on for the device, get autocast dtype gives its precision, bfloat16 in our training, and q, k and v are cast to it. The function's comment adds the reason: under autocast, it works in autocast's precision, as PyTorch's attention would, because rotary embeddings leave q and k in float32. Their angles are float32, so q times a cosine is too, while v comes from its linear layer in bfloat16. The test now checks the gradients too, so q k v makes a fourth tensor, a gradient for the output, and all four require gradients. Attend runs a function and returns a list of its output and the gradients of q, k and v for that gradient, from torch dot autograd dot grad, which hands the gradients back without storing them in the tensors. The float64 answer comes from PyTorch's attention on float64 copies, detached from the originals and requiring gradients of their own, with the gradient in float64 too. Theirs is PyTorch's attention in the test's precision, and ours is our kernel's. Then, for each of the four, the output and the three gradients, the test computes their error and ours against float64, holds ours to twice theirs, and the assertion's message says which of the four failed. Attention bench dot py now measures the backward pass too, as training runs both passes. Inputs makes q, k and v require gradients when we ask for backward, so autograd tracks them, and returns a gradient too, a random tensor shaped like the output. Run once runs attention, and with backward, it computes the gradients of q, k and v for that random output gradient, as training would. Then it clears them, so they don't pile up from one run to the next. Milliseconds now goes through run once, and its comment adds that the first run is also when Triton compiles the kernels, which takes a few seconds and isn't what we want to time. Peak memory goes through run once too, and now counts the memory on top of the output's gradient as well, while on the CPU, the child process gets the backward flag. Under the profiler, no grad keeps PyTorch from preparing a backward pass we wouldn't run. And flops counts the backward too, which does five products where the forward does two, so both passes together come to three and a half times the forward's operations. In main, the backward flag joins the others, the line printed at the start says which passes run, and the JSON notes it too, so the charts can tell the two kinds of measurement apart. In the loop, the gradient comes along with q, k and v, into the timing, and it's deleted with them before the memory is measured. Then model dot py, where causal self-attention gets an attend attribute, none by default, which can hold another function for causal attention over whole sequences, such as flash dot py's kernel. In forward, dropout p is the model's dropout in training, and zero otherwise, as before. When attend is set, dropout p is zero and S equals T, which means a whole sequence with no cache in front, it calls attend with q, k and v. Otherwise it calls PyTorch's attention, with the mask, dropout and is causal, as before, under its comment. Our kernel has no dropout, and with a cache, the queries are a few new tokens against all the keys, which it doesn't handle. GPT gets attend with, which goes through the blocks and sets attend in each block's attention. Its comment says where that applies: in place of PyTorch's attention where it can, whole sequences without dropout, while none goes back to PyTorch's. And pretrain DDP dot py gets an attention flag, pytorch, by hand or triton. By hand comes from attention dot py, while our kernel is imported when it's asked for, as Triton runs on Linux with NVIDIA GPUs, and the first line the run prints says which attention it uses.

34:52 On the 4090, the ten cases pass with the gradients too, so our backward pass is as accurate as PyTorch's, within that factor of two.

35:02 Now forwards and backwards together, as in training, for all three. By hand now runs out of memory at eight thousand tokens, as the backward pass needs its tables too, and at four thousand it holds about six point one gigabytes. Ours holds less than PyTorch's, about six hundred and seventy-five megabytes at sixteen thousand tokens against about seven hundred and seventy-four, and it's slower, at about forty-nine milliseconds against about thirty-eight.

35:30 With the backward pass, ours gets about seventy-seven percent of PyTorch's speed at sixteen thousand tokens. I'd put most of that down to our two backward kernels each computing the scores, where FlashAttention-2 computes them once, in a single kernel.

Training with it

35:44 Then we pretrain our bigger model, the ninety-one million parameters of the last episode, on the 4090, with a context of two thousand and forty-eight tokens, four times last time's, and twelve stories a step, for five minutes with each attention, starting with PyTorch's. Torchrun, with a process for the one GPU, runs pretrain DDP dot py, and it trains at about eighty-five thousand tokens a second, and peaks at about sixteen point two gigabytes. Then five minutes with ours, with the first run's seed, so both runs start from identical weights and get the stories in the order the first run got them. Ours trains at about eighty-two thousand tokens a second, about four percent slower, which gives it seven hundred and eighty-eight steps to PyTorch's eight hundred and twenty-seven, in about the memory PyTorch's took.

36:31 The held-out losses stay close, one point eight four with PyTorch's attention and one point eight four with ours at step seven hundred and fifty. They differ a little, as the two attentions round differently in bfloat16. Attention is a small part of a step's time at this length, next to the model's linear layers, so a backward pass at about three quarters of PyTorch's speed costs the whole step about four percent, at about eighty-two thousand tokens a second against eighty-five thousand.

36:58 And by hand, with those settings, grep picks the error out of the log: the run stops before its first step, as CUDA runs out of memory trying to allocate about nine gigabytes. The allocation is for the first evaluation, which scores forty-eight stories at once, a table of forty-eight stories by twelve heads by two thousand and forty-eight tokens by two thousand and forty-eight, in float32.

37:20 Scp copies the measurements and the runs' TensorBoard logs back to the Mac, for the charts, and vastai destroy instance gives the machine back, after under an hour, which cost about forty cents.

37:35 Last, the docs get a page on flash attention, with the online softmax, the code, how the kernel is tested, which attention PyTorch runs on each device, and all of today's measurements, on the Mac and on the 4090. It goes in the list on the docs' home page, after rotary position embeddings, and at the end of the navigation, after training on several GPUs. And the README has the command to measure attention.

Wrap-up

38:01 So, today we measured attention's table of scores, and on the CPU, attention by hand held about sixty-three hundred megabytes at eight thousand tokens, four times as much each time the length doubled. We computed the softmax a block at a time instead, with a running maximum and sum, which held about thirty megabytes, and was slow, written in Python. As a Triton kernel on an RTX 4090, at long lengths, it ran within a few percent of the FlashAttention-2 inside PyTorch, which PyTorch picks for our model's training on NVIDIA GPUs. With its backward pass, at about three quarters of PyTorch's speed, our training ran about four percent slower, with the held-out losses close together, while attention by hand didn't fit in the GPU's memory at that length.

38:46 Let me know in the comments which parts of the model you'd like us to write as kernels next, and which parts of the pipeline you'd like us to go deeper on. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.