Scaled_dot_product_attention

AlbertZeyer: <p>In a recent PyTorch version (since when exactly?), to use an efficient attention implementation, you can simply use <a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html" rel="noopener nofollow ugc"><code>torch.nn.functional.scaled_dot_product_attention</code></a>, right? As I understand, it would automatically use <a href="https://arxiv.org/abs/2307.08691" rel="noopener nofollow ugc">FlashAttention-2</a>:</p>
<blockquote>
<p>automatically select the most optimal implementation based on the inputs</p>
</blockquote>
<p>I’m not sure exactly what this means though. How exactly is the logic? In what cases would it select FlashAttention-2?</p>
<p>Also, as far as I know, FlashAttention-2 only works on more recent Nvidia GPUs but does not on older. (Since what kind GPU? I guess 1080 not? 2080?)</p>
<p>Alternatively, there is also <a href="https://github.com/facebookresearch/xformers" rel="noopener nofollow ugc">Memory-Efficient Attention</a>. How does this compare in speed? This works on a 1080?</p>
<p>What is the most efficient implementation for a 1080, or what would be a reasonable choice there? Should I just use <code>torch.nn.functional.scaled_dot_product_attention</code> or sth more custom?</p>
<p>I have seen many other custom implementations, for example in <a href="https://github.com/pytorch-labs/segment-anything-fast/blob/2ad6b94306dd90adbf250a818b6678619054fd12/segment_anything_fast/flash_4.py#L335" rel="noopener nofollow ugc">segment_anything_fast</a>. I’m not really sure if this is outdated (i.e. obsolete with recent PyTorch) or still makes sense. I think I also have seen a fast Triton implementation somewhere.</p>
<p>In some cases, I also need self-attention with relative positional encoding. As far as I understand, <code>torch.nn.functional.scaled_dot_product_attention</code> should support that. But I can imagine that not all implementations would support it. I also need both cross-attention and self-attention. (But it would be ok for me to use different implementations for each of those cases.)</p>
<p>Does <code>torch.nn.functional.scaled_dot_product_attention</code> work for training as well, i.e. the gradient is defined?</p>
<p>I saw that <code>torch.nn.MultiheadAttention</code> for some reason does not use the native attention function when training is enabled. Why?</p>

eqy: <p>The dispatching logic is here:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/ce503c1b40207dab770c28cbd4568cd9e105277b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp#L556">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/ce503c1b40207dab770c28cbd4568cd9e105277b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp#L556" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/ce503c1b40207dab770c28cbd4568cd9e105277b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp#L556" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/ce503c1b40207dab770c28cbd4568cd9e105277b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp#L556</a></h4>



    <pre class="onebox"><code class="lang-cpp">
      <ol class="start lines" start="546" style="counter-reset: li-counter 545 ;">
          <li>    }</li>
          <li>  }</li>
          <li></li>
          <li>  auto dprop = at::cuda::getCurrentDeviceProperties();</li>
          <li>  if (dprop-&gt;major &gt;= 8) {</li>
          <li>    return check_tensor_dtype(params, greater_than_or_equal_sm80_mem_efficient_dtypes, debug);</li>
          <li>  }</li>
          <li>  return check_tensor_dtype(params, less_than_sm80_mem_efficient_dtypes, debug);</li>
          <li>}</li>
          <li></li>
          <li class="selected">SDPBackend select_sdp_backend(sdp_params const&amp; kernel_params) {</li>
          <li>  // This function defines the priority order of the different sdp backends</li>
          <li>  // 1. Flash Attention</li>
          <li>  // 2. Mem Efficient Attention</li>
          <li>  // 3. Math fallback</li>
          <li>  auto&amp; ctx = at::globalContext();</li>
          <li>  if (!ctx.userEnabledMathSDP() &amp;&amp; !ctx.userEnabledFlashSDP() &amp;&amp;</li>
          <li>      !ctx.userEnabledMemEfficientSDP() &amp;&amp; !ctx.userEnabledCuDNNSDP()) {</li>
          <li>    return SDPBackend::error;</li>
          <li>  }</li>
          <li>  // Get ideal kernel ordering</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>On pre-sm80 GPUs like the 1080 (sm60) or 2080 (sm75) the mem efficient attention backend could be used (depending on shapes and other constraints) but flash attention requires sm80 and newer.</p>