Reinstall xformers on demand

pytorch-labs · Dec 6, 2024 · 5e601ff · 5e601ff
1 parent bffc713
commit 5e601ff
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/.github/workflows/_linux-test-h100.yml b/.github/workflows/_linux-test-h100.yml
@@ -28,6 +28,8 @@ jobs:
           filters: |
             fa:
               - 'submodules/flash-attention/**'
+            xformers:
+              - 'submodules/xformers/**'
       - name: Tune Nvidia GPU
         run: |
           sudo nvidia-smi -pm 1
@@ -39,6 +41,11 @@ jobs:
           . "${SETUP_SCRIPT}"
           python install.py --fa2
           python install.py --fa3
+      - name: Reinstall xformers (optional)
+        if: steps.submodules_changes.outputs.xformers == 'true'
+        run: |
+          . "${SETUP_SCRIPT}"
+          python install.py --xformers
       - name: Install Tritonbench
         run: |
           # speedup install and skip compile by reusing the docker .so files

diff --git a/submodules/xformers b/submodules/xformers
diff --git a/tritonbench/operators/decoding_attention/operator.py b/tritonbench/operators/decoding_attention/operator.py
@@ -27,7 +27,7 @@
 try:
     torch_lib_path = os.path.join(os.path.dirname(__file__), "lib")
     with add_ld_library_path(torch_lib_path):
-        from flash_attn_interface import flash_attn_func as flash_attn_v3
+        import flash_attn_interface as flash_attn_v3
 except (ImportError, IOError, AttributeError):
     try:
         from ai_codesign.gen_ai.flash_attention_v2.hopper import (
+2 −2		.github/actions/setup-build-cuda/action.yml
+0 −0		.github/compute_wheel_version.py
+0 −154		.github/workflows/conda.yml
+5 −0		.github/workflows/gpu_test_gh.yml
+2 −2		.github/workflows/rocm_build.yml
+4 −4		.github/workflows/wheels.yml
+1 −1		.github/workflows/wheels_build.yml
+1 −1		.github/workflows/wheels_upload_pip.yml
+8 −1		.github/workflows/win-build.yml
+7 −0		CHANGELOG.md
+2 −14		README.md
+1 −1		docs/source/xformers_mingpt.ipynb
+0 −184		examples/cifar_MetaFormer.py
+0 −242		examples/cifar_ViT.py
+0 −182		packaging/build_conda.py
+0 −44		packaging/xformers/meta.yaml
+0 −3		requirements-test.txt
+0 −90		tests/test_feedforward.py
+0 −19		tests/test_hydra_helper.py
+0 −59		tests/test_pickling.py
+0 −224		tests/test_reversible.py
+20 −12		xformers/csrc/swiglu/cuda/dual_gemm_silu_identity_mul.cu
+108 −0		xformers/csrc/swiglu/cuda/epilogue_lhs_activation_and_mul.h
+6 −12		xformers/csrc/swiglu/swiglu_packedw.cpp
+0 −1		xformers/factory/README.md
+0 −11		xformers/factory/__init__.py
+0 −237		xformers/factory/block_configs.py
+0 −358		xformers/factory/block_factory.py
+0 −36		xformers/factory/hydra_helper.py
+0 −313		xformers/factory/model_factory.py
+0 −293		xformers/factory/weight_init.py
+0 −7		xformers/helpers/__init__.py
+0 −124		xformers/helpers/hierarchical_configs.py
+0 −32		xformers/helpers/test_utils.py
+0 −55		xformers/helpers/timm_sparse_attention.py
+1 −0		xformers/ops/__init__.py
+1 −1		xformers/ops/_triton/rope_padded_kernels.py
+3 −3		xformers/ops/fmha/__init__.py
+1 −1		xformers/ops/fmha/_triton/splitk_kernels.py