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 14, 05:30 UTC
[compile] Add patched_fused_scaled_matmul_reduce_scatter (#26604)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Oct 11, 2025 at 05:44 UTC
a25f2adee9d66bf16128f2a8f399c558fd181647
3 files changed
+119
-6
tests/compile/test_async_tp.py
+22
-4
@@ -142,7 +142,7 @@ class TestScaledMMRSModel(_BaseScaledMMModel):
142
return [torch.ops.vllm.reduce_scatter.default]
143
144
def ops_in_model_after(self):
145
- return [torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter.default]
145
+ return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default]
146
147
148
class TestAGScaledMMModel(_BaseScaledMMModel):
@@ -195,7 +195,7 @@ class TestCutlassScaledMMRSModel(_BaseScaledMMModel):
195
return [torch.ops.vllm.reduce_scatter.default]
196
197
def ops_in_model_after(self):
198
- return [torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter.default]
198
+ return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default]
199
200
201
class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
@@ -243,9 +243,15 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
243
@pytest.mark.parametrize("seq_len", [16])
244
@pytest.mark.parametrize("hidden_size", [16])
245
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
246
+@pytest.mark.parametrize("dynamic", [True, False])
247
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
248
def test_async_tp_pass_replace(
248
- test_model: str, batch_size: int, seq_len: int, hidden_size: int, dtype: torch.dtype
249
+ test_model: str,
250
+ batch_size: int,
251
+ seq_len: int,
252
+ hidden_size: int,
253
+ dtype: torch.dtype,
254
+ dynamic: bool,
255
):
256
if (
257
test_model
@@ -269,7 +275,15 @@ def test_async_tp_pass_replace(
275
# torch.distributed and cuda
276
torch.multiprocessing.spawn(
277
fn,
272
- args=(num_processes, test_model, batch_size, seq_len, hidden_size, dtype),
278
+ args=(
279
+ num_processes,
280
+ test_model,
281
+ batch_size,
282
+ seq_len,
283
+ hidden_size,
284
+ dtype,
285
+ dynamic,
286
+ ),
287
nprocs=nprocs,
288
)
289
@@ -284,6 +298,7 @@ def async_tp_pass_on_test_model(
298
seq_len: int,
299
hidden_size: int,
300
dtype: torch.dtype,
301
+ dynamic: bool,
302
):
303
current_platform.seed_everything(0)
304
@@ -331,6 +346,9 @@ def async_tp_pass_on_test_model(
346
(batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False
347
)
348
349
+ if dynamic:
350
+ torch._dynamo.mark_dynamic(hidden_states, 0)
351
+
352
compiled_model = torch.compile(model, backend=backend)
353
compiled_model(hidden_states)
354
vllm/compilation/collective_fusion.py
+2
-2
@@ -172,7 +172,7 @@ class ScaledMMReduceScatterPattern(BasePattern):
172
# Calculate output shape: input @ mat2 with scatter_dim reduced
173
output_shape = [*input.shape[:-1], mat2.shape[1]]
174
scatter_dim = 0
175
- gemm_rs = torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(
175
+ gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
176
input,
177
mat2,
178
scale_a,
@@ -307,7 +307,7 @@ class CutlassScaledMMReduceScatterPattern(BasePattern):
307
# Calculate output shape: input @ mat2 with scatter_dim reduced
308
output_shape = [*input.shape[:-1], mat2.shape[1]]
309
scatter_dim = 0
310
- gemm_rs = torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(
310
+ gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
311
input,
312
mat2,
313
scale_a,
vllm/distributed/parallel_state.py
+95
@@ -37,6 +37,8 @@ from unittest.mock import patch
37
38
import torch
39
import torch.distributed
40
+import torch.distributed._functional_collectives as funcol
41
+import torch.distributed._symmetric_memory
42
from torch.distributed import Backend, ProcessGroup
43
from typing_extensions import deprecated
44
@@ -159,6 +161,90 @@ def all_gather_fake(
161
return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
162
163
164
+def patched_fused_scaled_matmul_reduce_scatter_fake(
165
+ A: torch.Tensor,
166
+ B: torch.Tensor,
167
+ A_scale: torch.Tensor,
168
+ B_scale: torch.Tensor,
169
+ reduce_op: str,
170
+ orig_scatter_dim: int,
171
+ scatter_dim_after_maybe_reshape: int,
172
+ group_name: str,
173
+ output_shape: list[int],
174
+ bias: torch.Tensor | None = None,
175
+ result_scale: torch.Tensor | None = None,
176
+ out_dtype: torch.dtype | None = None,
177
+ use_fast_accum: bool = False,
178
+) -> torch.Tensor:
179
+ # Copied from
180
+ # https://github.com/pytorch/pytorch/blob/50c338c2da905062449e4d9ac807832d1b5cd90e/torch/distributed/_symmetric_memory/__init__.py#L1189
181
+ if A_scale.numel() > 1:
182
+ if A_scale.shape[:-1] != A.shape[:-1]:
183
+ raise ValueError(
184
+ "For row-wise scaling, the leading dims of A_scale "
185
+ "must match the leading dims of A "
186
+ f"(A shape: {A.shape}, A_scale shape: {A_scale.shape})"
187
+ )
188
+ A_scale = A_scale.flatten(0, -2).contiguous()
189
+ elif A_scale.numel() != 1:
190
+ raise ValueError(
191
+ "Invalid A_scale shape "
192
+ f"(A shape: {A.shape}, A_scale shape: {A_scale.shape})"
193
+ )
194
+
195
+ C = torch._scaled_mm(
196
+ A.flatten(0, -2).contiguous(),
197
+ B,
198
+ A_scale,
199
+ B_scale,
200
+ bias,
201
+ result_scale,
202
+ out_dtype,
203
+ use_fast_accum,
204
+ )
205
+ C = C.view(*output_shape[:-1], B.shape[1])
206
+ res = funcol.reduce_scatter_tensor(
207
+ C,
208
+ reduce_op,
209
+ orig_scatter_dim, # need original scatter dim for 3D+ output tensor here
210
+ group_name,
211
+ )
212
+ res = funcol.wait_tensor(res)
213
+ return res
214
+
215
+
216
+def patched_fused_scaled_matmul_reduce_scatter(
217
+ A: torch.Tensor,
218
+ B: torch.Tensor,
219
+ A_scale: torch.Tensor,
220
+ B_scale: torch.Tensor,
221
+ reduce_op: str,
222
+ orig_scatter_dim: int,
223
+ scatter_dim_after_maybe_reshape: int,
224
+ group_name: str,
225
+ output_shape: list[int],
226
+ bias: torch.Tensor | None = None,
227
+ result_scale: torch.Tensor | None = None,
228
+ out_dtype: torch.dtype | None = None,
229
+ use_fast_accum: bool = False,
230
+) -> torch.Tensor:
231
+ return torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(
232
+ A,
233
+ B,
234
+ A_scale,
235
+ B_scale,
236
+ reduce_op,
237
+ orig_scatter_dim,
238
+ scatter_dim_after_maybe_reshape,
239
+ group_name,
240
+ output_shape,
241
+ bias,
242
+ result_scale,
243
+ out_dtype,
244
+ use_fast_accum,
245
+ )
246
+
247
+
248
if supports_custom_op():
249
direct_register_custom_op(
250
op_name="all_reduce",
@@ -178,6 +264,15 @@ if supports_custom_op():
264
fake_impl=all_gather_fake,
265
)
266
267
+ # TODO: Remove this once the pytorch fix
268
+ # (https://github.com/pytorch/pytorch/pull/165086) gets released,
269
+ # in either 2.9.1 or 2.10
270
+ direct_register_custom_op(
271
+ op_name="patched_fused_scaled_matmul_reduce_scatter",
272
+ op_func=patched_fused_scaled_matmul_reduce_scatter,
273
+ fake_impl=patched_fused_scaled_matmul_reduce_scatter_fake,
274
+ )
275
+
276
277
class GroupCoordinator:
278
"""