[XPU] Fix device assignment for DP external LB (#53037)
Signed-off-by: Tony Lin <tony.lin@intel.com>
Tony Lin committed
Sep 4, 2026 at 15:51 UTC
605ca45b917ac0614958c63851c31a3f1c82e20d
1 file changed
+63
-7
vllm/v1/worker/xpu_worker.py
+63
-7
@@ -47,18 +47,74 @@ class XPUWorker(Worker):
47
parallel_config.distributed_executor_backend
48
not in ("ray", "external_launcher")
49
and parallel_config.data_parallel_backend != "ray"
50
- and parallel_config.nnodes_within_dp == 1
50
+ and (
51
+ parallel_config.data_parallel_external_lb
52
+ or parallel_config.nnodes_within_dp == 1
53
+ )
54
):
55
dp_local_rank = parallel_config.data_parallel_rank_local
56
if dp_local_rank is None:
57
dp_local_rank = parallel_config.data_parallel_index
55
- tp_pp_world_size = (
56
- parallel_config.pipeline_parallel_size
57
- * parallel_config.tensor_parallel_size
58
- )
59
- self.local_rank += dp_local_rank * tp_pp_world_size
60
-
58
+ replica_world_size = parallel_config.world_size
59
visible_device_count = torch.accelerator.device_count()
60
+
61
+ if parallel_config.data_parallel_external_lb:
62
+ if parallel_config.nnodes_within_dp > 1:
63
+ # The replica spans nodes, so use its node-local shard.
64
+ replica_world_size = parallel_config.local_world_size
65
+ if replica_world_size > visible_device_count:
66
+ raise ValueError(
67
+ f"Local TP/PP/PCP replica size ({replica_world_size}) "
68
+ "exceeds the number of visible XPU devices "
69
+ f"({visible_device_count})."
70
+ )
71
+ local_dp_capacity = visible_device_count // replica_world_size
72
+ elif replica_world_size < visible_device_count:
73
+ # A node can host multiple complete TP/PP/PCP replicas.
74
+ local_dp_capacity = visible_device_count // replica_world_size
75
+ if visible_device_count % replica_world_size != 0:
76
+ logger.warning_once(
77
+ "XPU external LB cannot evenly divide "
78
+ "%d visible devices into TP/PP/PCP replicas of "
79
+ "size %d. This node can host %d complete DP "
80
+ "replicas, leaving %d visible devices unused.",
81
+ visible_device_count,
82
+ replica_world_size,
83
+ local_dp_capacity,
84
+ visible_device_count % replica_world_size,
85
+ )
86
+ elif replica_world_size == visible_device_count:
87
+ # A node hosts exactly one complete TP/PP/PCP replica.
88
+ local_dp_capacity = 1
89
+ logger.warning_once(
90
+ "XPU external LB sees exactly enough devices for one "
91
+ "TP/PP/PCP replica. This may be the intended "
92
+ "configuration, but it may also indicate that device "
93
+ "visibility is misconfigured. Every DP rank must see "
94
+ "all XPU devices on its node; consider removing "
95
+ "ZE_AFFINITY_MASK or setting it to expose the complete "
96
+ "device set."
97
+ )
98
+ else:
99
+ # The topology says single-node, but the replica does not fit.
100
+ raise ValueError(
101
+ f"TP/PP/PCP replica size ({replica_world_size}) exceeds "
102
+ f"the number of visible XPU devices ({visible_device_count}), "
103
+ "but nnodes_within_dp is 1. Configure the multi-node "
104
+ "topology, or ensure every DP rank can see all devices "
105
+ "on its node."
106
+ )
107
+ # Strip the node component off the global DP index to get this
108
+ # engine's slot on its own node. Assumes the launcher assigns
109
+ # DP ranks to nodes in contiguous blocks (node 0 gets ranks
110
+ # 0..capacity-1, and so on), which is what the usual sequential
111
+ # and one-pod-per-rank deployments do. A round-robin or
112
+ # unbalanced assignment would silently map two engines onto the
113
+ # same device.
114
+ dp_local_rank = parallel_config.data_parallel_index % local_dp_capacity
115
+
116
+ self.local_rank += dp_local_rank * replica_world_size
117
+
118
assert self.local_rank < visible_device_count, (
119
f"DP adjusted local rank {self.local_rank} is out of bounds. "
120
)