Hi,
When using self-attention, I found it’s common usage to use torch.einsum such as follows:
queries  =  torch.normal(0, 1, (b, h, q, d)).to('cuda') 
keys  =  torch.normal(0, 1, (b, h, q, d)).to('cuda')  # Because self-attention k == q
pre_softmax = torch.einsum("bhqd,bhkd->bhqk", queries, keys)

If my understanding is correct in that full self-attention example, we perform b*h*q*q operations.
If one wants to reduce the numbers of keys for each query to a fixed value k << q, then theoretically the number of operations reduces to b*h*q*k which should therefore reduce both time and memory complexity. Because one would have k separate keys for each query , the implementation using torch.einsum would look like the following:
queries  =  torch.normal(0, 1, (b, h, q, d)).to('cuda') 
keys  =  torch.normal(0, 1, (b, h, q, k, d)).to('cuda') 
pre_softmax = torch.einsum("bhqd,bhqkd->bhqk", queries, keys)

In practice, memory consumption does decrease by a factor q/k but not time complexity which becomes much larger as k grows.
On the other, if we simulate a scenario where the k keys are shared across all the q queries such that the code would look like the following :
queries  =  torch.normal(0, 1, (b, h, q, d)).to('cuda') 
keys  =  torch.normal(0, 1, (b, h, k, d)).to('cuda') 
pre_softmax = torch.einsum("bhqd,bhkd->bhqk", queries, keys)

We then obtain both linear improvement in time and memory as expected.
If we have a 1000 queries I plotted the following performance depending on k:

image819×545 32.2 KB

I was wondering what explained such difference in the implementation of einsum and I assume torch.bmm when the number of operations is the same?  Is there a known way to get around this issue?
Thanks in advance !
Here is gist  from a colab to replicate plot and with more example using also bmm directly.