A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 08:30 UTC
[Misc]Replace `cuda` hard code with `current_platform` in Ray (#14668)
Signed-off-by: noemotiovon <757486878@qq.com>
Chenguang Li committed
May 25, 2025 at 11:26 UTC
cebc22f3b674e58fa004625c6afe5a6ec6f2d9de
1 file changed
+2
-4
vllm/executor/ray_utils.py
+2
-4
@@ -87,9 +87,8 @@ try:
87
# TODO(swang): This is needed right now because Ray Compiled Graph
88
# executes on a background thread, so we need to reset torch's
89
# current device.
90
- import torch
90
if not self.compiled_dag_cuda_device_set:
92
- torch.cuda.set_device(self.worker.device)
91
+ current_platform.set_device(self.worker.device)
92
self.compiled_dag_cuda_device_set = True
93
94
output = self.worker._execute_model_spmd(execute_model_req,
@@ -113,8 +112,7 @@ try:
112
# Not needed
113
pass
114
else:
116
- import torch
117
- torch.cuda.set_device(self.worker.device)
115
+ current_platform.set_device(self.worker.device)
116
117
self.compiled_dag_cuda_device_set = True
118