Skip to content

vllm.model_executor.models.transformers.fuser

Fuser detection for the Transformers modeling backend.

Classes:

  • Fusers

    Mapping from module class and shape to fusers, for all fusable modules.

Functions:

  • get_fuser

    The first fuser_cls that applies to module, if one does.

  • get_fusers

    Every fuser that could apply to module's class (cached), in FUSERS order.

  • key

    Cache key for get_fusers. Considers module type and its immediate children.

Attributes:

  • FUSERS (tuple[type[BaseFuser], ...]) –

    Every fuser, in priority order: those that redefine the forward first, then those

FUSERS = (MLAFuser, GLUFuser, QKVFuser, PackedQKVFuser, MergedColumnParallelFuser, RMSNormFuser, AttentionFuser) module-attribute

Every fuser, in priority order: those that redefine the forward first, then those that leave it alone. A new fuser is added here.

Fusers

Bases: UserDict

Mapping from module class and shape to fusers, for all fusable modules.

Methods:

  • __getitem__

    The fusers this instance can take, in the order to apply them.

Source code in vllm/model_executor/models/transformers/fuser.py
class Fusers(UserDict):
    """Mapping from module class and shape to fusers, for all fusable modules."""

    def __init__(self, model: nn.Module, vllm_config: "VllmConfig"):
        self.vllm_config = vllm_config
        super().__init__({key(m): get_fusers(m) for m in model.modules()})

    def __getitem__(self, m: nn.Module) -> list[BaseFuser]:
        """The fusers this instance can take, in the order to apply them."""
        chosen: list[BaseFuser] = []
        redefined = False
        for fuser in self.data.get(key(m), ()):
            # Only one candidate may rewrite m's forward; take the first
            # (highest-priority) one that validates for this instance.
            if fuser.redefines_forward and redefined:
                continue
            if not fuser.validate(m, self.vllm_config):
                continue
            redefined = redefined or fuser.redefines_forward
            chosen.append(fuser)
        return chosen

__getitem__(m)

The fusers this instance can take, in the order to apply them.

Source code in vllm/model_executor/models/transformers/fuser.py
def __getitem__(self, m: nn.Module) -> list[BaseFuser]:
    """The fusers this instance can take, in the order to apply them."""
    chosen: list[BaseFuser] = []
    redefined = False
    for fuser in self.data.get(key(m), ()):
        # Only one candidate may rewrite m's forward; take the first
        # (highest-priority) one that validates for this instance.
        if fuser.redefines_forward and redefined:
            continue
        if not fuser.validate(m, self.vllm_config):
            continue
        redefined = redefined or fuser.redefines_forward
        chosen.append(fuser)
    return chosen

get_fuser(module, fuser_cls)

The first fuser_cls that applies to module, if one does.

Source code in vllm/model_executor/models/transformers/fuser.py
def get_fuser(module: nn.Module, fuser_cls: type[F]) -> "F | None":
    """The first `fuser_cls` that applies to `module`, if one does."""
    return next((f for f in get_fusers(module) if isinstance(f, fuser_cls)), None)

get_fusers(module)

Every fuser that could apply to module's class (cached), in FUSERS order.

Source code in vllm/model_executor/models/transformers/fuser.py
@cached(cache={}, key=key)
def get_fusers(module: nn.Module) -> list[BaseFuser]:
    """Every fuser that could apply to `module`'s class (cached), in `FUSERS` order."""
    # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a
    # leaf module (raw tensor math, no submodules). Nothing else can match, and
    # tracing is skipped for it.
    n_linear = sum(isinstance(c, nn.Linear) for c in module.children())
    is_leaf = next(module.children(), None) is None
    graph = trace(module) if n_linear >= 2 or is_leaf else None

    # Every redefining candidate that matches is cached, not just the first:
    # instances of this class can be shaped heterogeneously, so which one is
    # usable is decided per instance in Fusers.__getitem__, not here. A forward
    # can only be rewritten from a trace of it, so each starts fresh regardless.
    fusers: list[BaseFuser] = []
    for fuser_cls in FUSERS:
        if fuser_cls.redefines_forward and graph is None:
            continue
        if (fuser := fuser_cls.match(graph, module)) is None:
            continue
        if isinstance(fuser, RewriteFuser):
            try:
                fuser.update_forward(module)
            except Exception as exc:
                logger.debug(
                    "Attempted to fuse %s using %s but failed "
                    "to update its forward method: %s",
                    type(module),
                    fuser_cls.__name__,
                    exc,
                )
                continue
        fusers.append(fuser)

    # A norm we could not match structurally is left unfused; flag likely misses.
    if (
        module.__class__.__name__.endswith("RMSNorm")
        and graph is not None
        and not any(isinstance(fuser, RMSNormFuser) for fuser in fusers)
    ):
        logger.warning_once(
            "%s looks like an RMSNorm but its computation did not match the "
            "expected pattern, so it was left unfused.",
            module.__class__.__name__,
        )
    return fusers

key(module)

Cache key for get_fusers. Considers module type and its immediate children.

Source code in vllm/model_executor/models/transformers/fuser.py
def key(module: nn.Module) -> tuple:
    """Cache key for `get_fusers`. Considers module type and its immediate children."""
    return (type(module), tuple(name for name, _ in module.named_children()))