Repository: facebookresearch/xformers

Issue:
What is actually being calculated in memory_efficient_attention?
In the docs of the memory_efficient_attention function it says:

```
:Equivalent pytorch code:

    .. code-block:: python

        scale = 1 / query.shape[-1] ** 0.5
        query = query * scale
        attn = query @ key.transpose(-2, -1)
        if attn_bias is not None:
            attn = attn + attn_bias
        attn = attn.softmax(-1)
        attn = F.dropout(attn, p)
        return attn @ value
```

But when I create a minimal example this doesn't seem to give the same result as what's actually calculated:

```
import torch

from xformers.ops import memory_efficient_attention

def equiv_fn(q, k, v, attn_bias=None):
    scale = 1.0 / q.shape[-1] ** 0.5
    q = q * scale
    attn = q @ k.transpose(-2, -1)
    if attn_bias is not None:
        attn = attn + attn_bias
    attn = attn.softmax(-1)
    return attn @ v


B = 16
T = 50
HD = 4
NH = 12
q = torch.randn(B, T, NH, HD, device="cuda")
k = torch.randn(B, T, NH, HD, device="cuda")
v = torch.randn(B, T, NH, HD, device="cuda")

out_1 = memory_efficient_attention(q, k, v)
out_2 = equiv_fn(q, k, v)

print(torch.allclose(out_1, out_2)) # False
print(torch.mean(torch.abs(out_1 - out_2))) # gives 0.355
```

I realise they probably shouldn't agree completely but they are not even close so I don't think it's numerical. I also tried passing `scale=1.0 / math.sqrt(HD)` into memory_efficient_attention and it still doesn't match (or get close).

So what is actually being calculated internally?



Discussion:
OK, I was being a bit stupid, but the equivalent Pytorch code in the doc is wrong. It should be something like:

```
def equiv_fn(q, k, v, attn_bias=None):
    scale = 1.0 / q.shape[-1] ** 0.5
    q = q * scale
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    attn = q @ k.transpose(-2, -1)
    if attn_bias is not None:
        attn = attn + attn_bias
    attn = attn.softmax(-1)
    attn = attn @ v
    return attn.transpose(1, 2)
```

which then matches what memory_efficient_attention calculates.
> the equivalent Pytorch code in the doc is wrong

Good catch, we should probably fix that. I believe the "equivalent" code in the doc only works for 3d inputs (with no head dimension)
I would like to work on issue you mentioned . Could you assign me the same @danthe3rd  . Thankyou
Sure @ajinkyavbhandare - and thanks for looking into this :)

Resolving patch:
diff --git a/xformers/ops/fmha/__init__.py b/xformers/ops/fmha/__init__.py
--- a/xformers/ops/fmha/__init__.py
+++ b/xformers/ops/fmha/__init__.py
@@ -142,14 +142,18 @@ def memory_efficient_attention(
 
     .. code-block:: python
 
-        scale = 1 / query.shape[-1] ** 0.5
+        scale = 1.0 / query.shape[-1] ** 0.5
         query = query * scale
+        query = query.transpose(1, 2) 
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
         attn = query @ key.transpose(-2, -1)
         if attn_bias is not None:
             attn = attn + attn_bias
         attn = attn.softmax(-1)
-        attn = F.dropout(attn, p)
-        return attn @ value
+        attn = F.dropout(attn, p) 
+        attn = attn @ value
+        return attn.transpose(1, 2)
 
     :Examples:
 


Test patch: