chore: import upstream snapshot with attribution

vllm/model_executor/models/transformers/fuser.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fuser detection for the Transformers modeling backend.
+
+`get_fuser` traces a module class once (see `fx_utils`) and matches it against
+each concrete fuser in `fusers`; `Fusers` caches the result per class for a
+whole model. `base.recursive_replace` then applies the matched fuser per
+instance. RMSNorm-shaped modules the tracer cannot match are warned about.
+"""
+
+from collections import UserDict
+from typing import TYPE_CHECKING
+
+from cachetools import cached
+from torch import nn
+
+from vllm.logger import init_logger
+from vllm.model_executor.models.transformers.fusers import (
+    BaseFuser,
+    GLUFuser,
+    QKVFuser,
+    RMSNormFuser,
+    StackedFuser,
+)
+from vllm.model_executor.models.transformers.fx_utils import trace
+
+if TYPE_CHECKING:
+    from vllm.config.model import ModelConfig
+
+logger = init_logger(__name__)
+
+
+@cached(cache={}, key=type)
+def get_fuser(module: nn.Module) -> BaseFuser | None:
+    """The fuser for `type(module)` (cached per class), or `None` if no match."""
+    # 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
+    if n_linear < 2 and not is_leaf:
+        return None
+    if (graph := trace(module)) is None:
+        return None
+    for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser):
+        if (fuser := fuser_cls.match(graph, module)) is not None:
+            if isinstance(fuser, StackedFuser):
+                try:
+                    fuser.update_forward(module)
+                except Exception as exc:
+                    # An unrecognised source just means we cannot fuse here.
+                    logger.debug(
+                        "Could not rewrite %s for fusion: %s", type(module), exc
+                    )
+                    return None
+            return fuser
+    # A norm we could not match structurally is left unfused; flag likely misses.
+    if module.__class__.__name__.endswith("RMSNorm"):
+        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 None
+
+
+class Fusers(UserDict):
+    """Mapping from module class to fuser, for all fusable classes in a model."""
+
+    def __init__(self, model: nn.Module, model_config: "ModelConfig"):
+        self.model_config = model_config
+        super().__init__({type(m): get_fuser(m) for m in model.modules()})
+
+    def __getitem__(self, m: nn.Module) -> BaseFuser | None:
+        fuser = self.data.get(type(m))
+        if fuser is not None and fuser.validate(m, self.model_config):
+            return fuser
+        return None