[Core][Optimization] remove vllm-nccl (#5091)

This commit is contained in:
youkaichao
2024-05-28 22:13:52 -07:00
committed by GitHub
parent 616e600e0b
commit 5bd3c65072
7 changed files with 21 additions and 100 deletions

View File

@@ -358,11 +358,8 @@ def get_requirements() -> List[str]:
cuda_major, cuda_minor = torch.version.cuda.split(".")
modified_requirements = []
for req in requirements:
if "vllm-nccl-cu12" in 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")):
if ("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