| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import functools |
| 4 | from collections.abc import Callable |
| 5 | from typing import Any, ParamSpec, TypeVar |
| 6 | |
| 7 | from torch import fx as fx |
| 8 | |
| 9 | from vllm import envs |
| 10 | from vllm._aiter_ops import rocm_aiter_ops |
| 11 | from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass |
| 12 | from vllm.config import VllmConfig, set_current_vllm_config |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.system_utils import set_env_var |
| 16 | |
| 17 | from .ir.clone_elimination import UnsafeCloneEliminationPass |
| 18 | from .ir.lowering_pass import VllmIRLoweringPass |
| 19 | from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass |
| 20 | |
| 21 | if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled(): |
| 22 | from .fusion.allreduce_rms_fusion import ( |
| 23 | RocmAiterAllReduceFusionPass, |
| 24 | ) |
| 25 | from .fusion.rocm_aiter_fusion import ( |
| 26 | MLADualRMSNormFusionPass, |
| 27 | RocmAiterRMSNormQuantFusionPass, |
| 28 | RocmAiterSiluMulFp8GroupQuantFusionPass, |
| 29 | RocmAiterTritonAddRMSNormPadFusionPass, |
| 30 | ) |
| 31 | |
| 32 | if current_platform.is_cuda_alike() or current_platform.is_xpu(): |
| 33 | from .fusion.add_rms_fusion import ( |
| 34 | AddRMSNormFusionPass, |
| 35 | RMSNormReshapeFusionPass, |
| 36 | ) |
| 37 | from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass |
| 38 | from .fusion.sequence_parallelism import SequenceParallelismPass |
| 39 | from .utility.split_coalescing import SplitCoalescingPass |
| 40 | |
| 41 | if current_platform.is_cuda_alike(): |
| 42 | from .fusion.act_quant_fusion import ActivationQuantFusionPass |
| 43 | from .fusion.attn_quant_fusion import AttnQuantFusionPass |
| 44 | from .fusion.collective_fusion import AsyncTPPass |
| 45 | from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass |
| 46 | from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass |
| 47 | from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass |
| 48 | from .fusion.rms_quant_fusion import RMSNormQuantFusionPass |
| 49 | from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass |
| 50 | from .utility.scatter_split_replace import ScatterSplitReplacementPass |
| 51 | |
| 52 | if current_platform.is_cuda(): |
| 53 | from .fusion.allreduce_rms_fusion import AllReduceFusionPass |
| 54 | |
| 55 | if current_platform.is_xpu(): |
| 56 | from .fusion.act_quant_fusion import ActivationQuantFusionPass |
| 57 | from .fusion.rms_quant_fusion import RMSNormQuantFusionPass |
| 58 | |
| 59 | from .inductor_pass import ( |
| 60 | CustomGraphPass, |
| 61 | InductorPass, |
| 62 | get_pass_context, |
| 63 | ) |
| 64 | from .utility.fix_functionalization import FixFunctionalizationPass |
| 65 | from .utility.noop_elimination import NoOpEliminationPass |
| 66 | |
| 67 | logger = init_logger(__name__) |
| 68 | |
| 69 | P = ParamSpec("P") |
| 70 | R = TypeVar("R") |
| 71 | |
| 72 | |
| 73 | def with_pattern_match_debug(fn: Callable[P, R]) -> Callable[P, R]: |
| 74 | """Function decorator that turns on inductor pattern match debug |
| 75 | for the duration of the call. |
| 76 | Used to avoid logging builtin Inductor pattern matching. |
| 77 | """ |
| 78 | |
| 79 | @functools.wraps(fn) |
| 80 | def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: |
| 81 | if (debug_val := envs.VLLM_PATTERN_MATCH_DEBUG) is not None: |
| 82 | # optionally check rank here |
| 83 | with set_env_var("TORCHINDUCTOR_PATTERN_MATCH_DEBUG", debug_val): |
| 84 | return fn(*args, **kwargs) |
| 85 | return fn(*args, **kwargs) |
| 86 | |
| 87 | return wrapper |
| 88 | |
| 89 | |
| 90 | class PostGradPassManager(CustomGraphPass): # type: ignore[misc] |
| 91 | """The pass manager for post-grad passes. |
| 92 | It handles configuration, adding custom passes, and running passes. |
| 93 | It supports uuid for the Inductor code cache. That includes torch<2.6 |
| 94 | support using pickling (in .inductor_pass.CustomGraphPass). |
| 95 | |
| 96 | The order of the post-grad post-passes is: |
| 97 | 1. passes (constructor parameter) |
| 98 | 2. default passes (NoopEliminationPass, FusionPass) |
| 99 | 3. config["post_grad_custom_post_pass"] (if it exists) |
| 100 | 4. fix_functionalization |
| 101 | This way, all passes operate on a functionalized graph. |
| 102 | """ |
| 103 | |
| 104 | def __init__(self) -> None: |
| 105 | self.passes: list[InductorPass] = [] |
| 106 | |
| 107 | @with_pattern_match_debug |
| 108 | def __call__(self, graph: fx.Graph) -> None: |
| 109 | VllmInductorPass.dump_prefix = 0 # reset dump index |
| 110 | |
| 111 | compile_range = get_pass_context().compile_range |
| 112 | for pass_ in self.passes: |
| 113 | if pass_.is_applicable_for_range(compile_range): |
| 114 | pass_(graph) |
| 115 | VllmInductorPass.dump_prefix += 1 |
| 116 | else: |
| 117 | logger.debug("Skipping %s with compile range %s", pass_, compile_range) |
| 118 | |
| 119 | # perform the first post-cleanup before IR lowering to clean up fusion artifacts |
| 120 | # and make sure no dead IR ops are lowered. |
| 121 | self.post_cleanup(graph) |
| 122 | VllmInductorPass.dump_prefix += 1 |
| 123 | |
| 124 | # lowering before cleanup so DCE can clean up lowered ops. |
| 125 | # DCE handles mutating ops correctly as well. |
| 126 | self.ir_lowering(graph) |
| 127 | VllmInductorPass.dump_prefix += 1 |
| 128 | self.clone_elimination(graph) |
| 129 | VllmInductorPass.dump_prefix += 1 |
| 130 | |
| 131 | # clean up after lowering again |
| 132 | self.post_cleanup(graph) |
| 133 | VllmInductorPass.dump_prefix += 1 |
| 134 | |
| 135 | # always run fix_functionalization last |
| 136 | self.fix_functionalization(graph) |
| 137 | VllmInductorPass.dump_prefix = None # Cleanup index |
| 138 | |
| 139 | VllmPatternMatcherPass.log_match_summary() |
| 140 | |
| 141 | def configure(self, config: VllmConfig) -> None: |
| 142 | self.pass_config = config.compilation_config.pass_config |
| 143 | model_config = config.model_config |
| 144 | enable_transformers_norm_canonicalization = ( |
| 145 | ( |
| 146 | self.pass_config.fuse_act_padding |
| 147 | or self.pass_config.fuse_allreduce_rms |
| 148 | or self.pass_config.fuse_norm_quant |
| 149 | ) |
| 150 | and model_config is not None |
| 151 | and model_config.using_transformers_backend() |
| 152 | ) |
| 153 | |
| 154 | # Set the current vllm config to allow tracing CustomOp instances |
| 155 | with set_current_vllm_config(config, check_compile=False): |
| 156 | if self.pass_config.eliminate_noops: |
| 157 | self.passes += [NoOpEliminationPass(config)] |
| 158 | |
| 159 | if self.pass_config.enable_sp: |
| 160 | self.passes += [SequenceParallelismPass(config)] |
| 161 | if self.pass_config.fuse_gemm_comms: |
| 162 | self.passes += [AsyncTPPass(config)] |
| 163 | |
| 164 | if enable_transformers_norm_canonicalization: |
| 165 | self.passes += [AddRMSNormFusionPass(config)] |
| 166 | |
| 167 | if self.pass_config.fuse_act_padding and ( |
| 168 | rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled() |
| 169 | ): |
| 170 | # Run the more specific RMSNorm+router-pad fusion before |
| 171 | # AR+RMS, since both consume fused_add_rms_norm. |
| 172 | self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)] |
| 173 | |
| 174 | if self.pass_config.fuse_allreduce_rms: |
| 175 | if rocm_aiter_ops.is_enabled(): |
| 176 | self.passes += [RocmAiterAllReduceFusionPass(config)] |
| 177 | else: |
| 178 | self.passes += [AllReduceFusionPass(config)] |
| 179 | |
| 180 | if enable_transformers_norm_canonicalization: |
| 181 | # Let AR+RMS match before moving output reshapes ahead of the |
| 182 | # remaining RMSNorms, exposing them to RMS+Quant fusion. |
| 183 | self.passes += [RMSNormReshapeFusionPass(config)] |
| 184 | |
| 185 | if self.pass_config.fuse_norm_quant: |
| 186 | if ( |
| 187 | rocm_aiter_ops.is_enabled() |
| 188 | or rocm_aiter_ops.is_rdna_aiter_enabled() |
| 189 | ): |
| 190 | self.passes += [ |
| 191 | RocmAiterRMSNormQuantFusionPass(config), |
| 192 | ] |
| 193 | self.passes += [RMSNormQuantFusionPass(config)] |
| 194 | |
| 195 | if self.pass_config.fuse_act_quant: |
| 196 | self.passes += [ActivationQuantFusionPass(config)] |
| 197 | if ( |
| 198 | rocm_aiter_ops.is_enabled() |
| 199 | or rocm_aiter_ops.is_rdna_aiter_enabled() |
| 200 | ): |
| 201 | self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)] |
| 202 | |
| 203 | if self.pass_config.fuse_qk_norm_rope_kvcache: |
| 204 | self.passes += [SplitCoalescingPass(config)] |
| 205 | self.passes += [ScatterSplitReplacementPass(config)] |
| 206 | self.passes += [QkNormRopeKvCacheFusionPass(config)] |
| 207 | |
| 208 | if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled(): |
| 209 | self.passes += [MLADualRMSNormFusionPass(config)] |
| 210 | |
| 211 | if self.pass_config.fuse_rope_kvcache: |
| 212 | self.passes += [SplitCoalescingPass(config)] |
| 213 | self.passes += [ScatterSplitReplacementPass(config)] |
| 214 | self.passes += [RopeKVCacheFusionPass(config)] |
| 215 | |
| 216 | if self.pass_config.fuse_rope_kvcache_cat_mla: |
| 217 | self.passes += [MLARoPEKVCacheCatFusionPass(config)] |
| 218 | |
| 219 | if self.pass_config.fuse_attn_quant: |
| 220 | self.passes += [AttnQuantFusionPass(config)] |
| 221 | self.passes += [MLAAttnQuantFusionPass(config)] |
| 222 | |
| 223 | if self.pass_config.enable_qk_norm_rope_fusion: |
| 224 | self.passes += [SplitCoalescingPass(config)] |
| 225 | self.passes += [QKNormRoPEFusionPass(config)] |
| 226 | |
| 227 | self.ir_lowering = VllmIRLoweringPass(config) |
| 228 | self.clone_elimination = UnsafeCloneEliminationPass(config) |
| 229 | self.post_cleanup = PostCleanupPass(config) |
| 230 | self.fix_functionalization = FixFunctionalizationPass(config) |
| 231 | |
| 232 | def add(self, pass_: InductorPass) -> None: |
| 233 | assert isinstance(pass_, InductorPass) |
| 234 | self.passes.append(pass_) |
| 235 | |
| 236 | def uuid(self) -> str: |
| 237 | """The PostGradPassManager is set as a custom pass in the Inductor and |
| 238 | affects compilation caching. Its uuid depends on the UUIDs of all |
| 239 | dependent passes and the pass config. See InductorPass for more info. |
| 240 | """ |
| 241 | passes = [] |
| 242 | |
| 243 | state: dict[str, Any] = {"pass_config": self.pass_config.compute_hash()} |
| 244 | for pass_ in self.passes: |
| 245 | passes.append(pass_.uuid()) |
| 246 | |
| 247 | passes.append(self.post_cleanup.uuid()) |
| 248 | passes.append(self.ir_lowering.uuid()) |
| 249 | passes.append(self.clone_elimination.uuid()) |
| 250 | passes.append(self.post_cleanup.uuid()) |
| 251 | passes.append(self.fix_functionalization.uuid()) |
| 252 | |
| 253 | # Include the compile range in the uuid to ensure that inductor |
| 254 | # recompiles the graph for the new dynamic compile range. |
| 255 | state["compile_range"] = str(get_pass_context().compile_range) |
| 256 | state["passes"] = passes |
| 257 | return InductorPass.hash_dict(state) |