| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fuser detection for the Transformers modeling backend.""" |
| 4 | |
| 5 | from collections import UserDict |
| 6 | from typing import TYPE_CHECKING, TypeVar |
| 7 | |
| 8 | from cachetools import cached |
| 9 | from torch import nn |
| 10 | |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.models.transformers.fusers import ( |
| 13 | AttentionFuser, |
| 14 | BaseFuser, |
| 15 | GLUFuser, |
| 16 | MergedColumnParallelFuser, |
| 17 | MLAFuser, |
| 18 | PackedQKVFuser, |
| 19 | QKVFuser, |
| 20 | RewriteFuser, |
| 21 | RMSNormFuser, |
| 22 | ) |
| 23 | from vllm.model_executor.models.transformers.fx_utils import trace |
| 24 | |
| 25 | if TYPE_CHECKING: |
| 26 | from vllm.config import VllmConfig |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | F = TypeVar("F", bound=BaseFuser) |
| 31 | |
| 32 | |
| 33 | def key(module: nn.Module) -> tuple: |
| 34 | """Cache key for `get_fusers`. Considers module type and its immediate children.""" |
| 35 | return (type(module), tuple(name for name, _ in module.named_children())) |
| 36 | |
| 37 | |
| 38 | FUSERS: tuple[type[BaseFuser], ...] = ( |
| 39 | # Order these by priority, e.g. QKV would be preferred over MergedColumnParallel |
| 40 | MLAFuser, |
| 41 | GLUFuser, |
| 42 | QKVFuser, |
| 43 | PackedQKVFuser, |
| 44 | MergedColumnParallelFuser, |
| 45 | RMSNormFuser, |
| 46 | # Put fusers that don't redefine forward after this comment |
| 47 | AttentionFuser, |
| 48 | ) |
| 49 | """Every fuser, in priority order: those that redefine the forward first, then those |
| 50 | that leave it alone. A new fuser is added here.""" |
| 51 | |
| 52 | |
| 53 | @cached(cache={}, key=key) |
| 54 | def get_fusers(module: nn.Module) -> list[BaseFuser]: |
| 55 | """Every fuser that could apply to `module`'s class (cached), in `FUSERS` order.""" |
| 56 | # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a |
| 57 | # leaf module (raw tensor math, no submodules). Nothing else can match, and |
| 58 | # tracing is skipped for it. |
| 59 | n_linear = sum(isinstance(c, nn.Linear) for c in module.children()) |
| 60 | is_leaf = next(module.children(), None) is None |
| 61 | graph = trace(module) if n_linear >= 2 or is_leaf else None |
| 62 | |
| 63 | # Every redefining candidate that matches is cached, not just the first: |
| 64 | # instances of this class can be shaped heterogeneously, so which one is |
| 65 | # usable is decided per instance in Fusers.__getitem__, not here. A forward |
| 66 | # can only be rewritten from a trace of it, so each starts fresh regardless. |
| 67 | fusers: list[BaseFuser] = [] |
| 68 | for fuser_cls in FUSERS: |
| 69 | if fuser_cls.redefines_forward and graph is None: |
| 70 | continue |
| 71 | if (fuser := fuser_cls.match(graph, module)) is None: |
| 72 | continue |
| 73 | if isinstance(fuser, RewriteFuser): |
| 74 | try: |
| 75 | fuser.update_forward(module) |
| 76 | except Exception as exc: |
| 77 | logger.debug( |
| 78 | "Attempted to fuse %s using %s but failed " |
| 79 | "to update its forward method: %s", |
| 80 | type(module), |
| 81 | fuser_cls.__name__, |
| 82 | exc, |
| 83 | ) |
| 84 | continue |
| 85 | fusers.append(fuser) |
| 86 | |
| 87 | # A norm we could not match structurally is left unfused; flag likely misses. |
| 88 | if ( |
| 89 | module.__class__.__name__.endswith("RMSNorm") |
| 90 | and graph is not None |
| 91 | and not any(isinstance(fuser, RMSNormFuser) for fuser in fusers) |
| 92 | ): |
| 93 | logger.warning_once( |
| 94 | "%s looks like an RMSNorm but its computation did not match the " |
| 95 | "expected pattern, so it was left unfused.", |
| 96 | module.__class__.__name__, |
| 97 | ) |
| 98 | return fusers |
| 99 | |
| 100 | |
| 101 | def get_fuser(module: nn.Module, fuser_cls: type[F]) -> "F | None": |
| 102 | """The first `fuser_cls` that applies to `module`, if one does.""" |
| 103 | return next((f for f in get_fusers(module) if isinstance(f, fuser_cls)), None) |
| 104 | |
| 105 | |
| 106 | class Fusers(UserDict): |
| 107 | """Mapping from module class and shape to fusers, for all fusable modules.""" |
| 108 | |
| 109 | def __init__(self, model: nn.Module, vllm_config: "VllmConfig"): |
| 110 | self.vllm_config = vllm_config |
| 111 | super().__init__({key(m): get_fusers(m) for m in model.modules()}) |
| 112 | |
| 113 | def __getitem__(self, m: nn.Module) -> list[BaseFuser]: |
| 114 | """The fusers this instance can take, in the order to apply them.""" |
| 115 | chosen: list[BaseFuser] = [] |
| 116 | redefined = False |
| 117 | for fuser in self.data.get(key(m), ()): |
| 118 | # Only one candidate may rewrite m's forward; take the first |
| 119 | # (highest-priority) one that validates for this instance. |
| 120 | if fuser.redefines_forward and redefined: |
| 121 | continue |
| 122 | if not fuser.validate(m, self.vllm_config): |
| 123 | continue |
| 124 | redefined = redefined or fuser.redefines_forward |
| 125 | chosen.append(fuser) |
| 126 | return chosen |