Repository: facebookresearch/xformers

Issue:
blocksparse gives RuntimeError: CUDA: Error- illegal address when increase the block size
# 🐛 Bug

I'm not sure whether this is a bug or simply a restriction of triton. But when I follow you doc here https://github.com/facebookresearch/xformers/blob/main/HOWTO.md#blocksparseattention

The sample works fine but it does not work when I increase the block size. 

## Command

## To Reproduce

Steps to reproduce the behavior:

Simply replacing your hyperparameter like this 
BATCH = 1
HEADS = 16
SEQ = 8192
EMB = 64 * HEADS
BLOCK_SIZE = 512
DROPOUT = 0.1
should reproduce the error "RuntimeError: CUDA: Error- illegal address"


- PyTorch Version (e.g., 1.0): 1.10.2
- OS (e.g., Linux): Ubuntu 18.04
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.8
- CUDA/cuDNN version: 11.6
- GPU models and configuration: A100
- Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Ah 512 is not an option, it's 16/32/64 top of mind, I need to guard that! Thanks a lot for the report @xwhan , will fix asap
> Ah 512 is not an option, it's 16/32/64 top of mind, I need to guard that! Thanks a lot for the report @xwhan , will fix asap

I was wrong with 64, needs to be a multiple of 16 (because of tensor cores) but no obvious upper bound, except that you don't really win above a given size (you can repro a 512 block with multiple 64 blocks of course). I'll add an assert to make sure that users stay within reasonable bounds

Resolving patch:
diff --git a/xformers/components/attention/blocksparse.py b/xformers/components/attention/blocksparse.py
--- a/xformers/components/attention/blocksparse.py
+++ b/xformers/components/attention/blocksparse.py
@@ -52,6 +52,9 @@ class BlockSparseAttention(Attention):
         .. warning: for now, the sequence (context) length has to be a power of two. This constraint could
             be relaxed in the future.
 
+        .. warning: the block size has to be picked from [16, 32, 64]. Some speed is gained from bigger blocks.
+            It is of course possible to reproduce coarser patterns given these primitives, as the user sees fit.
+
         .. note: it is possible to pass a specific per batch mask in the forward call,
             but this will not lead to any speed up.
             Any constant sparsity pattern is better passed through the layout parameter.
@@ -76,7 +79,11 @@ def __init__(
                 layout = layout.unsqueeze(0).expand(num_heads, -1, -1)
                 logging.warning(f"New layout dimensions: {layout.shape}")
 
-            assert block_size >= 16, "Minimum block size is 16, for now at least"
+            assert block_size in (
+                16,
+                32,
+                64,
+            ), "Only block sizes in [16, 32, 64] are supported"
 
             super().__init__()
             self.attn_drop = torch.nn.Dropout(dropout, inplace=False)


Test patch:
diff --git a/tests/test_triton_blocksparse.py b/tests/test_triton_blocksparse.py
--- a/tests/test_triton_blocksparse.py
+++ b/tests/test_triton_blocksparse.py
@@ -139,7 +139,7 @@ def test_softmax(BLOCK, WIDTH, DTYPE):
 
 
 @pytest.mark.skipif(not _triton_available, reason="Triton requires a recent CUDA gpu")
-@pytest.mark.parametrize("block", [32])  # 16, 32,
+@pytest.mark.parametrize("block", [32, 43])  # 16, 32,
 def test_attention_fwd_bwd(
     block,
     input_scale=1.0,
@@ -177,45 +177,50 @@ def loss_fn(x):
     query.retain_grad()
     key.retain_grad()
     value.retain_grad()
-    block_sparse_attention = BlockSparseAttention(layout, block)
-    attn_out = block_sparse_attention(
-        att_mask=attn_mask, q=query, k=key, v=value, scale=scale
-    )
-
-    # ad hoc loss
-    loss = loss_fn(attn_out)
-    loss.backward()
-    grads = [query.grad, key.grad, value.grad]
-
-    # Torch version:
-    torch_q, torch_k, torch_v = [x.clone() for x in qkvs]
-    torch_q = torch_q / math.sqrt(head_dim)
-    attn_mask = 1e6 * (-1 + (attn_mask.reshape((1, 1, n_ctx, n_ctx)).cuda()))
-    torch_q.retain_grad()
-    torch_k.retain_grad()
-    torch_v.retain_grad()
-    scores = scale * torch.einsum("bhsd,bhtd->bhst", torch_q, torch_k)
-    scores = scores + attn_mask
-    probs = torch.softmax(scores, dim=-1)
-    torch_attn_out = torch.einsum("bhst,bhtd->bhsd", probs, torch_v)
-
-    # ad hoc loss
-    torch_loss = loss_fn(torch_attn_out)
-    torch_loss.backward()
-    torch_grads = [torch_q.grad, torch_k.grad, torch_v.grad]
-
-    # comparison
-    assert_almost_equal(
-        loss, torch_loss, err_msg=f"Triton loss {loss} and torch loss {torch_loss}"
-    )
+    if block not in [16, 32, 64]:
+        # Check that unsupported dimensions are caught
+        with pytest.raises(AssertionError):
+            _ = BlockSparseAttention(layout, block)
+    else:
+        block_sparse_attention = BlockSparseAttention(layout, block)
+        attn_out = block_sparse_attention(
+            att_mask=attn_mask, q=query, k=key, v=value, scale=scale
+        )
 
-    for g1, g2 in zip(grads, torch_grads):
+        # ad hoc loss
+        loss = loss_fn(attn_out)
+        loss.backward()
+        grads = [query.grad, key.grad, value.grad]
+
+        # Torch version:
+        torch_q, torch_k, torch_v = [x.clone() for x in qkvs]
+        torch_q = torch_q / math.sqrt(head_dim)
+        attn_mask = 1e6 * (-1 + (attn_mask.reshape((1, 1, n_ctx, n_ctx)).cuda()))
+        torch_q.retain_grad()
+        torch_k.retain_grad()
+        torch_v.retain_grad()
+        scores = scale * torch.einsum("bhsd,bhtd->bhst", torch_q, torch_k)
+        scores = scores + attn_mask
+        probs = torch.softmax(scores, dim=-1)
+        torch_attn_out = torch.einsum("bhst,bhtd->bhsd", probs, torch_v)
+
+        # ad hoc loss
+        torch_loss = loss_fn(torch_attn_out)
+        torch_loss.backward()
+        torch_grads = [torch_q.grad, torch_k.grad, torch_v.grad]
+
+        # comparison
         assert_almost_equal(
-            torch.norm(g1),
-            torch.norm(g2),
-            err_msg=f"Triton grad {torch.norm(g1).item()} and torch grad {torch.norm(g2).item()}",
+            loss, torch_loss, err_msg=f"Triton loss {loss} and torch loss {torch_loss}"
         )
 
+        for g1, g2 in zip(grads, torch_grads):
+            assert_almost_equal(
+                torch.norm(g1),
+                torch.norm(g2),
+                err_msg=f"Triton grad {torch.norm(g1).item()} and torch grad {torch.norm(g2).item()}",
+            )
+
 
 @pytest.mark.skipif(not _triton_available, reason="Triton requires a recent CUDA gpu")
 def test_blocksparse_attention_parity():