collective_fusion.py 43.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from importlib.util import find_spec
4
5
6
7

import torch
import torch._inductor.pattern_matcher as pm
import torch.fx as fx
8
from torch._higher_order_ops.auto_functionalize import auto_functionalized
9
10
11
12
from torch._inductor.pattern_matcher import PatternMatcherPass
from torch.distributed._symmetric_memory import enable_symm_mem_for_group

from vllm.config import VllmConfig
13
from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
14
from vllm.distributed.parallel_state import (
15
16
17
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
18
from vllm.logger import init_logger
19
20
21
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    kFp8StaticTensorSym,
)
22
from vllm.platforms import current_platform
23
from vllm.utils.torch_utils import direct_register_custom_op
24

25
from .inductor_pass import enable_fake_mode
26
from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
27
from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
28

29
30
FP8_DTYPE = current_platform.fp8_dtype()

31
if find_spec("flashinfer"):
32
33
    try:
        import flashinfer.comm as flashinfer_comm
34
35
36
37
38
39

        flashinfer_comm = (
            flashinfer_comm
            if hasattr(flashinfer_comm, "trtllm_allreduce_fusion")
            else None
        )
40
41
    except ImportError:
        flashinfer_comm = None
42
43
44
else:
    flashinfer_comm = None

45
46
logger = init_logger(__name__)

47
48
if hasattr(torch.ops._C, "scaled_fp4_quant"):
    STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
49

50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

class BasePattern:
    def __init__(self, dtype: torch.dtype, device: str):
        self.dtype = dtype
        self.device = device
        self.tp = get_tp_group()
        self.tp_size = get_tensor_model_parallel_world_size()


class GEMMReduceScatterPattern(BasePattern):
    def get_inputs(self):
        mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
        mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
        return [mul, mm_weight]

    def register(self, pm_pass: PatternMatcherPass):
        def pattern(mul: torch.Tensor, mm_weight: torch.Tensor):
            mm = torch.ops.aten.mm.default(mul, mm_weight)
            reduce_scatter = torch.ops.vllm.reduce_scatter.default(
                mm,
                dim=0,
                world_size=self.tp_size,
72
73
                group_name=self.tp.unique_name,
            )
74
75
76
77
78
79
80
81
82
83
84
85
86
            return reduce_scatter

        def replacement(mul: torch.Tensor, mm_weight: torch.Tensor):
            gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
                mul,
                mm_weight,
                "avg",
                scatter_dim=0,
                group_name=self.tp.device_group.group_name,
            )

            return gemm_rs

87
88
89
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107


class AllGatherGEMMPattern(BasePattern):
    def get_inputs(self):
        x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
        weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)

        return [x, weight]

    def register(self, pm_pass: PatternMatcherPass):
        def pattern(
            x: torch.Tensor,
            weight: torch.Tensor,
        ) -> tuple[torch.Tensor, torch.Tensor]:
            all_gather = torch.ops.vllm.all_gather.default(
                x,
                dim=0,
                world_size=self.tp_size,
108
109
                group_name=self.tp.unique_name,
            )
110
111
112
113

            return torch.ops.aten.mm.default(all_gather, weight)

        def replacement(
114
115
            x: torch.Tensor, weight: torch.Tensor
        ) -> tuple[torch.Tensor, torch.Tensor]:
116
117
118
119
120
121
122
123
            ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
                x,
                [weight],
                gather_dim=0,
                group_name=self.tp.device_group.group_name,
            )
            return mm_outputs

124
125
126
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
127
128


129
130
131
class ScaledMMReduceScatterPattern(BasePattern):
    def get_inputs(self):
        input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
132
133
134
135
136
        mm_weight = (
            torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
            .contiguous()
            .transpose(0, 1)
        )
137
138
139
140
141
        scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
        scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
        return [input, mm_weight, scale_a, scale_b]

    def register(self, pm_pass: PatternMatcherPass):
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
        def pattern(
            input: torch.Tensor,
            mat2: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
        ) -> torch.Tensor:
            scaled_mm = torch.ops.aten._scaled_mm.default(
                input,
                mat2=mat2,
                scale_a=scale_a,
                scale_b=scale_b,
                bias=None,
                scale_result=None,
                out_dtype=self.dtype,
            )
157
158
159
160
            reduce_scatter = torch.ops.vllm.reduce_scatter.default(
                scaled_mm,
                dim=0,
                world_size=self.tp_size,
161
162
                group_name=self.tp.unique_name,
            )
163
164
            return reduce_scatter

165
166
167
168
169
170
        def replacement(
            input: torch.Tensor,
            mat2: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
        ) -> torch.Tensor:
171
172
173
            # Calculate output shape: input @ mat2 with scatter_dim reduced
            output_shape = [*input.shape[:-1], mat2.shape[1]]
            scatter_dim = 0
174
            gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
175
176
177
178
179
                input,
                mat2,
                scale_a,
                scale_b,
                "avg",
180
181
182
183
184
185
186
187
                scatter_dim,  # orig_scatter_dim
                scatter_dim,  # scatter_dim_after_maybe_reshape
                self.tp.device_group.group_name,
                output_shape,
                None,  # bias
                None,  # result_scale
                self.dtype,  # out_dtype
                False,  # use_fast_accum
188
189
190
191
            )

            return gemm_rs

192
193
194
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
195
196
197
198
199


class AllGatherScaledMMPattern(BasePattern):
    def get_inputs(self):
        x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
200
201
202
203
204
        weight = (
            torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
            .contiguous()
            .transpose(0, 1)
        )
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220

        s1 = x.shape[0] * self.tp_size

        scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
        scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)

        return [x, weight, scale_a, scale_b]

    def register(self, pm_pass: PatternMatcherPass):
        def pattern(
            x: torch.Tensor,
            weight: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
        ) -> torch.Tensor:
            all_gather = torch.ops.vllm.all_gather.default(
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
                x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
            )

            return torch.ops.aten._scaled_mm.default(
                all_gather,
                mat2=weight,
                scale_a=scale_a,
                scale_b=scale_b,
                bias=None,
                scale_result=None,
                out_dtype=self.dtype,
            )

        def replacement(
            x: torch.Tensor,
            weight: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
        ) -> torch.Tensor:
240
241
242
243
244
245
246
247
248
249
250
251
252
253
            ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul(  # noqa
                x,
                [weight],
                scale_a,
                [scale_b],
                gather_dim=0,
                biases=[None],
                result_scales=[None],
                out_dtypes=[self.dtype],
                use_fast_accum=[False],
                group_name=self.tp.device_group.group_name,
            )
            return mm_outputs

254
255
256
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
257
258
259
260
261


class CutlassScaledMMReduceScatterPattern(BasePattern):
    def get_inputs(self):
        input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
262
263
264
265
266
        mm_weight = (
            torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
            .contiguous()
            .transpose(0, 1)
        )
267
268
269
        scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
        scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)

270
        cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype)
271
272
273
        return [input, mm_weight, scale_a, scale_b, cutlass_mm_output]

    def register(self, pm_pass: PatternMatcherPass):
274
275
276
277
278
279
280
        def pattern(
            input: torch.Tensor,
            weight: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
            cutlass_mm_output: torch.Tensor,
        ) -> torch.Tensor:
281
282
283
284
285
286
287
            cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
                torch.ops._C.cutlass_scaled_mm.default,
                out=cutlass_mm_output,
                a=input,
                b=weight,
                a_scales=scale_a,
                b_scales=scale_b,
288
289
                bias=None,
            )
290
291
292
293
294

            reduce_scatter = torch.ops.vllm.reduce_scatter.default(
                cutlass_scaled_mm[1],
                dim=0,
                world_size=self.tp_size,
295
296
                group_name=self.tp.unique_name,
            )
297
298
            return reduce_scatter

299
300
301
302
303
304
305
        def replacement(
            input: torch.Tensor,
            mat2: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
            cutlass_mm_output: torch.Tensor,
        ) -> torch.Tensor:
306
307
308
            # Calculate output shape: input @ mat2 with scatter_dim reduced
            output_shape = [*input.shape[:-1], mat2.shape[1]]
            scatter_dim = 0
309
            gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
310
311
312
313
314
                input,
                mat2,
                scale_a,
                scale_b,
                "avg",
315
316
317
318
319
320
321
322
                scatter_dim,  # orig_scatter_dim
                scatter_dim,  # scatter_dim_after_maybe_reshape
                self.tp.device_group.group_name,
                output_shape,
                None,  # bias
                None,  # result_scale
                self.dtype,  # out_dtype
                False,  # use_fast_accum
323
324
325
326
            )

            return gemm_rs

327
328
329
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
330
331
332
333
334


class AllGatherCutlassScaledMMPattern(BasePattern):
    def get_inputs(self):
        x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
335
336
337
338
339
        weight = (
            torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
            .contiguous()
            .transpose(0, 1)
        )
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359

        s1 = x.shape[0] * self.tp_size

        scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
        scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)

        s2 = weight.shape[1]
        output = torch.empty([s1, s2], device=self.device, dtype=self.dtype)

        return [x, weight, scale_a, scale_b, output]

    def register(self, pm_pass: PatternMatcherPass):
        def pattern(
            x: torch.Tensor,
            weight: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
            output: torch.Tensor,
        ) -> torch.Tensor:
            all_gather = torch.ops.vllm.all_gather.default(
360
361
                x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
            )
362
363
364
365
366
367
368
369

            cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
                torch.ops._C.cutlass_scaled_mm.default,
                out=output,
                a=all_gather,
                b=weight,
                a_scales=scale_a,
                b_scales=scale_b,
370
371
                bias=None,
            )
372
373
            return cutlass_scaled_mm[1]

374
375
376
377
378
379
380
        def replacement(
            x: torch.Tensor,
            weight: torch.Tensor,
            scale_a: torch.Tensor,
            scale_b: torch.Tensor,
            output: torch.Tensor,
        ) -> torch.Tensor:
381
382
383
384
385
386
387
388
389
390
391
392
393
394
            ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul(  # noqa
                x,
                [weight],
                scale_a,
                [scale_b],
                gather_dim=0,
                biases=[None],
                result_scales=[None],
                out_dtypes=[self.dtype],
                use_fast_accum=[False],
                group_name=self.tp.device_group.group_name,
            )
            return mm_outputs

395
396
397
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
398
399


400
class AsyncTPPass(VllmPatternMatcherPass):
401
    @enable_fake_mode
402
403
404
405
406
407
    def __init__(self, config: VllmConfig):
        super().__init__(config)

        # Enable symmetric memory for the TP process group
        enable_symm_mem_for_group(get_tp_group().device_group.group_name)
        self.patterns: PatternMatcherPass = PatternMatcherPass(
408
409
410
            pass_name="async_tp_pass"
        )
        GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.patterns)
411

412
        AllGatherGEMMPattern(self.model_dtype, self.device).register(self.patterns)
413

414
415
416
417
        # These fusions are enabled only for bfloat16 models because
        # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
        # only supports bfloat16 as the output dtype.
        if self.model_dtype == torch.bfloat16:
418
419
420
421
422
423
            ScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
                self.patterns
            )
            AllGatherScaledMMPattern(self.model_dtype, self.device).register(
                self.patterns
            )
424

425
426
427
428
429
430
            CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
                self.patterns
            )
            AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
                self.patterns
            )
431

432
433
        self.dump_patterns(config, self.patterns)

434
435
436
437
438
439
440
441
442
    def is_applicable(self, shape: int | None) -> bool:
        # This pass is applied on top of the sequence parallelism pass.
        # It inherits the same applicability condition as `SequenceParallelismPass`.
        # See `SequenceParallelismPass.is_applicable` for more details.
        if (
            not self.compilation_config.splitting_ops
            or self.compilation_config.use_inductor_graph_partition
        ):
            return True
443
444
445
        tp_size = get_tensor_model_parallel_world_size()
        return shape is not None and shape % tp_size == 0

446
    @VllmInductorPass.time_and_log
447
    def __call__(self, graph: fx.Graph):
448
449
        self.matched_count = self.patterns.apply(graph)
        logger.debug("Replaced %s patterns", self.matched_count)
450
451


452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
# Max size of the input tensor per world size per device capability
# to use flashinfer fused allreduce
FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = {
    90: {
        2: 64,  # 64MB
        4: 2,  # 2MB
        8: 0.5,  # 0.5MB
    },
    100: {
        2: 64,  # 64MB
        4: 32,  # 32MB
        8: 1,  # 1MB
    },
}

# Max size of the input tensor per world size per device capability
# to use flashinfer one shot fused allreduce
# OneShot max size is at most 64MB / world size (FlashInfer restriction)
_FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {
    90: {
        2: 32,  # 32MB
        4: 2,  # 2MB
        8: 0.5,  # 0.5MB
    },
    100: {
        2: 32,  # 32MB
        4: 4,  # 4MB
        8: 1,  # 1MB
    },
}


484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
if flashinfer_comm is not None:
    _FI_WORKSPACE_TENSOR = None
    MiB = 1024 * 1024

    def call_trtllm_fused_allreduce_norm(
        allreduce_in: torch.Tensor,
        residual: torch.Tensor,
        rms_gamma: torch.Tensor,
        rms_eps: float,
        world_rank: int,
        world_size: int,
        launch_with_pdl: bool,
        trigger_completion_at_end: bool,
        fp32_acc: bool,
        max_token_num: int,
499
        pattern_code: int,
500
501
502
503
        norm_out: torch.Tensor | None = None,
        quant_out: torch.Tensor | None = None,
        scale_out: torch.Tensor | None = None,
        scale_factor: torch.Tensor | None = None,
504
    ) -> None:
505
506
507
        num_tokens, hidden_size = allreduce_in.shape
        element_size = allreduce_in.element_size()
        current_tensor_size = num_tokens * hidden_size * element_size
508
509
510
511
512
513
514
515
516
517
518
519
520
521

        if num_tokens <= max_token_num:
            device_capability = current_platform.get_device_capability().to_int()
            # Get one shot input size limit for the current world size
            # for the current device capability
            max_one_shot_size_mb = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
                device_capability, {}
            ).get(world_size, None)
            # Use one shot if no max size for one shot is specified
            use_oneshot = (
                max_one_shot_size_mb is None
                or current_tensor_size <= max_one_shot_size_mb * MiB
            )

522
523
524
            assert _FI_WORKSPACE_TENSOR is not None, (
                "Flashinfer must be enabled when using flashinfer"
            )
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
            if norm_out is None:
                norm_out = allreduce_in
                residual_out = residual
            else:
                # return residual_out as allreduce_out with zeroed residual_in
                # as flashinfer does not support rms_norm
                # and allreduce_out together
                residual_out = allreduce_in
            # For the sizes that are smaller than the max size,
            # we only use flashinfer one shot allreduce
            flashinfer_comm.trtllm_allreduce_fusion(
                allreduce_in=allreduce_in,
                token_num=allreduce_in.shape[0],
                residual_in=residual,
                residual_out=residual_out,
                norm_out=norm_out,
                rms_gamma=rms_gamma,
                rms_eps=rms_eps,
                world_rank=world_rank,
                world_size=world_size,
                hidden_dim=allreduce_in.shape[-1],
                workspace_ptrs=_FI_WORKSPACE_TENSOR,
                launch_with_pdl=launch_with_pdl,
548
                use_oneshot=use_oneshot,
549
550
                trigger_completion_at_end=trigger_completion_at_end,
                fp32_acc=fp32_acc,
551
                pattern_code=pattern_code,
552
                allreduce_out=None,
553
554
555
                quant_out=quant_out,
                scale_out=scale_out,
                # in vllm we only support swizzled layout
556
                layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
557
                scale_factor=scale_factor,
558
559
560
            )
        else:
            allreduce_out = tensor_model_parallel_all_reduce(allreduce_in)
561
            if scale_factor is not None and scale_out is None:
562
563
564
                # Do fused rms norm static fp8 quant fused op
                if norm_out is None:
                    torch.ops._C.fused_add_rms_norm_static_fp8_quant(
565
566
567
568
569
570
571
                        quant_out,
                        allreduce_out,
                        residual,
                        rms_gamma,
                        scale_factor,
                        rms_eps,
                    )
572
573
                else:
                    torch.ops._C.rms_norm_static_fp8_quant(
574
575
                        quant_out, allreduce_out, rms_gamma, scale_factor, rms_eps
                    )
576
            else:
577
                if norm_out is None:
578
579
580
                    torch.ops._C.fused_add_rms_norm(
                        allreduce_out, residual, rms_gamma, rms_eps
                    )
581
582
                    norm_out = allreduce_out
                else:
583
                    torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma, rms_eps)
584
585
586
587
                if scale_factor is not None and scale_out is not None:
                    torch.ops._C.scaled_fp4_quant(
                        quant_out, norm_out, scale_out, scale_factor
                    )
588
            if scale_factor is None or norm_out is not None:
589
                # we need to return allreduce output
590
591
592
                # in cases of non quant fused AR + RMS norm
                # and fused AR + RMS norm + quant without fused add
                allreduce_in.copy_(allreduce_out)
593
594

    def call_trtllm_fused_allreduce_norm_fake(
595
596
597
598
599
600
601
602
603
604
605
        allreduce_in: torch.Tensor,
        residual: torch.Tensor,
        rms_gamma: torch.Tensor,
        rms_eps: float,
        world_rank: int,
        world_size: int,
        launch_with_pdl: bool,
        trigger_completion_at_end: bool,
        fp32_acc: bool,
        max_token_num: int,
        pattern_code: int,
606
607
608
609
        norm_out: torch.Tensor | None = None,
        quant_out: torch.Tensor | None = None,
        scale_out: torch.Tensor | None = None,
        scale_factor: torch.Tensor | None = None,
610
    ) -> None:
611
612
613
614
615
616
617
618
619
        pass

    direct_register_custom_op(
        op_name="flashinfer_trtllm_fused_allreduce_norm",
        op_func=call_trtllm_fused_allreduce_norm,
        mutates_args=[
            "allreduce_in",
            "residual",
            "norm_out",
620
621
            "quant_out",
            "scale_out",
622
623
624
625
        ],
        fake_impl=call_trtllm_fused_allreduce_norm_fake,
    )
    flashinfer_trtllm_fused_allreduce_norm = (
626
627
        torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
    )
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658


class FlashInferFusedAllReduceParams:
    """Parameters for FlashInfer fused allreduce operations."""

    def __init__(
        self,
        rank: int,
        world_size: int,
        use_fp32_lamport: bool = False,
        max_token_num: int = 1024,
    ):
        self.rank = rank
        self.world_size = world_size
        self.use_fp32_lamport = use_fp32_lamport
        self.trigger_completion_at_end = True
        self.launch_with_pdl = True
        self.fp32_acc = True
        self.max_token_num = max_token_num

    def get_trtllm_fused_allreduce_kwargs(self):
        return {
            "world_rank": self.rank,
            "world_size": self.world_size,
            "launch_with_pdl": self.launch_with_pdl,
            "trigger_completion_at_end": self.trigger_completion_at_end,
            "fp32_acc": self.fp32_acc,
            "max_token_num": self.max_token_num,
        }


659
660
class AllReduceRMSNormPattern(BasePattern):
    """
661
    This pattern replaces the allreduce + rms norm (without residual)
662
663
664
    with fused flashinfer implementation.
    Applies to allreduce + rmsnorm before attn in the first Transformer block.
    """
665
666
667
668
669
670
671
672
673
674
675

    def __init__(
        self,
        epsilon: float,
        dtype: torch.dtype,
        device: str,
        allreduce_params: FlashInferFusedAllReduceParams,
    ):
        super().__init__(dtype, device)
        self.epsilon = epsilon
        self.allreduce_params = allreduce_params
676
        self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
677
678

    def get_inputs(self):
679
        input, weight = self.rmsnorm_matcher.inputs()
680

681
682
        # input goes through allreduce first, always 16-bit
        return [input.to(self.dtype), weight]
683
684

    def register(self, pm_pass: PatternMatcherPass):
685
        def pattern(input: torch.Tensor, weight: torch.Tensor):
686
            allreduce_output = tensor_model_parallel_all_reduce(input)
687
            rms = self.rmsnorm_matcher(allreduce_output, weight)
688

689
690
691
            return rms, allreduce_output

        def replacement(input: torch.Tensor, weight: torch.Tensor):
692
            residual = torch.zeros_like(input)
693
            rms_result = torch.empty_like(input)
694
            allreduce = auto_functionalized(
695
                flashinfer_trtllm_fused_allreduce_norm,
696
697
698
                allreduce_in=input,
                residual=residual,
                norm_out=rms_result,
699
700
                quant_out=None,
                scale_out=None,
701
702
                rms_gamma=weight,
                rms_eps=self.epsilon,
703
                pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
704
705
                **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
            )
706
            # rms_result, allreduce_in
707
708
            return allreduce[3], allreduce[1]

709
710
711
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
712
713
714


class AllReduceFusedAddRMSNormPattern(BasePattern):
715
    """
716
    This pattern replaces the allreduce + rms norm (with residual)
717
718
719
    with fused flashinfer implementation.
    Applies to o_proj + rmsnorm after attn and mlp + rmsnorm before attn.
    """
720
721
722
723
724
725
726
727
728
729
730

    def __init__(
        self,
        epsilon: float,
        dtype: torch.dtype,
        device: str,
        allreduce_params: FlashInferFusedAllReduceParams,
    ):
        super().__init__(dtype, device)
        self.epsilon = epsilon
        self.allreduce_params = allreduce_params
731
        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
732
733

    def get_inputs(self):
734
735
736
737
        input, residual, weight = self.rmsnorm_matcher.inputs()

        # input goes through allreduce first, always 16-bit
        return [residual, input.to(self.dtype), weight]
738
739

    def register(self, pm_pass: PatternMatcherPass):
740
        def pattern(residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor):
741
            allreduce_output = tensor_model_parallel_all_reduce(input)
742
743
            rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
            return rms, residual
744

745
746
747
        def replacement(
            residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
        ):
748
            allreduce = auto_functionalized(
749
                flashinfer_trtllm_fused_allreduce_norm,
750
751
                allreduce_in=input,
                residual=residual,
752
753
754
                norm_out=None,
                quant_out=None,
                scale_out=None,
755
756
                rms_gamma=weight,
                rms_eps=self.epsilon,
757
                pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
758
759
                **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
            )
760
            # allreduce_in, residual
761
762
            return allreduce[1], allreduce[2]

763
764
765
        pm.register_replacement(
            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
        )
766

767
768
769
770
771
772
773
774
775
776
777
778
        # Same pattern, but only return the output and not residual
        # (helpful for end of graph where residual is not used again)
        first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0]

        pm.register_replacement(
            first_return_only(pattern),
            first_return_only(replacement),
            self.get_inputs(),
            pm.fwd_only,
            pm_pass,
        )

779

780
781
class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
    """
782
    This pattern replaces the allreduce + rms norm (without residual)
783
    + static fp8 quant with fused flashinfer implementation.
784
    Applies to allreduce + rmsnorm + quant before attn
785
786
787
    in the first Transformer block.
    """

788
789
790
791
792
793
794
    def __init__(
        self,
        epsilon: float,
        dtype: torch.dtype,
        device: str,
        allreduce_params: FlashInferFusedAllReduceParams,
    ):
795
796
797
798
        super().__init__(dtype, device)
        self.epsilon = epsilon
        self.allreduce_params = allreduce_params
        self.quant_dtype = torch.float8_e4m3fn
799
800
        self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
        self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
801
802
803

    def register(self, pm_pass: PatternMatcherPass):
        def get_inputs():
804
805
806
807
808
            input, weight = self.rmsnorm_matcher.inputs()
            _, scale = self.quant_matcher.inputs()

            # input goes through allreduce first, always 16-bit
            return [input.to(self.dtype), weight, scale]
809
810
811
812
813
814
815

        def pattern(
            input: torch.Tensor,
            weight: torch.Tensor,
            scale: torch.Tensor,
        ):
            all_reduce = tensor_model_parallel_all_reduce(input)
816
817
818
            rms = self.rmsnorm_matcher(all_reduce, weight)
            quant, _ = self.quant_matcher(rms, scale)
            return quant, all_reduce
819

820
        def replacement(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):
821
            residual = torch.zeros_like(input)
822
823
            result_rms = torch.empty_like(input)
            result_quant = torch.empty_like(input, dtype=self.quant_dtype)
824
825
826
827
828
            allreduce = auto_functionalized(
                flashinfer_trtllm_fused_allreduce_norm,
                allreduce_in=input,
                residual=residual,
                norm_out=result_rms,
829
                quant_out=result_quant,
830
831
832
                scale_out=None,
                rms_gamma=weight,
                rms_eps=self.epsilon,
833
834
835
836
                # We don't use norm_out afterwards
                pattern_code=(
                    flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant
                ),
837
838
839
840
841
842
843
                scale_factor=scale,
                **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
            )

            # quant_out, allreduce_output
            return allreduce[4], allreduce[1]

844
845
846
        pm.register_replacement(
            pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
        )
847
848
849
850
851
852


class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
    """
    This pattern replaces the allreduce + rms norm (with residual)
    + static fp8 quant with fused flashinfer implementation.
853
    Applies to o_proj + rmsnorm after attn + quant and
854
855
856
    mlp + rmsnorm + quant before attn.
    """

857
858
859
860
861
862
863
    def __init__(
        self,
        epsilon: float,
        dtype: torch.dtype,
        device: str,
        allreduce_params: FlashInferFusedAllReduceParams,
    ):
864
865
866
867
868
        super().__init__(dtype, device)
        self.epsilon = epsilon
        self.allreduce_params = allreduce_params
        self.quant_dtype = torch.float8_e4m3fn

869
870
871
        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
        self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)

872
873
    def register(self, pm_pass: PatternMatcherPass):
        def get_inputs():
874
875
            input, residual, weight = self.rmsnorm_matcher.inputs()
            _, scale = self.quant_matcher.inputs()
876

877
878
            # input goes through allreduce first, always 16-bit
            return [residual, input.to(self.dtype), weight, scale]
879
880
881
882
883
884
885
886

        def pattern(
            residual: torch.Tensor,
            input: torch.Tensor,
            weight: torch.Tensor,
            scale: torch.Tensor,
        ):
            allreduce_output = tensor_model_parallel_all_reduce(input)
887
888
            rms, res = self.rmsnorm_matcher(allreduce_output, weight, residual)
            quant, _ = self.quant_matcher(rms, scale)
889

890
            return quant, res
891

892
893
894
895
896
897
        def replacement(
            residual: torch.Tensor,
            input: torch.Tensor,
            weight: torch.Tensor,
            scale: torch.Tensor,
        ):
898
            result_quant = torch.empty_like(input, dtype=self.quant_dtype)
899
900
901
902
903
            allreduce = auto_functionalized(
                flashinfer_trtllm_fused_allreduce_norm,
                allreduce_in=input,
                residual=residual,
                norm_out=None,
904
                quant_out=result_quant,
905
906
907
                scale_out=None,
                rms_gamma=weight,
                rms_eps=self.epsilon,
908
909
910
911
                # We don't use norm_out afterwards
                pattern_code=(
                    flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant
                ),
912
913
914
                scale_factor=scale,
                **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
            )
915
            # quant_out, rms_norm_residual
916
917
            return allreduce[4], allreduce[2]

918
919
920
        pm.register_replacement(
            pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
        )
921
922
923
924


class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
    """
925
    This pattern replaces the allreduce + rms norm (without residual)
926
    + static nvfp4 quant with fused flashinfer implementation.
927
    Applies to allreduce + rmsnorm + quant before attn
928
929
930
    in the first Transformer block.
    """

931
932
933
934
935
936
937
    def __init__(
        self,
        epsilon: float,
        dtype: torch.dtype,
        device: str,
        allreduce_params: FlashInferFusedAllReduceParams,
    ):
938
939
940
        super().__init__(dtype, device)
        self.epsilon = epsilon
        self.allreduce_params = allreduce_params
941
        self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
942
943
944

    def register(self, pm_pass: PatternMatcherPass):
        def get_inputs():
945
946
947
948
949
            input = torch.empty([1, 16, 16], device=self.device, dtype=self.dtype)
            quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)
            input_global_scale = torch.empty(
                [1, 1], device=self.device, dtype=torch.float32
            )
950
            weight = torch.empty([16], device=self.device, dtype=self.dtype)
951
            output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
952

953
            return [input, quant_result, weight, input_global_scale, output_scale]
954
955
956
957
958
959
960
961
962

        def pattern(
            input: torch.Tensor,
            quant_result: torch.Tensor,
            weight: torch.Tensor,
            input_global_scale: torch.Tensor,
            output_scale: torch.Tensor,
        ):
            all_reduce = tensor_model_parallel_all_reduce(input)
963
            rms = self.rmsnorm_matcher(all_reduce, weight)
964
965
966
            quant_out_tuple = auto_functionalized(
                STATIC_FP4_QUANT_OP,
                output=quant_result,
967
                input=rms,
968
                output_scale=output_scale,
969
970
                input_scale=input_global_scale,
            )
971
972
973
974

            # quant_out, allreduce_output, output_scale
            return quant_out_tuple[1], all_reduce, quant_out_tuple[2]

975
976
977
978
979
980
981
        def replacement(
            input: torch.Tensor,
            quant_result: torch.Tensor,
            weight: torch.Tensor,
            input_global_scale: torch.Tensor,
            output_scale: torch.Tensor,
        ):
982
            residual = torch.zeros_like(input)
983
            result_rms = torch.empty_like(input)
984
985
986
987
988
989
990
991
992
            allreduce = auto_functionalized(
                flashinfer_trtllm_fused_allreduce_norm,
                allreduce_in=input,
                residual=residual,
                norm_out=result_rms,
                quant_out=quant_result,
                scale_out=output_scale,
                rms_gamma=weight,
                rms_eps=self.epsilon,
993
994
995
996
                # We don't use norm_out afterwards
                pattern_code=(
                    flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant
                ),
997
998
999
1000
1001
1002
1003
                scale_factor=input_global_scale,
                **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
            )

            # quant_out, allreduce_output, output_scale
            return allreduce[4], allreduce[1], allreduce[5]

1004
1005
1006
        pm.register_replacement(
            pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
        )
1007
1008
1009
1010
1011
1012


class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
    """
    This pattern replaces the allreduce + rms norm (with residual)
    + static nvfp4 quant with fused flashinfer implementation.
1013
    Applies to o_proj + rmsnorm after attn + quant and
1014
1015
1016
    mlp + rmsnorm + quant before attn.
    """

1017
1018
1019
1020
1021
1022
1023
    def __init__(
        self,
        epsilon: float,
        dtype: torch.dtype,
        device: str,
        allreduce_params: FlashInferFusedAllReduceParams,
    ):
1024
1025
1026
        super().__init__(dtype, device)
        self.epsilon = epsilon
        self.allreduce_params = allreduce_params
1027
        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
1028
1029
1030
1031
1032

    def register(self, pm_pass: PatternMatcherPass):
        def get_inputs():
            input = torch.empty([16, 16], device=self.device, dtype=self.dtype)

1033
1034
1035
1036
1037
1038
1039
            residual = torch.empty([16, 16], device=self.device, dtype=self.dtype)
            weight = torch.empty([16, 16], device=self.device, dtype=self.dtype)
            quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)
            input_global_scale = torch.empty(
                [1, 1], device=self.device, dtype=torch.float32
            )
            output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049

            return [
                quant_result,
                residual,
                input,
                output_scale,
                weight,
                input_global_scale,
            ]

1050
1051
1052
1053
1054
1055
1056
1057
        def pattern(
            quant_result: torch.Tensor,
            residual: torch.Tensor,
            input: torch.Tensor,
            output_scale: torch.Tensor,
            weight: torch.Tensor,
            input_global_scale: torch.Tensor,
        ):
1058
            allreduce_output = tensor_model_parallel_all_reduce(input)
1059
            rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
1060
1061
1062
            quant_out_tuple = auto_functionalized(
                STATIC_FP4_QUANT_OP,
                output=quant_result,
1063
                input=rms,
1064
                output_scale=output_scale,
1065
1066
                input_scale=input_global_scale,
            )
1067
1068

            # quant_out, allreduce_output, output_scale
1069
            return quant_out_tuple[1], residual, quant_out_tuple[2]
1070

1071
1072
1073
1074
1075
1076
1077
1078
        def replacement(
            quant_result: torch.Tensor,
            residual: torch.Tensor,
            input: torch.Tensor,
            output_scale: torch.Tensor,
            weight: torch.Tensor,
            input_global_scale: torch.Tensor,
        ):
1079
1080
1081
1082
1083
1084
1085
1086
1087
            allreduce = auto_functionalized(
                flashinfer_trtllm_fused_allreduce_norm,
                allreduce_in=input,
                residual=residual,
                norm_out=None,
                quant_out=quant_result,
                scale_out=output_scale,
                rms_gamma=weight,
                rms_eps=self.epsilon,
1088
1089
1090
1091
                # We don't use norm_out afterwards
                pattern_code=(
                    flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant
                ),
1092
1093
1094
1095
1096
1097
                scale_factor=input_global_scale,
                **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
            )
            # quant_out, rms_norm_residual, output_scale
            return allreduce[4], allreduce[2], allreduce[5]

1098
1099
1100
        pm.register_replacement(
            pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
        )
1101
1102


1103
class AllReduceFusionPass(VllmPatternMatcherPass):
1104
    def __init__(self, config: VllmConfig):
1105
1106
1107
1108
1109
1110
        super().__init__(config)
        self.disabled = True
        self.tp_size = get_tensor_model_parallel_world_size()
        if self.tp_size <= 1:
            return
        self.patterns: PatternMatcherPass = PatternMatcherPass(
1111
1112
            pass_name="all_reduce_fusion_pass"
        )
1113
1114
1115
1116
1117
1118
1119
1120
        if config.model_config is None:
            return
        self.hidden_dim = config.model_config.get_hidden_size()
        self.group = get_tp_group().device_group
        rank = get_tensor_model_parallel_rank()
        use_fp32_lamport = self.model_dtype == torch.float32
        if flashinfer_comm is None:
            logger.warning(
1121
                "Flashinfer is not installed or comm module not found, "
1122
1123
                "skipping allreduce fusion pass"
            )
1124
            return
1125
1126
1127
1128
1129
        max_size = config.compilation_config.pass_config.flashinfer_max_size(
            self.tp_size
        )
        if max_size is None:
            # Flashinfer doesn't support current world size
1130
            logger.warning(
1131
                "Flashinfer allreduce fusion is not supported for world size %s",
1132
1133
1134
                self.tp_size,
            )
            return
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
        element_size = 4 if use_fp32_lamport else 2
        self.max_token_num = max_size // (self.hidden_dim * element_size)
        # take the min to save workspace size and we'll never use more
        # than max_num_batched_tokens anyways
        self.max_token_num = min(
            self.max_token_num, config.scheduler_config.max_num_batched_tokens
        )
        logger.debug_once(
            f"Flashinfer max size: {max_size // (1024 * 1024)} MB,"
            "Maximal number of tokens used by "
            f"Flashinfer Allreduce Fusion: {self.max_token_num}",
            scope="global",
1147
        )
1148

1149
1150
1151
1152
        self.ipc_handles, workspace_tensor = (
            flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
                tp_rank=rank,
                tp_size=self.tp_size,
1153
                max_token_num=self.max_token_num,
1154
1155
1156
                hidden_dim=self.hidden_dim,
                group=self.group,
                use_fp32_lamport=use_fp32_lamport,
1157
1158
            )
        )
1159
1160
1161
1162
1163
1164
1165

        global _FI_WORKSPACE_TENSOR
        _FI_WORKSPACE_TENSOR = workspace_tensor
        self.allreduce_params = FlashInferFusedAllReduceParams(
            rank=rank,
            world_size=self.tp_size,
            use_fp32_lamport=use_fp32_lamport,
1166
            max_token_num=self.max_token_num,
1167
        )
1168

1169
        self.register_patterns()
1170
        self.dump_patterns(config, self.patterns)
1171
1172
1173

    @enable_fake_mode
    def register_patterns(self):
1174
        for epsilon in [1e-5, 1e-6]:
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
            AllReduceFusedRMSNormStaticQuantFP8Pattern(
                epsilon,
                self.model_dtype,
                self.device,
                self.allreduce_params,
            ).register(self.patterns)
            AllReduceFusedAddRMSNormStaticQuantFP8Pattern(
                epsilon,
                self.model_dtype,
                self.device,
                self.allreduce_params,
            ).register(self.patterns)
            if current_platform.has_device_capability(100):
                AllReduceFusedRMSNormStaticQuantNVFP4Pattern(
                    epsilon,
                    self.model_dtype,
                    self.device,
                    self.allreduce_params,
                ).register(self.patterns)
                AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(
                    epsilon,
                    self.model_dtype,
                    self.device,
                    self.allreduce_params,
                ).register(self.patterns)
            AllReduceRMSNormPattern(
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
                epsilon,
                self.model_dtype,
                self.device,
                self.allreduce_params,
            ).register(self.patterns)
            AllReduceFusedAddRMSNormPattern(
                epsilon,
                self.model_dtype,
                self.device,
                self.allreduce_params,
            ).register(self.patterns)

1213
1214
1215
1216
            # WARNING: This is a hack to clear the pattern matcher cache
            # and allow multiple values of epsilon.
            torch._inductor.pattern_matcher._seen_patterns.clear()

1217
1218
        self.disabled = False

1219
    @VllmInductorPass.time_and_log
1220
1221
    def __call__(self, graph: fx.Graph):
        if self.disabled:
1222
            logger.debug("AllReduceFusionPass disabled")
1223
            return
1224
1225
1226

        self.matched_count = self.patterns.apply(graph)
        logger.debug("Replaced %s patterns", self.matched_count)
1227
1228

    def __del__(self):
1229
        if getattr(self, "disabled", True):
1230
1231
            return
        if flashinfer_comm is not None:
1232
            flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(
1233
1234
                self.ipc_handles, self.group
            )