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
[BugFix][torch.compile] Fix fused_scaled_matmul_reduce_scatter signature for PyTorch 2.8 (#26038)
Signed-off-by: jasonlizhengjian <jasonlizhengjian@gmail.com> Signed-off-by: <> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Jason Li committed
Oct 10, 2025 at 10:42 UTC
f4ba2061cf412365405d1b2c4239fe3f5981271c
2 files changed
+24
-8
.buildkite/test-pipeline.yaml
+2
-2
@@ -400,8 +400,6 @@ steps:
400
- pytest -v -s compile/test_fusion_attn.py
401
- pytest -v -s compile/test_functionalization.py
402
- pytest -v -s compile/test_silu_mul_quant_fusion.py
403
- - pytest -v -s compile/test_sequence_parallelism.py
404
- - pytest -v -s compile/test_async_tp.py
403
- pytest -v -s compile/test_fusion_all_reduce.py
404
- pytest -v -s compile/test_decorator.py
405
- pytest -v -s compile/test_noop_elimination.py
@@ -1093,6 +1091,8 @@ steps:
1091
working_dir: "/vllm-workspace/"
1092
num_gpus: 2
1093
commands:
1094
+ - pytest -v -s tests/compile/test_async_tp.py
1095
+ - pytest -v -s tests/compile/test_sequence_parallelism.py
1096
- pytest -v -s tests/distributed/test_context_parallel.py
1097
- CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048
1098
vllm/compilation/collective_fusion.py
+22
-6
@@ -169,15 +169,23 @@ class ScaledMMReduceScatterPattern(BasePattern):
169
scale_a: torch.Tensor,
170
scale_b: torch.Tensor,
171
) -> torch.Tensor:
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(
176
input,
177
mat2,
178
scale_a,
179
scale_b,
180
"avg",
178
- scatter_dim=0,
179
- out_dtype=self.dtype,
180
- group_name=self.tp.device_group.group_name,
181
+ scatter_dim, # orig_scatter_dim
182
+ scatter_dim, # scatter_dim_after_maybe_reshape
183
+ self.tp.device_group.group_name,
184
+ output_shape,
185
+ None, # bias
186
+ None, # result_scale
187
+ self.dtype, # out_dtype
188
+ False, # use_fast_accum
189
)
190
191
return gemm_rs
@@ -296,15 +304,23 @@ class CutlassScaledMMReduceScatterPattern(BasePattern):
304
scale_b: torch.Tensor,
305
cutlass_mm_output: torch.Tensor,
306
) -> torch.Tensor:
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(
311
input,
312
mat2,
313
scale_a,
314
scale_b,
315
"avg",
305
- scatter_dim=0,
306
- out_dtype=self.dtype,
307
- group_name=self.tp.device_group.group_name,
316
+ scatter_dim, # orig_scatter_dim
317
+ scatter_dim, # scatter_dim_after_maybe_reshape
318
+ self.tp.device_group.group_name,
319
+ output_shape,
320
+ None, # bias
321
+ None, # result_scale
322
+ self.dtype, # out_dtype
323
+ False, # use_fast_accum
324
)
325
326
return gemm_rs