code: def get_flash_attention_extensions(
      cuda_version: int, extra_compile_args
  ):
      # XXX: Not supported on windows for cuda<12
      # https://github.com/Dao-AILab/flash-attention/issues/345
      if platform.system() != "Linux" and cuda_version < 1200:
          return []
      # Figure out default archs to target
      DEFAULT_ARCHS_LIST = ""
      if cuda_version >= 1108:
          DEFAULT_ARCHS_LIST = "8.0;8.6;9.0"
      elif cuda_version > 1100:
          DEFAULT_ARCHS_LIST = "8.0;8.6"
      elif cuda_version == 1100:
          DEFAULT_ARCHS_LIST = "8.0"
      else:
          return []
  
      if os.getenv("XFORMERS_DISABLE_FLASH_ATTN", "0") != "0":
          return []
  
      # Supports `9.0`, `9.0+PTX`, `9.0a+PTX` etc...
      PARSE_CUDA_ARCH_RE = re.compile(
          r"(?P<major>[0-9]+)\.(?P<minor>[0-9])(?P<suffix>[a-zA-Z]{0,1})(?P<ptx>\+PTX){0,1}"
      )
      archs_list = os.environ.get(
          "TORCH_CUDA_ARCH_LIST", DEFAULT_ARCHS_LIST
      )
      nvcc_archs_flags = []
      for arch in archs_list.replace(" ", ";").split(";"):
          match = PARSE_CUDA_ARCH_RE.match(arch)
          assert match is not None, f"Invalid sm version: {arch}"
          num = 10 * int(match.group("major")) + int(
              match.group("minor")
          )
          # Need at least Sm80
          if num < 80:
              continue
          # Sm90 requires nvcc 11.8+
          if num >= 90 and cuda_version < 1108:
              continue
          suffix = match.group("suffix")
          nvcc_archs_flags.append(
              f"-gencode=arch=compute_{num}{suffix},code=sm_{num}{suffix}"
          )
          if match.group("ptx") is not None:
              nvcc_archs_flags.append(
                  f"-gencode=arch=compute_{num}{suffix},code=compute_{num}{suffix}"
              )
      if not nvcc_archs_flags:
          return []
  
      nvcc_windows_flags = []
      if platform.system() == "Windows":
          nvcc_windows_flags = ["-Xcompiler", "/permissive-"]
  
      flash_root = os.path.join(
          this_dir, "third_party", "flash-attention"
      )
      cutlass_inc = os.path.join(
          flash_root, "csrc", "cutlass", "include"
      )
      if not os.path.exists(flash_root) or not os.path.exists(
          cutlass_inc
      ):
          raise RuntimeError(
              "flashattention submodule not found. Did you forget "
              "to run `git submodule update --init --recursive` ?"
          )
  
      sources = ["csrc/flash_attn/flash_api.cpp"]
      for f in glob.glob(
          os.path.join(flash_root, "csrc", "flash_attn", "src", "*.cu")
      ):
          sources.append(str(Path(f).relative_to(flash_root)))
      return [
          CUDAExtension(
              name="xformers._C_flashattention",
              sources=[
                  os.path.join(flash_root, path) for path in sources
              ],
              extra_compile_args={
                  **extra_compile_args,
                  "nvcc": extra_compile_args.get("nvcc", [])
                  + [
                      "-O3",
                      "-std=c++17",
                      "-U__CUDA_NO_HALF_OPERATORS__",
                      "-U__CUDA_NO_HALF_CONVERSIONS__",
                      "-U__CUDA_NO_HALF2_OPERATORS__",
                      "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                      "--expt-relaxed-constexpr",
                      "--expt-extended-lambda",
                      "--use_fast_math",
                      "--ptxas-options=-v",
                  ]
                  + nvcc_archs_flags
                  + nvcc_windows_flags
                  + get_extra_nvcc_flags_for_build_type(),
              },
              include_dirs=[
                  p.absolute()
                  for p in [
                      Path(flash_root) / "csrc" / "flash_attn",
                      Path(flash_root) / "csrc" / "flash_attn" / "src",
                      Path(flash_root) / "csrc" / "cutlass" / "include",
                  ]
              ],
          )
      ]
  
imports: import datetime
  import distutils.command.clean
  import glob
  import json
  import os
  import platform
  import re
  import shlex
  import shutil
  import subprocess
  import sys
  from pathlib import Path
  from typing import List
  import setuptools
  import torch
  from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CppExtension, CUDAExtension
label_desc: Use image node_id 8 for calling a global function with example usage: get_flash_attention_extensions(cuda_version, extra_compile_args) and returns: list, list, list, list, list