Repository: facebookresearch/xformers

Issue:
[feat] Make the batch dimension broadcastable
# 🚀 Feature
Support different key/query/value batch sizes, if 1 or the same value

## Motivation
(popped up when working on [salina](https://github.com/facebookresearch/salina)) PyTorch MHA handles key/query/value with 1 to B batch sizes, will broadcast if in need, We don´t do that and hard crash instead, so the users have to align dimensions from the outside

## Pitch

<!-- A clear and concise description of what you want to happen. -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
I can try to take this. I'll submit a PR later today!
Maybe to make sure I understand the ask, I see a few ways to interpret this:

1. xFormers doesn't support unbatched (2D) QKVs currently
2. xFormers doesn't support batched, but varying batch sizes between QKV (i.e., `Q.size(0) != K.size(0) != V.size(0)`)
3. xFormers doesn't support >=4D QKVs (like [torchtext](https://github.com/pytorch/text/blob/main/torchtext/nn/modules/multiheadattention.py#L12-L16))

For the scope of a relevant PR, which aspects would you like to address immediately?
sorry for the delay, and thanks for raising that @erip ! I did not explain this issue well enough, sorry about that, it was a brain dump at the time as I was having a look at salina. 2. is the issue I had in mind (one of the batch dims is 1 and we can broadcast it towards the others), but 1. would make sense also, as it's a one liner really. 

The 4D case is basically handled within the MHA if I understood torchtext case correctly (the MHA is the one generating the head dimension), I think that it's simpler not to expose that as it would have consequences on all the attention implementations. The attentions sit below the MHA right now, so they're guaranteed to always get the dimensions in a consistent way, which simplifies them a great deal I believe (and some of them are complicated enough already :D)

Resolving patch:
diff --git a/xformers/components/multi_head_dispatch.py b/xformers/components/multi_head_dispatch.py
--- a/xformers/components/multi_head_dispatch.py
+++ b/xformers/components/multi_head_dispatch.py
@@ -143,6 +143,11 @@ def forward(
         self._check(value, "value")
         self._check(key, "key")
 
+        max_batch = max((query.shape[0], key.shape[0], value.shape[0]))
+        query, key, value = map(
+            lambda x: x.expand(max_batch, -1, -1), [query, key, value]
+        )
+
         B, S_Q, _ = query.size()  # Batch x Sequence x Embedding (latent)
         _, S_K, _ = key.size()  # K, Q's sequence length could differ
 


Test patch:
diff --git a/tests/test_attentions.py b/tests/test_attentions.py
--- a/tests/test_attentions.py
+++ b/tests/test_attentions.py
@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 
+from typing import Tuple
+
 import pytest
 import torch
 
@@ -253,6 +255,41 @@ def test_different_kq_dimensions(
     assert res.shape == torch.Size([BATCH, seq_q, MODEL])
 
 
+@pytest.mark.parametrize("heads", [1, 4])
+@pytest.mark.parametrize("attention_name", ATTENTION_REGISTRY.keys())
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize(
+    "batch_sizes",
+    [
+        (1, BATCH, BATCH),
+        (BATCH, 1, BATCH),
+        (BATCH, BATCH, 1),
+        (1, 1, BATCH),
+        (BATCH, 1, 1),
+        (1, BATCH, 1),
+    ],
+)
+def test_broadcast_batch_dimension(
+    attention_name: str,
+    heads: int,
+    device: torch.device,
+    batch_sizes: Tuple[int, int, int],
+):
+    Q_BATCH, K_BATCH, V_BATCH = batch_sizes
+    multi_head = _get_multihead(attention_name, 0.0, 0.0, False, heads, device)
+
+    if multi_head.attention.requires_same_k_q_dimensions:
+        # pyre-fixme[29]: The library function `pytest.skip` is not supported by Pyre.
+        pytest.skip(f"{attention_name} does not support different k, q dimensions yet.")
+
+    q = torch.rand((Q_BATCH, SEQ, MODEL), device=device)
+    k = torch.rand((K_BATCH, SEQ, MODEL), device=device)
+    v = torch.rand((V_BATCH, SEQ, MODEL), device=device)
+
+    res = multi_head(query=q, key=k, value=v)
+    assert res.shape == torch.Size([BATCH, SEQ, MODEL])
+
+
 @pytest.mark.parametrize("heads", [1, 4])
 @pytest.mark.parametrize("attention_name", ["scaled_dot_product", "favor"])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA gpu")