Repository: facebookresearch/xformers

Issue:
Blocksparse crashes in the encoder benchmark
# 🐛 Bug
Running `python3 xformers/benchmarks/benchmark_encoder.py --activations relu  --plot -emb 256 -bs 32 -heads 16 -mlp MLP -a blocksparse` with latest Triton crashes on an illegal memory access.

The dedicated blocksparse matmul benchmark runs, same for the unit test

## Expected behavior
Well, not crashing

## Environment
```
PyTorch version: 1.9.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8 (64-bit runtime)
Python platform: Linux-5.4.0-52-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: 
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 450.80.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] torch==1.9.1+cu111
[pip3] torch-tb-profiler==0.2.1
[pip3] torchaudio==0.9.1
[pip3] torchvision==0.10.1+cu111
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] pytorch-sphinx-theme      0.0.24                   pypi_0    pypi
[conda] torch                     1.9.1+cu111              pypi_0    pypi
[conda] torch-tb-profiler         0.2.1                    pypi_0    pypi
[conda] torchaudio                0.9.1                    pypi_0    pypi
[conda] torchvision               0.10.1+cu111             pypi_0    pypi
```
## Additional context
Looks like something which can happen depending on the layout, not super clear why


Discussion:
not a stride issue, by default our q and k do have a strange pattern (due to the head splitting), but passing in a dummy tensor with standard decreasing strides leads to the same mem access crash

Resolving patch:
diff --git a/xformers/benchmarks/benchmark_encoder.py b/xformers/benchmarks/benchmark_encoder.py
--- a/xformers/benchmarks/benchmark_encoder.py
+++ b/xformers/benchmarks/benchmark_encoder.py
@@ -220,7 +220,7 @@ def instantiate_xformer(
         "layout": torch.eye(
             sequence_length // block_size,
             sequence_length // block_size,
-            dtype=torch.int,
+            dtype=torch.long,
         )
         .unsqueeze(0)
         .expand(heads, -1, -1),
@@ -325,12 +325,12 @@ def plot(args, results: List[Dict[str, Any]]):
         "-emb", "--embedding_dim", nargs="+", default=[64, 128, 256], type=int
     )
     parser.add_argument(
-        "-sl", "--sequence_length", nargs="+", default=[512, 768, 1024], type=int
+        "-sl", "--sequence_length", nargs="+", default=[512, 1024], type=int
     )
     parser.add_argument("-bs", "--batch_size", nargs="+", default=[8, 16, 32], type=int)
     parser.add_argument("-heads", "--heads", nargs="+", default=[8, 16], type=int)
 
-    parser.add_argument("-fp16", "--pytorch_amp", nargs="+", default=[False], type=bool)
+    parser.add_argument("-fp16", "--pytorch_amp", nargs="+", default=[True], type=bool)
     parser.add_argument("-causal", "--causal", nargs="+", default=[False], type=bool)
     parser.add_argument("-plot", "--plot", action="store_true", default=False)
     parser.add_argument(
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
@@ -5,6 +5,7 @@
 
 
 import logging
+import math
 from dataclasses import dataclass
 from typing import Optional
 
@@ -22,6 +23,14 @@
         from triton.ops.blocksparse import softmax as blocksparse_softmax
 
         from xformers.triton.softmax import MaskType
+        from xformers.triton.utils import gpu_capabilities_older_than_70
+
+        # Blocksparse requires Tensor cores
+        if gpu_capabilities_older_than_70():
+            logging.warning(
+                "Blocksparse is not available: the current GPU does not expose Tensor cores"
+            )
+            _use_triton = False
 
     except ImportError as e:
         logging.warning(
@@ -47,6 +56,9 @@ class BlockSparseAttention(Attention):
         .. warning: the layout is assumed to have the dimensions [heads, seq, seq].
             If some dimensions are missing, we assume that the same layout is to be used across heads.
 
+        .. warning: for now, the sequence (context) length has to be a power of two. This constraint could
+            be relaxed in the future.
+
         .. 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.
@@ -141,12 +153,30 @@ def forward(
                 att_mask is None or att_mask.dim() == 2
             ), "The attention mask is constant across heads, expected dimensions are [seq x seq]"
 
-            # Self-attend: (B, nh, S, hs) x (B, nh, hs, S) -> (B, nh, S, S)
-            # When the computations are block sparse, the matrix types change along the way:
-            # - (sparse) attention matrix = (dense) Kt * (dense) Q
             assert (
                 q.shape[-2] == k.shape[-2]
             ), "Blocksparse requires the same dimensions for K and Q for now"
+
+            assert (
+                q.shape[-2] == self.layout.shape[-2] * self.block_size
+            ), "Actual sequence size and layout are inconsistent"
+            assert (
+                k.shape[-2] == self.layout.shape[-2] * self.block_size
+            ), "Actual sequence size and layout are inconsistent"
+
+            assert math.log(
+                q.shape[-2], 2
+            ).is_integer(), (
+                "For now blocksparse only works on power-of-two sequence lengths"
+            )
+
+            # Blocksparse only works on fp16
+            q_dtype = q.dtype
+            q, k, v = q.half(), k.half(), v.half()
+
+            # Self-attend: (B, nh, S, hs) x (B, nh, hs, S) -> (B, nh, S, S)
+            # When the computations are block sparse, the matrix types change along the way:
+            # - (sparse) attention matrix = (dense) Kt * (dense) Q
             sparse_att_mat = self.sparse_dot_sdd(q, k)
 
             # - softmax on the sparse attention matrix
@@ -161,5 +191,4 @@ def forward(
 
             # - then (dense) attention is (sparse) attention matrix * dense (value)
             a = self.sparse_dot_dsd(sparse_att_mat, v)
-
-            return a
+            return a.to(q_dtype)


Test patch: