From 171775f306a333a9cf105bfd533bf3e113d401d9 Mon Sep 17 00:00:00 2001 From: Liwen <53441624+li-liwen@users.noreply.github.com> Date: Sat, 28 Mar 2026 08:27:11 +0000 Subject: [PATCH] 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 --- benchmarks/kernels/benchmark_moe.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 515406aa9..65bc38c6c 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -627,9 +627,8 @@ class BenchmarkWorker: need_device_guard = True with ( - torch.accelerator.device_index(self.device_id) - if need_device_guard - else nullcontext() + # Ray restricts each worker to one GPU; use local index 0 + torch.accelerator.device_index(0) if need_device_guard else nullcontext() ): for idx, config in enumerate(tqdm(search_space)): try: