FlashAttention 2: making Transformers 800% faster w/o approximation - with Tri Dao of Together AI

**Alessio Fanelli** (0:10)
Hi, everyone. Welcome to the Latent Space Podcast. This is Alessio, Partner and CTO in Residence at Decibel Partners.
Today, we have Nuswix, because he's in Singapore, so it's a one-on-one discussion with Tri Dao. Welcome.

**Tri Dao** (0:24)
Hi, everyone.
I'm Tri Dao. Excited to be here.

**Alessio Fanelli** (0:27)
Tri just completed his PhD at Stanford a month ago. You might not remember his name, but he's one of the main authors in the FlashAttention paper, which is one of the seminal work in the Transformers era.
He's got a lot of interest from efficient Transformer training and inference, long-range sequence model, a lot of interesting stuff. And now you're going to be an assistant professor in CS at Princeton next year.

**Tri Dao** (0:51)
Yeah, that's right.

**Alessio Fanelli** (0:52)
And in the meantime, just to get a low-pressure thing, you're a chief scientist at Together as well, which is the company behind RedPajama.

**Tri Dao** (1:00)
Yeah. So I just joined this week, actually, and it's been really exciting.

**Alessio Fanelli** (1:04)
So what's something that is not on the internet that people should know about you?

**Tri Dao** (1:09)
Let's see.
When I started college, I was going to be an economist. So I was fully on board. I was going to major in economics. But the first week I was at Stanford undergrad. I took a few math classes and I immediately decided that I was going to be a math major.
And that kind of changed the course of my career. So now I'm doing math, computer science, AI research.

**Alessio Fanelli** (1:32)
At a similar thing, I started with physics.
And then I took a programming course. And I was like, I got to do computer science. I don't want to do physics.
So FlashAttention is definitely... Everybody is using this. Everybody loves it. You just released FlashAttention 2 last week.

**Tri Dao** (1:48)
Yeah, early this week on Monday.

**Alessio Fanelli** (1:49)
Yeah, AI time.

**Tri Dao** (1:51)
Things move fast.

**Alessio Fanelli** (1:52)
Yeah, four days ago is one week ago in AI.
So maybe let's run through some of the FlashAttention highlights, some of the innovation there. And then we can dive into FlashAttention 2 So, the core improvement in FlashAttention is that traditional attention is a quadratic sequence line. So it's M to the 2 FlashAttention is linear, which obviously helps with scaling some of these models.

**Tri Dao** (2:17)
There are two factors there. So, of course, the goal has been to make attention go faster or more memory-efficient. And ever since attention became popular in 2017 with the transformer paper, lots and lots of folks have been working on this. And a lot of approaches have been focusing on approximating attention. The goal is you want to scale to longer sequences. There are tons of applications where you want to do that. But scaling to longer sequences is difficult because attention scales quadratically in sequence length on both runtime and memory, as you mentioned.
So instead of trying to approximate attention, we were trying to figure out, can we do the same computation and maybe be more memory-efficient? So in the end, we ended up being, the memory is linear in sequence length.
In terms of computation, it's still quadratic, but we managed to make it much more hardware-friendly. And as a result, we do get wall clock speed up on the order of 2 to 4x, which really helps because that just means that you will be able to train with 2 to 4x longer sequence length for the same cost without doing any approximations. As a result, lots of folks have been using this. Because I think it's available in a lot of libraries that train, do language model training or fine-tuning.

**Alessio Fanelli** (3:32)
And the approximation thing is important because this is an exact thing versus like a sparse. So maybe explain a little bit the difference there.

**Tri Dao** (3:40)
For sure, for sure.
Attention, essentially you compute pairwise similarity between every single element in a sequence against each other. So there's been other approaches where, instead of doing all that kind of pairwise computation, you only compute similarity for some pairs of elements in the sequence. So you don't do kind of quadratic number of comparison. And this can be seen as some form of sparsity. Essentially, you're ignoring some of the elements. When you write down the matrix, you essentially say, okay, I'm going to pretend they're zero.
That has some benefits in terms of runtime and memory. But the trade-off is that it tends to do worse in terms of quality because you're essentially approximating or ignoring some elements. And I personally have worked on this as well for a few years. But when we talk to practitioners who actually train models, especially at large scale, they say, tend not to use these approximate attention methods.