@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e6b1a5e5e3

[XPU] Enable XPU eplb (#44987)

Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Signed-off-by: Qiming Zhang <qiming1.zhang@intel.com>

Qiming Zhang committed Sep 16, 2026 at 22:19 UTC e6b1a5e5e3fce6b777c5f63398f87c7c911a40b3

16 files changed +267 -109

tests/conftest.py

+1 -1

@@ -935,7 +935,7 @@ def hf_runner():

935

936

937 def _default_block_size() -> int:

938 - if torch.xpu.is_available():

938 + if current_platform.is_xpu():

939 return 64

940 if current_platform.is_cpu():

941 return 128

tests/distributed/eplb_utils.py

+6 -2

@@ -13,6 +13,7 @@ from vllm.config import VllmConfig, set_current_vllm_config

13 from vllm.distributed.parallel_state import (

14 init_distributed_environment,

15 )

16 +from vllm.platforms import current_platform

17 from vllm.utils.system_utils import update_environment_variables

19 mp.set_start_method("spawn", force=True)

@@ -73,13 +74,16 @@ def distributed_run(fn, world_size, *args):

74 def set_env_vars_and_device(env: dict[str, str]) -> None:

75 update_environment_variables(env)

76 local_rank = os.environ["LOCAL_RANK"]

76 - device = torch.device(f"cuda:{local_rank}")

77 + accelerator_type = torch.accelerator.current_accelerator().type

78 + device = torch.device(f"{accelerator_type}:{local_rank}")

79 torch.accelerator.set_device_index(device)

81 + backend = "xccl" if current_platform.is_xpu() else "nccl"

82 +

83 # Create a minimal vllm config for init_distributed_environment

84 vllm_config = VllmConfig()

85 with set_current_vllm_config(vllm_config):

82 - init_distributed_environment()

86 + init_distributed_environment(backend=backend)

87 atexit.register(_destroy_process_group_if_initialized)

88 # Ensure each worker process has the same random seed

89 random.seed(42)

tests/distributed/test_elastic_ep.py

@@ -10,6 +10,8 @@ from concurrent.futures import ThreadPoolExecutor

10 import pytest

11 import requests

13 +from vllm.platforms import current_platform

14 +

15 from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k

16 from ..utils import RemoteOpenAIServer, multi_gpu_test

@@ -233,6 +235,8 @@ def test_elastic_ep_scaling(

235

236 if not has_nixl():

237 pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")

238 + if current_platform.is_xpu():

239 + pytest.skip("Elastic EP with NIXL is not supported on XPU")

240 if all2all_backend == "nixl_ep" and not has_nixl_ep():

241 pytest.skip("NIXL EP is not installed")

242

tests/distributed/test_eplb_algo.py

+4 -1

@@ -10,6 +10,7 @@ from vllm.distributed.eplb.eplb_state import (

10 compute_logical_maps,

11 )

12 from vllm.distributed.eplb.policy.default import DefaultEplbPolicy

13 +from vllm.platforms import current_platform

16 def test_eplb_load_stats_reduce_across_ranks():

@@ -263,11 +264,13 @@ def test_global_load_balance_fallback():

264 assert torch.sum(logcnt) == num_replicas

265

266

266 -@pytest.mark.parametrize("device", ["cpu", "cuda"])

267 +@pytest.mark.parametrize("device", ["cpu", "cuda", "xpu"])

268 def test_device_compatibility(device):

269 """Test device compatibility."""

270 if device == "cuda" and not torch.cuda.is_available():

271 pytest.skip("CUDA not available")

272 + if device == "xpu" and not current_platform.is_xpu():

273 + pytest.skip("XPU not available")

274

275 weight = torch.tensor([[10, 20, 30, 40]], device=device)

276 num_replicas = 6

tests/distributed/test_eplb_events.py

+18 -11

@@ -4,15 +4,17 @@

4 import threading

5 import time

7 +import pytest

8 import torch

9 -from vllm.distributed.eplb.eplb_utils import CpuGpuEvent

10 +from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream

11 +from vllm.platforms import current_platform

14 def test_wait_blocks_until_record():

15 event = CpuGpuEvent()

14 - record_stream = torch.cuda.Stream()

15 - wait_stream = torch.cuda.Stream()

16 + record_stream = torch.Stream()

17 + wait_stream = torch.Stream()

18 wait_returned = threading.Event()

20 def waiter():

@@ -33,8 +35,8 @@ def test_wait_blocks_until_record():

36 def test_reuse_across_multiple_cycles():

37 wrapper = CpuGpuEvent()

36 - record_stream = torch.cuda.Stream()

37 - wait_stream = torch.cuda.Stream()

38 + record_stream = torch.Stream()

39 + wait_stream = torch.Stream()

40 NUM_CYCLES = 8

41 completed_cycles = []

42 barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)]

@@ -56,13 +58,18 @@ def test_reuse_across_multiple_cycles():

58 assert len(completed_cycles) == NUM_CYCLES

59 -def test_producer_consumer():

60 - """This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared

61 - GPU tensor on multiple CPU threads.

61 +@pytest.mark.parametrize("device", ["cuda", "xpu"])

62 +def test_producer_consumer(device):

63 + """This test uses the CpuGpuEvent to synchronize reads and writes to/from

64 + a shared GPU tensor on multiple CPU threads.

65 """

63 - worker_stream = torch.cuda.Stream()

66 + if device == "cuda" and not torch.cuda.is_available():

67 + pytest.skip("CUDA not available")

68 + if device == "xpu" and not current_platform.is_xpu():

69 + pytest.skip("XPU not available")

70 + worker_stream = torch.Stream()

71 # Create a single element counter that will be shared between two threads

65 - buf = torch.zeros(1, device="cuda")

72 + buf = torch.zeros(1, device=device)

73 NUM_ROUNDS = 5

75 ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)]

@@ -76,7 +83,7 @@ def test_producer_consumer():

83 if i > 0:

84 events[i - 1].wait(stream=worker_stream)

79 - with torch.cuda.stream(worker_stream):

86 + with device_stream(worker_stream):

87 buf.fill_(float(i + 1))

89 worker_stream.synchronize()

tests/distributed/test_eplb_execute.py

+78 -29

@@ -23,6 +23,7 @@ from vllm.distributed.parallel_state import (

23 ensure_model_parallel_initialized,

24 get_tp_group,

25 )

26 +from vllm.platforms import current_platform

28 from .eplb_utils import distributed_run, set_env_vars_and_device

@@ -208,10 +209,14 @@ def verify_redundant_experts_have_same_weights(

209 ep_rank: int,

210 world_size: int,

211 num_local_experts: int,

212 + cpu_group,

213 ) -> bool:

214 """Verify that all replicas of the same logical expert have the same weights."""

215 num_layers = len(expert_weights)

216 total_physical_experts = world_size * num_local_experts

217 + # XPU: gloo does not support XPU tensors, so gather on CPU then move back.

218 + # CUDA: gloo supports CUDA tensors directly, keep original device.

219 + use_cpu_gather = current_platform.is_xpu()

220

221 ok = True

222 for layer in range(num_layers):

@@ -219,31 +224,35 @@ def verify_redundant_experts_have_same_weights(

224 all_weights: list[torch.Tensor] = []

225

226 for weight_idx, hidden_size in enumerate(hidden_sizes):

227 + orig_device = expert_weights[layer][weight_idx].device

228 +

229 + if use_cpu_gather:

230 + local_weights = expert_weights[layer][weight_idx].cpu()

231 + gather_device = "cpu"

232 + else:

233 + local_weights = expert_weights[layer][weight_idx]

234 + gather_device = orig_device

235 +

236 # Create tensor to store all expert weights

237 # Shape: [total_physical_experts, hidden_size]

238 gathered_weights = torch.zeros(

239 total_physical_experts,

240 hidden_size,

227 - device=expert_weights[layer][weight_idx].device,

228 - dtype=expert_weights[layer][weight_idx].dtype,

241 + device=gather_device,

242 + dtype=local_weights.dtype,

243 )

244

231 - # Use all_gather to collect expert weights from current node

232 - # expert_weights[layer][weight_idx] shape:

233 - # [num_local_experts, hidden_size]

234 - local_weights = expert_weights[layer][

235 - weight_idx

236 - ] # [num_local_experts, hidden_size]

237 -

245 # Split tensor along dim 0 into a list for all_gather

246 gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0)

247

248 torch.distributed.all_gather(

242 - # Output list: each element corresponds to one rank's weights

249 list(gathered_weights_list),

244 - local_weights, # Input: current rank's local weights

250 + local_weights,

251 + group=cpu_group,

252 )

253

254 + if use_cpu_gather:

255 + gathered_weights = gathered_weights.to(orig_device)

256 all_weights.append(gathered_weights)

257

258 # Verify that all replicas of the same logical expert have the same

@@ -294,9 +303,23 @@ def verify_redundant_experts_have_same_weights(

303 return ok

304

305

297 -def assert_verification_synced(local_ok: bool, msg: str) -> None:

298 - ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32)

299 - torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)

306 +def assert_verification_synced(local_ok: bool, msg: str, cpu_group) -> None:

307 + if current_platform.is_xpu():

308 + # XPU: default group backend (gloo) does not support XPU tensors;

309 + # use CPU tensor with cpu_group instead.

310 + ok_tensor = torch.tensor(

311 + [1 if local_ok else 0], device="cpu", dtype=torch.int32

312 + )

313 + torch.distributed.all_reduce(

314 + ok_tensor, op=torch.distributed.ReduceOp.MIN, group=cpu_group

315 + )

316 + else:

317 + # CUDA: nccl (default group) supports device tensors directly.

318 + device = torch.accelerator.current_accelerator()

319 + ok_tensor = torch.tensor(

320 + [1 if local_ok else 0], device=device, dtype=torch.int32

321 + )

322 + torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)

323 assert bool(ok_tensor.item()), msg

324

325

@@ -336,8 +359,10 @@ def _test_async_transfer_layer_without_mtp_worker(

359

360 ep_group_coordinator = get_tp_group()

361 ep_group = ep_group_coordinator.device_group

362 + cpu_group = ep_group_coordinator.cpu_group

363 ep_rank = torch.distributed.get_rank()

340 - device = torch.device(f"cuda:{ep_rank}")

364 + accelerator_type = torch.accelerator.current_accelerator().type

365 + device = torch.device(f"{accelerator_type}:{ep_rank}")

366

367 total_physical_experts = world_size * num_local_experts

368 hidden_sizes = [16, 32]

@@ -376,7 +401,7 @@ def _test_async_transfer_layer_without_mtp_worker(

401 new_indices_cpu = new_indices.cpu()

402

403 expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

379 - cuda_stream = torch.cuda.Stream(device=device)

404 + stream = torch.Stream(device=device)

405

406 communicator = create_eplb_communicator_or_raise(

407 group_coordinator=ep_group_coordinator,

@@ -384,7 +409,7 @@ def _test_async_transfer_layer_without_mtp_worker(

409 expert_weights=expert_weights,

410 expert_buffer=expert_buffer,

411 )

387 - communicator.set_stream(cuda_stream)

412 + communicator.set_stream(stream)

413

414 for layer_idx in range(num_layers):

415 transfer_metadata = transfer_layer(

@@ -394,10 +419,10 @@ def _test_async_transfer_layer_without_mtp_worker(

419 expert_weights_buffer=expert_buffer,

420 ep_group=ep_group,

421 communicator=communicator,

397 - cuda_stream=cuda_stream,

422 + stream=stream,

423 layer_idx=layer_idx,

424 )

400 - cuda_stream.synchronize()

425 + stream.synchronize()

426 move_from_buffer(

427 expert_weights=expert_weights[layer_idx],

428 expert_weights_buffers=expert_buffer,

@@ -421,6 +446,7 @@ def _test_async_transfer_layer_without_mtp_worker(

446 ep_rank,

447 world_size,

448 num_local_experts,

449 + cpu_group=cpu_group,

450 )

451 and local_ok

452 )

@@ -428,6 +454,7 @@ def _test_async_transfer_layer_without_mtp_worker(

454 local_ok,

455 "Async transfer verification failed on at least one rank. "

456 "See logs for details.",

457 + cpu_group=cpu_group,

458 )

459

460

@@ -452,9 +479,11 @@ def _test_rearrange_expert_weights_with_redundancy(

479 )

480

481 ep_group_coordinator = get_tp_group()

455 - ep_group = ep_group_coordinator.cpu_group

482 + ep_group = ep_group_coordinator.device_group

483 + cpu_group = ep_group_coordinator.cpu_group

484 ep_rank = torch.distributed.get_rank()

457 - device = torch.device(f"cuda:{ep_rank}")

485 + accelerator_type = torch.accelerator.current_accelerator().type

486 + device = torch.device(f"{accelerator_type}:{ep_rank}")

487

488 # Test parameters

489 total_physical_experts = world_size * num_local_experts

@@ -523,12 +552,14 @@ def _test_rearrange_expert_weights_with_redundancy(

552 ep_rank,

553 world_size,

554 num_local_experts,

555 + cpu_group=cpu_group,

556 )

557 and local_ok

558 )

559 assert_verification_synced(

560 local_ok,

561 "Rearrange verification failed on at least one rank. See logs for details.",

562 + cpu_group=cpu_group,

563 )

564

565

@@ -556,7 +587,7 @@ def _test_rearrange_expert_weights_with_redundancy(

587 ],

588 )

589 @pytest.mark.parametrize(

559 - "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]

590 + "eplb_communicator", ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl"]

591 )

592 def test_rearrange_expert_weights_with_redundancy(

593 world_size,

@@ -568,6 +599,12 @@ def test_rearrange_expert_weights_with_redundancy(

599 """Test the functionality of rearranging expert weights with redundancy."""

600 if eplb_communicator == "nixl" and not has_nixl():

601 pytest.skip("NIXL is not available")

602 + if eplb_communicator == "nixl" and current_platform.is_xpu():

603 + pytest.skip("NIXL does not support XPU")

604 + if eplb_communicator in ("torch_nccl", "pynccl") and not torch.cuda.is_available():

605 + pytest.skip(f"{eplb_communicator} requires CUDA")

606 + if eplb_communicator == "torch_xccl" and not current_platform.is_xpu():

607 + pytest.skip("torch_xccl requires XPU")

608 if torch.accelerator.device_count() < world_size:

609 pytest.skip(f"Need at least {world_size} GPUs to run the test")

610 distributed_run(

@@ -592,9 +629,11 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

629 )

630

631 ep_group_coordinator = get_tp_group()

595 - ep_group = ep_group_coordinator.cpu_group

632 + ep_group = ep_group_coordinator.device_group

633 + cpu_group = ep_group_coordinator.cpu_group

634 ep_rank = torch.distributed.get_rank()

597 - device = torch.device(f"cuda:{ep_rank}")

635 + accelerator_type = torch.accelerator.current_accelerator().type

636 + device = torch.device(f"{accelerator_type}:{ep_rank}")

637

638 num_layers = 2

639 num_local_experts = 2

@@ -623,9 +662,10 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

662 original_weights.append(layer_copy)

663

664 expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

665 + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl"

666 communicator = create_eplb_communicator_or_raise(

667 group_coordinator=ep_group_coordinator,

628 - backend="torch_nccl",

668 + backend=default_backend,

669 expert_weights=expert_weights,

670 expert_buffer=expert_buffer,

671 )

@@ -657,6 +697,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

697 assert_verification_synced(

698 local_ok,

699 "No-change EPLB verification failed on at least one rank.",

700 + cpu_group=cpu_group,

701 )

702

703

@@ -666,7 +707,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

707 (2, 2, 2, 3),

708 ],

709 )

669 -@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "nixl"])

710 +@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "torch_xccl", "nixl"])

711 def test_async_transfer_layer_without_mtp(

712 world_size: int,

713 num_layers: int,

@@ -677,6 +718,10 @@ def test_async_transfer_layer_without_mtp(

718 """Exercise async EPLB transfer path without MTP/spec decode."""

719 if eplb_communicator == "nixl" and not has_nixl():

720 pytest.skip("NIXL is not available")

721 + if eplb_communicator == "nixl" and current_platform.is_xpu():

722 + pytest.skip("NIXL does not support XPU")

723 + if eplb_communicator == "torch_xccl" and not current_platform.is_xpu():

724 + pytest.skip("torch_xccl requires XPU")

725 if torch.accelerator.device_count() < world_size:

726 pytest.skip(f"Need at least {world_size} GPUs to run the test")

727

@@ -715,9 +760,11 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

760 )

761

762 ep_group_coordinator = get_tp_group()

718 - ep_group = ep_group_coordinator.cpu_group

763 + ep_group = ep_group_coordinator.device_group

764 + cpu_group = ep_group_coordinator.cpu_group

765 ep_rank = torch.distributed.get_rank()

720 - device = torch.device(f"cuda:{ep_rank}")

766 + accelerator_type = torch.accelerator.current_accelerator().type

767 + device = torch.device(f"{accelerator_type}:{ep_rank}")

768

769 num_layers = 1

770 num_local_experts = 2

@@ -753,9 +800,10 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

800 original_weights.append(layer_copy)

801

802 expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

803 + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl"

804 communicator = create_eplb_communicator_or_raise(

805 group_coordinator=ep_group_coordinator,

758 - backend="torch_nccl",

806 + backend=default_backend,

807 expert_weights=expert_weights,

808 expert_buffer=expert_buffer,

809 )

@@ -788,6 +836,7 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

836 assert_verification_synced(

837 local_ok,

838 "Profile-mode EPLB verification failed on at least one rank.",

839 + cpu_group=cpu_group,

840 )

841

842

tests/distributed/test_eplb_fused_moe_layer.py

+6 -3

@@ -17,6 +17,7 @@ from vllm.distributed.parallel_state import (

17 get_tp_group,

18 )

19 from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner

20 +from vllm.platforms import current_platform

22 from .eplb_utils import distributed_run, set_env_vars_and_device

@@ -81,7 +82,7 @@ def make_fused_moe_layer(

82 )

83 re = fml.routed_experts

84 - device = torch.device(f"cuda:{rank}")

85 + device = torch.device(f"{torch.accelerator.current_accelerator().type}:{rank}")

87 from functools import partial

@@ -190,13 +191,14 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):

191 vllm_config = VllmConfig()

192 vllm_config.parallel_config.tensor_parallel_size = world_size

193 vllm_config.parallel_config.enable_expert_parallel = True

194 + vllm_config.parallel_config.enable_eplb = True

195

196 with set_current_vllm_config(vllm_config):

197 ensure_model_parallel_initialized(

198 tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

199 )

200

199 - ep_group = get_tp_group().cpu_group

201 + ep_group = get_tp_group().device_group

202 ep_rank = torch.distributed.get_rank()

203

204 fml_layers = [

@@ -216,9 +218,10 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):

218 shuffled_indices[lidx] = torch.randperm(test_config.num_experts)

219

220 expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]]

221 + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl"

222 communicator = create_eplb_communicator(

223 group_coordinator=get_eplb_group(),

221 - backend="torch_nccl",

224 + backend=default_backend,

225 expert_weights=rank_expert_weights,

226 expert_buffer=expert_buffer,

227 )

tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.modelopt import (

24 ModelOptNvFp4Config,

25 ModelOptNvFp4FusedMoE,

26 )

27 +from vllm.platforms import current_platform

29 from .eplb_utils import distributed_run, set_env_vars_and_device

@@ -270,6 +271,8 @@ def test_eplb_fml(

271 ):

272 if torch.accelerator.device_count() < world_size:

273 pytest.skip(f"Need at least {world_size} GPUs to run the test")

274 + if current_platform.is_xpu():

275 + pytest.skip("NVFP4 quantization and flashinfer backends are CUDA-only")

276

277 num_local_experts = num_experts // world_size

278 num_topk = 4

vllm/config/parallel.py

+12 -5

@@ -38,7 +38,9 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]

38 DataParallelBackend = Literal["ray", "mp"]

39 EPLBPolicyOption = Literal["default"]

40 DCPCommBackend = Literal["ag_rs", "a2a"]

41 -EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"]

41 +EPLBCommunicatorBackend = Literal[

42 + "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl"

43 +]

44 All2AllBackend = Literal[

45 "naive",

46 "pplx",

@@ -94,9 +96,11 @@ class EPLBConfig:

96 Backend for EPLB expert weight communication:

97 - "torch_nccl": Use torch.distributed on the device process group

98 - "torch_gloo": Use torch.distributed gloo with CPU staging

99 + - "torch_xccl": Use torch.distributed XCCL device P2P on XPU

100 - "nixl": Use NIXL with staged send/recv buffers

101 - "pynccl": Use PyNccl send/recv

99 - - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")

102 + - None: Auto-select backend ("torch_xccl" on XPU, prefers "nixl"

103 + on CUDA, falls back to "torch_gloo")

104 """

105

106 @model_validator(mode="after")

@@ -525,10 +529,10 @@ class ParallelConfig:

529 )

530

531 if self.enable_eplb:

528 - if not current_platform.is_cuda_alike():

532 + if not current_platform.is_cuda_alike() and not current_platform.is_xpu():

533 raise ValueError(

534 "Expert parallelism load balancing is only supported on "

531 - "CUDA devices or ROCm devices now."

535 + "CUDA devices or ROCm devices or XPU devices now."

536 )

537 if not self.enable_expert_parallel:

538 raise ValueError("enable_expert_parallel must be True to use EPLB.")

@@ -1034,7 +1038,10 @@ class ParallelConfig:

1038 # See https://github.com/pytorch/pytorch/issues/174288

1039 from vllm.distributed.nixl_utils import is_nixl_available

1040

1037 - if is_nixl_available():

1041 + if current_platform.is_xpu():

1042 + # On XPU, use the device-native XCCL P2P backend.

1043 + self.eplb_config.communicator = "torch_xccl"

1044 + elif is_nixl_available():

1045 self.eplb_config.communicator = "nixl"

1046 elif self.enable_elastic_ep:

1047 self.eplb_config.communicator = "pynccl"

vllm/distributed/eplb/async_worker.py

+14 -14

@@ -10,7 +10,7 @@ import torch

10 from vllm.distributed.parallel_state import get_eplb_group

11 from vllm.logger import init_logger

13 -from .eplb_utils import CpuGpuEvent

13 +from .eplb_utils import CpuGpuEvent, device_stream

14 from .rebalance_execute import AsyncEplbLayerResult, transfer_layer

16 if TYPE_CHECKING:

@@ -24,17 +24,17 @@ def start_async_worker(

24 is_profile: bool = False,

25 ) -> threading.Thread:

26 rank = get_eplb_group().device_group.rank()

27 - device_index = state.cuda_device_index

27 + device_index = state.device_index

28 assert state.is_async

30 def thread_target() -> None:

31 assert device_index is not None

32 torch.accelerator.set_device_index(device_index)

33 - cuda_stream = torch.cuda.Stream(device=device_index)

33 + stream = torch.Stream(device=torch.accelerator.current_accelerator())

34 try:

35 transfer_run_periodically(

36 state=state,

37 - cuda_stream=cuda_stream,

37 + stream=stream,

38 is_profile=is_profile,

39 )

40 except Exception as exc: # pragma: no cover - diagnostic path

@@ -49,13 +49,13 @@ def run_rebalance_experts(

49 model_state: "EplbModelState",

50 eplb_state: "EplbState",

51 physical_to_logical_map_cpu: torch.Tensor,

52 - cuda_stream: torch.cuda.Stream,

52 + stream: torch.Stream,

53 ) -> torch.Tensor:

54 assert model_state.eplb_stats is not None

55 eplb_stats = model_state.eplb_stats

57 # Move the global expert load window to CPU for computation.

58 - with torch.cuda.stream(cuda_stream):

58 + with device_stream(stream):

59 global_expert_load_window = eplb_stats.global_expert_load_window.cpu()

60 # Compute new expert mappings for the model

61 new_physical_to_logical_map = eplb_state.policy.rebalance_experts(

@@ -73,11 +73,11 @@ def run_rebalance_experts(

74 def transfer_run_periodically(

75 state: "EplbState",

76 - cuda_stream: torch.cuda.Stream,

76 + stream: torch.Stream,

77 is_profile: bool = False,

78 ) -> None:

79 while True:

80 - state.rearrange_event.wait(stream=cuda_stream)

80 + state.rearrange_event.wait(stream=stream)

82 eplb_group = get_eplb_group().device_group

83 eplb_cpu_group = get_eplb_group().cpu_group

@@ -87,16 +87,16 @@ def transfer_run_periodically(

87 for model_state in state.model_states.values():

88 layer_idx = 0

89 # Set the async worker's CUDA stream on the communicator

90 - model_state.communicator.set_stream(cuda_stream)

90 + model_state.communicator.set_stream(stream)

91 num_layers = model_state.model.num_moe_layers

93 # Snapshot the physical_to_logical_map (synchronized with

94 # rearrange_event) and copy it to CPU

95 - with torch.cuda.stream(cuda_stream):

95 + with device_stream(stream):

96 physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu()

98 new_physical_to_logical_map = run_rebalance_experts(

99 - model_state, state, physical_to_logical_map_cpu, cuda_stream

99 + model_state, state, physical_to_logical_map_cpu, stream

100 )

101

102 # Execute one EPLB layer transfer per model forward pass. Each iteration

@@ -132,13 +132,13 @@ def transfer_run_periodically(

132 communicator=model_state.communicator,

133 ep_group=eplb_group,

134 is_profile=is_profile,

135 - cuda_stream=cuda_stream,

135 + stream=stream,

136 layer_idx=layer_idx,

137 )

138

139 # Wait until all writes to expert_buffer have finished before making the

140 # AsyncEplbLayerResult visible to the main thread.

141 - cuda_stream.synchronize()

141 + stream.synchronize()

142

143 # This event guarantees that expert_buffer will not be overwritten by

144 # subsequent iterations of this loop until the main thread has consumed

@@ -155,6 +155,6 @@ def transfer_run_periodically(

155 # Block this thread until the main thread and main stream

156 # finish copying model_state.expert_buffer into

157 # model_state.model.expert_weights[layer_idx]

158 - consumed_event.wait(stream=cuda_stream)

158 + consumed_event.wait(stream=stream)

159 assert model_state.pending_result is None

160 layer_idx += 1

vllm/distributed/eplb/eplb_communicator.py

+78 -17

@@ -22,6 +22,7 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

22 from vllm.distributed.device_communicators.pynccl_wrapper import (

23 ncclDataTypeEnum,

24 )

25 +from vllm.distributed.eplb.eplb_utils import device_stream

26 from vllm.distributed.parallel_state import (

27 GroupCoordinator,

28 get_pp_group,

@@ -87,8 +88,8 @@ class EplbCommunicator(ABC):

88 communication buffers."""

89 return True

90 - def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:

91 - self._cuda_stream = cuda_stream

91 + def set_stream(self, stream: torch.Stream | None) -> None:

92 + self._stream = stream

94 def _log_initialized(self) -> None:

95 if is_local_first_rank():

@@ -101,10 +102,10 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator):

102 def __init__(

103 self,

104 ep_group: ProcessGroup,

104 - cuda_stream: torch.cuda.Stream | None = None,

105 + stream: torch.Stream | None = None,

106 ) -> None:

107 self._ep_group = ep_group

107 - self._cuda_stream = cuda_stream

108 + self._stream = stream

109 self._p2p_ops: list[P2POp] = []

110 self._log_initialized()

111

@@ -144,7 +145,7 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator):

145 if not self._p2p_ops:

146 return

147 try:

147 - with torch.cuda.stream(self._cuda_stream):

148 + with device_stream(self._stream):

149 reqs = batch_isend_irecv(self._p2p_ops)

150 for req in reqs:

151 req.wait()

@@ -158,10 +159,10 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

159 def __init__(

160 self,

161 cpu_group: ProcessGroup,

161 - cuda_stream: torch.cuda.Stream | None = None,

162 + stream: torch.Stream | None = None,

163 ) -> None:

164 self._cpu_group = cpu_group

164 - self._cuda_stream = cuda_stream

165 + self._stream = stream

166 self._ops: list[tuple[str, torch.Tensor, int]] = []

167 self._log_initialized()

168

@@ -217,7 +218,7 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

218 recv_staging.append((tensor, cpu_tensor))

219

220 try:

220 - with torch.cuda.stream(self._cuda_stream):

221 + with device_stream(self._stream):

222 build_ops()

223 finally:

224 self._ops.clear()

@@ -225,10 +226,10 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

226 # Wait for all D2H copies to finish

227 # before issuing gloo batch_isend_irecv operations.

228 with gpu_sync_allowed():

228 - if self._cuda_stream is not None:

229 - self._cuda_stream.synchronize()

229 + if self._stream is not None:

230 + self._stream.synchronize()

231 else:

231 - torch.cuda.current_stream().synchronize()

232 + torch.accelerator.current_stream().synchronize()

233

234 reqs = batch_isend_irecv(p2p_ops)

235 for req in reqs:

@@ -236,11 +237,69 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

237

238 if not recv_staging:

239 return

239 - with torch.cuda.stream(self._cuda_stream):

240 + with device_stream(self._stream):

241 for dst_tensor, cpu_tensor in recv_staging:

242 dst_tensor.copy_(cpu_tensor, non_blocking=True)

243

244

245 +class TorchDistXCCLStagedEplbCommunicator(EplbCommunicator):

246 + """EPLB communicator using XCCL device-to-device P2P on XPU."""

247 +

248 + def __init__(

249 + self,

250 + ep_group: ProcessGroup,

251 + stream: torch.Stream | None = None,

252 + ) -> None:

253 + self._ep_group = ep_group

254 + self._stream = stream

255 + self._ops: list[tuple[str, torch.Tensor, int]] = []

256 + self._log_initialized()

257 +

258 + def add_send(

259 + self,

260 + tensors: list[torch.Tensor],

261 + dst_rank: int,

262 + expert_id: int, # unused by this backend

263 + ) -> None:

264 + for tensor in tensors:

265 + self._ops.append(("send", tensor, dst_rank))

266 +

267 + def add_recv(

268 + self,

269 + tensors: list[torch.Tensor],

270 + src_rank: int,

271 + expert_id: int, # unused by this backend

272 + ) -> None:

273 + for tensor in tensors:

274 + self._ops.append(("recv", tensor, src_rank))

275 +

276 + def execute(self) -> None:

277 + if not self._ops:

278 + return

279 +

280 + p2p_ops: list[P2POp] = []

281 + try:

282 + for op, tensor, peer_rank in self._ops:

283 + send_or_recv = (

284 + torch.distributed.isend if op == "send" else torch.distributed.irecv

285 + )

286 + p2p_ops.append(

287 + P2POp(

288 + send_or_recv,

289 + tensor,

290 + peer_rank,

291 + self._ep_group,

292 + )

293 + )

294 + finally:

295 + self._ops.clear()

296 +

297 + with device_stream(self._stream):

298 + reqs = batch_isend_irecv(p2p_ops)

299 + for req in reqs:

300 + req.wait()

301 +

302 +

303 class NixlEplbCommunicator(EplbCommunicator):

304 """EPLB communicator backed by NIXL READ transfers."""

305

@@ -343,7 +402,7 @@ class NixlEplbCommunicator(EplbCommunicator):

402 uid = uuid.uuid4().hex[:8]

403 return f"eplb-{self._rank}{pp_suffix}-{uid}"

404

346 - def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:

405 + def set_stream(self, stream: torch.Stream | None) -> None:

406 pass

407

408 def add_send(

@@ -602,10 +661,10 @@ class PyNcclEplbCommunicator(EplbCommunicator):

661 def __init__(

662 self,

663 pynccl_comm: PyNcclCommunicator,

605 - cuda_stream: torch.cuda.Stream | None = None,

664 + stream: torch.Stream | None = None,

665 ) -> None:

666 self._pynccl_comm = pynccl_comm

608 - self._cuda_stream = cuda_stream

667 + self._stream = stream

668 self._group_started = False

669 self._log_initialized()

670

@@ -622,7 +681,7 @@ class PyNcclEplbCommunicator(EplbCommunicator):

681 ) -> None:

682 self._ensure_group_started()

683 for tensor in tensors:

625 - self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)

684 + self._pynccl_comm.send(tensor, dst_rank, stream=self._stream)

685

686 def add_recv(

687 self,

@@ -632,7 +691,7 @@ class PyNcclEplbCommunicator(EplbCommunicator):

691 ) -> None:

692 self._ensure_group_started()

693 for tensor in tensors:

635 - self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)

694 + self._pynccl_comm.recv(tensor, src_rank, stream=self._stream)

695

696 def execute(self) -> None:

697 if self._group_started:

@@ -750,6 +809,8 @@ def create_eplb_communicator(

809 return TorchDistGlooStagedEplbCommunicator(

810 cpu_group=group_coordinator.cpu_group,

811 )

812 + elif backend == "torch_xccl":

813 + return TorchDistXCCLStagedEplbCommunicator(ep_group=torch_group)

814 elif backend == "torch_nccl":

815 return TorchDistNcclEplbCommunicator(ep_group=torch_group)

816 elif backend == "pynccl":

vllm/distributed/eplb/eplb_state.py

+11 -11

@@ -203,9 +203,9 @@ class EplbModelState:

203 """

204 EPLB stats for the model.

205 """

206 - cuda_device_index: int | None

206 + device_index: int | None

207 """

208 - CUDA device index for the async EPLB worker thread.

208 + Device index for the async EPLB worker thread.

209 """

210 communicator: EplbCommunicator

211 """

@@ -286,14 +286,14 @@ class EplbState:

286 """

287 Background thread handling async transfers.

288 """

289 - self.cuda_device_index: int | None = None

289 + self.device_index: int | None = None

290 """

291 - CUDA device index for the async EPLB worker thread.

291 + Device index for the async EPLB worker thread.

292 """

293 - if self.device.type == "cuda":

294 - self.cuda_device_index = self.device.index

295 - if self.cuda_device_index is None and torch.cuda.is_available():

296 - self.cuda_device_index = torch.accelerator.current_device_index()

293 + if self.device.type in ("cuda", "xpu"):

294 + self.device_index = self.device.index

295 + if self.device_index is None and torch.accelerator.is_available():

296 + self.device_index = torch.accelerator.current_device_index()

297

298 @staticmethod

299 def build_initial_global_physical_to_logical_map(

@@ -504,7 +504,7 @@ class EplbState:

504 expert_buffer=expert_buffer,

505 rebalanced=False,

506 eplb_stats=None,

507 - cuda_device_index=self.cuda_device_index,

507 + device_index=self.device_index,

508 communicator=communicator,

509 num_unpadded_tokens_tensors=num_unpadded_tokens_tensors,

510 )

@@ -763,8 +763,8 @@ class EplbState:

763 is_main_rank = ep_rank == 0

764 if is_main_rank:

765 if not self.is_async or is_profile:

766 - start_event = torch.cuda.Event(enable_timing=True)

767 - end_event = torch.cuda.Event(enable_timing=True)

766 + start_event = torch.Event(enable_timing=True)

767 + end_event = torch.Event(enable_timing=True)

768 start_event.record()

769 logger.info(

770 "Rearranging experts %s %s...",

vllm/distributed/eplb/eplb_utils.py

+21 -4

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Utility functions for EPLB (Expert Parallel Load Balancing)."""

5 +import contextlib

6 import os

7 import threading

@@ -13,6 +14,22 @@ from vllm.logger import init_logger

14 logger = init_logger(__name__)

17 +@contextlib.contextmanager

18 +def device_stream(stream: torch.Stream | None):

19 + """Platform-agnostic context manager that activates *stream* as the

20 + current accelerator stream for the duration of the ``with`` block.

21 + A no-op when *stream* is ``None``."""

22 + if stream is None:

23 + yield

24 + return

25 + prev = torch.accelerator.current_stream()

26 + torch.accelerator.set_stream(stream)

27 + try:

28 + yield

29 + finally:

30 + torch.accelerator.set_stream(prev)

31 +

32 +

33 class CpuGpuEvent:

34 """Combines a CUDA event with a CPU threading event to enforce record->wait

35 ordering across two threads.

@@ -30,10 +47,10 @@ class CpuGpuEvent:

47 """

49 def __init__(self):

33 - self._event = torch.cuda.Event()

50 + self._event = torch.Event()

51 self._recorded = threading.Event()

36 - def wait(self, stream: torch.cuda.Stream | None = None):

53 + def wait(self, stream: torch.Stream | None = None):

54 """Blocks the calling thread until record finishes. Used to guarantee that the

55 record kernel is called before wait.

@@ -43,7 +60,7 @@ class CpuGpuEvent:

60 self._event.wait(stream)

61 self._recorded.clear()

46 - def record(self, stream: torch.cuda.Stream | None = None):

63 + def record(self, stream: torch.Stream | None = None):

64 """Unblocks the waiting thread after calling event.record().

66 Should only be called by the main thread.

@@ -53,7 +70,7 @@ class CpuGpuEvent:

70 "CpuGpuEvent.record() called before the previous event was "

71 "consumed by wait()"

72 )

56 - self._event = torch.cuda.Event()

73 + self._event = torch.Event()

74 self._event.record(stream)

75 self._recorded.set()

vllm/distributed/eplb/rebalance_execute.py

+9 -10

@@ -13,7 +13,7 @@ import torch

13 from torch.distributed import ProcessGroup, all_gather

15 from vllm.distributed.eplb.eplb_communicator import EplbCommunicator

16 -from vllm.distributed.eplb.eplb_utils import CpuGpuEvent

16 +from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream

17 from vllm.logger import init_logger

18 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

@@ -173,7 +173,7 @@ def move_to_buffer(

173 new_indices: np.ndarray,

174 expert_weights: Sequence[torch.Tensor],

175 expert_weights_buffers: Sequence[torch.Tensor],

176 - cuda_stream: torch.cuda.Stream | None,

176 + stream: torch.Stream | None,

177 ep_rank: int,

178 communicator: EplbCommunicator,

179 layer_idx: int = 0,

@@ -183,12 +183,12 @@ def move_to_buffer(

183 Args:

184 num_local_experts: Number of local experts.

185 old_indices: (num_experts_total,) ndarray of current (old)

186 - global-to-local expert assignments.

186 + global-to-locals expert assignments.

187 new_indices: (num_experts_total,) ndarray of desired (new)

188 global-to-local assignments after rebalance.

189 expert_weights: Original expert weights for the layer.

190 expert_weights_buffers: Intermediate buffers (one per tensor).

191 - cuda_stream: CUDA stream for async copies (can be None for sync mode).

191 + stream: CUDA/XPU stream for async copies (can be None for sync mode).

192 ep_rank: Rank of this process in expert parallel group.

193 communicator: EplbCommunicator instance for P2P communication.

194 layer_idx: Index of the MoE layer being transferred.

@@ -261,10 +261,9 @@ def move_to_buffer(

261 expert = new_local_expert_ids[dst]

262 src_local = expert_to_src_map.get(expert, -1)

263 if src_local != -1:

264 - with torch.cuda.stream(cuda_stream):

264 + with device_stream(stream):

265 for w, b in zip(expert_weights, expert_weights_buffers):

266 b[dst].copy_(w[src_local], non_blocking=True)

267 -

267 communicator.set_transfer_context(old_indices, layer_idx)

268

269 # 2. Post sends

@@ -430,7 +429,7 @@ def transfer_layer(

429 ep_group: ProcessGroup,

430 communicator: EplbCommunicator,

431 is_profile: bool = False,

433 - cuda_stream: torch.cuda.Stream | None = None,

432 + stream: torch.Stream | None = None,

433 rank_mapping: dict[int, int] | None = None,

434 layer_idx: int = 0,

435 ) -> TransferMetadata:

@@ -451,7 +450,7 @@ def transfer_layer(

450 is_profile (bool): If `True`, do not perform any actual weight copy.

451 This is used during profile run, where we only perform dummy

452 communications to reserve enough memory for the buffers.

454 - cuda_stream: CUDA stream for async copies (can be None for sync mode).

453 + stream: CUDA stream for async copies (can be None for sync mode).

454 rank_mapping: Optional rank mapping for elastic expert parallelism.

455 layer_idx: Index of the MoE layer being transferred.

456

@@ -499,7 +498,7 @@ def transfer_layer(

498 new_indices=new_layer_indices_np,

499 expert_weights=expert_weights,

500 expert_weights_buffers=expert_weights_buffer,

502 - cuda_stream=cuda_stream,

501 + stream=stream,

502 ep_rank=ep_group.rank(),

503 communicator=communicator,

504 layer_idx=layer_idx,

@@ -601,7 +600,7 @@ def rearrange_expert_weights_inplace(

600 new_indices=new_global_expert_indices_cpu[layer_idx],

601 expert_weights=expert_weights[layer_idx],

602 expert_weights_buffers=weights_buffer,

604 - cuda_stream=None,

603 + stream=None,

604 ep_rank=ep_rank,

605 communicator=communicator,

606 layer_idx=layer_idx,

vllm/model_executor/layers/fused_moe/experts/xpu_moe.py

@@ -171,6 +171,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):

171 num_experts=self.moe_config.num_local_experts,

172 ep_rank=self.moe_config.ep_rank,

173 ep_size=self.moe_config.ep_size,

174 + expert_map=expert_map,

175 gemm1_clamp_limit=self.gemm1_clamp_limit,

176 )

177 assert self.fused_moe_impl is not None

vllm/model_executor/layers/fused_moe/router/base_router.py

+1 -1

@@ -13,7 +13,7 @@ from vllm.platforms import current_platform

13 from vllm.triton_utils import tl, triton

14 from vllm.v1.worker.ubatching import dbo_current_ubatch_id

16 -if current_platform.is_cuda_alike():

16 +if current_platform.is_cuda_alike() or current_platform.is_xpu():

18 @triton.jit

19 def _eplb_map_and_record_i32_kernel(