We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Can't deepcopy an xformer model with triton 2 update
# 🐛 Bug

Given the recent triton2 update an xformer model cannot be deep-copied. This is an important requirement for numerous tasks including EMA (without knowledge of the generating class / hyper-params).

## xformers ViT-B Config

```yaml
reversible: False
block_type: "encoder"
num_layers: 12
dim_model: 768
layer_norm_style: "pre"

multi_head_config:
  num_heads: 12
  residual_dropout: 0.1  # (1) tried without this, (2) swapping this for DropPath, (3) with regular dropout
  use_rotary_embeddings: False

  attention:
    name: "scaled_dot_product"
    dropout: 0.0
    causal: False

feedforward_config:
  name: "MLP"
  dropout: 0.0
  activation: "gelu"
  hidden_layer_multiplier: 4
```

## To reproduce

```python
from copy import deepcopy

with open(transfomer_config_file, "rb") as fileptr:
    self.model_config = yaml.load(fileptr, Loader=yaml.FullLoader)

model = xFormer.from_config(xFormerConfig([self.model_config]))
deepcopy(model)
```

Error is:
```bash
TypeError: cannot pickle 'PyCapsule' object
```

--- 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/dropout.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 # CREDITS: This comes almost as-is from the Triton dropout tutorial
8 # https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py
9 
10 from typing import Optional
11 
12 import torch
13 import triton
14 from torch.cuda.amp import custom_bwd, custom_fwd
15 
16 from xformers.components.activations import Activation, build_activation
17 from xformers.triton.k_activations import (
18     get_triton_activation_bwd_kernel,
19     get_triton_activation_kernel,
20 )
21 from xformers.triton.k_dropout import k_dropout_bw, k_dropout_fw
22 
23 GROUP_M = 32
24 BLOCK_M = GROUP_M // 4
25 BLOCK_N = 128
26 
27 
28 # Helper to handle the SPMD launch grid and error cases
29 class _dropout(torch.autograd.Function):
30     @staticmethod
31     @custom_fwd(cast_inputs=torch.float16)
32     def forward(ctx, x, p, bias, activation, activation_grad, trainable_bias):
33         # Soft-flatten an hypothetical 3rd dimension
34         x_ = x.reshape(-1, x.shape[-1]).contiguous()
35         y = torch.empty_like(x_)
36         M, N = x_.shape
37 
38         assert bias is None or (bias.dtype == x.dtype and bias.shape[0] == N)
39         assert p > 0.0
40 
41         def grid(meta):
42             # NOTE: We use Triton Philox random number generator, which optimally generates 4 blocks for
43             # a given seed and offsets. "BLOCK_M" here describes the size of one of these blocks
44             # but we need to take this factor of 4 into account when scheduling all the kernels
45             return (
46                 triton.cdiv(M, meta["BLOCK_M"] * 4),
47                 triton.cdiv(N, meta["BLOCK_N"]),
48             )
49 
50         N_BLOCK_N = triton.cdiv(N, BLOCK_N)
51 
52         # Generate one seed per sample
53         # seed max is int32 max for positive numbers: 2**16
54         seeds = torch.randint(65536, (N_BLOCK_N,), device=x.device).to(torch.int32)
55 
56         # fmt: off
57         bias_ptr = bias if bias is not None else x_  # Possibly not being used
58 
59         k_dropout_fw[grid](
60             y, x_,
61             bias_ptr,
62             seeds,
63             y.stride(0),
64             M, N,
65             p,
66             x.dtype == torch.float16,
67             USE_BIAS=bias is not None,
68             ACTIVATION=activation,
69             BLOCK_M=BLOCK_M,
70             BLOCK_N=BLOCK_N,
71         )
72         # fmt: on
73 
74         if activation is not None:
75             ctx.save_for_backward(seeds, bias, x)
76         else:
77             ctx.save_for_backward(seeds, bias, None)
78 
79         ctx.trainable_bias = bias is not None and trainable_bias
80         ctx.activation_grad = activation_grad
81         ctx.p = p
82 
83         return y.reshape_as(x)
84 
85     @staticmethod
86     @custom_bwd
87     def backward(ctx, grad_out):
88         (seeds, bias, inputs) = ctx.saved_tensors
89 
90         # Soft-flatten an hypothetical 3rd dimension
91         grad_out_ = grad_out.reshape(-1, grad_out.shape[-1]).contiguous()
92         grad_in = torch.empty_like(grad_out_)
93 
94         M, N = grad_out_.shape
95 
96         # Optional inputs to compute the activation contribution to the gradient
97         assert inputs is not None or ctx.activation_grad is None
98 
99         if inputs is None:
100             inputs = grad_out_
101         elif inputs.ndim > 2:
102             inputs = inputs.reshape(-1, N)
103 
104         # We split the problem in tiles:
105         # - over M there will be a follow up reduction
106         # - over M, we go by 4 tiles at at time (consequence of the random number generation)
107         # - over N we compromise in between trying to use as much memory paralellism as possible,
108         # (fill in the warps, there are 32 threads per warps, and 4 warps default), and not being too
109         # big because of register spilling
110         N_BLOCKS_M = triton.cdiv(M, GROUP_M)
111 
112         if ctx.trainable_bias:
113             grad_bias = torch.empty(
114                 (
115                     N_BLOCKS_M,
116                     N,
117                 ),
118                 device=grad_in.device,
119                 dtype=grad_in.dtype,
120             )
121 
122         else:
123             grad_bias = grad_in  # will not be used
124 
125         def grid(meta):
126             # NOTE: We use Triton Philox random number generator, which optimally generates 4 blocks for
127             # a given seed and offsets. "BLOCK_M" here describes the size of one of these blocks
128             # but we need to take this factor of 4 into account when scheduling all the kernels
129             return (
130                 triton.cdiv(M, meta["BLOCK_M"] * 4),
131                 triton.cdiv(N, meta["BLOCK_N"]),
132             )
133 
134         # fmt: off
135         k_dropout_bw[grid](
136             grad_in, grad_bias, grad_out_,
137             inputs, bias if bias is not None else inputs,
138             seeds,
139             grad_out_.stride(0), inputs.stride(0),
140             M, N,
141             ctx.p,
142             grad_in.dtype == torch.float16,
143             USE_BIAS=bias is not None,
144             ACTIVATION_GRAD=ctx.activation_grad,
145             TRAINABLE_BIAS=ctx.trainable_bias,
146             BLOCK_M=BLOCK_M,
147             BLOCK_N=BLOCK_N,
148         )
149         # fmt: on
150 
151         return (
152             grad_in.reshape_as(grad_out),
153             None,
154             torch.sum(grad_bias, dim=0) if ctx.trainable_bias else None,
155             None,
156             None,
157             None,
158         )
159 
160 
161 def dropout(
162     x: torch.Tensor,
163     p: float,
164     bias: Optional[torch.Tensor] = None,
165     activation: Optional[Activation] = None,
166 ):
167     """
168     Apply dropout on the input tensor.
169     Optionally add a bias, the computation will be fused.
170     """
171 
172     assert p <= 1.0 and p >= 0.0
173 
174     if p == 1.0:
175         return torch.zeros_like(x)
176 
177     # Micro optim, skip dropout
178     if p == 0.0:
179         x = x + bias if bias is not None else x
180         if activation is not None:
181             activation_fn = build_activation(activation)
182             return activation_fn(x)
183         return x
184 
185     # The normal triton enabled codepath
186     act_kernel = get_triton_activation_kernel(activation)
187     act_grad_kernel = get_triton_activation_bwd_kernel(activation)
188     return _dropout.apply(
189         x,
190         float(p),
191         bias,
192         act_kernel,
193         act_grad_kernel,
194         bias is not None and bias.requires_grad,
195     )
196 
197 
198 class FusedDropoutBias(torch.nn.Module):
199     def __init__(
200         self,
201         p: float,
202         bias_shape: Optional[int],
203         activation: Optional[Activation] = None,
204     ) -> None:
205         super().__init__()
206 
207         self.p = float(p)
208 
209         assert (
210             self.p < 1.0
211         ), f"We don't want to drop all the values, most probably p={p} is not properly set"
212 
213         self.activation_type = activation
214         self.bias = (
215             torch.zeros(bias_shape, requires_grad=True)
216             if bias_shape is not None
217             else None
218         )
219         self.activation = get_triton_activation_kernel(activation)
220         self.pytorch_activation = build_activation(self.activation_type)
221         self.activation_grad = get_triton_activation_bwd_kernel(activation)
222 
223     def forward(self, x: torch.Tensor) -> torch.Tensor:
224         # Convenience, catch a possible type or device mismatch
225         if self.bias is not None:
226             self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore
227 
228         # Train/inference
229         p = self.p if self.training else 0.0
230 
231         # This kernel is slower than pytorch for small buffers, bypassing it in that case
232         perf_check = x.shape[-1] > 512
233 
234         # Catch a non-cuda setup, fallback to pytorch
235         if not x.is_cuda or not perf_check or p == 0.0:
236             x = x + self.bias if self.bias is not None else x
237             x = self.pytorch_activation(x)
238             return torch.nn.functional.dropout(x, p) if p > 0.0 else x
239 
240         # The normal, Triton-backed path
241         return _dropout.apply(
242             x, p, self.bias, self.activation, self.activation_grad, True
243         )
244
```
--- 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/dropout.py b/xformers/triton/dropout.py
--- a/xformers/triton/dropout.py
+++ b/xformers/triton/dropout.py
@@ -7,7 +7,7 @@
 # CREDITS: This comes almost as-is from the Triton dropout tutorial
 # https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py
 
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 import triton
@@ -196,6 +196,11 @@
 
 
 class FusedDropoutBias(torch.nn.Module):
+    """
+    A layer which fuses the computation of Dropout(Activation(x))
+    in a single GPU kernel
+    """
+
     def __init__(
         self,
         p: float,
@@ -216,15 +221,24 @@
             if bias_shape is not None
             else None
         )
-        self.activation = get_triton_activation_kernel(activation)
-        self.pytorch_activation = build_activation(self.activation_type)
-        self.activation_grad = get_triton_activation_bwd_kernel(activation)
+
+        self.activation: Optional[Any] = None
+        self.activation_grad: Optional[Any] = None
+        self.activation_pytorch: Optional[Any] = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Convenience, catch a possible type or device mismatch
         if self.bias is not None:
             self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore
 
+        # Lazy init (helps with pickling)
+        if self.activation is None:
+            self.activation = get_triton_activation_kernel(self.activation_type)
+            self.pytorch_activation = build_activation(self.activation_type)
+            self.activation_grad = get_triton_activation_bwd_kernel(
+                self.activation_type
+            )
+
         # Train/inference
         p = self.p if self.training else 0.0

{"golden_diff": "diff --git a/xformers/triton/dropout.py b/xformers/triton/dropout.py\n--- a/xformers/triton/dropout.py\n+++ b/xformers/triton/dropout.py\n@@ -7,7 +7,7 @@\n # CREDITS: This comes almost as-is from the Triton dropout tutorial\n # https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py\n \n-from typing import Optional\n+from typing import Any, Optional\n \n import torch\n import triton\n@@ -196,6 +196,11 @@\n \n \n class FusedDropoutBias(torch.nn.Module):\n+    \"\"\"\n+    A layer which fuses the computation of Dropout(Activation(x))\n+    in a single GPU kernel\n+    \"\"\"\n+\n     def __init__(\n         self,\n         p: float,\n@@ -216,15 +221,24 @@\n             if bias_shape is not None\n             else None\n         )\n-        self.activation = get_triton_activation_kernel(activation)\n-        self.pytorch_activation = build_activation(self.activation_type)\n-        self.activation_grad = get_triton_activation_bwd_kernel(activation)\n+\n+        self.activation: Optional[Any] = None\n+        self.activation_grad: Optional[Any] = None\n+        self.activation_pytorch: Optional[Any] = None\n \n     def forward(self, x: torch.Tensor) -> torch.Tensor:\n         # Convenience, catch a possible type or device mismatch\n         if self.bias is not None:\n             self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore\n \n+        # Lazy init (helps with pickling)\n+        if self.activation is None:\n+            self.activation = get_triton_activation_kernel(self.activation_type)\n+            self.pytorch_activation = build_activation(self.activation_type)\n+            self.activation_grad = get_triton_activation_bwd_kernel(\n+                self.activation_type\n+            )\n+\n         # Train/inference\n         p = self.p if self.training else 0.0\n", "issue": "Can't deepcopy an xformer model with triton 2 update\n# \ud83d\udc1b Bug\r\n\r\nGiven the recent triton2 update an xformer model cannot be deep-copied. This is an important requirement for numerous tasks including EMA (without knowledge of the generating class / hyper-params).\r\n\r\n## xformers ViT-B Config\r\n\r\n```yaml\r\nreversible: False\r\nblock_type: \"encoder\"\r\nnum_layers: 12\r\ndim_model: 768\r\nlayer_norm_style: \"pre\"\r\n\r\nmulti_head_config:\r\n  num_heads: 12\r\n  residual_dropout: 0.1  # (1) tried without this, (2) swapping this for DropPath, (3) with regular dropout\r\n  use_rotary_embeddings: False\r\n\r\n  attention:\r\n    name: \"scaled_dot_product\"\r\n    dropout: 0.0\r\n    causal: False\r\n\r\nfeedforward_config:\r\n  name: \"MLP\"\r\n  dropout: 0.0\r\n  activation: \"gelu\"\r\n  hidden_layer_multiplier: 4\r\n```\r\n\r\n## To reproduce\r\n\r\n```python\r\nfrom copy import deepcopy\r\n\r\nwith open(transfomer_config_file, \"rb\") as fileptr:\r\n    self.model_config = yaml.load(fileptr, Loader=yaml.FullLoader)\r\n\r\nmodel = xFormer.from_config(xFormerConfig([self.model_config]))\r\ndeepcopy(model)\r\n```\r\n\r\nError is:\r\n```bash\r\nTypeError: cannot pickle 'PyCapsule' object\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\n# CREDITS: This comes almost as-is from the Triton dropout tutorial\n# https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py\n\nfrom typing import Optional\n\nimport torch\nimport triton\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\nfrom xformers.components.activations import Activation, build_activation\nfrom xformers.triton.k_activations import (\n    get_triton_activation_bwd_kernel,\n    get_triton_activation_kernel,\n)\nfrom xformers.triton.k_dropout import k_dropout_bw, k_dropout_fw\n\nGROUP_M = 32\nBLOCK_M = GROUP_M // 4\nBLOCK_N = 128\n\n\n# Helper to handle the SPMD launch grid and error cases\nclass _dropout(torch.autograd.Function):\n    @staticmethod\n    @custom_fwd(cast_inputs=torch.float16)\n    def forward(ctx, x, p, bias, activation, activation_grad, trainable_bias):\n        # Soft-flatten an hypothetical 3rd dimension\n        x_ = x.reshape(-1, x.shape[-1]).contiguous()\n        y = torch.empty_like(x_)\n        M, N = x_.shape\n\n        assert bias is None or (bias.dtype == x.dtype and bias.shape[0] == N)\n        assert p > 0.0\n\n        def grid(meta):\n            # NOTE: We use Triton Philox random number generator, which optimally generates 4 blocks for\n            # a given seed and offsets. \"BLOCK_M\" here describes the size of one of these blocks\n            # but we need to take this factor of 4 into account when scheduling all the kernels\n            return (\n                triton.cdiv(M, meta[\"BLOCK_M\"] * 4),\n                triton.cdiv(N, meta[\"BLOCK_N\"]),\n            )\n\n        N_BLOCK_N = triton.cdiv(N, BLOCK_N)\n\n        # Generate one seed per sample\n        # seed max is int32 max for positive numbers: 2**16\n        seeds = torch.randint(65536, (N_BLOCK_N,), device=x.device).to(torch.int32)\n\n        # fmt: off\n        bias_ptr = bias if bias is not None else x_  # Possibly not being used\n\n        k_dropout_fw[grid](\n            y, x_,\n            bias_ptr,\n            seeds,\n            y.stride(0),\n            M, N,\n            p,\n            x.dtype == torch.float16,\n            USE_BIAS=bias is not None,\n            ACTIVATION=activation,\n            BLOCK_M=BLOCK_M,\n            BLOCK_N=BLOCK_N,\n        )\n        # fmt: on\n\n        if activation is not None:\n            ctx.save_for_backward(seeds, bias, x)\n        else:\n            ctx.save_for_backward(seeds, bias, None)\n\n        ctx.trainable_bias = bias is not None and trainable_bias\n        ctx.activation_grad = activation_grad\n        ctx.p = p\n\n        return y.reshape_as(x)\n\n    @staticmethod\n    @custom_bwd\n    def backward(ctx, grad_out):\n        (seeds, bias, inputs) = ctx.saved_tensors\n\n        # Soft-flatten an hypothetical 3rd dimension\n        grad_out_ = grad_out.reshape(-1, grad_out.shape[-1]).contiguous()\n        grad_in = torch.empty_like(grad_out_)\n\n        M, N = grad_out_.shape\n\n        # Optional inputs to compute the activation contribution to the gradient\n        assert inputs is not None or ctx.activation_grad is None\n\n        if inputs is None:\n            inputs = grad_out_\n        elif inputs.ndim > 2:\n            inputs = inputs.reshape(-1, N)\n\n        # We split the problem in tiles:\n        # - over M there will be a follow up reduction\n        # - over M, we go by 4 tiles at at time (consequence of the random number generation)\n        # - over N we compromise in between trying to use as much memory paralellism as possible,\n        # (fill in the warps, there are 32 threads per warps, and 4 warps default), and not being too\n        # big because of register spilling\n        N_BLOCKS_M = triton.cdiv(M, GROUP_M)\n\n        if ctx.trainable_bias:\n            grad_bias = torch.empty(\n                (\n                    N_BLOCKS_M,\n                    N,\n                ),\n                device=grad_in.device,\n                dtype=grad_in.dtype,\n            )\n\n        else:\n            grad_bias = grad_in  # will not be used\n\n        def grid(meta):\n            # NOTE: We use Triton Philox random number generator, which optimally generates 4 blocks for\n            # a given seed and offsets. \"BLOCK_M\" here describes the size of one of these blocks\n            # but we need to take this factor of 4 into account when scheduling all the kernels\n            return (\n                triton.cdiv(M, meta[\"BLOCK_M\"] * 4),\n                triton.cdiv(N, meta[\"BLOCK_N\"]),\n            )\n\n        # fmt: off\n        k_dropout_bw[grid](\n            grad_in, grad_bias, grad_out_,\n            inputs, bias if bias is not None else inputs,\n            seeds,\n            grad_out_.stride(0), inputs.stride(0),\n            M, N,\n            ctx.p,\n            grad_in.dtype == torch.float16,\n            USE_BIAS=bias is not None,\n            ACTIVATION_GRAD=ctx.activation_grad,\n            TRAINABLE_BIAS=ctx.trainable_bias,\n            BLOCK_M=BLOCK_M,\n            BLOCK_N=BLOCK_N,\n        )\n        # fmt: on\n\n        return (\n            grad_in.reshape_as(grad_out),\n            None,\n            torch.sum(grad_bias, dim=0) if ctx.trainable_bias else None,\n            None,\n            None,\n            None,\n        )\n\n\ndef dropout(\n    x: torch.Tensor,\n    p: float,\n    bias: Optional[torch.Tensor] = None,\n    activation: Optional[Activation] = None,\n):\n    \"\"\"\n    Apply dropout on the input tensor.\n    Optionally add a bias, the computation will be fused.\n    \"\"\"\n\n    assert p <= 1.0 and p >= 0.0\n\n    if p == 1.0:\n        return torch.zeros_like(x)\n\n    # Micro optim, skip dropout\n    if p == 0.0:\n        x = x + bias if bias is not None else x\n        if activation is not None:\n            activation_fn = build_activation(activation)\n            return activation_fn(x)\n        return x\n\n    # The normal triton enabled codepath\n    act_kernel = get_triton_activation_kernel(activation)\n    act_grad_kernel = get_triton_activation_bwd_kernel(activation)\n    return _dropout.apply(\n        x,\n        float(p),\n        bias,\n        act_kernel,\n        act_grad_kernel,\n        bias is not None and bias.requires_grad,\n    )\n\n\nclass FusedDropoutBias(torch.nn.Module):\n    def __init__(\n        self,\n        p: float,\n        bias_shape: Optional[int],\n        activation: Optional[Activation] = None,\n    ) -> None:\n        super().__init__()\n\n        self.p = float(p)\n\n        assert (\n            self.p < 1.0\n        ), f\"We don't want to drop all the values, most probably p={p} is not properly set\"\n\n        self.activation_type = activation\n        self.bias = (\n            torch.zeros(bias_shape, requires_grad=True)\n            if bias_shape is not None\n            else None\n        )\n        self.activation = get_triton_activation_kernel(activation)\n        self.pytorch_activation = build_activation(self.activation_type)\n        self.activation_grad = get_triton_activation_bwd_kernel(activation)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        # Convenience, catch a possible type or device mismatch\n        if self.bias is not None:\n            self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore\n\n        # Train/inference\n        p = self.p if self.training else 0.0\n\n        # This kernel is slower than pytorch for small buffers, bypassing it in that case\n        perf_check = x.shape[-1] > 512\n\n        # Catch a non-cuda setup, fallback to pytorch\n        if not x.is_cuda or not perf_check or p == 0.0:\n            x = x + self.bias if self.bias is not None else x\n            x = self.pytorch_activation(x)\n            return torch.nn.functional.dropout(x, p) if p > 0.0 else x\n\n        # The normal, Triton-backed path\n        return _dropout.apply(\n            x, p, self.bias, self.activation, self.activation_grad, True\n        )\n", "path": "xformers/triton/dropout.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\n# CREDITS: This comes almost as-is from the Triton dropout tutorial\n# https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py\n\nfrom typing import Any, Optional\n\nimport torch\nimport triton\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\nfrom xformers.components.activations import Activation, build_activation\nfrom xformers.triton.k_activations import (\n    get_triton_activation_bwd_kernel,\n    get_triton_activation_kernel,\n)\nfrom xformers.triton.k_dropout import k_dropout_bw, k_dropout_fw\n\nGROUP_M = 32\nBLOCK_M = GROUP_M // 4\nBLOCK_N = 128\n\n\n# Helper to handle the SPMD launch grid and error cases\nclass _dropout(torch.autograd.Function):\n    @staticmethod\n    @custom_fwd(cast_inputs=torch.float16)\n    def forward(ctx, x, p, bias, activation, activation_grad, trainable_bias):\n        # Soft-flatten an hypothetical 3rd dimension\n        x_ = x.reshape(-1, x.shape[-1]).contiguous()\n        y = torch.empty_like(x_)\n        M, N = x_.shape\n\n        assert bias is None or (bias.dtype == x.dtype and bias.shape[0] == N)\n        assert p > 0.0\n\n        def grid(meta):\n            # NOTE: We use Triton Philox random number generator, which optimally generates 4 blocks for\n            # a given seed and offsets. \"BLOCK_M\" here describes the size of one of these blocks\n            # but we need to take this factor of 4 into account when scheduling all the kernels\n            return (\n                triton.cdiv(M, meta[\"BLOCK_M\"] * 4),\n                triton.cdiv(N, meta[\"BLOCK_N\"]),\n            )\n\n        N_BLOCK_N = triton.cdiv(N, BLOCK_N)\n\n        # Generate one seed per sample\n        # seed max is int32 max for positive numbers: 2**16\n        seeds = torch.randint(65536, (N_BLOCK_N,), device=x.device).to(torch.int32)\n\n        # fmt: off\n        bias_ptr = bias if bias is not None else x_  # Possibly not being used\n\n        k_dropout_fw[grid](\n            y, x_,\n            bias_ptr,\n            seeds,\n            y.stride(0),\n            M, N,\n            p,\n            x.dtype == torch.float16,\n            USE_BIAS=bias is not None,\n            ACTIVATION=activation,\n            BLOCK_M=BLOCK_M,\n            BLOCK_N=BLOCK_N,\n        )\n        # fmt: on\n\n        if activation is not None:\n            ctx.save_for_backward(seeds, bias, x)\n        else:\n            ctx.save_for_backward(seeds, bias, None)\n\n        ctx.trainable_bias = bias is not None and trainable_bias\n        ctx.activation_grad = activation_grad\n        ctx.p = p\n\n        return y.reshape_as(x)\n\n    @staticmethod\n    @custom_bwd\n    def backward(ctx, grad_out):\n        (seeds, bias, inputs) = ctx.saved_tensors\n\n        # Soft-flatten an hypothetical 3rd dimension\n        grad_out_ = grad_out.reshape(-1, grad_out.shape[-1]).contiguous()\n        grad_in = torch.empty_like(grad_out_)\n\n        M, N = grad_out_.shape\n\n        # Optional inputs to compute the activation contribution to the gradient\n        assert inputs is not None or ctx.activation_grad is None\n\n        if inputs is None:\n            inputs = grad_out_\n        elif inputs.ndim > 2:\n            inputs = inputs.reshape(-1, N)\n\n        # We split the problem in tiles:\n        # - over M there will be a follow up reduction\n        # - over M, we go by 4 tiles at at time (consequence of the random number generation)\n        # - over N we compromise in between trying to use as much memory paralellism as possible,\n        # (fill in the warps, there are 32 threads per warps, and 4 warps default), and not being too\n        # big because of register spilling\n        N_BLOCKS_M = triton.cdiv(M, GROUP_M)\n\n        if ctx.trainable_bias:\n            grad_bias = torch.empty(\n                (\n                    N_BLOCKS_M,\n                    N,\n                ),\n                device=grad_in.device,\n                dtype=grad_in.dtype,\n            )\n\n        else:\n            grad_bias = grad_in  # will not be used\n\n        def grid(meta):\n            # NOTE: We use Triton Philox random number generator, which optimally generates 4 blocks for\n            # a given seed and offsets. \"BLOCK_M\" here describes the size of one of these blocks\n            # but we need to take this factor of 4 into account when scheduling all the kernels\n            return (\n                triton.cdiv(M, meta[\"BLOCK_M\"] * 4),\n                triton.cdiv(N, meta[\"BLOCK_N\"]),\n            )\n\n        # fmt: off\n        k_dropout_bw[grid](\n            grad_in, grad_bias, grad_out_,\n            inputs, bias if bias is not None else inputs,\n            seeds,\n            grad_out_.stride(0), inputs.stride(0),\n            M, N,\n            ctx.p,\n            grad_in.dtype == torch.float16,\n            USE_BIAS=bias is not None,\n            ACTIVATION_GRAD=ctx.activation_grad,\n            TRAINABLE_BIAS=ctx.trainable_bias,\n            BLOCK_M=BLOCK_M,\n            BLOCK_N=BLOCK_N,\n        )\n        # fmt: on\n\n        return (\n            grad_in.reshape_as(grad_out),\n            None,\n            torch.sum(grad_bias, dim=0) if ctx.trainable_bias else None,\n            None,\n            None,\n            None,\n        )\n\n\ndef dropout(\n    x: torch.Tensor,\n    p: float,\n    bias: Optional[torch.Tensor] = None,\n    activation: Optional[Activation] = None,\n):\n    \"\"\"\n    Apply dropout on the input tensor.\n    Optionally add a bias, the computation will be fused.\n    \"\"\"\n\n    assert p <= 1.0 and p >= 0.0\n\n    if p == 1.0:\n        return torch.zeros_like(x)\n\n    # Micro optim, skip dropout\n    if p == 0.0:\n        x = x + bias if bias is not None else x\n        if activation is not None:\n            activation_fn = build_activation(activation)\n            return activation_fn(x)\n        return x\n\n    # The normal triton enabled codepath\n    act_kernel = get_triton_activation_kernel(activation)\n    act_grad_kernel = get_triton_activation_bwd_kernel(activation)\n    return _dropout.apply(\n        x,\n        float(p),\n        bias,\n        act_kernel,\n        act_grad_kernel,\n        bias is not None and bias.requires_grad,\n    )\n\n\nclass FusedDropoutBias(torch.nn.Module):\n    \"\"\"\n    A layer which fuses the computation of Dropout(Activation(x))\n    in a single GPU kernel\n    \"\"\"\n\n    def __init__(\n        self,\n        p: float,\n        bias_shape: Optional[int],\n        activation: Optional[Activation] = None,\n    ) -> None:\n        super().__init__()\n\n        self.p = float(p)\n\n        assert (\n            self.p < 1.0\n        ), f\"We don't want to drop all the values, most probably p={p} is not properly set\"\n\n        self.activation_type = activation\n        self.bias = (\n            torch.zeros(bias_shape, requires_grad=True)\n            if bias_shape is not None\n            else None\n        )\n\n        self.activation: Optional[Any] = None\n        self.activation_grad: Optional[Any] = None\n        self.activation_pytorch: Optional[Any] = None\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        # Convenience, catch a possible type or device mismatch\n        if self.bias is not None:\n            self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore\n\n        # Lazy init (helps with pickling)\n        if self.activation is None:\n            self.activation = get_triton_activation_kernel(self.activation_type)\n            self.pytorch_activation = build_activation(self.activation_type)\n            self.activation_grad = get_triton_activation_bwd_kernel(\n                self.activation_type\n            )\n\n        # Train/inference\n        p = self.p if self.training else 0.0\n\n        # This kernel is slower than pytorch for small buffers, bypassing it in that case\n        perf_check = x.shape[-1] > 512\n\n        # Catch a non-cuda setup, fallback to pytorch\n        if not x.is_cuda or not perf_check or p == 0.0:\n            x = x + self.bias if self.bias is not None else x\n            x = self.pytorch_activation(x)\n            return torch.nn.functional.dropout(x, p) if p > 0.0 else x\n\n        # The normal, Triton-backed path\n        return _dropout.apply(\n            x, p, self.bias, self.activation, self.activation_grad, True\n        )\n", "path": "xformers/triton/dropout.py"}]}