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, 22:35 UTC
[Feature] Add async tensor parallelism for scaled mm (#20155)
Signed-off-by: cascade812 <cascade812@outlook.com>
cascade committed
Jul 30, 2025 at 14:23 UTC
287f527f5403bb42a32136cf6c802faeb92a09ef
3 files changed
+381
-8
tests/compile/test_async_tp.py
+138
-5
@@ -22,6 +22,8 @@ from ..utils import (compare_two_settings, create_new_process_for_each_test,
22
multi_gpu_test)
23
from .backend import TestBackend
24
25
+FP8_DTYPE = current_platform.fp8_dtype()
26
+
27
prompts = [
28
"Hello, my name is",
29
"The president of the United States is",
@@ -32,9 +34,10 @@ prompts = [
34
35
class TestMMRSModel(torch.nn.Module):
36
35
- def __init__(self, hidden_size=16):
37
+ def __init__(self, hidden_size=16, dtype=torch.float16):
38
super().__init__()
39
self.hidden_size = hidden_size
40
+ self.dtype = dtype
41
self.gate_proj = torch.nn.Parameter(torch.empty(
42
(self.hidden_size * 2, hidden_size)),
43
requires_grad=False)
@@ -64,9 +67,10 @@ class TestMMRSModel(torch.nn.Module):
67
68
class TestAGMMModel(torch.nn.Module):
69
67
- def __init__(self, hidden_size=16):
70
+ def __init__(self, hidden_size=16, dtype=torch.float16):
71
super().__init__()
72
self.hidden_size = hidden_size
73
+ self.dtype = dtype
74
self.weight = torch.nn.Parameter(torch.empty(
75
(hidden_size, hidden_size)),
76
requires_grad=False)
@@ -91,8 +95,125 @@ class TestAGMMModel(torch.nn.Module):
95
return [torch.ops.symm_mem.fused_all_gather_matmul.default]
96
97
98
+class _BaseScaledMMModel(torch.nn.Module):
99
+
100
+ def __init__(self, hidden_size=16, dtype=torch.float16):
101
+ super().__init__()
102
+ self.hidden_size = hidden_size
103
+ self.dtype = dtype
104
+ self.weight = torch.empty([hidden_size, hidden_size], dtype=FP8_DTYPE)\
105
+ .contiguous().transpose(0, 1)
106
+
107
+ # Initialize scale_b for _scaled_mm.
108
+ self.scale_b = torch.ones(1, self.hidden_size, dtype=torch.float32)
109
+
110
+
111
+class TestScaledMMRSModel(_BaseScaledMMModel):
112
+
113
+ def forward(self, input: torch.Tensor):
114
+ """
115
+ Forward pass implementing the scaled_mm + reduce scatter in the FX graph
116
+
117
+ """
118
+ fp8_input = input.to(FP8_DTYPE)
119
+ scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32)
120
+ scaled_mm = torch._scaled_mm(fp8_input,
121
+ self.weight,
122
+ scale_a=scale_a,
123
+ scale_b=self.scale_b,
124
+ out_dtype=self.dtype)
125
+ reduce_scatter = tensor_model_parallel_reduce_scatter(scaled_mm, dim=0)
126
+ return reduce_scatter
127
+
128
+ def ops_in_model_before(self):
129
+ return [torch.ops.vllm.reduce_scatter.default]
130
+
131
+ def ops_in_model_after(self):
132
+ return [torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter.default]
133
+
134
+
135
+class TestAGScaledMMModel(_BaseScaledMMModel):
136
+
137
+ def forward(self, input: torch.Tensor):
138
+ """
139
+ Forward pass implementing the all gather + scaled_mm in the FX graph
140
+ """
141
+ # Reshape input
142
+ fp8_input = input.to(FP8_DTYPE)
143
+ all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0)
144
+
145
+ scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32)
146
+ scaled_mm = torch._scaled_mm(all_gather,
147
+ self.weight,
148
+ scale_a=scale_a,
149
+ scale_b=self.scale_b,
150
+ out_dtype=self.dtype)
151
+ return scaled_mm
152
+
153
+ def ops_in_model_before(self):
154
+ return [torch.ops.vllm.all_gather.default]
155
+
156
+ def ops_in_model_after(self):
157
+ return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default]
158
+
159
+
160
+class TestCutlassScaledMMRSModel(_BaseScaledMMModel):
161
+
162
+ def forward(self, input: torch.Tensor):
163
+ """
164
+ Forward pass implementing the cutlass_scaled_mm + reduce scatter
165
+ in the FX graph
166
+
167
+ """
168
+ fp8_input = input.to(FP8_DTYPE)
169
+ scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32)
170
+ mm_out = torch.empty((fp8_input.shape[0], self.weight.shape[1]),
171
+ dtype=self.dtype,
172
+ device=input.device)
173
+ torch.ops._C.cutlass_scaled_mm(mm_out, fp8_input, self.weight, scale_a,
174
+ self.scale_b, None)
175
+ reduce_scatter = tensor_model_parallel_reduce_scatter(mm_out, dim=0)
176
+ return reduce_scatter
177
+
178
+ def ops_in_model_before(self):
179
+ return [torch.ops.vllm.reduce_scatter.default]
180
+
181
+ def ops_in_model_after(self):
182
+ return [torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter.default]
183
+
184
+
185
+class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
186
+
187
+ def forward(self, input: torch.Tensor):
188
+ """
189
+ Forward pass implementing the all gather + cutlass_scaled_mm
190
+ in the FX graph
191
+ """
192
+ # Reshape input
193
+ fp8_input = input.to(FP8_DTYPE)
194
+ all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0)
195
+
196
+ scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32)
197
+
198
+ mm_out = torch.empty((all_gather.shape[0], self.weight.shape[1]),
199
+ dtype=self.dtype,
200
+ device=all_gather.device)
201
+ torch.ops._C.cutlass_scaled_mm(mm_out, all_gather, self.weight,
202
+ scale_a, self.scale_b, None)
203
+ return mm_out
204
+
205
+ def ops_in_model_before(self):
206
+ return [torch.ops.vllm.all_gather.default]
207
+
208
+ def ops_in_model_after(self):
209
+ return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default]
210
+
211
+
212
@multi_gpu_test(num_gpus=2)
95
-@pytest.mark.parametrize("test_model", [TestMMRSModel, TestAGMMModel])
213
+@pytest.mark.parametrize("test_model", [
214
+ TestMMRSModel, TestAGMMModel, TestScaledMMRSModel, TestAGScaledMMModel,
215
+ TestCutlassScaledMMRSModel, TestAGCutlassScaledMMModel
216
+])
217
@pytest.mark.parametrize("batch_size", [8])
218
@pytest.mark.parametrize("seq_len", [16])
219
@pytest.mark.parametrize("hidden_size", [16])
@@ -101,6 +222,14 @@ class TestAGMMModel(torch.nn.Module):
222
reason="Only test on CUDA")
223
def test_async_tp_pass_replace(test_model: str, batch_size: int, seq_len: int,
224
hidden_size: int, dtype: torch.dtype):
225
+ if test_model in (TestScaledMMRSModel, TestAGScaledMMModel,
226
+ TestCutlassScaledMMRSModel,
227
+ TestAGCutlassScaledMMModel) and dtype == torch.float16:
228
+ pytest.skip(
229
+ "Only bf16 high precision output types are supported for " \
230
+ "per-token (row-wise) scaling"
231
+ )
232
+
233
num_processes = 2
234
235
def run_torch_spawn(fn, nprocs):
@@ -155,7 +284,8 @@ def async_tp_pass_on_test_model(local_rank: int, world_size: int,
284
async_tp_pass = AsyncTPPass(vllm_config)
285
backend = TestBackend(async_tp_pass)
286
158
- model = test_model_cls(hidden_size)
287
+ model = test_model_cls(hidden_size,
288
+ dtype) # Pass dtype to model constructor
289
290
hidden_states = torch.randn((batch_size * seq_len, hidden_size),
291
dtype=dtype,
@@ -174,7 +304,10 @@ def async_tp_pass_on_test_model(local_rank: int, world_size: int,
304
305
306
@create_new_process_for_each_test()
177
-@pytest.mark.parametrize("model_id", ["meta-llama/Llama-3.2-1B-Instruct"])
307
+@pytest.mark.parametrize("model_id", [
308
+ "meta-llama/Llama-3.2-1B-Instruct",
309
+ "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"
310
+])
311
@pytest.mark.parametrize("tp_size", [2])
312
@pytest.mark.parametrize("async_tp_enabled", [True])
313
@pytest.mark.parametrize("distributed_backend", ["mp"])
vllm/compilation/collective_fusion.py
+242
-2
@@ -15,10 +15,13 @@ from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
15
from vllm.distributed.parallel_state import (
16
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
17
from vllm.logger import init_logger
18
+from vllm.platforms import current_platform
19
from vllm.utils import direct_register_custom_op
20
21
from .vllm_inductor_pass import VllmInductorPass
22
23
+FP8_DTYPE = current_platform.fp8_dtype()
24
+
25
if find_spec("flashinfer"):
26
try:
27
import flashinfer.comm as flashinfer_comm
@@ -28,7 +31,6 @@ if find_spec("flashinfer"):
31
flashinfer_comm = None
32
else:
33
flashinfer_comm = None
31
-from vllm.platforms import current_platform
34
35
logger = init_logger(__name__)
36
@@ -118,6 +120,230 @@ class AllGatherGEMMPattern(BasePattern):
120
pm.fwd_only, pm_pass)
121
122
123
+class ScaledMMReduceScatterPattern(BasePattern):
124
+
125
+ def get_inputs(self):
126
+ input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
127
+ mm_weight = torch.empty([16, 16], device=self.device,
128
+ dtype=FP8_DTYPE).contiguous().transpose(0, 1)
129
+ scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
130
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
131
+ return [input, mm_weight, scale_a, scale_b]
132
+
133
+ def register(self, pm_pass: PatternMatcherPass):
134
+
135
+ def pattern(input: torch.Tensor, mat2: torch.Tensor,
136
+ scale_a: torch.Tensor,
137
+ scale_b: torch.Tensor) -> torch.Tensor:
138
+ scaled_mm = torch.ops.aten._scaled_mm.default(input,
139
+ mat2=mat2,
140
+ scale_a=scale_a,
141
+ scale_b=scale_b,
142
+ bias=None,
143
+ scale_result=None,
144
+ out_dtype=self.dtype)
145
+ reduce_scatter = torch.ops.vllm.reduce_scatter.default(
146
+ scaled_mm,
147
+ dim=0,
148
+ world_size=self.tp_size,
149
+ group_name=self.tp.unique_name)
150
+ return reduce_scatter
151
+
152
+ def replacement(input: torch.Tensor, mat2: torch.Tensor,
153
+ scale_a: torch.Tensor,
154
+ scale_b: torch.Tensor) -> torch.Tensor:
155
+ gemm_rs = torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(
156
+ input,
157
+ mat2,
158
+ scale_a,
159
+ scale_b,
160
+ "avg",
161
+ scatter_dim=0,
162
+ out_dtype=self.dtype,
163
+ group_name=self.tp.device_group.group_name,
164
+ )
165
+
166
+ return gemm_rs
167
+
168
+ pm.register_replacement(pattern, replacement, self.get_inputs(),
169
+ pm.fwd_only, pm_pass)
170
+
171
+
172
+class AllGatherScaledMMPattern(BasePattern):
173
+
174
+ def get_inputs(self):
175
+ x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
176
+ weight = torch.empty([16, 16], device=self.device,
177
+ dtype=FP8_DTYPE).contiguous().transpose(0, 1)
178
+
179
+ s1 = x.shape[0] * self.tp_size
180
+
181
+ scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
182
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
183
+
184
+ return [x, weight, scale_a, scale_b]
185
+
186
+ def register(self, pm_pass: PatternMatcherPass):
187
+
188
+ def pattern(
189
+ x: torch.Tensor,
190
+ weight: torch.Tensor,
191
+ scale_a: torch.Tensor,
192
+ scale_b: torch.Tensor,
193
+ ) -> torch.Tensor:
194
+ all_gather = torch.ops.vllm.all_gather.default(
195
+ x,
196
+ dim=0,
197
+ world_size=self.tp_size,
198
+ group_name=self.tp.unique_name)
199
+
200
+ return torch.ops.aten._scaled_mm.default(all_gather,
201
+ mat2=weight,
202
+ scale_a=scale_a,
203
+ scale_b=scale_b,
204
+ bias=None,
205
+ scale_result=None,
206
+ out_dtype=self.dtype)
207
+
208
+ def replacement(x: torch.Tensor, weight: torch.Tensor,
209
+ scale_a: torch.Tensor,
210
+ scale_b: torch.Tensor) -> torch.Tensor:
211
+ ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
212
+ x,
213
+ [weight],
214
+ scale_a,
215
+ [scale_b],
216
+ gather_dim=0,
217
+ biases=[None],
218
+ result_scales=[None],
219
+ out_dtypes=[self.dtype],
220
+ use_fast_accum=[False],
221
+ group_name=self.tp.device_group.group_name,
222
+ )
223
+ return mm_outputs
224
+
225
+ pm.register_replacement(pattern, replacement, self.get_inputs(),
226
+ pm.fwd_only, pm_pass)
227
+
228
+
229
+class CutlassScaledMMReduceScatterPattern(BasePattern):
230
+
231
+ def get_inputs(self):
232
+ input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
233
+ mm_weight = torch.empty([16, 16], device=self.device,
234
+ dtype=FP8_DTYPE).contiguous().transpose(0, 1)
235
+ scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
236
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
237
+
238
+ cutlass_mm_output = torch.empty([16, 16],
239
+ device=self.device,
240
+ dtype=self.dtype)
241
+ return [input, mm_weight, scale_a, scale_b, cutlass_mm_output]
242
+
243
+ def register(self, pm_pass: PatternMatcherPass):
244
+
245
+ def pattern(input: torch.Tensor, weight: torch.Tensor,
246
+ scale_a: torch.Tensor, scale_b: torch.Tensor,
247
+ cutlass_mm_output: torch.Tensor) -> torch.Tensor:
248
+ cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
249
+ torch.ops._C.cutlass_scaled_mm.default,
250
+ out=cutlass_mm_output,
251
+ a=input,
252
+ b=weight,
253
+ a_scales=scale_a,
254
+ b_scales=scale_b,
255
+ bias=None)
256
+
257
+ reduce_scatter = torch.ops.vllm.reduce_scatter.default(
258
+ cutlass_scaled_mm[1],
259
+ dim=0,
260
+ world_size=self.tp_size,
261
+ group_name=self.tp.unique_name)
262
+ return reduce_scatter
263
+
264
+ def replacement(input: torch.Tensor, mat2: torch.Tensor,
265
+ scale_a: torch.Tensor, scale_b: torch.Tensor,
266
+ cutlass_mm_output: torch.Tensor) -> torch.Tensor:
267
+ gemm_rs = torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(
268
+ input,
269
+ mat2,
270
+ scale_a,
271
+ scale_b,
272
+ "avg",
273
+ scatter_dim=0,
274
+ out_dtype=self.dtype,
275
+ group_name=self.tp.device_group.group_name,
276
+ )
277
+
278
+ return gemm_rs
279
+
280
+ pm.register_replacement(pattern, replacement, self.get_inputs(),
281
+ pm.fwd_only, pm_pass)
282
+
283
+
284
+class AllGatherCutlassScaledMMPattern(BasePattern):
285
+
286
+ def get_inputs(self):
287
+ x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
288
+ weight = torch.empty([16, 16], device=self.device,
289
+ dtype=FP8_DTYPE).contiguous().transpose(0, 1)
290
+
291
+ s1 = x.shape[0] * self.tp_size
292
+
293
+ scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
294
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
295
+
296
+ s2 = weight.shape[1]
297
+ output = torch.empty([s1, s2], device=self.device, dtype=self.dtype)
298
+
299
+ return [x, weight, scale_a, scale_b, output]
300
+
301
+ def register(self, pm_pass: PatternMatcherPass):
302
+
303
+ def pattern(
304
+ x: torch.Tensor,
305
+ weight: torch.Tensor,
306
+ scale_a: torch.Tensor,
307
+ scale_b: torch.Tensor,
308
+ output: torch.Tensor,
309
+ ) -> torch.Tensor:
310
+ all_gather = torch.ops.vllm.all_gather.default(
311
+ x,
312
+ dim=0,
313
+ world_size=self.tp_size,
314
+ group_name=self.tp.unique_name)
315
+
316
+ cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
317
+ torch.ops._C.cutlass_scaled_mm.default,
318
+ out=output,
319
+ a=all_gather,
320
+ b=weight,
321
+ a_scales=scale_a,
322
+ b_scales=scale_b,
323
+ bias=None)
324
+ return cutlass_scaled_mm[1]
325
+
326
+ def replacement(x: torch.Tensor, weight: torch.Tensor,
327
+ scale_a: torch.Tensor, scale_b: torch.Tensor,
328
+ output: torch.Tensor) -> torch.Tensor:
329
+ ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
330
+ x,
331
+ [weight],
332
+ scale_a,
333
+ [scale_b],
334
+ gather_dim=0,
335
+ biases=[None],
336
+ result_scales=[None],
337
+ out_dtypes=[self.dtype],
338
+ use_fast_accum=[False],
339
+ group_name=self.tp.device_group.group_name,
340
+ )
341
+ return mm_outputs
342
+
343
+ pm.register_replacement(pattern, replacement, self.get_inputs(),
344
+ pm.fwd_only, pm_pass)
345
+
346
+
347
class AsyncTPPass(VllmInductorPass):
348
349
def __init__(self, config: VllmConfig):
@@ -133,6 +359,20 @@ class AsyncTPPass(VllmInductorPass):
359
AllGatherGEMMPattern(self.model_dtype,
360
self.device).register(self.patterns)
361
362
+ # These fusions are enabled only for bfloat16 models because
363
+ # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
364
+ # only supports bfloat16 as the output dtype.
365
+ if self.model_dtype == torch.bfloat16:
366
+ ScaledMMReduceScatterPattern(self.model_dtype,
367
+ self.device).register(self.patterns)
368
+ AllGatherScaledMMPattern(self.model_dtype,
369
+ self.device).register(self.patterns)
370
+
371
+ CutlassScaledMMReduceScatterPattern(
372
+ self.model_dtype, self.device).register(self.patterns)
373
+ AllGatherCutlassScaledMMPattern(
374
+ self.model_dtype, self.device).register(self.patterns)
375
+
376
def is_applicable_for_shape(self, shape: Optional[int]) -> bool:
377
# only do replace for specific shapes
378
tp_size = get_tensor_model_parallel_world_size()
@@ -142,7 +382,7 @@ class AsyncTPPass(VllmInductorPass):
382
self.begin()
383
self.dump_graph(graph, "before_async_tp_pass")
384
count = self.patterns.apply(graph)
145
- logger.debug("Replaced %s patterns", count)
385
+ logger.debug("Replaced %s patterns with async TP pass.", count)
386
self.dump_graph(graph, "after_async_tp_pass")
387
self.end_and_log()
388
vllm/compilation/sequence_parallelism.py
+1
-1
@@ -477,6 +477,6 @@ class SequenceParallelismPass(VllmInductorPass):
477
self.begin()
478
self.dump_graph(graph, "before_sequence_parallelism_pass")
479
count = self.patterns.apply(graph)
480
- logger.debug("Replaced %s patterns", count)
480
+ logger.debug("Replaced %s patterns with sequence parallelism", count)
481
self.dump_graph(graph, "after_sequence_parallelism_pass")
482
self.end_and_log()