Using F.scaled_dot_product_attention gives the error

whatisslove11: <p>Hello,</p>
<p>I’m trying to run the ‘FlashAttention’ variant of the F.scaled_dot_product_attention by the code below:</p>
<pre><code class="lang-auto">import torch

B = 4
H = 12
N = 2**12
D = 1024

dtype = torch.float32
device = "cuda"

q = torch.rand(B, H, N, D, dtype=dtype, device=device)
k = torch.rand(B, H, N, D, dtype=dtype, device=device)
v = torch.rand(B, H, N, D, dtype=dtype, device=device)

m = torch.ones((N, N), dtype=torch.bool, device=device).triu(N - N + 1)
m = m.float().masked_fill(m, float("-inf"))

with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
    out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=m)
</code></pre>
<p>I tried to run this code in two machines:</p>
<ol>
<li>Google Colab with Tesla T4</li>
<li>my own PC with 4080</li>
</ol>
<p>I found this topic: <a href="https://discuss.pytorch.org/t/using-f-scaled-dot-product-attention-gives-the-error-runtimeerror-no-available-kernel-aborting-execution/180900/1" class="inline-onebox">Using F.scaled_dot_product_attention gives the error RuntimeError: No available kernel. Aborting execution</a> and tried to install nightly version by the command <code>pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121</code> from the comments. Either I used wrong command to install right version or there is new problem for using FlashAttention from pytorch</p>
<p>I will be grateful for your answers!</p>

soulitzer: <p>Do you have a error/stack trace? If it’s the same error as the post, i.e., “No available kernel”, there should’ve been warnings being printed telling you why FlashAttention isn’t being run. If so, what do those warnings say?</p>

whatisslove11: <p>Hello!<br>
Whole error code + pytorch version is shown in the screenshot below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/9/b999ed56da95360e519a6d936504730e5d27747f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b999ed56da95360e519a6d936504730e5d27747f" title="erorr"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b999ed56da95360e519a6d936504730e5d27747f_2_660x500.jpeg" alt="erorr" data-base62-sha1="qtU7oJylvuds0JGkXeOne8cM68T" width="660" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b999ed56da95360e519a6d936504730e5d27747f_2_660x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b999ed56da95360e519a6d936504730e5d27747f_2_990x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/9/b999ed56da95360e519a6d936504730e5d27747f.jpeg 2x" data-dominant-color="EEE7E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">erorr</span><span class="informations">1081×818 166 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, when I used enable_mem_efficent=True instead of enable_flash=True, there is no errors</p>

whatisslove11: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> sorry for your tagging, maybe do you have any ideas how to solve this problem?</p>

soulitzer: <p>Flash attention is not run because it does not support attn_mask</p>