[XPU][UT]Fix test_mnnvl_alltoall device and distributed backend (#57591)
Signed-off-by: yisustc <yi.sheng@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
YiSheng5 committed
Sep 20, 2026 at 10:53 UTC
1c3ef2ad4cabfbc2e5d74494b0b2bf8c3cb3f936
2 files changed
+11
-3
tests/distributed/test_mnnvl_alltoall.py
+6
-3
@@ -14,6 +14,7 @@ import torch
14
import torch.multiprocessing as mp
15
16
from vllm.distributed import get_ep_group
17
+from vllm.platforms import current_platform
18
from vllm.utils.flashinfer import (
19
has_flashinfer_nvlink_one_sided,
20
has_flashinfer_nvlink_two_sided,
@@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port
24
25
from ..utils import init_test_distributed_environment
26
27
+DEVICE = current_platform.device_type
28
+
29
# ---------------------------------------------------------------------------
30
# Helpers
31
# ---------------------------------------------------------------------------
@@ -483,7 +486,7 @@ def _args_dispatch_combine_worker(rank, world_size):
486
from vllm.forward_context import get_forward_context
487
488
cpu_group = get_ep_group().cpu_group
486
- device = torch.device(f"cuda:{rank}")
489
+ device = torch.device(f"{DEVICE}:{rank}")
490
491
hidden_size = 64
492
tokens_per_rank = 16
@@ -624,7 +627,7 @@ def _two_sided_data_worker(rank, world_size):
627
# Use DP group because MnnvlMoe workspace allocation calls get_dp_group()
628
# internally and requires dp_size == ep_size.
629
cpu_group = get_dp_group().cpu_group
627
- device = torch.device(f"cuda:{rank}")
630
+ device = torch.device(f"{DEVICE}:{rank}")
631
num_gpus = torch.accelerator.device_count()
632
633
hidden_size = 128
@@ -769,7 +772,7 @@ def _one_sided_data_worker(rank, world_size):
772
from vllm.forward_context import get_forward_context
773
774
cpu_group = get_dp_group().cpu_group
772
- device = torch.device(f"cuda:{rank}")
775
+ device = torch.device(f"{DEVICE}:{rank}")
776
777
hidden_size = 256
778
tokens_per_rank = 32
tests/utils.py
+5
@@ -1427,8 +1427,10 @@ def init_test_distributed_environment(
1427
get_current_vllm_config_or_none,
1428
set_current_vllm_config,
1429
)
1430
+ from vllm.platforms import current_platform
1431
1432
distributed_init_method = f"tcp://localhost:{distributed_init_port}"
1433
+ backend = current_platform.dist_backend
1434
1435
if data_parallel_size > 1:
1436
# For DP we need to set a common DP master port
@@ -1449,6 +1451,7 @@ def init_test_distributed_environment(
1451
rank=rank % tp_pp_world,
1452
distributed_init_method=distributed_init_method,
1453
local_rank=local_rank if local_rank >= 0 else rank,
1454
+ backend=backend,
1455
)
1456
ensure_model_parallel_initialized(tp_size, pp_size)
1457
return
@@ -1460,6 +1463,7 @@ def init_test_distributed_environment(
1463
rank=rank,
1464
distributed_init_method=distributed_init_method,
1465
local_rank=local_rank,
1466
+ backend=backend,
1467
)
1468
ensure_model_parallel_initialized(tp_size, pp_size)
1469
else:
@@ -1470,6 +1474,7 @@ def init_test_distributed_environment(
1474
rank=rank,
1475
distributed_init_method=distributed_init_method,
1476
local_rank=local_rank,
1477
+ backend=backend,
1478
)
1479
ensure_model_parallel_initialized(tp_size, pp_size)
1480