Add support for Flash Attention for AMD/ROCm #112997

chauhang · 2023-11-05T21:07:53Z

🚀 The feature, motivation and pitch

Enable support for Flash Attention Memory Efficient and SDPA kernels for AMD GPUs.

At present using these gives below warning with latest nightlies (torch==2.2.0.dev20231105+rocm5.6, pytorch-triton-rocm==2.1.0+34f8189eae):

model.py:187: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:253.)
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
model.py:187: UserWarning: 1Torch was not compiled with memory efficient attention. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:291.)
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)

ROCm already has an implementation of Tri's FA here: https://github.com/ROCmSoftwarePlatform/flash-attention/tree/flash_attention_for_rocm2#amd-gpurocm-support

Alternatives

User's have to manually install the ROCm version of FA and use that in their code, vs using the native PyTorch APIs.

Additional context

The ROCM build currently has the FA related flags turned off by default: https://github.com/pytorch/pytorch/blob/main/CMakeLists.txt#L741-L750

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang

hongxiayang · 2023-11-06T22:21:40Z

We are working on it internally. We will update upstream after that.

This pull requests add initial Flash Attention support for AMD/ROCM platform. It added a specialized Triton repository/branch as a compile-time dependency for Flash Attention math library on AMD/ROCM. This triton submodule is not used at runtime and will not be shipped to the final pytorch package. We have the plan to release this specialized Triton as a separate project. Know limitations: - [ ] Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`. - [ ] Only supports power of two sequence lengths. - [ ] No support for varlen APIs. - [ ] Only support head dimension 16,32,64,128. - [ ] Performance is still being optimized. Fixes #112997 Pull Request resolved: #114309 Approved by: https://github.com/jeffdaily, https://github.com/malfet --------- Co-authored-by: Joseph Groenenboom <[email protected]>

This pull requests add initial Flash Attention support for AMD/ROCM platform. It added a specialized Triton repository/branch as a compile-time dependency for Flash Attention math library on AMD/ROCM. This triton submodule is not used at runtime and will not be shipped to the final pytorch package. We have the plan to release this specialized Triton as a separate project. Know limitations: - [ ] Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`. - [ ] Only supports power of two sequence lengths. - [ ] No support for varlen APIs. - [ ] Only support head dimension 16,32,64,128. - [ ] Performance is still being optimized. Fixes pytorch#112997 Pull Request resolved: pytorch#114309 Approved by: https://github.com/jeffdaily, https://github.com/malfet --------- Co-authored-by: Joseph Groenenboom <[email protected]>

Note about the Updates: This PR: 1. skips more flash attention related UTs on MI200 2. Fix additional ATen compiling errors after hipification 3. Fix the author "root" of a specific commit 4. Includes the patch from Nikita in favor of block level static initialization. CAVEAT: This revised PR has a commit that modifies the CI to force its running on MI200 nodes. That specific commit must be reverted before merge. Original PR (#114309) Note: This pull requests add initial Flash Attention support for AMD/ROCM platform. It added a specialized Triton repository/branch as a compile-time dependency for Flash Attention math library on AMD/ROCM. This triton submodule is not used at runtime and will not be shipped to the final pytorch package. We have the plan to release this specialized Triton as a separate project. Know limitations: - Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`. - Only supports power of two sequence lengths. - No support for varlen APIs. - Only support head dimension 16,32,64,128. - Performance is still being optimized. Fixes #112997 Pull Request resolved: #115981 Approved by: https://github.com/malfet

This patch addresses the major limitations in our previous [PR #115981](#115981) through the new dedicated repository [AOTriton](https://github.com/ROCm/aotriton) - [x] Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`). * MI300X is supported. More architectures will be added once Triton support them. - [x] Only supports power of two sequence lengths. * Now it support arbitrary sequence length - [ ] No support for varlen APIs. * varlen API will be supported in the next release of AOTriton - [x] Only support head dimension 16,32,64,128. * Now it support arbitrary head dimension <= 256 - [x] Performance is still being optimized. * Kernel is selected according to autotune information from Triton. Other improvements from AOTriton include * Allow more flexible Tensor storage layout * More flexible API This is a more extensive fix to #112997 Pull Request resolved: #121561 Approved by: https://github.com/malfet, https://github.com/atalman

chauhang · 2024-03-24T03:11:21Z

@jeffdaily @malfet Are the torch nightlies for rocm getting built with the correct flags for SDPA support? I am still getting this error with the latest torch==2.4.0.dev20240323+rocm6.0

UserWarning: 1Torch was not compiled with memory efficient attention. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:515.)
  output = nn.functional.scaled_dot_product_attention(

[pip3] pytorch-triton-rocm==3.0.0+0a22a91d04
[pip3] torch==2.4.0.dev20240323+rocm6.0
[conda] pytorch-triton-rocm       3.0.0+0a22a91d04          pypi_0    pypi
[conda] torch                     2.4.0.dev20240323+rocm6.0          pypi_0    pypi

xinyazhang · 2024-03-26T00:32:53Z

Are the torch nightlies for rocm getting built with the correct flags for SDPA support? I am still getting this error with the latest torch==2.4.0.dev20240323+rocm6.0

Hi @chauhang , If you're referring to memory efficient attention this is expected. This PR only adds Flash Attention to PyTorch and we are working on Memory efficient attention support now.

jithunnair-amd · 2024-03-27T03:18:51Z

@chauhang Since this issue is regarding FA support for ROCm in nightlies, and PR #121561 has been reverted, I'm reopening this issue.

This patch addresses the major limitations in our previous [PR #115981](#115981) through the new dedicated repository [AOTriton](https://github.com/ROCm/aotriton) - [x] Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`). * MI300X is supported. More architectures will be added once Triton support them. - [x] Only supports power of two sequence lengths. * Now it support arbitrary sequence length - [ ] No support for varlen APIs. * varlen API will be supported in future release of AOTriton - [x] Only support head dimension 16,32,64,128. * Now it support arbitrary head dimension <= 256 - [x] Performance is still being optimized. * Kernel is selected according to autotune information from Triton. Other improvements from AOTriton include * Allow more flexible Tensor storage layout * More flexible API This is a more extensive fix to #112997 Pull Request resolved: #121561 Approved by: https://github.com/huydhn

This patch addresses the major limitations in our previous [PR pytorch#115981](pytorch#115981) through the new dedicated repository [AOTriton](https://github.com/ROCm/aotriton) - [x] Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`). * MI300X is supported. More architectures will be added once Triton support them. - [x] Only supports power of two sequence lengths. * Now it support arbitrary sequence length - [ ] No support for varlen APIs. * varlen API will be supported in future release of AOTriton - [x] Only support head dimension 16,32,64,128. * Now it support arbitrary head dimension <= 256 - [x] Performance is still being optimized. * Kernel is selected according to autotune information from Triton. Other improvements from AOTriton include * Allow more flexible Tensor storage layout * More flexible API This is a more extensive fix to pytorch#112997 Pull Request resolved: pytorch#121561 Approved by: https://github.com/huydhn

jinsong-mao · 2024-05-22T06:34:38Z

@xinyazhang @jithunnair-amd Hi the default flash_attention backward is using the kernel from aotriton bwd_kernel_dk_dv now, https://github.com/ROCm/aotriton/blob/6035fb4aa021d39a18adce883e059d5b99a8192c/v2python/rules/flash/bwd_kernel_dk_dv.py#L7. it's very low efficiency, any optimization configuration or usage should us do to improve the performance?

jeffdaily · 2024-05-22T15:31:09Z

CC @groenenboomj

maxall41 · 2024-10-27T19:15:53Z

@xinyazhang Any updates on implementing memory efficient attention?

xinyazhang · 2024-10-28T16:38:30Z

@xinyazhang Any updates on implementing memory efficient attention?

ME and FA are fundamentally the same algorithm but providing different feature sets and on ROCm they are both implemented in AOTrition. The support is already included in PyTorch 2.4 #124885

xinyazhang · 2024-10-28T16:40:57Z

it's very low efficiency, any optimization configuration or usage should us do to improve the performance?

AOTriton 0.7.1 https://github.com/ROCm/aotriton/releases/tag/0.7.1b made some improvements to the bwd kernel, which is ABI compatible (and only compatible) with AOTriton 0.7b and can be used as drop-in replacement for official PyTorch 2.5 wheel.

hongxiayang · 2024-12-03T19:28:29Z

close it as this is supported now

pytorch-bot bot added ciflow/rocm module: rocm AMD GPU support for Pytorch labels Nov 5, 2023

github-project-automation bot added this to PyTorch on ROCm Nov 5, 2023

colesbury added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Nov 6, 2023

hongxiayang moved this to Todo in PyTorch on ROCm Nov 7, 2023

Charmandrigo mentioned this issue Nov 7, 2023

Adding support for FlashAttention (flash-attn) optimization AUTOMATIC1111/stable-diffusion-webui#11902

Closed

4 tasks

xinyazhang mentioned this issue Nov 21, 2023

Initial Flash Attention support on ROCM #114309

Merged

5 tasks

malfet closed this as completed in #114309 Dec 14, 2023

github-project-automation bot moved this from Todo to Done in PyTorch on ROCm Dec 14, 2023

xinyazhang mentioned this issue Dec 16, 2023

Re-add initial Flash Attention support on ROCM #115981

Closed

xinyazhang mentioned this issue Mar 9, 2024

Add Flash Attention support on ROCM #121561

Closed

5 tasks

jithunnair-amd reopened this Mar 27, 2024

jeroen-mostert mentioned this issue Aug 5, 2024

Patches from 6.1.1 to 6.1.2, watch for regressions lamikr/rocm_sdk_builder#125

Closed

TNT3530 mentioned this issue Aug 9, 2024

CUDA OOM and possible solution -- diffusers cli_demo.py with Nvidia 3090 24GB THUDM/CogVideo#92

Closed

2 tasks

jin-eld mentioned this issue Sep 6, 2024

Question: How to run on AMD GPUs? How to enable multi GPU use? THUDM/CogVideo#223

Closed

OrenLeung mentioned this issue Sep 8, 2024

[rocm] Unusable torch.ops.aten._scaled_dot_product_flash_attention_backward at 9.6TFLOPs #135431

Open

hongxiayang closed this as completed Dec 3, 2024

darren-amd mentioned this issue Dec 3, 2024

When would you support USE_FLASH_ATTENTION compile? ROCm/pytorch#1252

Closed

taylding-amd mentioned this issue Dec 3, 2024

Add support for memory efficient attention for AMD/ROCm ROCm/pytorch#1390

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for Flash Attention for AMD/ROCm #112997

Add support for Flash Attention for AMD/ROCm #112997

chauhang commented Nov 5, 2023 •

edited by pytorch-bot bot

Loading

hongxiayang commented Nov 6, 2023

chauhang commented Mar 24, 2024 •

edited

Loading

xinyazhang commented Mar 26, 2024

jithunnair-amd commented Mar 27, 2024

jinsong-mao commented May 22, 2024

jeffdaily commented May 22, 2024

maxall41 commented Oct 27, 2024

xinyazhang commented Oct 28, 2024

xinyazhang commented Oct 28, 2024

hongxiayang commented Dec 3, 2024

Add support for Flash Attention for AMD/ROCm #112997

Add support for Flash Attention for AMD/ROCm #112997

Comments

chauhang commented Nov 5, 2023 • edited by pytorch-bot bot Loading

🚀 The feature, motivation and pitch

Alternatives

Additional context

hongxiayang commented Nov 6, 2023

chauhang commented Mar 24, 2024 • edited Loading

xinyazhang commented Mar 26, 2024

jithunnair-amd commented Mar 27, 2024

jinsong-mao commented May 22, 2024

jeffdaily commented May 22, 2024

maxall41 commented Oct 27, 2024

xinyazhang commented Oct 28, 2024

xinyazhang commented Oct 28, 2024

hongxiayang commented Dec 3, 2024

chauhang commented Nov 5, 2023 •

edited by pytorch-bot bot

Loading

chauhang commented Mar 24, 2024 •

edited

Loading