chore: import upstream snapshot with attribution

vllm/model_executor/models/transformers/fusers/base.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Base classes for the Transformers backend fusers."""
+
+import types
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, ClassVar
+
+from torch import fx, nn
+
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
+
+if TYPE_CHECKING:
+    from vllm.config.model import ModelConfig
+    from vllm.model_executor.layers.quantization import QuantizationConfig
+
+
+@dataclass
+class BaseFuser(ABC):
+    """A detected fusion and how to apply it.
+
+    `match` analyses the module *class* once (cached, see `get_fuser`); `fuse`
+    then applies the fusion to an instance in `recursive_replace`, returning the
+    module to install in its place.
+    """
+
+    @abstractmethod
+    def info(self, name: str) -> str:
+        """A human-readable description of the fusion at `name`, for logging."""
+
+    @classmethod
+    @abstractmethod
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
+        """Match the pattern in `graph`, returning a fuser if found."""
+
+    @abstractmethod
+    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
+        """Whether this fuser can be applied to this `module` instance."""
+
+    @abstractmethod
+    def fuse(
+        self,
+        module: nn.Module,
+        prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
+    ) -> nn.Module:
+        """Apply the fusion to an already-validated `module`, returning the
+        module to install in its place (mutated in place, or freshly built)."""
+
+    def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
+        """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes
+        (none unless it stacks weights)."""
+        return {}
+
+    @property
+    def packed_modules_mapping(self) -> dict[str, list[str]]:
+        """`packed_modules_mapping` entries this fuser contributes (none unless
+        it stacks weights)."""
+        return {}
+
+
+@dataclass
+class StackedFuser(BaseFuser):
+    """A fuser that merges sibling projections into one stacked linear and
+    rewrites the forward to call it.
+
+    `match` and `update_forward` analyse the class once; `fuse` builds the merged
+    submodule and binds the compiled forward on an instance in place, so it keeps
+    its class and any attribute the fusion does not consume.
+    """
+
+    merged_name: ClassVar[str]
+    """Attribute name of the merged module created by `update_attrs`."""
+    merged_cls: ClassVar[str]
+    """Name of the vLLM class the merged projection becomes (for logging)."""
+
+    source_cls: str
+    """Class of the HF module the fused projections belonged to (for logging)."""
+
+    fused_forward: Callable = field(init=False, repr=False)
+    """The compiled rewritten forward, set by `update_forward`."""
+
+    def info(self, name: str) -> str:
+        sources = " + ".join(shard for shard, _ in self.shards)
+        return (
+            f"Fused: {sources} ({name}: {self.source_cls}) -> "
+            f"{self.merged_name} ({self.merged_cls})"
+        )
+
+    @property
+    @abstractmethod
+    def shards(self) -> list[tuple[str, ShardId]]:
+        """Each projection's original name and its shard id in the merged module.
+
+        Source for both `orig_to_new_stacked` and `packed_modules_mapping`."""
+
+    def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
+        """`WeightsMapper.orig_to_new_stacked` entries for one fused instance.
+
+        Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname
+        so only this exact layer is remapped, never a same-named projection
+        elsewhere (e.g. an unfused MoE expert's `gate_proj`)."""
+        merged = maybe_prefix(prefix, self.merged_name)
+        return {
+            maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards
+        }
+
+    @property
+    def packed_modules_mapping(self) -> dict[str, list[str]]:
+        """`{merged_name: [projection names]}` so quantization can unpack the
+        fused layer into its per-shard configs."""
+        return {self.merged_name: [name for name, _ in self.shards]}
+
+    @abstractmethod
+    def update_forward(self, module: nn.Module) -> None:
+        """Rewrite and compile `type(module)`'s forward source.
+
+        Raises if the source does not admit the rewrite (fusion is then skipped).
+        """
+
+    @abstractmethod
+    def update_attrs(
+        self,
+        module: nn.Module,
+        prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
+    ) -> None:
+        """Replace `module`'s submodules with the merged module."""
+
+    def fuse(
+        self,
+        module: nn.Module,
+        prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
+    ) -> nn.Module:
+        """Fuse an already-validated `module` in place (see `Fusers.__getitem__`).
+
+        Builds the merged submodule and binds the compiled forward."""
+        self.update_attrs(module, prefix, model_config, quant_config)
+        module.forward = types.MethodType(self.fused_forward, module)
+        return module