prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Flash attention (v1) is broken
  Defaulting window_size_left to "-1" breaks flash attention v1 since v1 only works if window_size_left is 0 (line 83).
  
  https://github.com/huggingface/text-generation-inference/blob/5ba53d44a18983a4de32d122f4cb46f4a17d9ef6/server/text_generation_server/utils/flash_attn.py#L60C7-L60C7
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `server/text_generation_server/utils/flash_attn.py`
  Content:
  ```
  1 import os
  2 import torch
  3 
  4 from loguru import logger
  5 
  6 if os.getenv("USE_FLASH_ATTENTION", "").lower() == "false":
  7     raise ImportError("`USE_FLASH_ATTENTION` is false.")
  8 
  9 if not torch.cuda.is_available():
  10     raise ImportError("CUDA is not available")
  11 
  12 major, minor = torch.cuda.get_device_capability()
  13 is_sm75 = major == 7 and minor == 5
  14 is_sm8x = major == 8 and minor >= 0
  15 is_sm90 = major == 9 and minor == 0
  16 
  17 HAS_FLASH_ATTN = False
  18 HAS_FLASH_ATTN_V2 = False
  19 try:
  20     try:
  21         import flash_attn_2_cuda
  22     except ImportError:
  23         raise ImportError(
  24             "Flash Attention V2 is not installed.\n"
  25             "Use the official Docker image (ghcr.io/huggingface/text-generation-inference:latest) "
  26             "or install flash attention v2 with `cd server && make install install-flash-attention-v2`"
  27         )
  28     if not (is_sm8x or is_sm90):
  29         raise ImportError(
  30             f"GPU with CUDA capability {major} {minor} is not supported for "
  31             "Flash Attention V2"
  32         )
  33     HAS_FLASH_ATTN_V2 = True
  34 except ImportError as e:
  35     try:
  36         import flash_attn_cuda
  37     except ImportError:
  38         raise ImportError(
  39             "Flash Attention is not installed.\n"
  40             "Use the official Docker image (ghcr.io/huggingface/text-generation-inference:latest) "
  41             "or install flash attention with `cd server && make install install-flash-attention`"
  42         ) from e
  43 
  44     if not (is_sm75 or is_sm8x or is_sm90):
  45         raise ImportError(
  46             f"GPU with CUDA capability {major} {minor} is not supported"
  47         ) from e
  48     logger.warning(f"Unable to use Flash Attention V2: {e}")
  49     HAS_FLASH_ATTN = True
  50 
  51 
  52 def attention(
  53     q,
  54     k,
  55     v,
  56     out,
  57     cu_seqlens,
  58     max_s,
  59     softmax_scale,
  60     window_size_left=-1,
  61 ):
  62     if HAS_FLASH_ATTN_V2:
  63         return flash_attn_2_cuda.varlen_fwd(
  64             q,
  65             k,
  66             v,
  67             out,
  68             cu_seqlens,
  69             cu_seqlens,
  70             max_s,
  71             max_s,
  72             0.0,
  73             softmax_scale,
  74             False,
  75             True,
  76             window_size_left,
  77             0,
  78             False,
  79             None,
  80         )
  81 
  82     if HAS_FLASH_ATTN:
  83         if window_size_left != 0:
  84             raise NotImplementedError(
  85                 "window_size_left is only available with flash attn v2"
  86             )
  87 
  88         # Flash attention v1 requires q, k and v to have the same number of heads
  89         if k.shape[1] != q.shape[1]:
  90             # MQA expand
  91             if k.shape[1] == 1:
  92                 k = k.expand(-1, q.shape[1], -1)
  93             # Grouped attention reshape
  94             else:
  95                 original_shape = k.shape
  96                 k = (
  97                     k.unsqueeze(2)
  98                     .expand(-1, -1, q.shape[1] // k.shape[1], -1)
  99                     .reshape(original_shape[0], -1, original_shape[2])
  100                 )
  101         if v.shape[1] != q.shape[1]:
  102             # MQA expand
  103             if v.shape[1] == 1:
  104                 v = v.expand(-1, q.shape[1], -1)
  105             # Grouped attention reshape
  106             else:
  107                 original_shape = v.shape
  108                 v = (
  109                     v.unsqueeze(2)
  110                     .expand(-1, -1, q.shape[1] // v.shape[1], -1)
  111                     .reshape(original_shape[0], -1, original_shape[2])
  112                 )
  113 
  114         return flash_attn_cuda.fwd(
  115             q,
  116             k,
  117             v,
  118             out,
  119             cu_seqlens,
  120             cu_seqlens,
  121             max_s,
  122             max_s,
  123             0.0,
  124             softmax_scale,
  125             False,
  126             True,
  127             False,
  128             0,
  129             None,
  130         )
  131 
  132     raise NotImplementedError("flash attention is not installed")
  133
  ```
  --- 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/server/text_generation_server/utils/flash_attn.py b/server/text_generation_server/utils/flash_attn.py
  --- a/server/text_generation_server/utils/flash_attn.py
  +++ b/server/text_generation_server/utils/flash_attn.py
  @@ -80,7 +80,7 @@
           )
   
       if HAS_FLASH_ATTN:
  -        if window_size_left != 0:
  +        if window_size_left != -1:
               raise NotImplementedError(
                   "window_size_left is only available with flash attn v2"
               )
  
verification_info: {"golden_diff": "diff --git a/server/text_generation_server/utils/flash_attn.py b/server/text_generation_server/utils/flash_attn.py\n--- a/server/text_generation_server/utils/flash_attn.py\n+++ b/server/text_generation_server/utils/flash_attn.py\n@@ -80,7 +80,7 @@\n         )\n \n     if HAS_FLASH_ATTN:\n-        if window_size_left != 0:\n+        if window_size_left != -1:\n             raise NotImplementedError(\n                 \"window_size_left is only available with flash attn v2\"\n             )\n", "issue": "Flash attention (v1) is broken\nDefaulting window_size_left to \"-1\" breaks flash attention v1 since v1 only works if window_size_left is 0 (line 83).\r\n\r\nhttps://github.com/huggingface/text-generation-inference/blob/5ba53d44a18983a4de32d122f4cb46f4a17d9ef6/server/text_generation_server/utils/flash_attn.py#L60C7-L60C7\n", "before_files": [{"content": "import os\nimport torch\n\nfrom loguru import logger\n\nif os.getenv(\"USE_FLASH_ATTENTION\", \"\").lower() == \"false\":\n    raise ImportError(\"`USE_FLASH_ATTENTION` is false.\")\n\nif not torch.cuda.is_available():\n    raise ImportError(\"CUDA is not available\")\n\nmajor, minor = torch.cuda.get_device_capability()\nis_sm75 = major == 7 and minor == 5\nis_sm8x = major == 8 and minor >= 0\nis_sm90 = major == 9 and minor == 0\n\nHAS_FLASH_ATTN = False\nHAS_FLASH_ATTN_V2 = False\ntry:\n    try:\n        import flash_attn_2_cuda\n    except ImportError:\n        raise ImportError(\n            \"Flash Attention V2 is not installed.\\n\"\n            \"Use the official Docker image (ghcr.io/huggingface/text-generation-inference:latest) \"\n            \"or install flash attention v2 with `cd server && make install install-flash-attention-v2`\"\n        )\n    if not (is_sm8x or is_sm90):\n        raise ImportError(\n            f\"GPU with CUDA capability {major} {minor} is not supported for \"\n            \"Flash Attention V2\"\n        )\n    HAS_FLASH_ATTN_V2 = True\nexcept ImportError as e:\n    try:\n        import flash_attn_cuda\n    except ImportError:\n        raise ImportError(\n            \"Flash Attention is not installed.\\n\"\n            \"Use the official Docker image (ghcr.io/huggingface/text-generation-inference:latest) \"\n            \"or install flash attention with `cd server && make install install-flash-attention`\"\n        ) from e\n\n    if not (is_sm75 or is_sm8x or is_sm90):\n        raise ImportError(\n            f\"GPU with CUDA capability {major} {minor} is not supported\"\n        ) from e\n    logger.warning(f\"Unable to use Flash Attention V2: {e}\")\n    HAS_FLASH_ATTN = True\n\n\ndef attention(\n    q,\n    k,\n    v,\n    out,\n    cu_seqlens,\n    max_s,\n    softmax_scale,\n    window_size_left=-1,\n):\n    if HAS_FLASH_ATTN_V2:\n        return flash_attn_2_cuda.varlen_fwd(\n            q,\n            k,\n            v,\n            out,\n            cu_seqlens,\n            cu_seqlens,\n            max_s,\n            max_s,\n            0.0,\n            softmax_scale,\n            False,\n            True,\n            window_size_left,\n            0,\n            False,\n            None,\n        )\n\n    if HAS_FLASH_ATTN:\n        if window_size_left != 0:\n            raise NotImplementedError(\n                \"window_size_left is only available with flash attn v2\"\n            )\n\n        # Flash attention v1 requires q, k and v to have the same number of heads\n        if k.shape[1] != q.shape[1]:\n            # MQA expand\n            if k.shape[1] == 1:\n                k = k.expand(-1, q.shape[1], -1)\n            # Grouped attention reshape\n            else:\n                original_shape = k.shape\n                k = (\n                    k.unsqueeze(2)\n                    .expand(-1, -1, q.shape[1] // k.shape[1], -1)\n                    .reshape(original_shape[0], -1, original_shape[2])\n                )\n        if v.shape[1] != q.shape[1]:\n            # MQA expand\n            if v.shape[1] == 1:\n                v = v.expand(-1, q.shape[1], -1)\n            # Grouped attention reshape\n            else:\n                original_shape = v.shape\n                v = (\n                    v.unsqueeze(2)\n                    .expand(-1, -1, q.shape[1] // v.shape[1], -1)\n                    .reshape(original_shape[0], -1, original_shape[2])\n                )\n\n        return flash_attn_cuda.fwd(\n            q,\n            k,\n            v,\n            out,\n            cu_seqlens,\n            cu_seqlens,\n            max_s,\n            max_s,\n            0.0,\n            softmax_scale,\n            False,\n            True,\n            False,\n            0,\n            None,\n        )\n\n    raise NotImplementedError(\"flash attention is not installed\")\n", "path": "server/text_generation_server/utils/flash_attn.py"}], "after_files": [{"content": "import os\nimport torch\n\nfrom loguru import logger\n\nif os.getenv(\"USE_FLASH_ATTENTION\", \"\").lower() == \"false\":\n    raise ImportError(\"`USE_FLASH_ATTENTION` is false.\")\n\nif not torch.cuda.is_available():\n    raise ImportError(\"CUDA is not available\")\n\nmajor, minor = torch.cuda.get_device_capability()\nis_sm75 = major == 7 and minor == 5\nis_sm8x = major == 8 and minor >= 0\nis_sm90 = major == 9 and minor == 0\n\nHAS_FLASH_ATTN = False\nHAS_FLASH_ATTN_V2 = False\ntry:\n    try:\n        import flash_attn_2_cuda\n    except ImportError:\n        raise ImportError(\n            \"Flash Attention V2 is not installed.\\n\"\n            \"Use the official Docker image (ghcr.io/huggingface/text-generation-inference:latest) \"\n            \"or install flash attention v2 with `cd server && make install install-flash-attention-v2`\"\n        )\n    if not (is_sm8x or is_sm90):\n        raise ImportError(\n            f\"GPU with CUDA capability {major} {minor} is not supported for \"\n            \"Flash Attention V2\"\n        )\n    HAS_FLASH_ATTN_V2 = True\nexcept ImportError as e:\n    try:\n        import flash_attn_cuda\n    except ImportError:\n        raise ImportError(\n            \"Flash Attention is not installed.\\n\"\n            \"Use the official Docker image (ghcr.io/huggingface/text-generation-inference:latest) \"\n            \"or install flash attention with `cd server && make install install-flash-attention`\"\n        ) from e\n\n    if not (is_sm75 or is_sm8x or is_sm90):\n        raise ImportError(\n            f\"GPU with CUDA capability {major} {minor} is not supported\"\n        ) from e\n    logger.warning(f\"Unable to use Flash Attention V2: {e}\")\n    HAS_FLASH_ATTN = True\n\n\ndef attention(\n    q,\n    k,\n    v,\n    out,\n    cu_seqlens,\n    max_s,\n    softmax_scale,\n    window_size_left=-1,\n):\n    if HAS_FLASH_ATTN_V2:\n        return flash_attn_2_cuda.varlen_fwd(\n            q,\n            k,\n            v,\n            out,\n            cu_seqlens,\n            cu_seqlens,\n            max_s,\n            max_s,\n            0.0,\n            softmax_scale,\n            False,\n            True,\n            window_size_left,\n            0,\n            False,\n            None,\n        )\n\n    if HAS_FLASH_ATTN:\n        if window_size_left != -1:\n            raise NotImplementedError(\n                \"window_size_left is only available with flash attn v2\"\n            )\n\n        # Flash attention v1 requires q, k and v to have the same number of heads\n        if k.shape[1] != q.shape[1]:\n            # MQA expand\n            if k.shape[1] == 1:\n                k = k.expand(-1, q.shape[1], -1)\n            # Grouped attention reshape\n            else:\n                original_shape = k.shape\n                k = (\n                    k.unsqueeze(2)\n                    .expand(-1, -1, q.shape[1] // k.shape[1], -1)\n                    .reshape(original_shape[0], -1, original_shape[2])\n                )\n        if v.shape[1] != q.shape[1]:\n            # MQA expand\n            if v.shape[1] == 1:\n                v = v.expand(-1, q.shape[1], -1)\n            # Grouped attention reshape\n            else:\n                original_shape = v.shape\n                v = (\n                    v.unsqueeze(2)\n                    .expand(-1, -1, q.shape[1] // v.shape[1], -1)\n                    .reshape(original_shape[0], -1, original_shape[2])\n                )\n\n        return flash_attn_cuda.fwd(\n            q,\n            k,\n            v,\n            out,\n            cu_seqlens,\n            cu_seqlens,\n            max_s,\n            max_s,\n            0.0,\n            softmax_scale,\n            False,\n            True,\n            False,\n            0,\n            None,\n        )\n\n    raise NotImplementedError(\"flash attention is not installed\")\n", "path": "server/text_generation_server/utils/flash_attn.py"}]}