Flash Attention

EnricoShippole: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>I just wanted to confirm what is the best way to ensure that only the new Flash Attention in PyTorch 2.0 is being used for scaled dot product attention:</p>
<p>For example:</p>
<pre><code class="lang-python"># pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale
with torch.backends.cuda.sdp_kernel(
    enable_flash=True, 
    enable_math=False, 
    enable_mem_efficient=False
):
    out = F.scaled_dot_product_attention(
        q, k, v,
        attn_mask = mask,
        dropout_p = flash_attn_dropout, 
        is_causal = causal, 
        scale = scale
    )
</code></pre>
<p>I greatly appreciate your help.</p>
<p>Thank you,</p>
<p>Enrico</p>

ptrblck: <p><code>torch.backends.cuda.enable_flash_sdp</code> is not a context manager and you could use <code>with torch.backends.cuda.sdp_kernel</code> instead as sen here:</p>
<pre><code class="lang-python">print(torch.backends.cuda.flash_sdp_enabled())
# True
print(torch.backends.cuda.mem_efficient_sdp_enabled())
# True
print(torch.backends.cuda.math_sdp_enabled())
# True

with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    print(torch.backends.cuda.flash_sdp_enabled())
    # True
    print(torch.backends.cuda.mem_efficient_sdp_enabled())
    # False
    print(torch.backends.cuda.math_sdp_enabled())
    # False
</code></pre>

EnricoShippole: <p>Thank you for verifying it should be used as such:</p>
<pre><code class="lang-python">with torch.backends.cuda.sdp_kernel(
    enable_flash=True, 
    enable_math=False, 
    enable_mem_efficient=False
):
</code></pre>
<p>Appreciation as always.</p>
<p>Best,</p>
<p>Enrico</p>

kchoi: <p>Hi <a class="mention" href="/u/enricoshippole">@EnricoShippole</a>, did you get any performance improvement by this change to force flash attention only?</p>

EnricoShippole: <p>Hi <a class="mention" href="/u/kchoi">@kchoi</a> ,</p>
<p>You should see both improvements in speed and memory consumption. You can check out some recent small baseline models I trained here: <a href="https://github.com/conceptofmind/PaLM" class="inline-onebox" rel="noopener nofollow ugc">GitHub - conceptofmind/PaLM: An open-source implementation of Google's PaLM models</a></p>
<p>Thank you,</p>
<p>Enrico</p>

kchoi: <p>thanks a lot!</p>
<p>i compared <code>true, false, false</code> (=force to use flash attention) vs <code>false, true, true</code>. as you said, i expected the former should be faster, but but it’s slightly slower (second/iteration is like 7 vs 6).</p>
<p>i also ran <a class="mention" href="/u/ptrblck">@ptrblck</a> 's code snippet to get the same results when printing the lines.</p>
<p>i’m using torch2.0, a100-80gb. is there anything i might be missing?</p>

EnricoShippole: <p>That is interesting. It should be both much faster and more memory-optimized on an A100 (80GB) due to the increased bandwidth. Have you opened an issue on the PyTorch github with the benchmarks for your tests?</p>

EnricoShippole: <p>I have been speaking to a few different peers and they are noticing results similar to yours. I will have to test the Triton version by Tri Dao too.</p>

kchoi: <p>that’s interesting! i didn’t open an issue there (yet). just fyi (or for anyone), the training was with deepspeed 2 and 3, batch size per gpu 2 something, but n_head is like 40 hence the number – i believe, <code>n_head * batch_size</code>, that matters – should be large enough.</p>

EnricoShippole: <p>Can you try with a head dim of 128?</p>

ahmdtaha: <p>I wrote the following toy snippet to eval flash-attention speed up. The code outputs</p>
<pre><code class="lang-auto">Flash attention took 0.0018491744995117188 seconds
Standard attention took 0.6876699924468994 seconds
</code></pre>
<p>Notice the following<br>
1- I am using float16 on cuda, because flash-attention supports float16 and bfloat16<br>
2- Flash-attention aggregates multiple operations into a single fused-kernel. Thus, more operations leads to more savings. In my code snippet, I am doing matmul, softmax, dropout only. I believe further speed up can be gained by adding the mask operation as well.<br>
3- Flash-attention can support longer sequences that standard attention can’t. For instance, my GPU can perform flash-attention with seq_len=4096, but throws OOM error with standard attention</p>
<pre><code class="lang-auto">import time
import torch
import torch.nn.functional as F


bz = 32
seq_len = 2048
dims = 64
n_heads = 8
q = torch.randn(bz, n_heads, seq_len, dims, dtype=torch.float16).cuda()
k = torch.randn(bz, n_heads, seq_len, dims, dtype=torch.float16).cuda()
v = torch.randn(bz, n_heads, seq_len, dims, dtype=torch.float16).cuda()

dropout_rate = 0.2
num_trials = 10

with torch.backends.cuda.sdp_kernel(
    enable_flash=True, enable_math=False, enable_mem_efficient=False
):
    start = time.time()
    for i in range(num_trials):
        out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_rate)
    end = time.time()
    print('Flash attention took {} seconds'.format(end - start))


start = time.time()
for i in range(num_trials):
    attn = q @ k.transpose(-2, -1)
    attn = attn.softmax(dim=-1)
    attn = F.dropout(attn, p=dropout_rate, training=True)
    x = (attn @ v).transpose(1, 2)  # .reshape(bz, seq_len, n_heads*dims)
end = time.time()
print('Standard attention took {} seconds'.format(end - start))

</code></pre>

ptrblck: <p>CUDA kernels are executed asynchronously so you would need to synchronize the code before starting and stopping the host timers. Otherwise you would profile the dispatching, kernel launches, or implicit syncs making your profile invalid.</p>

ahmdtaha: <p>Thanks for catching this issue. I updated the code accordingly. Please let me know if you see other mistakes.</p>
<p>I also switched the order of standard and flash attention evaluations as a sanity check.<br>
The current output is</p>
<pre><code class="lang-auto">Standard attention took 0.8632566928863525 seconds for 10 trials
Flash attention took 0.07728338241577148 seconds for 10 trials
</code></pre>
<p>The updated code snippet is</p>
<pre><code class="lang-auto">import time
import torch
import torch.nn.functional as F


bz = 32
seq_len = 2048
dims = 64
n_heads = 8
q = torch.randn(bz, n_heads, seq_len, dims, dtype=torch.float16).cuda()
k = torch.randn(bz, n_heads, seq_len, dims, dtype=torch.float16).cuda()
v = torch.randn(bz, n_heads, seq_len, dims, dtype=torch.float16).cuda()

dropout_rate = 0.2
num_trials = 10


torch.cuda.synchronize()
start = time.time()
for i in range(num_trials):
    attn = q @ k.transpose(-2, -1)
    attn = attn.softmax(dim=-1)
    attn = F.dropout(attn, p=dropout_rate, training=True)
    x = (attn @ v).transpose(1, 2)  # .reshape(bz, seq_len, n_heads*dims)
torch.cuda.synchronize()
end = time.time()
print('Standard attention took {} seconds for {} trials'.format(end - start, num_trials))

with torch.backends.cuda.sdp_kernel(
    enable_flash=True, enable_math=False, enable_mem_efficient=False
):
    torch.cuda.synchronize()
    start = time.time()
    for i in range(num_trials):
        out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_rate)
    torch.cuda.synchronize()
    end = time.time()
    print('Flash attention took {} seconds for {} trials'.format(end - start, num_trials))

</code></pre>

yysirs: <pre><code class="lang-auto">with torch.backends.cuda.sdp_kernel(
            enable_flash=True, enable_math=False, enable_mem_efficient=False
        ):
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, dropout_p=self.dropout if self.training else 0.0, is_causal = False)
</code></pre>
<p>I try run code, but error, the error info here:</p>
<pre><code class="lang-auto">&lt;string&gt;:1: UserWarning: Memory efficient kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.h:545.)
&lt;string&gt;:1: UserWarning: Memory Efficient attention has been runtime disabled. (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.h:338.)
&lt;string&gt;:1: UserWarning: Flash attention kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.h:547.)
&lt;string&gt;:1: UserWarning: Both fused kernels do not support non-null attn_mask. (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.h:191.)
Traceback (most recent call last):
  File "&lt;string&gt;", line 1, in &lt;module&gt;
RuntimeError: No available kernel.  Aborting execution.
</code></pre>
<p>but, when I remove atten_mask paramters, it work.</p>
<pre><code class="lang-auto">with torch.backends.cuda.sdp_kernel(
            enable_flash=True, enable_math=False, enable_mem_efficient=False
        ):
       out = F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout if self.training else 0.0, is_causal = False)
</code></pre>
<p>attention_mask shape is  [bz, seq_len. target_len ,src_len]</p>

yysirs: <p>if I want to add <code>atten_mask</code> parameter, what shuold I do?</p>

Jackmin801: <p>I know its been forever.<br>
But only the <code>math</code> and <code>meff</code> kernel supports the <code>attn_mask</code> parameter.</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/96099#issuecomment-1498388194">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/96099#issuecomment-1498388194" 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="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/96099#issuecomment-1498388194" target="_blank" rel="noopener nofollow ugc">Add support for ALiBi/relative positional biases to the fast path for Transformers</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2023-03-06" data-time="15:01:54" data-timezone="UTC">03:01PM - 06 Mar 23 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2023-08-02" data-time="00:04:55" data-timezone="UTC">12:04AM - 02 Aug 23 UTC</span>
        </div>

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

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: transformer/mha
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🚀 The feature, motivation and pitch

Recent work in transformer architecture<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">s such as T5 [1] and Attention with Linear Biases (ALiBi) [2] have shown that moving the positional encodings from the word embedding layer and directly into the self-attention computation improves the capacity of models to extrapolate to longer sentences while keeping lower perplexity scores. Both T5 and ALiBi add biases to the logits resulting from the QK multiplication, turning the softmax operation into something like this:

weights = softmax(QK^T + B), where B can either be a constant pre-defined matrix, like in ALiBi, or a learnable bias, like in T5. In both cases, the values of B are dependent on the relative distance between tokens, with the farthest tokens usually having the highest penalties.

Usually, this is implemented through modification of the attention mask to add these extra weights on top, changing it from a boolean yes/no attention to a weighted kind of model.

Performance-wise, our experiments are summarized in the graph below, with @daviswer explaining: "Here we train 3 AliBi decoders (390M params) with different sequence lengths (2048, 1024, 512) but perform validation at sequence length 2048 all the time. This allows us to track the degree of overfitting to sequence length over time. We then do the same for 3 equivalent models using T5 relative positional bias (rpb) instead of AliBi. While val perplexity of both variants suffers greatly over time with diverging train/val sequence lengths, the degree of increase is worse for rpb, indicating that while generalization to longer sequence lengths is a difficult problem, AliBi is indeed better at it than rpb." From the ALiBi paper we also know that both perform better than the traditional word embedding-level positional encodings.

![loss_curve (1)](https://user-images.githubusercontent.com/919977/223146953-d460bf47-29fe-482f-8f85-13376c686bab.png)

As of right now, the BetterTransformers path in PyTorch only supports two models for the attention mask: none at all, or the one used for causal LMs. We propose modifying the current implementation so that more complex masks like the ones required for relative positional biases and ALiBi can also be supported.

cc: @daviswer @supriyogit @raghukiran1224 @mudhakar @mayank31398 @cpuhrsch @HamidShojanazeri

[1] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (https://arxiv.org/abs/1910.10683)

[2] Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (https://arxiv.org/abs/2108.12409)

### Alternatives

_No response_

### Additional context

_No response_

cc @jbschlosser @bhosmer @cpuhrsch @erichan1</span></p>
  </div>

  </article>

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

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


Ezra_Chua: <p>Noob question. Why cant you set everything to True? like so. Wouldn’t that make it more memory efficient?<br>
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):</p>

MinSnz: <p><code>flash</code> / <code>math</code> / <code>mem_efficient</code> are different backends. By setting all to <code>True</code>, you are letting PyTorch choose the most favorable one; and by setting one to <code>True</code>, you are forcing a backend and letting it fail if not available. Usually you want to force flash attention for the best speed and check why it may fail.</p>

Nader_El_Zik: <p>Hello please i need your answer i did this code on GPU colab T4 and i get this error still:<br>
RuntimeError: No available kernel. Aborting execution.</p>