Using F.scaled_dot_product_attention gives the error RuntimeError: No available kernel. Aborting execution

username: sree_harsha
cooked: <p>Hello,</p>
<p>I’m trying to run the ‘FlashAttention’ variant of the F.scaled_dot_product_attention.<br>
Config = namedtuple(‘FlashAttentionConfig’, [‘enable_flash’, ‘enable_math’, ‘enable_mem_efficient’])’<br>
self.cuda_config = Config(True, False, False)<br>
with torch.backends.cuda.sdp_kernel(**self.cuda_config._asdict()):<br>
x = F.scaled_dot_product_attention(q, k, v)</p>
<p>I am on A100-SXM,<br>
Tried running this with</p>
<ul>
<li>CUDA version 12.0 and PyTorch 2.1.0.dev20230526+cu121</li>
<li>CUDA 11.7 and PyTorch 2.0.1</li>
</ul>
<p>I see no references to this error in general, and not sure what I’m doing wrong?</p>
<p>It works just fine with<br>
Config(False, True, True) → which uses math and memory efficient attention. But I would preferably use Flash Attention.</p>
<p>Would appreciate any help in this context!</p>

username: ptrblck
cooked: <p>Could you check if you are seeing the same issue in the recent nightly binary, please?</p>

username: sree_harsha
cooked: <p>Thank you, its working now. It says ‘reduction over non-contiguous data’ frequently but it seems to be working.</p>
<p>If you don’t mind I have two follow up questions:</p>
<ol>
<li>
<p>Does it hold that the Flash Attention implementation available in PyTorch is only usable with float16/bfloat16 like the original repo’s implementation? Or would it work with Float32 as well?</p>
</li>
<li>
<p>Is this fully compatible with torch.compile?</p>
</li>
</ol>
<p>For those interested, this error no longer occurs after updating to PyTorch nightly version 2.1.0.dev20230527+cu121.</p>

username: kyegomez
cooked: <p>How you can pip install the nightly release? <code>2.1.0.dev20230527+cu121</code></p>

username: sree_harsha
cooked: <p>Just install the latest nightly</p>
<p><code>pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121</code></p>

username: trusira
cooked: <p>I get the same error with Pytorch 2.3.0+cu121 built from source. What could be the possible reason?<br>
I’m running on a cluster of Tesla V100-SXM2-32GB with CUDA version: 12.4.</p>
<pre><code class="lang-auto">with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    attn_out = F.scaled_dot_product_attention(xq, keys, values)
</code></pre>

username: ptrblck
cooked: <p>Based in <a href="https://github.com/pytorch/pytorch/blob/b5fdbc1a9fb94e80f0addd836fee358d2446ea88/aten/src/ATen/native/transformers/cuda/flash_attn/flash_api.cpp#L359-L363">this code</a> Volta GPUs are not supported.<br>
I would generally not recommend trying to force a specific algorithm, but let PyTorch select the fastest one for the used device.</p>

username: Varun_Gumma
cooked: <p>I get the same error with Pytorch 2.3.0, with A100. I read somewhere that <code>flash_attn</code> does not support <code>attn_mask</code>. Could you please help me here?</p>

username: ptrblck
cooked: <p>Are you also trying to call a specific algorithm explicitly? Does it work, if you let PyTorch pick the right one?</p>

username: Varun_Gumma
cooked: <p>Yes, it works if I get Pytorch choose the best algorithm. But flash attention alone does seem to work as it does not support an attention mask separately.</p>