Repository: facebookresearch/xformers

Issue:
Pre-build wheels don't have xformers._flash_attn
# 🐛 Bug

When I install the wheel from https://pypi.org/project/xformers/0.0.16rc392/ the module xformers._flash_attn does not exist.

This mean importing xformers.ops fails.

```File /usr/local/lib/python3.8/dist-packages/xformers/ops/__init__.py:8
      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.
      6 import torch
----> 8 from .fmha import (
      9     AttentionMask,
     10     AttentionOp,
     11     AttentionOpBase,
     12     AttentionOpDispatch,
     13     LowerTriangularMask,
     14     MemoryEfficientAttentionCutlassFwdFlashBwOp,
     15     MemoryEfficientAttentionCutlassOp,
     16     MemoryEfficientAttentionFlashAttentionOp,
     17     MemoryEfficientAttentionOp,
     18     MemoryEfficientAttentionTritonFwdFlashBwOp,
     19     TritonFlashAttentionOp,
     20     memory_efficient_attention,
     21     memory_efficient_attention_backward,
     22     memory_efficient_attention_forward,
     23     memory_efficient_attention_forward_requires_grad,
     24 )
     25 from .swiglu_op import (
     26     SwiGLU,
     27     SwiGLUEagerOp,
   (...)
     32     swiglu,
     33 )
     34 from .unbind import get_stack_strides, stack_or_none, unbind

File /usr/local/lib/python3.8/dist-packages/xformers/ops/fmha/__init__.py:10
      6 from typing import Any, Optional, Tuple, Type, Union
      8 import torch
---> 10 from . import cutlass, flash, small_k, triton
     11 from .common import (
     12     AttentionBwOpBase,
     13     AttentionFwOpBase,
   (...)
     22     bmk2bmhk,
     23 )
     24 from .dispatch import _dispatch_bw, _dispatch_fw

File /usr/local/lib/python3.8/dist-packages/xformers/ops/fmha/triton.py:16
     13 from ..common import register_operator
     15 if TYPE_CHECKING or _is_triton_available():
---> 16     from ..._flash_attn.flash_attn_triton import (
     17         _flash_attn_backward,
     18         _flash_attn_forward,
     19     )
     21     triton_flash_backward = _flash_attn_backward
     22     triton_flash_forward = _flash_attn_forward

ModuleNotFoundError: No module named 'xformers._flash_attn'
```
## Expected behavior

Wheels should be compiled with Cuda support

## Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

You can run the script with:

```bash
# For security purposes, please check the contents of collect_env.py before running it.
python -m torch.utils.collect_env
```

```
$ python -m torch.utils.collect_env 
Collecting environment information...
PyTorch version: 1.13.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 8.4.0-3ubuntu2) 8.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 14 2022, 12:59:47)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-1022-aws-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A10G
Nvidia driver version: 510.73.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] ema-pytorch==0.0.10
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.4
[pip3] torch==1.13.0
[pip3] torchvision==0.13.1+cu116
[conda] No relevant packages
```



Discussion:
Do you have an idea @AbdBarho ?

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -62,7 +62,7 @@ def write_version_file(version: str):
             f.write(f'git_tag = "{tag}"\n')
 
 
-def symlink_package(name: str, path: Path) -> None:
+def symlink_package(name: str, path: Path, is_building_wheel: bool) -> None:
     cwd = Path(__file__).parent
     path_from = cwd / path
     path_to = os.path.join(cwd, *name.split("."))
@@ -78,7 +78,9 @@ def symlink_package(name: str, path: Path) -> None:
         pass
     # OSError: [WinError 1314] A required privilege is not held by the client
     # Windows requires special permission to symlink. Fallback to copy
-    use_symlink = os.name != "nt"
+    # When building wheels for linux 3.7 and 3.8, symlinks are not included
+    # So we force a copy, see #611
+    use_symlink = os.name != "nt" and not is_building_wheel
     if use_symlink:
         os.symlink(src=path_from, dst=path_to)
     else:
@@ -296,13 +298,17 @@ def run(self):
                 version = f.readline().strip()
             version += get_local_version_suffix()
         write_version_file(version)
+
+    is_building_wheel = "bdist_wheel" in sys.argv
     # Embed a fixed version of flash_attn
     # NOTE: The correct way to do this would be to use the `package_dir`
     # parameter in `setuptools.setup`, but this does not work when
     # developing in editable mode
     # See: https://github.com/pypa/pip/issues/3160 (closed, but not fixed)
     symlink_package(
-        "xformers._flash_attn", Path("third_party") / "flash-attention" / "flash_attn"
+        "xformers._flash_attn",
+        Path("third_party") / "flash-attention" / "flash_attn",
+        is_building_wheel,
     )
     setuptools.setup(
         name="xformers",


Test patch: