We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Follow up on #141 with a proper unit test
# 🐛 Bug
Make sure that #141 cannot happen anymore, add a small unit test to guard parity

Follow up on #141 with a proper unit test
# 🐛 Bug
Make sure that #141 cannot happen anymore, add a small unit test to guard parity


--- 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/triton/softmax.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 from enum import Enum
9 from typing import Optional
10 
11 import torch
12 import triton
13 from torch.cuda.amp import custom_bwd, custom_fwd
14 
15 from xformers.triton.k_softmax import _softmax, _softmax_backward
16 
17 # CREDITS: This is adapted from the vanilla Triton example. See https://openai.com/blog/triton/
18 # and https://triton-lang.org/getting-started/tutorials/02-fused-softmax.html
19 
20 
21 _triton_registered_overflow = False
22 _triton_registered_warnings = False
23 _triton_softmax_fp16_enabled = False  # NOTE: PyTorch keeps softmax as fp32
24 
25 
26 class MaskType(str, Enum):
27     ADD = "add"
28     MUL = "mul"
29 
30 
31 # Helper to handle the SPMD launch grid and error cases
32 class _softmax_triton(torch.autograd.Function):
33     @staticmethod
34     @custom_fwd(cast_inputs=torch.float16 if _triton_softmax_fp16_enabled else None)
35     def forward(ctx, x, mask, log_outputs, causal):
36         """
37         Fused softmax implementation, using the Triton programming model.
38         This only supports a reduction over the last dimension for now
39         """
40 
41         # Handle 2D/3D tensors
42         x_ = x.unsqueeze(0) if x.ndim == 2 else x
43 
44         if not x_.is_contiguous():
45             x_ = x_.contiguous()
46 
47         y = torch.empty_like(x_)
48         assert (
49             y.stride(2) == 1 and x_.stride(2) == 1
50         ), f"{x.shape} - {x_.shape} - {x_.stride()}"
51 
52         # SPMD launch grid
53         grid_2d = (
54             x_.shape[0],
55             x_.shape[1],
56         )
57 
58         # enqueue GPU kernel
59         use_mask = True
60         if mask is None:
61             #  placeholder, will not be used
62             mask = x_
63             use_mask = False
64         else:
65             # Make sure that the mask is binary
66             assert mask.dtype == x.dtype, "An additive mask is requested"
67 
68         _softmax[grid_2d](
69             y,
70             x_,
71             mask,
72             y.stride(0),
73             y.stride(1),
74             x_.stride(0),
75             x_.stride(1),
76             mask.stride(0),
77             x_.shape[2],
78             log=log_outputs,
79             use_mask=use_mask,
80             causal=causal,
81         )
82 
83         ctx.save_for_backward(y)
84         ctx.log_outputs = log_outputs
85         ctx.causal = causal
86         return y.reshape_as(x)
87 
88     @staticmethod
89     @custom_bwd
90     def backward(ctx, grad_out):
91         (out,) = ctx.saved_tensors
92 
93         # Handle 2D/3D tensors
94         grad_out_ = grad_out.unsqueeze(0) if grad_out.ndim == 2 else grad_out
95 
96         # SPMD launch grid
97         grid_2d = (
98             grad_out_.shape[0],
99             grad_out_.shape[1],
100         )
101 
102         depth = triton.next_power_of_2(grad_out_.shape[2])
103         grad_in = torch.empty_like(
104             out
105         )  # torch.zeros is measurably slower, we'll zero out in the kernel
106 
107         # Make sure that the tensor are contiguous
108         grad_in, grad_out, out = map(lambda x: x.contiguous(), [grad_in, grad_out, out])
109 
110         # fmt: off
111         _softmax_backward[grid_2d](
112             grad_in, grad_out_, out,
113             grad_in.stride(0), grad_in.stride(1),
114             grad_out_.stride(0), grad_out_.stride(1),
115             out.stride(0), out.stride(1),
116             out.shape[2],
117             depth=depth,
118             log=ctx.log_outputs,
119             causal=ctx.causal
120         )
121         # fmt: on
122         return grad_in.reshape_as(grad_out), None, None, None
123 
124 
125 def softmax(
126     x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False
127 ) -> torch.Tensor:
128     r"""Applies the Softmax function to an 3-dimensional input Tensor
129     rescaling them so that the elements of the n-dimensional output Tensor
130     lie in the range [0,1] and sum to 1.
131 
132     Softmax is defined as:
133 
134     .. math::
135         \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
136 
137     .. warning: softmax is computed on the last dimension of the input tensor.
138 
139 
140     Args:
141         x: input tensor.
142         mask: optional mask, its application will be fused to the softmax computation if triton is used
143         causal: optional performance optimization, if triton is used and the attention is causal
144 
145     Returns:
146         a Tensor of the same dimension and shape as the input with
147         values in the range [0, 1] and sum to 1
148     """
149     return _softmax_dispatch(x, log=False, mask=mask, causal=causal)
150 
151 
152 def log_softmax(
153     x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False
154 ) -> torch.Tensor:
155     r"""Applies the :math:`\log(\text{Softmax}(x))` function to an 3-dimensional
156     input Tensor. The LogSoftmax formulation can be simplified as:
157 
158     .. math::
159         \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)
160 
161     Args:
162         x: input tensor.
163 
164     Returns:
165         a Tensor of the same dimension and shape as the input with
166         values in the range [-inf, 0)
167     """
168     return _softmax_dispatch(x, log=True, mask=mask, causal=causal)
169 
170 
171 def _softmax_dispatch(
172     x: torch.Tensor, log: bool, mask: Optional[torch.Tensor], causal: bool = False
173 ) -> torch.Tensor:
174     # Triton is used if
175     # - CUDA
176     # - there's enough data to make it faster than pytorch. This could change over time, Triton is improving
177     # - there was no previous failure
178 
179     global _triton_registered_overflow
180     global _triton_registered_warnings
181 
182     try:
183         if torch.cuda.is_available() and x.is_cuda and not _triton_registered_overflow:
184             return _softmax_triton.apply(x, mask, log, causal)
185     except (triton.code_gen.OutOfResources, RuntimeError) as e:
186         # Catch cases where the current GPU does not have enough registers to hold a full tensor line
187         # fallback to PyTorch's implementation, which streams the tensor in and out
188         _triton_registered_overflow = True
189         logging.warning(
190             "Triton softmax kernel register spillover or invalid image caught."
191             "Deactivating this kernel, please file an issue int the xFormers repository"
192         )
193         logging.warning(e)
194 
195     if causal and not _triton_registered_warnings:
196         logging.warning(
197             "Triton softmax could not be used. \
198                 The causal flags is being passed but it does not provide any benefit with PyTorch softmax."
199         )
200         _triton_registered_warnings = True
201 
202     if mask is not None:
203         x += mask
204 
205     if log:
206         return torch.log_softmax(x, dim=-1)
207     else:
208         return torch.softmax(x, dim=-1)
209
```
--- 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/triton/softmax.py b/xformers/triton/softmax.py
--- a/xformers/triton/softmax.py
+++ b/xformers/triton/softmax.py
@@ -200,7 +200,7 @@
         _triton_registered_warnings = True
 
     if mask is not None:
-        x += mask
+        x = x + mask
 
     if log:
         return torch.log_softmax(x, dim=-1)

{"golden_diff": "diff --git a/xformers/triton/softmax.py b/xformers/triton/softmax.py\n--- a/xformers/triton/softmax.py\n+++ b/xformers/triton/softmax.py\n@@ -200,7 +200,7 @@\n         _triton_registered_warnings = True\n \n     if mask is not None:\n-        x += mask\n+        x = x + mask\n \n     if log:\n         return torch.log_softmax(x, dim=-1)\n", "issue": "Follow up on #141 with a proper unit test\n# \ud83d\udc1b Bug\r\nMake sure that #141 cannot happen anymore, add a small unit test to guard parity\r\n\nFollow up on #141 with a proper unit test\n# \ud83d\udc1b Bug\r\nMake sure that #141 cannot happen anymore, add a small unit test to guard parity\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\nfrom enum import Enum\nfrom typing import Optional\n\nimport torch\nimport triton\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\nfrom xformers.triton.k_softmax import _softmax, _softmax_backward\n\n# CREDITS: This is adapted from the vanilla Triton example. See https://openai.com/blog/triton/\n# and https://triton-lang.org/getting-started/tutorials/02-fused-softmax.html\n\n\n_triton_registered_overflow = False\n_triton_registered_warnings = False\n_triton_softmax_fp16_enabled = False  # NOTE: PyTorch keeps softmax as fp32\n\n\nclass MaskType(str, Enum):\n    ADD = \"add\"\n    MUL = \"mul\"\n\n\n# Helper to handle the SPMD launch grid and error cases\nclass _softmax_triton(torch.autograd.Function):\n    @staticmethod\n    @custom_fwd(cast_inputs=torch.float16 if _triton_softmax_fp16_enabled else None)\n    def forward(ctx, x, mask, log_outputs, causal):\n        \"\"\"\n        Fused softmax implementation, using the Triton programming model.\n        This only supports a reduction over the last dimension for now\n        \"\"\"\n\n        # Handle 2D/3D tensors\n        x_ = x.unsqueeze(0) if x.ndim == 2 else x\n\n        if not x_.is_contiguous():\n            x_ = x_.contiguous()\n\n        y = torch.empty_like(x_)\n        assert (\n            y.stride(2) == 1 and x_.stride(2) == 1\n        ), f\"{x.shape} - {x_.shape} - {x_.stride()}\"\n\n        # SPMD launch grid\n        grid_2d = (\n            x_.shape[0],\n            x_.shape[1],\n        )\n\n        # enqueue GPU kernel\n        use_mask = True\n        if mask is None:\n            #  placeholder, will not be used\n            mask = x_\n            use_mask = False\n        else:\n            # Make sure that the mask is binary\n            assert mask.dtype == x.dtype, \"An additive mask is requested\"\n\n        _softmax[grid_2d](\n            y,\n            x_,\n            mask,\n            y.stride(0),\n            y.stride(1),\n            x_.stride(0),\n            x_.stride(1),\n            mask.stride(0),\n            x_.shape[2],\n            log=log_outputs,\n            use_mask=use_mask,\n            causal=causal,\n        )\n\n        ctx.save_for_backward(y)\n        ctx.log_outputs = log_outputs\n        ctx.causal = causal\n        return y.reshape_as(x)\n\n    @staticmethod\n    @custom_bwd\n    def backward(ctx, grad_out):\n        (out,) = ctx.saved_tensors\n\n        # Handle 2D/3D tensors\n        grad_out_ = grad_out.unsqueeze(0) if grad_out.ndim == 2 else grad_out\n\n        # SPMD launch grid\n        grid_2d = (\n            grad_out_.shape[0],\n            grad_out_.shape[1],\n        )\n\n        depth = triton.next_power_of_2(grad_out_.shape[2])\n        grad_in = torch.empty_like(\n            out\n        )  # torch.zeros is measurably slower, we'll zero out in the kernel\n\n        # Make sure that the tensor are contiguous\n        grad_in, grad_out, out = map(lambda x: x.contiguous(), [grad_in, grad_out, out])\n\n        # fmt: off\n        _softmax_backward[grid_2d](\n            grad_in, grad_out_, out,\n            grad_in.stride(0), grad_in.stride(1),\n            grad_out_.stride(0), grad_out_.stride(1),\n            out.stride(0), out.stride(1),\n            out.shape[2],\n            depth=depth,\n            log=ctx.log_outputs,\n            causal=ctx.causal\n        )\n        # fmt: on\n        return grad_in.reshape_as(grad_out), None, None, None\n\n\ndef softmax(\n    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False\n) -> torch.Tensor:\n    r\"\"\"Applies the Softmax function to an 3-dimensional input Tensor\n    rescaling them so that the elements of the n-dimensional output Tensor\n    lie in the range [0,1] and sum to 1.\n\n    Softmax is defined as:\n\n    .. math::\n        \\text{Softmax}(x_{i}) = \\frac{\\exp(x_i)}{\\sum_j \\exp(x_j)}\n\n    .. warning: softmax is computed on the last dimension of the input tensor.\n\n\n    Args:\n        x: input tensor.\n        mask: optional mask, its application will be fused to the softmax computation if triton is used\n        causal: optional performance optimization, if triton is used and the attention is causal\n\n    Returns:\n        a Tensor of the same dimension and shape as the input with\n        values in the range [0, 1] and sum to 1\n    \"\"\"\n    return _softmax_dispatch(x, log=False, mask=mask, causal=causal)\n\n\ndef log_softmax(\n    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False\n) -> torch.Tensor:\n    r\"\"\"Applies the :math:`\\log(\\text{Softmax}(x))` function to an 3-dimensional\n    input Tensor. The LogSoftmax formulation can be simplified as:\n\n    .. math::\n        \\text{LogSoftmax}(x_{i}) = \\log\\left(\\frac{\\exp(x_i) }{ \\sum_j \\exp(x_j)} \\right)\n\n    Args:\n        x: input tensor.\n\n    Returns:\n        a Tensor of the same dimension and shape as the input with\n        values in the range [-inf, 0)\n    \"\"\"\n    return _softmax_dispatch(x, log=True, mask=mask, causal=causal)\n\n\ndef _softmax_dispatch(\n    x: torch.Tensor, log: bool, mask: Optional[torch.Tensor], causal: bool = False\n) -> torch.Tensor:\n    # Triton is used if\n    # - CUDA\n    # - there's enough data to make it faster than pytorch. This could change over time, Triton is improving\n    # - there was no previous failure\n\n    global _triton_registered_overflow\n    global _triton_registered_warnings\n\n    try:\n        if torch.cuda.is_available() and x.is_cuda and not _triton_registered_overflow:\n            return _softmax_triton.apply(x, mask, log, causal)\n    except (triton.code_gen.OutOfResources, RuntimeError) as e:\n        # Catch cases where the current GPU does not have enough registers to hold a full tensor line\n        # fallback to PyTorch's implementation, which streams the tensor in and out\n        _triton_registered_overflow = True\n        logging.warning(\n            \"Triton softmax kernel register spillover or invalid image caught.\"\n            \"Deactivating this kernel, please file an issue int the xFormers repository\"\n        )\n        logging.warning(e)\n\n    if causal and not _triton_registered_warnings:\n        logging.warning(\n            \"Triton softmax could not be used. \\\n                The causal flags is being passed but it does not provide any benefit with PyTorch softmax.\"\n        )\n        _triton_registered_warnings = True\n\n    if mask is not None:\n        x += mask\n\n    if log:\n        return torch.log_softmax(x, dim=-1)\n    else:\n        return torch.softmax(x, dim=-1)\n", "path": "xformers/triton/softmax.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\nfrom enum import Enum\nfrom typing import Optional\n\nimport torch\nimport triton\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\nfrom xformers.triton.k_softmax import _softmax, _softmax_backward\n\n# CREDITS: This is adapted from the vanilla Triton example. See https://openai.com/blog/triton/\n# and https://triton-lang.org/getting-started/tutorials/02-fused-softmax.html\n\n\n_triton_registered_overflow = False\n_triton_registered_warnings = False\n_triton_softmax_fp16_enabled = False  # NOTE: PyTorch keeps softmax as fp32\n\n\nclass MaskType(str, Enum):\n    ADD = \"add\"\n    MUL = \"mul\"\n\n\n# Helper to handle the SPMD launch grid and error cases\nclass _softmax_triton(torch.autograd.Function):\n    @staticmethod\n    @custom_fwd(cast_inputs=torch.float16 if _triton_softmax_fp16_enabled else None)\n    def forward(ctx, x, mask, log_outputs, causal):\n        \"\"\"\n        Fused softmax implementation, using the Triton programming model.\n        This only supports a reduction over the last dimension for now\n        \"\"\"\n\n        # Handle 2D/3D tensors\n        x_ = x.unsqueeze(0) if x.ndim == 2 else x\n\n        if not x_.is_contiguous():\n            x_ = x_.contiguous()\n\n        y = torch.empty_like(x_)\n        assert (\n            y.stride(2) == 1 and x_.stride(2) == 1\n        ), f\"{x.shape} - {x_.shape} - {x_.stride()}\"\n\n        # SPMD launch grid\n        grid_2d = (\n            x_.shape[0],\n            x_.shape[1],\n        )\n\n        # enqueue GPU kernel\n        use_mask = True\n        if mask is None:\n            #  placeholder, will not be used\n            mask = x_\n            use_mask = False\n        else:\n            # Make sure that the mask is binary\n            assert mask.dtype == x.dtype, \"An additive mask is requested\"\n\n        _softmax[grid_2d](\n            y,\n            x_,\n            mask,\n            y.stride(0),\n            y.stride(1),\n            x_.stride(0),\n            x_.stride(1),\n            mask.stride(0),\n            x_.shape[2],\n            log=log_outputs,\n            use_mask=use_mask,\n            causal=causal,\n        )\n\n        ctx.save_for_backward(y)\n        ctx.log_outputs = log_outputs\n        ctx.causal = causal\n        return y.reshape_as(x)\n\n    @staticmethod\n    @custom_bwd\n    def backward(ctx, grad_out):\n        (out,) = ctx.saved_tensors\n\n        # Handle 2D/3D tensors\n        grad_out_ = grad_out.unsqueeze(0) if grad_out.ndim == 2 else grad_out\n\n        # SPMD launch grid\n        grid_2d = (\n            grad_out_.shape[0],\n            grad_out_.shape[1],\n        )\n\n        depth = triton.next_power_of_2(grad_out_.shape[2])\n        grad_in = torch.empty_like(\n            out\n        )  # torch.zeros is measurably slower, we'll zero out in the kernel\n\n        # Make sure that the tensor are contiguous\n        grad_in, grad_out, out = map(lambda x: x.contiguous(), [grad_in, grad_out, out])\n\n        # fmt: off\n        _softmax_backward[grid_2d](\n            grad_in, grad_out_, out,\n            grad_in.stride(0), grad_in.stride(1),\n            grad_out_.stride(0), grad_out_.stride(1),\n            out.stride(0), out.stride(1),\n            out.shape[2],\n            depth=depth,\n            log=ctx.log_outputs,\n            causal=ctx.causal\n        )\n        # fmt: on\n        return grad_in.reshape_as(grad_out), None, None, None\n\n\ndef softmax(\n    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False\n) -> torch.Tensor:\n    r\"\"\"Applies the Softmax function to an 3-dimensional input Tensor\n    rescaling them so that the elements of the n-dimensional output Tensor\n    lie in the range [0,1] and sum to 1.\n\n    Softmax is defined as:\n\n    .. math::\n        \\text{Softmax}(x_{i}) = \\frac{\\exp(x_i)}{\\sum_j \\exp(x_j)}\n\n    .. warning: softmax is computed on the last dimension of the input tensor.\n\n\n    Args:\n        x: input tensor.\n        mask: optional mask, its application will be fused to the softmax computation if triton is used\n        causal: optional performance optimization, if triton is used and the attention is causal\n\n    Returns:\n        a Tensor of the same dimension and shape as the input with\n        values in the range [0, 1] and sum to 1\n    \"\"\"\n    return _softmax_dispatch(x, log=False, mask=mask, causal=causal)\n\n\ndef log_softmax(\n    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False\n) -> torch.Tensor:\n    r\"\"\"Applies the :math:`\\log(\\text{Softmax}(x))` function to an 3-dimensional\n    input Tensor. The LogSoftmax formulation can be simplified as:\n\n    .. math::\n        \\text{LogSoftmax}(x_{i}) = \\log\\left(\\frac{\\exp(x_i) }{ \\sum_j \\exp(x_j)} \\right)\n\n    Args:\n        x: input tensor.\n\n    Returns:\n        a Tensor of the same dimension and shape as the input with\n        values in the range [-inf, 0)\n    \"\"\"\n    return _softmax_dispatch(x, log=True, mask=mask, causal=causal)\n\n\ndef _softmax_dispatch(\n    x: torch.Tensor, log: bool, mask: Optional[torch.Tensor], causal: bool = False\n) -> torch.Tensor:\n    # Triton is used if\n    # - CUDA\n    # - there's enough data to make it faster than pytorch. This could change over time, Triton is improving\n    # - there was no previous failure\n\n    global _triton_registered_overflow\n    global _triton_registered_warnings\n\n    try:\n        if torch.cuda.is_available() and x.is_cuda and not _triton_registered_overflow:\n            return _softmax_triton.apply(x, mask, log, causal)\n    except (triton.code_gen.OutOfResources, RuntimeError) as e:\n        # Catch cases where the current GPU does not have enough registers to hold a full tensor line\n        # fallback to PyTorch's implementation, which streams the tensor in and out\n        _triton_registered_overflow = True\n        logging.warning(\n            \"Triton softmax kernel register spillover or invalid image caught.\"\n            \"Deactivating this kernel, please file an issue int the xFormers repository\"\n        )\n        logging.warning(e)\n\n    if causal and not _triton_registered_warnings:\n        logging.warning(\n            \"Triton softmax could not be used. \\\n                The causal flags is being passed but it does not provide any benefit with PyTorch softmax.\"\n        )\n        _triton_registered_warnings = True\n\n    if mask is not None:\n        x = x + mask\n\n    if log:\n        return torch.log_softmax(x, dim=-1)\n    else:\n        return torch.softmax(x, dim=-1)\n", "path": "xformers/triton/softmax.py"}]}