[RFC]: Continue on Device agnostic API abstraction to current_platform.XXX

[xuechendi]

co-author with @jikunshang

Motivation.

This RFC is aiming to reuse GPUWorker and GPUModelRunner for any GPGPU devices, such as CUDA, ROCM and Intel GPU(aka: XPU).

• By doing so, we can remove redundant duplication by adding a new XXXWorker/XXXModelRunner and derive from GPUWorker/GPUModelRunner
• Any feature implemented in GPUWorker/GPUModelRunner such as logitsProcessor, samplingOutput optimization, spec_decode can be shared to all GPGPU hardware.

Status & Challenge

• Previous RFC from Huawei has made significant work done through - [RFC]: Make device agnostic for diverse hardware support #9268

• Currently, GPUWorker and GPUModelRunner is assumed that it will only be used by CUDA and RocM, so hard-coded to cuda API will be used in above two files. Ex:torch.cuda.XXX or tensor.to('cuda')

Proposed Change.

1. Add abstract API into platforms/interface.py and implement in cuda.py, rocm.py, xpu.py.
2. update any tensor.to('cuda') or tensor.cuda() to use tensor.to(current_platform.device).
3. Add a skip check in case of API mismatch.
4. Add static check to PR pre_commit to indicate future contributor to use current_platform instead calling torch.cuda directly.

Plan

1. Add abstract API

• torch.cuda.empty_cache -> current_platform.empty_cache
• torch.cuda.set_device -> current_platform.set_device
• torch.cuda.reset_peak_memory_stats -> current_platform.reset_peak_memory_stats
• torch.cuda.mem_get_info -> current_platform.mem_get_info
• torch.cuda.memory_stats -> current_platform.memory_stats
• torch.cuda.memory_reserved -> current_platform.memory_reserved
• torch.cuda.synchronize -> current_platform.synchronize

2. abstract device / dist_device

• tensor.to('cuda') -> tensor.to(current_platform.device)
• use current_platform.dist_backend for init_worker_distributed_environment. Done in [Refactor]Abstract Platform Interface for Distributed Backend and Add xccl Support for Intel XPU #19410

3. add skip check

• add current_platform.is_graph_mode_supported()
• TBD

4. Add static check

• TBD

** PR list **
#19410 => Merged
#20751 => Open

** additional **

torch APIs support list in torch.cuda and torch.xpu - By Q2 25

no.	api	cuda	xpu
1	current_device	✅	✅
2	current_stream	✅	✅
3	device	✅	✅
4	device_count	✅	✅
5	get_device_capability	✅	✅
6	get_device_name	✅	✅
7	get_device_properties	✅	✅
8	init	✅	✅
9	is_available	✅	✅
10	is_initialized	✅	✅
11	set_device	✅	✅
12	set_stream	✅	✅
13	stream	✅	✅
14	synchronize	✅	✅
15	manual_seed	✅	✅
16	manual_seed_all	✅	✅
17	Stream	✅	✅
18	Event	✅	✅
19	empty_cache	✅	✅
20	mem_get_info	✅	✅
21	memory_stats, memory_stats_as_nested_dict	✅	✅
22	memory_allocated	✅	✅
23	max_memory_allocated	✅	✅
24	memory_reserved	✅	✅
25	reset_peak_memory_stats	✅	✅

Not matched APIs

no.	api	cuda	xpu
1	Cudart	✅	enable with alternative API
2	is_current_stream_capturing	✅	enable with alternative API
3	graph_pool_handle	✅	check and skip
4	CUDAGraph	✅	check and skip
5	Graph	✅	check and skip
6	CUDAPluggableAllocator	✅	check and skip
7	nvtx.range_push	✅	check and skip
8	nvtx.range_pop	✅	check and skip
9	nvtx.range	✅	check and skip
10	_lazy_init	✅	enable with alternative API
11	_is_compiled	✅	enable with alternative API
12	_device_count_amdsmi	✅	enable with alternative API
13	_device_count_nvml	✅	enable with alternative API
14	tunnable	✅	enable with alternative API

Feedback Period.

No response

CC List.

@jikunshang @simon-mo @WoosukKwon @youkaichao @gshtras

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[dvrogozh]

Add abstract API

This API resembles the effort on pytorch side to create device abstraction API under torch.accelerate. See https://docs.pytorch.org/docs/stable/accelerator.html#module-torch.accelerator. There is an ongoing associated effort in pytorch examples to demonstrate an API usage, see for example:

• Use torch.accelerator API in Fast Neural Style example pytorch/examples#1327
• Use torch.accelerator API in mnist examples pytorch/examples#1334
• and other similar PRs in the history of the named repo

As an alternative proposal, instead of creating similar API on vLLM side, maybe it's better to contribute to the missing parts of torch.accelerate and its stability?

[jikunshang]

an early draft PR is here #20721

[dvrogozh]

the suggestion you made is to replace torch.cuda to torch.accelerator.XXX

Yes, that is correct

have data to prove the performance is on-par with torch.accelerator vs torch.cuda or torch.accelerator vs torch.xpu

torch.accelerator is a tiny wrapper similar to torch.cuda and torch.xpu. It does not contain any fundamentally different implementation. Thus you should expect on par performance out of the box. Of course some evaluation needs to be done, but I believe it might be quite limited

we have need to convince CUDA and ROCM people to agree of using torch.accelerator

I hope to hear opinion from vLLM maintainers also on this.

[dvrogozh]

For now, the proposal is to simply reuse GPUWorker/GPUModelRunner for any GPGPU.

Actually you are making a bigger proposal to implement an abstract API in the vLLM which is essentially the same as torch.accelerator API. It's better to align with the vLLM community on the approach as it will be additional effort to first introduce abstraction here in vLLM and then replace this abstraction with the torch.accelerator. It might be more efficient to go with torch.accelerator right away.

To consider the simpler solution, then simple if-else logic right in the GPUWorker/GPUModelRunner might be enough. For example, as I did propose to implement earlier today before this RFC was opened:

• [XPU] dispatch xpu/cuda specific calls in the model runner #20698

[jikunshang]

submit first one to refactor basic calls here #20751

[xuechendi]

@gshtras , we want to abstract torch.cuda API to interface.py; May you take a look of the RFC, thanks so much
Here is also the first PR sent by kunshang #20751