@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,013 lines 33.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 from contextlib import suppress
6
7 import torch
8 import torch._inductor.pattern_matcher as pm
9 import torch.distributed.distributed_c10d as c10d
10 import torch.fx as fx
11 from torch._inductor.pattern_matcher import PatternMatcherPass
12 from torch.distributed._symmetric_memory import enable_symm_mem_for_group
13
14 from vllm.config import VllmConfig
15 from vllm.config.utils import Range
16 from vllm.distributed import get_tp_group
17 from vllm.distributed.parallel_state import (
18 get_tensor_model_parallel_world_size,
19 )
20 from vllm.logger import init_logger
21 from vllm.platforms import current_platform
22 from vllm.utils.torch_utils import direct_register_custom_op
23
24 from ..inductor_pass import enable_fake_mode
25 from ..vllm_inductor_pass import (
26 VllmFusionPatternMatcherPass,
27 VllmInductorPass,
28 VllmPatternMatcherPass,
29 VllmPatternReplacement,
30 )
31
32 FP8_DTYPE = current_platform.fp8_dtype()
33
34 logger = init_logger(__name__)
35
36
37 def _flashinfer_scaled_mm_out(
38 A: torch.Tensor,
39 B: torch.Tensor,
40 *,
41 scale_a: torch.Tensor,
42 scale_b: torch.Tensor,
43 out: torch.Tensor,
44 bias: torch.Tensor | None = None,
45 scale_result: torch.Tensor | None = None,
46 out_dtype: torch.dtype | None = None,
47 use_fast_accum: bool = False,
48 ) -> None:
49 # Import lazily to avoid a circular import during module initialization
50 # when docs or other tooling import the pass without FlashInfer.
51 from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm_out
52
53 assert bias is None, "FlashInfer symm_mem adapter does not support bias"
54 assert scale_result is None, (
55 "FlashInfer symm_mem adapter does not support result scaling"
56 )
57 assert not use_fast_accum, (
58 "FlashInfer symm_mem adapter does not support use_fast_accum"
59 )
60 assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, (
61 "FlashInfer symm_mem adapter expects 2D inputs and output"
62 )
63 assert scale_a.numel() == 1 and scale_b.numel() == 1, (
64 "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales"
65 )
66
67 flashinfer_scaled_fp8_mm_out(
68 A,
69 B,
70 scale_a,
71 scale_b,
72 out=out,
73 out_dtype=out_dtype or out.dtype,
74 )
75
76
77 def _flashinfer_fp4_mm_out(
78 A: torch.Tensor,
79 B: torch.Tensor,
80 *,
81 scale_a: torch.Tensor,
82 scale_b: torch.Tensor,
83 out: torch.Tensor,
84 alpha: torch.Tensor,
85 out_dtype: torch.dtype | None = None,
86 use_8x4_sf_layout: bool = False,
87 backend: str = "cutlass",
88 ) -> None:
89 from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm_out
90
91 assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, (
92 "FlashInfer FP4 symm_mem adapter expects 2D inputs and output"
93 )
94 flashinfer_scaled_fp4_mm_out(
95 A,
96 B,
97 scale_a,
98 scale_b,
99 alpha,
100 out=out,
101 out_dtype=out_dtype or out.dtype,
102 use_8x4_sf_layout=use_8x4_sf_layout,
103 backend=backend,
104 )
105
106
107 def fused_flashinfer_scaled_matmul_reduce_scatter_fake(
108 A: torch.Tensor,
109 B: torch.Tensor,
110 A_scale: torch.Tensor,
111 B_scale: torch.Tensor,
112 reduce_op: str,
113 orig_scatter_dim: int,
114 scatter_dim_after_maybe_reshape: int,
115 group_name: str,
116 output_shape: list[int],
117 out_dtype: torch.dtype | None = None,
118 ) -> torch.Tensor:
119 world_size = c10d._resolve_process_group(group_name).size()
120 result_shape = list(output_shape)
121 result_shape[orig_scatter_dim] //= world_size
122 return torch.empty(
123 result_shape,
124 dtype=out_dtype or torch.bfloat16,
125 device=A.device,
126 )
127
128
129 def fused_flashinfer_scaled_matmul_reduce_scatter(
130 A: torch.Tensor,
131 B: torch.Tensor,
132 A_scale: torch.Tensor,
133 B_scale: torch.Tensor,
134 reduce_op: str,
135 orig_scatter_dim: int,
136 scatter_dim_after_maybe_reshape: int,
137 group_name: str,
138 output_shape: list[int],
139 out_dtype: torch.dtype | None = None,
140 ) -> torch.Tensor:
141 assert orig_scatter_dim == 0 and scatter_dim_after_maybe_reshape == 0, (
142 "FlashInfer symm_mem adapter currently only supports scatter_dim=0"
143 )
144 world_size = c10d._resolve_process_group(group_name).size()
145 assert A.ndim == 2 and B.ndim == 2, "FlashInfer symm_mem adapter expects 2D inputs"
146 assert A.is_contiguous(), "FlashInfer symm_mem adapter expects contiguous A"
147 assert A_scale.numel() == 1 and B_scale.numel() == 1, (
148 "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales"
149 )
150 assert A.shape[0] % world_size == 0, (
151 "FlashInfer symm_mem adapter expects M divisible by world size"
152 )
153
154 kwargs = {
155 "scale_b": B_scale,
156 "bias": None,
157 "scale_result": None,
158 "out_dtype": out_dtype,
159 "use_fast_accum": False,
160 }
161 return torch.distributed._symmetric_memory._fused_scaled_matmul_reduce_scatter_impl(
162 mm_out_op=_flashinfer_scaled_mm_out,
163 A=A,
164 B=B,
165 A_scale=A_scale,
166 kwargs=kwargs,
167 out_dtype=out_dtype,
168 reduce_op=reduce_op,
169 orig_scatter_dim=orig_scatter_dim,
170 scatter_dim_after_maybe_reshape=scatter_dim_after_maybe_reshape,
171 group_name=group_name,
172 output_shape=output_shape,
173 )
174
175
176 def fused_all_gather_flashinfer_scaled_matmul_fake(
177 A_shard: torch.Tensor,
178 B: torch.Tensor,
179 A_scale: torch.Tensor,
180 B_scale: torch.Tensor,
181 gather_dim: int,
182 group_name: str,
183 out_dtype: torch.dtype | None = None,
184 ) -> torch.Tensor:
185 world_size = c10d._resolve_process_group(group_name).size()
186 output_shape = list(A_shard.shape)
187 output_shape[gather_dim] *= world_size
188 output_shape[-1] = B.shape[1]
189 return torch.empty(
190 output_shape,
191 dtype=out_dtype or torch.bfloat16,
192 device=A_shard.device,
193 )
194
195
196 def fused_all_gather_flashinfer_scaled_matmul(
197 A_shard: torch.Tensor,
198 B: torch.Tensor,
199 A_scale: torch.Tensor,
200 B_scale: torch.Tensor,
201 gather_dim: int,
202 group_name: str,
203 out_dtype: torch.dtype | None = None,
204 ) -> torch.Tensor:
205 assert gather_dim == 0, (
206 "FlashInfer symm_mem adapter currently only supports gather_dim=0"
207 )
208 _, outputs = torch.distributed._symmetric_memory._fused_all_gather_matmul_impl(
209 mm_out_op=_flashinfer_scaled_mm_out,
210 A_shard=A_shard,
211 Bs=[B],
212 A_scale=A_scale,
213 kwargs_list=[
214 {
215 "scale_b": B_scale,
216 "bias": None,
217 "scale_result": None,
218 "out_dtype": out_dtype,
219 "use_fast_accum": False,
220 }
221 ],
222 out_dtypes=[out_dtype],
223 gather_dim=gather_dim,
224 group_name=group_name,
225 return_A=False,
226 )
227 return outputs[0]
228
229
230 def fused_all_gather_flashinfer_fp4_matmul_fake(
231 A_shard: torch.Tensor,
232 B: torch.Tensor,
233 A_scale_shard: torch.Tensor,
234 B_scale: torch.Tensor,
235 alpha: torch.Tensor,
236 gather_dim: int,
237 group_name: str,
238 out_dtype: torch.dtype | None = None,
239 view_a_scale_as_fp8: bool = False,
240 use_8x4_sf_layout: bool = False,
241 backend: str = "cutlass",
242 ) -> torch.Tensor:
243 world_size = c10d._resolve_process_group(group_name).size()
244 output_shape = list(A_shard.shape)
245 output_shape[gather_dim] *= world_size
246 output_shape[-1] = B.shape[1]
247 return torch.empty(
248 output_shape,
249 dtype=out_dtype or torch.bfloat16,
250 device=A_shard.device,
251 )
252
253
254 def fused_all_gather_flashinfer_fp4_matmul(
255 A_shard: torch.Tensor,
256 B: torch.Tensor,
257 A_scale_shard: torch.Tensor,
258 B_scale: torch.Tensor,
259 alpha: torch.Tensor,
260 gather_dim: int,
261 group_name: str,
262 out_dtype: torch.dtype | None = None,
263 view_a_scale_as_fp8: bool = False,
264 use_8x4_sf_layout: bool = False,
265 backend: str = "cutlass",
266 ) -> torch.Tensor:
267 assert gather_dim == 0, (
268 "FlashInfer FP4 symm_mem adapter currently only supports gather_dim=0"
269 )
270 assert A_shard.ndim == 2 and A_scale_shard.ndim == 2 and B.ndim == 2, (
271 "FlashInfer FP4 symm_mem adapter expects 2D inputs"
272 )
273 if view_a_scale_as_fp8:
274 A_scale_shard = A_scale_shard.view(torch.float8_e4m3fn)
275
276 group = c10d._resolve_process_group(group_name)
277 world_size = group.size()
278 output = A_shard.new_empty(
279 A_shard.shape[0] * world_size,
280 B.shape[1],
281 dtype=out_dtype or torch.bfloat16,
282 )
283 output_shards = output.chunk(world_size)
284
285 A = A_shard.new_empty(A_shard.shape[0] * world_size, A_shard.shape[1])
286 A_scale = A_scale_shard.new_empty(
287 A_scale_shard.shape[0] * world_size,
288 A_scale_shard.shape[1],
289 )
290
291 def fp4_shard_consumer(shards: list[torch.Tensor], rank: int) -> None:
292 _flashinfer_fp4_mm_out(
293 shards[0],
294 B,
295 scale_a=shards[1],
296 scale_b=B_scale,
297 alpha=alpha,
298 out=output_shards[rank],
299 out_dtype=out_dtype,
300 use_8x4_sf_layout=use_8x4_sf_layout,
301 backend=backend,
302 )
303
304 torch.distributed._symmetric_memory._pipelined_multi_all_gather_and_consume(
305 [A_shard, A_scale_shard],
306 fp4_shard_consumer,
307 [A, A_scale],
308 group_name,
309 False,
310 )
311 return output
312
313
314 direct_register_custom_op(
315 op_name="fused_flashinfer_scaled_matmul_reduce_scatter",
316 op_func=fused_flashinfer_scaled_matmul_reduce_scatter,
317 fake_impl=fused_flashinfer_scaled_matmul_reduce_scatter_fake,
318 )
319
320 direct_register_custom_op(
321 op_name="fused_all_gather_flashinfer_scaled_matmul",
322 op_func=fused_all_gather_flashinfer_scaled_matmul,
323 fake_impl=fused_all_gather_flashinfer_scaled_matmul_fake,
324 )
325
326 direct_register_custom_op(
327 op_name="fused_all_gather_flashinfer_fp4_matmul",
328 op_func=fused_all_gather_flashinfer_fp4_matmul,
329 fake_impl=fused_all_gather_flashinfer_fp4_matmul_fake,
330 )
331
332
333 class BasePattern:
334 def __init__(self, dtype: torch.dtype, device: str | None) -> None:
335 self.dtype = dtype
336 self.device = device
337 self.tp = get_tp_group()
338 self.tp_size = get_tensor_model_parallel_world_size()
339
340
341 class GEMMReduceScatterPattern(BasePattern):
342 rocm_gemm = False
343
344 def get_inputs(self) -> list[torch.Tensor]:
345 mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
346 mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
347 return [mul, mm_weight]
348
349 def register(self, pm_pass: PatternMatcherPass) -> None:
350 def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
351 if self.rocm_gemm:
352 mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None)
353 else:
354 mm = torch.ops.aten.mm.default(mul, mm_weight)
355 reduce_scatter = torch.ops.vllm.reduce_scatter.default(
356 mm,
357 dim=0,
358 world_size=self.tp_size,
359 group_name=self.tp.unique_name,
360 )
361 return reduce_scatter
362
363 def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
364 gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
365 mul,
366 mm_weight.t() if self.rocm_gemm else mm_weight,
367 "sum",
368 scatter_dim=0,
369 group_name=self.tp.device_group.group_name,
370 )
371
372 return gemm_rs
373
374 pm.register_replacement(
375 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
376 )
377
378
379 class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern):
380 rocm_gemm = True
381
382
383 class AllGatherGEMMPattern(BasePattern):
384 rocm_gemm = False
385
386 def get_inputs(self) -> list[torch.Tensor]:
387 x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
388 weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
389
390 return [x, weight]
391
392 def register(self, pm_pass: PatternMatcherPass) -> None:
393 def pattern(
394 x: torch.Tensor,
395 weight: torch.Tensor,
396 ) -> torch.Tensor:
397 all_gather = torch.ops.vllm.all_gather.default(
398 x,
399 dim=0,
400 world_size=self.tp_size,
401 group_name=self.tp.unique_name,
402 )
403
404 if self.rocm_gemm:
405 return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None)
406 return torch.ops.aten.mm.default(all_gather, weight)
407
408 def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
409 ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
410 x,
411 [weight.t() if self.rocm_gemm else weight],
412 gather_dim=0,
413 group_name=self.tp.device_group.group_name,
414 )
415 return mm_outputs
416
417 pm.register_replacement(
418 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
419 )
420
421
422 class RocmAllGatherGEMMPattern(AllGatherGEMMPattern):
423 rocm_gemm = True
424
425
426 class ScaledMMReduceScatterPattern(BasePattern):
427 def get_inputs(self) -> list[torch.Tensor]:
428 input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
429 mm_weight = (
430 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
431 .contiguous()
432 .transpose(0, 1)
433 )
434 scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
435 scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
436 return [input, mm_weight, scale_a, scale_b]
437
438 def register(self, pm_pass: PatternMatcherPass) -> None:
439 def pattern(
440 input: torch.Tensor,
441 mat2: torch.Tensor,
442 scale_a: torch.Tensor,
443 scale_b: torch.Tensor,
444 ) -> torch.Tensor:
445 scaled_mm = torch.ops.aten._scaled_mm.default(
446 input,
447 mat2=mat2,
448 scale_a=scale_a,
449 scale_b=scale_b,
450 bias=None,
451 scale_result=None,
452 out_dtype=self.dtype,
453 )
454 reduce_scatter = torch.ops.vllm.reduce_scatter.default(
455 scaled_mm,
456 dim=0,
457 world_size=self.tp_size,
458 group_name=self.tp.unique_name,
459 )
460 return reduce_scatter
461
462 def replacement(
463 input: torch.Tensor,
464 mat2: torch.Tensor,
465 scale_a: torch.Tensor,
466 scale_b: torch.Tensor,
467 ) -> torch.Tensor:
468 # Calculate output shape: input @ mat2 with scatter_dim reduced
469 output_shape = [*input.shape[:-1], mat2.shape[1]]
470 scatter_dim = 0
471 gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
472 input,
473 mat2,
474 scale_a,
475 scale_b,
476 "sum",
477 scatter_dim, # orig_scatter_dim
478 scatter_dim, # scatter_dim_after_maybe_reshape
479 self.tp.device_group.group_name,
480 output_shape,
481 None, # bias
482 None, # result_scale
483 self.dtype, # out_dtype
484 False, # use_fast_accum
485 )
486
487 return gemm_rs
488
489 pm.register_replacement(
490 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
491 )
492
493
494 class AllGatherScaledMMPattern(BasePattern):
495 def get_inputs(self) -> list[torch.Tensor]:
496 x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
497 weight = (
498 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
499 .contiguous()
500 .transpose(0, 1)
501 )
502
503 s1 = x.shape[0] * self.tp_size
504
505 scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
506 scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
507
508 return [x, weight, scale_a, scale_b]
509
510 def register(self, pm_pass: PatternMatcherPass) -> None:
511 def pattern(
512 x: torch.Tensor,
513 weight: torch.Tensor,
514 scale_a: torch.Tensor,
515 scale_b: torch.Tensor,
516 ) -> torch.Tensor:
517 all_gather = torch.ops.vllm.all_gather.default(
518 x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
519 )
520
521 return torch.ops.aten._scaled_mm.default(
522 all_gather,
523 mat2=weight,
524 scale_a=scale_a,
525 scale_b=scale_b,
526 bias=None,
527 scale_result=None,
528 out_dtype=self.dtype,
529 )
530
531 def replacement(
532 x: torch.Tensor,
533 weight: torch.Tensor,
534 scale_a: torch.Tensor,
535 scale_b: torch.Tensor,
536 ) -> torch.Tensor:
537 ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
538 x,
539 [weight],
540 scale_a,
541 [scale_b],
542 gather_dim=0,
543 biases=[None],
544 result_scales=[None],
545 out_dtypes=[self.dtype],
546 use_fast_accum=[False],
547 group_name=self.tp.device_group.group_name,
548 )
549 return mm_outputs
550
551 pm.register_replacement(
552 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
553 )
554
555
556 class CutlassScaledMMReduceScatterPattern(BasePattern):
557 def get_inputs(self) -> list[torch.Tensor]:
558 input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
559 mm_weight = (
560 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
561 .contiguous()
562 .transpose(0, 1)
563 )
564 scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
565 scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
566
567 cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype)
568 return [input, mm_weight, scale_a, scale_b, cutlass_mm_output]
569
570 def register(self, pm_pass: PatternMatcherPass) -> None:
571 def pattern(
572 input: torch.Tensor,
573 weight: torch.Tensor,
574 scale_a: torch.Tensor,
575 scale_b: torch.Tensor,
576 cutlass_mm_output: torch.Tensor,
577 ) -> torch.Tensor:
578 cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
579 torch.ops._C.cutlass_scaled_mm.default,
580 out=cutlass_mm_output,
581 a=input,
582 b=weight,
583 a_scales=scale_a,
584 b_scales=scale_b,
585 bias=None,
586 )
587
588 reduce_scatter = torch.ops.vllm.reduce_scatter.default(
589 cutlass_scaled_mm[1],
590 dim=0,
591 world_size=self.tp_size,
592 group_name=self.tp.unique_name,
593 )
594 return reduce_scatter
595
596 def replacement(
597 input: torch.Tensor,
598 mat2: torch.Tensor,
599 scale_a: torch.Tensor,
600 scale_b: torch.Tensor,
601 cutlass_mm_output: torch.Tensor,
602 ) -> torch.Tensor:
603 # Calculate output shape: input @ mat2 with scatter_dim reduced
604 output_shape = [*input.shape[:-1], mat2.shape[1]]
605 scatter_dim = 0
606 gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
607 input,
608 mat2,
609 scale_a,
610 scale_b,
611 "sum",
612 scatter_dim, # orig_scatter_dim
613 scatter_dim, # scatter_dim_after_maybe_reshape
614 self.tp.device_group.group_name,
615 output_shape,
616 None, # bias
617 None, # result_scale
618 self.dtype, # out_dtype
619 False, # use_fast_accum
620 )
621
622 return gemm_rs
623
624 pm.register_replacement(
625 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
626 )
627
628
629 class AllGatherCutlassScaledMMPattern(BasePattern):
630 def get_inputs(self) -> list[torch.Tensor]:
631 x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
632 weight = (
633 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
634 .contiguous()
635 .transpose(0, 1)
636 )
637
638 s1 = x.shape[0] * self.tp_size
639
640 scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
641 scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
642
643 s2 = weight.shape[1]
644 output = torch.empty([s1, s2], device=self.device, dtype=self.dtype)
645
646 return [x, weight, scale_a, scale_b, output]
647
648 def register(self, pm_pass: PatternMatcherPass) -> None:
649 def pattern(
650 x: torch.Tensor,
651 weight: torch.Tensor,
652 scale_a: torch.Tensor,
653 scale_b: torch.Tensor,
654 output: torch.Tensor,
655 ) -> torch.Tensor:
656 all_gather = torch.ops.vllm.all_gather.default(
657 x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
658 )
659
660 cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
661 torch.ops._C.cutlass_scaled_mm.default,
662 out=output,
663 a=all_gather,
664 b=weight,
665 a_scales=scale_a,
666 b_scales=scale_b,
667 bias=None,
668 )
669 return cutlass_scaled_mm[1]
670
671 def replacement(
672 x: torch.Tensor,
673 weight: torch.Tensor,
674 scale_a: torch.Tensor,
675 scale_b: torch.Tensor,
676 output: torch.Tensor,
677 ) -> torch.Tensor:
678 ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
679 x,
680 [weight],
681 scale_a,
682 [scale_b],
683 gather_dim=0,
684 biases=[None],
685 result_scales=[None],
686 out_dtypes=[self.dtype],
687 use_fast_accum=[False],
688 group_name=self.tp.device_group.group_name,
689 )
690 return mm_outputs
691
692 pm.register_replacement(
693 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
694 )
695
696
697 class FlashInferBMMFP8ReduceScatterPattern(
698 BasePattern, VllmPatternReplacement[..., torch.Tensor]
699 ):
700 def get_inputs(self) -> list[torch.Tensor]:
701 a_2d = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
702 b_2d = (
703 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
704 .contiguous()
705 .transpose(0, 1)
706 )
707 a_scale = torch.empty([1], device=self.device, dtype=torch.float32)
708 b_scale = torch.empty([1], device=self.device, dtype=torch.float32)
709 return [a_2d, b_2d, a_scale, b_scale]
710
711 @property
712 def pattern(self) -> Callable[..., torch.Tensor]:
713 def _pattern(
714 a_2d: torch.Tensor,
715 b_2d: torch.Tensor,
716 a_scale: torch.Tensor,
717 b_scale: torch.Tensor,
718 ) -> torch.Tensor:
719 bmm = torch.ops.vllm.bmm_fp8.default(
720 torch.ops.aten.unsqueeze.default(a_2d, 0),
721 torch.ops.aten.unsqueeze.default(b_2d, 0),
722 a_scale,
723 b_scale,
724 self.dtype,
725 "auto",
726 )
727 output = torch.ops.aten.reshape.default(bmm, list(bmm.shape[1:]))
728 return torch.ops.vllm.reduce_scatter.default(
729 output,
730 dim=0,
731 world_size=self.tp_size,
732 group_name=self.tp.unique_name,
733 )
734
735 return _pattern
736
737 @property
738 def replacement(self) -> Callable[..., torch.Tensor]:
739 def _replacement(
740 a_2d: torch.Tensor,
741 b_2d: torch.Tensor,
742 a_scale: torch.Tensor,
743 b_scale: torch.Tensor,
744 ) -> torch.Tensor:
745 return torch.ops.vllm.fused_flashinfer_scaled_matmul_reduce_scatter.default(
746 a_2d,
747 b_2d,
748 a_scale,
749 b_scale,
750 "sum",
751 0,
752 0,
753 self.tp.device_group.group_name,
754 [a_2d.shape[0], b_2d.shape[1]],
755 self.dtype,
756 )
757
758 return _replacement
759
760
761 class FlashInferAllGatherBMMFP8Pattern(
762 BasePattern, VllmPatternReplacement[..., torch.Tensor]
763 ):
764 def get_inputs(self) -> list[torch.Tensor]:
765 a_shard_2d = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
766 b_2d = (
767 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
768 .contiguous()
769 .transpose(0, 1)
770 )
771 a_scale = torch.empty([1], device=self.device, dtype=torch.float32)
772 b_scale = torch.empty([1], device=self.device, dtype=torch.float32)
773 return [a_shard_2d, b_2d, a_scale, b_scale]
774
775 @property
776 def pattern(self) -> Callable[..., torch.Tensor]:
777 def _pattern(
778 a_shard_2d: torch.Tensor,
779 b_2d: torch.Tensor,
780 a_scale: torch.Tensor,
781 b_scale: torch.Tensor,
782 ) -> torch.Tensor:
783 all_gather = torch.ops.vllm.all_gather.default(
784 a_shard_2d,
785 dim=0,
786 world_size=self.tp_size,
787 group_name=self.tp.unique_name,
788 )
789 return torch.ops.vllm.bmm_fp8.default(
790 torch.ops.aten.unsqueeze.default(all_gather, 0),
791 torch.ops.aten.unsqueeze.default(b_2d, 0),
792 a_scale,
793 b_scale,
794 self.dtype,
795 "auto",
796 )
797
798 return _pattern
799
800 @property
801 def replacement(self) -> Callable[..., torch.Tensor]:
802 def _replacement(
803 a_shard_2d: torch.Tensor,
804 b_2d: torch.Tensor,
805 a_scale: torch.Tensor,
806 b_scale: torch.Tensor,
807 ) -> torch.Tensor:
808 fused = torch.ops.vllm.fused_all_gather_flashinfer_scaled_matmul.default(
809 a_shard_2d,
810 b_2d,
811 a_scale,
812 b_scale,
813 0,
814 self.tp.device_group.group_name,
815 self.dtype,
816 )
817 return torch.ops.aten.unsqueeze.default(fused, 0)
818
819 return _replacement
820
821
822 class FlashInferAllGatherFP4Pattern(
823 BasePattern, VllmPatternReplacement[..., torch.Tensor]
824 ):
825 def __init__(
826 self,
827 dtype: torch.dtype,
828 device: str | None,
829 backend: str,
830 use_8x4_sf_layout: bool,
831 a_scale_view: str,
832 ) -> None:
833 super().__init__(dtype, device)
834 self.backend = backend
835 self.use_8x4_sf_layout = use_8x4_sf_layout
836 self.a_scale_view = a_scale_view
837
838 def get_inputs(self) -> list[torch.Tensor]:
839 a_shard_2d = torch.empty([8, 8], device=self.device, dtype=torch.uint8)
840 b_2d = torch.empty([8, 16], device=self.device, dtype=torch.uint8)
841 a_scale_shard = torch.empty([128, 4], device=self.device, dtype=torch.int32)
842 b_scale = torch.empty([4, 128], device=self.device, dtype=torch.uint8)
843 alpha = torch.empty([], device=self.device, dtype=torch.float32)
844 return [
845 a_shard_2d,
846 b_2d,
847 a_scale_shard,
848 b_scale,
849 alpha,
850 ]
851
852 @property
853 def pattern(self) -> Callable[..., torch.Tensor]:
854 def _pattern(
855 a_shard_2d: torch.Tensor,
856 b_2d: torch.Tensor,
857 a_scale_shard: torch.Tensor,
858 b_scale: torch.Tensor,
859 alpha: torch.Tensor,
860 ) -> torch.Tensor:
861 all_gather_a = torch.ops.vllm.all_gather.default(
862 a_shard_2d,
863 dim=0,
864 world_size=self.tp_size,
865 group_name=self.tp.unique_name,
866 )
867 all_gather_a_scale = torch.ops.vllm.all_gather.default(
868 a_scale_shard,
869 dim=0,
870 world_size=self.tp_size,
871 group_name=self.tp.unique_name,
872 )
873 a_scale = all_gather_a_scale
874 if self.a_scale_view in ("float8", "float8_uint8"):
875 a_scale = torch.ops.aten.view.dtype(a_scale, torch.float8_e4m3fn)
876 if self.a_scale_view in ("uint8", "float8_uint8"):
877 a_scale = torch.ops.aten.view.dtype(a_scale, torch.uint8)
878 return torch.ops.vllm.flashinfer_mm_fp4.default(
879 all_gather_a,
880 b_2d,
881 a_scale,
882 b_scale,
883 alpha,
884 self.dtype,
885 self.use_8x4_sf_layout,
886 self.backend,
887 )
888
889 return _pattern
890
891 @property
892 def replacement(self) -> Callable[..., torch.Tensor]:
893 def _replacement(
894 a_shard_2d: torch.Tensor,
895 b_2d: torch.Tensor,
896 a_scale_shard: torch.Tensor,
897 b_scale: torch.Tensor,
898 alpha: torch.Tensor,
899 ) -> torch.Tensor:
900 return torch.ops.vllm.fused_all_gather_flashinfer_fp4_matmul.default(
901 a_shard_2d,
902 b_2d,
903 a_scale_shard,
904 b_scale,
905 alpha,
906 0,
907 self.tp.device_group.group_name,
908 self.dtype,
909 self.a_scale_view in ("float8", "float8_uint8"),
910 self.use_8x4_sf_layout,
911 self.backend,
912 )
913
914 return _replacement
915
916
917 class AsyncTPPass(VllmFusionPatternMatcherPass):
918 @enable_fake_mode
919 def __init__(self, config: VllmConfig) -> None:
920 super().__init__(config, pass_name="async_tp_pass")
921
922 enable_symm_mem_for_group(get_tp_group().device_group.group_name)
923 GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass)
924
925 AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass)
926 if current_platform.is_rocm() and self.model_dtype == torch.bfloat16:
927 RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register(
928 self.pm_pass
929 )
930 RocmAllGatherGEMMPattern(self.model_dtype, self.device).register(
931 self.pm_pass
932 )
933
934 # These fusions are enabled only for bfloat16 models because
935 # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
936 # only supports bfloat16 as the output dtype.
937 if self.model_dtype == torch.bfloat16:
938 ScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
939 self.pm_pass
940 )
941 AllGatherScaledMMPattern(self.model_dtype, self.device).register(
942 self.pm_pass
943 )
944 if hasattr(torch.ops._C, "cutlass_scaled_mm"):
945 CutlassScaledMMReduceScatterPattern(
946 self.model_dtype, self.device
947 ).register(self.pm_pass)
948 AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
949 self.pm_pass
950 )
951 with suppress(ImportError):
952 import vllm.utils.flashinfer # noqa: F401
953 if hasattr(torch.ops.vllm, "bmm_fp8"):
954 self.register(
955 FlashInferAllGatherBMMFP8Pattern(self.model_dtype, self.device)
956 )
957 self.register(
958 FlashInferBMMFP8ReduceScatterPattern(self.model_dtype, self.device)
959 )
960 if hasattr(torch.ops.vllm, "flashinfer_mm_fp4"):
961 for backend in ("cutlass", "cudnn"):
962 for a_scale_view in ("float8_uint8", "uint8"):
963 self.register(
964 FlashInferAllGatherFP4Pattern(
965 self.model_dtype,
966 self.device,
967 backend,
968 use_8x4_sf_layout=False,
969 a_scale_view=a_scale_view,
970 )
971 )
972 for use_8x4_sf_layout in (False, True):
973 for a_scale_view in ("float8",):
974 self.register(
975 FlashInferAllGatherFP4Pattern(
976 self.model_dtype,
977 self.device,
978 "trtllm",
979 use_8x4_sf_layout=use_8x4_sf_layout,
980 a_scale_view=a_scale_view,
981 )
982 )
983 self.register(
984 FlashInferAllGatherFP4Pattern(
985 self.model_dtype,
986 self.device,
987 "cute-dsl",
988 use_8x4_sf_layout=False,
989 a_scale_view="float8",
990 )
991 )
992 # NVFP4 reduce-scatter does not need scale communication: FP4
993 # scales are consumed by the local GEMM and only BF16 partial
994 # outputs are reduced. Keep this PR scoped to the all-gather
995 # path; reduce-scatter needs a dedicated FP4 producer rather
996 # than the existing FP8-style helper.
997
998 self.dump_patterns(config, self.pm_pass)
999
1000 def is_applicable_for_range(self, compile_range: Range) -> bool:
1001 # This pass is applied on top of the sequence parallelism pass,
1002 # which is only supported in fullgraph compilation mode.
1003 assert (
1004 self.compilation_config.use_inductor_graph_partition
1005 or not self.compilation_config.splitting_ops
1006 ), "AsyncTPPass requires full-graph compilation"
1007 return True
1008
1009 @VllmInductorPass.time_and_log
1010 def __call__(self, graph: fx.Graph) -> None:
1011 self.matched_count = self.pm_pass.apply(graph)
1012 VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
1013 logger.debug("Replaced %s patterns", self.matched_count)