[Setup] Enable TORCH_CUDA_ARCH_LIST for selecting target GPUs

[WoosukKwon]

Fixes #1070

The TORCH_CUDA_ARCH_LIST env variable is a standard place to specify the target GPUs one wants to build a PyTorch project for. This PR enables using the variable in our setup.py. This will be especially useful for those who build vLLM images for specific GPUs.

[WoosukKwon]

@zhuohan123 This PR is also ready for review.

[v1nc3nt27]

When I build this branch on an A10, it just stops after the initial request (shortened):

INFO 09-19 16:02:13 async_llm_engine.py:328] Received request cmpl-833d8e3bd8b84aa9a7a4675abd62a19b: prompt: ... 29914, 25580, 29962].

If I use branch merge_quant, it works fine.

All I do to test this in this and merge_quant branch is:

git clone https://github.com/vllm-project/vllm.git
git checkout fix-setup
cd vllm
pip install -e .
pip install fschat[webui,model_worker]
python3 -m vllm.entrypoints.openai.api_server --model abhinavkulkarni/meta-llama-Llama-2-13b-chat-hf-w4-g128-awq --quantization awq --max-num-batched-tokens 4096 --host 0.0.0.0 --port 8080

Is there a way to increase logging and find out what is happening? In the case where it doesn't work, GPU memory stays at around 18GB, in the other case it goes up to 23GB, so I don't think much is happening. There is no error coming back from the server to the client sending the request, it just stays open.

UPDATE:

Launching the default api server via python -m vllm.entrypoints.api_server --model abhinavkulkarni/meta-llama-Llama-2-13b-chat-hf-w4-g128-awq --quantization awq --max-num-batched-tokens 4096 --host 0.0.0.0 --port 8080 and using the example curl in the docs also returns no result but the logs show that prompt token ids are None:

INFO 09-19 16:18:56 async_llm_engine.py:372] Received request 49d601e1afbe4075987bea7877cb06ef: prompt: 'San Francisco is a', sampling params: SamplingParams(n=4, best_of=4, presence_penalty=0.0, frequency_penalty=0.0, temperature=0, top_p=1.0, top_k=-1, use_beam_search=True, length_penalty=1.0, early_stopping=False, stop=[], ignore_eos=False, max_tokens=7, logprobs=None), prompt token ids: None.

[WoosukKwon]

@v1nc3nt27 What is the merge_quant branch? Our repo doesn't have a branch with the name.

[v1nc3nt27]

@WoosukKwon sorry, it was this branch #1032

[zhuohan123]

discard does not raise an error if the element is not present in the set. Similar for the remove below.

Suggested change

	compute_capabilities.remove("8.6")
	compute_capabilities.discard("8.6")

[zhuohan123]

Also we might need to remove *+PTX

[WoosukKwon]

L80 and below is executed when compute_capabilities is empty. In this case, we add all SUPPORTED_ARCHS and remove some of them based on the user's CUDA version. So, I think using remove is more appropriate than discard, and we don't need to remove *PTX because SUPPORTED_ARCHS does not have it.

[zhuohan123]

Actually, should we print a warning for this comment?

[WoosukKwon]

Sounds good! Added.

[WoosukKwon]

@zhuohan123 Addressed your comments. PTAL.

[WoosukKwon]

L80 and below is executed when compute_capabilities is empty. In this case, we add all SUPPORTED_ARCHS and remove some of them based on the user's CUDA version. So, I think using remove is more appropriate than discard, and we don't need to remove *PTX because SUPPORTED_ARCHS does not have it.

[WoosukKwon]

Sounds good! Added.

[zhuohan123]

LGTM! Thanks for the contribution!