prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  [bug] Failing triton dropout test 
  # 🐛 Bug
  
  See https://app.circleci.com/pipelines/github/facebookresearch/xformers/212/workflows/8988c71c-84f5-4bd0-bd59-ac7d293c2370/jobs/398
  
  Not sure why this happens just now, looking into that
  
  ## Command
  can repro locally with ` pytest tests -k dropout -x -v `
  
  --- 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/k_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 import triton
  11 import triton.language as tl
  12 
  13 
  14 # fmt: off
  15 @triton.autotune(
  16     configs=[
  17         triton.Config({"BLOCK_SIZE" : 256}, num_warps=1),
  18         triton.Config({"BLOCK_SIZE" : 512}, num_warps=2),
  19         triton.Config({"BLOCK_SIZE" : 1024}, num_warps=4),
  20         triton.Config({"BLOCK_SIZE" : 2048}, num_warps=8),
  21         triton.Config({"BLOCK_SIZE" : 4096}, num_warps=8),
  22     ],
  23     key=["N"],
  24 )
  25 @triton.jit
  26 def k_dropout(
  27     Y, X, S,
  28     stride,
  29     N,
  30     p,
  31     **meta,
  32 ):
  33     """
  34     Apply dropout on an input tensor
  35     Y : Output (M, N)
  36     X : Input (M, N)
  37     S : Seeds (M,)
  38     p : dropout probability
  39     """
  40     # fmt: on
  41 
  42     # compute memory offsets of elements handled by this instance
  43     BLOCK_SIZE = meta["BLOCK_SIZE"]
  44     row = tl.program_id(axis=0)
  45     col = tl.program_id(axis=1)
  46     offsets = row * stride + col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
  47     mask = col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < N
  48 
  49     # load data from x
  50     x_ptrs = X + offsets
  51     x = tl.load(x_ptrs, mask=mask)
  52 
  53     # randomly prune it
  54     seed = S + row
  55     random = tl.rand(seed.to(tl.int32), offsets)
  56     x_keep = random > p
  57 
  58     # write-back
  59     zero = 0.
  60     zero = zero.to(x.dtype)
  61     output = tl.where(x_keep, (x / (1 - p)).to(x.dtype), zero)
  62     y_ptrs = Y + offsets
  63     tl.store(y_ptrs, output, mask=mask)
  64
  ```
  
  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 import torch
  11 import triton
  12 from torch.cuda.amp import custom_bwd, custom_fwd
  13 
  14 from xformers.triton.k_dropout import k_dropout
  15 
  16 
  17 # Helper to handle the SPMD launch grid and error cases
  18 class _dropout(torch.autograd.Function):
  19     @staticmethod
  20     @custom_fwd(cast_inputs=torch.float16)
  21     def forward(ctx, x, p):
  22         # Soft-flatten an hypothetical 3rd dimension
  23         x_ = x.reshape(-1, x.shape[-1])
  24         y = torch.empty_like(x_)
  25         _, N = x_.shape
  26 
  27         assert y.stride(-1) == 1 and x_.stride(-1) == 1
  28 
  29         # Generate one seed per sample
  30         # seed max is int32 max for positive numbers: 2**16
  31         seeds = torch.randint(65536, (x_.shape[0],), device=x.device).to(torch.int32)
  32 
  33         # SPMD launch grid
  34         def grid(meta):
  35             return (
  36                 x_.shape[0],
  37                 triton.cdiv(x_.shape[1], meta["BLOCK_SIZE"]),
  38             )
  39 
  40         k_dropout[grid](y, x_, seeds, y.stride(0), N, p)
  41 
  42         ctx.save_for_backward(seeds)
  43         ctx.p = p
  44 
  45         return y.reshape_as(x)
  46 
  47     @staticmethod
  48     @custom_bwd
  49     def backward(ctx, grad_out):
  50         (seeds,) = ctx.saved_tensors
  51 
  52         # Soft-flatten an hypothetical 3rd dimension
  53         grad_out_ = grad_out.reshape(-1, grad_out.shape[-1])
  54         grad_in = torch.empty_like(grad_out_)
  55         _, N = grad_out_.shape
  56 
  57         assert grad_in.stride(-1) == 1 and grad_out_.stride(-1) == 1
  58 
  59         # SPMD launch grid
  60         def grid(meta):
  61             return (
  62                 grad_out_.shape[0],
  63                 triton.cdiv(grad_out_.shape[1], meta["BLOCK_SIZE"]),
  64             )
  65 
  66         k_dropout[grid](grad_in, grad_out_, seeds, grad_out_.stride(0), N, ctx.p)
  67 
  68         return grad_in.reshape_as(grad_out), None
  69 
  70 
  71 def dropout(x: torch.Tensor, p: float):
  72     return _dropout.apply(x, p)
  73
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/xformers/triton/dropout.py b/xformers/triton/dropout.py
  --- a/xformers/triton/dropout.py
  +++ b/xformers/triton/dropout.py
  @@ -69,4 +69,7 @@
   
   
   def dropout(x: torch.Tensor, p: float):
  -    return _dropout.apply(x, p)
  +    if p > 0.0:
  +        return _dropout.apply(x, p)
  +
  +    return x
  diff --git a/xformers/triton/k_dropout.py b/xformers/triton/k_dropout.py
  --- a/xformers/triton/k_dropout.py
  +++ b/xformers/triton/k_dropout.py
  @@ -24,7 +24,7 @@
   )
   @triton.jit
   def k_dropout(
  -    Y, X, S,
  +    Y, X, SEEDS,
       stride,
       N,
       p,
  @@ -51,7 +51,7 @@
       x = tl.load(x_ptrs, mask=mask)
   
       # randomly prune it
  -    seed = S + row
  +    seed = SEEDS + row
       random = tl.rand(seed.to(tl.int32), offsets)
       x_keep = random > p
  
verification_info: {"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@@ -69,4 +69,7 @@\n \n \n def dropout(x: torch.Tensor, p: float):\n-    return _dropout.apply(x, p)\n+    if p > 0.0:\n+        return _dropout.apply(x, p)\n+\n+    return x\ndiff --git a/xformers/triton/k_dropout.py b/xformers/triton/k_dropout.py\n--- a/xformers/triton/k_dropout.py\n+++ b/xformers/triton/k_dropout.py\n@@ -24,7 +24,7 @@\n )\n @triton.jit\n def k_dropout(\n-    Y, X, S,\n+    Y, X, SEEDS,\n     stride,\n     N,\n     p,\n@@ -51,7 +51,7 @@\n     x = tl.load(x_ptrs, mask=mask)\n \n     # randomly prune it\n-    seed = S + row\n+    seed = SEEDS + row\n     random = tl.rand(seed.to(tl.int32), offsets)\n     x_keep = random > p\n", "issue": "[bug] Failing triton dropout test \n# \ud83d\udc1b Bug\r\n\r\nSee https://app.circleci.com/pipelines/github/facebookresearch/xformers/212/workflows/8988c71c-84f5-4bd0-bd59-ac7d293c2370/jobs/398\r\n\r\nNot sure why this happens just now, looking into that\r\n\r\n## Command\r\ncan repro locally with ` pytest tests -k dropout -x -v `\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\nimport triton\nimport triton.language as tl\n\n\n# fmt: off\n@triton.autotune(\n    configs=[\n        triton.Config({\"BLOCK_SIZE\" : 256}, num_warps=1),\n        triton.Config({\"BLOCK_SIZE\" : 512}, num_warps=2),\n        triton.Config({\"BLOCK_SIZE\" : 1024}, num_warps=4),\n        triton.Config({\"BLOCK_SIZE\" : 2048}, num_warps=8),\n        triton.Config({\"BLOCK_SIZE\" : 4096}, num_warps=8),\n    ],\n    key=[\"N\"],\n)\n@triton.jit\ndef k_dropout(\n    Y, X, S,\n    stride,\n    N,\n    p,\n    **meta,\n):\n    \"\"\"\n    Apply dropout on an input tensor\n    Y : Output (M, N)\n    X : Input (M, N)\n    S : Seeds (M,)\n    p : dropout probability\n    \"\"\"\n    # fmt: on\n\n    # compute memory offsets of elements handled by this instance\n    BLOCK_SIZE = meta[\"BLOCK_SIZE\"]\n    row = tl.program_id(axis=0)\n    col = tl.program_id(axis=1)\n    offsets = row * stride + col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    mask = col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < N\n\n    # load data from x\n    x_ptrs = X + offsets\n    x = tl.load(x_ptrs, mask=mask)\n\n    # randomly prune it\n    seed = S + row\n    random = tl.rand(seed.to(tl.int32), offsets)\n    x_keep = random > p\n\n    # write-back\n    zero = 0.\n    zero = zero.to(x.dtype)\n    output = tl.where(x_keep, (x / (1 - p)).to(x.dtype), zero)\n    y_ptrs = Y + offsets\n    tl.store(y_ptrs, output, mask=mask)\n", "path": "xformers/triton/k_dropout.py"}, {"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\nimport torch\nimport triton\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\nfrom xformers.triton.k_dropout import k_dropout\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):\n        # Soft-flatten an hypothetical 3rd dimension\n        x_ = x.reshape(-1, x.shape[-1])\n        y = torch.empty_like(x_)\n        _, N = x_.shape\n\n        assert y.stride(-1) == 1 and x_.stride(-1) == 1\n\n        # Generate one seed per sample\n        # seed max is int32 max for positive numbers: 2**16\n        seeds = torch.randint(65536, (x_.shape[0],), device=x.device).to(torch.int32)\n\n        # SPMD launch grid\n        def grid(meta):\n            return (\n                x_.shape[0],\n                triton.cdiv(x_.shape[1], meta[\"BLOCK_SIZE\"]),\n            )\n\n        k_dropout[grid](y, x_, seeds, y.stride(0), N, p)\n\n        ctx.save_for_backward(seeds)\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,) = ctx.saved_tensors\n\n        # Soft-flatten an hypothetical 3rd dimension\n        grad_out_ = grad_out.reshape(-1, grad_out.shape[-1])\n        grad_in = torch.empty_like(grad_out_)\n        _, N = grad_out_.shape\n\n        assert grad_in.stride(-1) == 1 and grad_out_.stride(-1) == 1\n\n        # SPMD launch grid\n        def grid(meta):\n            return (\n                grad_out_.shape[0],\n                triton.cdiv(grad_out_.shape[1], meta[\"BLOCK_SIZE\"]),\n            )\n\n        k_dropout[grid](grad_in, grad_out_, seeds, grad_out_.stride(0), N, ctx.p)\n\n        return grad_in.reshape_as(grad_out), None\n\n\ndef dropout(x: torch.Tensor, p: float):\n    return _dropout.apply(x, p)\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\nimport triton\nimport triton.language as tl\n\n\n# fmt: off\n@triton.autotune(\n    configs=[\n        triton.Config({\"BLOCK_SIZE\" : 256}, num_warps=1),\n        triton.Config({\"BLOCK_SIZE\" : 512}, num_warps=2),\n        triton.Config({\"BLOCK_SIZE\" : 1024}, num_warps=4),\n        triton.Config({\"BLOCK_SIZE\" : 2048}, num_warps=8),\n        triton.Config({\"BLOCK_SIZE\" : 4096}, num_warps=8),\n    ],\n    key=[\"N\"],\n)\n@triton.jit\ndef k_dropout(\n    Y, X, SEEDS,\n    stride,\n    N,\n    p,\n    **meta,\n):\n    \"\"\"\n    Apply dropout on an input tensor\n    Y : Output (M, N)\n    X : Input (M, N)\n    S : Seeds (M,)\n    p : dropout probability\n    \"\"\"\n    # fmt: on\n\n    # compute memory offsets of elements handled by this instance\n    BLOCK_SIZE = meta[\"BLOCK_SIZE\"]\n    row = tl.program_id(axis=0)\n    col = tl.program_id(axis=1)\n    offsets = row * stride + col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    mask = col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < N\n\n    # load data from x\n    x_ptrs = X + offsets\n    x = tl.load(x_ptrs, mask=mask)\n\n    # randomly prune it\n    seed = SEEDS + row\n    random = tl.rand(seed.to(tl.int32), offsets)\n    x_keep = random > p\n\n    # write-back\n    zero = 0.\n    zero = zero.to(x.dtype)\n    output = tl.where(x_keep, (x / (1 - p)).to(x.dtype), zero)\n    y_ptrs = Y + offsets\n    tl.store(y_ptrs, output, mask=mask)\n", "path": "xformers/triton/k_dropout.py"}, {"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\nimport torch\nimport triton\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\nfrom xformers.triton.k_dropout import k_dropout\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):\n        # Soft-flatten an hypothetical 3rd dimension\n        x_ = x.reshape(-1, x.shape[-1])\n        y = torch.empty_like(x_)\n        _, N = x_.shape\n\n        assert y.stride(-1) == 1 and x_.stride(-1) == 1\n\n        # Generate one seed per sample\n        # seed max is int32 max for positive numbers: 2**16\n        seeds = torch.randint(65536, (x_.shape[0],), device=x.device).to(torch.int32)\n\n        # SPMD launch grid\n        def grid(meta):\n            return (\n                x_.shape[0],\n                triton.cdiv(x_.shape[1], meta[\"BLOCK_SIZE\"]),\n            )\n\n        k_dropout[grid](y, x_, seeds, y.stride(0), N, p)\n\n        ctx.save_for_backward(seeds)\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,) = ctx.saved_tensors\n\n        # Soft-flatten an hypothetical 3rd dimension\n        grad_out_ = grad_out.reshape(-1, grad_out.shape[-1])\n        grad_in = torch.empty_like(grad_out_)\n        _, N = grad_out_.shape\n\n        assert grad_in.stride(-1) == 1 and grad_out_.stride(-1) == 1\n\n        # SPMD launch grid\n        def grid(meta):\n            return (\n                grad_out_.shape[0],\n                triton.cdiv(grad_out_.shape[1], meta[\"BLOCK_SIZE\"]),\n            )\n\n        k_dropout[grid](grad_in, grad_out_, seeds, grad_out_.stride(0), N, ctx.p)\n\n        return grad_in.reshape_as(grad_out), None\n\n\ndef dropout(x: torch.Tensor, p: float):\n    if p > 0.0:\n        return _dropout.apply(x, p)\n\n    return x\n", "path": "xformers/triton/dropout.py"}]}