[Misc] Use vllm-flash-attn instead of flash-attn (#4686)

This commit is contained in:
Woosuk Kwon
2024-05-08 13:15:34 -07:00
committed by GitHub
parent 230c4b38c1
commit 89579a201f
6 changed files with 16 additions and 31 deletions

View File

@@ -355,14 +355,18 @@ def get_requirements() -> List[str]:
if _is_cuda():
requirements = _read_requirements("requirements-cuda.txt")
cuda_major = torch.version.cuda.split(".")[0]
cuda_major, cuda_minor = torch.version.cuda.split(".")
modified_requirements = []
for req in requirements:
if "vllm-nccl-cu12" in req:
modified_requirements.append(
req.replace("vllm-nccl-cu12", f"vllm-nccl-cu{cuda_major}"))
else:
modified_requirements.append(req)
req = req.replace("vllm-nccl-cu12",
f"vllm-nccl-cu{cuda_major}")
elif ("vllm-flash-attn" in req
and not (cuda_major == "12" and cuda_minor == "1")):
# vllm-flash-attn is built only for CUDA 12.1.
# Skip for other versions.
continue
modified_requirements.append(req)
requirements = modified_requirements
elif _is_hip():
requirements = _read_requirements("requirements-rocm.txt")