hi everyone Welcome to our 54th session

of the main AI group exchange this week

we have three doubt from Stanford

University here with us to present his

work on Fresh attention fast and memory

efficient exact extension with IO

awareness

three is a PhD student in computer

science at Stanford co-advised by Chris

Ray and Stefano Ormond he works at the

intersection of machine learning and

systems his research interests include

sequence models with long-range memory

as well as structured matrices for

compact deep learning models his work

has received the ICM of 2022 outstanding

paper running up

thank you tree for joining us today

before we start could you let us know

how you want to take questions

yeah uh thank you see so I'm really

excited to be here uh so this is a

fairly informal talk I'm so happy to

take questions uh during during the talk

that's that's great

um and yeah thanks see and nadita too

for for organizing this I'm really

excited to be here okay so let's try to

make this session as interactive as

possible without further Ado let me help

me over to three

yeah

thanks

um yeah so we're I'm really excited to

hear uh talking about our recent work

um it's called flash attention and we

try to put all the main ideas in in the

title so uh what this is going to do is

um attention we're going to make it one

faster

two is going to be Memory efficient and

three is going to be exact so we don't

do any approximation and and the trick

is called i o awareness and I'll go into

a little bit of details

um so that although that high level

takeaway messages are in the title

all right so uh let me talk a little bit

more about the high level

um so what we're going to do is we're

going to use Iowa aware tiling

techniques to speed up attention and

attention is the the heart of

Transformers which is now the the

dominating uh dominant machine learning

architectures in many fields okay so

what does Iowa mean IO just means in and

out read some writes so the thing we're

going to exploit is that on the left I'm

showing a diagram of GPU memory in SRAM

and by exploiting SRAM which is way

faster but way smaller than the GPU

memory or hbm we're going to make things

go faster and the trick we're doing that

as I've shown in the the middle is

you're going to compute things by block

and that's called tiling and the out the

the upshot is on the right we're going

to make it faster we're going to save a

lot of memory and we're going to have

higher quality models

um so that those are the main main

takeaways

so before

um before I dive into the details so I

want to talk a little bit of about

motivation and why this could be

relevant to you know folks working in in

medical uh medical domain or medical AI

so the our motivating

um

uh idea is we want to model longer

sequences and and Chris uh my advisor

and a bunch of us put out blog post

that's Linked In the at the bottom here

about why we we care about long

sequences and there are three three main

reasons one is we want to reach new

capabilities so the anchoring

application is in natural language

processing where we want to model uh

books and plays and instruction manuals

not just paragraphs

um and the second reason is we want to

close the reality Gap and the anchoring

application here is computer vision

where we've seen that if you can model

images at higher resolution you get

better and more robust insight and high

resolution uh generally translates to

longer sequences

um and the third one which is probably

the most exciting is you want to open up

new new areas so things like stamp

series audio medical imaging video

things like that they're naturally

modeled as sequence of up to millions of

steps so that's why we care about long

sequences but the problem is that the

the dominant architecture machine

learning architecture nowadays the

Transformers uh they have trouble

scaling to Long sequences so now go into

why uh without uh the the problem we're

going to tackle is that we want to scale

Transformers to Long sequences

so that's the high level

um motivation and I'll go into three

examples uh in in some medical

applications why you should care about

uh Transformers and why should we care

about Transformer on Long sequences

the first this is in collaboration with

Armin at the uh the psychology

department so we're interested in

building for example Foundation models

for brain decoding from fmri

Readings so

um surprisingly you can use the same

Transformer model that has been used in

national language processing

um to build to build representations for

for brain decoding and arminess has a

great paper on this and uh we want to

make this faster we want to make this

more memory efficient so that we can

model let's say longer time scale when

you you do fmri decoding

um so that's one thing we're excited

about this is ongoing collaboration with

Armin the second application motivating

application is uh with Arjun who's a

future student in Subway and in medical

school where we want to model high

resolution MRI images so these are knee

MRIs and usually one of the tasks to the

task to care about is detection and

segmentation and as you see

um for Radiologists they usually care

about these tiny regions that if you

don't you can't Model A long high

resolution you can the models will tend

to miss these these regions so the goal

again is to model high resolution images

and that's like for for Transformers to

Long sequences

and the third application this is in

collaboration with Gautam

um uh is on megapixel and gigapixel uh

pathology which is a very similar

application that uh for example the the

talk last week

um yeah from Richard he also talked

about megapixel histo pathology so let's

say the the task here you might want to

predict tumor versus healthy tissue at

different image levels or you want to do

object detection for tumor lesion and

for for uh pathology images these are

usually stored at different resolutions

and you want to model and incorporate

information from different resolution we

be it at the style level is it be it at

the organelles level and things like

that so you need to be able to deal with

with a really high resolution images and

and transform Transformer nowadays

um they can't quite do this yet the

bunch of tricks you have to do but again

we we hope to make Transformer able to

model long sequences so that they can

deal with these high resolution images

in in these applications so those are

the three motivating applications uh why

we care about Transformers for for long

sequences

um so that's kind of the the motivation

um and then uh so the talk we have four

parts I'll talk about some background

then I'll talk about our our proposal

flash attention

um and then I'll I'll go into empirical

validation and finally a little bit on

on future directions and yeah feel free

to ask ask questions throughout the talk

I'll pass I'll also passed uh

periodically for uh for questions so uh

before we go into the background happy

you know I could pause here and see if

people have questions

yeah

all right so so let's let's get into the

the background uh

why Transformers have trouble modeling

long sequences

um so just many of you might be aware of

the Transformer architecture already

this is just so that we're on the same

page in terms of of notation so

Transformer is a machine learning

architecture where you stack a bunch of

these blocks or encoder or decoder the

two types and each block is each block

consists of two layers the attention

layer and the MLP layer and we'll focus

on the attention layer because that's

usually the bottleneck for uh for long

sequences

and attention it looks like this many of

you my might have seen this so this is

just so that we're on the same page in

terms of notation

um the inputs you would have the query q

the key K and the value V

and the steps are that you would take q

and K and you multiply them together

optionally you apply some masking uh

then you do a softmax to normalize

things and you'd apply a drop out

optionally and then you multiply that

with v to get an output o

and the the the thing that prevents

attention from scaling to Long sequences

is that naively attention has compute

and memory quadratic in sequence length

uh here's a good thing I didn't know

with with n

um and the reason is simple is because

you have to deal with these big n by n

matrices

um so uh you need to you know n Square

amount of memory to store it you need n

Square amount of compute to process it

so that's

um that's a simple answer and you can

think of M as let's say a thousand or

two thousand

um and then D the Hat Dimension is

usually around 64 or 128. so and and by

n really uh

uh as it was whether quadratic blow up

happens

so we want to we want to address the

address this so we're going to

understand what are the bottlenecks and

the bottleneck turns out to be not be uh

computation the battery is in memory

reads and writes

um so what do I mean by that is uh there

are a bunch of intermediate steps where

you have to deal with this n by n deal

with matrices of size n by n

and so

um that means a lot of times it's going

to be spend on just reading and writing

into memory

um and that's really slows you down as

I'll show you with some some profile

and so naively it requires a lot of

reads and rights from slow GPU memory

hbm you know it's just the technology

what you think of as sgpu memory

uh and there's a paper from mlc's last

year that we really like uh the called

Data movement is all you need where they

study Transformers and they they found

out that yeah the the bottleneck is is

data movement which is just reading and

writing data is actually not computation

and that's quite wasteful

um so we're very much inspired by that

paper and just as a preview for the

result so here I'm comparing the against

the the pi torch attention

implementation and our flash retention

implementation and they compute the same

thing you know we're not changing the

output and we're not changing the math

they compute the same thing but it turns

out Pi torch uh when you compute

attention most of the time is spend on

these memory intensive operations

masking soft mix and Dropout whereas the

matrix multiplication uh

counter-intuitively they don't take that

much time even though matrix

multiplication takes the most uh

floating Point operations

um is the memory reads and writes that

really approach you and so flash

attention is a way to compute the same

thing but with significantly reduced

memory reads and rights and we can go in

this case maybe six times faster

so to to understand how we do that we

need to understand a little bit about

how gpus works so this is a very brief

background on on GPU and here I'm

showing some diagrams let's say your

machine has a bunch of gpus

um and each GPU you can think of as

having a pool of memory would you know

if you type Nvidia SMI it shows you 32

gigs that's the hbm that's the amount of

GPU memory you have but when you do

computation uh what happens is that the

data is loaded the input is loaded from

the hbm it gets loaded to the compute

unit and the SRAM which is you can think

of as analogous to CPU cache it gets

loaded from hbm to the compute unit

all the computation happens there and

then the output gets written back to hbm

so it turns out this loading and writing

of

um of data to HPM actually is the is the

bottleneck and I'm I'm looking here at

the bottom uh a blog post from Horus

this is a part switch developer is an

excellent blog post on understanding

what are the bottlenecks for a lot of

these deep learning operations so it's a

great buff of um I highly recommend it

so okay so how do we so we see that

there is an asymmetry between uh the hbm

and the SRAM so hbm is usually quite

fast and it's quite large but SRAM is

even an order of magnitude faster but

it's let's say three orders of Mac into

smaller this is what I'm showing on the

on the right and these numbers are for

the a100 gpus but the gpus are similar

so the questions can we exploit this

this asymmetry in the the memory

hierarchy so that we can get speed up

and the answer is yes with with i o

awareness that means we want to account

for different accountable reads and

rights to different levels of memory so

we're going to try and reduce the reads

and rights to GPU hbm

so that's the main

objective we're going to reduce reads

and write

so that's uh that's it for the

background I'll pause here for for

questions

um either about attention or about you

know GPU architecture

so sorry I just have one quick question

since nobody has it like um

so when you do the metric multiplication

operation are you saving it as a

different variables or you are saving

just overwriting the same variables

it seems that you are just overriding

right

um yes uh so here we're dealing with the

an end by n Matrix so yeah you can I

guess stay from memory if you just write

to the same location yes okay

okay

um in practice uh when you train models

people don't do that because generally

for when you do training you need the

the the to store the intermediate value

so that you can compute gradient in the

backward pass so in practice most of the

attempts people don't write to the same

location but yeah it's a great idea you

can write to the same location maybe say

I don't know 4X memory or something yeah

right because I would imagine like in

the pie torch exactly um the pie touch

operation that you are showing for the

attention I think they are also not

really not only like time efficient but

they are not also like memory efficient

I don't know if you have any solutions

since you are just doing the metric of

multiplication this kind of like um

inverse functions I don't know if you

thought about like optimizing also the

memory at the same time of optimizing

the time

yeah so that's that's a great point so

as we show that you can be clever and

not you know not materialize these n by

n Matrix matrices at all

um so in the end you get something like

10x memory savings overall

um so one

um one thing about reading and writing

to the same location is that you still

need to you know in terms of memory yeah

you do save the total amount of memory

but in terms of time you still need to

do the same number of reads and writes

right so you might not get speed up but

you might do my save some some memory

yeah I agree

okay yeah

um could you could you explain why the

Dropout and softness actually takes uh

much more time than the yeah yeah so

this is a little bit surprising

um so turns out the

um

there most of the time is actually spent

on reading and writing into memory so

every time let's say you start with with

a uh a you you takes Q times K you get

the Matrix a and then you apply some

masking and what really happens is

actually you're not doing any any math

what you need what you do is you load

the value from hbm to the compute unit

the compute units set some of the values

to to zero or negative infinity and then

it writes them back down the the output

so it's uh uh it has very minimal

floating Point operations but the amount

of memory reads and wise is much larger

than Matrix or at least you know double

the major amount of memory written right

when you do matrix multiplication so

um this is a little bit surprising

because attention

um yeah most of the time you think

things are bottom act by computes if you

reduce the amount of computation you

would go faster but for attention it's a

little bit counterintuitive is that the

The Tam is taken by memory reads and

writes just the you know the matrix

multiplication turns out to be faster

than these uh masking or Dropout or

softmax operations

I see thanks

I had a good question but for Dropout if

you're doing random dropper you won't

need to like load anything from HPM

right

uh so you still do need to load things

from hbm and in the sense that because

for each GPU uh operations it's called

kernel the input starts out in HPM you

need to load it into hbm you set some of

those to be zero and you write it down

to hbm

um

yeah the architecture works I see yeah

so if we prevent that from going down to

hbm before doing the Dropout then that

can serious but I guess we still need to

do it for the gradients so yes yeah so

that's an excellent uh idea actually you

know people have tried it attention so

well studied that people try all all

ways to get all the ways you could think

of to optimize these things and yeah

they can combine like hey let's just

load you know let's say combine stop Max

and rubber you can load things in do

softmax and then just apply Dropout

before you write the the output out uh

so it is an excellent idea and people

have done that and they have seen some

speed up

um but again generally you need to

materialize at least one of these n by

end matrices and um that's still quite a

bit slower yeah correct but thanks yeah

so there are more optimized

implementation here we're comparing

against pythosh but uh you know the the

excellent people at Nvidia they they

have a bunch of different

implementations

um and we do compare against them

um to you know combine some of these

steps to reduce IELTS so uh from a

hardware perspective yeah it's very much

known that yeah you want to do we do

reads and writes and you can find some

of these operations

um but so far people haven't been able

to combine all of them

um and and as we show there there's a

little bit subtlety but but you you can

do it and and you can get the whole

thing much faster

yeah

yeah great questions

all right so so I'll go into details of

how we actually do this and the the one

trick is actually really simple

um so before that so uh as I explained

we want to reduce the memory reads and

rights to hbm and the answer usually is

that you want to compute things by

blocks so that it's small enough that it

fits into SRAM that you can load from

hbm to SRAM which is new local

computation on SRAM

um and so you you don't need to to

repeatedly read and write to to hbm but

there is there are two challenges one is

that the softmax as defined it needs the

whole row before you can do softmax by

by definition uh that's challenge number

one so if you want to do things by Block

it's not clear how you would do softmax

if you want access to one block instead

of the whole the whole row and Challenge

number two is that

um when you train models you usually

need these intermediate values so that

you can in the backward pass you can

compute gradients

um so if we want to save memory or we

want to do reads and write if we don't

save these large attention matrices how

do we compute the backward pass

um and we use two well-established

techniques to address these one is

called tiling which just means Thailand

just means doing things by block but you

need to restructure the the attention

algorithm a little bit uh to get the the

uh to get things to load from hbm to

SRAM and the compute attention kind of

locally uh it needs a little bit of a

trick to get the the right answer

um and the second Technique we use is

called recomputation so the answer is

um we don't store the attention Matrix

from the forward pass we recomputed in

the backward pass and turns out

breathing computation is actually faster

than than storing it and reading it

again so that's a little bit surprising

um and I'll go into details of some of

these techniques

um so the way we implement it is that we

have a few scooter kernel for so that we

have fine grain control of the memory

reads and rights so we can control where

to place

um what and we provide a pi torch

interface that that people can use

so let's go into detail of some of these

techniques

um the first one is tiling and this is

you just need a little bit of math this

is the only math in in the talk

um the trick is that you have a soft Max

over let's say a large attention Matrix

um

but you can break it up let's say you

can break it up into two pieces so if

you want to compute top Max of A1

concatenate with A2 you can actually

compute topics of A1 first then computer

soft Max of A2 and then you can rescale

them using these scalar Alpha and beta

and you you can you can work out exactly

what Alpha and beta are uh they're

they're the details in the paper but

they're they're pretty easy to to work

out so uh softmax you can actually break

it into smaller softmaxes so that you

can exploit locality and when you apply

to attention you're going to have

softmax with A1 A1 and A2 times the V

the value Matrix V1 and V2 and you can

decompose it exactly the same way so

you're gonna do you get soft Max of A1

times V1

so that's a local computation another

local conversation is you take stock Max

of 82 times V2 and is rescale the

ingredients and adding things together

you get the right answers so

mathematically is the same answer you're

just doing the computational let's say a

different order

and the nice thing is that you can

exploit locality because here we don't

need to load the whole of A1 A2 in we

just need to load A1 and V1 and then we

compute whatever local things we we need

so this is the the only trick that you

need to know

um and the algorithm is is relatively

simple

um so we're going to load inputs in

block by block from HB into SRAM so on

the diagram on the right you can hbm is

in green and SRAM is in Orange so you

can think of conceptually we're

Computing attention block by block and

for each block we're going to load

whatever inputs we need into from hbm to

SRAM we do whatever local computation

and then we're going to write down the

output to httpm so we never materialize

this big n by n Matrix

um

um and and with that you you get memory

saving and you get speed up as as I'll

show you in some of the benchmarks

um so Tyler that's the that's the idea

uh you analytically rewrite softmax a

little bit

um so that you can do local computation

block by block so that's the that's the

main idea

uh another quick question so in this

case the alphas and betas are not

constants right they depend on A1 and

attitude yes yes

um because yes that's right so when you

load in A1 and A2 you also need to

compute let's say as one scalar per row

um so it's depends it's dependent on on

the data you're right here so then you

can't do it by the blog because it also

depends on A1 E2 and it also depends on

everyone I need to provide

so Alpha only depends on A1 oh sorry

sorry yeah yeah no no you're right so

Alpha depends on both A1 and A2 but it

did depend on just one scalar they're

just normalizing uh scalars from uh from

each block so you can think of uh think

of it as you know you have a summary

statistics each block and then based on

the summary statistics which are really

small you can combine them yes

I see so you first read okay so you

compute some researchistics and okay

yeah that that makes sense that can be

easily done when you are Computing the

A1 and A2 right so yeah

yeah great great question

foreign

has been known in a while uh uh you know

tiling in general computing by block is

kind of the Cornerstone of of databases

Graphics

um numerical linear algebra and this

trick of decomposing softmax uh the

first time we saw it was in 2018 so

um so these are well established

techniques that were that we're using

but surprisingly these simple techniques

will allow us to go faster and more

memory efficient

okay so that's the first technique

tiling the second technique is

recomputation it conceptually is also

very simple so you don't store the the

uh attention Matrix from the

um from the forward pass in the back

request you recompute it and if you

already have these normalizing

normalization factors Alpha and beta you

can actually recompute it pretty quickly

since the inputs you need are going to

be loaded from SRAM anyway so in some

sense uh computation is is almost free

because you're already doing the reads

and write to load inputs in

um and again we don't store this big n

by a matrix everything is loaded in and

then compute on SRAM

um

so the the memory requirement is much

smaller we're much fewer reads but we do

encourage more flaps and here I'm

comparing uh the standard implementation

with flash attention

uh and this table and first of all in

terms of flops we actually incur more

flops it's about 12 more flops because

we recompute things in the backward pass

but in terms of hbm reads and writes so

so memory reads and writes we're about

nine times smaller and in terms of

runtime we're about in this case like

six times smaller so even though we

recompute things even though we in term

of lab surprisingly it actually speeds

uh things up We Begin Again the the

observation is that is attention is not

bottlenecked by uh computation is

bottlenecked by reads and rights so you

want to reduce it right even if you have

to incur more computation

so that's the the second technique

um and so those are you know the two

techniques that we use these are well

established

um but surprisingly

um attention's been out for five years

is one of the most studied uh problem in

machine learning is uh but surprisingly

no one has uh has done this

um so it's pretty cool you can put it

put together these these simple

techniques with a little bit of a Twist

and then you can you can make attention

go faster and more memory efficient and

even compared to the best uh the the

strongest Baseline we we compare against

which I'll go into yeah so I'll pass

here

um for for questions

another question I had was how do you

make sure that this Affinity From the

Block to the SRAM is maintained in like

the python implementation for your

flashlight yeah yeah so that's the

difficult part so we had to implement

this in Cuda so that we can decide we

can control when to load what things

um so that you know that actually took

significant obviously two students

writing Cuda for three months or

something but uh yeah but we're excited

about a couple of things one is that uh

we're collaborating with with compiler

researchers so that we can automate a

bunch of these things so the compilers

can do it too

um number two uh some people you know

some some other researchers really like

this so they they independently

re-implement it in in Jax or in Triton

Triton is a python-like high-level

language so it's not as painful to to

write

um

um so we're we're excited about that

hopefully this will be accessible to

researchers if you just want to use it

then we have a pi switch interface you

just call it and it's going to run

um if you want to Tinker with it then

some of these higher level uh

implementation in highly level languages

uh could be pretty useful

awesome yeah that's great thanks

okay so those are the techniques we use

to speed up attention again the goal is

so that we can make a lab Transformer to

model long sequences and and now we're

gonna uh yeah I'll show you some some

empirical validations on how this this

helps us with that goal

so we're going to Benchmark the

attention module we're going to show

that it is faster uh

even when you put in an end-to-end model

so you get faster end-to-end model

training or Transformers and again we're

not changing the model at all we're

competing the same thing which is faster

and the second thing is that we're going

to save a lot of memory so it allows to

model longer sequences and that leads to

higher quality Transformers

so first we when we brand Benchmark just

the attention module we see that in

terms of speed is about 2 to 4X speed up

compared to the standard implementation

from python hockey phase or whatever you

people commonly use

um so this is we have a chart on a100

where we get 2 to 4X speed up depending

on whether you have a masking or drop

out

um and in the paper we have more

comparisons with approximate attention

and sparse attention and so on we also

Benchmark and uh Benchmark on a bunch of

different gpus I'm not showing here all

the charts are on our GitHub repo yeah

generally you can we observe about two

to four x speed up uh of attention and

again attention is a very well studied

problem and surprisingly you can you can

get substantial speed up

um

uh even though tons of engineering

efforts has been poured into improving

uh attention implementation

so that says uh just a tension module

you get two to four x speed up and when

you uh when you measure memory by usage

which is the total memory you need uh

the nice thing about Flash attention is

memory is linear in sequence length it's

not quadratic

um so depending on sequence length

um where you can see either 10x or 20x

memory savings uh for the attention

layer loan so this allows to scale up to

64 sequence thing of 64k on on one a100

GPU and just for contacts uh some of the

largest Transformer models have context

line of like 2K just because of this

quadratic block so if you have memory

linear in sequence Lanes you can scale

to much longer sequences

um all right so now we're going to plug

flash attention into a model and we're

going to train end to end and we'll see

that it leads to faster end-to-end

training

um

so the first one uh we want to compare

is is bird large it's extremely commonly

used model

um and The Benchmark or the yeah The

Benchmark we're going to use is ammo

perf which is a highly optimized

standard Benchmark for training speed

um so uh this is a comparison where we

all the models starts from the same

checkpoint they train on the exact same

data same validation data and you just

measure how long it takes to reach a

certain accuracy

um so everything else the same the model

is the same just the different

implementations different algorithms

um how fast do they run and uh if you

use a commonly used hugging phase bird

um

then this takes like 55 minutes

um if you use the ml perf uh

implementation from Nvidia which is

which was the the world record

um

when we submitted the paper

um you know this is extremely optimized

teams of Engineers have worked on this

so they they're like two 2.5 x faster

than than hugging face uh but

surprisingly flash attention can be even

faster than than

um the previous ml perf record by 15 so

it was really excited uh exciting for us

so you know we're also 2x 2.2x faster

than hockey Facebook so if you're

training board and you're using hacking

face implementation you're leaving

around 2x performance on on the table

um

so the second model we can compare to is

gpd2 and we're training gpd2 small and

medium from scratch on this commonly

used data set open web text and again

again we compare against hugging phase

which is very commonly used and we

compare against the strongest Baseline

we could find which is Megatron LM which

again an implementation from Nvidia

that's been used to train some of the

largest language models

um and we see that we get the same

quality we in terms of training time in

terms of training time we're about three

or three point five x faster than Huggy

face

um so again if you're using hanging face

you're leaving about two 3x on on the

table I know the models are the same

we're not changing the model we're only

changing the the algorithm and the

implementation

so the upshot is you get faster training

for a bunch of these Transformer based

models

um

policy for for questions

so do you start analyzing also the

attention visualization for this model

just to see like if there is any

difference in the attention

output or I know that it is more or less

like mathematically it's a exactly the

same but I don't know if you have

um analyzed like how that engine output

looks like in the normal model which is

your flash attention

yeah so we haven't visualized things we

did some spot check where the outputs

are the same or up to numerical

Precision like yeah one e minus four or

something so the upwards are the same

yeah so I don't expect the attention

patterns to change

um if you want to change the attention

pattern then there's some other methods

where you say you know I want to add

this to the attention so that it biases

towards like local features or Global

features or something so we haven't

explored that yeah but that's uh that

could be interesting

okay so you didn't have any kind of like

linear oppression performance on the

attention and just to see like if there

is any change in attention so you just

analyze it's more like a training speed

wise right

yeah training speed wise

um the the training curve look identical

so they look like they lie on top of

each other again it's because we're not

changing the model right

all right so here we can change the

model and we can we can train uh

Transformers on longer sequences and

that generally leads to higher quality

so here's an example again gpt2 but we

increase the context length and usually

when you increase the context length it

goes slower

but generally it gets it gives you a

better model and this is measured by

perplexity where lower is is better and

surprisingly when we increase the

context length by 4K uh we get we're

still faster than Megatron and LM at 1K

and we achieve lower complexity

so this is this is really exciting and

in fact a bunch of you know a few

companies are looking to use Flash

attention to train uh language models

with longer context up to 8K

um you know so this stuff actually works

yeah so one longer context it leads to

better language models the second one

let's say is a classification model this

is a classification task where you have

a bunch of long documents these are

either physician notes uh that's mimic

and then the second one is a legal

reasoning task and they could the

sequences could be pretty long up to 10

10K and when we take a pre-trained model

and we we extend this context and and we

fine-tune on this classification task we

see that we immediately get some

improvement uh inequality by just simply

using longer longer contexts so yeah

this this makes sense

and so one one synthetic task that's

that's pretty cool is called path X this

is more of a benchmark test to see if

Transformers you know how how much you

can push transformer into the sequence

length and the task is looks like this

image where there are two dots and a

bunch of lines and you're supposed to

say whether the two lines are the two

dots are connected and in this case it

is not connected it's a classification

task uh so it's more of a benchmark

synthetic tasks you're not supposed to

use patches so you actually need to

process one pixel at a time your

sequence length is either 16k or 64k and

in the past people have tried and

um basically no Transformer can get

better than chance of a better Enchanted

performance either because they run out

of memory they too slow to train or

um they do some approximation and they

don't they you know it can't pick up

these signals at all and so what's

really exciting is Flash intention use

the first Transformer that could solve

better these tasks better than Jazz so

it should go up to sequence length 64k

or

um yeah 16k or 64k

so that's something we're pretty excited

about

um yeah policy for for questions before

going to uh future directions

what's the difference between flash

attention and blog Sports flash

attention here

yeah so I I didn't mention those um

since I have flash attention you can

also impose sparsity which is you know

so far I've talked about only doing

exact attention we're not approximately

I'm approximating anything for

um there's been a long line of work on

on attention approximation in one such

way to do it is impose some sparsity so

you if you similarly impose some

sparsity flash attention is again

compatible with that and that allows you

to uh reduce the computation

by exploiting sparsity at the cost of

some approximation

um so for the for the largest task path

256 where the sequence thing is 64k uh

faster tension would run but it would

take like a month to to train but block

sparse if you impose some sparsity it

finishes training in like three days or

something

um

so it's one way you can push the the

sequence thing even even longer by using

some approximation but for the most part

flash attention

um that it just does exact attention is

not approximating anything yeah I see

all right so let's go into some future

directions that we're pretty excited

about the first is just getting it into

more people's hand

um so the pi support developers are

integrating this into pytor

torch.nn.transformer uh the hugging face

folks are integrating into their models

as well and independently

um other researchers they were excited

about these ideas so they re-implemented

it in in Jax uh the openai folks

re-implemented in in Triton which is

this python-like language that's a

little bit higher level than 10 Cuda so

that's a little bit more friendly for

for tinkering

and second one there's some technical

work

um you see that the ideas are actually

simple so maybe compilers can support me

automatically do these things and so far

right now they are not able to because

they need a little bit of math rewrite

um but uh we're working we're

collaborating with with compiler

researchers so that some of these

techniques could be done automatically

and maybe in the future yeah your your

model will just get faster more memory

efficient automatically

uh and the third one is a high level uh

point which is we want to bring Hardware

we're thinking into other deep learning

operations Beyond just just attention

and we believe this is a being a big and

important idea and in fact it came from

things like databases where for decades

people have been trying to reduce the i

o and spend hugely successful

so the more AI question is what can we

do with longer sequences

um you know we're we're really happy

that several companies are looking in to

use Flash attention to train language

models with 8K context Lane which is

about four times longer than the the the

the largest models right now gp3 gp3 can

model of context have context thing

about 2k so people are excited about

pushing this further to get better

models another thing as I mentioned in

the beginning is we want to model highs

resolution images be it from MRI images

or pathology images maybe High high-res

images will give us a better insights or

better models

um so those those are the things we're

excited about

um so in summary flash attention it is

fast it is memory efficient and it is

exact so it doesn't do any approximation

to attention and the two key ideas are

tiling and recomputation and the upshot

is you get faster some model training or

you get better models with longer

sequences and with that uh General I'd

like to conclude so thanks so much for

for your attention and our code is is

open source um if you will flash

attention or you go to GitHub we have

links to everything

um and yeah we're

um we're excited to see

how modeling last sequences with flash

attention what what kind of new

capabilities is going to give us yeah

thanks for your attention

all right thanks Chi for the amazing

talk let's all give him a round of

virtual Applause

and is there any questions from the

audience

yeah have to to take questions

um as well yeah

when you mentioned uh high-res images

what sort of Matrix are you thinking or

like bit depth to

um so one application more collaborating

with real term is these pathology images

uh where they're either megapixel or

gigapixel so the that's like I think the

largest one they have is 50K by 50k with

a hundred channels and that is generally

just way too big to do anything with and

right now people just cut them into

patches and

and process one patch at a time so

um you know maybe we can do something

there

um

let's see if I can find these

um yeah

um for MRI I think the the resolutions

are like fifth uh 5 12 to 512 something

like that yeah that's about right even

256 occurs sometimes

yeah

um so again that's still larger than

what people you know current models can

support

foreign

what's the longest sequence lens you

have experimented with

um using flash attention

yeah so the longest we've tried is uh

either 16k or 64k and um

I think that corresponds to images of

like

two 256 if you do it by pixel if you do

patching like vit then

um you know that allows you that could

correspond to high resolution images

yeah so similar thing is about 16k to

64k which again is quite a bit longer

than the most common Transformers

yeah

um yeah I guess for the histopathology

images it can be a thousand hundred

sorry tens of thousands patches or yeah

even more than that so yeah yes I can

see that flash attention could be a very

good um

could be very useful here

yeah I guess like if you can do 64k then

yeah you can probably like you just like

take a square root of 64k and those many

patches it's still a huge context lens

for uh pathology which is for sure and

probably can work on whole MRI which is

completely that right

yeah hopefully hopefully we'll get there

this is in active research and if you're

interested please please talk to us uh

yeah we we this is an area we're really

excited about so

um yeah we'd love to talk to you

[Music]

is there any other questions

okay thank you so much I think yeah

let's think three again for the talk and

thank you everyone for joining us today

we will upload the recording of this

talk to our YouTube channel later today

and if you have follow questions feel

free to reach out to us on Twitter

directly

thank you thank you Trey thanks for

giving doctor bye