chore: import upstream snapshot with attribution

vllm/model_executor/models/transformers/__init__.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Wrapper around `transformers` models"""
+
+from typing import TYPE_CHECKING
+
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from vllm.model_executor.models.transformers.base import Base
+from vllm.model_executor.models.transformers.causal import CausalMixin
+from vllm.model_executor.models.transformers.legacy import LegacyMixin
+from vllm.model_executor.models.transformers.moe import MoEMixin
+from vllm.model_executor.models.transformers.multimodal import (
+    MultiModalDummyInputsBuilder,
+    MultiModalMixin,
+    MultiModalProcessingInfo,
+    MultiModalProcessor,
+)
+from vllm.model_executor.models.transformers.pooling import (
+    EmbeddingMixin,
+    SequenceClassificationMixin,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+if TYPE_CHECKING:
+    import torch
+
+    from vllm.model_executor.layers.attention import Attention
+
+
+def vllm_attention_forward(
+    # Transformers args
+    module: "torch.nn.Module",
+    query: "torch.Tensor",
+    key: "torch.Tensor",
+    value: "torch.Tensor",
+    attention_mask: "torch.Tensor",
+    # Transformers kwargs
+    scaling: float | None = None,
+    # vLLM kwargs
+    attention_instances: dict[int, "Attention"] | None = None,
+    **kwargs,
+):
+    self_attn = attention_instances[module.layer_idx]
+    if scaling is not None:
+        self_attn.impl.scale = float(scaling)
+    hidden = query.shape[-2]
+    query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+    query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))
+    return self_attn.forward(query, key, value), None
+
+
+ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward
+
+
+# Text only models
+class TransformersForCausalLM(CausalMixin, Base): ...
+
+
+class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
+
+
+# Multimodal models
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+class TransformersMultiModalMoEForCausalLM(
+    MoEMixin, MultiModalMixin, CausalMixin, Base
+): ...
+
+
+# Embedding models
+class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ...
+
+
+class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ...
+
+
+# Sequence classification models
+class TransformersForSequenceClassification(
+    SequenceClassificationMixin, LegacyMixin, Base
+): ...
+
+
+class TransformersMoEForSequenceClassification(
+    SequenceClassificationMixin, MoEMixin, Base
+): ...
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+class TransformersMultiModalForSequenceClassification(
+    SequenceClassificationMixin, MultiModalMixin, Base
+): ...
+
+
+def __getattr__(name: str):
+    """Handle imports of non-existent classes with a helpful error message."""
+    if name not in globals():
+        raise AttributeError(
+            "The Transformers modeling backend does not currently have a class to "
+            f"handle the requested model type: {name}. Please open an issue at "
+            "https://github.com/vllm-project/vllm/issues/new"
+        )
+    return globals()[name]

vllm/model_executor/models/transformers/base.py

@@ -0,0 +1,702 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend base class."""
+
+from collections.abc import Callable, Iterable
+from itertools import chain
+from operator import attrgetter
+from typing import TYPE_CHECKING
+
+import regex as re
+import torch
+import transformers
+from packaging.version import Version
+from torch import nn
+from transformers import AutoModel
+from transformers.conversion_mapping import (
+    WeightRenaming,
+    get_model_conversion_mapping,
+)
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config.utils import getattr_iter
+from vllm.distributed import get_pp_group, get_tp_group
+from vllm.distributed.utils import get_pp_indices
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention import (
+    Attention,
+    EncoderOnlyAttention,
+)
+from vllm.model_executor.layers.fused_moe import MoERunner
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.models.interfaces import (
+    SupportsEagle,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsPP,
+    SupportsQuant,
+)
+from vllm.model_executor.models.interfaces_base import VllmModel
+from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers
+from vllm.model_executor.models.transformers.utils import (
+    can_enable_torch_compile,
+    get_feature_request_tip,
+    init_on_device_without_buffers,
+    log_replacement,
+    replace_conv_class,
+    replace_linear_class,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    make_empty_intermediate_tensors_factory,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backend import AttentionType
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
+
+
+class ScaledVocabParallelEmbedding(VocabParallelEmbedding):
+    """`VocabParallelEmbedding` that scales its output."""
+
+    def __init__(self, *args, embed_scale: float, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.embed_scale = embed_scale
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        return super().forward(input_) * self.embed_scale
+
+
+class Base(
+    nn.Module,
+    VllmModel,
+    SupportsQuant,
+    SupportsLoRA,
+    SupportsPP,
+    SupportsEagle,
+    SupportsEagle3,
+):
+    embedding_modules = ["embed_tokens"]  # TODO transformers will have a util to get it
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        super().__init__()
+        logger.info("Using Transformers modeling backend.")
+
+        self.config = vllm_config.model_config.hf_config
+        self.text_config = self.config.get_text_config()
+        self.cache_config = vllm_config.cache_config
+        self.compilation_config = vllm_config.compilation_config
+        self.device_config = vllm_config.device_config
+        self.model_config = vllm_config.model_config
+        self.parallel_config = vllm_config.parallel_config
+        self.quant_config = vllm_config.quant_config
+
+        self.pp_group = get_pp_group()
+        self.tp_group = get_tp_group()
+
+        # Attrs for weight loading (see self.load_weights)
+        self.skip_prefixes: list[str] = []
+        """Skip loading weights whose qualname starts with these prefixes."""
+        self.skip_substrs: list[str] = []
+        """Skip loading weights whose qualname contains these substrings."""
+        self.ignore_unexpected_prefixes: list[str] = []
+        """Ignore unexpected weights whose qualname starts with these prefixes."""
+        self.ignore_unexpected_suffixes: list[str] = []
+        """Ignore unexpected weights whose qualname ends with these suffixes."""
+        self.packed_modules_mapping: dict[str, list[str]] = {}
+        """Fused module -> constituent projections, populated by `recursive_replace`
+        for the quantization machinery and loaders (e.g. bitsandbytes)."""
+
+        # Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
+        self._target_class: type[nn.Module] = nn.Module
+        """Target class for Eagle3 aux hidden state recording."""
+        self._layer_names: dict[int, str] = {}
+        """Mapping from layer index to layer name for Eagle3."""
+        self._output_aux_hidden_states_kwargs: dict[str, bool] = {}
+        """Kwargs to pass to model forward for Eagle3 aux hidden states."""
+
+        if self.quant_config:
+            quant_method_name = self.quant_config.get_name()
+            # Check for unsupported quantization methods.
+            if quant_method_name in ("mxfp4", "gpt_oss_mxfp4"):
+                raise NotImplementedError(
+                    "Transformers modeling backend does "
+                    "not support MXFP4 quantization yet."
+                )
+
+        self._patch_config()
+        from_config_kwargs = dict(
+            config=self.config,
+            dtype=self.model_config.dtype,
+            trust_remote_code=self.model_config.trust_remote_code,
+        )
+        self._decorate_for_torch_compile(**from_config_kwargs)
+        # Init on "meta" to delay allocating GPU tensors
+        with init_on_device_without_buffers("meta"):
+            self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs)
+
+        # Create weight name to module qualname mapper
+        self._create_hf_to_vllm_mapper()
+        # Remove layers not on this pipeline parallel rank
+        self.pipeline_parallel()
+        # Substitute remaining layers with vLLM's layers as needed
+        self.recursive_replace()
+        # Create attention instances for KV cache allocation
+        self.attention_instances = self.create_attention_instances()
+
+        # Input embeddings
+        input_embeddings = self.model.get_input_embeddings()
+        if not isinstance(input_embeddings, PPMissingLayer):
+            names = ("embedding_size", "hidden_size")
+            embedding_dim = getattr_iter(self.text_config, names, None)
+            assert embedding_dim is not None
+            embedding_kwargs = dict(
+                num_embeddings=self.text_config.vocab_size,
+                embedding_dim=embedding_dim,
+                org_num_embeddings=self.text_config.vocab_size,
+                quant_config=self.quant_config,
+            )
+            embed_scale = getattr(input_embeddings, "embed_scale", None)
+            if embed_scale is not None:
+                # Some models scale embeddings inside the input embedding layer
+                new_input_embeddings = ScaledVocabParallelEmbedding(
+                    **embedding_kwargs, embed_scale=float(embed_scale)
+                )
+            else:
+                new_input_embeddings = VocabParallelEmbedding(**embedding_kwargs)
+            self.model.set_input_embeddings(new_input_embeddings)
+
+        # Initialize any parameters that have not had their modules replaced
+        self.init_parameters(self.model)
+
+        # Pipeline parallel intermediate tensors
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states"], self.text_config.hidden_size
+        )
+
+    def _patch_config(self):
+        """
+        Patch the config to ensure that the model is created correctly:
+
+        - Sets the attention implementation to "vllm" so the attention instances from
+        `create_attention_instances` are used
+        - Sets the dtype to the default torch dtype set by vLLM because Transformers
+        uses the config dtype when creating the model
+        """
+        self.text_config._attn_implementation = "vllm"
+        self.config.dtype = torch.get_default_dtype()
+
+    def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]:
+        """
+        Get the decoder class from the model.
+
+        Args:
+            kwargs: The kwargs to create the model.
+
+        Returns:
+            The decoder class.
+        """
+        with torch.device("meta"):
+            model: PreTrainedModel = AutoModel.from_config(**kwargs)
+        decoder_cls = type(model.get_decoder())
+        logger.debug("Identified decoder class as: %s", decoder_cls)
+        del model
+        return decoder_cls
+
+    def _decorate_cls_for_torch_compile(
+        self,
+        cls: type["PreTrainedModel"],
+        dynamic_arg_dims: dict[str, int] | None,
+        enable_if: Callable[["VllmConfig"], bool],
+        is_encoder: bool,
+    ):
+        """
+        Decorate `cls` to indicate to vLLM that it supports torch compile.
+
+        Args:
+            cls: The PreTrainedModel class to decorate.
+            dynamic_arg_dims: A mapping from argument name to the dynamic dimensions
+                of the argument. If None, default dynamic arg dims will be used. See
+                [`support_torch_compile`][vllm.compilation.decorators.support_torch_compile]
+                for more details.
+            enable_if: A function which takes in the vLLM config and returns whether
+                torch compile should be enabled for this class.
+            is_encoder: Whether the class being decorated is an encoder.
+        """
+        logger.debug(
+            "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s",
+            cls.__name__,
+            "encoder" if is_encoder else "decoder",
+            dynamic_arg_dims,
+        )
+
+        support_torch_compile(
+            dynamic_arg_dims=dynamic_arg_dims,
+            enable_if=enable_if,
+            is_encoder=is_encoder,
+        )(cls)
+
+    def _decorate_for_torch_compile(self, **kwargs: dict):
+        """
+        Decorate the model's decoder class to indicate to vLLM that it supports torch
+        compile if `can_enable_torch_compile` is True.
+
+        Args:
+            kwargs: The kwargs to create the model, which are needed to get the decoder
+                class.
+        """
+        self._decorate_cls_for_torch_compile(
+            cls=self._get_decoder_cls(**kwargs),
+            # Applied to a PreTrainedModel so the batch dimension will exist
+            dynamic_arg_dims=dict[str, int](
+                input_ids=1,  # shape: [1, seq_len]
+                inputs_embeds=1,  # shape: [1, seq_len, hidden_size]
+                position_ids=-1,  # shape: [1, seq_len] or [3, 1, seq_len] for mrope
+            ),
+            enable_if=can_enable_torch_compile,
+            is_encoder=False,
+        )
+
+    def _create_hf_to_vllm_mapper(self):
+        """
+        Create a WeightsMapper to map checkpoint weight names to module qualnames.
+
+        This handles:
+
+        - Transformers weight renaming from `WeightRenaming`
+        - Checkpoints saved with a base model prefix that is not `model`
+        - Checkpoints saved with no base model prefix
+        - Any quantization config specific mappings
+        """
+        self.hf_to_vllm_mapper = WeightsMapper()
+        orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming
+        orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex
+
+        for mapping in get_model_conversion_mapping(self.model):
+            # Handle weights which have been renamed in Transformers
+            if isinstance(mapping, WeightRenaming):
+                orig_to_new_renaming.append(mapping)
+            # TODO: Handle WeightConverter to enable layer merging
+
+        # Handle unexpected weights which should be ignored
+        if self.model._keys_to_ignore_on_load_unexpected is not None:
+            for key in self.model._keys_to_ignore_on_load_unexpected:
+                orig_to_new_regex[re.compile(key)] = None
+
+        # Standardise base model prefix
+        bmp = self.model.base_model_prefix
+        expected_bmp = r"model.\1"
+        # Handle checkpoints saved with different base model prefix
+        if bmp and bmp != "model":
+            different_bmp_pattern = re.compile(rf"^{bmp}\.(.+)")
+            orig_to_new_regex[different_bmp_pattern] = expected_bmp
+        # Handle direct children of self.model which were saved without the model prefix
+        direct_children = chain(
+            self.model.named_children(),
+            self.model.named_parameters(recurse=False),
+            self.model.named_buffers(recurse=False),
+        )
+        model_children = "|".join(name for name, _ in direct_children)
+        missing_bmp_pattern = re.compile(rf"^(?!model\.)(({model_children}).*)")
+        orig_to_new_regex[missing_bmp_pattern] = expected_bmp
+        # Handle weights saved as direct children of self.model which no longer are
+        unexpected_bmp_pattern = re.compile(rf"^(model\.)((?!{model_children}).+)")
+        orig_to_new_regex[unexpected_bmp_pattern] = r"\2"
+        # Handle lm_head which was saved inside the base model
+        nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)")
+        orig_to_new_regex[nested_lm_head_pattern] = r"\2"
+
+        # Apply mapping to quantization config if needed
+        self._maybe_apply_model_mapping()
+
+    def _get_tie_word_embeddings(self):
+        """
+        Check if the model has tied word embeddings.
+        """
+        # Models created with Transformers v4 and v5 will store this in different places
+        tie_word_embeddings_v4 = getattr(self.text_config, "tie_word_embeddings", False)
+        tie_word_embeddings_v5 = getattr(self.config, "tie_word_embeddings", False)
+        return tie_word_embeddings_v4 or tie_word_embeddings_v5
+
+    def pipeline_parallel(self):
+        """
+        Apply the model's pipeline parallelization plan.
+        """
+        if self.pp_group.world_size <= 1:
+            return
+
+        if self.model.supports_pp_plan:
+            module = self.model
+            names = list(module._pp_plan.keys())
+        else:
+            module = self.model.get_decoder()
+            has_parameters = lambda m: next(m.parameters(), None) is not None
+            names = [n for n, c in module.named_children() if has_parameters(c)]
+            tip = get_feature_request_tip(
+                self.model_config.model, self.model_config.trust_remote_code
+            )
+            logger.warning(
+                "%s does not define a pipeline parallel plan. The Transformers "
+                "modeling backend will infer the split from the layers of %s in order "
+                "of declaration and keep parameter-free modules on every rank. This "
+                "may fail if the model's structure is non-standard. %s",
+                type(self.model),
+                type(module),
+                tip,
+            )
+
+        def attrsetter(attr: str) -> Callable[[object, object], None]:
+            """Set a possibly nested attribute, like the inverse of attrgetter."""
+            parent, _, name = attr.rpartition(".")
+
+            def setter(obj: object, value: object):
+                attr_parent = attrgetter(parent)(obj) if parent else obj
+                setattr(attr_parent, name, value)
+
+            return setter
+
+        module_lists = []
+        module_list_idx = None
+        for i, name in enumerate(names):
+            # attrgetter in case the module is nested (e.g. "text_model.layers")
+            if isinstance(attrgetter(name)(module), nn.ModuleList):
+                module_lists.append(name)
+                module_list_idx = i
+
+        if len(module_lists) > 1:
+            raise ValueError(
+                "Pipeline parallel of models with multiple `ModuleList`s "
+                "in the base model are not supported yet!"
+            )
+        if module_list_idx is None:
+            raise ValueError(f"Could not find `ModuleList` in {type(module)}")
+
+        # Layers before module list
+        for name in names[:module_list_idx]:
+            if self.pp_group.is_first_rank or (
+                self._get_tie_word_embeddings() and self.pp_group.is_last_rank
+            ):
+                continue
+            # attrsetter in case the module is nested (e.g. "text_model.embed_tokens")
+            attrsetter(name)(module, PPMissingLayer())
+
+        # Module list
+        start_layer, end_layer = get_pp_indices(
+            self.text_config.num_hidden_layers,
+            self.pp_group.rank_in_group,
+            self.pp_group.world_size,
+        )
+        layers_name = names[module_list_idx]
+        # attrgetter in case the module is nested (e.g. "text_model.layers")
+        layers = attrgetter(layers_name)(module)
+        for i in range(len(layers)):
+            if start_layer <= i and i < end_layer:
+                continue
+            layers[i] = PPMissingLayer()
+
+        # Layers after module list
+        for name in names[module_list_idx + 1 :]:
+            # Modules that should be on last rank
+            if not self.pp_group.is_last_rank:
+                # attrsetter in case the module is nested (e.g. "text_model.norm")
+                attrsetter(name)(module, PPMissingLayer())
+
+    def recursive_replace(self):
+        """Recursively replace modules in the model as needed.
+
+        Currently, this replaces:
+
+        - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul`
+        - Attention QKV projections with a fused `QKVParallelLinear` + split
+        - `nn.Linear` with vLLM's tensor parallel linear classes
+        - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`
+        - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`
+        """
+        tp_plan = self.model.tp_plan or {}
+
+        if not tp_plan and self.tp_group.world_size > 1:
+            tip = get_feature_request_tip(
+                self.model_config.model, self.model_config.trust_remote_code
+            )
+            logger.warning_once(
+                "%s does not define a tensor parallel plan. The Transformers modeling "
+                "backend will shard the model the best it can during graph fusion and "
+                "replicate the rest. This may be suboptimal or fail if the model does "
+                "not fuse cleanly. %s",
+                type(self.model),
+                tip,
+            )
+
+        # Prefix the patterns because we always start from `self.model`
+        tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()}
+        # Detect fusable patterns once per module class (cached, so this is cheap)
+        fusers = Fusers(self.model, self.model_config)
+
+        def register_fusion(fuser: BaseFuser, prefix: str):
+            """Register a fused layer's mappings just before it is built."""
+            orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)
+            self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
+
+            packed_modules_mapping = fuser.packed_modules_mapping
+            self.packed_modules_mapping.update(packed_modules_mapping)
+            if self.quant_config is not None:
+                self.quant_config.packed_modules_mapping.update(packed_modules_mapping)
+
+        def _recursive_replace(module: nn.Module, prefix: str):
+            for child_name, child_module in module.named_children():
+                new_module = child_module
+                qual_name = maybe_prefix(prefix, child_name)
+                if (
+                    isinstance(module, nn.ModuleList)
+                    and len(module) == self.text_config.num_hidden_layers
+                ):
+                    # Populate Eagle3 attrs
+                    self._target_class = type(child_module)
+                    layer_name = qual_name.removeprefix("model.")
+                    self._layer_names[int(child_name)] = layer_name
+                    # MTP weights should not be loaded into the base model
+                    num_hidden_layers = self.text_config.num_hidden_layers
+                    names = (
+                        "n_predict",  # Override from SpeculativeConfig
+                        "num_nextn_predict_layers",  # Most models
+                        "mtp_num_hidden_layers",  # Qwen 3.5
+                    )
+                    n_predict = getattr_iter(self.text_config, names, 0)
+                    for i in range(num_hidden_layers, num_hidden_layers + n_predict):
+                        mtp_prefix = f"{prefix}.{i}."
+                        if mtp_prefix not in self.ignore_unexpected_prefixes:
+                            self.ignore_unexpected_prefixes.append(mtp_prefix)
+                # Replace modules as needed
+                if isinstance(child_module, nn.Linear):
+                    generator = (p for p in tp_plan if re.match(p, qual_name))
+                    pattern = next(generator, None)
+                    # Some weight loaders expect all linear layers to inherit
+                    # LinearBase, so we set a default style which causes any
+                    # unspecified layers to be replaced with ReplicatedLinear
+                    style = tp_plan.get(pattern, "replicate")
+                    new_module = replace_linear_class(
+                        child_module, style, self.quant_config, prefix=qual_name
+                    )
+                elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)):
+                    new_module = replace_conv_class(child_module)
+                elif (fuser := fusers[child_module]) is not None:
+                    register_fusion(fuser, qual_name)
+                    new_module = fuser.fuse(
+                        child_module, qual_name, self.model_config, self.quant_config
+                    )
+                    logger.info_once(fuser.info(child_name))
+                    _recursive_replace(new_module, prefix=qual_name)
+                elif not isinstance(child_module, MoERunner):
+                    # MoERunner can contain aliases of shared experts and gates,
+                    # so we don't want to recurse into it and break weight loading.
+                    _recursive_replace(child_module, prefix=qual_name)
+
+                if new_module is not child_module:
+                    setattr(module, child_name, new_module)
+                    log_replacement(qual_name, child_module, new_module)
+
+        _recursive_replace(self.model, prefix="model")
+
+    def create_attention_instances(self) -> dict[int, Attention]:
+        """
+        Create `Attention` instances to inform KV cache allocation.
+        """
+        text_config = self.text_config
+
+        num_heads = self.model_config.get_num_attention_heads(self.parallel_config)
+        head_size = self.model_config.get_head_size()
+        num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
+        logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)
+
+        # In encoder models, the attention layers will have `is_causal=False`
+        is_encoder = lambda module: not getattr(module, "is_causal", True)
+        has_encoder = lambda model: any(is_encoder(m) for m in model.modules())
+        is_multimodal = lambda config: config != config.get_text_config()
+        # vLLM does not support encoder-decoder models, so if any encoder layer is
+        # found in a text only model, we assume the whole model is an encoder model
+        if has_encoder(self.model) and not is_multimodal(self.config):
+            self.check_version("5.0.0", "encoder models support")
+            attn_type = AttentionType.ENCODER_ONLY
+        else:
+            attn_type = AttentionType.DECODER
+
+        pp_rank = self.pp_group.rank_in_group
+        pp_size = self.pp_group.world_size
+        start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
+
+        attention_instances = {}
+        for i in range(start, end):
+            # Handle interleaved sliding window attention
+            per_layer_sliding_window = None
+            if (
+                hasattr(self.config, "layer_types")
+                and self.config.layer_types[i] == "sliding_attention"
+            ):
+                per_layer_sliding_window = self.config.sliding_window
+
+            attn_cls = (
+                EncoderOnlyAttention
+                if attn_type == AttentionType.ENCODER_ONLY
+                else Attention
+            )
+            attention_instances[i] = attn_cls(
+                num_heads=num_heads,
+                head_size=head_size,
+                # NOTE: We use Llama scale as default, if it's set by
+                # Transformers, it's updated in vllm_attention_forward
+                scale=head_size**-0.5,
+                num_kv_heads=num_kv_heads,
+                cache_config=self.cache_config,
+                quant_config=self.quant_config,
+                logits_soft_cap=logits_soft_cap,
+                per_layer_sliding_window=per_layer_sliding_window,
+                prefix=f"{i}.attn",
+                attn_type=attn_type,
+            )
+        return attention_instances
+
+    def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None):
+        """
+        If a `parameter` is on the `meta` device, then its parent
+        `module` is the original module created by:
+
+        ```python
+        with torch.device("meta"):
+            self.model: "PreTrainedModel" = AutoModel.from_config(...)
+        ```
+        """
+
+        def _init_parameters(module: nn.Module, dtype: torch.dtype | None):
+            for name, param in module.named_parameters(recurse=False):
+                if param.device == torch.device("meta"):
+                    new_param = nn.Parameter(
+                        torch.empty_like(
+                            param.data,
+                            dtype=dtype or self.model_config.dtype,
+                            device=self.device_config.device,
+                        )
+                    )
+                    setattr(module, name, new_param)
+            for child in module.children():
+                _init_parameters(child, dtype)
+
+        _init_parameters(module, dtype)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.get_input_embeddings()(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor | IntermediateTensors:
+        if not self.pp_group.is_first_rank:
+            assert intermediate_tensors is not None
+            input_ids = None
+            inputs_embeds = intermediate_tensors["hidden_states"]
+
+        # Add batch dimension before entering Transformers model
+        if input_ids is not None and input_ids.ndim == 1:
+            # [seq_len] -> [1, seq_len]
+            input_ids = input_ids[None, ...]
+        if inputs_embeds is not None and inputs_embeds.ndim == 2:
+            # [seq_len, hidden_size] -> [1, seq_len, hidden_size]
+            inputs_embeds = inputs_embeds[None, ...]
+        if positions.ndim == 1:
+            # [seq_len] -> [1, seq_len]
+            positions = positions[None, ...]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            use_cache=False,
+            position_ids=positions,
+            attention_instances=self.attention_instances,
+            return_dict=False,
+            **self._output_aux_hidden_states_kwargs,
+            **kwargs,
+        )
+
+        # Remove batch dimension after exiting Transformers model
+        hidden_states = outputs[0][0, ...]
+        if self._output_aux_hidden_states_kwargs:
+            aux_hidden_states = [x[0][0, ...] for x in outputs[1:]]
+
+        if not self.pp_group.is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states})
+
+        if self._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=self.skip_prefixes,
+            skip_substrs=self.skip_substrs,
+            ignore_unexpected_prefixes=self.ignore_unexpected_prefixes,
+            ignore_unexpected_suffixes=self.ignore_unexpected_suffixes,
+        )
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+    @staticmethod
+    def check_version(min_version: str, feature: str):
+        installed = Version(transformers.__version__)
+        required = Version(min_version)
+        if installed < required:
+            raise ImportError(
+                f"Transformers modeling backend requires transformers>={required} "
+                f"for {feature}, but got {installed}"
+            )
+
+    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
+        self.check_version("5.2.0", "Eagle3 support")
+        from transformers.utils.output_capturing import (
+            OutputRecorder,
+            maybe_install_capturing_hooks,
+        )
+
+        # The default value in PreTrainedModel is None
+        if self.model._can_record_outputs is None:
+            self.model._can_record_outputs = {}
+
+        target_class = self._target_class
+        for layer in layers:
+            # layer - 1 because we want the input to the layer
+            layer_name = self._layer_names[layer - 1]
+            layer_key = f"aux_hidden_state_{layer}"
+            aux_hidden_state_i = OutputRecorder(target_class, layer_name=layer_name)
+            self.model._can_record_outputs[layer_key] = aux_hidden_state_i
+            self._output_aux_hidden_states_kwargs[f"output_{layer_key}"] = True
+
+        # Ensure that the capture hooks are installed before dynamo traces the model
+        maybe_install_capturing_hooks(self.model)
+
+    def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        num_layers = self.text_config.num_hidden_layers
+        return (2, num_layers // 2, num_layers - 3)

vllm/model_executor/models/transformers/causal.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend mixin for causal language models."""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration
+from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
+
+if TYPE_CHECKING:
+    import torch
+
+    from vllm.config import VllmConfig
+
+
+class CausalMixin(VllmModelForTextGeneration):
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO
+        super(VllmModelForTextGeneration, self).__init__(
+            vllm_config=vllm_config, prefix=prefix
+        )
+
+        # Tell `Base.load_weights` to skip
+        # `lm_head` if the model has tied word embeddings
+        tie_word_embeddings = self._get_tie_word_embeddings()
+        if tie_word_embeddings:
+            self.skip_prefixes.append("lm_head.")
+
+        if self.pp_group.is_last_rank:
+            self.lm_head = ParallelLMHead(
+                self.text_config.vocab_size,
+                self.text_config.hidden_size,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+            if tie_word_embeddings:
+                self.lm_head = self.lm_head.tie_weights(
+                    self.model.get_input_embeddings()
+                )
+
+            logit_scale = getattr(self.text_config, "logit_scale", 1.0)
+            self.logits_processor = LogitsProcessor(
+                self.text_config.vocab_size, scale=logit_scale
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+
+    def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]:
+        """A thin wrapper around `Base.load_weights` to handle the lm_head bias."""
+
+        lm_head_bias = set()
+
+        def auto_load_lm_head_bias(weights):
+            for name, weight in weights:
+                if name.endswith("lm_head.bias") and self.pp_group.is_last_rank:
+                    self.lm_head._register_bias()
+                    self.lm_head.bias.weight_loader(self.lm_head.bias, weight)
+                    lm_head_bias.add(name)
+                else:
+                    yield name, weight
+
+        return super().load_weights(auto_load_lm_head_bias(weights)) | lm_head_bias
+
+    def compute_logits(self, hidden_states: "torch.Tensor") -> "torch.Tensor | None":
+        logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias)
+        return logits

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

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

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Concrete fusers for the Transformers modeling backend."""
+
+from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser
+from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
+from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
+from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser
+from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
+
+__all__ = [
+    "BaseFuser",
+    "StackedFuser",
+    "GLUFuser",
+    "MoEBlockFuser",
+    "QKVFuser",
+    "RMSNormFuser",
+]

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

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

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLU projection fuser: `act(gate(x)) * up(x)` -> a fused gate/up linear."""
+
+import ast
+import operator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+from torch import fx, nn
+from transformers.activations import ACT2CLS
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import (
+    _ACTIVATION_AND_MUL_REGISTRY,
+    get_act_and_mul_fn,
+)
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear
+from vllm.model_executor.models.transformers.fusers.base import StackedFuser
+from vllm.model_executor.models.transformers.fx_utils import (
+    compile_forward,
+    find_node,
+    is_linear,
+    peel,
+    recover_forward,
+    replace_expr,
+    single_self_call,
+)
+from vllm.model_executor.models.transformers.utils import (
+    log_replacement,
+    replace_linear_class,
+)
+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
+
+logger = init_logger(__name__)
+
+
+CLS2ACT: dict[type, list[str]] = {}
+for _act_name, _act_cls in ACT2CLS.items():
+    if isinstance(_act_cls, tuple):
+        _act_cls = _act_cls[0]
+    CLS2ACT.setdefault(_act_cls, []).append(_act_name)
+
+ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys())
+
+
+@dataclass
+class GLUFuser(StackedFuser):
+    """Fuser for the GLU pattern `act(gate(x)) * up(x)`."""
+
+    act_name: str
+    gate_name: str
+    up_name: str
+    down_name: str | None
+    merged_name: ClassVar[str] = "gate_up_proj"
+    merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
+
+    @property
+    def shards(self) -> list[tuple[str, ShardId]]:
+        return [(self.gate_name, 0), (self.up_name, 1)]
+
+    @classmethod
+    def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool:
+        """Is node `act(gate(x))` where `gate` is linear and `act` is not linear."""
+        return (
+            node.op == "call_module"
+            and not is_linear(node, module)
+            and len(node.args) == 1
+            and isinstance(node.args[0], fx.Node)
+            and is_linear(node.args[0], module)
+        )
+
+    @classmethod
+    def _get_glu_nodes(
+        cls, graph: fx.Graph, module: nn.Module
+    ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None:
+        """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
+        for mul in graph.nodes:
+            if (
+                mul.op == "call_function"
+                and mul.target == operator.mul
+                and len(mul.args) == 2
+                and all(isinstance(arg, fx.Node) for arg in mul.args)
+            ):
+                a, b = mul.args
+                if cls._is_act_of_gate(a, module) and is_linear(b, module):
+                    act, gate, up = a, a.args[0], b
+                elif cls._is_act_of_gate(b, module) and is_linear(a, module):
+                    act, gate, up = b, b.args[0], a
+                else:
+                    continue
+                if (
+                    all(len(args) == 1 for args in (gate.args, up.args))
+                    and isinstance(x := gate.args[0], fx.Node)
+                    and x is up.args[0]
+                ):
+                    return act, gate, up, mul
+        return None
+
+    @staticmethod
+    def _get_act_and_mul_name(act: nn.Module) -> str | None:
+        """Get the name of `act` if it has an `...AndMul` equivalent."""
+        for name in CLS2ACT.get(type(act), []):
+            if name in ACT_AND_MUL_NAMES:
+                return name
+        # nn.GELU is not in ACT2CLS, but could be in model code
+        if type(act) is nn.GELU:
+            return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu"
+        return None
+
+    @classmethod
+    def _get_act_and_mul(cls, act: nn.Module) -> nn.Module:
+        """Get the `...AndMul` equivalent of a Transformers activation module."""
+        if name := cls._get_act_and_mul_name(act):
+            return get_act_and_mul_fn(name)
+        raise ValueError(f"No AndMul equivalent for {type(act)}")
+
+    @classmethod
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None":
+        if (glu_nodes := cls._get_glu_nodes(graph, module)) is None:
+            return None
+        act_node, gate_node, up_node, mul_node = glu_nodes
+
+        gate = module.get_submodule(gate_node.target)
+        up = module.get_submodule(up_node.target)
+        # Shapes must be compatible for a single merged GEMM.
+        if gate.in_features == up.in_features and (gate.bias is None) == (
+            up.bias is None
+        ):
+            predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node
+            down_node = find_node(graph, predicate)
+            return cls(
+                source_cls=type(module).__name__,
+                act_name=act_node.target,
+                gate_name=gate_node.target,
+                up_name=up_node.target,
+                down_name=down_node.target if down_node is not None else None,
+            )
+        return None
+
+    def update_forward(self, module: nn.Module) -> None:
+        """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
+        funcdef, fn = recover_forward(type(module))
+        act_call = single_self_call(funcdef, self.act_name)
+        gate_call = single_self_call(funcdef, self.gate_name)
+        up_call = single_self_call(funcdef, self.up_name)
+        if act_call.args[0] is not gate_call:
+            raise ValueError("activation does not directly wrap the gate")
+        if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]):
+            raise ValueError("gate and up inputs are written differently")
+        muls = [
+            node
+            for node in ast.walk(funcdef)
+            if isinstance(node, ast.BinOp)
+            and isinstance(node.op, ast.Mult)
+            and {id(node.left), id(node.right)} == {id(act_call), id(up_call)}
+        ]
+        if len(muls) != 1:
+            raise ValueError("no multiply of the activation and up projection")
+
+        # act(gate(x)) * up(x) -> act(gate_up(x))
+        assert isinstance(gate_call.func, ast.Attribute)
+        gate_call.func.attr = self.merged_name
+        replace_expr(funcdef, muls[0], act_call)
+        self.fused_forward = compile_forward(funcdef, fn)
+
+    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
+        act = module.get_submodule(self.act_name)
+        if self._get_act_and_mul_name(act) is None:
+            logger.debug("No AndMul equivalent for %s; skipping fusion", type(act))
+            return False
+        return True
+
+    def update_attrs(
+        self,
+        module: nn.Module,
+        prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
+    ) -> None:
+        act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))
+        gate = module.get_submodule(self.gate_name)
+        up = module.get_submodule(self.up_name)
+        merged = MergedColumnParallelLinear(
+            input_size=gate.in_features,
+            output_sizes=[gate.out_features, up.out_features],
+            bias=gate.bias is not None,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, self.merged_name),
+            return_bias=False,
+        )
+        logger.debug(
+            "%s: %s, %s: %s -> %s: %s",
+            self.gate_name,
+            gate,
+            self.up_name,
+            up,
+            self.merged_name,
+            merged,
+        )
+        setattr(module, self.merged_name, merged)
+        setattr(module, self.act_name, act_fn)
+        # Drop the consumed submodules so their (meta) params are not expected.
+        delattr(module, self.gate_name)
+        delattr(module, self.up_name)
+        # If there is a down projection, we know it must be rowwise.
+        if self.down_name is not None:
+            down_prefix = maybe_prefix(prefix, self.down_name)
+            down = module.get_submodule(self.down_name)
+            new_down = replace_linear_class(
+                down, "rowwise", quant_config, prefix=down_prefix
+            )
+            setattr(module, self.down_name, new_down)
+            log_replacement(down_prefix, down, new_down)

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

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing."""
+
+import ast
+import inspect
+import textwrap
+import types
+from collections.abc import Iterator
+from dataclasses import dataclass
+from itertools import chain
+
+import torch
+from torch import fx, nn
+
+from vllm.distributed import tensor_model_parallel_all_gather
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.models.transformers.fx_utils import (
+    find_node,
+    is_op,
+    peel,
+    trace,
+)
+from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk
+
+
+def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
+    """`module`'s own state (i.e. named parameters and buffers)."""
+    return chain(module.named_parameters(), module.named_buffers())
+
+
+def _own_returns(node: ast.AST) -> Iterator[ast.Return]:
+    """`return` statements in `node`'s own scope, not in nested functions."""
+    stack = list(ast.iter_child_nodes(node))
+    while stack:
+        child = stack.pop()
+        if isinstance(child, ast.Return):
+            yield child
+        elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
+            stack.extend(ast.iter_child_nodes(child))
+
+
+def _returns_tuple(cls: type[nn.Module]) -> bool:
+    """Does `cls.forward()` return a tuple?"""
+    try:
+        source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward)))
+        forward = ast.parse(source).body[0]
+    except (OSError, SyntaxError, TypeError, IndexError):
+        return True
+    # Names bound to a tuple literal, e.g. `out = hidden, logits` then `return out`.
+    tuple_names = {
+        target.id
+        for node in ast.walk(forward)
+        if isinstance(node, ast.Assign) and isinstance(node.value, ast.Tuple)
+        for target in node.targets
+        if isinstance(target, ast.Name)
+    }
+
+    def yields_tuple(value: ast.expr | None) -> bool:
+        if isinstance(value, ast.Tuple):
+            return True
+        if isinstance(value, ast.Name):
+            return value.id in tuple_names
+        if isinstance(value, ast.IfExp):
+            return yields_tuple(value.body) or yields_tuple(value.orelse)
+        return False
+
+    return any(yields_tuple(ret.value) for ret in _own_returns(forward))
+
+
+def _is_scalar_gate(module: nn.Module) -> bool:
+    """A linear projecting to a single logit (the shared-expert sigmoid gate)."""
+    weight = getattr(module, "weight", None)
+    return (
+        isinstance(module, nn.Linear)
+        and weight is not None
+        and weight.ndim == 2
+        and weight.shape[0] == 1
+    )
+
+
+def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
+    """Returns the set of nodes reachable from `node` by following `key` edges."""
+    seen: set[fx.Node] = set()
+    stack = [node]
+    while stack:
+        n = stack.pop()
+        if n in seen:
+            continue
+        seen.add(n)
+        stack.extend(getattr(n, key))
+    return seen
+
+
+class SharedExpertMLP(nn.Module):
+    """Wraps an HF shared expert, applying the output gating it is paired with."""
+
+    def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None):
+        super().__init__()
+        self.shared_experts = shared_experts
+        self.gate = gate
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        out = self.shared_experts(hidden_states)
+        if self.gate is not None:
+            out = torch.sigmoid(self.gate(hidden_states)[0]) * out
+        return out
+
+
+def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
+    """Standard MoE block forward.
+
+    Routing and any shared experts are handled inside `self.experts: MoERunner`."""
+    orig_shape = hidden_states.shape
+    hidden_states = hidden_states.reshape(-1, orig_shape[-1])
+    num_tokens = hidden_states.shape[0]
+    is_sequence_parallel = self.experts.moe_config.is_sequence_parallel
+    if is_sequence_parallel:
+        hidden_states = sequence_parallel_chunk(hidden_states)
+    out = self.experts(hidden_states, router_logits=hidden_states)
+    if is_sequence_parallel:
+        out = tensor_model_parallel_all_gather(out, 0)[:num_tokens]
+    return out.reshape(orig_shape)
+
+
+@dataclass
+class MoEBlockFuser:
+    """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""
+
+    gate_name: str
+    scoring_func: str
+    shared_name: str | None
+    shared_gate_name: str | None
+
+    @staticmethod
+    def _match_router(gate: nn.Module) -> str | None:
+        """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`."""
+        if [name for name, _ in named_state(gate)] != ["weight"]:
+            return None
+        graph = trace(gate)
+        if graph is None:
+            return None
+        topk = find_node(graph, lambda n: is_op(n, "topk"))
+        if topk is None:
+            return None
+        # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear.
+        scorers = [
+            n
+            for n in _reaches(topk, "all_input_nodes")
+            if is_op(n, "softmax") or is_op(n, "sigmoid")
+        ]
+        if len(scorers) != 1:
+            return None
+        scorer = scorers[0]
+        if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")):
+            return None
+        return "softmax" if is_op(scorer, "softmax") else "sigmoid"
+
+    @staticmethod
+    def _match_shared_experts(
+        graph: fx.Graph, experts: str
+    ) -> tuple[str | None, str | None]:
+        """Detects the shared expert and its optional gate by dataflow."""
+        experts_predicate = lambda n: n.op == "call_module" and n.target == experts
+        if (experts_node := find_node(graph, experts_predicate)) is None:
+            return None, None
+        from_experts = _reaches(experts_node, "users")
+        for add in graph.nodes:
+            if not is_op(add, "add"):
+                continue
+            operands = [a for a in add.args if isinstance(a, fx.Node)]
+            # Exactly one side is the experts' output; the other is the shared path.
+            sides = [a in from_experts for a in operands]
+            if len(operands) != 2 or sides.count(True) != 1:
+                continue
+            cone = _reaches(operands[sides.index(False)], "all_input_nodes")
+            modules = [n for n in cone if n.op == "call_module" and n.target != experts]
+            # A sigmoid wrapping one of those modules marks the shared-expert gate.
+            gate = next(
+                (
+                    src
+                    for n in cone
+                    if is_op(n, "sigmoid")
+                    and isinstance(src := peel(n.args[0]), fx.Node)
+                    and src in modules
+                ),
+                None,
+            )
+            shared = [n for n in modules if n is not gate]
+            if len(shared) != 1:
+                return None, None
+            return shared[0].target, (gate.target if gate is not None else None)
+        return None, None
+
+    @classmethod
+    def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None":
+        # Standard MoE block returns a single tensor.
+        if _returns_tuple(type(moe_block)):
+            return None
+        # Router: the child that scores + top-k selects.
+        gate_name = scoring_func = None
+        for name, child in moe_block.named_children():
+            if name != experts_name and (func := cls._match_router(child)) is not None:
+                gate_name, scoring_func = name, func
+                break
+        if gate_name is None or scoring_func is None:
+            return None
+        # Shared expert: a child the block adds to the experts' output.
+        shared_name = shared_gate_name = None
+        others = [
+            n
+            for n, _ in moe_block.named_children()
+            if n not in {experts_name, gate_name}
+        ]
+        if others:
+            graph = trace(moe_block)
+            if graph is None:
+                return None
+            shared_name, shared_gate_name = cls._match_shared_experts(
+                graph, experts_name
+            )
+            if shared_gate_name is not None and not _is_scalar_gate(
+                getattr(moe_block, shared_gate_name)
+            ):
+                return None
+        # Fail closed: `rewrite_forward` runs only the experts and the detected
+        # shared expert, so any other stateful child would be dropped.
+        accounted = {experts_name, gate_name, shared_name, shared_gate_name}
+        for name, child in moe_block.named_children():
+            if name not in accounted and next(named_state(child), None) is not None:
+                return None
+        return cls(gate_name, scoring_func, shared_name, shared_gate_name)
+
+    def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
+        """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE."""
+        num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape
+        gate = ReplicatedLinear(
+            hidden_size,
+            num_experts,
+            bias=False,
+            prefix=maybe_prefix(prefix, self.gate_name),
+        )
+        setattr(moe_block, self.gate_name, gate)
+        return gate
+
+    def shared_experts(
+        self, moe_block: nn.Module, prefix: str
+    ) -> SharedExpertMLP | None:
+        """Build the HF shared expert (and its optional gate)
+        as a `SharedExpertMLP` for vLLM's fused MoE."""
+        if self.shared_name is None:
+            return None
+        shared_experts = getattr(moe_block, self.shared_name)
+        gate = None
+        if self.shared_gate_name is not None:
+            hf_gate = getattr(moe_block, self.shared_gate_name)
+            gate = ReplicatedLinear(
+                hf_gate.in_features,
+                hf_gate.out_features,
+                bias=hf_gate.bias is not None,
+                prefix=maybe_prefix(prefix, self.shared_gate_name),
+            )
+            setattr(moe_block, self.shared_gate_name, gate)
+        return SharedExpertMLP(shared_experts, gate)
+
+    def rewrite_forward(self, moe_block: nn.Module) -> None:
+        """Rewrite `moe_block.forward` to route through vLLM's fused MoE."""
+        moe_block.forward = types.MethodType(_moe_block_forward, moe_block)

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

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split."""
+
+import ast
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+from torch import fx, nn
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm.model_executor.models.transformers.fusers.base import StackedFuser
+from vllm.model_executor.models.transformers.fx_utils import (
+    compile_forward,
+    innermost_block,
+    is_linear,
+    recover_forward,
+    replace_expr,
+    single_self_call,
+)
+from vllm.model_executor.models.transformers.utils import (
+    log_replacement,
+    replace_linear_class,
+)
+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
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class QKVFuser(StackedFuser):
+    """Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
+
+    q_name: str
+    k_name: str
+    v_name: str
+    o_name: str | None
+    merged_name: ClassVar[str] = "qkv_proj"
+    merged_cls: ClassVar[str] = "QKVParallelLinear"
+
+    @property
+    def shards(self) -> list[tuple[str, ShardId]]:
+        return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")]
+
+    @classmethod
+    def _get_qkv_nodes(
+        cls, graph: fx.Graph, module: nn.Module
+    ) -> tuple[fx.Node, fx.Node, fx.Node] | None:
+        """Search `graph` for the QKV pattern `q(x), k(x), v(x)`."""
+        by_input: dict[fx.Node, list[fx.Node]] = {}
+        for node in graph.nodes:
+            if (
+                is_linear(node, module)
+                and len(node.args) == 1
+                and not node.kwargs
+                and isinstance(node.args[0], fx.Node)
+                and node.args[0].op == "placeholder"
+            ):
+                by_input.setdefault(node.args[0], []).append(node)
+        triples = [nodes for nodes in by_input.values() if len(nodes) == 3]
+        if len(triples) != 1:
+            return None
+
+        q_node, k_node, v_node = nodes = triples[0]
+        outs = [module.get_submodule(node.target).out_features for node in nodes]
+        if len(set(outs)) == 2:
+            # q is identified as the larger projection (GQA)
+            (q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1)
+            k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2)
+            if module.get_submodule(q_node.target).out_features != max(outs):
+                return None
+        elif len(set(outs)) != 1:
+            return None
+        return q_node, k_node, v_node
+
+    @classmethod
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
+        if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
+            return None
+        q, k, v = qkv_nodes
+        names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
+        attn_width = module.get_submodule(q.target).out_features
+        candidates = [
+            name
+            for name, child in module.named_children()
+            if isinstance(child, nn.Linear)
+            and name not in names.values()
+            and child.in_features == attn_width
+        ]
+        names["o_name"] = candidates[0] if len(candidates) == 1 else None
+        return cls(source_cls=type(module).__name__, **names)
+
+    def update_forward(self, module: nn.Module) -> None:
+        """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
+        funcdef, fn = recover_forward(type(module))
+        calls = [
+            single_self_call(funcdef, name)
+            for name in (self.q_name, self.k_name, self.v_name)
+        ]
+        arg_dumps = {ast.dump(call.args[0]) for call in calls}
+        if len(arg_dumps) != 1:
+            raise ValueError("projection inputs are written differently")
+        # The trace may be partial, so prove projection exclusivity in source:
+        # no other linear child may consume the same input (else the matched
+        # three may not be q, k and v)
+        other_linears = {
+            name
+            for name, child in module.named_children()
+            if isinstance(child, nn.Linear)
+        } - {self.q_name, self.k_name, self.v_name}
+        for node in ast.walk(funcdef):
+            if (
+                isinstance(node, ast.Call)
+                and isinstance(node.func, ast.Attribute)
+                and node.func.attr in other_linears
+                and any(ast.dump(arg) in arg_dumps for arg in node.args)
+            ):
+                raise ValueError("another linear consumes the same input")
+        blocks = [innermost_block(funcdef.body, call) for call in calls]
+        if any(found is None for found in blocks):
+            raise ValueError("projection calls not found in the function body")
+        if len({id(block) for block, _ in blocks}) != 1:
+            raise ValueError("projection calls are in different blocks")
+
+        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
+        names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
+        temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
+        if names & set(temps):
+            raise ValueError("fused temporaries would shadow existing names")
+        merged = f"self.{self.merged_name}"
+        sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
+        template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
+        assign = ast.parse(template).body[0]
+        arg = next(
+            node
+            for node in ast.walk(assign)
+            if isinstance(node, ast.Name) and node.id == "__arg__"
+        )
+        replace_expr(assign, arg, calls[0].args[0])
+        block, index = blocks[0]
+        ast.copy_location(assign, block[index])
+        block.insert(min(index for _, index in blocks), assign)
+        for call, temp in zip(calls, temps):
+            replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
+        self.fused_forward = compile_forward(funcdef, fn)
+
+    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
+        """Shapes must be compatible for a single merged, head-sharded GEMM."""
+        q = module.get_submodule(self.q_name)
+        k = module.get_submodule(self.k_name)
+        v = module.get_submodule(self.v_name)
+        head_size = model_config.get_head_size()
+        compatible = (
+            q.in_features == k.in_features == v.in_features
+            and len({proj.bias is None for proj in (q, k, v)}) == 1
+            and k.out_features == v.out_features
+            and q.out_features % head_size == 0
+            and k.out_features % head_size == 0
+        )
+        if not compatible:
+            logger.debug("%s is not compatible with QKV fusion", type(module))
+        return compatible
+
+    def update_attrs(
+        self,
+        module: nn.Module,
+        prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
+    ) -> None:
+        head_size = model_config.get_head_size()
+        q = module.get_submodule(self.q_name)
+        k = module.get_submodule(self.k_name)
+        merged = QKVParallelLinear(
+            hidden_size=q.in_features,
+            head_size=head_size,
+            total_num_heads=q.out_features // head_size,
+            total_num_kv_heads=k.out_features // head_size,
+            bias=q.bias is not None,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, self.merged_name),
+            return_bias=False,
+        )
+        logger.debug(
+            "%s: %s, %s: %s, %s: %s -> %s: %s",
+            self.q_name,
+            q,
+            self.k_name,
+            k,
+            self.v_name,
+            module.get_submodule(self.v_name),
+            self.merged_name,
+            merged,
+        )
+        setattr(module, self.merged_name, merged)
+        # Drop the consumed submodules so their (meta) params are not expected.
+        for name in (self.q_name, self.k_name, self.v_name):
+            delattr(module, name)
+        # If there is an output projection, we know it must be rowwise.
+        if self.o_name is not None:
+            o_proj_prefix = maybe_prefix(prefix, self.o_name)
+            o_proj = module.get_submodule(self.o_name)
+            new_o = replace_linear_class(
+                o_proj, "rowwise", quant_config, prefix=o_proj_prefix
+            )
+            setattr(module, self.o_name, new_o)
+            log_replacement(o_proj_prefix, o_proj, new_o)

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

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+from torch import fx, nn
+
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+)
+from vllm.distributed.parallel_state import model_parallel_is_initialized
+from vllm.distributed.utils import split_tensor_along_last_dim
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
+from vllm.model_executor.models.transformers.fusers.base import BaseFuser
+from vllm.model_executor.models.transformers.fx_utils import (
+    find_node,
+    forward_input_count,
+    is_op,
+    peel,
+    trace,
+)
+
+if TYPE_CHECKING:
+    from vllm.config.model import ModelConfig
+    from vllm.model_executor.layers.quantization import QuantizationConfig
+
+
+def _is_squared(node: object, x: fx.Node) -> bool:
+    """`x**2`, `x.square()` or `x * x`, through any dtype casts."""
+    node = peel(node)
+    if is_op(node, "pow"):
+        base, exp = node.args
+        return peel(base) is x and exp == 2
+    if is_op(node, "square"):
+        return peel(node.args[0]) is x
+    if is_op(node, "mul"):
+        a, b = node.args
+        return peel(a) is x and peel(b) is x
+    return False
+
+
+def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
+    """eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
+    add = peel(rsqrt.args[0])
+    if not is_op(add, "add"):
+        return None
+    consts = [a for a in add.args if isinstance(a, (int, float))]
+    nodes = [a for a in add.args if isinstance(a, fx.Node)]
+    if len(consts) != 1 or len(nodes) != 1:
+        return None
+    mean = peel(nodes[0])
+    if not is_op(mean, "mean"):
+        return None
+    if not _is_squared(mean.args[0], x):
+        return None
+    return float(consts[0])
+
+
+def _is_one_plus(node: object) -> bool:
+    """`1 + weight` in either operand order (marks a zero-centered weight)."""
+    node = peel(node)
+    if not is_op(node, "add"):
+        return False
+    return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
+
+
+def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool:
+    """Does the forward compute anything after `node` before returning?"""
+    output = find_node(graph, lambda n: n.op == "output")
+    if output is None or not output.args:
+        return False
+    return peel(output.args[0]) is not node
+
+
+class TPAwareNormMixin(nn.Module):
+    """Mixin for RMSNorms that reconstructs a TP-sharded input before normalizing."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if model_parallel_is_initialized():
+            self.tp_size = get_tensor_model_parallel_world_size()
+            self.tp_rank = get_tensor_model_parallel_rank()
+        else:
+            self.tp_size, self.tp_rank = 1, 0
+
+    def forward(
+        self, x: torch.Tensor, residual: torch.Tensor | None = None
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if self.tp_size > 1 and x.shape[-1] < (full := self.weight.shape[0]):
+            if x.shape[-1] * self.tp_size != full:
+                raise ValueError(
+                    f"Cannot gather norm of width {full}: a TP-sharded input of "
+                    f"width {x.shape[-1]} does not tile it evenly across "
+                    f"{self.tp_size} ranks (replicated or uneven sharding)."
+                )
+            x = tensor_model_parallel_all_gather(x.contiguous())
+            x = super().forward(x)
+            splits = split_tensor_along_last_dim(x, num_partitions=self.tp_size)
+            return splits[self.tp_rank]
+        return super().forward(x, residual)
+
+
+class TPAwareRMSNorm(TPAwareNormMixin, RMSNorm):
+    """`RMSNorm` that reconstructs a TP-sharded input before normalizing."""
+
+
+class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm):
+    """`GemmaRMSNorm` that reconstructs a TP-sharded input before normalizing."""
+
+
+@dataclass
+class RMSNormFuser(BaseFuser):
+    """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights."""
+
+    zero_centered: bool
+    """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
+    source_cls: str
+    """Class name of the norm this was matched from (for logging)."""
+
+    def info(self, name: str) -> str:
+        norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
+        return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)"
+
+    @classmethod
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
+        """Match a graph to the RMSNorm pattern, returning a fuser if found."""
+        if forward_input_count(type(module)) != 1:
+            return None
+        x = find_node(graph, lambda n: n.op == "placeholder")
+        if x is None:
+            return None
+        # Handle native torch `rms_norm` op.
+        rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
+        if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
+            if _has_trailing_compute(graph, rms_norm):
+                return None
+            return cls(zero_centered=False, source_cls=type(module).__name__)
+        # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern.
+        # The rsqrt over the mean-square variance is the spine of the norm.
+        rsqrt = None
+        for node in graph.nodes:
+            if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
+                rsqrt = node
+                break
+        if rsqrt is None:
+            return None
+        # The `x * rsqrt(...)` normalize multiply.
+        normalize = find_node(
+            graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
+        )
+        if normalize is None:
+            return None
+        # An optional later `weight * normalized` (or `(1 + weight) * normalized`).
+        tail, zero_centered = normalize, False
+        for node in graph.nodes:
+            if not is_op(node, "mul") or node is normalize:
+                continue
+            operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
+            if len(operands) == 2 and normalize in operands:
+                weight = next(o for o in operands if o is not normalize)
+                tail, zero_centered = node, _is_one_plus(weight)
+                break
+        # The norm must be the last compute in forward, or it is not a pure norm.
+        if _has_trailing_compute(graph, tail):
+            return None
+        return cls(zero_centered=zero_centered, source_cls=type(module).__name__)
+
+    @staticmethod
+    def _eps_from_graph(graph: fx.Graph) -> float | None:
+        """Extract the `eps` constant from the graph, if present."""
+        if (x := find_node(graph, lambda n: n.op == "placeholder")) is None:
+            return None
+        fused = find_node(graph, lambda n: is_op(n, "rms_norm"))
+        if fused is not None and fused.args and peel(fused.args[0]) is x:
+            args, kwargs = fused.args, fused.kwargs
+            eps = args[3] if len(args) > 3 else kwargs.get("eps")
+            return eps if isinstance(eps, (int, float)) else None
+        for node in graph.nodes:
+            if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
+                return eps
+        return None
+
+    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
+        return True
+
+    def fuse(
+        self,
+        module: nn.Module,
+        prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
+    ) -> nn.Module:
+        """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
+        weight = getattr(module, "weight", None)
+        hidden_size = (
+            weight.size(0) if weight is not None else model_config.get_hidden_size()
+        )
+        graph = trace(module)
+        eps = self._eps_from_graph(graph) if graph is not None else None
+        if eps is None:
+            # If eps not in graph, match torch behaviour.
+            dtype = weight.dtype if weight is not None else model_config.dtype
+            eps = torch.finfo(dtype).eps
+        if self.zero_centered:
+            return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
+        has_weight = weight is not None
+        return TPAwareRMSNorm(
+            hidden_size=hidden_size,
+            eps=eps,
+            has_weight=has_weight,
+            dtype=weight.dtype if has_weight else None,
+        )

vllm/model_executor/models/transformers/fx_utils.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""fx tracing and forward-source rewriting for the Transformers backend fusers.
+
+A small engine, independent of any particular pattern: trace a module's forward
+with `torch.fx` (tolerating a partial graph), inspect the resulting nodes, and
+rewrite the forward's *source* (AST) so only matched calls change while the rest
+stays live Python. `fusion.py` builds the concrete fusion patterns on top.
+"""
+
+import ast
+import inspect
+import operator
+import textwrap
+from collections.abc import Callable
+
+import torch
+from torch import fx, nn
+from torch.nn import functional as F
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def _infer_len(node: fx.Node) -> int | None:
+    """Concrete length of a proxy's value, inferred from its node chain.
+
+    Lets tracing pass through the shape unpacks and `*`-splats (e.g.
+    `(*input_shape, -1, head_dim)`) that precede the patterns in HF attention.
+    """
+    # `x.shape` has the rank of `x`, when known
+    if (
+        node.op == "call_function"
+        and node.target is getattr
+        and node.args[1] == "shape"
+        and (rank := _rank(node.args[0])) is not None
+    ):
+        return rank
+    # Slices of known-length values
+    if node.op == "call_function" and node.target is operator.getitem:
+        src_len = _infer_len(node.args[0])
+        index = node.args[1]
+        if src_len is not None and isinstance(index, slice):
+            return len(range(*index.indices(src_len)))
+    return None
+
+
+def _rank(node: fx.Node) -> int | None:
+    """The tensor rank of `node`'s value, if known."""
+    # vLLM always feeds the model [1, seq_len, hidden_size] hidden states
+    if node.op == "placeholder" and node.target == "hidden_states":
+        return 3
+    return None
+
+
+class _SizedProxy(fx.Proxy):
+    """Proxy whose `len` is inferred from the graph (see `_infer_len`)."""
+
+    def __len__(self) -> int:
+        length = _infer_len(self.node)
+        if length is None:
+            return super().__len__()
+        return length
+
+
+class _AllLeafTracer(fx.Tracer):
+    """Tracer that treats every submodule as a leaf.
+
+    Each child stays one `call_module` node, so matching sees the module's own
+    forward structure (activations aren't decomposed into e.g. `sigmoid * x`).
+    `iter` traces through the leading shape unpacks (see `_infer_len`); anything
+    else untraceable ends the trace early and the partial graph is matched.
+    """
+
+    def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:
+        return True
+
+    def proxy(self, node: fx.Node) -> fx.Proxy:
+        return _SizedProxy(node, self)
+
+    def iter(self, obj: fx.Proxy):
+        length = _infer_len(obj.node)
+        if length is None:
+            return super().iter(obj)
+        return iter([obj[i] for i in range(length)])
+
+
+def trace(module: nn.Module) -> fx.Graph | None:
+    """Trace `module.forward`, returning the partial graph on failure.
+
+    The graph is only evidence for matching, and the patterns sit at the top of
+    their forwards, so a trace that fails partway can still be matched."""
+    tracer = _AllLeafTracer()
+    try:
+        return tracer.trace(module)
+    except Exception as exc:
+        logger.debug("Could not fully trace %s: %s", type(module), exc)
+        return getattr(tracer, "graph", None)
+
+
+def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]:
+    """Parse the source of `cls.forward`, ready for rewriting."""
+    fn = inspect.unwrap(cls.forward)
+    if fn.__code__.co_freevars:
+        raise ValueError("forward is a closure")
+    tree = ast.parse(textwrap.dedent(inspect.getsource(fn)))
+    funcdef = tree.body[0]
+    if not isinstance(funcdef, ast.FunctionDef):
+        raise ValueError("source is not a plain function definition")
+    # `fn` is already unwrapped; don't re-apply its decorators
+    funcdef.decorator_list.clear()
+    # Annotations may not evaluate outside the defining module (e.g. with
+    # postponed evaluation); they're not needed at runtime
+    funcdef.returns = None
+    args = funcdef.args
+    for arg in (
+        *args.posonlyargs,
+        *args.args,
+        *args.kwonlyargs,
+        *filter(None, (args.vararg, args.kwarg)),
+    ):
+        arg.annotation = None
+    # Recompiling outside the class body would break name mangling
+    for node in ast.walk(funcdef):
+        name = getattr(node, "attr", None) or getattr(node, "id", None)
+        if name and name.startswith("__") and not name.endswith("__"):
+            raise ValueError(f"{name} would be name mangled")
+    return funcdef, fn
+
+
+def forward_input_count(cls: type[nn.Module]) -> int:
+    """The number of tensor inputs `cls.forward` declares, excluding `self` and
+    any `*args`/`**kwargs`. Read from the signature, so it is independent of
+    whether the trace completes (unlike counting placeholders)."""
+    try:
+        params = list(inspect.signature(cls.forward).parameters.values())[1:]
+    except (ValueError, TypeError):
+        return 1  # uninspectable: assume a single input and let matching decide
+    fixed = (
+        inspect.Parameter.POSITIONAL_ONLY,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        inspect.Parameter.KEYWORD_ONLY,
+    )
+    return sum(1 for p in params if p.kind in fixed)
+
+
+def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:
+    """Compile `funcdef` in `fn`'s module so tracebacks point at the source."""
+    module = ast.Module(body=[funcdef], type_ignores=[])
+    ast.fix_missing_locations(module)
+    ast.increment_lineno(module, fn.__code__.co_firstlineno - 1)
+    code = compile(module, fn.__code__.co_filename, "exec")
+    namespace: dict = {}
+    exec(code, fn.__globals__, namespace)
+    return namespace[funcdef.name]
+
+
+def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call:
+    """The unique `self.<name>(arg)` call in `funcdef`.
+
+    Raises unless `name` appears exactly once, as such a call, so the source
+    rewrite agrees with the fx match.
+    """
+    uses = [
+        node
+        for node in ast.walk(funcdef)
+        if isinstance(node, ast.Attribute) and node.attr == name
+    ]
+    if len(uses) != 1:
+        raise ValueError(f"{name} is referenced {len(uses)} times")
+    calls = [
+        node
+        for node in ast.walk(funcdef)
+        if isinstance(node, ast.Call)
+        and node.func is uses[0]
+        and len(node.args) == 1
+        and not isinstance(node.args[0], ast.Starred)
+        and not node.keywords
+    ]
+    if (
+        len(calls) != 1
+        or not isinstance(uses[0].value, ast.Name)
+        or uses[0].value.id != "self"
+    ):
+        raise ValueError(f"{name} is not a single-argument call on self")
+    return calls[0]
+
+
+def innermost_block(
+    block: list[ast.stmt], node: ast.AST
+) -> tuple[list[ast.stmt], int] | None:
+    """The innermost statement list containing `node`, and the index within."""
+    for index, stmt in enumerate(block):
+        if not any(child is node for child in ast.walk(stmt)):
+            continue
+        child_blocks = [
+            getattr(stmt, fld, None) for fld in ("body", "orelse", "finalbody")
+        ]
+        child_blocks += [h.body for h in getattr(stmt, "handlers", [])]
+        child_blocks += [c.body for c in getattr(stmt, "cases", [])]
+        for child_block in child_blocks:
+            if (
+                isinstance(child_block, list)
+                and child_block
+                and (found := innermost_block(child_block, node)) is not None
+            ):
+                return found
+        return block, index
+    return None
+
+
+def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None:
+    """Replace the expression `old` (by identity) with `new` within `module`."""
+
+    class _Replacer(ast.NodeTransformer):
+        def visit(self, node: ast.AST) -> ast.AST:
+            if node is old:
+                return new
+            return super().generic_visit(node)
+
+    _Replacer().visit(module)
+
+
+def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None:
+    """The first node in `graph` matching `predicate`, or `None`."""
+    return next((n for n in graph.nodes if predicate(n)), None)
+
+
+def is_linear(node: fx.Node, module: nn.Module) -> bool:
+    """Is node `nn.Linear.__call__()`."""
+    return node.op == "call_module" and isinstance(
+        module.get_submodule(node.target), nn.Linear
+    )
+
+
+_DTYPE_CASTS = frozenset({"to", "float", "double", "half", "bfloat16", "type_as"})
+
+
+def peel(node: object) -> object:
+    """Strip dtype-cast wrappers (`.to(...)`, `.float()`, `.type_as(...)`)."""
+    while (
+        isinstance(node, fx.Node)
+        and node.op == "call_method"
+        and node.target in _DTYPE_CASTS
+    ):
+        node = node.args[0]
+    return node
+
+
+def is_fn(node: object, target: Callable) -> bool:
+    """Is node `<target>()`."""
+    return (
+        isinstance(node, fx.Node)
+        and node.op == "call_function"
+        and node.target is target
+    )
+
+
+def is_method(node: object, name: str) -> bool:
+    """Is node `.<name>()`."""
+    return (
+        isinstance(node, fx.Node) and node.op == "call_method" and node.target == name
+    )
+
+
+def is_op(node: object, name: str) -> bool:
+    """
+    Is node `torch.<name>()`, `F.<name>()`, `operator.<name>()`, or `Tensor.<name>()`.
+    """
+    return any(
+        is_fn(node, getattr(module, name, None)) for module in (torch, F, operator)
+    ) or (hasattr(torch.Tensor, name) and is_method(node, name))

vllm/model_executor/models/transformers/legacy.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend mixin for legacy models."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+
+class LegacyMixin:
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+
+        # Skip unsupported/unwanted output embeddings layers
+        self.skip_prefixes.extend(
+            [
+                "model.lm_head.",
+                "model.predictions.",
+                "model.qa_outputs.",
+                "model.embeddings_project.",
+                "model.discriminator_predictions.",
+            ]
+        )
+
+        # Some encoder models have the position_ids buffer in the checkpoint.
+        # vLLM will always pass position_ids as an argument, so we skip loading
+        # the buffer if it exists
+        self.skip_substrs.append("position_ids")
+
+        # Some encoder models have the bias of the final classifier layer
+        # in the checkpoint. vLLM does not use this bias, so we skip loading
+        # it if it exists
+        self.skip_substrs.append("score.bias")
+
+        # roberta-like models an extra padding in positions.
+        # FIXME(Isotr0py): This is quite hacky for roberta edge case,
+        # we should find a better way to handle this.
+        self.is_roberta = "roberta" in self.text_config.model_type
+        self.padding_idx = self.text_config.pad_token_id
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if self.is_roberta:
+            # RoBERTa positions start at padding_idx + 1.
+            # Non-in-place add to avoid mutating the persistent GPU buffer --
+            # in-place += would accumulate on CUDA graph padding slots.
+            positions = positions + self.padding_idx + 1
+        return super().forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )

vllm/model_executor/models/transformers/moe.py

@@ -0,0 +1,335 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend mixin for Mixture of Experts (MoE) models."""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn as nn
+
+from vllm.config.utils import getattr_iter
+from vllm.distributed import get_dp_group, get_ep_group
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.custom_op import PluggableLayer
+from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts
+from vllm.model_executor.models.interfaces import MixtureOfExperts
+from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from .utils import log_replacement
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class TransformersMoEState:
+    topk_ids: torch.Tensor | None = None
+    is_sequence_parallel: bool = False
+
+
+# --8<-- [start:transformers_fused_moe]
+@PluggableLayer.register("transformers_fused_moe")
+class TransformersMoERunner(MoERunner):
+    """Custom FusedMoE for the Transformers modeling backend."""
+
+    # --8<-- [end:transformers_fused_moe]
+    def __init__(self, *args, moe_state: TransformersMoEState, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._moe_state = moe_state
+        self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """In Transformers `experts.forward` will have this signature.
+
+        We discard any extra kwargs because we cannot use them here."""
+        # Note: we need to forward through a custom op so the topk_ids
+        # can be transferred without interfering with cudagraphs.
+        return torch.ops.vllm.transformers_moe_forward(
+            hidden_states,
+            topk_ids.to(torch.int32),
+            topk_weights.to(torch.float32),
+            self.layer_name,
+        )
+
+    def _forward_super(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        return super().forward(hidden_states, topk_weights)
+
+
+def _transformers_moe_forward(
+    hidden_states: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    """Store the `topk_ids` in the layer and call the actual forward."""
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    self._moe_state.topk_ids = topk_ids
+    return self._forward_super(hidden_states, topk_weights)
+
+
+def _transformers_moe_forward_fake(
+    hidden_states: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
+direct_register_custom_op(
+    op_name="transformers_moe_forward",
+    op_func=_transformers_moe_forward,
+    mutates_args=["hidden_states"],
+    fake_impl=_transformers_moe_forward_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
+
+
+class TransformersRoutedExperts(RoutedExperts):
+    def get_expert_mapping(
+        self, include_fused: bool = False
+    ) -> list[tuple[str, str, int, str]]:
+        common_names = ("gate_proj", "down_proj", "up_proj")
+        common_map = super().get_expert_mapping(*common_names, include_fused)
+        mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused)
+        if not include_fused:
+            return common_map + mixtral_map
+        common_fused, common_unfused = common_map[:3], common_map[3:]
+        mixtral_fused, mixtral_unfused = mixtral_map[:3], mixtral_map[3:]
+        return common_fused + mixtral_fused + common_unfused + mixtral_unfused
+
+
+class MoEMixin(MixtureOfExperts):
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        self.check_version("5.0.0", "MoE models support")
+        # Skip MixtureOfExperts.__init__ and call the next class in MRO
+        super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix)
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ):
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for moe_block in self.mlp_layers:
+            moe_block.n_local_physical_experts = num_local_physical_experts
+            moe_block.n_physical_experts = num_physical_experts
+            moe_block.n_redundant_experts = self.num_redundant_experts
+            moe_block.experts.update_expert_map()
+
+    def recursive_replace(self):
+        """Initialize the MoE layers."""
+        experts_name = "experts"
+        text_config = self.text_config
+
+        # Positional arguments
+        num_experts = self.model_config.get_num_experts()
+        top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None)
+        assert top_k is not None
+        hidden_size = text_config.hidden_size
+        intermediate_size = getattr_iter(
+            text_config, ["moe_intermediate_size", "intermediate_size"], None
+        )
+        assert intermediate_size is not None
+
+        num_shared_experts = getattr_iter(
+            text_config,
+            [
+                "n_shared_experts",  # DeepSeek, Docs, GLM
+                "moe_num_shared_experts",  # Aria, Ernie
+            ],
+            0,
+        )
+
+        # Unused kwargs since we use custom_routing_function:
+        # - `scoring_func` and `e_score_correction_bias` only used for grouped
+        #    topk routing inside vLLM and are non-trivial to infer
+        #    and hard code `use_grouped_topk=False`
+        # - `renormalize` passed anyway because it's easy to infer
+        # - `num_expert_group` and `topk_group` used for inferring expert
+        #    placement strategy in FusedMoE
+        # - `apply_router_weight_on_input` is already applied in Transformers
+        renormalize = getattr(text_config, "norm_topk_prob", top_k > 1)
+        num_expert_group = getattr(text_config, "n_group", None)
+        topk_group = getattr(text_config, "topk_group", None)
+
+        # MoE activation function
+        activation = "silu"
+        wrapped_arch = self.config.architectures[0].lower()
+        if "gptoss" in wrapped_arch:
+            activation = "swigluoai"
+
+        # Expert parallel load balancing kwargs
+        enable_eplb = self.parallel_config.enable_eplb
+        num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
+
+        # MixtureOfExperts mixin settings
+        ep_size = get_ep_group().world_size
+
+        self.mlp_layers = []  # Used for MixtureOfExperts methods
+        self.moe_layers = []
+        self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
+        self.num_logical_experts = num_experts
+        self.num_physical_experts = num_experts + num_redundant_experts
+        self.num_local_physical_experts = self.num_physical_experts // ep_size
+        self.num_routed_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_redundant_experts = num_redundant_experts
+
+        # Recursively fuse MoE layers
+        def _recursive_replace(module: nn.Module, prefix: str):
+            for child_name, child_module in module.named_children():
+                qual_name = maybe_prefix(prefix, child_name)
+                # Naive implementations will have experts as ModuleList
+                is_modulelist = isinstance(child_module, nn.ModuleList)
+                # Packed implementations will have experts as 3D tensors of shapes like:
+                # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
+                # down_proj = (num_experts, intermediate_size, hidden_size)
+                params = list(child_module.parameters())
+                is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
+                if child_name == experts_name and (is_modulelist or is_3d):
+                    # Alias for readability
+                    moe_block = module
+                    experts = child_module
+                    # Class of the fused block (parent of gate/experts/shared)
+                    moe_block_cls = type(moe_block).__name__
+                    experts_cls = type(experts).__name__
+                    # Do the experts have biases
+                    has_bias = False
+                    for experts_param_name, _ in experts.named_parameters():
+                        if "bias" in experts_param_name:
+                            has_bias = True
+                            break
+                    # If the config does not specify num_shared_experts, but
+                    # the model has shared experts, we assume there is one.
+                    if self.num_shared_experts == 0:
+                        for moe_block_param_name, _ in moe_block.named_parameters():
+                            if "shared_expert" in moe_block_param_name:
+                                self.num_shared_experts = 1
+                                break
+
+                    kwargs: dict[str, Any] = dict(
+                        num_experts=num_experts,
+                        top_k=top_k,
+                        hidden_size=hidden_size,
+                        intermediate_size=intermediate_size,
+                        renormalize=renormalize,
+                        use_grouped_topk=False,
+                        quant_config=self.quant_config,
+                        prefix=qual_name,
+                        activation=activation,
+                        enable_eplb=enable_eplb,
+                        num_redundant_experts=num_redundant_experts,
+                        has_bias=has_bias,
+                        routed_experts_cls=TransformersRoutedExperts,
+                    )
+                    fuser = MoEBlockFuser.match(moe_block, experts_name)
+                    if self.num_expert_groups <= 1 and fuser is not None:
+                        # MoE block forward is fully replaced.
+                        # gate/router and shared expert (if any) runs in FusedMoE.
+                        kwargs |= dict(
+                            scoring_func=fuser.scoring_func,
+                            is_sequence_parallel=(
+                                self.parallel_config.use_sequence_parallel_moe
+                            ),
+                            gate=fuser.gate(moe_block, prefix),
+                            shared_experts=fuser.shared_experts(moe_block, prefix),
+                        )
+                        fuser.rewrite_forward(moe_block)
+                        routed = "gate + experts"
+                        if fuser.shared_name:
+                            routed += " + shared experts"
+                        logger.info_once(
+                            "Fused: %s (%s) -> FusedMoE (internal routing)",
+                            routed,
+                            moe_block_cls,
+                        )
+                    else:
+                        # MoE block forward is unmodified.
+                        # gate/router and shared expert (if any) runs in Transformers.
+                        # We then smuggle the topk_ids in using a custom op.
+                        moe_state = TransformersMoEState()
+
+                        def custom_routing_function(
+                            hidden_states: torch.Tensor,
+                            gating_output: torch.Tensor,
+                            topk: int,
+                            renormalize: bool,
+                            moe_state: TransformersMoEState,
+                        ):
+                            """Return `topk_weights` from `gating_output` and the
+                            `topk_ids` we stored in the layer earlier."""
+                            topk_weights = gating_output
+                            topk_ids = moe_state.topk_ids
+                            assert topk_ids is not None
+                            # Handle all gather in expert parallel
+                            if topk_ids.size(0) != hidden_states.size(0):
+                                dp_metadata = get_forward_context().dp_metadata
+                                sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+                                is_sp = moe_state.is_sequence_parallel
+                                group = get_ep_group() if is_sp else get_dp_group()
+                                assert sizes[group.rank_in_group] == topk_ids.shape[0]
+                                (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes)
+                            return topk_weights, topk_ids
+
+                        kwargs |= dict(
+                            num_expert_group=num_expert_group,
+                            topk_group=topk_group,
+                            custom_routing_function=partial(
+                                custom_routing_function, moe_state=moe_state
+                            ),
+                            runner_cls=TransformersMoERunner,
+                            runner_args={"moe_state": moe_state},
+                        )
+                        logger.info_once(
+                            "Fused: experts (%s) -> FusedMoE (external routing)",
+                            experts_cls,
+                        )
+                    fused_experts = FusedMoE(**kwargs)
+                    moe_block.experts = fused_experts
+                    log_replacement(qual_name, experts, fused_experts)
+                    # Update MixtureOfExperts mixin state
+                    self.mlp_layers.append(moe_block)
+                    self.moe_layers.append(fused_experts)
+                else:
+                    _recursive_replace(child_module, prefix=qual_name)
+
+        _recursive_replace(self.model, prefix="model")
+        self.num_moe_layers = len(self.moe_layers)
+        # Continue with the replacement of layers in Base
+        super().recursive_replace()

vllm/model_executor/models/transformers/multimodal.py

@@ -0,0 +1,521 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend mixin for multi-modal models."""
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+import torch
+from transformers import AutoModel
+
+from vllm.compilation.decorators import should_torch_compile_mm_encoder
+from vllm.config.utils import getattr_iter
+from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal
+from vllm.multimodal import MultiModalKwargsItems
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldConfig,
+    PlaceholderRange,
+)
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    MultiModalDataItems,
+)
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    ProcessorInputs,
+    TimingContext,
+)
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from transformers import BatchFeature, PreTrainedModel
+
+    from vllm.config import VllmConfig
+    from vllm.config.multimodal import BaseDummyOptions
+
+logger = init_logger(__name__)
+
+_MODALITY_TO_TOKEN_TYPE_ID = {"image": 1, "video": 2, "audio": 3}
+
+
+class MultiModalProcessingInfo(BaseProcessingInfo):
+    def get_supported_mm_limits(self):
+        return {"image": None}
+
+    def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
+        return {"image": self.get_max_image_tokens()}
+
+    def get_max_image_tokens(self) -> int:
+        width, height = self.get_max_image_size()
+        processor = self.get_hf_processor()
+        multimodal_config = self.ctx.model_config.multimodal_config
+        mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
+        mm_tokens = processor._get_num_multimodal_tokens(
+            image_sizes=([height, width],), **mm_processor_kwargs
+        )
+        image_tokens = mm_tokens["num_image_tokens"][0]
+        return image_tokens
+
+    def get_max_image_size(self):
+        return 10_000, 10_000  # hardcode for arbitrary very large size
+
+
+class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_images = mm_counts.get("image", 0)
+
+        processor = self.info.get_hf_processor()
+        if "gemma3" in processor.__class__.__name__.lower():
+            image_token = processor.boi_token
+        else:
+            image_token = getattr(processor, "image_token", "")
+        return image_token * num_images
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, "BaseDummyOptions"],
+    ) -> MultiModalDataDict:
+        num_images = mm_counts.get("image", 0)
+
+        target_width, target_height = self.info.get_max_image_size()
+
+        image_overrides = mm_options.get("image")
+
+        return {
+            "image": self._get_dummy_images(
+                width=target_width,
+                height=target_height,
+                num_images=num_images,
+                overrides=image_overrides,
+            ),
+        }
+
+
+class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ):
+        """
+        Given the original multi-modal items for this modality
+        and HF-processed data, output the updates to perform.
+
+        The information returned by this method is used to update token inputs
+        which bypass the HF processor. It is also used to update the output of
+        HF processor if the HF process does not apply prompt updates to text
+        inputs.
+
+        Moreover, this information is critical to determine the token positions
+        in order to construct  :class:`~vllm-multimodal.input.PlaceholderRange`
+        for each multi-modal item.
+        """
+        return None
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: "BatchFeature",
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        # HF Processors always return a mask but vLLM doesn't need it
+        hf_inputs.pop("attention_mask", None)
+        num_image_patches = hf_inputs.get("num_image_patches")
+        mm_fields = {
+            key: MultiModalFieldConfig.flat_from_sizes("image", num_image_patches)
+            for key in hf_inputs
+        }
+        mm_fields["image_embeds"] = MultiModalFieldConfig.flat_from_sizes(
+            "image", num_image_patches
+        )
+
+        # Keep these as batched, as they always have batch size as first dim
+        mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image")
+        mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image")
+        mm_fields["num_image_patches"] = MultiModalFieldConfig.batched(
+            "image", keep_on_cpu=True
+        )
+        return mm_fields
+
+    def _get_hf_mm_data(
+        self,
+        mm_items: MultiModalDataItems,
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        """
+        In contrast to the base class, this method always adds
+        `return_mm_token_type_ids` to the processor data
+        """
+        processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
+        processor_data["return_mm_token_type_ids"] = True
+        return processor_data, passthrough_data
+
+    def apply(
+        self,
+        inputs: ProcessorInputs,
+        timing_ctx: TimingContext,
+    ) -> MultiModalInput:
+        """
+        Process multi-modal inputs to be used in vLLM.
+
+        Apply HF Processor on prompt text and multi-modal data together,
+        outputting token IDs and processed tensors.
+        """
+        prompt = inputs.prompt
+        mm_items = inputs.mm_data_items
+        hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs
+        tokenization_kwargs = inputs.tokenization_kwargs
+
+        with timing_ctx.record("apply_hf_processor"):
+            hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+            if not isinstance(prompt, str):
+                # the prompt is the tokenized ids which is not supported
+                # by the hf_processor, which is why we would need to decode the ids
+                # into string
+                prompt = hf_processor.decode(prompt)
+
+            # Bypass cached processor and always apply to the full set of mm inputs
+            # NOTE: we can't just set caching=False because base class method
+            # transforms outputs to `MultiModalKwargs` which is not going to
+            # work for Transformers. We have a lot of logic tied to
+            # `mm_tokens_per_modality` below
+            prompt_ids, processed_data, _ = self._apply_hf_processor_text_mm(
+                prompt_text=prompt,
+                mm_items=mm_items,
+                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+                tokenization_kwargs=tokenization_kwargs,
+            )
+
+        # For gemma3 we check `token_type_ids` as the key
+        mm_token_type_ids = processed_data.pop("token_type_ids", None)
+        mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids)
+
+        # We can infer vLLM style placeholder from token type ids, if we split
+        # it for each input `mm_data`.
+        mm_positions = torch.where(mm_token_type_ids == 1)[1]
+        images = mm_items.get_items("image", ImageProcessorItems)
+        image_sizes = []
+        for item_idx in range(len(images)):
+            image_size = images.get_image_size(item_idx)
+            image_sizes.append((image_size.height, image_size.width))
+
+        mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens(
+            image_sizes=image_sizes,
+            **self.info.ctx.get_merged_mm_kwargs({}),
+        )
+
+        mm_placeholders = {}
+        split_sizes = mm_tokens_per_modality["num_image_tokens"]
+        if split_sizes:
+            chunked_mm_positions = torch.split(mm_positions, split_sizes)
+            mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()]
+            chunked_mm_tokens = torch.split(mm_tokens, split_sizes)
+            ranges = [
+                PlaceholderRange(
+                    offset=positions[0].item(),
+                    length=positions.shape[0],
+                    is_embed=(mm_tokens == hf_processor.image_token_id).bool(),
+                )
+                for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens)
+            ]
+            mm_placeholders = {"image": ranges}
+
+        processed_data["num_image_patches"] = torch.tensor(
+            mm_tokens_per_modality["num_image_patches"]
+        )
+        mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
+            processed_data,
+            self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs),
+        )
+
+        # Use overrides if provided; fallback to data-dependent hashing.
+        with timing_ctx.record("get_mm_hashes"):
+            mm_hashes = inputs.get_mm_hashes(self.info.model_id)
+
+        return mm_input(
+            prompt_token_ids=prompt_ids,
+            mm_kwargs=mm_kwargs,
+            mm_hashes=mm_hashes,
+            mm_placeholders=mm_placeholders,
+        )
+
+
+class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip SupportsMRoPE.__init__ and call the next class in MRO
+        super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
+
+    def _get_encoder_cls(
+        self, modality: str = "image", **kwargs: dict
+    ) -> type["PreTrainedModel"]:
+        """
+        Get the encoder class from the model.
+
+        Args:
+            kwargs: The kwargs to create the model.
+
+        Returns:
+            The encoder class.
+        """
+        with torch.device("meta"):
+            model: PreTrainedModel = AutoModel.from_config(**kwargs)
+        encoder_cls = type(model.get_encoder(modality=modality))
+        logger.debug("Identified encoder class as: %s", encoder_cls)
+        if type(model) is encoder_cls:
+            raise ValueError(
+                "Unable to infer vision encoder class from the model. "
+                "You must either: update the model so that "
+                "https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
+                " can detect the vision encoder correctly, or remove "
+                "'compile_mm_encoder'."
+            )
+        del model
+        return encoder_cls
+
+    def _decorate_for_torch_compile(self, **kwargs: dict):
+        """
+        Decorate the model's decoder and encoder classes to indicate to vLLM that they
+        support torch compile if `can_enable_torch_compile` and
+        `should_torch_compile_mm_encoder` are True respectively.
+
+        Args:
+            kwargs: The kwargs to create the model, which are needed to get the decoder
+                and encoder classes.
+        """
+        super()._decorate_for_torch_compile(**kwargs)
+        # Decorate the vision encoder model class to support torch compile if needed
+        if self.compilation_config.compile_mm_encoder:
+            self.check_version("5.0.0", "multimodal encoder compilation support")
+            logger.warning_once(
+                "Multimodal encoder compilation with the Transformers modeling backend "
+                "is an experimental feature. It relies on:\n"
+                "- The vision encoder being torch compilable.\n"
+                "- All vision encoder tensor inputs must be type hinted as either "
+                "`torch.Tensor` or `torch.FloatTensor`.\n"
+                "- The 0-th dimension of all tensor inputs to the vision encoder being "
+                "the dynamic dimension (i.e., sequence length or number of patches).\n"
+                "Please report any issues you encounter to help us improve it."
+            )
+            self._decorate_cls_for_torch_compile(
+                cls=self._get_encoder_cls(**kwargs),
+                # TODO: properly infer dynamic_arg_dims based on the encoder's forward
+                # method signature. Currently we assume dim 0 for all tensor inputs.
+                dynamic_arg_dims=None,
+                enable_if=should_torch_compile_mm_encoder,
+                is_encoder=True,
+            )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        # Positions shape handling for MRoPE models
+        if self.model_config.uses_mrope:
+            # [3, seq_len] -> [3, 1, seq_len]
+            positions = positions[:, None]
+        model_output = super().forward(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
+        return model_output
+
+    def get_language_model(self) -> torch.nn.Module:
+        """Transformers modeling backend multimodal classes do not contain a separate
+        vLLM language model class. Therefore, in order to return a language model vLLM
+        class, we use a wrapper to give `self` the same interface as a text model."""
+
+        # Exclude self and object
+        bases = self.__class__.mro()[1:-1]
+        # Keep only classes defined in `vllm.model_executor.models.transformers`
+        bases = [b for b in bases if ".transformers." in b.__module__]
+        # Exclude MultiModalMixin itself
+        bases = [b for b in bases if b is not MultiModalMixin]
+
+        class LanguageModel(*bases):
+            def __init__(self, multimodal_model):
+                # Don't call super().__init__() to avoid re-initialization
+                self.__dict__.update(multimodal_model.__dict__)
+
+            model = getattr_iter(self.model, ("language_model", "text_model"), None)
+
+        return LanguageModel(self)
+
+    def embed_multimodal(self, **kwargs):
+        pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None)
+        image_embeds: torch.Tensor | None = kwargs.pop("image_embeds", None)
+        # Model might use `image_patches` instead of `pixel_values`
+        if pixel_values is None:
+            pixel_values = kwargs.pop("image_patches", None)
+
+        if image_embeds is not None:
+            return image_embeds
+
+        if pixel_values is None:
+            return None
+
+        num_image_patches = kwargs.pop("num_image_patches")
+
+        if pixel_values is not None:
+            # ROCm: Force math SDP backend for vision encoder to avoid accuracy issues
+            # with flash_sdp and mem_efficient_sdp
+            if current_platform.is_rocm():
+                # TODO: [ROCm] Fix accuracy issues with flash backend
+                logger.debug(
+                    "ROCm platform detected. Forcing math SDP backend "
+                    "for vision encoder. Currently ROCm platform has "
+                    "accuracy issues with `flash_sdp` and"
+                    "`mem_efficient_sdp` backends. See issue: "
+                    "https://github.com/vllm-project/vllm/issues/30167"
+                )
+                with torch.nn.attention.sdpa_kernel(
+                    backends=[torch.nn.attention.SDPBackend.MATH]
+                ):
+                    vision_embeddings = self.model.get_image_features(
+                        pixel_values, **kwargs
+                    )
+            else:
+                vision_embeddings = self.model.get_image_features(
+                    pixel_values, **kwargs
+                )
+
+            # Transformers `v5`, `self.get_image_features` returns a tuple
+            # containing the features and optionally attentions/hidden_states
+            # After v5 is settled, we can enable qwen3-vl with several outputs
+            # from `self.get_image_features`
+            if isinstance(vision_embeddings, tuple):
+                vision_embeddings = vision_embeddings[0]
+            elif isinstance(vision_embeddings, dict):
+                vision_embeddings = vision_embeddings.pooler_output
+
+            if isinstance(vision_embeddings, torch.Tensor):
+                split_sizes = num_image_patches.flatten().tolist()
+                total_patches = sum(split_sizes)
+
+                # Flatten to 2D: [total_tokens, hidden_dim]
+                if vision_embeddings.ndim == 3:
+                    vision_embeddings = vision_embeddings.view(
+                        -1, vision_embeddings.shape[-1]
+                    )
+
+                total_tokens = vision_embeddings.shape[0]
+                if total_tokens == total_patches:
+                    # Direct match: num_image_patches are actual token counts
+                    # (e.g., Qwen2.5-VL style)
+                    token_split_sizes = split_sizes
+                elif total_patches > 0 and total_tokens % total_patches == 0:
+                    # Uniform expansion: each patch expands to N tokens
+                    # (e.g., Idefics3 style)
+                    tokens_per_patch = total_tokens // total_patches
+                    token_split_sizes = [s * tokens_per_patch for s in split_sizes]
+                elif total_patches > 0:
+                    # Mismatch (profiling with dummy data) - pad/truncate
+                    if total_tokens == 0:
+                        raise ValueError(
+                            "Vision encoder returned empty embeddings. "
+                            f"Expected {total_patches} patches from "
+                            f"num_image_patches={split_sizes}"
+                        )
+                    if total_tokens < total_patches:
+                        repeat_factor = (
+                            total_patches + total_tokens - 1
+                        ) // total_tokens
+                        vision_embeddings = vision_embeddings.repeat(repeat_factor, 1)
+                    vision_embeddings = vision_embeddings[:total_patches]
+                    token_split_sizes = split_sizes
+                else:
+                    return []
+
+                return list(torch.split(vision_embeddings, token_split_sizes, dim=0))
+
+            return vision_embeddings
+        else:
+            logger.debug(
+                "No pixel values or image embeddings provided for multimodal embedding."
+            )
+            return None
+
+    def get_mrope_input_positions(
+        self,
+        input_tokens: list[int],
+        mm_features: list[MultiModalFeatureSpec],
+    ) -> tuple[torch.Tensor, int]:
+        kwargs = MultiModalFeatureSpec.gather_kwargs(
+            mm_features,
+            {
+                "image_grid_thw",
+                "video_grid_thw",
+                "second_per_grid_ts",
+                "audio_feature_lengths",
+                "use_audio_in_video",
+            },
+        )
+        if any(v for k, v in kwargs.items() if k not in {"image_grid_thw"}):
+            raise NotImplementedError(
+                "Transformers modeling backend only supports images."
+            )
+
+        image_grid_thw = kwargs.get("image_grid_thw", [])
+        video_grid_thw = kwargs.get("video_grid_thw", [])
+
+        image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(
+            image_grid_thw
+        )
+        video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(
+            video_grid_thw
+        )
+
+        # `get_rope_index` doesn't always accept arbitrary `kwargs`
+        kwargs = {}
+        if not hasattr(self, "_get_rope_index_accepts_mm_token_type_ids"):
+            import inspect
+
+            sig = inspect.signature(self.model.get_rope_index)
+            params = sig.parameters
+            self._get_rope_index_accepts_mm_token_type_ids = (
+                "mm_token_type_ids" in params
+                or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+            )
+        if self._get_rope_index_accepts_mm_token_type_ids:
+            mm_token_type_ids = torch.zeros(len(input_tokens), dtype=torch.int)
+            for feature in mm_features:
+                position = feature.mm_position
+                offset, length = position.offset, position.length
+                mm_token_type_id = _MODALITY_TO_TOKEN_TYPE_ID[feature.modality]
+                mm_token_type_ids[offset : offset + length] = mm_token_type_id
+            kwargs["mm_token_type_ids"] = mm_token_type_ids.unsqueeze(0)
+
+        mrope_positions, mrope_position_delta = self.model.get_rope_index(
+            input_ids=torch.tensor(input_tokens).unsqueeze(0),
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            **kwargs,
+        )
+
+        mrope_positions = mrope_positions[:, 0]
+        mrope_position_delta = mrope_position_delta[0].item()
+
+        return mrope_positions, mrope_position_delta

vllm/model_executor/models/transformers/pooling.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend mixins for pooling models."""
+
+from typing import TYPE_CHECKING
+
+import torch
+from transformers import AutoModelForSequenceClassification
+
+from vllm.config.utils import getattr_iter
+from vllm.model_executor.layers.pooler import DispatchPooler
+from vllm.model_executor.models.interfaces import SupportsCrossEncoding
+from vllm.model_executor.models.interfaces_base import VllmModelForPooling
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+
+class EmbeddingMixin(VllmModelForPooling):
+    default_seq_pooling_type = "CLS"
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip VllmModelForPooling.__init__ and call the next class in MRO
+        super(VllmModelForPooling, self).__init__(
+            vllm_config=vllm_config, prefix=prefix
+        )
+
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        self.pooler = DispatchPooler.for_embedding(pooler_config)
+
+
+class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):
+    default_seq_pooling_type = "CLS"
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip VllmModelForPooling.__init__ and call the next class in MRO
+        super(VllmModelForPooling, self).__init__(
+            vllm_config=vllm_config, prefix=prefix
+        )
+
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        # Certain information about the model and classifier can only be
+        # inferred from the `ForSequenceClassification` class. Therefore, we
+        # instantiate it on the "meta" device to avoid allocating GPU memory.
+        with torch.device("meta"):
+            seq_cls_model = AutoModelForSequenceClassification.from_config(
+                self.config,
+                dtype=self.model_config.dtype,
+                trust_remote_code=self.model_config.trust_remote_code,
+            )
+
+        # When used for sequence classification, some models have their
+        # pooling layers removed. Make sure this is reflected in vLLM.
+        for module in seq_cls_model.modules():
+            if hasattr(module, "pooler") and module.pooler is None:
+                self.model.pooler = None
+                break
+
+        # Unlike `lm_head`, `classifier` is not always `nn.Linear`.
+        self.classifier = getattr_iter(seq_cls_model, ["classifier", "score"], None)
+        if self.classifier is None:
+            raise ValueError(
+                "Could not find `classifier` or `score` layer in the "
+                "`AutoModelForSequenceClassification` instance."
+            )
+        self.init_parameters(self.classifier, dtype=self.model_config.head_dtype)
+
+        class ClassifierWithReshape(self.classifier.__class__):
+            """
+            Token extraction has already been applied in `pooler.pooling`.
+            Add dim to match expected input shape of `classifier.forward`.
+            """
+
+            def forward(self, *args, **kwargs):
+                if len(args) > 0:
+                    args = (args[0].unsqueeze(1), *args[1:])
+                return super().forward(*args, **kwargs)
+
+        self.classifier.__class__ = ClassifierWithReshape
+
+        self.pooler = DispatchPooler.for_seq_cls(
+            pooler_config,
+            classifier=self.classifier,
+        )

vllm/model_executor/models/transformers/utils.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformers modeling backend utilities."""
+
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import torch
+from torch import nn
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.transformers_utils.config import is_rope_parameters_nested
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.model_executor.layers.quantization import QuantizationConfig
+
+
+logger = init_logger(__name__)
+
+
+# Copied from `accelerate`
+@contextmanager
+def init_on_device_without_buffers(device: torch.device):
+    """
+    A context manager under which models are initialized with all
+    parameters on the specified device. However buffers are not
+    initialized on specified device.
+
+    Args:
+        device (`torch.device`):
+            Device to initialize all parameters on.
+    """
+
+    old_register_parameter = nn.Module.register_parameter
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            kwargs["requires_grad"] = param.requires_grad
+            module._parameters[name] = param_cls(
+                module._parameters[name].to(device), **kwargs
+            )
+
+    tensor_constructors_to_patch = {}
+
+    def patch_tensor_constructor(fn):
+        def wrapper(*args, **kwargs):
+            kwargs["device"] = device
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        for torch_function_name in tensor_constructors_to_patch:
+            setattr(
+                torch,
+                torch_function_name,
+                patch_tensor_constructor(getattr(torch, torch_function_name)),
+            )
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        for (
+            torch_function_name,
+            old_torch_function,
+        ) in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)
+
+
+Style = Literal[
+    "colwise",
+    "rowwise",
+    "replicate",
+    "colwise_gather_output",
+    "rowwise_split_input",
+]
+
+
+def replace_linear_class(
+    linear: nn.Linear,
+    style: Style = "replicate",
+    quant_config: "QuantizationConfig | None" = None,
+    *,
+    prefix: str = "",
+) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:
+    """
+    Replace nn.Linear with one of vLLM's tensor parallel linear classes.
+
+    Args:
+        linear: `nn.Linear` to be replaced.
+        style: Tensor parallel style of the new linear, e.g. "colwise".
+        quant_config: Quantization config for the new linear.
+    Returns:
+        The new linear.
+    """
+
+    if not isinstance(style, str):
+        raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")
+
+    vllm_linear_cls, vllm_linear_kwargs = {
+        "colwise": (ColumnParallelLinear, {}),
+        "rowwise": (RowParallelLinear, {}),
+        "replicate": (ReplicatedLinear, {}),
+        "colwise_gather_output": (ColumnParallelLinear, {"gather_output": True}),
+        "rowwise_split_input": (RowParallelLinear, {"input_is_parallel": False}),
+    }.get(style, (ReplicatedLinear, {}))
+
+    return vllm_linear_cls(
+        input_size=linear.in_features,
+        output_size=linear.out_features,
+        bias=linear.bias is not None,
+        quant_config=quant_config,
+        prefix=prefix,
+        return_bias=False,
+        **vllm_linear_kwargs,
+    )
+
+
+TorchConv = nn.Conv2d | nn.Conv3d
+VllmConv = Conv2dLayer | Conv3dLayer
+
+
+def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
+    """Replace a Transformers Conv2d/Conv3d with vLLM's Conv2d/Conv3d.
+
+    Args:
+        conv: `nn.Conv2d` or `nn.Conv3d` to be replaced.
+    Returns:
+        The new `Conv2dLayer` or `Conv3dLayer`. If the conv module is not supported,
+        returns the original conv module.
+    """
+    # vLLM does not handle non-zero padding modes
+    if conv.padding_mode != "zeros":
+        return conv
+
+    vllm_conv_cls = {
+        nn.Conv2d: Conv2dLayer,
+        nn.Conv3d: Conv3dLayer,
+    }.get(type(conv))
+
+    if vllm_conv_cls is None:
+        return conv
+
+    return vllm_conv_cls(
+        in_channels=conv.in_channels,
+        out_channels=conv.out_channels,
+        kernel_size=conv.kernel_size,
+        stride=conv.stride,
+        padding=conv.padding,
+        dilation=conv.dilation,
+        groups=conv.groups,
+        bias=conv.bias is not None,
+        padding_mode=conv.padding_mode,
+        params_dtype=conv.weight.dtype,
+    )
+
+
+def recursive_replace_linear(
+    model: nn.Module,
+    quant_config: "QuantizationConfig | None",
+    prefix: str = "",
+):
+    """Recursively replace linear modules in the model as needed."""
+
+    def _recursive_replace(module: nn.Module, prefix: str):
+        for child_name, child_module in module.named_children():
+            new_module = child_module
+            qual_name = maybe_prefix(prefix, child_name)
+            # Replace modules as needed
+            if isinstance(child_module, nn.Linear):
+                style = "replicate"
+                new_module = replace_linear_class(
+                    child_module,
+                    style,
+                    quant_config,
+                    prefix=qual_name,
+                )
+            else:
+                _recursive_replace(child_module, prefix=qual_name)
+            if new_module is not child_module:
+                setattr(module, child_name, new_module)
+
+    _recursive_replace(model, prefix=prefix)
+
+
+def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module):
+    logger.debug("%s: %s -> %s", name, old_module, new_module)
+
+
+def get_feature_request_tip(
+    model: str,
+    trust_remote_code: bool,
+) -> str:
+    hf_url = f"a discussion at https://huggingface.co/{model}/discussions/new"
+    gh_url = "an issue at https://github.com/huggingface/transformers/issues/new/choose"
+    url = hf_url if trust_remote_code else gh_url
+    prefix = f"Please open {url} to request support for this feature. "
+    if Path(model).exists():
+        prefix = ""
+    doc_url = "https://docs.vllm.ai/en/latest/models/supported_models.html#writing-custom-models"
+    tip = f"See {doc_url} for instructions on how to add support yourself."
+    return f"{prefix}{tip}"
+
+
+def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool:
+    """
+    Callable to be passed to `@support_torch_compile`'s `enable_if` argument.
+
+    Defaults to `True` but is disabled in the following situations:
+
+    - The model uses dynamic rope scaling.
+    """
+    text_config = vllm_config.model_config.hf_config.get_text_config()
+    # Dynamic rope scaling is not compatible with torch.compile
+    rope_parameters: dict | None = getattr(text_config, "rope_parameters", None) or {}
+    if rope_parameters:
+        # Nest rope_parameters if not nested already to simplify logic
+        if not is_rope_parameters_nested(rope_parameters):
+            rope_parameters = {"": rope_parameters}
+        return all(rp["rope_type"] != "dynamic" for rp in rope_parameters.values())
+    return True