Does vllm support pytorch/xla ?

[dinghaodhd]

Anything you want to discuss about vllm.

Hi, We have adapt our harware to pytorch/xla, what should we do to run vllm via pytorch/xla?

[rkooo567]

I don't think it is working with pytorch/XLA now.

1. The core attention algorithm that supports paged-attention is written in Cuda (and there are some other custom cuda kernels).
2. https://github.com/vllm-project/vllm/blob/429284dc374bab79d4dfbb25053583901e6e5051/vllm/config.py#L520 there are only neuron and cuda device options.

I think to make this work;

• we should pass correct device settings. I don't know if xla requires additional changes other than providing a correct device, but if so, you need to change model_runner.py (preparing tensors) and general model path to make sure to use XLA config.
• The core attention should use the kernel that works with paged attention.

[richardliaw]

There will be some exploration done from Google side I believe (cc @allenwang28)

[allenwang28]

We are exploring TPU compatibility and support in https://github.com/vllm-project/vllm/issues/3620 through PyTorch/XLA.

but after this work, it should be clearer how to adapt PyTorch/XLA in vLLM for different hardware backends CC @WoosukKwon @miladm @shauheen

[yiakwy-xpu-ml-framework-team]

Google has made its efforts in supportin TPU. We will make efforts to utilize the canonical IR optimizattion ,auto tuning, auto fusion, layout optimization form xla and stableHLO in GPU devices.