FlashAttention on AMD MI250X not working with pytorch 2.2.0

username: jpata
cooked: <p>Hi,</p>
<p>Using the latest pytorch 2.2.0, installed from pip, running on AMD MI250X on a HPC system, I find that FlashAttention does not work:</p>
<pre><code class="lang-auto">/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:5476: UserWarning: 1Torch was not compiled with memory efficient attention. (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:320.)
  attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:5476: UserWarning: Memory efficient kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:416.)
  attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:5476: UserWarning: Flash attention kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/hip/sdp_utils.cpp:418.)
  attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
Epoch 1 train loop on rank=0:   0%|          | 0/45943 [00:01&lt;?, ?it/s]
...
    attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
RuntimeError: No available kernel. Aborting execution.
</code></pre>
<p>When I check the GCN arch and whether flash sdp is enabled, I get:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; print("GCN arch:", torch.cuda.get_device_properties('cuda').gcnArchName)
GCN arch: gfx90a:sramecc+:xnack-
&gt;&gt;&gt; print("FlashAttention available:", torch.backends.cuda.flash_sdp_enabled())
FlashAttention available: True
&gt;&gt;&gt; torch.__version__
'2.2.0+rocm5.7'
</code></pre>
<p>I see that FlashAttention support for ROCm was being worked on:</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/114309">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/114309" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/pytorch/pull/114309" target="_blank" rel="noopener nofollow ugc">Initial Flash Attention support on ROCM</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>ROCm:xinyazhang/up-fa-mathaot</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2023-11-21" data-time="22:25:39" data-timezone="UTC">10:25PM - 21 Nov 23 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/xinyazhang" target="_blank" rel="noopener nofollow ugc">
            <img alt="xinyazhang" src="https://avatars.githubusercontent.com/u/11429462?v=4" class="onebox-avatar-inline" width="20" height="20">
            xinyazhang
          </a>
        </div>

        <div class="lines" title="166 commits changed 14 files with 854 additions and 23 deletions">
          <a href="https://github.com/pytorch/pytorch/pull/114309/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+854</span>
            <span class="removed">-23</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">This pull requests add initial Flash Attention support for AMD/ROCM platform. It<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/114309" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden"> 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



cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
but it was reverted, and work was done to readd it, although the PR was closed before merging (?):</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/115981">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/115981" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/pytorch/pull/115981" target="_blank" rel="noopener nofollow ugc">Re-add initial Flash Attention support on ROCM</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>ROCm:xinyazhang/up-fa-mathaot</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2023-12-16" data-time="09:25:12" data-timezone="UTC">09:25AM - 16 Dec 23 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/xinyazhang" target="_blank" rel="noopener nofollow ugc">
            <img alt="xinyazhang" src="https://avatars.githubusercontent.com/u/11429462?v=4" class="onebox-avatar-inline" width="20" height="20">
            xinyazhang
          </a>
        </div>

        <div class="lines" title="13 commits changed 21 files with 885 additions and 38 deletions">
          <a href="https://github.com/pytorch/pytorch/pull/115981/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+885</span>
            <span class="removed">-38</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Note about the Updates:

This PR:
1. skips more flash attention related UTs o<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/115981" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">n 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 (https://github.com/pytorch/pytorch/pull/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

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>so I’m not sure if this is supposed to work yet or not with pytorch 2.2.0?<br>
Any AMD folks (<span class="mention">@xinyazhang</span> <a class="mention" href="/u/jithunnair-amd">@jithunnair-amd</a>) can confirm?</p>
<p>Thanks!</p>

username: jpata
cooked: <p>Answering myself in case anyone else stumbles here.</p>
<p>As far as I can tell, the pytorch wheels supplied at <a href="https://download.pytorch.org/whl/rocm5.7" rel="noopener nofollow ugc">https://download.pytorch.org/whl/rocm5.7</a> do not have Flash Attention for ROCM/AMD builtin so it’s not callable via <code>torch.nn.functional.scaled_dot_product_attention</code> directly.</p>
<p>I was successful in installing <code>ROCmSoftwarePlatform/flash-attention.git</code> on top of the docker image <code>rocm/pytorch:rocm6.0.2_ubuntu22.04_py3.10_pytorch_2.1.2</code> and thus replacing <code>torch.nn.MultiheadAttention -&gt; flash_attn.modules.mha.MHA</code>, but this means that extra steps still have to be taken for running Flash Attention on AMD and there is no full API compatibility as of 2.2.0.</p>
<p>Hopefully subsequent releases will have Flash Attention for AMD builtin… <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>