We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `xformers/components/attention/blocksparse.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
2 #
3 # This source code is licensed under the BSD license found in the
4 # LICENSE file in the root directory of this source tree.
5 
6 
7 import logging
8 import math
9 from dataclasses import dataclass
10 from typing import Optional
11 
12 import torch
13 
14 from xformers import _is_triton_available
15 from xformers.components.attention import Attention, AttentionConfig, register_attention
16 from xformers.components.attention.utils import bool_mask_to_additive
17 
18 _mask_type_warning = True
19 
20 if _is_triton_available:
21     from triton.ops.blocksparse import matmul as blocksparse_matmul  # type: ignore
22     from triton.ops.blocksparse import softmax as blocksparse_softmax  # type: ignore
23 
24     from xformers.triton.softmax import MaskType
25     from xformers.triton.utils import gpu_capabilities_older_than_70
26 
27     # Blocksparse requires Tensor cores
28     if gpu_capabilities_older_than_70():
29         logging.warning(
30             "Blocksparse is not available: the current GPU does not expose Tensor cores"
31         )
32         _is_triton_available = False
33 
34 
35 if _is_triton_available:
36 
37     @dataclass
38     class BlockSparseAttentionConfig(AttentionConfig):
39         layout: torch.Tensor  # The dimensions of the random features
40         block_size: int
41         dropout: float
42         num_heads: int
43 
44     @register_attention("blocksparse", BlockSparseAttentionConfig)
45     class BlockSparseAttention(Attention):
46         r"""
47         Thin wrap over the Triton blocksparse computations. The sparsity pattern is determined through the layout.
48 
49         .. warning: the layout is assumed to have the dimensions [heads, seq, seq].
50             If some dimensions are missing, we assume that the same layout is to be used across heads.
51 
52         .. warning: for now, the sequence (context) length has to be a power of two. This constraint could
53             be relaxed in the future.
54 
55         .. note: it is possible to pass a specific per batch mask in the forward call,
56             but this will not lead to any speed up.
57             Any constant sparsity pattern is better passed through the layout parameter.
58         """
59 
60         def __init__(
61             self,
62             layout: torch.Tensor,
63             block_size: int = 16,
64             dropout: float = 0.0,
65             num_heads: int = 1,  # optional, used to adapt the layout if in need
66             *args,
67             **kwargs,
68         ):
69             if layout.dim() == 2:
70                 logging.warning(
71                     "The layout passed is lacking a head dimension and a batch dimension"
72                 )
73                 logging.warning(
74                     "Now assuming that the same layout is to be used across all heads"
75                 )
76                 layout = layout.unsqueeze(0).expand(num_heads, -1, -1)
77                 logging.warning(f"New layout dimensions: {layout.shape}")
78 
79             assert block_size >= 16, "Minimum block size is 16, for now at least"
80 
81             super().__init__()
82             self.attn_drop = torch.nn.Dropout(dropout, inplace=False)
83 
84             # Pure blocksparse data
85             self.layout = layout
86             self.block_size = block_size
87 
88             # blocksparse operators
89             self.sparse_dot_sdd = blocksparse_matmul(
90                 self.layout,
91                 self.block_size,
92                 "sdd",
93                 trans_a=False,
94                 trans_b=True,
95             )
96             self.sparse_dot_dsd = blocksparse_matmul(
97                 self.layout,
98                 self.block_size,
99                 "dsd",
100                 trans_a=False,
101                 trans_b=False,
102             )
103             self.sparse_softmax = blocksparse_softmax(self.layout, self.block_size)
104 
105             # make sure that the head dimension is not folded down with the batch
106             self.requires_head_dimension = True
107 
108             # key padding mask and attention mask must be passed in separately
109             self.requires_separate_masks = True
110 
111             self.requires_same_k_q_dimensions = True
112 
113         def update_mask_type(self, mask: torch.Tensor):
114             global _mask_type_warning
115             if _mask_type_warning:
116                 logging.warning(
117                     "Mask has to be additive. Fixing that but this slows things down"
118                 )
119             mask = bool_mask_to_additive(mask)
120 
121         def forward(
122             self,
123             q: torch.Tensor,
124             k: torch.Tensor,
125             v: torch.Tensor,
126             att_mask: Optional[torch.Tensor] = None,
127             key_padding_mask: Optional[torch.Tensor] = None,
128             scale: float = 1.0,
129             *args,
130             **kwargs,
131         ) -> torch.Tensor:
132             r"""
133             att_mask            A 2D attention mask. The dtype must be the same as q. An additive mask is expected,
134                                 meaning float values using "-inf" to mask values.
135             key_padding_mask    A mask with size (batch size x sequence length). The dtype must be the same as q.
136                                 An additive mask is expected, meaning float values using "-inf" to mask values
137             """
138 
139             # NOTE:
140             # The attention mask will be taken into account when computing the softmax
141             # meaning that non-masked values which are present in the initial blocksparse layout will be computed.
142             # If blocks are to be constantly masked, better perf would thus be reached by signalling them out in the
143             # initial attention setup
144 
145             if att_mask is not None and att_mask.dtype == torch.bool:
146                 self.update_mask_type(att_mask)
147             if key_padding_mask is not None and key_padding_mask.dtype == torch.bool:
148                 self.update_mask_type(key_padding_mask)
149 
150             assert (
151                 att_mask is None or att_mask.dim() == 2
152             ), "The attention mask is constant across heads, expected dimensions are [seq x seq]"
153 
154             assert (
155                 q.shape[-2] == k.shape[-2]
156             ), "Blocksparse requires the same dimensions for K and Q for now"
157 
158             assert (
159                 q.shape[-2] == self.layout.shape[-2] * self.block_size
160             ), "Actual sequence size and layout are inconsistent"
161             assert (
162                 k.shape[-2] == self.layout.shape[-2] * self.block_size
163             ), "Actual sequence size and layout are inconsistent"
164 
165             assert math.log(
166                 q.shape[-2], 2
167             ).is_integer(), (
168                 "For now blocksparse only works on power-of-two sequence lengths"
169             )
170 
171             # Blocksparse only works on fp16
172             q_dtype = q.dtype
173             q, k, v = q.half(), k.half(), v.half()
174 
175             if att_mask is not None:
176                 att_mask = att_mask.half()
177 
178             if key_padding_mask is not None:
179                 key_padding_mask = key_padding_mask.half()
180 
181             # Self-attend: (B, nh, S, hs) x (B, nh, hs, S) -> (B, nh, S, S)
182             # When the computations are block sparse, the matrix types change along the way:
183             # - (sparse) attention matrix = (dense) Kt * (dense) Q
184             q = q / math.sqrt(q.size(-1))
185             sparse_att_mat = self.sparse_dot_sdd(q, k)
186 
187             # - softmax on the sparse attention matrix
188             sparse_att_mat = self.sparse_softmax(
189                 sparse_att_mat,
190                 scale=scale,
191                 key_padding_mask=key_padding_mask,
192                 attn_mask=att_mask,
193                 key_padding_mask_mode=MaskType.ADD,
194                 attn_mask_mode=MaskType.ADD,
195             )
196 
197             sparse_att_mat = self.attn_drop(sparse_att_mat)
198 
199             # - then (dense) attention is (sparse) attention matrix * dense (value)
200             a = self.sparse_dot_dsd(sparse_att_mat, v)
201             return a.to(q_dtype)
202
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         .. 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 @@
                 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)

{"golden_diff": "diff --git a/xformers/components/attention/blocksparse.py b/xformers/components/attention/blocksparse.py\n--- a/xformers/components/attention/blocksparse.py\n+++ b/xformers/components/attention/blocksparse.py\n@@ -52,6 +52,9 @@\n         .. warning: for now, the sequence (context) length has to be a power of two. This constraint could\n             be relaxed in the future.\n \n+        .. warning: the block size has to be picked from [16, 32, 64]. Some speed is gained from bigger blocks.\n+            It is of course possible to reproduce coarser patterns given these primitives, as the user sees fit.\n+\n         .. note: it is possible to pass a specific per batch mask in the forward call,\n             but this will not lead to any speed up.\n             Any constant sparsity pattern is better passed through the layout parameter.\n@@ -76,7 +79,11 @@\n                 layout = layout.unsqueeze(0).expand(num_heads, -1, -1)\n                 logging.warning(f\"New layout dimensions: {layout.shape}\")\n \n-            assert block_size >= 16, \"Minimum block size is 16, for now at least\"\n+            assert block_size in (\n+                16,\n+                32,\n+                64,\n+            ), \"Only block sizes in [16, 32, 64] are supported\"\n \n             super().__init__()\n             self.attn_drop = torch.nn.Dropout(dropout, inplace=False)\n", "issue": "blocksparse gives RuntimeError: CUDA: Error- illegal address when increase the block size\n# \ud83d\udc1b Bug\r\n\r\nI'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\r\n\r\nThe sample works fine but it does not work when I increase the block size. \r\n\r\n## Command\r\n\r\n## To Reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\nSimply replacing your hyperparameter like this \r\nBATCH = 1\r\nHEADS = 16\r\nSEQ = 8192\r\nEMB = 64 * HEADS\r\nBLOCK_SIZE = 512\r\nDROPOUT = 0.1\r\nshould reproduce the error \"RuntimeError: CUDA: Error- illegal address\"\r\n\r\n\r\n- PyTorch Version (e.g., 1.0): 1.10.2\r\n- OS (e.g., Linux): Ubuntu 18.04\r\n- How you installed PyTorch (`conda`, `pip`, source): pip\r\n- Build command you used (if compiling from source):\r\n- Python version: 3.8\r\n- CUDA/cuDNN version: 11.6\r\n- GPU models and configuration: A100\r\n- Any other relevant information:\r\n\r\n## Additional context\r\n\r\n<!-- Add any other context about the problem here. -->\r\n\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport logging\nimport math\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\n\nfrom xformers import _is_triton_available\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\nfrom xformers.components.attention.utils import bool_mask_to_additive\n\n_mask_type_warning = True\n\nif _is_triton_available:\n    from triton.ops.blocksparse import matmul as blocksparse_matmul  # type: ignore\n    from triton.ops.blocksparse import softmax as blocksparse_softmax  # type: ignore\n\n    from xformers.triton.softmax import MaskType\n    from xformers.triton.utils import gpu_capabilities_older_than_70\n\n    # Blocksparse requires Tensor cores\n    if gpu_capabilities_older_than_70():\n        logging.warning(\n            \"Blocksparse is not available: the current GPU does not expose Tensor cores\"\n        )\n        _is_triton_available = False\n\n\nif _is_triton_available:\n\n    @dataclass\n    class BlockSparseAttentionConfig(AttentionConfig):\n        layout: torch.Tensor  # The dimensions of the random features\n        block_size: int\n        dropout: float\n        num_heads: int\n\n    @register_attention(\"blocksparse\", BlockSparseAttentionConfig)\n    class BlockSparseAttention(Attention):\n        r\"\"\"\n        Thin wrap over the Triton blocksparse computations. The sparsity pattern is determined through the layout.\n\n        .. warning: the layout is assumed to have the dimensions [heads, seq, seq].\n            If some dimensions are missing, we assume that the same layout is to be used across heads.\n\n        .. warning: for now, the sequence (context) length has to be a power of two. This constraint could\n            be relaxed in the future.\n\n        .. note: it is possible to pass a specific per batch mask in the forward call,\n            but this will not lead to any speed up.\n            Any constant sparsity pattern is better passed through the layout parameter.\n        \"\"\"\n\n        def __init__(\n            self,\n            layout: torch.Tensor,\n            block_size: int = 16,\n            dropout: float = 0.0,\n            num_heads: int = 1,  # optional, used to adapt the layout if in need\n            *args,\n            **kwargs,\n        ):\n            if layout.dim() == 2:\n                logging.warning(\n                    \"The layout passed is lacking a head dimension and a batch dimension\"\n                )\n                logging.warning(\n                    \"Now assuming that the same layout is to be used across all heads\"\n                )\n                layout = layout.unsqueeze(0).expand(num_heads, -1, -1)\n                logging.warning(f\"New layout dimensions: {layout.shape}\")\n\n            assert block_size >= 16, \"Minimum block size is 16, for now at least\"\n\n            super().__init__()\n            self.attn_drop = torch.nn.Dropout(dropout, inplace=False)\n\n            # Pure blocksparse data\n            self.layout = layout\n            self.block_size = block_size\n\n            # blocksparse operators\n            self.sparse_dot_sdd = blocksparse_matmul(\n                self.layout,\n                self.block_size,\n                \"sdd\",\n                trans_a=False,\n                trans_b=True,\n            )\n            self.sparse_dot_dsd = blocksparse_matmul(\n                self.layout,\n                self.block_size,\n                \"dsd\",\n                trans_a=False,\n                trans_b=False,\n            )\n            self.sparse_softmax = blocksparse_softmax(self.layout, self.block_size)\n\n            # make sure that the head dimension is not folded down with the batch\n            self.requires_head_dimension = True\n\n            # key padding mask and attention mask must be passed in separately\n            self.requires_separate_masks = True\n\n            self.requires_same_k_q_dimensions = True\n\n        def update_mask_type(self, mask: torch.Tensor):\n            global _mask_type_warning\n            if _mask_type_warning:\n                logging.warning(\n                    \"Mask has to be additive. Fixing that but this slows things down\"\n                )\n            mask = bool_mask_to_additive(mask)\n\n        def forward(\n            self,\n            q: torch.Tensor,\n            k: torch.Tensor,\n            v: torch.Tensor,\n            att_mask: Optional[torch.Tensor] = None,\n            key_padding_mask: Optional[torch.Tensor] = None,\n            scale: float = 1.0,\n            *args,\n            **kwargs,\n        ) -> torch.Tensor:\n            r\"\"\"\n            att_mask            A 2D attention mask. The dtype must be the same as q. An additive mask is expected,\n                                meaning float values using \"-inf\" to mask values.\n            key_padding_mask    A mask with size (batch size x sequence length). The dtype must be the same as q.\n                                An additive mask is expected, meaning float values using \"-inf\" to mask values\n            \"\"\"\n\n            # NOTE:\n            # The attention mask will be taken into account when computing the softmax\n            # meaning that non-masked values which are present in the initial blocksparse layout will be computed.\n            # If blocks are to be constantly masked, better perf would thus be reached by signalling them out in the\n            # initial attention setup\n\n            if att_mask is not None and att_mask.dtype == torch.bool:\n                self.update_mask_type(att_mask)\n            if key_padding_mask is not None and key_padding_mask.dtype == torch.bool:\n                self.update_mask_type(key_padding_mask)\n\n            assert (\n                att_mask is None or att_mask.dim() == 2\n            ), \"The attention mask is constant across heads, expected dimensions are [seq x seq]\"\n\n            assert (\n                q.shape[-2] == k.shape[-2]\n            ), \"Blocksparse requires the same dimensions for K and Q for now\"\n\n            assert (\n                q.shape[-2] == self.layout.shape[-2] * self.block_size\n            ), \"Actual sequence size and layout are inconsistent\"\n            assert (\n                k.shape[-2] == self.layout.shape[-2] * self.block_size\n            ), \"Actual sequence size and layout are inconsistent\"\n\n            assert math.log(\n                q.shape[-2], 2\n            ).is_integer(), (\n                \"For now blocksparse only works on power-of-two sequence lengths\"\n            )\n\n            # Blocksparse only works on fp16\n            q_dtype = q.dtype\n            q, k, v = q.half(), k.half(), v.half()\n\n            if att_mask is not None:\n                att_mask = att_mask.half()\n\n            if key_padding_mask is not None:\n                key_padding_mask = key_padding_mask.half()\n\n            # Self-attend: (B, nh, S, hs) x (B, nh, hs, S) -> (B, nh, S, S)\n            # When the computations are block sparse, the matrix types change along the way:\n            # - (sparse) attention matrix = (dense) Kt * (dense) Q\n            q = q / math.sqrt(q.size(-1))\n            sparse_att_mat = self.sparse_dot_sdd(q, k)\n\n            # - softmax on the sparse attention matrix\n            sparse_att_mat = self.sparse_softmax(\n                sparse_att_mat,\n                scale=scale,\n                key_padding_mask=key_padding_mask,\n                attn_mask=att_mask,\n                key_padding_mask_mode=MaskType.ADD,\n                attn_mask_mode=MaskType.ADD,\n            )\n\n            sparse_att_mat = self.attn_drop(sparse_att_mat)\n\n            # - then (dense) attention is (sparse) attention matrix * dense (value)\n            a = self.sparse_dot_dsd(sparse_att_mat, v)\n            return a.to(q_dtype)\n", "path": "xformers/components/attention/blocksparse.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport logging\nimport math\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\n\nfrom xformers import _is_triton_available\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\nfrom xformers.components.attention.utils import bool_mask_to_additive\n\n_mask_type_warning = True\n\nif _is_triton_available:\n    from triton.ops.blocksparse import matmul as blocksparse_matmul  # type: ignore\n    from triton.ops.blocksparse import softmax as blocksparse_softmax  # type: ignore\n\n    from xformers.triton.softmax import MaskType\n    from xformers.triton.utils import gpu_capabilities_older_than_70\n\n    # Blocksparse requires Tensor cores\n    if gpu_capabilities_older_than_70():\n        logging.warning(\n            \"Blocksparse is not available: the current GPU does not expose Tensor cores\"\n        )\n        _is_triton_available = False\n\n\nif _is_triton_available:\n\n    @dataclass\n    class BlockSparseAttentionConfig(AttentionConfig):\n        layout: torch.Tensor  # The dimensions of the random features\n        block_size: int\n        dropout: float\n        num_heads: int\n\n    @register_attention(\"blocksparse\", BlockSparseAttentionConfig)\n    class BlockSparseAttention(Attention):\n        r\"\"\"\n        Thin wrap over the Triton blocksparse computations. The sparsity pattern is determined through the layout.\n\n        .. warning: the layout is assumed to have the dimensions [heads, seq, seq].\n            If some dimensions are missing, we assume that the same layout is to be used across heads.\n\n        .. warning: for now, the sequence (context) length has to be a power of two. This constraint could\n            be relaxed in the future.\n\n        .. warning: the block size has to be picked from [16, 32, 64]. Some speed is gained from bigger blocks.\n            It is of course possible to reproduce coarser patterns given these primitives, as the user sees fit.\n\n        .. note: it is possible to pass a specific per batch mask in the forward call,\n            but this will not lead to any speed up.\n            Any constant sparsity pattern is better passed through the layout parameter.\n        \"\"\"\n\n        def __init__(\n            self,\n            layout: torch.Tensor,\n            block_size: int = 16,\n            dropout: float = 0.0,\n            num_heads: int = 1,  # optional, used to adapt the layout if in need\n            *args,\n            **kwargs,\n        ):\n            if layout.dim() == 2:\n                logging.warning(\n                    \"The layout passed is lacking a head dimension and a batch dimension\"\n                )\n                logging.warning(\n                    \"Now assuming that the same layout is to be used across all heads\"\n                )\n                layout = layout.unsqueeze(0).expand(num_heads, -1, -1)\n                logging.warning(f\"New layout dimensions: {layout.shape}\")\n\n            assert block_size in (\n                16,\n                32,\n                64,\n            ), \"Only block sizes in [16, 32, 64] are supported\"\n\n            super().__init__()\n            self.attn_drop = torch.nn.Dropout(dropout, inplace=False)\n\n            # Pure blocksparse data\n            self.layout = layout\n            self.block_size = block_size\n\n            # blocksparse operators\n            self.sparse_dot_sdd = blocksparse_matmul(\n                self.layout,\n                self.block_size,\n                \"sdd\",\n                trans_a=False,\n                trans_b=True,\n            )\n            self.sparse_dot_dsd = blocksparse_matmul(\n                self.layout,\n                self.block_size,\n                \"dsd\",\n                trans_a=False,\n                trans_b=False,\n            )\n            self.sparse_softmax = blocksparse_softmax(self.layout, self.block_size)\n\n            # make sure that the head dimension is not folded down with the batch\n            self.requires_head_dimension = True\n\n            # key padding mask and attention mask must be passed in separately\n            self.requires_separate_masks = True\n\n            self.requires_same_k_q_dimensions = True\n\n        def update_mask_type(self, mask: torch.Tensor):\n            global _mask_type_warning\n            if _mask_type_warning:\n                logging.warning(\n                    \"Mask has to be additive. Fixing that but this slows things down\"\n                )\n            mask = bool_mask_to_additive(mask)\n\n        def forward(\n            self,\n            q: torch.Tensor,\n            k: torch.Tensor,\n            v: torch.Tensor,\n            att_mask: Optional[torch.Tensor] = None,\n            key_padding_mask: Optional[torch.Tensor] = None,\n            scale: float = 1.0,\n            *args,\n            **kwargs,\n        ) -> torch.Tensor:\n            r\"\"\"\n            att_mask            A 2D attention mask. The dtype must be the same as q. An additive mask is expected,\n                                meaning float values using \"-inf\" to mask values.\n            key_padding_mask    A mask with size (batch size x sequence length). The dtype must be the same as q.\n                                An additive mask is expected, meaning float values using \"-inf\" to mask values\n            \"\"\"\n\n            # NOTE:\n            # The attention mask will be taken into account when computing the softmax\n            # meaning that non-masked values which are present in the initial blocksparse layout will be computed.\n            # If blocks are to be constantly masked, better perf would thus be reached by signalling them out in the\n            # initial attention setup\n\n            if att_mask is not None and att_mask.dtype == torch.bool:\n                self.update_mask_type(att_mask)\n            if key_padding_mask is not None and key_padding_mask.dtype == torch.bool:\n                self.update_mask_type(key_padding_mask)\n\n            assert (\n                att_mask is None or att_mask.dim() == 2\n            ), \"The attention mask is constant across heads, expected dimensions are [seq x seq]\"\n\n            assert (\n                q.shape[-2] == k.shape[-2]\n            ), \"Blocksparse requires the same dimensions for K and Q for now\"\n\n            assert (\n                q.shape[-2] == self.layout.shape[-2] * self.block_size\n            ), \"Actual sequence size and layout are inconsistent\"\n            assert (\n                k.shape[-2] == self.layout.shape[-2] * self.block_size\n            ), \"Actual sequence size and layout are inconsistent\"\n\n            assert math.log(\n                q.shape[-2], 2\n            ).is_integer(), (\n                \"For now blocksparse only works on power-of-two sequence lengths\"\n            )\n\n            # Blocksparse only works on fp16\n            q_dtype = q.dtype\n            q, k, v = q.half(), k.half(), v.half()\n\n            if att_mask is not None:\n                att_mask = att_mask.half()\n\n            if key_padding_mask is not None:\n                key_padding_mask = key_padding_mask.half()\n\n            # Self-attend: (B, nh, S, hs) x (B, nh, hs, S) -> (B, nh, S, S)\n            # When the computations are block sparse, the matrix types change along the way:\n            # - (sparse) attention matrix = (dense) Kt * (dense) Q\n            q = q / math.sqrt(q.size(-1))\n            sparse_att_mat = self.sparse_dot_sdd(q, k)\n\n            # - softmax on the sparse attention matrix\n            sparse_att_mat = self.sparse_softmax(\n                sparse_att_mat,\n                scale=scale,\n                key_padding_mask=key_padding_mask,\n                attn_mask=att_mask,\n                key_padding_mask_mode=MaskType.ADD,\n                attn_mask_mode=MaskType.ADD,\n            )\n\n            sparse_att_mat = self.attn_drop(sparse_att_mat)\n\n            # - then (dense) attention is (sparse) attention matrix * dense (value)\n            a = self.sparse_dot_dsd(sparse_att_mat, v)\n            return a.to(q_dtype)\n", "path": "xformers/components/attention/blocksparse.py"}]}