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 9, 23:00 UTC
[XPU]Support AgRsAll2AllManager on XPU device (#32654)
Signed-off-by: yisheng <yi.sheng@intel.com>
YiSheng5 committed
Jan 20, 2026 at 22:27 UTC
13f6630a9ea78bee4bd80bb6e842e55e374eec9a
1 file changed
+130
-7
vllm/distributed/device_communicators/xpu_communicator.py
+130
-7
@@ -23,23 +23,146 @@ class XpuCommunicator(DeviceCommunicatorBase):
23
):
24
super().__init__(cpu_group, device, device_group, unique_name)
25
if self.use_all2all:
26
- if self.all2all_backend != "naive": # type: ignore[has-type]
27
- logger.warning(
28
- "`%s` all2all manager is not supported on XPU. "
29
- "Falling back to `naive` all2all manager for XPU.",
30
- self.all2all_backend, # type: ignore[has-type]
31
- )
32
- self.all2all_backend = "naive"
26
if self.all2all_backend == "naive":
27
from .all2all import NaiveAll2AllManager
28
29
self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
30
logger.info("Using naive all2all manager.")
31
32
+ elif self.all2all_backend == "allgather_reducescatter":
33
+ from .all2all import AgRsAll2AllManager
34
+
35
+ self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
36
+ logger.info("Using AgRs manager on XPU device.")
37
+
38
+ else: # type: ignore[has-type]
39
+ logger.warning(
40
+ "`%s` all2all manager is not supported on XPU. "
41
+ "Falling back to AgRs manager for XPU, "
42
+ "which is the Default backend",
43
+ self.all2all_backend, # type: ignore[has-type]
44
+ )
45
+ from .all2all import AgRsAll2AllManager
46
+
47
+ self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
48
+ logger.info("Using AgRs manager on XPU device.")
49
+
50
def all_reduce(self, input_) -> torch.Tensor:
51
dist.all_reduce(input_, group=self.device_group)
52
return input_
53
54
+ def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
55
+ world_size = self.world_size
56
+
57
+ if dim < 0:
58
+ # Convert negative dim to positive.
59
+ dim += input_.dim()
60
+
61
+ # Note: This will produce an incorrect answer if we don't make
62
+ # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
63
+ input_tensor = input_.movedim(0, dim).contiguous()
64
+
65
+ assert input_tensor.shape[0] % world_size == 0
66
+ chunk_size = input_tensor.shape[0] // world_size
67
+ output_shape = (chunk_size,) + input_tensor.shape[1:]
68
+
69
+ output = torch.empty(
70
+ output_shape, dtype=input_tensor.dtype, device=input_tensor.device
71
+ )
72
+
73
+ dist.reduce_scatter_tensor(output, input_tensor)
74
+
75
+ # Reshape before returning
76
+ return output.movedim(0, dim).contiguous()
77
+
78
+ def reduce_scatterv(
79
+ self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
80
+ ):
81
+ world_size = self.world_size
82
+
83
+ if dim < 0:
84
+ # Convert negative dim to positive.
85
+ dim += input_.dim()
86
+
87
+ # Note: This will produce an incorrect answer if we don't make
88
+ # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
89
+ input_tensor = input_.movedim(0, dim).contiguous()
90
+
91
+ if sizes is not None:
92
+ assert len(sizes) == world_size
93
+ assert input_tensor.shape[0] == sum(sizes)
94
+ chunk_size = sizes[self.rank_in_group]
95
+ else:
96
+ assert input_tensor.shape[0] % world_size == 0
97
+ chunk_size = input_tensor.shape[0] // world_size
98
+ output_shape = (chunk_size,) + input_tensor.shape[1:]
99
+
100
+ output = torch.empty(
101
+ output_shape, dtype=input_tensor.dtype, device=input_tensor.device
102
+ )
103
+ if sizes is not None and sizes.count(sizes[0]) != len(sizes):
104
+ # if inputs shape in different ranks is not the same using reduce_scatter
105
+ input_splits = list(input_tensor.split(sizes, dim=0))
106
+ dist.reduce_scatter(output, input_splits)
107
+ else:
108
+ dist.reduce_scatter_tensor(output, input_tensor)
109
+ # Reshape before returning
110
+ return output.movedim(0, dim).contiguous()
111
+
112
+ def all_gatherv(
113
+ self,
114
+ input_: torch.Tensor | list[torch.Tensor],
115
+ dim: int = 0,
116
+ sizes: list[int] | None = None,
117
+ ):
118
+ if dim != 0:
119
+ raise NotImplementedError("only dim 0 all-gatherv is supported")
120
+ world_size = self.world_size
121
+
122
+ # 'sizes' is not needed if all inputs in the same group have the same
123
+ # shape
124
+ if sizes is not None and all(s == sizes[0] for s in sizes):
125
+ sizes = None
126
+
127
+ def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
128
+ input_size = input_.size()
129
+ if sizes is not None:
130
+ assert len(sizes) == world_size
131
+ assert input_.shape[dim] == sizes[self.rank_in_group], (
132
+ f"{input_.shape[dim]} != {sizes[self.rank_in_group]}"
133
+ )
134
+ output_size = (sum(sizes),) + input_size[1:]
135
+ else:
136
+ output_size = (input_size[0] * world_size,) + input_size[1:]
137
+ # Allocate output tensor.
138
+ output_tensor = torch.empty(
139
+ output_size, dtype=input_.dtype, device=input_.device
140
+ )
141
+
142
+ if sizes is not None:
143
+ all_gather_list = []
144
+ for size in sizes:
145
+ all_gather_list.append(
146
+ torch.empty(
147
+ (size,) + input_.shape[1:],
148
+ dtype=input_.dtype,
149
+ device=input_.device,
150
+ )
151
+ )
152
+ dist.all_gather(all_gather_list, input_)
153
+ output_tensor = torch.cat(all_gather_list, dim=0)
154
+ else:
155
+ dist.all_gather([output_tensor], input_)
156
+ return output_tensor
157
+
158
+ if isinstance(input_, torch.Tensor):
159
+ return _all_gather_single(input_, sizes)
160
+
161
+ output_list = []
162
+ for inp in input_:
163
+ output_list.append(_all_gather_single(inp, sizes=sizes))
164
+ return output_list
165
+
166
def gather(
167
self, input_: torch.Tensor, dst: int = 0, dim: int = -1
168
) -> torch.Tensor | None: