Scaled_dot_product_attention - correct input shape

username: jws
cooked: <p>The documentation of <code>scaled_dot_product_attention</code> suggests the following dimensions for inputs:</p>
<ul>
<li>query: <code>(N,...,L,E)</code></li>
<li>key: <code>(N,...,S,E)</code></li>
<li>value: <code>(N,...,S,Ev)</code><br>
So these are <em>three</em> dimensional. However, when running either the flash or the memory efficient implementation, I get exceptions along the lines of <code>Both fused kernels requires query, key and value to be 4 dimensional, but got Query dim: 3, Key dim: 3, Value dim: 3 instead.</code></li>
</ul>
<p>Question: What are the four dimensions? Testing suggests strongly that it’s <code>batch, heads, sequence, embedding</code> which also agrees with various code examples.</p>
<p>However, the closest to actual documentation I find is the documentation for xformers, which suggests <code>batch, sequence, heads, embedding</code> (<a href="https://facebookresearch.github.io/xformers/components/ops.html#module-xformers.ops" class="inline-onebox" rel="noopener nofollow ugc">xFormers optimized operators | xFormers 0.0.24 documentation</a>).</p>

username: shivammehta007
cooked: <aside class="quote no-group" data-username="jws" data-post="1" data-topic="194396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jws/48/66534_2.png" class="avatar"> jws:</div>
<blockquote>
<p>query: <code>(N,...,L,E)</code></p>
</blockquote>
</aside>
<p>The representation contains three dots <code>...</code> which suggests you can add multiple dimensions in between and since the function works on the last two it doesn’t matter for that purpose.</p>
<aside class="quote no-group" data-username="jws" data-post="1" data-topic="194396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jws/48/66534_2.png" class="avatar"> jws:</div>
<blockquote>
<p>Question: What are the four dimensions? Testing suggests strongly that it’s <code>batch, heads, sequence, embedding</code> which also agrees with various code examples.</p>
</blockquote>
</aside>
<p>It is correct, the general order is this, where the number of heads can be any value. Perhaps for some reason you have any other dimension it still doesn’t matter the function will be able to compute the attention scores for the <code>L, S, E</code> dimensions. Example:</p>
<pre data-code-wrap="python"><code class="lang-python">&gt;&gt;&gt; k = torch.randn(3, 4, 25, 64) # b, h, L, E
&gt;&gt;&gt; q = torch.randn(3, 4, 28, 64) # b, h, S, E
&gt;&gt;&gt; v = torch.randn(3, 4, 28, 64) # b, h, S, E
&gt;&gt;&gt; x1 = F.scaled_dot_product_attention(k, q, v)
&gt;&gt;&gt; x1.shape
# torch.Size([3, 4, 25, 64])
</code></pre>
<p>But even if you add more dimensions before the last two dimensions the computation would not change.</p>
<pre data-code-wrap="python"><code class="lang-python">&gt;&gt;&gt; x2 = F.scaled_dot_product_attention(k.view(3, 2, 2, 25, 64), q.view(3, 2, 2, 28, 64), v.view(3, 2, 2, 28, 64))
&gt;&gt;&gt; x2.shape
torch.Size([3, 2, 2, 25, 64])

&gt;&gt;&gt; (x1.flatten() == x2.flatten()).all()
# True
</code></pre>
<p>Hope this helps!</p>

username: jws
cooked: <p>Thanks. I think we have a slight misunderstanding. My point was, that at least for the torch version I have (<code>'2.1.0+cu121'</code>) <em>some</em> of the kernels behind <code>scaled_dot_product_attention</code> mandate 4 dimensional inputs.</p>
<p>Sample code. When selecting memory efficient attention (V100 GPU), using anything except four-dim. inputs leads to a runtime error.</p>
<pre><code class="lang-auto">import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

assert torch.cuda.is_available()
device = torch.cuda.current_device()

torch.manual_seed(42)

mu, sigma = 100., 1.0

# Four-dim. tensors -- works fine
q = torch.normal(mu, sigma, size=(4, 6, 8, 2), device=device)
k = torch.normal(mu, sigma, size=(4, 6, 8, 2), device=device)
v = torch.normal(mu, sigma, size=(4, 6, 8, 2), device=device)

with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
    attn = F.scaled_dot_product_attention(q, k, v)
print('Success for 4D')

# Three-dim. tensors -- yields runtime error
q = torch.normal(mu, sigma, size=(4, 6, 8), device=device)
k = torch.normal(mu, sigma, size=(4, 6, 8), device=device)
v = torch.normal(mu, sigma, size=(4, 6, 8), device=device)

with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
    attn = F.scaled_dot_product_attention(q, k, v)
</code></pre>
<p>Error message:</p>
<pre><code class="lang-auto">/tmp/ipykernel_746318/3205699521.py:28: UserWarning: Memory efficient kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:367.)
  attn = F.scaled_dot_product_attention(q, k, v)
/tmp/ipykernel_746318/3205699521.py:28: UserWarning: Both fused kernels requires query, key and value to be 4 dimensional, but got Query dim: 3, Key dim: 3, Value dim: 3 instead. (Triggered internally at ../aten/src/ATen/native/transformers/sdp_utils_cpp.h:273.)
  attn = F.scaled_dot_product_attention(q, k, v)
/tmp/ipykernel_746318/3205699521.py:28: UserWarning: Flash attention kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:369.)
  attn = F.scaled_dot_product_attention(q, k, v)
/tmp/ipykernel_746318/3205699521.py:28: UserWarning: Flash attention has been runtime disabled. (Triggered internally at ../aten/src/ATen/native/transformers/sdp_utils_cpp.h:425.)
  attn = F.scaled_dot_product_attention(q, k, v)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[6], line 28
     25 v = torch.normal(mu, sigma, size=(4, 6, 8), device=device)
     27 with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
---&gt; 28     attn = F.scaled_dot_product_attention(q, k, v)

RuntimeError: No available kernel.  Aborting execution.
</code></pre>

username: nscorley
cooked: <p>Hi <a class="mention" href="/u/jws">@jws</a> ,</p>
<p>Did you ever find a solution to this problem? I’m encountering the same issue, although I have 5 dimensional query, key, and value vectors (vs. 4).</p>
<p>Thanks!</p>