Bringing Triton to AMD GPUs|Jason Furmanek, Lixun Zhang (AMD)

AMD back end bringing Triton to AMD GPU's.

Again, I'll get the presentation set up and we can get started.

Everybody, listen and I from Andy when I'm talking about how we brought into MDG views over the past 612 months.

Here we go, OK, so yeah, so we've been working on this for about a year.

Deep speed is important to us and supporting our customers.

It really kind of kicked into another gear once Pytorch 2 picked it up, right?

So we have spent a lot of work on it, so this is just a quick update of where the back end is, and then I'll kind of take you through the story of what all we did to bring it.

But so the code is out there, so we have the third party PR 1983.

There was merged so we have our third party submodule link.

In the open ad trading repository now, which is great So you can have another slide on this on how to check it out, but it's all in there.

The actual development happens in our repository, though, so that some modules of link to our Rockham software platform Triton to a specific commit level there.

Yeah, so that's yeah, so that's the development.

I have another chart on that in a minute.

I'll talk about so but the so hardware support we have.

We focus on our CD net CDN GPU's which are big data center ones.

We do have experimental support for the RDNA as well, so that they both should work.

There's different warp sizes on both of those, so we have auto detection for the hardware implemented in there now, so we.

Train is able to find out what hardware were on and do the proper thing there.

So this is, but CDNA is our lead development platform on there.

We do have full support on PyTorch 2 for all the apps that are lowered to Triton.

Working on the hightorch21 at the moment, so all that all that works at the moment, which is which is great.

We have to change much of the front ends like almost at all right, but we did add this device type of hip.

If you're calling directly into the compile, and you wanted to specifically call out the devices that have so.

You know beyond the needs of dependencies reference for deep speed or even Pytorch right does we had to.

We had to kind of see does this map well to our hardware, right?

Is this worth doing and the answer is absolutely yes.

I was looking for some supporting evidence of this and I'd have to look for fills, original paper and even the docs directory out in the GitHub train.

GitHub out there is a good source of information on just kind of how things map over, right?

So the original design tenets of, you know, kernel level blocks and locality and efficiency and high occupancy all just translate really well to AMD GPUs.

So so in in it that's, you know that was just green light.

Let's go ahead and get this optimized and working on our hardware.

So Beyond that, beyond the design though.

Yeah, what are all the what are the pieces here, right?

Is there something here that's going to take a long time to port?

So we had to go through and look at everything.

So first and foremost, LVM MLIR that those are big pieces, right?

We have MD's already made huge investments there obviously, right?

We use that in all sorts of different places today, so.

So you know, so beyond that, like is what's higher up the chart and GPU dialects just have some Nvidia specifics in there, but While we were looking through this and porting it, there's we didn't really even touch that very much, so we didn't have to add any new OPS.

Although all the OPS just kind of, you know, work the same way as they do on Nvidia.

They lowering the dot op is a little bit different because we want to have.

We have our matrix cores versus the tensor core so that there's no topic on that.

I'll save that for a minute, but even at that layer we were looking pretty good so.

So beyond that, so and then we've got all the different layouts, right?

So our shared memory on AMD GPUs is the design is very similar to the Nvidia design.

So all that, the shared memory layout Mostly just kind of worked for us as well.

We do have different sizes of shared memory.

We in ours we have slightly less shared memory, but we have more general purpose registers.

We want to prioritize maybe our general purpose registers more than and try to try to be more conservative on the usage of the shared memory in that.

On is evident on some of the pipelining passes that we're talking about today already so.

And we'll talk about more of that in a second.

So then the other thing then is the inline PTX right?

So the PTX builder is was resonated build a GCN builder, So we've got our GCN builder that would just build in line assembly for our MDGC and assembly so that that was 1 piece we didn't have to go and change, but all this was to map towards our current hardware that we've got in the bottom there.

So it all looked pretty good, so we had to come up with a plan on getting all this working and then see if there's anything else that we were missing that we need to work on.

So the first thing was, let's get our empty kernel working right?

This was kind of the bulldozing effort to get everything.

Just just we can get a all the pipe all the.

We had to, you know, print that doesn't a big deal.

We could parameterize that we so we've got.

We implemented parameterize and we called the MD terminologies way front, so we front warp sizes interchangeable.

So we've got that that plumbed in and we have we can have auto detection.

So on our DNS we have 32 on the CD and then we have 64 and then the kernel launcher code we had to build new template code for the kernel launcher and then that was pretty much that we need to do for the for the empty kernel.

Get all the OPS working and this was where we need to get that GCM builder in.

There's a couple other things that we had to work out.

Some of the we had to disable the existing pipelining paths, for example.

And then you know, so that was probably 1 of the larger efforts.

Just kind of getting through all that, right?

So Once we had all that function working, we got to let's make this thing faster, right?

Let's implement our matrix cores within Triton.

So we shouldn't talk about that in just a minute, but Triton has the MMA layout, which is a big important piece and get performance, right?

We initially we were kind of editing that we ended up having an alternate layout.

The MFMA layout we are instructions are in the from the instructions so.

That just kind of slots alongside MMA B2 and V1 and V2 and now V3.

So that that brought us up to speed on our matrix core support and.

The original or that we so in which had our flash attention working well now and it's and so we started doing more performance tuning and more optimization, which is lead us where we're at right now, which is trying to maximize our occupancy, trying to maximize our performance and seeing what all we can do to really kind of just tune this all the way up to our distinct hardware.

And we should talk more about that in a little bit as well, but this mostly has a new pipe lining.

We were doing a streaming pipe lining, which is easier on our shared memory amounts, and then we also have the.

Swizzling which the most of the swizzling algorithms Mostly we can reuse, but the instructions are different, so we've got that in and we also have.

A piece where we can edit the number of wavefronts per EU that we can maximize our occupancy on the hardware.

OK, so before I turn over to show you some more details, I want to talk about the little bit about the third party back end that I mentioned at the beginning.

So this is merged now that there's the commit here where we're pinned to at the moment.

This does predate some of the performance stuff we're talking about today, but we're working on a synchronization effort that brings in a lot of the commits for that were committed upstream.

We're taking a little bit broke some of our stuff, so we're working through that.

So what once we get that all synchronized up, we're going to update that commit and then we'll have a kind of everything on the top layer and have that commit from the front end be the fastest 1 that we've got.

There's also on this chart on how to build it, right?

So there's a environment variable we use trading code gen India.

And then you can just check out recursive and you'll and then build like normal and then you'll get the MD back end enabled.

And thanks to the Intel folks that so we this was kind of a collaborative effort between everybody to try and OK, let's get.

We've got this the Triton GP Rockham, which is and there's a lot of duplicate code there, so we'll have to work together to try and reduce that duplication as we go forward so.

Is it going to do some of the more lower level details on the port.

Let's talk about optimizations we have done on the AMD GPU.

So first, let's talk about the AMD matrix course, which is a counterpart of NVIDIA's Tensor course.

It supports a full range of mixed precision gem operations, which is achieved by variations of MFMA instructions.

So the key to harnessing the MFMA instruction is its data layout, which refers to the distribution of elements among threads in a warp.

So this is an example of data layout for MFMA 32 by 32 by 8, flow 16 instruction.

So as you can see, in general, each thread is assigned on the segment for elements.

And for input A, the elements in the same segment is in the same row, and consecutive threads are assigned segments along the M dimension.

For input B and output D, elements in the same segment are in the same column, and the consecutive threads are assigned segments along the M dimension.

For input B and output D, elements in the same segment are in the same column, and the consecutive threads are assigned segments along the N dimension.

With the knowledge of data layout, now let's talk about the problem of chain dot.

So this problem refers to, we want to use the result of the first dot as the input to the second dot.

And which input to the second dot makes a difference here.

So if we're talking about input B to the second dot, we are good to go because input B and output D have the same data layout.

To put it another way, there's no need for redistribution of elements among threads.

However, if we're talking about input A to a second dot, we have to go to LDS to rearrange the data among threads.

So let's see how we can avoid this redistribution of the second case, which is what we have in flash attention.

Now you can see the elements from matrix A are passed to the MFA instruction as input B, and elements from matrix B are passed as input A.

Okay, so because of the way how the hardware interprets elements from threads, Now we are computing B transpose times A transpose.

Mathematically, we got D transpose with MFMA output layout.

MFMA output layout on D matrix is the same as transposed MFMA output layout on D because threads are held in the same elements.

So in this way, the output Is ready to be used as input to the 2nd dot and we achieve transposition without transposing anything in the hardware, so no overhead.

Software pipelining, this is a tricky part, so we are developing a 2 stage software pipeliner with 1 buffer in LDS and the other buffer in VGPR.

So the compute always start with data in LDS.

It's never start with data in the buffer in VGPR.

So the compute includes reading data from LDS and MFMA instructions.

Okay, so we don't use the buffer in which PR as a regular buffer, because this is used as a stream buffer to wait for data from global memory.

So well, we perform compute for the current iteration of the loop.

It's loaded, it's being loaded from global memory into VGPR and we always do the competition with data in LDS.

So in this way we achieve some kind of 2 buffer software pipelining with only 1 buffer in LDS.

So we save LDS space at a price of higher register pressure.

Okay, So we also have implemented more optimizations on AMVGPU.

So in the context of FlashAttention forward kernel, we keep Qtensor in register instead of loading it from global memory or LDS at every iteration of the loop.

We also arrange our, we also arrange warps vertically in the CTA so that we can.

Reorder pointer arithmetic so that we can generate global load instructions with maximum vectorization list.

Because the hardware is not doing a very good job on vectorization.

We are also experimenting with this waves per EU attribute, which limits the number of registers to increase the occupancy.

We see great performance improvement when setting this attribute to 3 for flash attention forward curl, which is weird.

So because if you don't set it, it's a it will be 2, so occupancy will be 2.

But if we set it to 3, we have some register spilling, but we are increasing the occupancy so we have better performance.

On gem performance, So this is kind of internal comparison of Triton and the Rock Plus on dense jams with different M and K values.

So as you can see, we are still lagging behind ROC+.

We are optimizing on this kind of stuff all the time.

And then we have some to-do optimizations.

So first, we plan to store the result of MFMA instructions of directly to global memory instead of going to LDS to rearrange the data.

And also we try to control, we plan to control the warp scheduling at final granularity.

So this is a comparison between MI250 and A100 for a flash attention forward kernel with 4 batches, 48 heads, head dimension equals to 64, and 0 dropout, and the causal equals to false.

Very specific, and with hardware specifications listed.

So we got slightly higher, better numbers compared to A100, and we also have some optimizations like we are trying to enabling which only enable software pipelining for flush attention and we also expect the performance.

Hey for all the flash attention optimizations you just talk about are they require any changes to the tutorial or try to knob or they are just compiler optimizations.

So the pointer arithmetic stuff requires changes in the Python level.

We try to push everything to the backend compiler, and everything else is in the compiler itself.

Is there any intuition as to why it's so much slower than robust for small MNK?

So basically when I'm OK small we have the workload is small so the world is small.

Last parallelism I know our jam colon is not optimized, so we do expect to high latency like in a inter warp manner.

So when the workload is small, we don't have that many parallelism from neighboring warps, so it's worse.

So after we can optimize gem kernel, We do expect we don't have to rely that heavily on interwarp parallelism stuff.

That's also explains why we are now working, we're not performing better numbers on gem, but we got better numbers on flash retention.

The rug blast is like using a different algorithms there, or why is it so much higher

But like the rug blast is not having worse performance, right?

Like is it using a different approach or?

We're still trying to learn something from Rock Plus and grab their kind of optimizations and try in Triton because we have to tweak it in the Triton system and how to make it work in China system.

That's the challenge here basically, so it's not the same.

Next we have Intel with the Intel XPU back end for Triton by Aiken.


Speaker 0: So for our next talk we have the. AMD back end bringing Triton to AMD GPU's. Again, I'll get the presentation set up and we can get started.

Speaker 1: Was doing most. So. Okay, I can pour. Is there a clicker? Is there a clicker?

Everybody, listen and I from Andy when I'm talking about how we brought into MDG views over the past 612 months. I'll just get the next chart here. Oops. Here we go, OK, so yeah, so we've been working on this for about a year. We originally started looking at Triton.

It was a dependency of deep speed. Deep speed is important to us and supporting our customers. It really kind of kicked into another gear once Pytorch 2 picked it up, right? So. So we have spent a lot of work on it, so this is just a quick update of where the back end is, and then I'll kind of take you through the story of what all we did to bring it.

But so the code is out there, so we have the third party PR 1983. There was merged so we have our third party submodule link. In the open ad trading repository now, which is great So you can have another slide on this on how to check it out, but it's all in there. The actual development happens in our repository, though, so that some modules of link to our Rockham software platform Triton to a specific commit level there. So.

Yeah, so that's yeah, so that's the development. I have another chart on that in a minute. I'll talk about so but the so hardware support we have. We focus on our CD net CDN GPU's which are big data center ones. We do have experimental support for the RDNA as well, so that they both should work.

There's different warp sizes on both of those, so we have auto detection for the hardware implemented in there now, so we. Train is able to find out what hardware were on and do the proper thing there. So this is, but CDNA is our lead development platform on there. We do have full support on PyTorch 2 for all the apps that are lowered to Triton. Working on the hightorch21 at the moment, so all that all that works at the moment, which is which is great.

We did add. We have to change much of the front ends like almost at all right, but we did add this device type of hip. If you're calling directly into the compile, and you wanted to specifically call out the devices that have so. OK, so scoping the port. You know beyond the needs of dependencies reference for deep speed or even Pytorch right does we had to.

We had to kind of see does this map well to our hardware, right? Is this worth doing and the answer is absolutely yes. And what was building this chart? I was looking for some supporting evidence of this and I'd have to look for fills, original paper and even the docs directory out in the GitHub train. GitHub out there is a good source of information on just kind of how things map over, right?

So the original design tenets of, you know, kernel level blocks and locality and efficiency and high occupancy all just translate really well to AMD GPUs. So so in in it that's, you know that was just green light. Let's go ahead and get this optimized and working on our hardware. So Beyond that, beyond the design though. Yeah, what are all the what are the pieces here, right?

What you know? Is there something here that's going to take a long time to port? What are we going to do? So we had to go through and look at everything. So first and foremost, LVM MLIR that those are big pieces, right?

We have MD's already made huge investments there obviously, right? A rock and tool chain is LVM based. We have a rock DL dialect and MIR. We've already. We've already been there.

We've already enabled all that. It works really great. We use that in all sorts of different places today, so. So you know, so beyond that, like is what's higher up the chart and GPU dialects just have some Nvidia specifics in there, but While we were looking through this and porting it, there's we didn't really even touch that very much, so we didn't have to add any new OPS. Although all the OPS just kind of, you know, work the same way as they do on Nvidia.

They lowering the dot op is a little bit different because we want to have. We have our matrix cores versus the tensor core so that there's no topic on that. I'll save that for a minute, but even at that layer we were looking pretty good so. So beyond that, so and then we've got all the different layouts, right? So our shared memory on AMD GPUs is the design is very similar to the Nvidia design.

So all that, the shared memory layout Mostly just kind of worked for us as well. We do have different sizes of shared memory. We in ours we have slightly less shared memory, but we have more general purpose registers. So at this point it came down to. A tuning more of a tuning thing, right?

We want to. We want to prioritize maybe our general purpose registers more than and try to try to be more conservative on the usage of the shared memory in that. On is evident on some of the pipelining passes that we're talking about today already so. And we'll talk about more of that in a second. So then the other thing then is the inline PTX right?

We don't have a PTX so we had to. So the PTX builder is was resonated build a GCN builder, So we've got our GCN builder that would just build in line assembly for our MDGC and assembly so that that was 1 piece we didn't have to go and change, but all this was to map towards our current hardware that we've got in the bottom there. Those are the specs. So it all looked pretty good, so we had to come up with a plan on getting all this working and then see if there's anything else that we were missing that we need to work on. So the first thing was, let's get our empty kernel working right?

This was kind of the bulldozing effort to get everything. Just just we can get a all the pipe all the. We had to, you know, print that doesn't a big deal. We could parameterize that we so we've got. We implemented parameterize and we called the MD terminologies way front, so we front warp sizes interchangeable.

So we've got that that plumbed in and we have we can have auto detection. So on our DNS we have 32 on the CD and then we have 64 and then the kernel launcher code we had to build new template code for the kernel launcher and then that was pretty much that we need to do for the for the empty kernel. We're often running. Get all the OPS working and this was where we need to get that GCM builder in. There's a couple other things that we had to work out.

Some of the we had to disable the existing pipelining paths, for example. And then you know, so that was probably 1 of the larger efforts. Just kind of getting through all that, right? So we've got the plumbing in. We've got all the apps working.

All the tests working. We had to do dot op via. Just kind of vector FMA's originally. So Once we had all that function working, we got to let's make this thing faster, right? Let's implement our matrix cores within Triton.

So we shouldn't talk about that in just a minute, but Triton has the MMA layout, which is a big important piece and get performance, right? We initially we were kind of editing that we ended up having an alternate layout. The MFMA layout we are instructions are in the from the instructions so. That just kind of slots alongside MMA B2 and V1 and V2 and now V3. So we'll have MFA layout as well.

So that that brought us up to speed on our matrix core support and. The original or that we so in which had our flash attention working well now and it's and so we started doing more performance tuning and more optimization, which is lead us where we're at right now, which is trying to maximize our occupancy, trying to maximize our performance and seeing what all we can do to really kind of just tune this all the way up to our distinct hardware. And we should talk more about that in a little bit as well, but this mostly has a new pipe lining. We were doing a streaming pipe lining, which is easier on our shared memory amounts, and then we also have the. Swizzling which the most of the swizzling algorithms Mostly we can reuse, but the instructions are different, so we've got that in and we also have.

A piece where we can edit the number of wavefronts per EU that we can maximize our occupancy on the hardware. OK, so before I turn over to show you some more details, I want to talk about the little bit about the third party back end that I mentioned at the beginning. So this is merged now that there's the commit here where we're pinned to at the moment. This does predate some of the performance stuff we're talking about today, but we're working on a synchronization effort that brings in a lot of the commits for that were committed upstream. The past month, right?

That all the hopper stuff. So there's so much new code there. We're taking a little bit broke some of our stuff, so we're working through that. So what once we get that all synchronized up, we're going to update that commit and then we'll have a kind of everything on the top layer and have that commit from the front end be the fastest 1 that we've got. So look forward to that.

So here's some. There's also on this chart on how to build it, right? So there's a environment variable we use trading code gen India. And then you can just check out recursive and you'll and then build like normal and then you'll get the MD back end enabled. So that's pretty cool.

And thanks to the Intel folks that so we this was kind of a collaborative effort between everybody to try and OK, let's get. Layer, so we've got in. It is the same for us too. We've got this the Triton GP Rockham, which is and there's a lot of duplicate code there, so we'll have to work together to try and reduce that duplication as we go forward so. OK, let's go over the shun.

Is it going to do some of the more lower level details on the port.

Speaker 2: OK, so. Hear me? No. What about now?

Speaker 1: What is this 1? Yeah. That's not work. I don't know.

Speaker 2: OK. Let's talk about optimizations we have done on the AMD GPU. The MIS series.

Speaker 1: Second, please. So.

Speaker 2: AMD GPU, the MIS series. So first, let's talk about the AMD matrix course, which is a counterpart of NVIDIA's Tensor course. It supports a full range of mixed precision gem operations, which is achieved by variations of MFMA instructions. So the key to harnessing the MFMA instruction is its data layout, which refers to the distribution of elements among threads in a warp. So this is an example of data layout for MFMA 32 by 32 by 8, flow 16 instruction.

So as you can see, in general, each thread is assigned on the segment for elements. And for input A, the elements in the same segment is in the same row, and consecutive threads are assigned segments along the M dimension. For input B and output D, elements in the same segment are in the same column, and the consecutive threads are assigned segments along the M dimension. For input B and output D, elements in the same segment are in the same column, and the consecutive threads are assigned segments along the N dimension. With the knowledge of data layout, now let's talk about the problem of chain dot.

So this problem refers to, we want to use the result of the first dot as the input to the second dot. And which input to the second dot makes a difference here. So if we're talking about input B to the second dot, we are good to go because input B and output D have the same data layout. To put it another way, there's no need for redistribution of elements among threads. However, if we're talking about input A to a second dot, we have to go to LDS to rearrange the data among threads.

So this is not good. So let's see how we can avoid this redistribution of the second case, which is what we have in flash attention. At software level, we swap the operands. Now you can see the elements from matrix A are passed to the MFA instruction as input B, and elements from matrix B are passed as input A. Okay, so because of the way how the hardware interprets elements from threads, Now we are computing B transpose times A transpose.

Mathematically, we got D transpose with MFMA output layout. Now is the interesting part. MFMA output layout on D matrix is the same as transposed MFMA output layout on D because threads are held in the same elements. So in this way, the output Is ready to be used as input to the 2nd dot and we achieve transposition without transposing anything in the hardware, so no overhead. Software pipelining, this is a tricky part, so we are developing a 2 stage software pipeliner with 1 buffer in LDS and the other buffer in VGPR.

So the compute always start with data in LDS. It's never start with data in the buffer in VGPR. So the compute includes reading data from LDS and MFMA instructions. Okay, so we don't use the buffer in which PR as a regular buffer, because this is used as a stream buffer to wait for data from global memory. So well, we perform compute for the current iteration of the loop.

Elements from the next iteration. It's loaded, it's being loaded from global memory into VGPR and we always do the competition with data in LDS. So in this way we achieve some kind of 2 buffer software pipelining with only 1 buffer in LDS. So we save LDS space at a price of higher register pressure. Okay, So we also have implemented more optimizations on AMVGPU.

So first, we try to keep Qtensor. So in the context of FlashAttention forward kernel, we keep Qtensor in register instead of loading it from global memory or LDS at every iteration of the loop. We also arrange our, we also arrange warps vertically in the CTA so that we can. Do role reduction more efficiently? We.

Reorder pointer arithmetic so that we can generate global load instructions with maximum vectorization list. Because the hardware is not doing a very good job on vectorization. We are also experimenting with this waves per EU attribute, which limits the number of registers to increase the occupancy. We see great performance improvement when setting this attribute to 3 for flash attention forward curl, which is weird. So because if you don't set it, it's a it will be 2, so occupancy will be 2.

But if we set it to 3, we have some register spilling, but we are increasing the occupancy so we have better performance. OK, performance. On gem performance, So this is kind of internal comparison of Triton and the Rock Plus on dense jams with different M and K values. So as you can see, we are still lagging behind ROC+. And we have larger gaps with small gems.

And this is basically our to-do work. We are optimizing on this kind of stuff all the time. And then we have some to-do optimizations. So first, we plan to store the result of MFMA instructions of directly to global memory instead of going to LDS to rearrange the data. And also we try to control, we plan to control the warp scheduling at final granularity.

Flash attention, forward kernel. So this is a comparison between MI250 and A100 for a flash attention forward kernel with 4 batches, 48 heads, head dimension equals to 64, and 0 dropout, and the causal equals to false. Very specific, and with hardware specifications listed. So we got slightly higher, better numbers compared to A100, and we also have some optimizations like we are trying to enabling which only enable software pipelining for flush attention and we also expect the performance. The gem optimizations will help.

Flash attention for kernel. That's it questions.

Speaker 1: No questions.

Speaker 3: Hey for all the flash attention optimizations you just talk about are they require any changes to the tutorial or try to knob or they are just compiler optimizations.

Speaker 2: Good question. So the pointer arithmetic stuff requires changes in the Python level. This is something we are working on. We try to push everything to the backend compiler, and everything else is in the compiler itself.

Speaker 4: Is there any intuition as to why it's so much slower than robust for small MNK?

Speaker 2: Yeah, yeah, yeah, I have some. So basically when I'm OK small we have the workload is small so the world is small. We have. Last parallelism I know our jam colon is not optimized, so we do expect to high latency like in a inter warp manner. So when the workload is small, we don't have that many parallelism from neighboring warps, so it's worse.

So after we can optimize gem kernel, We do expect we don't have to rely that heavily on interwarp parallelism stuff. That's also explains why we are now working, we're not performing better numbers on gem, but we got better numbers on flash retention. So. Yeah.

Speaker 4: The rug blast is like using a different algorithms there, or why is it so much higher

Speaker 2: so I can say again?

Speaker 4: But like the rug blast is not having worse performance, right? Like is it using a different approach or? Why is it fostered centric and?

Speaker 2: Oh, I think it well. We're still working on well. So our gem kernel is not optimized yet. We're still trying to learn something from Rock Plus and grab their kind of optimizations and try in Triton because we have to tweak it in the Triton system and how to make it work in China system. So that's the work.

That's the challenge here basically, so it's not the same. They are different. And we're trying to catch up. OK.

Speaker 0: Thanks, Lucian and Jason. Next we have Intel with the Intel XPU back end for Triton by Aiken. I'll.

See all Triton transcripts on Youtube

AMD