Revert "Separate attention backends (#3005)"

This reverts commit 2daf23a.

.gitignore

@@ -184,6 +184,3 @@ _build/
 
 # Benchmark dataset
 *.json
-
-# Third-party Python packages.
-vllm/thirdparty_files/

setup.py

@@ -3,7 +3,6 @@
 import os
 import re
 import subprocess
-import sys
 import warnings
 from pathlib import Path
 from typing import List, Set
@@ -15,8 +14,6 @@
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME, ROCM_HOME
 
 ROOT_DIR = os.path.dirname(__file__)
-# This is a temporary directory to store third-party packages.
-THIRDPARTY_SUBDIR = "vllm/thirdparty_files"
 
 # If you are developing the C++ backend of vLLM, consider building vLLM with
 # `python setup.py develop` since it will give you incremental builds.
@@ -327,46 +324,8 @@ def get_torch_arch_list() -> Set[str]:
                     "nvcc": NVCC_FLAGS_PUNICA,
                 },
             ))
-
-    # Download the FlashAttention package.
-    # Adapted from https://github.com/ray-project/ray/blob/f92928c9cfcbbf80c3a8534ca4911de1b44069c0/python/setup.py#L518-L530
-    flash_attn_version = "2.5.6"
-    install_dir = os.path.join(ROOT_DIR, THIRDPARTY_SUBDIR)
-    subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            "-q",
-            f"--target={install_dir}",
-            "einops",  # Dependency of flash-attn.
-            f"flash-attn=={flash_attn_version}",
-            "--no-dependencies",  # Required to avoid re-installing torch.
-        ],
-        env=dict(os.environ, CC="gcc"),
-    )
-
-    # Copy the FlashAttention package into the vLLM package after build.
-    class build_ext(BuildExtension):
-
-        def run(self):
-            super().run()
-            target_dir = os.path.join(self.build_lib, THIRDPARTY_SUBDIR)
-            if not os.path.exists(target_dir):
-                os.makedirs(target_dir)
-            self.copy_tree(install_dir, target_dir)
-
-    class BinaryDistribution(setuptools.Distribution):
-
-        def has_ext_modules(self):
-            return True
-
-else:
-    build_ext = BuildExtension
-    BinaryDistribution = setuptools.Distribution
-    if _is_neuron():
-        neuronxcc_version = get_neuronxcc_version()
+elif _is_neuron():
+    neuronxcc_version = get_neuronxcc_version()
 
 vllm_extension_sources = [
     "csrc/cache_kernels.cu",
@@ -509,7 +468,6 @@ def get_requirements() -> List[str]:
     python_requires=">=3.8",
     install_requires=get_requirements(),
     ext_modules=ext_modules,
-    cmdclass={"build_ext": build_ext} if not _is_neuron() else {},
-    distclass=BinaryDistribution,
+    cmdclass={"build_ext": BuildExtension} if not _is_neuron() else {},
     package_data=package_data,
 )

tests/kernels/test_prefix_prefill.py

@@ -3,7 +3,7 @@
 import time
 
 import torch
-from vllm.model_executor.layers.attention.ops.prefix_prefill import (
+from vllm.model_executor.layers.triton_kernel.prefix_prefill import (
     context_attention_fwd)
 from xformers import ops as xops
 from xformers.ops.fmha.attn_bias import BlockDiagonalCausalFromBottomRightMask

vllm/__init__.py

@@ -1,28 +1,12 @@
 """vLLM: a high-throughput and memory-efficient inference engine for LLMs"""
 
-
-# Adapted from https://github.com/ray-project/ray/blob/f92928c9cfcbbf80c3a8534ca4911de1b44069c0/python/ray/__init__.py#L11
-def _configure_system():
-    import os
-    import sys
-
-    # Importing flash-attn.
-    thirdparty_files = os.path.join(os.path.abspath(os.path.dirname(__file__)),
-                                    "thirdparty_files")
-    sys.path.insert(0, thirdparty_files)
-
-
-_configure_system()
-# Delete configuration function.
-del _configure_system
-
-from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs  # noqa: E402
-from vllm.engine.async_llm_engine import AsyncLLMEngine  # noqa: E402
-from vllm.engine.llm_engine import LLMEngine  # noqa: E402
-from vllm.engine.ray_utils import initialize_cluster  # noqa: E402
-from vllm.entrypoints.llm import LLM  # noqa: E402
-from vllm.outputs import CompletionOutput, RequestOutput  # noqa: E402
-from vllm.sampling_params import SamplingParams  # noqa: E402
+from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.engine.llm_engine import LLMEngine
+from vllm.engine.ray_utils import initialize_cluster
+from vllm.entrypoints.llm import LLM
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import SamplingParams
 
 __version__ = "0.3.3"