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 5, 21:00 UTC
Fix Device Index for ROCm Ray Workers in MoE Benchmark (#38108)
Signed-off-by: Liwen <53441624+li-liwen@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Liwen committed
Mar 28, 2026 at 08:27 UTC
171775f306a333a9cf105bfd533bf3e113d401d9
1 file changed
+2
-3
benchmarks/kernels/benchmark_moe.py
+2
-3
@@ -627,9 +627,8 @@ class BenchmarkWorker:
627
need_device_guard = True
628
629
with (
630
- torch.accelerator.device_index(self.device_id)
631
- if need_device_guard
632
- else nullcontext()
630
+ # Ray restricts each worker to one GPU; use local index 0
631
+ torch.accelerator.device_index(0) if need_device_guard else nullcontext()
632
):
633
for idx, config in enumerate(tqdm(search_space)):
634
try: